{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numba 0.53.0 Release Demo\n",
    "\n",
    "This notebook contains a demonstration of new features present in the 0.53.0 release of Numba. Whilst release notes are produced as part of the [CHANGE_LOG](https://github.com/numba/numba/blob/release0.53/CHANGE_LOG), there's nothing like seeing code in action!\n",
    "\n",
    "This release contains a few new features. In this notebook, the new CPU target features are demonstrated. The [CUDA target](https://numba.pydata.org/numba-doc/latest/cuda/index.html) also gained a lot of new features in 0.53.0 and [@gmarkall](https://github.com/gmarkall) has created a [demo notebook](https://mybinder.org/v2/gh/numba/numba-examples/master?filepath=notebooks%2FNumba_053_CUDA_Release_Demo.ipynb) especially for these!\n",
    "\n",
    "Key internal changes:\n",
    "\n",
    "- Support for Python 3.9 ([@stuartarchibald](https://github.com/stuartarchibald)).\n",
    "- Function sub-typing ([@luk-f-a](https://github.com/luk-f-a)).\n",
    "- Initial support for dynamic gufuncs (i.e. from `@guvectorize`) ([@guilhermeleobas](https://github.com/guilhermeleobas)).\n",
    "- Parallel Accelerator (`@njit(parallel=True)`) now supports Fortran ordered arrays ([@DrTodd13](https://github.com/DrTodd13) and [@sklam](https://github.com/sklam)).\n",
    "\n",
    "Intel also kindly sponsored research and development that lead to two new features for profiling the compiler:\n",
    "\n",
    "- Exposing LLVM compilation pass timings for diagnostic purposes ([@sklam](https://github.com/sklam)).\n",
    "- An event system for broadcasting compiler events ([@sklam](https://github.com/sklam)).\n",
    "\n",
    "Demonstrations of these compiler profiling features can be found in a [different notebook](https://mybinder.org/v2/gh/numba/numba-examples/master?filepath=notebooks%2FNumba_053_profiling_the_compiler.ipynb).\n",
    "\n",
    "\n",
    "This notebook will focus on:\n",
    "\n",
    "- [Function sub-typing](#Function-sub-typing)\n",
    "- [Dynamic gufunc](#Dynamic-GUFuncs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function sub-typing\n",
    "\n",
    "[@luk-f-a](https://github.com/luk-f-a) added function subtyping support to allow efficient passing of functions as arguments. It allows a function to be converted to its subtype. In detail, a jit function is a generic function. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "assert numba.version_info.short >= (0, 53)\n",
    "from numba.core import types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.njit\n",
    "def identity(x):\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`identity()` is a generic function that accepts any type. \n",
    "\n",
    "Let's define a 1-arity function type that takes and returns a `intp`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionType[int64(int64)]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define a function type that takes an int and returns an int.\n",
    "fn_sig = types.FunctionType(types.intp(types.intp))\n",
    "fn_sig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`fn_sig` is a subtype of `identity()`. Therefore, we can use `identity()` in any place that expects a `fn_sig`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a function that takes `(fn_sig, intp)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.njit((fn_sig, types.intp))\n",
    "def invoke_callback(callback, arg):\n",
    "    return callback(arg)\n",
    "\n",
    "# Disable compilation\n",
    "invoke_callback.disable_compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(FunctionType[int64(int64)], int64)]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "invoke_callback.signatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use `identity` for the first argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "123"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "invoke_callback(identity, 123)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `invoke_callback()` function can take any function that can be cast to the subtype `fn_sig`. We will define two more functions that are compatible with `fn_sig`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cb_add_one 124\n",
      "cb_twice 246\n"
     ]
    }
   ],
   "source": [
    "@numba.njit\n",
    "def cb_add_one(x):\n",
    "    return x + 1\n",
    "\n",
    "@numba.njit\n",
    "def cb_twice(x):\n",
    "    return x * 2\n",
    "\n",
    "print('cb_add_one', invoke_callback(cb_add_one, 123))\n",
    "print('cb_twice', invoke_callback(cb_twice, 123))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No new signature needs to be compiled:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(FunctionType[int64(int64)], int64)]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "invoke_callback.signatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Without declaring the expected function type, Numba would specialize to the given function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# of compiled version 0\n",
      "call identity() = 123\n",
      "# of compiled version 1\n",
      "call cb_add_one() = 124\n",
      "# of compiled version 2\n",
      "signatures [(type(CPUDispatcher(<function identity at 0x11dce0b80>)), int64), (type(CPUDispatcher(<function cb_add_one at 0x11f048310>)), int64)]\n"
     ]
    }
   ],
   "source": [
    "@numba.njit\n",
    "def invoke_callback_generic(callback, arg):\n",
    "    return callback(arg)\n",
    "\n",
    "\n",
    "print(f\"# of compiled version {len(invoke_callback_generic.signatures)}\")\n",
    "print(\"call identity() =\", invoke_callback_generic(identity, 123))\n",
    "print(f\"# of compiled version {len(invoke_callback_generic.signatures)}\")\n",
    "print(\"call cb_add_one() =\", invoke_callback_generic(cb_add_one, 123))\n",
    "print(f\"# of compiled version {len(invoke_callback_generic.signatures)}\")\n",
    "print(\"signatures\", invoke_callback_generic.signatures)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One other advantage to function subtyping is that first-class functions are not locked as previously required. We can request new subtypes from the functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's look at the compiled versions of `cb_add_one()` (there's just one):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(int64,)]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cb_add_one.signatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's make a list that contains functions of the signature `float64(float64)` and add our callback functions to it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_fn_sig = types.FunctionType(types.float64(types.float64))\n",
    "lst = numba.typed.List.empty_list(new_fn_sig)\n",
    "lst.append(cb_add_one)\n",
    "lst.append(cb_twice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A new float version of `cb_add_one()` is compiled as requested during the `.append()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(int64,), (float64,)]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cb_add_one.signatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use the list of functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[11.234, 10.468]\n",
      "[array([1.234, 2.234, 3.234, 4.234]), array([0.468, 1.468, 2.468, 3.468])]\n"
     ]
    }
   ],
   "source": [
    "@numba.njit\n",
    "def many_callbacks(cblist, arg, incr):\n",
    "    return [cb(arg) + incr for cb in cblist]\n",
    "\n",
    "print(many_callbacks(lst, 0.234, 10))\n",
    "print(many_callbacks(lst, 0.234, np.arange(4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamic GUFuncs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[@guilhermeleobas](https://github.com/guilhermeleobas) added dynamic [gufunc](https://numba.readthedocs.io/en/stable/user/vectorize.html#dynamic-generalized-universal-functions) support, allowing gufuncs to be used without pre-defining their accepted function types."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Previously, gufuncs must be defined with a fix set of function types. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float64 version [0. 2. 4. 6. 8.]\n",
      "complex128 version [0.+1.j 2.+1.j 4.+1.j 6.+1.j 8.+1.j]\n"
     ]
    }
   ],
   "source": [
    "@numba.guvectorize([(types.float64, types.float64[:]), \n",
    "                    (types.complex128, types.complex128[:])], \"()->()\")\n",
    "def static_twice(inp, out):\n",
    "    out[()] = inp * 2\n",
    "\n",
    "    \n",
    "inp = np.arange(5, dtype=np.float64)\n",
    "print(f\"{inp.dtype} version\", static_twice(inp))\n",
    "inp2 = inp + .5j\n",
    "print(f\"{inp2.dtype} version\", static_twice(inp2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With dynamic gufuncs, we can omit the function types. New compilation is triggered dynamically as needed. However, due to a limitation of type inference, the output argument must be specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float64 version [0. 2. 4. 6. 8.]\n",
      "complex128 version [0.+1.j 2.+1.j 4.+1.j 6.+1.j 8.+1.j]\n",
      "int64 version [0 2 4 6 8]\n"
     ]
    }
   ],
   "source": [
    "# creates a dynamic gufunc by omitting the function types\n",
    "@numba.guvectorize(\"()->()\")\n",
    "def dynamic_twice(inp, out):\n",
    "    out[()] = inp * 2\n",
    "\n",
    "\n",
    "# Use the gufunc with different dtypes\n",
    "out1 = np.zeros_like(inp)\n",
    "dynamic_twice(inp, out1)\n",
    "print(f\"{inp.dtype} version\", out1)\n",
    "\n",
    "inp2 = inp + 0.5j\n",
    "out2 = np.zeros_like(inp2)\n",
    "dynamic_twice(inp2, out2)\n",
    "\n",
    "print(f\"{inp2.dtype} version\", out2)\n",
    "\n",
    "\n",
    "inp3 = inp.astype(np.intp)\n",
    "out3 = np.zeros_like(inp3)\n",
    "dynamic_twice(inp3, out3)\n",
    "print(f\"{inp3.dtype} version\", out3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "\n",
    "## See Also\n",
    "\n",
    "- [0.53 CUDA Demo](https://mybinder.org/v2/gh/numba/numba-examples/master?filepath=notebooks%2FNumba_053_CUDA_Release_Demo.ipynb)\n",
    "- [0.53 Compiler Profiling Demo](https://mybinder.org/v2/gh/numba/numba-examples/master?filepath=notebooks%2FNumba_053_profiling_the_compiler.ipynb)"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}