05_Specialisation_Speed.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speed by specialisation\n",
    "\n",
    "In the previous notebooks we discussed Julia's elaborate type system and how multiple dispatch allows functions to decide on the method to dispatch really from looking at all the arguments. Intuitively **more specialised code**, i.e. code that is allowed to exploit known structure, **is faster.** In this notebook we will explore how Julia is able to exploit this while retaining generic code ... thanks to multiple dispatch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The speed of Julia\n",
    "\n",
    "Now let us first try to address the question *Is Julia fast?*\n",
    "\n",
    "To some extend this is a bit of a mismatching question, since one is able to write slow code in any language ... so let's try do address something else instead: *Can Julia be fast?*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple example: Sums\n",
    "\n",
    "Let's compare for a simple example ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function mysum(v)\n",
    "    result = zero(eltype(v))\n",
    "    for i in 1:length(v)\n",
    "        result += v[i]\n",
    "    end\n",
    "    result\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = randn(10^7);   # Large vector of numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using BenchmarkTools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile some C code and call it from julia ...\n",
    "using Libdl\n",
    "code = \"\"\"\n",
    "#include <stddef.h>\n",
    "double c_sum(size_t n, double *v) {\n",
    "    double accu = 0.0;\n",
    "    for (size_t i = 0; i < n; ++i) {\n",
    "        accu += v[i];\n",
    "    }\n",
    "    return accu;\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "# Compile to a shared library (with fast maths and machine-specific)\n",
    "const Clib = tempname()\n",
    "open(`gcc -fPIC -O3 -march=native -xc -shared -ffast-math -o $(Clib * \".\" * Libdl.dlext) -`, \"w\") do f\n",
    "    print(f, code) \n",
    "end\n",
    "\n",
    "# define a Julia function that calls the C function:\n",
    "c_sum(v::Array{Float64}) = ccall((\"c_sum\", Clib), Float64, (Csize_t, Ptr{Float64}), length(v), v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench = @benchmark c_sum($v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = Dict()\n",
    "times[\"C (naive)\"] = minimum(bench.times) / 1e6\n",
    "times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An explicitly and fully vectorised C++ code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile some C code and call it from julia ...\n",
    "using Libdl\n",
    "code = \"\"\"\n",
    "#include <stddef.h>\n",
    "#include <immintrin.h>\n",
    "\n",
    "double c_sum_vector(size_t n, double *v)\n",
    "{\n",
    "    size_t i;\n",
    "    double result;\n",
    "    double tmp[4] __attribute__ ((aligned(64)));\n",
    "\n",
    "    __m256d sums1 = _mm256_setzero_pd();\n",
    "    __m256d sums2 = _mm256_setzero_pd();\n",
    "    for ( i = 0; i + 7 < n; i += 8 )\n",
    "    {\n",
    "        sums1 = _mm256_add_pd( sums1, _mm256_loadu_pd(v+i  ) );\n",
    "        sums2 = _mm256_add_pd( sums2, _mm256_loadu_pd(v+i+4) );\n",
    "    }\n",
    "    _mm256_store_pd( tmp, _mm256_add_pd(sums1, sums2) );\n",
    "\n",
    "    return tmp[0] + tmp[1] + tmp[2] + tmp[3]; \n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "# Compile to a shared library (with fast maths and machine-specific)\n",
    "const Clib = tempname()\n",
    "open(`gcc -fPIC -O2 -march=native -xc -shared -ffast-math -o $(Clib * \".\" * Libdl.dlext) -`, \"w\") do f\n",
    "    print(f, code) \n",
    "end\n",
    "\n",
    "# define a Julia function that calls the C function:\n",
    "c_sum(v::Array{Float64}) = ccall((\"c_sum_vector\", Clib), Float64, (Csize_t, Ptr{Float64}), length(v), v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench = @benchmark c_sum($v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times[\"C (vectorised)\"] = minimum(bench.times) / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How does our version do?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench = @benchmark mysum($v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times[\"Julia (naive)\"] = minimum(bench.times) / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But unlike C we have not yet tried all tricks!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function fastsum(v)\n",
    "    result = zero(eltype(v))\n",
    "    @simd for i in 1:length(v)    # @simd enforces vectorisation in the loop\n",
    "        @inbounds result += v[i]  # @inbounds suppresses bound checks\n",
    "    end\n",
    "    result\n",
    "end\n",
    "\n",
    "# Still nicely readable code ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench = @benchmark fastsum($v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times[\"Julia (simd)\"] = minimum(bench.times) / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare with python ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using PyCall\n",
    "numpysum(v) = pyimport(\"numpy\").sum(v)\n",
    "\n",
    "bench = @benchmark numpysum($v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times[\"Numpy\"] = minimum(bench.times) / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "py\"\"\"\n",
    "def py_sum(A):\n",
    "    s = 0.0\n",
    "    for a in A:\n",
    "        s += a\n",
    "    return s\n",
    "\"\"\"\n",
    "\n",
    "pysum = py\"py_sum\"\n",
    "\n",
    "bench = @benchmark pysum($v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times[\"Python (naive)\"] = minimum(bench.times) / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Overview ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A more complicated example: Vandermonde matrices\n",
    "(modified from [Steven's Julia intro](https://web.mit.edu/18.06/www/Fall17/1806/julia/Julia-intro.pdf))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{align}V=\\begin{bmatrix}1&\\alpha _{1}&\\alpha _{1}^{2}&\\dots &\\alpha _{1}^{n-1}\\\\1&\\alpha _{2}&\\alpha _{2}^{2}&\\dots &\\alpha _{2}^{n-1}\\\\1&\\alpha _{3}&\\alpha _{3}^{2}&\\dots &\\alpha _{3}^{n-1}\\\\\\vdots &\\vdots &\\vdots &\\ddots &\\vdots \\\\1&\\alpha _{m}&\\alpha _{m}^{2}&\\dots &\\alpha _{m}^{n-1}\\end{bmatrix}\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using PyCall\n",
    "np = pyimport(\"numpy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.vander(1:5, increasing=true)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[The source code for this function](https://github.com/numpy/numpy/blob/v1.16.1/numpy/lib/twodim_base.py#L475-L563) calls `np.multiply.accumulate` [which is implemented in C](https://github.com/numpy/numpy/blob/deea4983aedfa96905bbaee64e3d1de84144303f/numpy/core/src/umath/ufunc_object.c#L3678). However, this code doesn't actually perform the computation, it basically only checks types and stuff. The actual kernel is [implemented here](https://github.com/numpy/numpy/blob/deea4983aedfa96905bbaee64e3d1de84144303f/numpy/core/src/umath/loops.c.src#L1742), which is not even C code, but only a template that gets transformed to type-specific kernels. Still only a limited set of types `Float64`, `Float32`, and so forth are supported."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a type-generic Julia implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function vander(x::AbstractVector{T}) where T\n",
    "    m = length(x)\n",
    "    V = Matrix{T}(undef, m, m)\n",
    "    for j = 1:m\n",
    "        V[j,1] = one(x[j])\n",
    "    end\n",
    "    for i= 2:m\n",
    "        for j = 1:m\n",
    "            V[j,i] = x[j] * V[j,i-1]\n",
    "            end\n",
    "        end\n",
    "    return V\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vander(1:5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quick speed comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using BenchmarkTools\n",
    "using Plots\n",
    "ns = exp10.(range(1, 4, length=30));\n",
    "\n",
    "tnp = Float64[]\n",
    "tjl = Float64[]\n",
    "for n in ns\n",
    "    x = collect(1:n)\n",
    "    push!(tnp, @belapsed np.vander($x) samples=3 evals=1)\n",
    "    push!(tjl, @belapsed vander($x)    samples=3 evals=1)\n",
    "end\n",
    "plot(ns, tnp./tjl, m=:circle, xscale=:log10, yscale=:log10, ylims=[1, Inf],\n",
    "     xlab=\"matrix size\", ylab=\"NumPy time / Julia time\", legend=:false)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notably the clean and concise Julia implementation is **faster than numpy**'s C implementation for small matrices and **as fast** for large matrix sizes. Still it works for **arbitrary types**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vander(Int32[4, 8, 16, 32])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This includes non-numerical types ... the only assumption is that the type induces a multiplicative group, i.e. has a `one` function to yield the identity element and an apropriate `*` defined.\n",
    "\n",
    "A rather unusual one is `String`, which works since `one(String) == \"\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vander([\"this\", \"is\", \"a\", \"test\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How does this work?\n",
    "\n",
    "<img src=\"img/from_source_to_native.png\" />\n",
    "\n",
    " \n",
    "- AST = Abstract Syntax Tree\n",
    "- SSA = Static Single Assignment\n",
    "- [LLVM](https://de.wikipedia.org/wiki/LLVM) = Low Level Virtual Machine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Julia compiles **ahead of time**, i.e. Julia waits until the first function call for a particular set of input types is issued and then compiles the full call stack all the way down to machine code specialising as much as possible along the way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = [1, 3, 5]  # Vector{Int64}\n",
    "\n",
    "@time mysum(x)\n",
    "@time mysum(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we saw if thes change, Julia compiles a new specialization of the function!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see what happens under the hood, we have a bunch of inspection macros:\n",
    "* The AST after parsing (**`@macroexpand`**)\n",
    "* The AST after lowering (**`@code_typed`**, **`@code_warntype`**)\n",
    "* The AST after type inference and optimization (**`@code_lowered`**)\n",
    "* The LLVM IR (**`@code_llvm`**)\n",
    "* The assembly machine code (**`@code_native`**)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_typed mysum(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_lowered mysum(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_llvm mysum(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can remove the comments (lines starting with `;` using `debuginfo=:none`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_llvm debuginfo=:none mysum(1.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_native debuginfo=:none mysum(1.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare this to `Float64` input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_native debuginfo=:none mysum(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Types and code specialisation\n",
    "\n",
    "As a rule of thumb: The Julia compiler only uses the types to specialise, not the values. Therefore using special types that match the possible assumptions most closely is crucial.\n",
    "\n",
    "As an example we consider the determinant of a 2x2 matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fast custom implementation ... note: No type annotation\n",
    "det_custom(M) = @inbounds (M[1, 1] * M[2, 2] - M[1, 2] * M[2, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M = randn(2, 2)\n",
    "@btime det_custom(M);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra\n",
    "@btime det(M);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that is a drastic difference ... but we actually did not tell the compiler anything about the size of `M`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "typeof(M)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's help along ... by using a static array type:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using StaticArrays\n",
    "\n",
    "S = @SMatrix randn(2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@btime det(S);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Got it! Now ... what happens under the hood:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_typed debuginfo=:none det(S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "in other words\n",
    "```\n",
    "%2  = S[1, 1]\n",
    "%4  = S[2, 2]\n",
    "%5  = S[1, 1] * S[2, 2]\n",
    "%7  = S[2, 1]\n",
    "%9  = S[1, 2]\n",
    "%10 = S[2, 1] * S[1, 2]\n",
    "%11 = S[1, 1] * S[2, 2] - S[2, 1] * S[1, 2]\n",
    "```\n",
    "exactly what we had hand-optimised before!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now one should add here that part of the magic has not been the compiler being smart, but much rather the `StaticArrays` package, which has a number of carfully optimised implementations for standard operations:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "============================================\n",
    "    Benchmarks for 3×3 Float64 matrices\n",
    "============================================\n",
    "Matrix multiplication               -> 5.9x speedup\n",
    "Matrix multiplication (mutating)    -> 1.8x speedup\n",
    "Matrix addition                     -> 33.1x speedup\n",
    "Matrix addition (mutating)          -> 2.5x speedup\n",
    "Matrix determinant                  -> 112.9x speedup\n",
    "Matrix inverse                      -> 67.8x speedup\n",
    "Matrix symmetric eigendecomposition -> 25.0x speedup\n",
    "Matrix Cholesky decomposition       -> 8.8x speedup\n",
    "Matrix LU decomposition             -> 6.1x speedup\n",
    "Matrix QR decomposition             -> 65.0x speedup\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### More details\n",
    "- https://github.com/JuliaCI/BenchmarkTools.jl\n",
    "- https://github.com/JuliaArrays/StaticArrays.jl/\n",
    "- https://docs.julialang.org/en/v1/devdocs/compiler/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Takeaways\n",
    "\n",
    "- Julia can be fast\n",
    "- Code is compiled down to the metal\n",
    "- Stages can be inspected using macros like `@code_warntype`\n",
    "- Typing is crucial to exploit specialised structure\n",
    "- Type annotations are irrelevant for performance"
   ]
  }
 ],
 "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": 4
}