euroscipy_pycon2019.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Make your numerical Python code fly at transonic 🚀 speed! \n",
    "\n",
    "# Transonic https://transonic.readthedocs.io\n",
    "\n",
    "- A new package\n",
    "\n",
    "- A unified modern Python API for Python/Numpy accelerators\n",
    "\n",
    "- A thin layer between developers and Pythran / Cython / Numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div align=\"middle\"> EuroScipy 2019 (4 September 2019, Bilbao) </div>\n",
    "\n",
    "<div align=\"middle\">\n",
    "    <a href=\"http://tiny.cc/euroscipy2019-transonic\">http://tiny.cc/euroscipy2019-transonic</a>    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Few Transonic code examples\n",
    "\n",
    "These codes can be accelerated with Pythran, Cython and Numba."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Ahead-of-time compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from transonic import boost\n",
    "\n",
    "T0 = \"int[:, :]\"\n",
    "T1 = \"int[:]\"\n",
    "\n",
    "@boost\n",
    "def row_sum(arr: T0, columns: T1):\n",
    "    return arr.T[columns].sum(0)\n",
    "\n",
    "@boost(boundscheck=False, wraparound=False)\n",
    "def row_sum_loops(arr: T0, columns: T1):\n",
    "    # locals type annotations are used only for Cython\n",
    "    i: int\n",
    "    j: int\n",
    "    sum_: int\n",
    "    res: \"int[]\" = np.empty(arr.shape[0], dtype=arr.dtype)\n",
    "    for i in range(arr.shape[0]):\n",
    "        sum_ = 0\n",
    "        for j in range(columns.shape[0]):\n",
    "            sum_ += arr[i, columns[j]]\n",
    "        res[i] = sum_\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Just-in-time compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from transonic import jit\n",
    "\n",
    "def add(a, b):\n",
    "    return a + b\n",
    "\n",
    "@jit\n",
    "def func(a, b):\n",
    "    return np.exp(a) * b * add(a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Time annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from transonic import Type, NDim, Array, boost\n",
    "\n",
    "T = Type(int, float, np.complex128)\n",
    "N = NDim(1, 2, 3)\n",
    "\n",
    "A = Array[T, N]\n",
    "A1 = Array[np.float32, N + 1]\n",
    "\n",
    "@boost\n",
    "def compute(a: A, b: A, c: T, d: A1):\n",
    "    ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### `inline` functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transonic import boost\n",
    "\n",
    "T = int\n",
    "\n",
    "@boost(inline=True)\n",
    "def add(a: T, b: T) -> T:\n",
    "    return a + b\n",
    "\n",
    "@boost\n",
    "def use_add(n: int = 10000):\n",
    "    _: int\n",
    "    for _ in range(n):\n",
    "        tmp = add(tmp, 1)\n",
    "    return tmp\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Accelerate methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transonic import boost\n",
    "\n",
    "@boost\n",
    "class MyClass:\n",
    "    attr: int\n",
    "    \n",
    "    @boost\n",
    "    def numerical_kernel(self, arg: int):\n",
    "        return self.attr + arg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Proper compilation needed for high efficiency !\n",
    "\n",
    "Not like CPython: `compile(...)` to (high level) virtual machine instructions (with nearly no optimization)\n",
    "\n",
    "## Compilation to machine instructions\n",
    "\n",
    "*One needs to write code that can be well optimized by a compiler!*\n",
    "\n",
    "- Just-in-time (`@jit`)\n",
    "\n",
    "  Has to be fast (warm up), can be hardware specific\n",
    "\n",
    "- Ahead-of-time (`@boost`)\n",
    "\n",
    "  Can be slow, hardware specific or more general to distribute binaries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "## First step for Python: Transpilation\n",
    "\n",
    "From one language to another language (for example Python to C++, or Cython to C)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Many tools to compile Python"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "<div align=\"middle\">\n",
    "<img src=\"./tag_clouds.png\"  style=\"width: 60%\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Compiled at which levels?\n",
    "\n",
    "- programs (Nuitka)\n",
    "\n",
    "- slowest loops (PyPy)\n",
    "\n",
    "- modules (Cython, Pythran)\n",
    "\n",
    "- user-defined functions / methods (Numba, Transonic)\n",
    "\n",
    "- blocks of code (Transonic)\n",
    "\n",
    "- expressions (Numexp)\n",
    "\n",
    "- call compiled functions (Numpy / Python)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Cython (module-level) AOT complier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "- Langage: superset of Python\n",
    "\n",
    "- A great mix of Python / C / CPython C API! \n",
    "\n",
    "  Very powerfull but a tool for experts!\n",
    "\n",
    "- Easy to study where the interpreter is used (`cython --annotate`).\n",
    "\n",
    "- Very mature\n",
    "\n",
    "- Very efficient for C-like code (explicit loops, \"low level\")\n",
    "\n",
    "- Now able to use Pythran internally..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "My experience: large Cython extensions difficult to maintain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Numba: (per-method) JIT for Python-Numpy code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "- Very simple to use (just add few decorators) 🙂"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit\n",
    "\n",
    "@jit\n",
    "def myfunc(x):\n",
    "    return x**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "- \"nopython\" mode (fast and no GIL) 🙂\n",
    "\n",
    "- Also a \"python\" mode 🙂\n",
    "\n",
    "- GPU and Cupy 😀\n",
    "\n",
    "- Methods (of classes) 🙂\n",
    "\n",
    "- Only JIT 🙁\n",
    "    - Sometimes not as much efficient as it could be 🙁\n",
    "      (sometimes slower than Pythran / Julia / C++)\n",
    "\n",
    "- Not good to optimize high-level NumPy code 🙁"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Pythran: AOT compiler for modules using Python-Numpy\n",
    "\n",
    "Transpiles Python to efficient C++"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "- Good to optimize *high-level NumPy code* 😎\n",
    "\n",
    "- Extensions never use the Python interpreter (pure C++ ⇒ no GIL) 🙂\n",
    "\n",
    "- Can produce C++ that can be used without Python\n",
    "\n",
    "- Usually **very efficient** (sometimes faster than Julia)\n",
    "\n",
    "    - High and low level optimizations\n",
    "    \n",
    "      (Python optimizations and C++ compilation)\n",
    "\n",
    "    - SIMD 🤩 (with [xsimd](https://github.com/QuantStack/xsimd)) \n",
    "\n",
    "    - Understand OpenMP instructions 🤗 !\n",
    "\n",
    "- Can [use and make PyCapsules](https://serge-sans-paille.github.io/pythran-stories/the-capsule-corporation.html) (functions operating in the native word) 🙂"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "### High level transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "from black import format_str, FileMode\n",
    "from pythran.toolchain import generate_py\n",
    "import gast as ast\n",
    "import astunparse\n",
    "\n",
    "def print_optimized(src):    \n",
    "    optimized_py = generate_py(\"bar\", src)\n",
    "    tree = ast.parse(optimized_py)\n",
    "    for node in tree.body:\n",
    "        if isinstance(node, ast.FunctionDef):\n",
    "            fdef = node\n",
    "            fdef.body = [node for node in fdef.body[:-1] if not isinstance(node, ast.Pass)] + [fdef.body[-1]]\n",
    "    optimized_code = astunparse.unparse(tree)\n",
    "    print(format_str(optimized_code, mode=FileMode(line_length=82)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# calcul of range\n",
    "print_optimized(\"\"\"\n",
    "def f(x):\n",
    "    y = 1 if x else 2\n",
    "    return y == 3\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# inlining\n",
    "print_optimized(\"\"\"\n",
    "def foo(a):\n",
    "    return  a + 1\n",
    "def bar(b, c):\n",
    "    return foo(b), foo(2 * c)\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# unroll loops\n",
    "print_optimized(\"\"\"\n",
    "def foo():\n",
    "    ret = 0\n",
    "    for i in range(1, 3):\n",
    "        for j in range(1, 4):\n",
    "            ret += i * j\n",
    "    return ret\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# constant propagation\n",
    "print_optimized(\"\"\"\n",
    "def fib(n):\n",
    "    return n if n< 2 else fib(n-1) + fib(n-2)\n",
    "    \n",
    "def bar(): \n",
    "    return [fib(i) for i in [1, 2, 8, 20]]\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# advanced transformations\n",
    "print_optimized(\"\"\"\n",
    "import numpy as np\n",
    "def wsum(v, w, x, y, z):\n",
    "    return sum(np.array([v, w, x, y, z]) * (.1, .2, .3, .2, .1))\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Pythran: AOT compiler for module using Python-Numpy\n",
    "\n",
    "- Compile only full modules (⇒ refactoring needed 🙁)\n",
    "\n",
    "- Only \"nopython\" mode\n",
    "\n",
    "    * limited to a subset of Python\n",
    "    \n",
    "        - only homogeneous list / dict 🤷‍♀️\n",
    "        - no methods (of classes) 😢 and user-defined class\n",
    "    \n",
    "    * limited to few extension packages (Numpy + bits of Scipy)\n",
    "    \n",
    "    * pythranized functions can't call Python functions\n",
    "\n",
    "- No JIT: need types (written manually in comments)\n",
    "\n",
    "- Lengthy ⌛️ and memory intensive compilations\n",
    "\n",
    "- Debugging 🐜 Pythran requires C++ skills\n",
    "\n",
    "- No GPU (maybe with [OpenMP 4](https://www.openmp.org/updates/openmp-accelerator-support-gpus/)?)\n",
    "\n",
    "- Intel compilers unable to compile Pythran C++11 👎\n",
    "\n",
    "- Small community, only 1 core-dev"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# First conclusions\n",
    "\n",
    "- Python great language & ecosystem for sciences & data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Transonic example: a unique code accelerated with 3 backends\n",
    "\n",
    "### Benchmark ahead-of-time compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from transonic import boost\n",
    "\n",
    "T0 = \"int[:, :]\"\n",
    "T1 = \"int[:]\"\n",
    "\n",
    "@boost\n",
    "def row_sum(arr: T0, columns: T1):\n",
    "    return arr.T[columns].sum(0)\n",
    "\n",
    "@boost\n",
    "def row_sum_loops(arr: T0, columns: T1):\n",
    "    # locals type annotations are used only for Cython\n",
    "    i: int\n",
    "    j: int\n",
    "    sum_: int\n",
    "    res: \"int[]\" = np.empty(arr.shape[0], dtype=arr.dtype)\n",
    "    for i in range(arr.shape[0]):\n",
    "        sum_ = 0\n",
    "        for j in range(columns.shape[0]):\n",
    "            sum_ += arr[i, columns[j]]\n",
    "        res[i] = sum_\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Transonic example: a unique code accelerated with 3 backends\n",
    "\n",
    "### Benchmark ahead-of-time compilation\n",
    "\n",
    "```\n",
    "TRANSONIC_BACKEND=\"python\" python row_sum_boost.py\n",
    "Python\n",
    "row_sum              1.38 s\n",
    "row_sum_loops        108.57 s\n",
    "\n",
    "TRANSONIC_BACKEND=\"cython\" python row_sum_boost.py\n",
    "Cython\n",
    "row_sum              1.32 s\n",
    "row_sum_loops        0.38 s\n",
    "\n",
    "```\n",
    "\n",
    "```\n",
    "TRANSONIC_BACKEND=\"numba\" python row_sum_boost.py\n",
    "Numba\n",
    "row_sum              1.16 s\n",
    "row_sum_loops        0.27 s\n",
    "\n",
    "TRANSONIC_BACKEND=\"pythran\" python row_sum_boost.py\n",
    "Pythran\n",
    "row_sum              0.76 s\n",
    "row_sum_loops        0.27 s\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Transonic example: a unique code accelerated with 3 backends\n",
    "\n",
    "### Benchmark Just-in-time compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from transonic import jit\n",
    "\n",
    "@jit(native=True, xsimd=True)\n",
    "def row_sum(arr, columns):\n",
    "    return arr.T[columns].sum(0)\n",
    "\n",
    "@jit(native=True, xsimd=True)\n",
    "def row_sum_loops(arr, columns):\n",
    "    res = np.empty(arr.shape[0], dtype=arr.dtype)\n",
    "    for i in range(arr.shape[0]):\n",
    "        sum_ = 0\n",
    "        for j in range(columns.shape[0]):\n",
    "            sum_ += arr[i, columns[j]]\n",
    "        res[i] = sum_\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Transonic example: a unique code accelerated with 3 backends\n",
    "\n",
    "### Benchmark Just-in-time compilation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "TRANSONIC_BACKEND=\"cython\" python row_sum_jit.py\n",
    "Cython\n",
    "row_sum              1.28 s\n",
    "row_sum_loops        11.94 s\n",
    "\n",
    "TRANSONIC_BACKEND=\"numba\" python row_sum_jit.py\n",
    "Numba\n",
    "row_sum              1.14 s\n",
    "row_sum_loops        0.28 s\n",
    "\n",
    "```\n",
    "\n",
    "```\n",
    "TRANSONIC_BACKEND=\"pythran\" python row_sum_jit.py\n",
    "Pythran\n",
    "row_sum              0.76 s\n",
    "row_sum_loops        0.28 s\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## NotImplemented\n",
    "\n",
    "### dataclass (`numba.jitclass` and Cython `cdef class`) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transonic import dataclass, boost\n",
    "\n",
    "@dataclass\n",
    "class MyStruct:\n",
    "    attr: int\n",
    "\n",
    "    def compute(self, arg: int):\n",
    "        return self.attr + arg\n",
    "\n",
    "    def modify(self, arg: int):\n",
    "        self.attr = arg\n",
    "\n",
    "@boost\n",
    "def func(o: MyStruct, a: int):\n",
    "    o.modify(a)\n",
    "    return o\n",
    "\n",
    "@boost\n",
    "def func1(o: MyStruct, a: int):\n",
    "    return o.compute(a)"
   ]
  }
 ],
 "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}