sparsity_benchmarks/SpMM.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sparse dense matmul with PopSparse on IPU\n",
    "\n",
    "<small>Copyright (c) 2023 Graphcore Ltd.</small>\n",
    "\n",
    "This tutorial will walk you through how to use PopSparse and Poptorch for sparse dense matmul (SpMM).\n",
    "The detailed benchmark results and discussion can be found in the paper:\n",
    "\n",
    "[PopSparse: Accelerated block sparse matrix multiplication on IPU](https://arxiv.org/abs/2303.16999)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The SpMM can be written as:\n",
    "\\begin{align*}\n",
    "    Y &= (M \\odot W) * X \\,,\\\\\n",
    "    M_{ij} &= \\hat{M}_{\\lfloor i/b \\rfloor, \\lfloor j/b \\rfloor} \\,,\n",
    "\\end{align*}\n",
    "where $\\odot$ denotes elementwise multiplication, $*$ for inner product, $Y \\in \\mathbb{R}^{m \\times n}$, $X \\in \\mathbb{R}^{k \\times n}$ are dense output and input matrices respectively. The sparse weight matrix $(M \\odot W)$ is defined via $M \\in \\mathbb{B}^{m \\times k}$ ($\\mathbb{B} = \\{0, 1\\}$), a mask that represents the *sparsity pattern*, itself derived from $\\hat{M} \\in \\mathbb{B}^{\\lceil m/b \\rceil \\times \\lceil k/b \\rceil}$, a block mask and $W \\in \\mathbb{R}^{m \\times k}$ defines weight values.\n",
    "\n",
    "In this formulation, $(M \\odot W)$ has a block-sparse structure, where contiguous square blocks of weights of shape $b \\times b$ contain all the non-zero elements, given *block-size* $b$. The dimensions $m$ and $k$ are referred to as output and input *feature size* and  $n$ is referred to as *batch size*, corresponding to their typical role in weight-sparse neural network computation.\n",
    "\n",
    "We define *density*, $d = \\sum_{ij} M_{ij} / (m \\cdot k)$, where the standard term *sparsity* corresponds to $(1 - d)$. We use floating point operation (FLOP) count to quantify the amount of useful arithmetic work in an SpMM as: $2\\cdot m\\cdot k\\cdot n\\cdot d$. Note that this only counts operations performed on non-zero values and does not depend on block size $b$."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we demonstrate the SpMM with PopSparse on IPU with static (main focus) and dynamic sparsity; elementwise (block size = 1) and block sparse. We compare the performance of sparse operations with its equivalent dense calculation. We start with an introduction using small matrices and show the comparison using large matrices at the end.\n",
    "\n",
    "**Note** that the throughput numbers reported in the paper is run on IPU Bow machine. If you are running the notebook on a classic machine, the performance will be $\\sim$ 30% less. You can use `gc-monitor` on the command line to check the IPU clock frequency when there is a matmul running on IPU. The Bow IPU should be 1850MHz and classic is 1325MHz."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing modules and functions"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install modules in requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -q -r requirements.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing modules and PopTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import poptorch\n",
    "import time\n",
    "import tabulate\n",
    "import matplotlib.pyplot as pyplot\n",
    "from typing import Callable"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build the SpMM custom ops (static and dynamic) based on PopSparse libray so that we could use them in the PopTorch models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "make: Entering directory '/nethome/zhiyil/git/sw_apps_sandbox/sparsity_demo/poptorch_static_sparse'\n",
      "make: 'libpoptorch_static_sparse_op.so' is up to date.\n",
      "make: Leaving directory '/nethome/zhiyil/git/sw_apps_sandbox/sparsity_demo/poptorch_static_sparse'\n",
      "make: Entering directory '/nethome/zhiyil/git/sw_apps_sandbox/sparsity_demo/poptorch_dynamic_sparse'\n",
      "make: 'build/libpoptorch_dynamic_sparse_op.so' is up to date.\n",
      "make: Leaving directory '/nethome/zhiyil/git/sw_apps_sandbox/sparsity_demo/poptorch_dynamic_sparse'\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "make -C poptorch_static_sparse\n",
    "make -C poptorch_dynamic_sparse\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing the SpMM custom ops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import poptorch_static_sparse.poptorch_static_sparse_op as static_sparse_op\n",
    "import poptorch_dynamic_sparse.poptorch_dynamic_sparse_op as dynamic_sparse_op\n",
    "from poptorch_static_sparse.poptorch_static_sparse_op import magnitude_prune "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example with small matrices\n",
    "Let's generate two random matrices and prune the left hand side matrix to be a sparse matrix:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the hyperparameters of the sparse and dense matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 1 # block size of the sparse matrix b\n",
    "shape = [5, 5] # block row and column of the sparse matrix mb * kb (m = mb * b, k = kb * b)\n",
    "batch_size = 4 # batch size of the dense matrix n\n",
    "density = 0.5 # density = (1 - sparsity) of the sparse matrix, d\n",
    "dtype = torch.float\n",
    "num_loops = 2000 # loop count in the poptprch for_loop to run for timing\n",
    "num_exe = 2 # number of execution, fixed to be 2: first time for warming up (compilation etc.), second time for timing\n",
    "seed = 1999\n",
    "torch.manual_seed(seed);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We generate the left hand side (lhs) and right hand side (rhs) of the matrix multiplication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_lhs_rhs_matrix(shape, block_size, batch_size, dtype):\n",
    "    lhs = torch.randn((shape[0] * block_size, shape[1] * block_size), dtype=dtype)\n",
    "    lhs /= lhs.shape[1] ** 0.5 \n",
    "    rhs = torch.randn((shape[1] * block_size, batch_size), dtype=dtype)\n",
    "    return lhs, rhs\n",
    "lhs, dense = generate_lhs_rhs_matrix(shape, block_size, batch_size, dtype)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We prune the lhs matrix and get our sparse matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "sparse = magnitude_prune(lhs.to(dtype), block_size, density)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could have a look at the original left hand side matrix in dense form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1257,  0.1839,  0.1565,  0.5638,  0.1385],\n",
       "        [-0.6040, -0.0577,  0.0542, -0.1277, -0.1356],\n",
       "        [-0.1255, -0.4487,  0.2926, -0.4057, -0.5863],\n",
       "        [ 0.4124,  0.1301, -0.2613, -0.1388, -0.4322],\n",
       "        [-0.0492,  0.6213, -0.6681,  0.3652, -0.3116]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lhs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pruned lhs, which shows in sparse format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(indices=tensor([[0, 1, 2, 2, 2, 2, 3, 3, 4, 4, 4, 4],\n",
       "                       [3, 0, 1, 2, 3, 4, 0, 4, 1, 2, 3, 4]]),\n",
       "       values=tensor([[[ 0.5638]],\n",
       "\n",
       "                      [[-0.6040]],\n",
       "\n",
       "                      [[-0.4487]],\n",
       "\n",
       "                      [[ 0.2926]],\n",
       "\n",
       "                      [[-0.4057]],\n",
       "\n",
       "                      [[-0.5863]],\n",
       "\n",
       "                      [[ 0.4124]],\n",
       "\n",
       "                      [[-0.4322]],\n",
       "\n",
       "                      [[ 0.6213]],\n",
       "\n",
       "                      [[-0.6681]],\n",
       "\n",
       "                      [[ 0.3652]],\n",
       "\n",
       "                      [[-0.3116]]]),\n",
       "       size=(5, 5, 1, 1), nnz=12, layout=torch.sparse_coo)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sparse"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The sparse matrix in dense format (we can spot the zeros after pruning):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000,  0.0000,  0.0000,  0.5638,  0.0000],\n",
       "        [-0.6040,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000, -0.4487,  0.2926, -0.4057, -0.5863],\n",
       "        [ 0.4124,  0.0000,  0.0000,  0.0000, -0.4322],\n",
       "        [ 0.0000,  0.6213, -0.6681,  0.3652, -0.3116]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sparse_dense_matrix = static_sparse_op.block_coo_to_dense(sparse)\n",
    "sparse_dense_matrix"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The right hand side dense matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.3168, -0.7362, -2.2414, -0.6063],\n",
       "        [ 0.2821, -0.8316,  0.2504,  1.9486],\n",
       "        [-0.0178,  0.5444, -0.5707, -1.3101],\n",
       "        [-0.6667,  0.5962, -0.4849,  1.1735],\n",
       "        [ 0.0036, -0.3892, -0.6853, -0.4562]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dense"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create PyTorch models for the sparse(static and dynamic) SpMM and dense matmul."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DemoSparseModel(torch.nn.Module):\n",
    "    \"\"\"A tiny demo model for SpMM with PopSparse.\"\"\"\n",
    "\n",
    "    def __init__(self, sparse: torch.Tensor, custom_op: Callable, loop_count: int):\n",
    "        super().__init__()\n",
    "        self.sparse = sparse\n",
    "        self.loop_count = loop_count\n",
    "        self.custom_op = custom_op\n",
    "\n",
    "    def forward(self, input: torch.Tensor)  -> torch.Tensor:\n",
    "        if self.loop_count > 1:\n",
    "            return poptorch.for_loop(\n",
    "                self.loop_count, lambda dense: self.custom_op(self.sparse, dense), [input]\n",
    "            )[0]\n",
    "        return self.custom_op(self.sparse, input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DemoDenseModel(torch.nn.Module):\n",
    "    \"\"\"A tiny demo model for dense matmul.\"\"\"\n",
    "\n",
    "    def __init__(self, sparse_dense_matrix: torch.Tensor, loop_count: int):\n",
    "        super().__init__()\n",
    "        self.sparse_dense_matrix = sparse_dense_matrix\n",
    "        self.loop_count = loop_count\n",
    "\n",
    "    def forward(self, input: torch.Tensor)  -> torch.Tensor:\n",
    "        if self.loop_count > 1:\n",
    "            return poptorch.for_loop(\n",
    "                self.loop_count, lambda dense: [torch.matmul(self.sparse_dense_matrix.to(device=\"ipu\"), dense)], [input]\n",
    "            )[0]\n",
    "        return torch.matmul(self.sparse_dense_matrix.to(device=\"ipu\"), input)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run the SpMM using the torch gather scatter method (gather_scatter) on CPU and PopSparse SpMM custom ops static_sparse_op and dynamic_sparse_op (inference_model) on IPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = {}\n",
    "durations = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'duration_gs': 0.00024115920066833496}\n"
     ]
    }
   ],
   "source": [
    "if dtype != torch.half:\n",
    "    for i in range(num_loops+num_exe):\n",
    "        if i == num_exe:\n",
    "            start_gs = time.time()\n",
    "        outputs[\"gather_scatter\"] = static_sparse_op.block_coo_spmm_gs(\n",
    "            sparse, dense, mode=\"sparse_dense\"\n",
    "        )\n",
    "    durations[\"duration_gs\"] = (time.time() - start_gs) / num_loops\n",
    "print(durations)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create the static_sparse_model with the custom op `static_sparse_op.block_coo_spmm_ipu` and convert the model to be an PopTorch inference model using `poptorch.inferenceModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:03<00:00]\n"
     ]
    }
   ],
   "source": [
    "static_sparse_model = DemoSparseModel(sparse, static_sparse_op.block_coo_spmm_ipu, loop_count=1)\n",
    "options = poptorch.Options()\n",
    "options.Precision.setPartialsType(torch.half)\n",
    "ipu_model_static_sparse = poptorch.inferenceModel(static_sparse_model, options)\n",
    "\n",
    "for i in range(num_exe):\n",
    "    if i == 0:\n",
    "        out = ipu_model_static_sparse(dense) # run for graph compilation\n",
    "    elif i == num_exe - 1:\n",
    "        start_spmm = time.time()\n",
    "        outputs[\"inference_model_static\"] = ipu_model_static_sparse(dense) # run without graph compilation\n",
    "durations[\"duration_spmm_static\"] = (time.time() - start_spmm) / num_loops    "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same with dynamic SpMM, we create the dynamic_sparse_model with the custom op `dynamic_sparse_op.dynamic_spmm_ipu` and convert it to PopTorch inference model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:05<00:00]\n"
     ]
    }
   ],
   "source": [
    "dynamic_sparse_model = DemoSparseModel(sparse, dynamic_sparse_op.dynamic_spmm_ipu, loop_count=1)\n",
    "options = poptorch.Options()\n",
    "options.Precision.setPartialsType(torch.half)\n",
    "ipu_model_dynamic_sparse = poptorch.inferenceModel(dynamic_sparse_model, options)\n",
    "\n",
    "for i in range(num_exe):\n",
    "    if i == 0:\n",
    "        out = ipu_model_dynamic_sparse(dense) # run once for graph compilation\n",
    "    elif i == num_exe - 1:\n",
    "        start_spmm = time.time()\n",
    "        outputs[\"inference_model_dynamic\"] = ipu_model_dynamic_sparse(dense) # run num_loops without graph compilation\n",
    "durations[\"duration_spmm_dynamic\"] = (time.time() - start_spmm) / num_loops\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also run the dense one as a reference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:03<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'duration_gs': 0.00024115920066833496, 'duration_spmm_static': 4.029273986816406e-07, 'duration_spmm_dynamic': 2.7751922607421873e-07, 'duration_dense': 2.4974346160888674e-07}\n"
     ]
    }
   ],
   "source": [
    "sparse_dense_matrix = static_sparse_op.block_coo_to_dense(sparse)\n",
    "dense_model = DemoDenseModel(sparse_dense_matrix, loop_count=1)\n",
    "options = poptorch.Options()\n",
    "ipu_dense_model = poptorch.inferenceModel(dense_model, options)\n",
    "\n",
    "for i in range(num_exe):\n",
    "    if i == 0:\n",
    "        out = ipu_dense_model(dense) # run for graph compilation\n",
    "    elif i == num_exe - 1:\n",
    "        start_spmm = time.time()\n",
    "        outputs[\"inference_dense\"] = ipu_dense_model(dense) # run without graph compilation\n",
    "durations[\"duration_dense\"] = (time.time() - start_spmm) / num_loops\n",
    "print(durations)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could examine the output from the gather scatter method and from the PopTorch inference models to make sure they give equal output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3758,  0.3361, -0.2734,  0.6615],\n",
       "        [-0.1913,  0.4446,  1.3538,  0.3662],\n",
       "        [ 0.1365,  0.5188,  0.3191, -1.4664],\n",
       "        [ 0.1291, -0.1354, -0.6283, -0.0529],\n",
       "        [-0.0575, -0.5414,  0.5733,  2.6567]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# output from CPU gather scatter\n",
    "outputs[\"gather_scatter\"] # shape seems to be different with inference model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3758,  0.3361, -0.2734,  0.6615],\n",
       "        [-0.1913,  0.4446,  1.3538,  0.3662],\n",
       "        [ 0.1365,  0.5188,  0.3191, -1.4664],\n",
       "        [ 0.1291, -0.1354, -0.6283, -0.0529],\n",
       "        [-0.0575, -0.5414,  0.5733,  2.6567]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# inference model output from static sparse\n",
    "outputs[\"inference_model_static\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3758,  0.3361, -0.2734,  0.6615],\n",
       "        [-0.1913,  0.4446,  1.3538,  0.3662],\n",
       "        [ 0.1365,  0.5188,  0.3191, -1.4664],\n",
       "        [ 0.1291, -0.1354, -0.6283, -0.0529],\n",
       "        [-0.0575, -0.5414,  0.5733,  2.6567]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# inference model output from dynamic sparse\n",
    "outputs[\"inference_model_dynamic\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3758,  0.3361, -0.2734,  0.6615],\n",
       "        [-0.1913,  0.4446,  1.3538,  0.3662],\n",
       "        [ 0.1365,  0.5188,  0.3191, -1.4664],\n",
       "        [ 0.1291, -0.1354, -0.6283, -0.0529],\n",
       "        [-0.0575, -0.5414,  0.5733,  2.6567]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# inference model output from dense\n",
    "outputs[\"inference_dense\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could also comapre the ouputs with the expected dense matmul calculated on CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.3758,  0.3361, -0.2734,  0.6615],\n",
      "        [-0.1913,  0.4446,  1.3538,  0.3662],\n",
      "        [ 0.1365,  0.5188,  0.3191, -1.4664],\n",
      "        [ 0.1291, -0.1354, -0.6283, -0.0529],\n",
      "        [-0.0575, -0.5414,  0.5733,  2.6567]])\n"
     ]
    }
   ],
   "source": [
    "expected_output = static_sparse_op.block_coo_to_dense(sparse).float() @ dense.float()\n",
    "print(expected_output)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize pruned matrix:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To have more visualization on the pruning process, we could use a slightly larger matrix and plot the matrix before and after pruning:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 4\n",
    "shape = [3, 5]\n",
    "batch_size = 4\n",
    "density = 0.5\n",
    "dtype = torch.float"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "lhs, dense = generate_lhs_rhs_matrix(shape, block_size, batch_size, dtype) # generate lhs and rhs\n",
    "sparse = magnitude_prune(lhs.to(dtype), block_size, density) # prune lhs\n",
    "sparse_dense = static_sparse_op.block_coo_to_dense(sparse) # convert lhs back to dense form"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could see that with the pruned sparse matrix, almost half (density 0.5) of the 4 x 4 blocks are zeros. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f435cc041f0>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA9EAAAEpCAYAAACUbqEuAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy88F64QAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAkbUlEQVR4nO3de3SdZZ0v8F+atkkpacqtaWOvVK5tqVw7BR0VuigdhPboEXA6DiADLiwqMjNCz7JUQI2oi8URmcJ4hOJSbs4awJERBioFlZZCWxRQawu1BKEtoE16IUlJ3vMHk4yhSfsGnt29k3w+a+21yM6b737ePnvnx3fvZKcsy7IsAAAAgD0aUOwFAAAAQG+hRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOA4u9gLdra2uLl19+OaqqqqKsrKzYywGAyLIstm7dGrW1tTFggOefUzDvASglPZn1JVeiX3755RgzZkyxlwEAu6ivr4/Ro0cXexl9gnkPQCnKM+tLrkRXVVVFRMSUsxZE+aDKZLmN49I/yz30fa8nz2x64oDkmRER1S+0Js98c5/0r8a8Nn1n8sx4M/06BzSlz9xv4p+TZ+5YXpj7U8UJf0qeuf03+yXPbK3MkmeOWJE+88//a0fyzPMOeyJ5ZkTE/3v45OSZgxrSP54qtqTNa21pit8tvrpjRvHutf9bjlnwpRhQmW7e93cHPZX+e1R/nvfl2/3kSWoDWtL/P7l5n15/nPc9mfUlV6Lbf6SrfFBllA9ON1TLK9M/YMv3qUifWVGY/5EYOCh9ic4GF6CcDilPnlmQEl2AtxPoTfenQqy1EP8TnRVgqA4clD6zfJ+25JmV+xbm23sh9qm8AE9KlQ9OHhkR4ceOE2r/txxQWalEJ1SI71H9et63KtGpDRiQ/vuoeZ9ef573eWa97wwAAACQkxINAAAAOSnRAAAAkFPBSvSNN94Y48ePj8rKypg2bVqsWLGiUDcFABSBWQ9Af1SQEn3XXXfFZZddFgsXLoxVq1bF1KlTY+bMmbF58+ZC3BwAsJeZ9QD0VwUp0dddd11ceOGFcf7558eRRx4ZN910U+yzzz5xyy23FOLmAIC9zKwHoL9KXqJbWlpi5cqVMWPGjP+5kQEDYsaMGbFs2bJdjm9ubo7GxsZOFwCgdPV01keY9wD0HclL9GuvvRatra1RU1PT6fqamprYuHHjLsfX1dVFdXV1x2XMmDGplwQAJNTTWR9h3gPQdxT93bnnz58fDQ0NHZf6+vpiLwkASMy8B6CvGJg68MADD4zy8vLYtGlTp+s3bdoUI0eO3OX4ioqKqKioSL0MAKBAejrrI8x7APqO5K9EDx48OI499thYsmRJx3VtbW2xZMmSmD59euqbAwD2MrMegP4s+SvRERGXXXZZnHvuuXHcccfFCSecENdff31s3749zj///ELcHACwl5n1APRXBSnRZ599drz66qtx5ZVXxsaNG+N973tfPPDAA7u8AQkA0DuZ9QD0VwUp0RERl1xySVxyySWFigcAisysB6A/Kvq7cwMAAEBvoUQDAABATko0AAAA5FSw34l+t6Z/+qmo2HdQsrx7fnFCsqx2La9XJc+8eO6DyTMjIn52Qvo3enn+yqnJMw87+JXkmX9sqE6e+ca69JlD/yV95o73JY+MiIh9b0m/1iGDs+SZO+ZuSZ4ZR7Ykjzxo0YHJM79z8qnJMyMiPvKhlckz7//t5OSZ2bohSfNam5LGQcEMu/+Z5Jn9et5vG5Y8s7876GnzPiXzPt2878ms90o0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE4Di72A7qyuOzoGDqpMlrfPkemfLxjy9ODkmY+MOyx5ZkTEq387Lnlm26AseeaaP4xKnjn9sOeTZy5vTHffbPfSKenvT1UvJI+MiIiX35/+8VT787bkmaePey555gM3vD955o4pZckzx/9Hc/LMiIinHj82eWbFIenvT2+Mb0ma1/ZG2jwolFf/dmryzP4871c8Pyx5Zn9X3mzep2Tep5vPPZn1XokGAACAnJRoAAAAyEmJBgAAgJyUaAAAAMhJiQYAAICclGgAAADIKXmJrquri+OPPz6qqqpixIgRMWfOnFizZk3qmwEAisSsB6A/S16iH3300Zg3b14sX748Hnroodi5c2eceuqpsX379tQ3BQAUgVkPQH82MHXgAw880OnjxYsXx4gRI2LlypXx13/916lvDgDYy8x6APqz5CX67RoaGiIiYv/99+/y883NzdHc3NzxcWNjY6GXBAAktKdZH2HeA9B3FPSNxdra2uLSSy+Nk046KSZPntzlMXV1dVFdXd1xGTNmTCGXBAAklGfWR5j3APQdBS3R8+bNi2effTbuvPPObo+ZP39+NDQ0dFzq6+sLuSQAIKE8sz7CvAeg7yjYj3Nfcskl8ZOf/CQee+yxGD16dLfHVVRUREVFRaGWAQAUSN5ZH2HeA9B3JC/RWZbFZz/72bjnnnti6dKlMWHChNQ3AQAUkVkPQH+WvETPmzcvbr/99rjvvvuiqqoqNm7cGBER1dXVMWTIkNQ3BwDsZWY9AP1Z8t+JXrRoUTQ0NMSHPvShGDVqVMflrrvuSn1TAEARmPUA9GcF+XFuAKDvMusB6M8K+u7cAAAA0Jco0QAAAJCTEg0AAAA5FezvRL9b295THuWDy5PlNR3UliyrXUt1WfLMUQPSrzMiYtDHNifPbHvhwOSZNT9Lf5fccnn6zLHHp7tvtqs/NXlktBboT7IetDJ95pD7ViTP/OHM6ckzx5y1MXnmftdUJ89s3adA397Tf9uLpprW5JkHLBuUNK+1pTVeSpoIhWHeJ3Z5+sj+zrxPy7xPN+97Muu9Eg0AAAA5KdEAAACQkxINAAAAOSnRAAAAkJMSDQAAADkp0QAAAJCTEg0AAAA5KdEAAACQkxINAAAAOSnRAAAAkJMSDQAAADkp0QAAAJCTEg0AAAA5KdEAAACQkxINAAAAOSnRAAAAkJMSDQAAADkp0QAAAJCTEg0AAAA5KdEAAACQkxINAAAAOQ0s9gK601ZeFmUDy5Ll7ftiuqx2J35yVfLMZ/80KnlmRETzfSOSZ37xcz9OnnnT72Ynz3x91nuTZ07+zDPJM3fcNCV55msnNyXPjIh4o6YieeafjpyePLNm/KvJM5vuGJk8c9NntyfPHH1zljwzIqLqhfRrfW3qvskzG2fsSJrXtqMp4vtJI6EgzPu0Pvj+9PP+mYLM++bkmRER5a+kn/fPf+uvkmeWp/2WHxERLz9Tkz50TvrI3mRAc/o+tuWIdP+/09aUP8sr0QAAAJCTEg0AAAA5KdEAAACQkxINAAAAOSnRAAAAkJMSDQAAADkVvER//etfj7Kysrj00ksLfVMAQBGY9QD0JwUt0U8++WTcfPPNcdRRRxXyZgCAIjHrAehvClait23bFnPnzo3vfve7sd9++xXqZgCAIjHrAeiPClai582bF6effnrMmDFjt8c1NzdHY2NjpwsAUPryzvoI8x6AvmNgIULvvPPOWLVqVTz55JN7PLauri6uuuqqQiwDACiQnsz6CPMegL4j+SvR9fX18fnPfz5++MMfRmVl5R6Pnz9/fjQ0NHRc6uvrUy8JAEiop7M+wrwHoO9I/kr0ypUrY/PmzXHMMcd0XNfa2hqPPfZYfOc734nm5uYoLy/v+FxFRUVUVFSkXgYAUCA9nfUR5j0AfUfyEn3KKafEM8880+m6888/Pw4//PC4/PLLdxmqAEDvYtYD0J8lL9FVVVUxefLkTtcNHTo0DjjggF2uBwB6H7MegP6soH8nGgAAAPqSgrw799stXbp0b9wMAFAkZj0A/YVXogEAACAnJRoAAAByUqIBAAAgp73yO9HvxLZj34gB+2TJ8gavG5Isq91PV01JnnnQ44XZkpHnbUieecu1ZybPLPvon5Jnbn98/+SZjzw5KXnmfoOTR0a2ozD3p31fLEueuc+rbckzd7wyIn3mxHTfl9rt3JT++9Mfzky/RxERA5rS31Frnky/99X/8nrSvDfbmmN90kQojJFnm/cpmffpZ8nWCennKOxtXokGAACAnJRoAAAAyEmJBgAAgJyUaAAAAMhJiQYAAICclGgAAADISYkGAACAnJRoAAAAyEmJBgAAgJyUaAAAAMhJiQYAAICclGgAAADISYkGAACAnJRoAAAAyEmJBgAAgJyUaAAAAMhJiQYAAICclGgAAADISYkGAACAnJRoAAAAyGlgsRfQnYP+qyIGDqpIltcyNEuW1e7oU9clz1y+89DkmRERVa3pt/q9F/0ueebtEx5JnnnIby9Onjn8ufTPP+2/pil55qAd6R5Df+mkLyxPnvkf/zUteeaw59M/7qfNeC555vJHJiXPLG8qS54ZETHl1DXJM//4vurkmWs+dmDSvLYdTREXJI2Egmg275My79PP+/seOSF5JuxtXokGAACAnJRoAAAAyEmJBgAAgJyUaAAAAMhJiQYAAICclGgAAADIqSAl+o9//GP83d/9XRxwwAExZMiQmDJlSjz11FOFuCkAoAjMegD6q+R/TPDPf/5znHTSSfHhD384fvrTn8ZBBx0Ua9eujf322y/1TQEARWDWA9CfJS/R1157bYwZMyZuvfXWjusmTJiQ+mYAgCIx6wHoz5L/OPePf/zjOO644+LjH/94jBgxIo4++uj47ne/2+3xzc3N0djY2OkCAJSuns76CPMegL4jeYl+4YUXYtGiRXHIIYfEgw8+GBdffHF87nOfi9tuu63L4+vq6qK6urrjMmbMmNRLAgAS6umsjzDvAeg7kpfotra2OOaYY+JrX/taHH300XHRRRfFhRdeGDfddFOXx8+fPz8aGho6LvX19amXBAAk1NNZH2HeA9B3JC/Ro0aNiiOPPLLTdUcccUS8+OKLXR5fUVERw4YN63QBAEpXT2d9hHkPQN+RvESfdNJJsWbNmk7X/f73v49x48alvikAoAjMegD6s+Ql+gtf+EIsX748vva1r8W6devi9ttvj3/913+NefPmpb4pAKAIzHoA+rPkJfr444+Pe+65J+64446YPHlyXHPNNXH99dfH3LlzU98UAFAEZj0A/VnyvxMdEfGRj3wkPvKRjxQiGgAoAWY9AP1V8leiAQAAoK9SogEAACAnJRoAAAByKsjvRKfw6vERAyrT5Q3eUpYu7L/97vURyTNj2M70mRFR/+R7kme+sP+byTMn//iI5Jn7bcySZ75+dPrMHbUJ7/D/7bgZv02eGRHxRN3xyTP3qU3/GG08ZXvyzGU/n5Q8c+KPGpNnbppenTwzImLlk4ckz9x3Q/rnc7NxbWnzmpLGQcGY92mZ9+nnffxV+kjY27wSDQAAADkp0QAAAJCTEg0AAAA5KdEAAACQkxINAAAAOSnRAAAAkJMSDQAAADkp0QAAAJCTEg0AAAA5KdEAAACQkxINAAAAOSnRAAAAkJMSDQAAADkp0QAAAJCTEg0AAAA5KdEAAACQkxINAAAAOSnRAAAAkJMSDQAAADkp0QAAAJDTwGIvoDvDf1MW5YPLkuW9Pr0lWVa7ptf3TZ5Z8WJF8syIiFNPfzJ55gMPH5c8c8ekpuSZ7zl1c/LMN/9tTPLMnUPT3d/bPf3jI5NnRkRkh6TPPGzW2uSZv1o1MXnmIce9mDxzfcvY5JkHf+3p5JkREVuvPip55ptDkkfGe496KWnem9ubI/3OQ3ptg7LkmeVby5NnNh3UVoDM5JFRvSb9bO7v8/6oE3rJvD+qPnnm+sfTz/vWyvSPefbMK9EAAACQkxINAAAAOSnRAAAAkJMSDQAAADkp0QAAAJCTEg0AAAA5JS/Rra2tsWDBgpgwYUIMGTIkJk6cGNdcc01kmbdfB4C+wKwHoD9L/neir7322li0aFHcdtttMWnSpHjqqafi/PPPj+rq6vjc5z6X+uYAgL3MrAegP0teoh9//PGYPXt2nH766RERMX78+LjjjjtixYoVqW8KACgCsx6A/iz5j3OfeOKJsWTJkvj9738fERG/+tWv4he/+EXMmjWry+Obm5ujsbGx0wUAKF09nfUR5j0AfUfyV6KvuOKKaGxsjMMPPzzKy8ujtbU1vvrVr8bcuXO7PL6uri6uuuqq1MsAAAqkp7M+wrwHoO9I/kr03XffHT/84Q/j9ttvj1WrVsVtt90W3/rWt+K2227r8vj58+dHQ0NDx6W+vj71kgCAhHo66yPMewD6juSvRP/zP/9zXHHFFXHOOedERMSUKVNiw4YNUVdXF+eee+4ux1dUVERFRUXqZQAABdLTWR9h3gPQdyR/JXrHjh0xYEDn2PLy8mhra0t9UwBAEZj1APRnyV+JPuOMM+KrX/1qjB07NiZNmhSrV6+O6667Lj71qU+lvikAoAjMegD6s+Ql+oYbbogFCxbEZz7zmdi8eXPU1tbGpz/96bjyyitT3xQAUARmPQD9WfISXVVVFddff31cf/31qaMBgBJg1gPQnyX/nWgAAADoq5RoAAAAyEmJBgAAgJyS/050KlsPjhhQmS6vrDxLF/bf9ls2OHnmn6e0Js+MiHhg7ZHJM8f95I3kmRvmpd+nlrby5JnNw8uSZ7bsl/7chxy6JXlmRETFj4cnz/zVUxOTZx60Knlk7H/8juSZ21em/7NAv7tucvLMiIiBjenv++XNySPjD4+PSZrX1tSUNA/oHcz74ckzzfu0Xjkp/X2UPfNKNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkNLDYC+jOm0PaYsCQtmR5g/dpSZbV7k/HpX8OouyNwjyvUTYgS575wv8ekjyz4tn05z9kxqvJM8f928bkmZXf25o8c8vVY5NnRkTUf+qN5JltjYOTZ276YPr7ffzLYckjt08swPeSltbkmRERA5rTZ26bmH6tR056MWnezu0t8ULSRKA3MO/N+5QKMe8j0vUl8vNKNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5NTjEv3YY4/FGWecEbW1tVFWVhb33ntvp89nWRZXXnlljBo1KoYMGRIzZsyItWvXplovAFBgZj0AdK/HJXr79u0xderUuPHGG7v8/De+8Y349re/HTfddFM88cQTMXTo0Jg5c2Y0NTW968UCAIVn1gNA93r8d6JnzZoVs2bN6vJzWZbF9ddfH1/60pdi9uzZERHx/e9/P2pqauLee++Nc845592tFgAoOLMeALqX9Hei169fHxs3bowZM2Z0XFddXR3Tpk2LZcuWpbwpAKAIzHoA+rsevxK9Oxs3boyIiJqamk7X19TUdHzu7Zqbm6O5ubnj48bGxpRLAgASeiezPsK8B6DvKPq7c9fV1UV1dXXHZcyYMcVeEgCQmHkPQF+RtESPHDkyIiI2bdrU6fpNmzZ1fO7t5s+fHw0NDR2X+vr6lEsCABJ6J7M+wrwHoO9IWqInTJgQI0eOjCVLlnRc19jYGE888URMnz69y6+pqKiIYcOGdboAAKXpncz6CPMegL6jx78TvW3btli3bl3Hx+vXr4+nn3469t9//xg7dmxceuml8ZWvfCUOOeSQmDBhQixYsCBqa2tjzpw5KdcNABSIWQ8A3etxiX7qqafiwx/+cMfHl112WUREnHvuubF48eL44he/GNu3b4+LLrootmzZEu9///vjgQceiMrKynSrBgAKxqwHgO71uER/6EMfiizLuv18WVlZXH311XH11Ve/q4UBAMVh1gNA94r+7twAAADQWyjRAAAAkJMSDQAAADkp0QAAAJBTj99YbG8Z+2BrDBzYmixv89FVybLafeCjv0qe+cL/OTx5ZkTECx9P/46pNU8mj4xXT3sjeeaap8cmz5w4oil55q9XjEieecyX1ybPjIj4w9PvTZ45dcofkmeeM3JF8sz/+8jZyTOr/5Due127ij8X5jnSwy/4bfLMlf91ZPLMl+8anzSvtSX9Yx4ofW+OSP/3zM37PyTP7M/z/o2RZckz2TOvRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOSjQAAADkpEQDAABATko0AAAA5KREAwAAQE5KNAAAAOSkRAMAAEBOA4u9gLfLsiwiIt58sylpbmtzW9K8iIiWbS3JM1Ofd7u2N8qTZ7a2pH8Opm1Hc/LMaMqSRxZin9oKsPU7t6e/j0ZEtL2RfrGFWOuOba3JM1t3pj/3bGf6+2ghHp8Rhdmn1qb0/6atLWn/TVtb3lpj+4zi3Wv/t2wrwP5DKua9eZ9SIeZ9W1NZ8sz+qn0e5Zn1ZVmJ/R/BSy+9FGPGjCn2MgBgF/X19TF69OhiL6NPMO8BKEV5Zn3Jlei2trZ4+eWXo6qqKsrKdv/MSmNjY4wZMybq6+tj2LBhe2mFheWcSl9fO58I59RbOKfiybIstm7dGrW1tTFggN+ESiHvvO8t95GecE69g3MqfX3tfCKcUzH1ZNaX3I9zDxgwoMfP8g8bNqykN+SdcE6lr6+dT4Rz6i2cU3FUV1cXewl9Sk/nfW+4j/SUc+odnFPp62vnE+GciiXvrPd0OgAAAOSkRAMAAEBOvbpEV1RUxMKFC6OioqLYS0nGOZW+vnY+Ec6pt3BO9Ed98T7inHoH51T6+tr5RDin3qLk3lgMAAAASlWvfiUaAAAA9iYlGgAAAHJSogEAACAnJRoAAAByKvkSfeONN8b48eOjsrIypk2bFitWrNjt8T/60Y/i8MMPj8rKypgyZUr853/+515a6Z7V1dXF8ccfH1VVVTFixIiYM2dOrFmzZrdfs3jx4igrK+t0qays3Esr3rMvf/nLu6zv8MMP3+3XlPIeRUSMHz9+l3MqKyuLefPmdXl8qe3RY489FmeccUbU1tZGWVlZ3HvvvZ0+n2VZXHnllTFq1KgYMmRIzJgxI9auXbvH3J4+FlPa3Tnt3LkzLr/88pgyZUoMHTo0amtr4+///u/j5Zdf3m3mO7nvprSnfTrvvPN2Wd9pp522x9xS3aeI6PJxVVZWFt/85je7zSz2PrF3mPWlNUfezqwvzT0y70t/3pv1fXfWl3SJvuuuu+Kyyy6LhQsXxqpVq2Lq1Kkxc+bM2Lx5c5fHP/744/GJT3wiLrjggli9enXMmTMn5syZE88+++xeXnnXHn300Zg3b14sX748Hnroodi5c2eceuqpsX379t1+3bBhw+KVV17puGzYsGEvrTifSZMmdVrfL37xi26PLfU9ioh48sknO53PQw89FBERH//4x7v9mlLao+3bt8fUqVPjxhtv7PLz3/jGN+Lb3/523HTTTfHEE0/E0KFDY+bMmdHU1NRtZk8fi6nt7px27NgRq1atigULFsSqVavi3//932PNmjVx5pln7jG3J/fd1Pa0TxERp512Wqf13XHHHbvNLOV9iohO5/LKK6/ELbfcEmVlZfGxj31st7nF3CcKz6x/SynNka6Y9aW3R+Z96c97s74Pz/qshJ1wwgnZvHnzOj5ubW3Namtrs7q6ui6PP+uss7LTTz+903XTpk3LPv3pTxd0ne/U5s2bs4jIHn300W6PufXWW7Pq6uq9t6geWrhwYTZ16tTcx/e2PcqyLPv85z+fTZw4MWtra+vy86W8RxGR3XPPPR0ft7W1ZSNHjsy++c1vdly3ZcuWrKKiIrvjjju6zenpY7GQ3n5OXVmxYkUWEdmGDRu6Paan991C6uqczj333Gz27Nk9yult+zR79uzs5JNP3u0xpbRPFIZZX9pzJMvM+iwr/T0y70t/3pv13SuVPeqJkn0luqWlJVauXBkzZszouG7AgAExY8aMWLZsWZdfs2zZsk7HR0TMnDmz2+OLraGhISIi9t9//90et23bthg3blyMGTMmZs+eHc8999zeWF5ua9eujdra2jj44INj7ty58eKLL3Z7bG/bo5aWlvjBD34Qn/rUp6KsrKzb40p9j9qtX78+Nm7c2GkPqqurY9q0ad3uwTt5LBZbQ0NDlJWVxfDhw3d7XE/uu8WwdOnSGDFiRBx22GFx8cUXx+uvv97tsb1tnzZt2hT3339/XHDBBXs8ttT3iXfOrP8fpT5HzPrS36O/ZN53VspzxKx/SynvUVdKtkS/9tpr0draGjU1NZ2ur6mpiY0bN3b5NRs3buzR8cXU1tYWl156aZx00kkxefLkbo877LDD4pZbbon77rsvfvCDH0RbW1uceOKJ8dJLL+3F1XZv2rRpsXjx4njggQdi0aJFsX79+vjABz4QW7du7fL43rRHERH33ntvbNmyJc4777xujyn1PfpL7f/OPdmDd/JYLKampqa4/PLL4xOf+EQMGzas2+N6et/d20477bT4/ve/H0uWLIlrr702Hn300Zg1a1a0trZ2eXxv26fbbrstqqqq4qMf/ehujyv1feLdMevfUupzxKwv/T16O/P+f5TyHDHr31LKe9SdgcVeQH81b968ePbZZ/f48/7Tp0+P6dOnd3x84oknxhFHHBE333xzXHPNNYVe5h7NmjWr47+POuqomDZtWowbNy7uvvvuXM86lbrvfe97MWvWrKitre32mFLfo/5k586dcdZZZ0WWZbFo0aLdHlvq991zzjmn47+nTJkSRx11VEycODGWLl0ap5xyShFXlsYtt9wSc+fO3eMb85T6PsHumPW9g1nf+/SVeW/Wv6WU96g7JftK9IEHHhjl5eWxadOmTtdv2rQpRo4c2eXXjBw5skfHF8sll1wSP/nJT+KRRx6J0aNH9+hrBw0aFEcffXSsW7euQKt7d4YPHx6HHnpot+vrLXsUEbFhw4Z4+OGH4x/+4R969HWlvEft/8492YN38lgshvaBumHDhnjooYd2+6x0V/Z03y22gw8+OA488MBu19db9iki4uc//3msWbOmx4+tiNLfJ3rGrO9aKc+RCLM+ovT3yLzvXinPEbP+LaW8R+1KtkQPHjw4jj322FiyZEnHdW1tbbFkyZJOzwT+penTp3c6PiLioYce6vb4vS3LsrjkkkvinnvuiZ/97GcxYcKEHme0trbGM888E6NGjSrACt+9bdu2xfPPP9/t+kp9j/7SrbfeGiNGjIjTTz+9R19Xyns0YcKEGDlyZKc9aGxsjCeeeKLbPXgnj8W9rX2grl27Nh5++OE44IADepyxp/tusb300kvx+uuvd7u+3rBP7b73ve/FscceG1OnTu3x15b6PtEzZn3XSnmORJj1EaW/R+Z990p5jpj1bynlPepQ3Pc1270777wzq6ioyBYvXpz95je/yS666KJs+PDh2caNG7Msy7JPfvKT2RVXXNFx/C9/+cts4MCB2be+9a3st7/9bbZw4cJs0KBB2TPPPFOsU+jk4osvzqqrq7OlS5dmr7zySsdlx44dHce8/Zyuuuqq7MEHH8yef/75bOXKldk555yTVVZWZs8991wxTmEX//iP/5gtXbo0W79+ffbLX/4ymzFjRnbggQdmmzdvzrKs9+1Ru9bW1mzs2LHZ5ZdfvsvnSn2Ptm7dmq1evTpbvXp1FhHZddddl61evbrjnSu//vWvZ8OHD8/uu+++7Ne//nU2e/bsbMKECdkbb7zRkXHyySdnN9xwQ8fHe3osFvOcWlpasjPPPDMbPXp09vTTT3d6bDU3N3d7Tnu67xbznLZu3Zr90z/9U7Zs2bJs/fr12cMPP5wdc8wx2SGHHJI1NTV1e06lvE/tGhoasn322SdbtGhRlxmltk8UnllfenPk7cz60twj8770571Z33dnfUmX6CzLshtuuCEbO3ZsNnjw4OyEE07Ili9f3vG5D37wg9m5557b6fi77747O/TQQ7PBgwdnkyZNyu6///69vOLuRUSXl1tvvbXjmLef06WXXtpx/jU1Ndnf/M3fZKtWrdr7i+/G2WefnY0aNSobPHhw9p73vCc7++yzs3Xr1nV8vrftUbsHH3wwi4hszZo1u3yu1PfokUce6fJ+1r7mtra2bMGCBVlNTU1WUVGRnXLKKbuc57hx47KFCxd2um53j8VC2905rV+/vtvH1iOPPNLtOe3pvlvMc9qxY0d26qmnZgcddFA2aNCgbNy4cdmFF164y4DsTfvU7uabb86GDBmSbdmypcuMUtsn9g6zvrTmyNuZ9aW5R+Z96c97s77vzvqyLMuyd/oqNgAAAPQnJfs70QAAAFBqlGgAAADISYkGAACAnJRoAAAAyEmJBgAAgJyUaAAAAMhJiQYAAICclGgAAADISYkGAACAnJRoAAAAyEmJBgAAgJyUaAAAAMjp/wOTW0ZN17CF/gAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 1200x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, (ax0, ax1) = pyplot.subplots(1, 2, figsize=(12, 4))\n",
    "ax0.imshow(lhs)\n",
    "ax1.imshow(sparse_dense)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sparse dense matmul with large matrices"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we generate much larger matrices to get some throughput numbers in TFLOPS:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 1\n",
    "shape = [4096, 4096]\n",
    "batch_size = 512\n",
    "density = 1/64\n",
    "dtype = torch.float"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "lhs, dense = generate_lhs_rhs_matrix(shape, block_size, batch_size, dtype) # generate lhs and rhs\n",
    "sparse = magnitude_prune(lhs.to(dtype), block_size, density) # prune lhs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the model with PyTorch gather scatter on cpu, dense multiplication on IPU and PopSparse SpMM on IPU. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "durations_large_matrix = {}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "gather scatter on cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gather_scatter_cpu(sparse, dense):\n",
    "    if dtype != torch.half:\n",
    "        for i in range(num_loops+num_exe):\n",
    "            if i == num_exe:\n",
    "                start_gs = time.time()\n",
    "            output = static_sparse_op.block_coo_spmm_gs(\n",
    "                sparse, dense, mode=\"sparse_dense\"\n",
    "            )\n",
    "        duration = (time.time() - start_gs) / num_loops\n",
    "    return output, duration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, duration = gather_scatter_cpu(sparse, dense)\n",
    "durations_large_matrix[\"duration_gs\"] = duration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "convert sparse matrix to dense format and run matrix multiplication on IPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dense_matmul_ipu(sparse, dense):\n",
    "    sparse_dense_matrix = static_sparse_op.block_coo_to_dense(sparse)\n",
    "    static_sparse_model = DemoDenseModel(sparse_dense_matrix, loop_count=num_loops)\n",
    "    options = poptorch.Options()\n",
    "    options.Precision.setPartialsType(torch.half)\n",
    "    ipu_model_static_sparse = poptorch.inferenceModel(static_sparse_model, options)\n",
    "\n",
    "    for i in range(num_exe):\n",
    "        if i == 0:\n",
    "            out = ipu_model_static_sparse(dense) # run for graph compilation\n",
    "        elif i == num_exe - 1:\n",
    "            start_spmm = time.time()\n",
    "            output = ipu_model_static_sparse(dense) # run without graph compilation\n",
    "    duration = (time.time() - start_spmm) / num_loops \n",
    "    return output, duration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:27<00:00]\n"
     ]
    }
   ],
   "source": [
    "_, duration = dense_matmul_ipu(sparse, dense)\n",
    "durations_large_matrix[\"duration_dense\"] = duration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SpMM on IPU"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the functions to run static and dynamic SpMM on IPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sparse_matmul_static_ipu(sparse, dense):\n",
    "    static_sparse_model = DemoSparseModel(sparse, static_sparse_op.block_coo_spmm_ipu, loop_count=num_loops)\n",
    "    options = poptorch.Options()\n",
    "    options.Precision.setPartialsType(torch.half)\n",
    "    ipu_model_static_sparse = poptorch.inferenceModel(static_sparse_model, options)\n",
    "\n",
    "    for i in range(num_exe):\n",
    "        if i == 0:\n",
    "            output = ipu_model_static_sparse(dense) # run for graph compilation\n",
    "        elif i == num_exe - 1:\n",
    "            start_spmm = time.time()\n",
    "            output = ipu_model_static_sparse(dense) # run without graph compilation\n",
    "    duration = (time.time() - start_spmm) / num_loops \n",
    "    return output, duration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sparse_matmul_dynamic_ipu(sparse, dense):\n",
    "    static_sparse_model = DemoSparseModel(sparse, dynamic_sparse_op.dynamic_spmm_ipu, loop_count=num_loops)\n",
    "    options = poptorch.Options()\n",
    "    options.Precision.setPartialsType(torch.half)\n",
    "    ipu_model_static_sparse = poptorch.inferenceModel(static_sparse_model, options)\n",
    "\n",
    "    for i in range(num_exe):\n",
    "        if i == 0:\n",
    "            output = ipu_model_static_sparse(dense) # run for graph compilation\n",
    "        elif i == num_exe - 1:\n",
    "            start_spmm = time.time()\n",
    "            output = ipu_model_static_sparse(dense) # run without graph compilation\n",
    "    duration = (time.time() - start_spmm) / num_loops \n",
    "    return output, duration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sparse_matmul_static_ipu_profiles(sparse, dense, num_loops):\n",
    "    static_sparse_model = DemoSparseModel(sparse, static_sparse_op.block_coo_spmm_ipu, loop_count=num_loops)\n",
    "    options = poptorch.Options()\n",
    "    options.Precision.setPartialsType(torch.half)\n",
    "    ipu_model_static_sparse = poptorch.inferenceModel(static_sparse_model, options)\n",
    "    output = ipu_model_static_sparse(dense)\n",
    "    return output"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we run the static SpMM and save the duration to the dict of durations_large_matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:19<00:00]\n"
     ]
    }
   ],
   "source": [
    "_, duration = sparse_matmul_static_ipu(sparse, dense)\n",
    "durations_large_matrix[\"duration_spmm_static\"] = duration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dynamic SpMM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:11<00:00]\n"
     ]
    }
   ],
   "source": [
    "_, duration = sparse_matmul_dynamic_ipu(sparse, dense)\n",
    "durations_large_matrix[\"duration_spmm_dynamic\"] = duration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we could compare the throughput with all the three models. We convert the FLOP to TFLOP to better quantify the results. Also note that for dense throughput calculation, we include only non-zero flops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of TFLOP: 0.000268435456\n"
     ]
    }
   ],
   "source": [
    "def calculate_tflop(density, shape, block_size, batch_size):\n",
    "    nnz = int(density * shape[0] * block_size * shape[1] * block_size)\n",
    "    num_blocks = int(nnz / (block_size * block_size))\n",
    "    tflop = 2.0 * batch_size * block_size * block_size * num_blocks / 10**12\n",
    "    return tflop\n",
    "tflop = calculate_tflop(density, shape, block_size, batch_size)\n",
    "print(f\"Number of TFLOP: {tflop}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we could calculate the FLOPS values for gather scatter on cpu, static and dynamic SpMM as well as dense matmul on IPU. We also show the comparison between static_sparse vs. dense and dynamic_sparse vs. dense. From the results we could see that the static sparse gain substantial advantage ($\\sim$ 3 times the dense regarding FLOPS) wih the test case. The dynamic sparse is also slightly faster than dense ($\\sim$ 1.5 times)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "--------------  ----------  -------\n",
      "Gather scatter  0.00125717  TFLOP/s\n",
      "Dense matmul    0.81948     TFLOP/s\n",
      "Static SpMM     3.13993     TFLOP/s\n",
      "Dynamic SpMM    1.30075     TFLOP/s\n",
      "--------------  ----------  -------\n",
      "Static sparse/dense throughput on IPU: 3.83\n",
      "Dynamic sparse/dense throughput on IPU: 1.59\n"
     ]
    }
   ],
   "source": [
    "flops_gs = (tflop / durations_large_matrix[\"duration_gs\"])\n",
    "flops_dense = (tflop / durations_large_matrix[\"duration_dense\"])\n",
    "flops_spmm_static = (tflop / durations_large_matrix[\"duration_spmm_static\"])\n",
    "flops_spmm_dynamic = (tflop / durations_large_matrix[\"duration_spmm_dynamic\"])\n",
    "sparse_dense_tput_ratio_static = flops_spmm_static / flops_dense\n",
    "sparse_dense_tput_ratio_dynamic = flops_spmm_dynamic / flops_dense\n",
    "\n",
    "print(f\"Throughput for different operations:\")\n",
    "print(tabulate.tabulate([(\"Gather scatter\", flops_gs, \"TFLOP/s\"), (\"Dense matmul\", flops_dense, \"TFLOP/s\"),\n",
    "                         (\"Static SpMM\", flops_spmm_static, \"TFLOP/s\"), (\"Dynamic SpMM\", flops_spmm_dynamic, \"TFLOP/s\")]))\n",
    "\n",
    "print(f\"Static sparse/dense throughput on IPU: {sparse_dense_tput_ratio_static :.3g}\")\n",
    "print(f\"Dynamic sparse/dense throughput on IPU: {sparse_dense_tput_ratio_dynamic :.3g}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Micro benchmark with dense matmul and SpMM on IPU"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Block size = 1, data type: half"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 1 # block size of the sparse matrix\n",
    "shape = [4096,4096] # block row and column of the sparse matrix\n",
    "batch_size = 1024 # batch size of the dense matrix\n",
    "dtype = torch.half"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function for benchmarking and printing some throughput and comparison numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "durations_large_matrix_block = {\"duration_dense\": None, \"duration_spmm_static\": None}\n",
    "def benchmark_and_print(block_size, shape, batch_size, density, dtype):\n",
    "      # print basic matrices info\n",
    "      print(f\"m = {shape[0]*block_size}, k = {shape[1]*block_size}, n = {batch_size}, \"\n",
    "            f\"b = {block_size}, d = {density}, dtype = {dtype}\")\n",
    "      # generate matrices  \n",
    "      lhs, dense = generate_lhs_rhs_matrix(shape, block_size, batch_size, dtype)\n",
    "      sparse = magnitude_prune(lhs.to(dtype), block_size, density)\n",
    "      # run dense matmul if dense was not run before\n",
    "      if not durations_large_matrix_block[\"duration_dense\"]:\n",
    "            _, duration = dense_matmul_ipu(sparse, dense)\n",
    "            durations_large_matrix_block[\"duration_dense\"] = duration\n",
    "      # run SpMM\n",
    "      _, duration = sparse_matmul_static_ipu(sparse, dense)\n",
    "      durations_large_matrix_block[\"duration_spmm_static\"] = duration\n",
    "      # calculate flops\n",
    "      tflop = calculate_tflop(density, shape, block_size, batch_size)\n",
    "      flops_spmm_static = (tflop / durations_large_matrix_block[\"duration_spmm_static\"])\n",
    "      flops_dense = (tflop / durations_large_matrix_block[\"duration_dense\"])\n",
    "      sparse_dense_tput_ratio_static = flops_spmm_static / flops_dense\n",
    "      # print throughput and sparse/dense throughput ratio\n",
    "      print(f\"Throughput for different operations:\")\n",
    "      print(tabulate.tabulate([(\"Dense matmul\", flops_dense, \"TFLOP/s\"),\n",
    "                              (\"Static SpMM\", flops_spmm_static, \"TFLOP/s\")]))\n",
    "      print(f\"Static sparse/dense throughput on IPU: {sparse_dense_tput_ratio_static :.3g}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bnechmark for static SpMM with various densities."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/64**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 1024, b = 1, d = 0.015625, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:21<00:00]\n",
      "Graph compilation: 100%|██████████| 100/100 [00:19<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  3.08973  TFLOP/s\n",
      "Static SpMM   6.16198  TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 1.99\n"
     ]
    }
   ],
   "source": [
    "density = 1/64\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/32**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 1024, b = 1, d = 0.03125, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:21<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  6.17946  TFLOP/s\n",
      "Static SpMM   8.20761  TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 1.33\n"
     ]
    }
   ],
   "source": [
    "density = 1/32\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/16**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 1024, b = 1, d = 0.0625, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:24<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  12.3589  TFLOP/s\n",
      "Static SpMM   10.2345  TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 0.828\n"
     ]
    }
   ],
   "source": [
    "density = 1/16\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype) # reference: 10.50"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/8**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 1024, b = 1, d = 0.125, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:28<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  24.7178  TFLOP/s\n",
      "Static SpMM   11.4647  TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 0.464\n"
     ]
    }
   ],
   "source": [
    "density = 1/8\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype) # reference: 11.77"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With block size = 1, static sparse is faster than dense matmul when density is low: for example, density = 1/64 and 1/32. \n",
    "Next, let's have a look at the case with large block size."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Block size = 16, data type = half"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 16 # block size of the sparse matrix\n",
    "shape = [256, 256] # block row and column of the sparse matrix\n",
    "batch_size = 2048 # batch size of the dense matrix\n",
    "dtype = torch.half\n",
    "durations_large_matrix_block = {\"duration_dense\": None, \"duration_spmm_static\": None}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/64**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 2048, b = 16, d = 0.015625, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:35<00:00]\n",
      "Graph compilation: 100%|██████████| 100/100 [00:40<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  3.29656  TFLOP/s\n",
      "Static SpMM   5.1096   TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 1.55\n"
     ]
    }
   ],
   "source": [
    "density = 1/64\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/32**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 2048, b = 16, d = 0.03125, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:38<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  6.59312  TFLOP/s\n",
      "Static SpMM   9.06285  TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 1.37\n"
     ]
    }
   ],
   "source": [
    "density = 1/32\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/16**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 2048, b = 16, d = 0.0625, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:39<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  13.1862  TFLOP/s\n",
      "Static SpMM   18.0305  TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 1.37\n"
     ]
    }
   ],
   "source": [
    "density = 1/16\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**density = 1/8**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m = 4096, k = 4096, n = 2048, b = 16, d = 0.125, dtype = torch.float16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:38<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Throughput for different operations:\n",
      "------------  -------  -------\n",
      "Dense matmul  26.3725  TFLOP/s\n",
      "Static SpMM   32.6905  TFLOP/s\n",
      "------------  -------  -------\n",
      "Static sparse/dense throughput on IPU: 1.24\n"
     ]
    }
   ],
   "source": [
    "density = 1/8\n",
    "benchmark_and_print(block_size, shape, batch_size, density, dtype)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could see that with block size = 16, static sparse with all densities chosen are faster than dense matmul."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Profiling the SpMM on IPU"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the `POPLAR_ENGINE_OPTIONS` to profile the SpMM and turn off profiling after it is done. The profiles are saved to spmm_profile directory. You can use PopVision to visualize the memory profiles and execution traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: POPLAR_ENGINE_OPTIONS={\"autoReport.all\":\"true\", \"autoReport.directory\":\"./spmm_profile\"}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Graph compilation: 100%|██████████| 100/100 [00:22<00:00]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: POPLAR_ENGINE_OPTIONS={\"autoReport.all\":\"false\"}\n"
     ]
    }
   ],
   "source": [
    "%env POPLAR_ENGINE_OPTIONS={\"autoReport.all\":\"true\", \"autoReport.directory\":\"./spmm_profile\"}\n",
    "output = sparse_matmul_static_ipu_profiles(sparse, dense, 5)\n",
    "%env POPLAR_ENGINE_OPTIONS={\"autoReport.all\":\"false\"}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Final thoughts\n",
    "\n",
    "This notebook demonstrates fast elementwise and block sparsity using PopSparse for inference on IPU. We hope that this work shows that block sparsity is a promising method for achieving practical acceleration of sparse models and will motivate further investigation into effective block sparse pruning algorithms in the research community.\n",
    "\n",
    "If you are actively pursuing research in sparsity that could benefit from IPU support please consider signing up for our [academic program](https://www.graphcore.ai/academics) where we support researchers to achieve great things using IPUs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}