posts/solving-sparse-linear-systems-in-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a Pytorch solver for sparse linear systems\n",
    "Not so long ago, I implemented a [wrapper library in PyTorch to solve sparse linear systems](https://github.com/flaport/torch_sparse_solve) on the CPU using the [SuiteSparse](https://github.com/DrTimothyAldenDavis/SuiteSparse) routines under the hood. My goal is to eventually integrate this sparse solver into my photonic circuit simulator, [Photontorch](https://github.com/flaport/photontorch). However, for now, I thought it would be instructive to go over the steps I took to implement both the forward pass and the backward pass of such a custom PyTorch function.\n",
    "\n",
    "So, with that in mind, the goal of the discussion below is to define a custom PyTorch function that solves the sparse linear system of the form\n",
    "$$\n",
    "\\begin{align}\n",
    "Ax &= b\n",
    "\\end{align}\n",
    "$$\n",
    "Where $A$ is sparse and $x$ and $b$ are dense."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A very simple system of equations: Scalar x Scalar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by implementing the $0D$ case. When $A$ is a $0D$ matrix (i.e. just a normal number or scalar). this system of equations has the following solution:\n",
    "\n",
    "\\begin{align*}\n",
    "    x &= \\frac{b}{A}\n",
    "\\end{align*}\n",
    "\n",
    "Implementing custom functions (with custom forward **and** backward) in PyTorch is usually done by subclassing `torch.autograd.Function` and defining the `forward` and `backward` method for your derived class. For example in the case of our simple linear system with the above solution we can write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.autograd import Function\n",
    "\n",
    "class Solve(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, A, b):\n",
    "        if A.ndim > 0 and not (A.ndim==1 and A.shape[0]==1):\n",
    "            raise ValueError(\"A should be 0D\")\n",
    "        ctx.save_for_backward(A, b)\n",
    "        return b / A\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad):\n",
    "        A, b = ctx.saved_tensors\n",
    "        return (-b/A**2)*grad, (1/A)*grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How this works is as follows: the forward pass performs the operation *and* stores the relevant matrices in the context object `ctx` for later. This way these matrices can later on be used again during the gradient calculation in the backward pass. In the backward pass, the gradient of the loss $L$ must be calculated with respect to all the inputs of the forward pass ($A$ and $b$ in this case), given the gradients of the loss with respect to the output of the forward pass.\n",
    "\n",
    "This means that in this case we need to find the following two transformations to describe the backward pass:\n",
    "$$\n",
    "\\begin{align*}\n",
    "    \\frac{\\partial L}{\\partial x} \\rightarrow \\frac{\\partial L}{\\partial A} \\\\\n",
    "    \\frac{\\partial L}{\\partial x} \\rightarrow \\frac{\\partial L}{\\partial b} \n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "By applying the chain rule, we have:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "    \\frac{\\partial L}{\\partial A} = \\frac{\\partial L}{\\partial x}\\frac{\\partial x}{\\partial A} &&\n",
    "    \\frac{\\partial L}{\\partial b} = \\frac{\\partial L}{\\partial x}\\frac{\\partial x}{\\partial b}\n",
    "\\end{align*} \n",
    "$$\n",
    "\n",
    "Hence defining the backward pass comes down to determining $\\frac{\\partial x}{\\partial A}$ and $\\frac{\\partial x}{\\partial b}$:\n",
    "\n",
    "$$\n",
    "\\begin{align*} \n",
    "    \\frac{\\partial x}{\\partial A} = - \\frac{b}{A^2} &&\n",
    "    \\frac{\\partial x}{\\partial b} = \\frac{1}{A} \n",
    "\\end{align*}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These gradients are simple enough to know that we implemented them correctly. However, Pytorch also offers a `gradcheck` function to check wether we implemented the backward pass correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch.autograd import gradcheck\n",
    "A = torch.randn(1, requires_grad=True).squeeze()\n",
    "b = torch.randn(1, requires_grad=True).squeeze()\n",
    "gradcheck(Solve.apply, [A.double(), b.double()]) # gradcheck requires double precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we used `Solve.apply` to use the function during `gradcheck`. It is common in PyTorch to actually alias this to a lowercase function name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "solve = Solve.apply"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A more general system: Matrix × Vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok, what we implemented above was quite trivial, but mostly served as an example on how the forward and backward of custom functions are implemented in PyTorch. Now that's behind us we can go on to more interesting domains... general systems where A is a square, $2D$ matrix acting on a $1D$ vector.\n",
    "\n",
    "As we know, the solution to the system of equation becomes in this case:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "    x = A^{-1} \\, b\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, calculating $A^{-1}$ is an expensive operation and should be avoided. Let's say, for the sake of our story that we want to use `scipy.linalg.solve`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.linalg import solve as scipy_solve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sidenote: Pytorch actually has a `torch.solve` function, which (in contrast to `scipy.linalg.solve`) works on CUDA GPUs as well. Hence in 99% of the cases this is the function you'll want. However, we go along here with `scipy.linalg.solve` as hopefully we'll learn something from writing the PyTorch wrapper. At the end of this post, we'll then implement our own sparse solver in C++."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Solve(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, A, b):\n",
    "        if A.ndim != 2 or (A.shape[0] != A.shape[1]):\n",
    "            raise ValueError(\"A should be a square 2D matrix.\")\n",
    "        A_np = A.data.numpy()\n",
    "        b_np = b.data.numpy()\n",
    "        x_np = scipy_solve(A_np, b_np)\n",
    "        x = torch.tensor(x_np, requires_grad=True)\n",
    "        ctx.save_for_backward(A, b, x)\n",
    "        return x\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad):\n",
    "        A, b, x = ctx.saved_tensors\n",
    "        gradb = Solve.apply(A.T, grad)\n",
    "        gradA = -gradb[:,None] * x[None,:]\n",
    "        return gradA, gradb\n",
    "\n",
    "solve = Solve.apply"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Where we used the following expressions for the gradients:\n",
    "$$\n",
    "\\begin{align*}\n",
    "    \\frac{\\partial L}{\\partial b} \n",
    "        &= \\frac{\\partial L}{\\partial x_i} \\frac{\\partial x_i}{\\partial b_j}  \\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_i} \\frac{\\partial}{\\partial b_k} ( A^{-1}_{ij} b_j )  \\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_i} A^{-1}_{ij} \\frac{\\partial b_j}{\\partial b_k}\\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_i} A^{-1}_{ij} \\delta_{jk}\\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_i} A^{-1}_{ik} \\\\\n",
    "        &= \\big(A^{-1}\\big)^{T} \\frac{\\partial L}{\\partial x} \\\\\n",
    "        &= \\mathrm{solve}\\big( A^T\\,,\\,\\,  \\frac{\\partial L}{\\partial x} \\big) \\\\\\\\\n",
    "\\frac{\\partial L}{\\partial A} \n",
    "        &= \\frac{\\partial L}{\\partial x_i} \\frac{\\partial x_i}{\\partial A_{mn}}  \\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_i} \\frac{\\partial}{\\partial A_{mn}} ( A^{-1}_{ij} b_j ) \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial x_i} A^{-1}_{ij} \\frac{\\partial A_{jk}}{\\partial A_{mn}} A^{-1}_{kl} b_l \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial x_i} A^{-1}_{ij} \\delta_{jm} \\delta_{kn} A^{-1}_{kl} b_l \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial x_i} A^{-1}_{im} A^{-1}_{nl} b_l \\\\\n",
    "        &= -\\left(\\big(A^{-1}\\big)^{T} \\frac{\\partial L}{\\partial x}\\right)\\otimes\\left( A^{-1} b \\right) \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial b} \\otimes x\n",
    "\\end{align*}\n",
    "$$\n",
    "where we used the einstein summation convention during the derivations. We also used the following identity:\n",
    "$$\n",
    "\\frac{\\partial (A^{-1})}{\\partial p} = - A^{-1} \\frac{\\partial A}{\\partial p} A^{-1}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check that the backward pass was implemented correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = torch.randn(3,3, requires_grad=True)\n",
    "b = torch.randn(3, requires_grad=True)\n",
    "gradcheck(solve, [A.double(), b.double()]) # gradcheck requires double precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Even more general: Matrix × Matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In fact, `scipy.linalg.solve` also supports Matrix × Matrix systems, hence extending the backward method of the above solver is not so difficult:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Solve(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, A, b):\n",
    "        if A.ndim != 2 or (A.shape[0] != A.shape[1]):\n",
    "            raise ValueError(\"A should be a square 2D matrix.\")\n",
    "        A_np = A.data.numpy()\n",
    "        b_np = b.data.numpy()\n",
    "        x_np = scipy_solve(A_np, b_np)\n",
    "        x = torch.tensor(x_np, requires_grad=True)\n",
    "        ctx.save_for_backward(A, b, x)\n",
    "        return x\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad):\n",
    "        A, b, x = ctx.saved_tensors\n",
    "        gradb = Solve.apply(A.T, grad)\n",
    "        gradA = -gradb @ x.T\n",
    "        return gradA, gradb\n",
    "\n",
    "solve = Solve.apply"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Where we used a slightly different derivation as above, taking into account that $x$ and $b$ are $m\\times n$ matrices (and $A$ remains an $m\\times m$ matrix):\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "    \\frac{\\partial L}{\\partial b}\n",
    "        &= \\frac{\\partial L}{\\partial x_{ij}} \\frac{\\partial x_{ij}}{\\partial b_{kl}}  \\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_{ij}} A^{-1}_{ik} \\frac{\\partial b_{kj}}{\\partial b_{lm}}\\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_{ij}} A^{-1}_{ik} \\delta_{kl}\\delta_{jm}\\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_{im}} A^{-1}_{il}\\\\\n",
    "        &= \\big(A^{-1}\\big)^{T} \\frac{\\partial L}{\\partial x} \\\\\n",
    "        &= \\mathrm{solve}\\big( A^T\\,,\\,\\,  \\frac{\\partial L}{\\partial x} \\big) \\\\\\\\\n",
    "\\frac{\\partial L}{\\partial A} \n",
    "        &= \\frac{\\partial L}{\\partial x_{ij}} \\frac{\\partial x_{ij}}{\\partial A_{mn}}  \\\\\n",
    "        &= \\frac{\\partial L}{\\partial x_{ij}} \\frac{\\partial}{\\partial A_{mn}} ( A^{-1}_{ik} b_{kj} ) \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial x_{ij}} A^{-1}_{ik} \\frac{\\partial A_{kl}}{\\partial A_{mn}} A^{-1}_{lp} b_{pj} \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial x_{ij}} A^{-1}_{ik} \\delta_{km} \\delta_{ln} A^{-1}_{lp} b_{pj} \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial x_{ij}} A^{-1}_{im} A^{-1}_{np} b_{pj} \\\\\n",
    "        &= -\\left(\\big(A^{-1}\\big)^{T} \\frac{\\partial L}{\\partial x}\\right)\\left( A^{-1} b \\right)^T \\\\\n",
    "        &= -\\frac{\\partial L}{\\partial b} x^T\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = torch.randn(3,3, requires_grad=True)\n",
    "b = torch.randn(3,2)\n",
    "gradcheck(solve, [A.double(), b.double()]) # gradcheck requires double precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Most general: Batched Matrix × Batched Matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above Matrix $\\times$ Matrix system is easily extended to a batched version (with the batch dimension being the first dimension):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Solve(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, A, b):\n",
    "        if A.ndim != 3 or (A.shape[1] != A.shape[2]):\n",
    "            raise ValueError(\"A should be a batch of square 2D matrices with shape (b, m, m)\")\n",
    "        A_np = A.data.numpy()\n",
    "        b_np = b.data.numpy()\n",
    "        x_np = np.stack([scipy_solve(A_np[i], b_np[i]) for i in range(A.shape[0])], 0)\n",
    "        x = torch.tensor(x_np, requires_grad=True)\n",
    "        ctx.save_for_backward(A, b, x)\n",
    "        return x\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad):\n",
    "        A, b, x = ctx.saved_tensors\n",
    "        gradb = Solve.apply(A.transpose(-1,-2), grad)\n",
    "        gradA = -torch.bmm(gradb, x.transpose(-1,-2))\n",
    "        return gradA, gradb\n",
    "\n",
    "solve = Solve.apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = torch.randn(4,3,3, requires_grad=True)\n",
    "b = torch.randn(4,3,2, requires_grad=True)\n",
    "gradcheck(solve, [A.double(), b.double()]) # gradcheck requires double precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making a C++ extension\n",
    "\n",
    "Let's now go on to make the above module in C++. However, we won't do this with the imported scipy function. For now, we'll just use `torch.inverse` when we need it. Just know, that after we get this simple C++ PyTorch extension to work, we'll swap out `torch.inverse` for our own C++ solver.\n",
    "\n",
    "We'll start by making a file called `solve.cpp`, which includes the following two headers:\n",
    "```cpp\n",
    "#include <torch/extension.h>\n",
    "#include <ATen/ATen.h>\n",
    "#include <vector>\n",
    "```\n",
    "\n",
    "Generally speaking, `torch/extension.h` implements equivalent C++ functions to what `torch` offers in python, while `ATen/ATen.h` offers Python Tensor *methods* as C++ *functions*. The `vector` header is needed when returning more than one tensor (as we'll do in the backward).\n",
    "\n",
    "Hence the batched matrix x batched matrix forward can be implemented in C++ as follows (remember we'll swap out `torch::inverse` for an actual solver later):\n",
    "```cpp\n",
    "torch::Tensor solve_forward(torch::Tensor A, torch::Tensor b){\n",
    "  auto result = torch::zeros_like(b);\n",
    "  for (int i = 0; i < at::size(b, 0); i++){\n",
    "      result[i] = torch::mm(torch::inverse(A[i]), b[i]); // we'll use an actual solver later.\n",
    "  }\n",
    "  return result;\n",
    "}\n",
    "```\n",
    " \n",
    "Implementing the backward pass can also be done relatively easily:\n",
    "```cpp\n",
    "std::vector<torch::Tensor> solve_backward(torch::Tensor grad, torch::Tensor A, torch::Tensor b, torch::Tensor x){\n",
    "    auto gradb = at::transpose(solve_forward(at::transpose(A, -1, -2), grad), -1, -2);\n",
    "    auto gradA = -torch::bmm(at::transpose(gradb, -1, -2), at::transpose(x, -1, -2));\n",
    "    return {gradA, gradb};\n",
    "}\n",
    "```\n",
    "Notice that we're returning a *vector* of tensors here, hence why we needed to include `vector` earlier.\n",
    "\n",
    "Finally, we need to register the C++ functions defined as python functions. This is done with pybind:\n",
    "```cpp\n",
    "PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n",
    "  m.def(\"forward\", &solve_forward, \"solve forward\");\n",
    "  m.def(\"backward\", &solve_backward, \"solve backward\");\n",
    "}\n",
    "```\n",
    "\n",
    "Here the macro `TORCH_EXTENSION_NAME` will be replaced during compilation to the name of the torch extension defined in the `setup.py` file. The `setup.py` file looks as follows:\n",
    "```python\n",
    "from setuptools import setup, Extension\n",
    "from torch.utils import cpp_extension\n",
    "\n",
    "solve_cpp = Extension(\n",
    "    name=\"solve_cpp\",\n",
    "    sources=[\"solve.cpp\"],\n",
    "    include_dirs=cpp_extension.include_paths(),\n",
    "    library_dirs=cpp_extension.library_paths(),\n",
    "    extra_compile_args=[],\n",
    "    libraries=[\n",
    "        \"c10\",\n",
    "        \"torch\",\n",
    "        \"torch_cpu\",\n",
    "        \"torch_python\",\n",
    "    ],\n",
    "    language=\"c++\",\n",
    ")\n",
    "\n",
    "setup(\n",
    "    name=\"solve\",\n",
    "    ext_modules=[solve_cpp],\n",
    "    cmdclass={\"build_ext\": cpp_extension.BuildExtension},\n",
    ")\n",
    "```\n",
    "\n",
    "The C++ extension can now be compiled as follows:\n",
    "```bash\n",
    "python setup.py install\n",
    "```\n",
    "Which will create a python executable (`.so` file on linux, `.pyd` file on windows) in your python's `site-packages` folder (i.e. it will be in your python path).\n",
    "\n",
    "The only thing that's left is creating a thin wrapper in python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch # always import torch BEFORE your custom torch extension\n",
    "import solve_cpp # the custom torch extension we just created\n",
    "\n",
    "class Solve(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, A, b):\n",
    "        if A.ndim != 3 or (A.shape[1] != A.shape[2]):\n",
    "            raise ValueError(\"A should be a batch of square 2D matrices with shape (b, m, m)\")\n",
    "        if b.ndim != 3:\n",
    "            raise ValueError(\"b should be a batch of matrices with shape (b, m, n)\")\n",
    "        x = solve_cpp.forward(A, b)\n",
    "        ctx.save_for_backward(A, b, x)\n",
    "        return x\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad):\n",
    "        A, b, x = ctx.saved_tensors\n",
    "        gradA, gradb = solve_cpp.backward(grad, A, b, x)\n",
    "        return gradA, gradb\n",
    "\n",
    "solve = Solve.apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = torch.randn(4,3,3, requires_grad=True)\n",
    "b = torch.randn(4,3,2)\n",
    "gradcheck(solve, [A.double(), b.double()]) # gradcheck requires double precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Towards a sparse solver for CPU tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OK, so now we have a representation for the backward pass and we know how to make a Pytorch C++ extension. Let's go on to make an actual sparse solver."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To create such a sparse solver, we'll wrap the [SuiteSparse](https://github.com/DrTimothyAldenDavis/SuiteSparse) routines in our C++ extension. In particular, we'll wrap the [KLU algorithm](https://ufdc.ufl.edu/UFE0011721/00001) provided by this library. The KLU sparse linear system solver is a very efficient solver for sparse matrices that arise from circuit simulation netlists. This means it will be most efficient for very sparse systems with often only one element per column of the (connection-)matrix. Obviously, it can tackle different linear systems of equations as well, but that's what it was originally intended for."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use the KLU algorithm, we'll have to add the `klu` header from the `SuiteSparse` library to the three headers we had in the C++ extension before.\n",
    "```c++\n",
    "#include <torch/extension.h>\n",
    "#include <ATen/ATen.h>\n",
    "#include <vector>\n",
    "#include <klu.h>\n",
    "```\n",
    "\n",
    "Note that with the [Anaconda distribution](https://www.anaconda.com/) you can simply do a\n",
    "```sh\n",
    "conda install suitesparse\n",
    "```\n",
    "to pull in all the SuiteSparse C++ libraries."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, before we can go on to solving the sparse system with the KLU algorithm, there's one more hurdle to overcome:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sparse COO -> Sparse CSC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The KLU algorithm expects the sparse matrix to be in [CSC-format](https://en.wikipedia.org/wiki/Sparse_matrix#Compressed_sparse_column_(CSC_or_CCS)) rather than [COO-format](https://en.wikipedia.org/wiki/Sparse_matrix#Coordinate_list_(COO)), the standard representation used in PyTorch. To do this conversion, we can have a look at how scipy does this conversion, and do something similar for our COO-pytorch tensors:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```c++\n",
    "\n",
    "std::vector<at::Tensor> _coo_to_csc(int ncol, at::Tensor Ai, at::Tensor Aj, at::Tensor Ax) {\n",
    "    int nnz = at::size(Ax, 0);\n",
    "    at::TensorOptions options = at::TensorOptions().dtype(torch::kInt32).device(at::device_of(Ai));\n",
    "    at::Tensor Bp = at::zeros(ncol+1, options);\n",
    "    at::Tensor Bi = at::zeros_like(Ai);\n",
    "    at::Tensor Bx = at::zeros_like(Ax);\n",
    "\n",
    "    int* ai = Ai.data_ptr<int>();\n",
    "    int* aj = Aj.data_ptr<int>();\n",
    "    double* ax = Ax.data_ptr<double>();\n",
    "\n",
    "    int* bp = Bp.data_ptr<int>();\n",
    "    int* bi = Bi.data_ptr<int>();\n",
    "    double* bx = Bx.data_ptr<double>();\n",
    "\n",
    "    //compute number of non-zero entries per row of A\n",
    "    for (int n = 0; n < nnz; n++) {\n",
    "        bp[aj[n]] += 1;\n",
    "    }\n",
    "\n",
    "    //cumsum the nnz per row to get Bp\n",
    "    int cumsum = 0;\n",
    "    int temp = 0;\n",
    "    for(int j = 0; j < ncol; j++) {\n",
    "        temp = bp[j];\n",
    "        bp[j] = cumsum;\n",
    "        cumsum += temp;\n",
    "    }\n",
    "    bp[ncol] = nnz;\n",
    "\n",
    "    //write Ai, Ax into Bi, Bx\n",
    "    int col = 0;\n",
    "    int dest = 0;\n",
    "    for(int n = 0; n < nnz; n++) {\n",
    "        col = aj[n];\n",
    "        dest = bp[col];\n",
    "        bi[dest] = ai[n];\n",
    "        bx[dest] = ax[n];\n",
    "        bp[col] += 1;\n",
    "    }\n",
    "\n",
    "    int last = 0;\n",
    "    for(int i = 0; i <= ncol; i++) {\n",
    "        temp = bp[i];\n",
    "        bp[i] = last;\n",
    "        last = temp;\n",
    "    }\n",
    "\n",
    "    return {Bp, Bi, Bx};\n",
    "}\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we're converting the PyTorch CPU tensor to a C-array (by asking for its pointer). This means that this conversion will be CPU-only. However, performing this conversion on native pytorch tensors would be **a lot** slower.\n",
    "\n",
    "This function returns three pytorch tensors: `Bp`: the column pointers, `Bi`: the indices in each column and `Bx`: the values of the sparse tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### KLU Solver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using these three vectors, one can define a KLU solver, by wrapping the KLU routines as follows:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```c++\n",
    "void _klu_solve(at::Tensor Ap, at::Tensor Ai, at::Tensor Ax, at::Tensor b) {\n",
    "    int ncol = at::size(Ap, 0) - 1;\n",
    "    int nb = at::size(b, 0);\n",
    "    int* ap = Ap.data_ptr<int>();\n",
    "    int* ai = Ai.data_ptr<int>();\n",
    "    double* ax = Ax.data_ptr<double>();\n",
    "    double* bb = b.data_ptr<double>();\n",
    "    klu_symbolic* Symbolic;\n",
    "    klu_numeric* Numeric;\n",
    "    klu_common Common;\n",
    "    klu_defaults(&Common);\n",
    "    Symbolic = klu_analyze(ncol, ap, ai, &Common);\n",
    "    Numeric = klu_factor(ap, ai, ax, Symbolic, &Common);\n",
    "    klu_solve(Symbolic, Numeric, ncol, nb/ncol, bb, &Common);\n",
    "    klu_free_symbolic(&Symbolic, &Common);\n",
    "    klu_free_numeric(&Numeric, &Common);\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the KLU algorithms comes down to first doing a symbolic analyzation and factorization of the sparse matrix `A` (i.e. `Ap`, `Ai` and `Ax`), probably to determine the sparsity pattern after which the system is solved with `klu_solve`. Note that this is an inplace operation on b, i.e. after solving, the solution `x` will be in the `b` tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Updated Forward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can update the forward method by using our `_klu_solve` wrapper in stead of `torch::inverse`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```c++\n",
    "at::Tensor solve_forward(at::Tensor A, at::Tensor b) {\n",
    "    int p = at::size(b, 0);\n",
    "    int m = at::size(b, 1);\n",
    "    int n = at::size(b, 2);\n",
    "    at::Tensor bflat = at::clone(at::reshape(at::transpose(b, 1, 2), {p, m*n}));\n",
    "    at::Tensor Ax = at::reshape(A._values(), {p, -1});\n",
    "    at::Tensor Ai = at::reshape(at::_cast_Int(A._indices()[1]), {p, -1});\n",
    "    at::Tensor Aj = at::reshape(at::_cast_Int(A._indices()[2]), {p, -1});\n",
    "    for (int i = 0; i < p; i++) {\n",
    "        std::vector<at::Tensor> Ap_Ai_Ax = _coo_to_csc(m, Ai[i], Aj[i], Ax[i]);\n",
    "        _klu_solve(Ap_Ai_Ax[0], Ap_Ai_Ax[1], Ap_Ai_Ax[2], bflat[i]); // result will be in bflat\n",
    "    }\n",
    "    return at::transpose(bflat.view({p,n,m}), 1, 2);\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Updated Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `setup.py` for this extension also becomes a bit more complex, as the SuiteSparse libraries need to be included:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "import os\n",
    "import glob\n",
    "from setuptools import setup, Extension\n",
    "from torch.utils import cpp_extension\n",
    "\n",
    "libroot = os.path.dirname(os.path.dirname(os.__file__))\n",
    "if os.name == \"nt\":  # Windows\n",
    "    suitesparse_lib = os.path.join(libroot, \"Library\", \"lib\")\n",
    "    suitesparse_include = os.path.join(libroot, \"Library\", \"include\", \"suitesparse\")\n",
    "else:  # Linux / Mac OS\n",
    "    suitesparse_lib = os.path.join(os.path.dirname(libroot), \"lib\")\n",
    "    suitesparse_include = os.path.join(os.path.dirname(libroot), \"include\")\n",
    "\n",
    "torch_sparse_solve_cpp = Extension(\n",
    "    name=\"torch_sparse_solve_cpp\",\n",
    "    sources=[\"torch_sparse_solve.cpp\"],\n",
    "    include_dirs=[*cpp_extension.include_paths(), suitesparse_include],\n",
    "    library_dirs=[*cpp_extension.library_paths(), suitesparse_lib],\n",
    "    extra_compile_args=[],\n",
    "    libraries=[\n",
    "        \"c10\",\n",
    "        \"torch\",\n",
    "        \"torch_cpu\",\n",
    "        \"torch_python\",\n",
    "        \"klu\",\n",
    "        \"btf\",\n",
    "        \"amd\",\n",
    "        \"colamd\",\n",
    "        \"suitesparseconfig\",\n",
    "    ],\n",
    "    language=\"c++\",\n",
    ")\n",
    "\n",
    "setup(\n",
    "    name=\"torch_sparse_solve\",\n",
    "    ext_modules=[torch_sparse_solve_cpp],\n",
    "    cmdclass={\"build_ext\": cpp_extension.BuildExtension},\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Updated python wrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can update the python wrapper to include the newly built C++ extension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_sparse_solve_cpp import solve_forward, solve_backward\n",
    "\n",
    "class Solve(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, A, b):\n",
    "        if A.ndim != 3 or (A.shape[1] != A.shape[2]) or not A.is_sparse:\n",
    "            raise ValueError(\n",
    "                \"'A' should be a batch of square 2D sparse matrices with shape (b, m, m).\"\n",
    "            )\n",
    "        if b.ndim != 3:\n",
    "            raise ValueError(\"'b' should be a batch of matrices with shape (b, m, n).\")\n",
    "        if not A.dtype == torch.float64:\n",
    "            raise ValueError(\"'A' should be a sparse float64 tensor (for now). Please first convert to float64.\")\n",
    "        if not b.dtype == torch.float64:\n",
    "            raise ValueError(\"'b' should be a float64 tensor (for now). Please first convert to float64\")\n",
    "\n",
    "        x = solve_forward(A, b)\n",
    "        ctx.save_for_backward(A, b, x)\n",
    "        return x\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad):\n",
    "        A, b, x = ctx.saved_tensors\n",
    "        gradA, gradb = solve_backward(grad, A, b, x)\n",
    "        return gradA, gradb\n",
    "    \n",
    "solve = Solve.apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask = torch.tensor([[[1, 0, 0], [1, 1, 0], [0, 0, 1]]], dtype=torch.float64)\n",
    "A = mask * torch.randn(4, 3, 3, dtype=torch.float64)\n",
    "Asp = A.to_sparse()\n",
    "Asp.requires_grad_()\n",
    "b = torch.randn(4, 3, 2, dtype=torch.float64, requires_grad=True)\n",
    "gradcheck(solve, [Asp, b], check_sparse_nnz=True) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## That's it!\n",
    "\n",
    "Those were the steps I went through creating my first PyTorch C++ extension. Please check it out on GitHub: https://github.com/flaport/torch_sparse_solve and consider giving it a star 😉"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}