07_Linear_Algebra_Profiling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "initial-movement",
   "metadata": {},
   "source": [
    "# Linear Algebra, Precision and Profiling\n",
    "\n",
    "Julias generic way of implementing algorithms often makes it easy to explore different storage schemes, elevated or reduced precision or to try acceleration hardware like a GPU. I want to present a few illustrating examples on a real-world iterative algorithm to show you how little effort is needed to give these things a try in Julia. I will also show in one example how one can track performance by profiling and understand what should be done to improve an algorithm at hand."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "883d1c41",
   "metadata": {},
   "source": [
    "## Linear Algebra\n",
    "\n",
    "For dense and sparse arrays, all important linear algebra routines are available in the `LinearAlgebra`. This includes common tasks such as\n",
    "- `qr` (also pivoted)\n",
    "- `cholesky` (also pivoted)\n",
    "- `eigen`, `eigvals`, `eigvecs` (compute eigenpairs, values, vectors)\n",
    "- `factorize` (for computing matrix factorisations)\n",
    "- `inv` (invert a matrix)\n",
    "\n",
    "All these methods are both implemented for generic matrices (all `AbstractMatrices`) and specialised for specific kinds. For example `factorize` is intended to compute a clever factorisation for solving linear systems. What it does depends on the matrix properties:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4589840",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra\n",
    "using SparseArrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5d91451",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Random real matrix -> will do an LU\n",
    "A = randn(10, 10)\n",
    "@show typeof(factorize(A))\n",
    "\n",
    "# Real-symmetric matrix ->  will do a Bunch-Kaufman\n",
    "Am = Symmetric(A + A')\n",
    "@show typeof(factorize(Am))\n",
    "\n",
    "# Symmetric tridiagonal -> will do a LDLt\n",
    "Am = SymTridiagonal(A + A')\n",
    "@show typeof(factorize(Am))\n",
    "\n",
    "# Random sparse matrix -> will do sparse LU\n",
    "S = sprandn(50, 50, 0.3)\n",
    "@show typeof(factorize(S))\n",
    "\n",
    "# ... and so on ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeedb7c2",
   "metadata": {},
   "source": [
    "The all share a common interface, such that an algorithm like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9f31e67",
   "metadata": {},
   "outputs": [],
   "source": [
    "function solve_many(A, xs)\n",
    "    F = factorize(A)\n",
    "    [F \\ rhs for rhs in xs]\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e201cf42",
   "metadata": {},
   "source": [
    "will automatically work for sparse arrays and dense arrays and is furthermore independent of the floating-point type.\n",
    "\n",
    "##### More details\n",
    "- https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ba0c729",
   "metadata": {},
   "source": [
    "## Use case: A generic Davidson\n",
    "\n",
    "Let's try this in a more realistic algorithm.\n",
    "A simple Davidson algorithm can be implemented quite concisely:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "general-computer",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra\n",
    "\n",
    "qrortho(X::Array)   = Array(qr(X).Q)\n",
    "qrortho(X, Y)       = qrortho(X - Y * Y'X)\n",
    "\n",
    "function rayleigh_ritz(X::Array, AX::Array, N)\n",
    "    F = eigen(Hermitian(X'AX))\n",
    "    F.values[1:N], F.vectors[:,1:N]\n",
    "end\n",
    "\n",
    "function davidson(A, SS::AbstractArray; tol=1e-5, maxsubspace=8size(SS, 2), verbose=true)\n",
    "    m = size(SS, 2)\n",
    "    for i in 1:100\n",
    "        Ass = A * SS\n",
    "        rvals, rvecs = rayleigh_ritz(SS, Ass, m)\n",
    "        Ax = Ass * rvecs\n",
    "\n",
    "        R = Ax - SS * rvecs * Diagonal(rvals)\n",
    "        if norm(R) < tol\n",
    "            return rvals, SS * rvecs\n",
    "        end\n",
    "\n",
    "        verbose && println(i, \"  \", size(SS, 2), \"  \", norm(R))\n",
    "\n",
    "        # Use QR to orthogonalise the subspace.\n",
    "        if size(SS, 2) + m > maxsubspace\n",
    "            SS = qrortho([SS*rvecs R])\n",
    "        else\n",
    "            SS = qrortho([SS       R])\n",
    "        end\n",
    "    end\n",
    "    error(\"not converged.\")\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "accompanied-homework",
   "metadata": {},
   "outputs": [],
   "source": [
    "nev = 2\n",
    "A = randn(50, 50); A = A + A' + 5I;\n",
    "\n",
    "# Generate two random orthogonal guess vectors\n",
    "x0 = qrortho(randn(size(A, 2), nev))\n",
    "\n",
    "# Run the problem\n",
    "davidson(A, x0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "competent-homework",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mixed precision!\n",
    "using GenericLinearAlgebra\n",
    "\n",
    "λ, v = davidson(Matrix{Float32}(A), Float32.(x0), tol=1e-3)\n",
    "println()\n",
    "λ, v = davidson(Matrix{Float64}(A), v, tol=1e-13)\n",
    "println()\n",
    "λ, v = davidson(Matrix{BigFloat}(A), v, tol=1e-25)\n",
    "λ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b242c18f",
   "metadata": {},
   "outputs": [],
   "source": [
    "using SparseArrays\n",
    "nev = 2\n",
    "spA = sprandn(100, 100, 0.3); spA = spA + spA' + 2I\n",
    "\n",
    "spA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "separate-moore",
   "metadata": {},
   "outputs": [],
   "source": [
    "spx0 = randn(size(spA, 2), nev)\n",
    "spx0 = Array(qr(spx0).Q)\n",
    "\n",
    "davidson(spA, spx0, tol=1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tribal-chase",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ... runs with GPUs !\n",
    "using CUDA\n",
    "\n",
    "qrortho(X::CuArray) = CuArray(qr(X).Q)\n",
    "function rayleigh_ritz(X::CuArray, AX::CuArray, N)\n",
    "    values, vectors = CUDA.CUSOLVER.syevd!('V', 'U', X'AX)\n",
    "    values[1:N], vectors[:, 1:N]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "centered-elite",
   "metadata": {},
   "outputs": [],
   "source": [
    "davidson(cu(A), cu(x0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a02a6dd",
   "metadata": {},
   "source": [
    "... but actually the performance is overall not that good out of the box, because we're doing a lot of copying and elementwise access in our naive algorithm, which is especially bad for the GPU version."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d0a6e28",
   "metadata": {},
   "source": [
    "## Iterative methods\n",
    "\n",
    "Instead of implementing iterative methods such as the Davidson diagonalisation ourselves, we can also build upon existing packages for standard linear algebra, such as [IterativeSolvers.jl](https://github.com/JuliaLinearAlgebra/IterativeSolvers.jl), [KrylovKit.jl](https://github.com/Jutho/KrylovKit.jl), [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl).\n",
    "\n",
    "For example instead of hand-coding a Davidson, we could use IterativeSolvers' LOBPCG implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d822e8d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "using IterativeSolvers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2689d0e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "largest = false\n",
    "IterativeSolvers.lobpcg(A, largest, x0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aba4370b",
   "metadata": {},
   "source": [
    "which works seamlessly with GPUs as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2932ce33",
   "metadata": {},
   "outputs": [],
   "source": [
    "largest = false\n",
    "IterativeSolvers.lobpcg(cu(A), largest, cu(x0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2181b3d5",
   "metadata": {},
   "source": [
    "## Profiling and timing measurements\n",
    "\n",
    "Now how would one go about improving this piece of code?\n",
    "\n",
    "The best way forward is to obtain an idea *where* the computational time is spent and then think where we could *locally* improve. We already saw the `@btime` macro (from [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl) for getting accurate timing measurements on single instructions. Let's see what other options there are.\n",
    "\n",
    "For our tests we will use this piece of code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f12fa526",
   "metadata": {},
   "outputs": [],
   "source": [
    "function myfunction(n)\n",
    "    for i = 1:n\n",
    "        A = randn(100,100,20)\n",
    "        m = maximum(A)\n",
    "        Am = mapslices(sum, A; dims=2)\n",
    "        B = A[:,:,5]\n",
    "        Bsort = mapslices(sort, B; dims=1)\n",
    "        b = rand(100)\n",
    "        C = B.*b\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8be8465c",
   "metadata": {},
   "source": [
    "### Profiling\n",
    "\n",
    "To profile this piece of code we will use Julia's builtin `Profile` package in combination with `ProfileView` as a grapical viewer. Some Julia editors (like VSCode) also have their own plugins to integrate with Julia's profiling capabilities, so worth to look out for this in your favourite editor!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b758fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Profile\n",
    "using ProfileView"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e032d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run once to compile everything ... this should be ignored\n",
    "ProfileView.@profview myfunction(1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a79dbdd2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ProfileView.@profview myfunction(10);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0adef51",
   "metadata": {},
   "source": [
    "**Note:** ProfileView does not always work so well with Jupyter Notebooks or Jupyterlab (but it's great from the REPL). An alternative is ProfileSVG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1577a951",
   "metadata": {},
   "outputs": [],
   "source": [
    "using ProfileSVG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea77234c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ProfileSVG.@profview myfunction(1)\n",
    "ProfileSVG.@profview myfunction(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d26af1a1",
   "metadata": {},
   "source": [
    "So how should one interpret this?\n",
    "- The horizontal direction is the time spent.\n",
    "- The vertical diretcion is the depth of the call stack.\n",
    "\n",
    "What do we learn:\n",
    "- The `mapslices` calls are clearly the most expensive parts of the function we should worry most\n",
    "- The first call (`mapslices(sum, A; dims=2)`) is more expensive as it works on more data than `mapslices(sort, B; dims=1)`\n",
    "- There is a stack of calls to functions in sort.jl on the right. This is because in Julia sorting is implemented recursively (sort functions call themselves)\n",
    "\n",
    "It is worth noting that red is a special colour in these graphs, highlighting a runtime dispatch, which can be an indicator for a type instability (more details and how to cure this in [Projects/03_Performance_Engineering.ipynb](Projects/03_Performance_Engineering.ipynb).\n",
    "\n",
    "For more details, take a look at the [ProfileView](https://github.com/timholy/ProfileView.jl) website."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5589da60",
   "metadata": {},
   "source": [
    "### High-level timings in TimerOutputs.jl\n",
    "\n",
    "[TimerOutputs.jl](https://github.com/KristofferC/TimerOutputs.jl) is great package to get a rough overview where time is spent. The idea is to annotate the code with simple tags, where timings are taken while the code is running. This is not for free, but if done at a high level cheap enough to be \"always on\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d218782",
   "metadata": {},
   "outputs": [],
   "source": [
    "using TimerOutputs\n",
    "\n",
    "const to = TimerOutput()\n",
    "function annotated_function(n)\n",
    "    @timeit to \"loop\" for i = 1:n\n",
    "        @timeit to \"initialisation\" A = randn(100,100,20)\n",
    "        m = maximum(A)\n",
    "        @timeit to \"mapslices on A\" Am = mapslices(sum, A; dims=2)\n",
    "        B = A[:,:,5]\n",
    "        @timeit to \"mapslices on B\" Bsort = mapslices(sort, B; dims=1)\n",
    "        b = rand(100)\n",
    "        C = B.*b\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b5db0b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "reset_timer!(to)\n",
    "annotated_function(10)\n",
    "to"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6050b458",
   "metadata": {},
   "source": [
    "##### More details\n",
    "- https://docs.julialang.org/en/v1/manual/profile/\n",
    "- https://github.com/kimikage/ProfileSVG.jl\n",
    "- https://github.com/JuliaCI/BenchmarkTools.jl\n",
    "- https://github.com/KristofferC/TimerOutputs.jl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc0c746f",
   "metadata": {},
   "source": [
    "## Takeaways\n",
    "\n",
    "- Julia has extensive builtin support for standard dense linear algebra (SVG, Diagonalisation, Linear systems)\n",
    "- Plenty of packages complement this by iterative methods\n",
    "- Profiling and timing can take place at various levels (From benchmarking individual statements to profiling a whole code)"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.7.2",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}