PELS.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c539a60d-4daf-482e-be9a-8c66f9a7692f",
   "metadata": {},
   "source": [
    "# PELS tutorial hands on"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9c8928-233a-4535-8792-2834608441e2",
   "metadata": {},
   "source": [
    "## Let us understand the system\n",
    "\n",
    "* We query the system topology using LIKWID tool.\n",
    "* Note that only 1 NUMA domain (16 cores) of CPU is available.\n",
    "* Attached is also one A100 GPU.\n",
    "* Watch out for the available cores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "332686c9-5f4b-4899-8f09-59cb57e03a3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!likwid-topology"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d751e94b-5c95-49f4-a7d1-f3dc8bdba371",
   "metadata": {},
   "source": [
    "## CPU main memory bandwidth\n",
    "\n",
    "* Main memory bandwidth is the most important hardware characteristic relevant to SpMV like kernels.\n",
    "* Let us measure the bandwidth using LIKWID tool.\n",
    "* Note that bandwidth of only 1 NUMA domain (16 cores) is measured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30ae5e2a-1b7d-44f4-b3f9-1481df7911e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!likwid-bench -t load_avx -w N:4GB:16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a80131f-d9b4-48d2-ad69-238fdda54dc4",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Let us set up our environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e57f9e1-5a48-4a8d-ba3b-013f37f330b8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "#On cluster we might not have internet to download packages\n",
    "#so we set path to pre-installed Python packages\n",
    "sys.path.insert(0, \"/home/hpc/unrz/unrz002h/installkit/lib/python3.9/site-packages\") \n",
    "\n",
    "import os\n",
    "os.environ[\"USE_RACE\"]=\"1\" #Compile RACE too\n",
    "os.environ[\"OMP_SCHEDULE\"]=\"static\"\n",
    "os.environ[\"OMP_NUM_THREADS\"]=\"16\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f01389c-5a1e-4f7c-a640-d97cac942fcf",
   "metadata": {},
   "source": [
    "## SpMV benchmarking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fbc259d-338f-4c66-9c9d-98a1f6b80b77",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kernels import *\n",
    "import numpy as np\n",
    "import numba\n",
    "import gc\n",
    "from numpy.linalg import norm\n",
    "from scipy.sparse import *\n",
    "from scipy.io import mmread\n",
    "from sellcs import sellcs_matrix\n",
    "from matrix_generator import create_matrix\n",
    "import time\n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "\n",
    "def spmv_bench(matrix, mat_fmt='CRS', arrow=False):\n",
    "    ## **Note:** The Python garbage collector (gc)\n",
    "    ## can kill the performance of the C kernels\n",
    "    ## for some obscure reason (possibly a conflict\n",
    "    ## between Numba/LLVM and other compilers like GCC).\n",
    "    ## For the pure Python/numba/cuda kernels, this is not\n",
    "    ## the case, but if you are facing obvious performnace\n",
    "    ## problems with the C kernels, you may want to disalbe\n",
    "    ## garbage collection:\n",
    "    gc.disable()\n",
    "\n",
    "    rand_seed=1\n",
    "    if rand_seed is not None:\n",
    "        np.random.seed(rand_seed)\n",
    "\n",
    "    A = create_matrix(matrix, imbal=arrow)\n",
    "    #plt.spy(A)\n",
    "    N = A.shape[0]\n",
    "    #print \"nrows = %d, nnz = %d, nnzr = %.2f\" % (N, A.nnz, A.nnz/N)\n",
    "    print(\"nnz = \", A.nnz, \"nrows =\", N, \"nnzr =\", round(A.nnz/N,2))\n",
    "    if 'SELL' in mat_fmt:\n",
    "        C_=int(mat_fmt.split('-')[1])\n",
    "        if C_ > 256:\n",
    "            print('C greater than 256. Setting to maximum possible value 256')\n",
    "            C_ = 256\n",
    "        sigma_=int(mat_fmt.split('-')[2])\n",
    "        print('Matrix format: SELL-%d-%d'%(C_,sigma_))\n",
    "        A = sellcs_matrix(A_csr=A, C=C_, sigma=sigma_)\n",
    "    else:\n",
    "        print('Matrix format: CSR')\n",
    "    \n",
    "    x=np.random.rand(N)\n",
    "    y=np.random.rand(N)\n",
    "\n",
    "    if available_gpus()>0:\n",
    "        type = 'gpu'\n",
    "    else:\n",
    "        type = 'cpu'\n",
    "\n",
    "    #Transfer to device in case of GPU\n",
    "    if type=='gpu':\n",
    "        y = to_device(y)\n",
    "        A = to_device(A)\n",
    "        x = to_device(x)\n",
    "        \n",
    "    #warm-up, determine iterations for 1 sec\n",
    "    iter=10\n",
    "    start = perf_counter()\n",
    "    for it in range(1, iter+1):\n",
    "        spmv(A,x,y)\n",
    "    end = perf_counter()\n",
    "    iter=int(iter/(end-start))\n",
    "    iter=max(1, iter)\n",
    "\n",
    "    reset_counters()\n",
    "    start = perf_counter()\n",
    "    #Run SpMV for iter times\n",
    "    for it in range(1, iter+1):\n",
    "        spmv(A,x,y)\n",
    "    end = perf_counter()\n",
    "    \n",
    "    print(\"Performance GFlop/s = \", 2*iter*A.nnz*1e-9/(end-start))\n",
    "    perf_report(type)\n",
    "\n",
    "    gc.enable()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaf58a3f-023e-4f73-8bd9-7725702ceede",
   "metadata": {},
   "source": [
    "## SpMV on CPU\n",
    "Let's run SpMV on CPU.\n",
    "* The matrix used is a simple 3D Laplacian of dimension 240x240x240.\n",
    "* Observe the achieved performance.\n",
    "* Benchmark also outputs bandwidth estimate based on minimum memory transfer assumption. \n",
    "* Does it match roofline predictions?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d8eb670-8411-435f-882a-c38451c26e95",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"USE_CPU\"]=\"True\"\n",
    "spmv_bench('Laplace240x240x240')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f99581e-a0b7-4d3e-baa0-4055ce46c5dd",
   "metadata": {},
   "source": [
    "# Sparse matrix data formats on CPU\n",
    "Does changing from CRS to SELL-C-sigma data storage format change performance on CPU? Why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f1b567c-a5b9-46b0-99a4-96a8b59db79a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"USE_CPU\"]=\"True\"\n",
    "spmv_bench('Laplace240x240x240', mat_fmt='SELL-32-1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "408158d2-5776-43dd-be03-36e5c0b1c83a",
   "metadata": {
    "tags": []
   },
   "source": [
    "## SpMV  on GPU\n",
    "* Observe the performance of SpMV on GPU with CSR format (which is the default)\n",
    "* Do we hit roofline limit?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2edec253-54cd-4c3a-8d77-5f4194468f82",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"USE_CPU\"]=\"False\"\n",
    "spmv_bench('Laplace240x240x240')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "218ff5e9-61be-48c3-a9f6-0e5150a3ab0b",
   "metadata": {},
   "source": [
    "## Sparse matrix data formats on GPU\n",
    "* Does SELL-C-sigma improve GPU performance?\n",
    "* Try increasing the chunk size (C) parameter (maximum possible value is 256) and observe what hapens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "decd43af-4a32-4a46-8834-88130bae4478",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"USE_CPU\"]=\"False\"\n",
    "spmv_bench('Laplace240x240x240', mat_fmt='SELL-8-1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee7f382e-63a2-4ffe-907c-873a0fba8655",
   "metadata": {},
   "source": [
    "## Is SELL-C-sigma good for all types of matrix?\n",
    "* Try using SELL-C-sigma on a modified Laplacian matrix with an arrow!\n",
    "* Is the performance similar to a plain Laplacian matrix? Why?\n",
    "* Hint: Observe the sparsity pattern of the two matrices. See next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2763094-386b-49a6-9a50-e5a7073aa13b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.subplots(1,2, figsize=(14,6))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "A = create_matrix('Laplace20x20x20', imbal=False)\n",
    "plt.spy(A)\n",
    "plt.subplot(1, 2, 2)\n",
    "A_arrow = create_matrix('Laplace20x20x20', imbal=True)\n",
    "plt.spy(A_arrow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f80b9df-7b84-4515-a319-0dc63f218cf2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"USE_CPU\"]=\"False\"\n",
    "spmv_bench('Laplace240x240x240', mat_fmt='SELL-8-1', arrow=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eae552b1-be64-447c-b9a3-a66c36c46b35",
   "metadata": {},
   "source": [
    "# PCG solver\n",
    "* The PCG solver is implemented using the cg_solve function. Have a look.\n",
    "* Notice the kernels used for implementing PCG solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19c5b6c5-5c17-4110-9ea1-439e90e51654",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kernels import *\n",
    "import numpy as np\n",
    "import numba\n",
    "import gc\n",
    "from numpy.linalg import norm\n",
    "from scipy.sparse import *\n",
    "from scipy.io import mmread\n",
    "from sellcs import sellcs_matrix\n",
    "from poly_op import *\n",
    "\n",
    "\n",
    "from matrix_generator import create_matrix\n",
    "\n",
    "def cg_solve(A, b, x0, tol, maxit, verbose=False, x_ex=None):\n",
    "    '''\n",
    "    x, tol, iter = cg_solve(A, b, x0, tol, maxit)\n",
    "    Where A is an spd scipy.sparse.csr_matrix, b and x0 are numpy.array's of size A.shpae[0],\n",
    "    tol is the convergence tolerance and maxit the maximum number of iterations.\n",
    "    '''\n",
    "    x = clone(b)\n",
    "    r = clone(b)\n",
    "    p = clone(b)\n",
    "    q = clone(b)\n",
    "\n",
    "    err_norms=[]\n",
    "    res_norms=[]\n",
    "    \n",
    "    if x_ex is not None:\n",
    "        print('PerfWarning: providing the exact solution x_ex results in additional operations to calculate and print the error norm.')\n",
    "        err = clone(b)\n",
    "\n",
    "    tol2 = tol*tol\n",
    "\n",
    "\n",
    "    axpby(1.0,x0,0.0,x)\n",
    "\n",
    "    #r = A*x\n",
    "    if hasattr(A, 'apply'):\n",
    "        A.apply(x, r)\n",
    "    else:\n",
    "        spmv(A, x, r)\n",
    "    #r = b - r\n",
    "    axpby(1.0, b, -1.0, r)\n",
    "    #p = r\n",
    "    axpby(1.0, r, 0.0, p)\n",
    "\n",
    "    # rho = <r, r>\n",
    "    rho = dot(r,r);\n",
    "    rho_old = 1.0\n",
    "    if verbose:\n",
    "        print('%d\\t%e'%(0, np.sqrt(rho)))\n",
    "\n",
    "\n",
    "    for iter in range(maxit+1):\n",
    "\n",
    "        # check stop criteria\n",
    "        if rho < tol2:\n",
    "            break;\n",
    "\n",
    "        # q = A*p\n",
    "        if hasattr(A, 'apply'):\n",
    "            A.apply(p, q)\n",
    "        else:\n",
    "            spmv(A, p, q)\n",
    "\n",
    "        pq = dot(p,q)\n",
    "        alpha = rho / pq\n",
    "        # x = x+alpha*p\n",
    "        axpby(alpha, p, 1.0, x)\n",
    "        # r = r - alpha*q\n",
    "        axpby(-alpha, q, 1.0, r)\n",
    "\n",
    "        rho_old = rho\n",
    "        rho = dot(r, r)\n",
    "\n",
    "        # Calculate absolute error -----------------------------------\n",
    "        curr_res_norm=np.sqrt(rho)\n",
    "        res_norms.append(curr_res_norm)\n",
    "        if x_ex is not None:\n",
    "            if hasattr(A, 'unprec_sol'):\n",
    "                A.unprec_sol(x, err)\n",
    "            else:\n",
    "                axpby(1.0, x, 0.0, err)\n",
    "            axpby(-1.0, x_ex, 1.0, err)\n",
    "            curr_err_norm=np.sqrt(dot(err, err))\n",
    "            err_norms.append(curr_err_norm)    \n",
    "            if verbose:\n",
    "                print('%d\\t%e\\t%e'%(iter+1, curr_res_norm, curr_err_norm))\n",
    "        else:   \n",
    "            if verbose:\n",
    "                print('%d\\t%e'%(iter+1, curr_res_norm))\n",
    "        # ------------------------------------------------------------\n",
    "        beta = rho / rho_old\n",
    "        # p = r+beta*p\n",
    "        axpby (1.0, r, beta, p)\n",
    "\n",
    "    iter_count = iter\n",
    "    final_residual = np.sqrt(rho)\n",
    "\n",
    "    return x, final_residual, iter_count, res_norms, err_norms\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def bench(matrix, maxit, mat_fmt='CSR', poly_k=0, tol=1e-3, use_RACE=False, cache_size=60, rand_seed=1, print_perf=False):\n",
    "    ## **Note:** The Python garbage collector (gc)\n",
    "    ## can kill the performance of the C kernels\n",
    "    ## for some obscure reason (possibly a conflict\n",
    "    ## between Numba/LLVM and other compilers like GCC).\n",
    "    ## For the pure Python/numba/cuda kernels, this is not\n",
    "    ## the case, but if you are facing obvious performnace\n",
    "    ## problems with the C kernels, you may want to disalbe\n",
    "    ## garbage collection:\n",
    "    gc.disable()\n",
    "\n",
    "    if rand_seed is not None:\n",
    "        np.random.seed(rand_seed)\n",
    "\n",
    "    #if args.matfile != 'None':\n",
    "    #    if args.matgen!='None':\n",
    "    #        print('got both -matfile and -matgen, the latter will be ignored.')\n",
    "    #    if not args.matfile.endswith(('.mm','.mtx','.mm.gz','.mtx.gz')):\n",
    "    #        raise(ValueError('Expecting MatrixMarket file with extension .mm[.gz] or .mtx[.gz]'))\n",
    "    #    A = csr_matrix(mmread(args.matfile))\n",
    "    A = create_matrix(matrix)\n",
    "    N = A.shape[0]\n",
    "\n",
    "    x_ex=np.random.rand(N)\n",
    "\n",
    "    b=A*x_ex\n",
    "\n",
    "    x0 = np.zeros(N,dtype='float64')\n",
    "\n",
    "    print('norm of rhs: %e'%(norm(b)))\n",
    "    print('rel. residual of given solution: %e'%(norm(A*x_ex-b)/norm(b)))\n",
    "\n",
    "    sigma=1\n",
    "\n",
    "    A_csr = A # we may need it for creating the preconditioner\n",
    "              # in case the user wants a SELL-C-sigma matrix.\n",
    "\n",
    "    print(\"nnz = \", A.nnz, \"nrows =\", N, \"nnzr =\", A.nnz/N)\n",
    "    if 'SELL' in mat_fmt:\n",
    "        C_=int(mat_fmt.split('-')[1])\n",
    "        if C_ > 256:\n",
    "            print('C greater than 256. Setting to maximum possible value 256')\n",
    "            C_ = 256\n",
    "        sigma_=int(mat_fmt.split('-')[2])\n",
    "        print('Matrix format: SELL-%d-%d'%(C_,sigma_))\n",
    "        A = sellcs_matrix(A_csr=A_csr, C=C_, sigma=sigma_)\n",
    "        b = b[A.permute]\n",
    "        A_csr = A_csr[A.permute[:,None], A.permute]\n",
    "    else:\n",
    "        print('Matrix format: CSR')\n",
    "\n",
    "    if available_gpus()>0:\n",
    "        type = 'gpu'\n",
    "    else:\n",
    "        type = 'cpu'\n",
    "\n",
    "    print('Will run on '+type)\n",
    "    \n",
    "    numa=True\n",
    "    \n",
    "    if type=='gpu':\n",
    "        x0 = to_device(x0)\n",
    "        b  = to_device(b)\n",
    "        A  = to_device(A)\n",
    "        x_ex = to_device(x_ex)\n",
    "\n",
    "    #Just run all the kernels one time. Just for catching any errors\n",
    "    compile_all() \n",
    "\n",
    "    A_prec = A\n",
    "    b_prec = b\n",
    "\n",
    "    # we want to make sure what we measure during CG in total\n",
    "    # is consistent with the sum of the kernel calls and their\n",
    "    # runtime as predicted by the roofline model, so reset all\n",
    "    # counters and timers:\n",
    "    reset_counters()\n",
    "\n",
    "    t0 = perf_counter()\n",
    "\n",
    "    t0_pre = perf_counter()\n",
    "    if poly_k>0:\n",
    "        # building preconditioners typically requires a certain format,\n",
    "        # in our case, the poly_op class uses scipy functions tril and triu,\n",
    "        # which are not implemented by the sellcs_matrix class.\n",
    "        print(\"calling poly_op \", poly_k)\n",
    "        A_prec = poly_op(A_csr, poly_k, use_RACE, cache_size)\n",
    "        if A_prec.mpkHandle != None:\n",
    "            b = b[A_prec.permute]\n",
    "            #A_csr = A_csr[A_prec.permute[:,None], A_prec.permute]\n",
    "            A_csr = permute_csr(A_csr, A_prec.permute)\n",
    "        if 'SELL' in mat_fmt:\n",
    "            # note: If A was originally sorted by row-length (sigma>1), use the same\n",
    "            # sorting for L and U to avoid intermittent permutation by setting sigma=1.\n",
    "            # There still seems to be some kind of bug, though, because the number of\n",
    "            # iterations will increase with poly_k>0 and sigma>1. Hence this warning.\n",
    "            C_=int(mat_fmt.split('-')[1])\n",
    "            sigma_=int(mat_fmt.split('-')[2])\n",
    "            A_prec.L = to_device(sellcs_matrix(A_csr=A_prec.L, C=C_, sigma=1))\n",
    "            A_prec.U = to_device(sellcs_matrix(A_csr=A_prec.U, C=C_, sigma=1))\n",
    "        b_prec = copy(b)\n",
    "        A_prec.prec_rhs(b, b_prec)\n",
    "\n",
    "    \n",
    "    #hw_string = type\n",
    "    #if type=='cpu':\n",
    "    #    hw_string+=' ('+str(numba.threading_layer())+', '+str(numba.get_num_threads())+' threads)'\n",
    "    #print('Hardware: '+hw_string)\n",
    "    \n",
    "    printerr=True\n",
    "    x_ex_in = None\n",
    "    if printerr:\n",
    "        x_ex_in = x_ex  \n",
    "        if poly_k>0:\n",
    "            if A_prec.mpkHandle != None:\n",
    "                x_ex_in = x_ex_in[A_prec.permute]\n",
    "\n",
    "    t1_pre = perf_counter()\n",
    "\n",
    "    t0_soln = perf_counter()\n",
    "    x_prec, relres, iter, res_norms, err_norms = cg_solve(A_prec,b_prec,x0,tol,maxit, x_ex=x_ex_in)\n",
    "\n",
    "    t1_soln = perf_counter()\n",
    "\n",
    "    if poly_k>0:\n",
    "        x = clone(x_prec)\n",
    "        A_prec.unprec_sol(x_prec, x)\n",
    "    else:\n",
    "        x = x_prec\n",
    "\n",
    "    t1 = perf_counter()\n",
    "    t_pre = t1_pre-t0_pre\n",
    "    t_soln = t1_soln-t0_soln\n",
    "    t_CG = t1-t0\n",
    "    gc.enable()\n",
    "\n",
    "    x = to_host(x)\n",
    "\n",
    "    rel_err_norms=err_norms/norm(x_ex)\n",
    "    rel_res_norms=res_norms/norm(b)\n",
    "    \n",
    "    print('number of CG iterations: %d'%(iter))\n",
    "    res = np.empty_like(x)\n",
    "    spmv(A_csr,x,res)\n",
    "    res=b-res\n",
    "    print('Residual of computed solution: %e'%(norm(res)))\n",
    "\n",
    "    if 'SELL' in mat_fmt and sigma>1:\n",
    "        x = x[A.unpermute]\n",
    "\n",
    "    if poly_k>0:\n",
    "        if A_prec.mpkHandle != None:\n",
    "            x = x[A_prec.unpermute]\n",
    "\n",
    "    print('Error of computed solution: %e'%(norm(x-x_ex)))\n",
    "\n",
    "    if print_perf:\n",
    "        perf_report(type)\n",
    "    if poly_k>0:\n",
    "        print('Total time for constructing precon: %g seconds.'%(t_pre))\n",
    "        print('Total time for solving: %g seconds.'%(t_soln))\n",
    "    print('Total time for CG: %g seconds.'%(t_CG))\n",
    "    \n",
    "    return rel_res_norms, rel_err_norms, iter, t_pre, t_soln, t_CG\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13c4ad05-e5d2-4f7d-b7fa-13bfa77cf789",
   "metadata": {},
   "source": [
    "## Let us run a simple example on CPU\n",
    "* CG solver is run for 50 iterations.\n",
    "* At the end of the solver, the performance of each kernel used in CG is reported.\n",
    "    * What is the most costly component?\n",
    "    * Does the performance of SpMV match with the measured bandwidth?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19312b3d-41c5-4e2f-b069-46cbd077f53d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"USE_CPU\"]=\"True\"\n",
    "rel_res_in_CG, rel_err, iter, t_pre, t_soln, t_CG = bench('Laplace240x240x240', 50, print_perf=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e3ad6dd-bd5f-4568-86d5-8f12ffad8d0d",
   "metadata": {},
   "source": [
    "# Polynomial preconditioner\n",
    "\n",
    "Given a matrix $A$ and an integer $k$, this operator constructs the splitting (cf. Jacobi iteration): $$D^{-1/2}AD^{-1/2} = A_d = I - (L+U)$$\n",
    "\n",
    "A simple polynomial preconditioner can then be defined as\n",
    "        $$w = M_1 A_d M_2 v,$$ where $M_1$ and $M_2$ approximately solve\n",
    "        $(I-L)v=w$ and $(I-U)v=w$, respectively, using the truncated Neumann series: \n",
    "        $$v \\approx \\sum_{i=0}^k L^{i} w.$$\n",
    "        \n",
    "==> The 'apply' function implements a preconditioned matrix-vector product\n",
    "        $$\\sum_{i=0}^kL^i A_d \\sum_{j=0}^kU^j v = w$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4d3c596-d784-41e8-b2f3-758be2ad3538",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Preconditioner and its effects\n",
    "* Now we will experiment with the polynomial preconditioner.\n",
    "* Let us stick with the CPU from now on.\n",
    "* We will run the solver till the residual reaches 1e-8. Note: The entire run may take five minutes.\n",
    "* Observe what happens as the degree of the polynomial increases.\n",
    "* How does it effect the CG iterations? How about runtime?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe999517-0fc0-4423-adcf-bcff4dbd7d8f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "degrees=[]\n",
    "iters=[]\n",
    "times=[]\n",
    "plt.subplots(1,3, figsize=(14,6))\n",
    "\n",
    "os.environ[\"USE_CPU\"]=\"True\"\n",
    "plt.subplot(1, 3, 1)\n",
    "\n",
    "for deg in [0,1,2]:\n",
    "    print(\"#################### \", \"Running with polynomial degree \", deg, \" ####################\")\n",
    "    rel_res_in_CG, rel_err, iter, t_pre, t_soln, t_CG = bench('Laplace240x240x240', 10000, poly_k=deg, tol=1e-8)\n",
    "    degrees.append(deg)\n",
    "    iters.append(iter)\n",
    "    times.append(t_CG)\n",
    "    plt.plot(rel_err, label=\"degree=\"+str(deg))\n",
    "\n",
    "    \n",
    "plt.xlabel(\"Iterations\")\n",
    "plt.ylabel(\"Relative error\")\n",
    "plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.xlabel(\"Polynomial degree\")\n",
    "plt.ylabel(\"Iterations to convergence\")\n",
    "plt.plot(degrees, iters)\n",
    "\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.xlabel(\"Polynomial degree\")\n",
    "plt.ylabel(\"Time to convergence\")\n",
    "plt.plot(degrees, times)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "435ed3b2-4a17-43ad-9029-38f7d88ba1f1",
   "metadata": {},
   "source": [
    "# Cache blocking using RACE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e203ef-d311-499a-9551-e58729ef9278",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kernels import *\n",
    "import numpy as np\n",
    "import numba\n",
    "import gc\n",
    "from numpy.linalg import norm\n",
    "from scipy.sparse import *\n",
    "from scipy.io import mmread\n",
    "from sellcs import sellcs_matrix\n",
    "from matrix_generator import create_matrix\n",
    "import time\n",
    "\n",
    "\n",
    "def mpk_bench(matrix, power=1, cacheSize=60):\n",
    "    ## **Note:** The Python garbage collector (gc)\n",
    "    ## can kill the performance of the C kernels\n",
    "    ## for some obscure reason (possibly a conflict\n",
    "    ## between Numba/LLVM and other compilers like GCC).\n",
    "    ## For the pure Python/numba/cuda kernels, this is not\n",
    "    ## the case, but if you are facing obvious performnace\n",
    "    ## problems with the C kernels, you may want to disalbe\n",
    "    ## garbage collection:\n",
    "    gc.disable()\n",
    "\n",
    "    rand_seed=1\n",
    "    if rand_seed is not None:\n",
    "        np.random.seed(rand_seed)\n",
    "\n",
    "    A = create_matrix(matrix)\n",
    "    N = A.shape[0]\n",
    "    print(\"nnz = \", A.nnz, \"nrows =\", N, \"nnzr =\", A.nnz/N)\n",
    "    \n",
    "    x_in=np.random.rand(N)\n",
    "    x=np.copy(x_in)\n",
    "    x_RACE=np.copy(x_in)\n",
    "    \n",
    "    y=np.zeros(N)\n",
    "    y_RACE=np.zeros(N)\n",
    "        \n",
    "    #Set up MPK\n",
    "    mpk_handle=mpk_setup(A, power, cacheSize, False)\n",
    "    perm=mpk_get_perm(mpk_handle, N)\n",
    "    #Permute input vector according to RACE permutation\n",
    "    x_RACE=x_RACE[perm]\n",
    "    \n",
    "    #warm-up, determine iterations for 1 sec\n",
    "    iter=10\n",
    "    start = perf_counter()\n",
    "    for it in range(1, iter+1):\n",
    "        for p in range(1,power+1):\n",
    "            spmv(A,x,y)\n",
    "            #interchange x and y\n",
    "            tmp = y\n",
    "            y = x\n",
    "            x = tmp\n",
    "    end = perf_counter()\n",
    "    iter=int(iter/(end-start))\n",
    "    iter=max(1, iter)\n",
    "\n",
    "    reset_counters()\n",
    "    start = perf_counter()\n",
    "    #Run SpMV-based MPK for iter times\n",
    "    for it in range(1, iter+1):\n",
    "        for p in range(1,power+1):\n",
    "            spmv(A,x,y)\n",
    "            #interchange x and y\n",
    "            tmp = y\n",
    "            y = x\n",
    "            x = tmp\n",
    "    end = perf_counter()\n",
    "    \n",
    "    print(\"Performance SpMV-based MPK = \", 2*iter*power*A.nnz*1e-9/(end-start), \" GFlop/s\")\n",
    "    \n",
    "    reset_counters()\n",
    "    start = perf_counter()\n",
    "    #Run RACE-based MPK for iter times\n",
    "    for it in range(1, iter+1):\n",
    "        mpk(mpk_handle,power,x_RACE,y_RACE)\n",
    "    end = perf_counter()\n",
    "    \n",
    "    print(\"Performance RACE-based MPK = \", 2*iter*power*A.nnz*1e-9/(end-start), \" GFlop/s\")\n",
    "\n",
    "    ## Validation\n",
    "    #Rest all vectors\n",
    "    y=np.zeros(N)\n",
    "    x=np.copy(x_in)\n",
    "    x_RACE=np.copy(x_in)\n",
    "    x_RACE=x_RACE[perm]\n",
    "    y_RACE=np.zeros(N)\n",
    "    \n",
    "    #Run SpMV-based once\n",
    "    for p in range(1,power+1):\n",
    "        spmv(A,x,y)\n",
    "        #interchange x and y\n",
    "        tmp=y\n",
    "        y=x\n",
    "        x=tmp\n",
    "    y=x\n",
    "\n",
    "    #Run RACE-based once\n",
    "    mpk(mpk_handle, power, x_RACE, y_RACE)\n",
    "\n",
    "    y=y[perm] #permute y to RACE permute space for comparison\n",
    "    err=y_RACE-y\n",
    "    err_norm = np.dot(err,err)\n",
    "    if (err_norm<1e-10):\n",
    "        print(\"Validation success\")\n",
    "    else:\n",
    "        print(\"Err norm = \", err_norm)\n",
    "    \n",
    "    mpk_free(mpk_handle)\n",
    "    gc.enable()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33d43d5f-4f42-45a4-ba7c-9b9ea0d78ae2",
   "metadata": {},
   "source": [
    "## Performance of RACE MPK\n",
    "* Observe the performance sppedup obtained by using RACE's cache blocking\n",
    "* Try changing the 'power' and 'cacheSize' input parameters to RACE. How does the performance change?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f82f6aa3-1142-4861-973a-161b70358c64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mpk_bench('Laplace240x240x240', power=1, cacheSize=60)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85c0baf5-cdba-4cfb-8124-b368d788c169",
   "metadata": {
    "tags": []
   },
   "source": [
    "# RACE and polynomial preconditioner\n",
    "\n",
    "==> The 'apply' function implements a preconditioned matrix-vector product\n",
    "        $$\\sum_{i=0}^kL^i A_d \\sum_{j=0}^kU^j v = w$$\n",
    "\n",
    "* With RACE we can reuse $L$ and $U$ across $k$ iteration from cache. **Can we do better?**\n",
    "* Let us experiment with the efficacy of this method.\n",
    "* Try playing with the following knobs:\n",
    "    * use_RACE: True or False to switch on or off RACE.\n",
    "    * poly_k: The degree of polynomial.\n",
    "    * cache_size: The cache size (in Mbyte) for which RACE will perform cache blocking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "146cad43-cd1f-45a1-ab58-289306ab065b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"USE_CPU\"]=\"True\"\n",
    "rel_res_in_CG, rel_err, iter, t_pre, t_soln, t_CG = bench('Laplace240x240x240', 50, poly_k=2, tol=1e-3, use_RACE=True, cache_size=60, print_perf=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7832cce9-5919-4106-bf2c-841fbf3eac01",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Does RACE improve the time to solution for polynomial preconditioner?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "581de5ba-331f-42ec-8df9-3674b7b93039",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt \n",
    "degrees=[]\n",
    "iters=[]\n",
    "times=[]\n",
    "plt.subplots(1,3, figsize=(14,6))\n",
    "\n",
    "os.environ[\"USE_CPU\"]=\"True\"\n",
    "plt.subplot(1, 3, 1)\n",
    "for deg in [0,1,2]:\n",
    "    print(\"#################### \", \"Running with polynomial degree \", deg, \" ####################\")\n",
    "    rel_res_in_CG, rel_err, iter, t_pre, t_soln, t_CG = bench('Laplace240x240x240', 10000, poly_k=deg, tol=1e-8, use_RACE=True, cache_size=60)\n",
    "    degrees.append(deg)\n",
    "    iters.append(iter)\n",
    "    times.append(t_CG)\n",
    "    plt.plot(rel_err, label=\"degree=\"+str(deg))\n",
    "\n",
    "    \n",
    "plt.xlabel(\"Iterations\")\n",
    "plt.ylabel(\"Relative error\")\n",
    "plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.xlabel(\"Polynomial degree\")\n",
    "plt.ylabel(\"Iterations to convergence\")\n",
    "plt.plot(degrees, iters)\n",
    "\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.xlabel(\"Polynomial degree\")\n",
    "plt.ylabel(\"Time to convergence\")\n",
    "plt.plot(degrees, times)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed855d1-f09b-4890-9320-aaa9ce61de99",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}