Let's Write an LLVM Specializer for Python! (Stephen Diehl).ipynb

{
 "metadata": {
  "name": "Let's Write an LLVM Specializer for Python! (Stephen Diehl)"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<center>\n",
      "<h1>Let's Write an LLVM Specializer for Python!</h1>\n",
      "</center>\n",
      "<br/>\n",
      "<center>\n",
      "**Stephen Diehl** ([@smdiehl](https://twitter.com/smdiehl))\n",
      "</center>\n",
      "\n",
      "-------\n"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<img src=\"files/images/logo_med.png\">\n",
      "\n",
      "This is the followup to my talk [**LLVM Optimized Python**](http://dev.stephendiehl.com/cfa/slides.pdf) at the **Harvard-Smithsonian Center for Astrophysics**, we'll do the deep dive that I didn't have time for. We're going to build a single module Numba-like compiler for Python. It won't be nearly as featureful or complete, but should demonstrate how you can go about building your own little LLVM specializer for a subset of Python or your own custom DSL expression compiler; and integrating it with the standard NumPy/SciPy stack for whatever scientific computing domain you work. The full source for this project [is available on Github](https://github.com/sdiehl/numpile) and comes in at 1000 lines for the whole specializer, very tiny!\n",
      "\n",
      "There's a whole slew of interesting domains where this kind of on-the-fly specializing compiler can be used:\n",
      "\n",
      "* Computation kernels for MapReduce\n",
      "* Financial backtesting\n",
      "* Dense linear algebra\n",
      "* Image processing\n",
      "* Data pipeline hotspots\n",
      "* [Speeding up SQL queries](http://tuulos.github.io/sf-python-meetup-sep-2013/#/22) in PostgreSQL\n",
      "* [Molecular dynamics](http://combichem.blogspot.com/2013/04/fun-with-numba-numpy-and-f2py.html)\n",
      "* Compiling UDFs for [Cloudera Impala](http://www.slideshare.net/urilaserson/numbacompiled-python-udfs-for-impala-impala-meetup-52014?related=1)\n",
      "\n",
      "Python is great for rapid development and high-level thinking, but is slow due to too many level of indirection, hashmap lookups, broken parallelism,slow garbage collector, and boxed PyObject types. With LLVM we can keep writing high-level code and not sacrifice performance.\n",
      "\n",
      "You will need ``python``, ``llvm``, ``llvmpy``, ``numpy`` and a bit of time. The best way to get all of these is to install [Anaconda](https://store.continuum.io/cshop/anaconda/) maintained by my good friend Ilan. Don't add any more entropy to the universe by compiling NumPy from source, just use Anaconda.\n"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import llvm.ee as le\n",
      "import llvm.core as lc\n",
      "\n",
      "int_type   = lc.Type.int()\n",
      "float_type = lc.Type.double()\n",
      "void_type  = lc.Type.void()\n",
      "\n",
      "def func(name, module, rettype, argtypes):\n",
      "    func_type   = lc.Type.function(rettype, argtypes, False)\n",
      "    lfunc       = lc.Function.new(module, func_type, name)\n",
      "    entry_block = lfunc.append_basic_block(\"entry\")\n",
      "    builder     = lc.Builder.new(entry_block)\n",
      "    return (lfunc, builder)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We'll create the toplevel LLVM module which will hold all of our definitions. When we call Python's ``__repr__`` function it will print out the LLVM IR to the module."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "mod = lc.Module.new('mymodule')\n",
      "print(mod)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "; ModuleID = 'mymodule'\n",
        "\n"
       ]
      }
     ],
     "prompt_number": 3
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We now create the builder function which we'll use to populate the basic block structure of the module. Again, when we call Python's __repr__ it will print out the LLVM IR for the function definition this time."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "(fn, builder) = func('main', mod, int_type, [])\n",
      "print(fn)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "define i32 @main() {\n",
        "entry:\n",
        "}\n",
        "\n"
       ]
      }
     ],
     "prompt_number": 4
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "And we just create a constant integer and use the builder to emit a `ret` instruction to return the current `entry`  basic block and yield the constant value."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "value = lc.Constant.int(int_type, 42)\n",
      "block = builder.ret(value)\n",
      "print(mod)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "; ModuleID = 'mymodule'\n",
        "\n",
        "define i32 @main() {\n",
        "entry:\n",
        "  ret i32 42\n",
        "}\n",
        "\n"
       ]
      }
     ],
     "prompt_number": 5
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print(mod.to_native_assembly())"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\t.file\t\"mymodule\"\n",
        "\t.text\n",
        "\t.globl\tmain\n",
        "\t.align\t16, 0x90\n",
        "\t.type\tmain,@function\n",
        "main:\n",
        "\t.cfi_startproc\n",
        "\tmovl\t$42, %eax\n",
        "\tret\n",
        ".Ltmp0:\n",
        "\t.size\tmain, .Ltmp0-main\n",
        "\t.cfi_endproc\n",
        "\n",
        "\n",
        "\t.section\t\".note.GNU-stack\",\"\",@progbits\n",
        "\n"
       ]
      }
     ],
     "prompt_number": 6
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "So that's pretty swell, we have a way to interactively generate machine code at runtime in Python! Now we'll use the LLVM JIT to actually actually execute the code and interchange values between the CPython runtime and the LLVM JIT."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "tm = le.TargetMachine.new(features='', cm=le.CM_JITDEFAULT)\n",
      "eb = le.EngineBuilder.new(mod)\n",
      "jit = eb.create(tm)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 7
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "ret = jit.run_function(fn, [])\n",
      "print(ret.as_int())"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "42\n"
       ]
      }
     ],
     "prompt_number": 8
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**That's pretty cool!** We've just created machine code on the fly and executed it inside of LLVM JIT *inside* CPython yielding value that we can work with in regular python. If the wheels aren't spinning in your head about what you can do with this awesome power, then they should be!\n",
      "\n",
      "So now let's set about building a tiny pipeline to support a custom ``autojit`` decorator. At the end we should be able to specialize the following dot product into efficient machine code."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "@autojit\n",
      "def dot(a, b):\n",
      "    c = 0\n",
      "    n = a.shape[0]\n",
      "    for i in range(n):\n",
      "       c += a[i]*b[i]\n",
      "    return c"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 9
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# LLVM Primer\n",
      "\n",
      "<img src=\"files/images/DragonMedium.png\">\n",
      "\n",
      "LLVM is the engine that drives our effort. It is a modern compiler framework and intermediate representation language together with toolchain for manipulating and optimizing this language.\n",
      "\n",
      "**Basic Types**\n",
      "\n",
      "LLVM types are your typical machine types plus pointers, structs, vectors and arrays.\n",
      "\n",
      "```\n",
      "i1 1                   ; boolean bit\n",
      "i32 299792458          ; integer\n",
      "float 7.29735257e-3    ; single precision\n",
      "double 6.62606957e-34  ; double precision\n",
      "```\n",
      "\n",
      "```\n",
      "{float, i64}           ; structure\n",
      "{float, {double, i3}}  ; nested structure\n",
      "<{float, [2 x i3]}>    ; packed structure\n",
      "```\n",
      "\n",
      "```\n",
      "[10 x float]           ; Array of 10 floats\n",
      "[10 x [20 x i32]]      ; Array of 10 arrays of 20 integers.\n",
      "```\n",
      "    \n",
      "``` \n",
      "<8 x float>            ; Vector of width 8 of floats\n",
      "```\n",
      "    \n",
      "```\n",
      "float*                 ; Pointer to a float\n",
      "[25 x float]*          ; Pointer to an array\n",
      "```\n",
      "\n",
      "**Instructions**\n",
      "\n",
      "All instructions are assignment to a unique virtual register. In SSA (Single Static Assignment) a register is never assigned to more than once.\n",
      "\n",
      "```\n",
      "%result = add i32 10, 20\n",
      "```\n",
      "\n",
      "Symbols used in an LLVM module are either global or local. Global symbols begin\n",
      "with ``@`` and local symbols begin with ``%``.\n",
      "\n",
      "The numerical instructions are:\n",
      "\n",
      "* ``add``  : Integer addition \n",
      "* ``fadd`` : Floating point addition\n",
      "* ``sub``  : Integer subtraction\n",
      "* ``fsub`` : Floating point subtraction\n",
      "* ``mul``  : Integer multiplication \n",
      "* ``fmul`` : Floating point multiplication\n",
      "* ``udiv`` : Unsigned integer quotient\n",
      "* ``sdiv`` : Signed integer quotient\n",
      "* ``fdiv`` : Floating point quotient\n",
      "* ``urem`` : Unsigned integer remainder \n",
      "* ``srem`` : Signed integer remainder \n",
      "* ``frem`` : Floating point integer remainder \n",
      "    \n",
      "**Memory**\n",
      "    \n",
      "LLVM uses the traditional load/store model:\n",
      "\n",
      "* ``load``: Load a typed value from a given reference\n",
      "* ``store``: Store a typed value in a given reference\n",
      "* ``alloca``: Allocate a pointer to memory on the virtual stack\n",
      "\n",
      "```\n",
      "%ptr = alloca i32\n",
      "store i32 3, i32* %ptr\n",
      "%val = load i32* %ptr\n",
      "```\n",
      "\n",
      "**Functions**\n",
      "    \n",
      "Functions are defined by as a collection of basic blocks, a return type and argument types. Function names must be unique in the module.\n",
      "\n",
      "```\n",
      "define i32 @add(i32 %a, i32 %b) {\n",
      "  %1 = add i32 %a, %b\n",
      "  ret i32 %1\n",
      "}\n",
      "```\n",
      "    \n",
      "**Basic Blocks**\n",
      "    \n",
      "The function is split across basic blocks which hold sequences of instructions and a terminator instruction which either returns or jumps to another local basic block.\n",
      "\n",
      "```\n",
      "define i1 @foo() {\n",
      "entry:\n",
      "  br label %next\n",
      "next:\n",
      "  br label %return\n",
      "return:\n",
      "  ret i1 0\n",
      "}\n",
      "```\n",
      "\n",
      "\n",
      "**Return**\n",
      "\n",
      "A function must have a terminator, one of such instructions is a ``ret`` which returns a value to the stack."
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "```\n",
      "define i1 @foo() {\n",
      "  ret i1 0\n",
      "}\n",
      "```"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<img src=\"files/images/ret.png\">"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**Unconditional Branch**\n",
      "\n",
      "An unconditional branch jumps unconditionally to a labeled basic block.\n",
      "\n",
      "```\n",
      "define i1 @foo() {\n",
      "start:\n",
      "  br label %next\n",
      "next:\n",
      "  br label %return\n",
      "return:\n",
      "  ret i1 0\n",
      "}\n",
      "```"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<img src=\"files/images/branch.png\">"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**Conditional Branch**\n",
      "\n",
      "```\n",
      "define i32 @foo() {\n",
      "start:\n",
      "  br i1 true, label %left, label %right\n",
      "left:\n",
      "  ret i32 10\n",
      "right:\n",
      "  ret i32 20\n",
      "}\n",
      "```"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<img src=\"files/images/cbranch.png\">"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**Phi**\n",
      "\n",
      "Phi nodes yield a value that depends on the operand corresponding to their predecessor basic block. These are used for implementing loops in SSA.\n",
      "\n",
      "```\n",
      "define i32 @foo() {\n",
      "start:\n",
      "  br i1 true, label %left, label %right\n",
      "left:\n",
      "  %plusOne = add i32 0, 1\n",
      "  br label %merge\n",
      "right:\n",
      "  br label %merge\n",
      "merge:\n",
      "  %join = phi i32 [ %plusOne, %left], [ -1, %right]\n",
      "  ret i32 %join\n",
      "}\n",
      "```"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<img src=\"files/images/phi.png\">"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**Switch**\n",
      "\n",
      "Switch statements are like switch statements in C, and can be used to build jump tables.\n",
      "\n",
      "```\n",
      "define i32 @foo(i32 %a) {\n",
      "entry:\n",
      "  switch i32 %a, label %default [ i32 0, label %f\n",
      "                                  i32 1, label %g\n",
      "                                  i32 2, label %h ]\n",
      "f:\n",
      "  ret i32 1\n",
      "g:\n",
      "  ret i32 2\n",
      "h:\n",
      "  ret i32 3\n",
      "default:\n",
      "  ret i32 0\n",
      "}\n",
      "```"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<img src=\"files/images/switch.png\">"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**Loops**\n",
      "\n",
      "Loops are written in terms of conditional branches and phi nodes.\n",
      "\n",
      "For example the translation of the following C code:\n",
      "\n",
      "```\n",
      "int count(int n) \n",
      "{\n",
      "  int i = 0;\n",
      "  while(i < n) \n",
      "  {\n",
      "    i++;\n",
      "  }\n",
      "  return i;\n",
      "}\n",
      "```\n",
      "\n",
      "Into LLVM:\n",
      "\n",
      "```\n",
      "define i32 @count(i32 %n) {\n",
      "entry:\n",
      "   br label %loop\n",
      "loop:\n",
      "   %i = phi i32 [ 1, %entry ], [ %nextvar, %loop ]\n",
      "   %nextvar = add i32 %i, 1\n",
      "   %cmptmp = icmp ult i32 %i, %n\n",
      "   %booltmp = zext i1 %cmptmp to i32\n",
      "   %loopcond = icmp ne i32 %booltmp, 0\n",
      "   br i1 %loopcond, label %loop, label %afterloop\n",
      "afterloop:\n",
      "   ret i32 %i\n",
      "}\n",
      "```"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "<img src=\"files/images/for.png\">\n",
      "\n",
      "**Toolchain**\n",
      "\n",
      "The command line utilities for LLVM can be used to transform IR to and from various forms and run optimizations over it. Everything we can do from the C++ API or llvmpy can also be done from the command line.\n",
      "\n",
      "```\n",
      "$ llc example.ll -o example.s             # compile\n",
      "$ lli example.ll                          # execute\n",
      "$ opt -S example.bc -o example.ll         # to assembly\n",
      "$ opt example.ll -o example.bc            # to bitcode\n",
      "$ opt -O3 example.ll -o example.opt.ll -S # run optimizer\n",
      "$ opt -view-cfg module.ll                 # view control flow graph\n",
      "```\n",
      "\n",
      "\n",
      "And that's basically all you need to know about LLVM. Also get used to segfaulting the Python interpreter a lot when using llvmpy."
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Python AST\n",
      "\n"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Python's internal AST is accessible from within the Python interpreter. Really the only time you'd ever use this module is if you're doing crazy metaprogramming, which is what we're about to do! Ostensibly we're going to be taking an arbitrary function introspecting it's AST and then mapping it into another syntax called the *Core* which we'll endow with a different (C-like) semantics on top of as well as doing type inference on the logic to make the AST explicitly typed.\n",
      "\n",
      "* https://greentreesnakes.readthedocs.org\n",
      "* https://docs.python.org/2/library/ast.html\n",
      "\n",
      "Hat tip to John Riehl for this pretty printing technique."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import ast\n",
      "import pprint\n",
      "\n",
      "def ast2tree(node, include_attrs=True):\n",
      "    def _transform(node):\n",
      "        if isinstance(node, ast.AST):\n",
      "            fields = ((a, _transform(b))\n",
      "                      for a, b in ast.iter_fields(node))\n",
      "            if include_attrs:\n",
      "                attrs = ((a, _transform(getattr(node, a)))\n",
      "                         for a in node._attributes\n",
      "                         if hasattr(node, a))\n",
      "                return (node.__class__.__name__, dict(fields), dict(attrs))\n",
      "            return (node.__class__.__name__, dict(fields))\n",
      "        elif isinstance(node, list):\n",
      "            return [_transform(x) for x in node]\n",
      "        elif isinstance(node, str):\n",
      "            return repr(node)\n",
      "        return node\n",
      "    if not isinstance(node, ast.AST):\n",
      "        raise TypeError('expected AST, got %r' % node.__class__.__name__)\n",
      "    return _transform(node)\n",
      "\n",
      "\n",
      "def pformat_ast(node, include_attrs=False, **kws):\n",
      "    return pprint.pformat(ast2tree(node, include_attrs), **kws)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 10
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "So if we feed this function a source string, the ast module will go off an pares it into the AST and we'll get this nicely presented nested-dict for it's field structure. In fact we'll use the ``ast.Node`` for our custom AST just so that we can reuse this pretty printer."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "source = \"\"\"\n",
      "def f(x): \n",
      "    return f(x+1)\n",
      "\"\"\"\n",
      "\n",
      "print(pformat_ast(ast.parse(source)))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "('Module',\n",
        " {'body': [('FunctionDef',\n",
        "            {'args': ('arguments',\n",
        "                      {'args': [('Name',\n",
        "                                 {'ctx': ('Param', {}), 'id': \"'x'\"})],\n",
        "                       'defaults': [],\n",
        "                       'kwarg': None,\n",
        "                       'vararg': None}),\n",
        "             'body': [('Return',\n",
        "                       {'value': ('Call',\n",
        "                                  {'args': [('BinOp',\n",
        "                                             {'left': ('Name',\n",
        "                                                       {'ctx': ('Load',\n",
        "                                                                {}),\n",
        "                                                        'id': \"'x'\"}),\n",
        "                                              'op': ('Add', {}),\n",
        "                                              'right': ('Num',\n",
        "                                                        {'n': 1})})],\n",
        "                                   'func': ('Name',\n",
        "                                            {'ctx': ('Load', {}),\n",
        "                                             'id': \"'f'\"}),\n",
        "                                   'keywords': [],\n",
        "                                   'kwargs': None,\n",
        "                                   'starargs': None})})],\n",
        "             'decorator_list': [],\n",
        "             'name': \"'f'\"})]})\n"
       ]
      }
     ],
     "prompt_number": 11
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Core Language\n",
      "\n",
      "First we'll need to bring in a few libraries, pretty standard fare standard library stuff. And a few LLVM libraries, more on this later.\n"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from __future__ import print_function\n",
      "\n",
      "import sys\n",
      "import ast\n",
      "import types\n",
      "import ctypes\n",
      "import inspect\n",
      "import pprint\n",
      "import string\n",
      "import numpy as np\n",
      "from itertools import tee, izip\n",
      "\n",
      "from textwrap import dedent\n",
      "from collections import deque, defaultdict\n",
      "\n",
      "import llvm.core as lc\n",
      "import llvm.passes as lp\n",
      "import llvm.ee as le\n",
      "from llvm.core import Module, Builder, Function, Type, Constant\n",
      "\n",
      "DEBUG = False"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 12
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Our Core language will be a simple expression language with 12 terms that we will condense a subset of the much larger Python AST into.\n",
      "\n",
      "```\n",
      "e : var                 (Variable)\n",
      "  | n = e               (Assignment)\n",
      "  | return e            (Return)\n",
      "  | loop n e e [e]      (Loop Construct)\n",
      "  | %int                (Integer)\n",
      "  | %float              (Float)\n",
      "  | %bool               (Boolean)\n",
      "  | e {e}               (Variadic Application)\n",
      "  | function n {e} [e]  (Variadic Function)\n",
      "  | prim n              (Primop)\n",
      "  | index e e           (Array indexing)\n",
      "  | noop                (Noop)\n",
      "```\n",
      "\n",
      "Our core language will have two forms, one is untyped and the other has all named expressions (``n``) annotated with an attached type field."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class Var(ast.AST):\n",
      "    _fields = [\"id\", \"type\"]\n",
      "\n",
      "    def __init__(self, id, type=None):\n",
      "        self.id = id\n",
      "        self.type = type\n",
      "\n",
      "class Assign(ast.AST):\n",
      "    _fields = [\"ref\", \"val\", \"type\"]\n",
      "\n",
      "    def __init__(self, ref, val, type=None):\n",
      "        self.ref = ref\n",
      "        self.val = val\n",
      "        self.type = type\n",
      "\n",
      "class Return(ast.AST):\n",
      "    _fields = [\"val\"]\n",
      "\n",
      "    def __init__(self, val):\n",
      "        self.val = val\n",
      "\n",
      "class Loop(ast.AST):\n",
      "    _fields = [\"var\", \"begin\", \"end\", \"body\"]\n",
      "\n",
      "    def __init__(self, var, begin, end, body):\n",
      "        self.var = var\n",
      "        self.begin = begin\n",
      "        self.end = end\n",
      "        self.body = body\n",
      "\n",
      "class App(ast.AST):\n",
      "    _fields = [\"fn\", \"args\"]\n",
      "\n",
      "    def __init__(self, fn, args):\n",
      "        self.fn = fn\n",
      "        self.args = args\n",
      "\n",
      "class Fun(ast.AST):\n",
      "    _fields = [\"fname\", \"args\", \"body\"]\n",
      "\n",
      "    def __init__(self, fname, args, body):\n",
      "        self.fname = fname\n",
      "        self.args = args\n",
      "        self.body = body\n",
      "\n",
      "class LitInt(ast.AST):\n",
      "    _fields = [\"n\"]\n",
      "\n",
      "    def __init__(self, n, type=None):\n",
      "        self.n = n\n",
      "        self.type = type\n",
      "\n",
      "class LitFloat(ast.AST):\n",
      "    _fields = [\"n\"]\n",
      "\n",
      "    def __init__(self, n, type=None):\n",
      "        self.n = n\n",
      "        self.type = None\n",
      "\n",
      "class LitBool(ast.AST):\n",
      "    _fields = [\"n\"]\n",
      "\n",
      "    def __init__(self, n):\n",
      "        self.n = n\n",
      "\n",
      "class Prim(ast.AST):\n",
      "    _fields = [\"fn\", \"args\"]\n",
      "\n",
      "    def __init__(self, fn, args):\n",
      "        self.fn = fn\n",
      "        self.args = args\n",
      "\n",
      "class Index(ast.AST):\n",
      "    _fields = [\"val\", \"ix\"]\n",
      "\n",
      "    def __init__(self, val, ix):\n",
      "        self.val = val\n",
      "        self.ix = ix\n",
      "\n",
      "class Noop(ast.AST):\n",
      "    _fields = []"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 13
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "In similar fashion we have a very simple type system. Our function type is variadic, it takes a tuple of arguments to a single output. \n",
      "\n",
      "```\n",
      "t : a            (Type Variable)\n",
      "  | C {t}        (Named Constructor)\n",
      "  | t            (Type Application)\n",
      "  | [t] -> t     (Function type)\n",
      "```\n",
      "\n",
      "\n",
      "The basic constructors will simply be the machine types. By default we will map Python's integer to ``int64`` and floating point to ``double``. Python's integer type is an  arbitrary precision integer, whereas LLVM is a machine integer so obviously there are different semantics."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class TVar(object):\n",
      "    def __init__(self, s):\n",
      "        self.s = s\n",
      "\n",
      "    def __hash__(self):\n",
      "        return hash(self.s)\n",
      "\n",
      "    def __eq__(self, other):\n",
      "        if isinstance(other, TVar):\n",
      "            return (self.s == other.s)\n",
      "        else:\n",
      "            return False\n",
      "\n",
      "    def __str__(self):\n",
      "        return self.s\n",
      "    __repr__ = __str__\n",
      "\n",
      "class TCon(object):\n",
      "    def __init__(self, s):\n",
      "        self.s = s\n",
      "\n",
      "    def __eq__(self, other):\n",
      "        if isinstance(other, TCon):\n",
      "            return (self.s == other.s)\n",
      "        else:\n",
      "            return False\n",
      "\n",
      "    def __hash__(self):\n",
      "        return hash(self.s)\n",
      "\n",
      "    def __str__(self):\n",
      "        return self.s\n",
      "    __repr__ = __str__\n",
      "\n",
      "class TApp(object):\n",
      "    def __init__(self, a, b):\n",
      "        self.a = a\n",
      "        self.b = b\n",
      "\n",
      "    def __eq__(self, other):\n",
      "        if isinstance(other, TApp):\n",
      "            return (self.a == other.a) & (self.b == other.b)\n",
      "        else:\n",
      "            return False\n",
      "\n",
      "    def __hash__(self):\n",
      "        return hash((self.a, self.b))\n",
      "\n",
      "    def __str__(self):\n",
      "        return str(self.a) + \" \" + str(self.b)\n",
      "    __repr__ = __str__\n",
      "\n",
      "class TFun(object):\n",
      "    def __init__(self, argtys, retty):\n",
      "        assert isinstance(argtys, list)\n",
      "        self.argtys = argtys\n",
      "        self.retty = retty\n",
      "\n",
      "    def __eq__(self, other):\n",
      "        if isinstance(other, TFun):\n",
      "            return (self.argtys == other.argtys) & (self.retty == other.retty)\n",
      "        else:\n",
      "            return False\n",
      "\n",
      "    def __str__(self):\n",
      "        return str(self.argtys) + \" -> \" + str(self.retty)\n",
      "    __repr__ = __str__\n",
      "\n",
      "def ftv(x):\n",
      "    if isinstance(x, TCon):\n",
      "        return set()\n",
      "    elif isinstance(x, TApp):\n",
      "        return ftv(x.a) | ftv(x.b)\n",
      "    elif isinstance(x, TFun):\n",
      "        return reduce(set.union, map(ftv, x.argtys)) | ftv(x.retty)\n",
      "    elif isinstance(x, TVar):\n",
      "        return set([x])\n",
      "\n",
      "def is_array(ty):\n",
      "    return isinstance(ty, TApp) and ty.a == TCon(\"Array\")"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 14
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "int32 = TCon(\"Int32\")\n",
      "int64 = TCon(\"Int64\")\n",
      "float32 = TCon(\"Float\")\n",
      "double64 = TCon(\"Double\")\n",
      "void = TCon(\"Void\")\n",
      "array = lambda t: TApp(TCon(\"Array\"), t)\n",
      "\n",
      "array_int32 = array(int32)\n",
      "array_int64 = array(int64)\n",
      "array_double64 = array(double64)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 15
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The Python to Core translator is a fairly unremarkable ``NodeVisitor`` class. It recursively descends through the Python AST compressing it into our Core form. For our example application this is obviously only a very small subset of the entire AST, and a lot of cases are missing. We are going to support basic loops, arithmetic with addition and multiplication, numeric literals, and array indexing. "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class PythonVisitor(ast.NodeVisitor):\n",
      "\n",
      "    def __init__(self):\n",
      "        pass\n",
      "\n",
      "    def __call__(self, source):\n",
      "        if isinstance(source, types.ModuleType):\n",
      "            source = dedent(inspect.getsource(source))\n",
      "        if isinstance(source, types.FunctionType):\n",
      "            source = dedent(inspect.getsource(source))\n",
      "        if isinstance(source, types.LambdaType):\n",
      "            source = dedent(inspect.getsource(source))\n",
      "        elif isinstance(source, (str, unicode)):\n",
      "            source = dedent(source)\n",
      "        else:\n",
      "            raise NotImplementedError\n",
      "\n",
      "        self._source = source\n",
      "        self._ast = ast.parse(source)\n",
      "        return self.visit(self._ast)\n",
      "\n",
      "    def visit_Module(self, node):\n",
      "        body = map(self.visit, node.body)\n",
      "        return body[0]\n",
      "\n",
      "    def visit_Name(self, node):\n",
      "        return Var(node.id)\n",
      "\n",
      "    def visit_Num(self, node):\n",
      "        if isinstance(node.n, float):\n",
      "            return LitFloat(node.n)\n",
      "        else:\n",
      "            return LitInt(node.n)\n",
      "\n",
      "    def visit_Bool(self, node):\n",
      "        return LitBool(node.n)\n",
      "\n",
      "    def visit_Call(self, node):\n",
      "        name = self.visit(node.func)\n",
      "        args = map(self.visit, node.args)\n",
      "        keywords = map(self.visit, node.keywords)\n",
      "        return App(name, args)\n",
      "\n",
      "    def visit_BinOp(self, node):\n",
      "        op_str = node.op.__class__\n",
      "        a = self.visit(node.left)\n",
      "        b = self.visit(node.right)\n",
      "        opname = primops[op_str]\n",
      "        return Prim(opname, [a, b])\n",
      "\n",
      "    def visit_Assign(self, node):\n",
      "        targets = node.targets\n",
      "\n",
      "        assert len(node.targets) == 1\n",
      "        var = node.targets[0].id\n",
      "        val = self.visit(node.value)\n",
      "        return Assign(var, val)\n",
      "\n",
      "    def visit_FunctionDef(self, node):\n",
      "        stmts = list(node.body)\n",
      "        stmts = map(self.visit, stmts)\n",
      "        args = map(self.visit, node.args.args)\n",
      "        res = Fun(node.name, args, stmts)\n",
      "        return res\n",
      "\n",
      "    def visit_Pass(self, node):\n",
      "        return Noop()\n",
      "\n",
      "    def visit_Lambda(self, node):\n",
      "        args = self.visit(node.args)\n",
      "        body = self.visit(node.body)\n",
      "\n",
      "    def visit_Return(self, node):\n",
      "        val = self.visit(node.value)\n",
      "        return Return(val)\n",
      "\n",
      "    def visit_Attribute(self, node):\n",
      "        if node.attr == \"shape\":\n",
      "            val = self.visit(node.value)\n",
      "            return Prim(\"shape#\", [val])\n",
      "        else:\n",
      "            raise NotImplementedError\n",
      "\n",
      "    def visit_Subscript(self, node):\n",
      "        if isinstance(node.ctx, ast.Load):\n",
      "            if node.slice:\n",
      "                val = self.visit(node.value)\n",
      "                ix = self.visit(node.slice.value)\n",
      "                return Index(val, ix)\n",
      "        elif isinstance(node.ctx, ast.Store):\n",
      "            raise NotImplementedError\n",
      "\n",
      "    def visit_For(self, node):\n",
      "        target = self.visit(node.target)\n",
      "        stmts = map(self.visit, node.body)\n",
      "        if node.iter.func.id in {\"xrange\", \"range\"}:\n",
      "            args = map(self.visit, node.iter.args)\n",
      "        else:\n",
      "            raise Exception(\"Loop must be over range\")\n",
      "\n",
      "        if len(args) == 1:   # xrange(n)\n",
      "            return Loop(target, LitInt(0, type=int32), args[0], stmts)\n",
      "        elif len(args) == 2:  # xrange(n,m)\n",
      "            return Loop(target, args[0], args[1], stmts)\n",
      "\n",
      "    def visit_AugAssign(self, node):\n",
      "        if isinstance(node.op, ast.Add):\n",
      "            ref = node.target.id\n",
      "            value = self.visit(node.value)\n",
      "            return Assign(ref, Prim(\"add#\", [Var(ref), value]))\n",
      "        if isinstance(node.op, ast.Mul):\n",
      "            ref = node.target.id\n",
      "            value = self.visit(node.value)\n",
      "            return Assign(ref, Prim(\"mult#\", [Var(ref), value]))\n",
      "        else:\n",
      "            raise NotImplementedError\n",
      "\n",
      "    def generic_visit(self, node):\n",
      "        raise NotImplementedError"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 16
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "So if we define a very simple function like:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def add(a,b):\n",
      "    return a + b"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 17
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "There are several builtin \"primops\" which are simply functions which have a direct mapping to some function lower in the pipeline. \n",
      "\n",
      "* ``add#`` : Generic addition (integral, floating point)\n",
      "* ``mult#`` : Generic multiplication (integral, floating point)\n",
      "* ``shape#`` : Shape extraction for NumPy ndarrays."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "primops = {ast.Add: \"add#\", ast.Mult: \"mult#\"}"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 18
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "And run our transformer over it with:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "transformer = PythonVisitor()\n",
      "core = transformer(add)\n",
      "print(pformat_ast(core))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "('Fun',\n",
        " {'args': [('Var', {'id': \"'a'\", 'type': None}),\n",
        "           ('Var', {'id': \"'b'\", 'type': None})],\n",
        "  'body': [('Return',\n",
        "            {'val': ('Prim',\n",
        "                     {'args': [('Var', {'id': \"'a'\", 'type': None}),\n",
        "                               ('Var', {'id': \"'b'\", 'type': None})],\n",
        "                      'fn': \"'add#'\"})})],\n",
        "  'fname': \"'add'\"})\n"
       ]
      }
     ],
     "prompt_number": 19
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "For a more complex function consider:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def count(n):\n",
      "    a = 0\n",
      "    for i in range(0, n):\n",
      "        a += i\n",
      "    return a\n",
      "\n",
      "transformer = PythonVisitor()\n",
      "core = transformer(count)\n",
      "print(pformat_ast(core))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "('Fun',\n",
        " {'args': [('Var', {'id': \"'n'\", 'type': None})],\n",
        "  'body': [('Assign',\n",
        "            {'ref': \"'a'\", 'type': None, 'val': ('LitInt', {'n': 0})}),\n",
        "           ('Loop',\n",
        "            {'begin': ('LitInt', {'n': 0}),\n",
        "             'body': [('Assign',\n",
        "                       {'ref': \"'a'\",\n",
        "                        'type': None,\n",
        "                        'val': ('Prim',\n",
        "                                {'args': [('Var',\n",
        "                                           {'id': \"'a'\",\n",
        "                                            'type': None}),\n",
        "                                          ('Var',\n",
        "                                           {'id': \"'i'\",\n",
        "                                            'type': None})],\n",
        "                                 'fn': \"'add#'\"})})],\n",
        "             'end': ('Var', {'id': \"'n'\", 'type': None}),\n",
        "             'var': ('Var', {'id': \"'i'\", 'type': None})}),\n",
        "           ('Return', {'val': ('Var', {'id': \"'a'\", 'type': None})})],\n",
        "  'fname': \"'count'\"})\n"
       ]
      }
     ],
     "prompt_number": 20
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Type Inference\n",
      "\n",
      "For type inference we wish to take our untyped AST and overlay types deduced from two sources\n",
      "\n",
      "* Types intrinsic to the operations in use\n",
      "* User input types\n",
      "\n",
      "To do this we will use a very traditional method of constraint based unification for type reconstruction. We will walk our AST generating a constraint set of equality relations between types (written as ``a ~ b``), which will give rise to a large constraint problem we will solve when given a set of input types for arguments. Whenever we  don't know the type of an expression we will place a fresh free type variable in it's place and solve for it when given more information.\n",
      "\n",
      "There are four possible outcomes:\n",
      "\n",
      "* The types are correctly determined.\n",
      "* The types are underdetermined.\n",
      "* The types is polymorphic.\n",
      "* The types are inconsistent.\n",
      "\n",
      "The case where the function is polymorphic implies that there are free type variables remaining in the toplevel type. For instance we might have a type like:\n",
      "\n",
      "```[Array a, Array a] -> a```\n",
      "\n",
      "Which just means that the logic is independent of the type of the element of the arrays, and can operate polymorphicly over any element type. This is good for code reuse and implies we get a whole family of functions supposing that our compiler knows how to lower ``a``.\n",
      "\n",
      "* The types are underdetermined. Implies that the constraints induced by usage are too lax to fully determine every subexpression. In this case an explicit annotation is needed.\n",
      "\n",
      "* The type inconsistent. This will happen where there is no solution that would satisfy the given constraints. For example trying to a call function with signature:\n",
      "\n",
      "```[a,a] -> a```\n",
      "\n",
      "Over the types ``[Int64, Double]`` has no solution since there can be no solution where ``Int64 ~ Double``."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def naming():\n",
      "    k = 0\n",
      "    while True:\n",
      "        for a in string.ascii_lowercase:\n",
      "            yield (\"'\"+a+str(k)) if (k > 0) else (a)\n",
      "        k = k+1\n",
      "\n",
      "class TypeInfer(object):\n",
      "\n",
      "    def __init__(self):\n",
      "        self.constraints = []\n",
      "        self.env = {}\n",
      "        self.names = naming()\n",
      "\n",
      "    def fresh(self):\n",
      "        return TVar('$' + next(self.names))  # New meta type variable.\n",
      "\n",
      "    def visit(self, node):\n",
      "        name = \"visit_%s\" % type(node).__name__\n",
      "        if hasattr(self, name):\n",
      "            return getattr(self, name)(node)\n",
      "        else:\n",
      "            return self.generic_visit(node)\n",
      "\n",
      "    def visit_Fun(self, node):\n",
      "        arity = len(node.args)\n",
      "        self.argtys = [self.fresh() for v in node.args]\n",
      "        self.retty = TVar(\"$retty\")\n",
      "        for (arg, ty) in zip(node.args, self.argtys):\n",
      "            arg.type = ty\n",
      "            self.env[arg.id] = ty\n",
      "        map(self.visit, node.body)\n",
      "        return TFun(self.argtys, self.retty)\n",
      "\n",
      "    def visit_Noop(self, node):\n",
      "        return None\n",
      "\n",
      "    def visit_LitInt(self, node):\n",
      "        tv = self.fresh()\n",
      "        node.type = tv\n",
      "        return tv\n",
      "\n",
      "    def visit_LitFloat(self, node):\n",
      "        tv = self.fresh()\n",
      "        node.type = tv\n",
      "        return tv\n",
      "\n",
      "    def visit_Assign(self, node):\n",
      "        ty = self.visit(node.val)\n",
      "        if node.ref in self.env:\n",
      "            # Subsequent uses of a variable must have the same type.\n",
      "            self.constraints += [(ty, self.env[node.ref])]\n",
      "        self.env[node.ref] = ty\n",
      "        node.type = ty\n",
      "        return None\n",
      "\n",
      "    def visit_Index(self, node):\n",
      "        tv = self.fresh()\n",
      "        ty = self.visit(node.val)\n",
      "        ixty = self.visit(node.ix)\n",
      "        self.constraints += [(ty, array(tv)), (ixty, int32)]\n",
      "        return tv\n",
      "\n",
      "    def visit_Prim(self, node):\n",
      "        if node.fn == \"shape#\":\n",
      "            return array(int32)\n",
      "        elif node.fn == \"mult#\":\n",
      "            tya = self.visit(node.args[0])\n",
      "            tyb = self.visit(node.args[1])\n",
      "            self.constraints += [(tya, tyb)]\n",
      "            return tyb\n",
      "        elif node.fn == \"add#\":\n",
      "            tya = self.visit(node.args[0])\n",
      "            tyb = self.visit(node.args[1])\n",
      "            self.constraints += [(tya, tyb)]\n",
      "            return tyb\n",
      "        else:\n",
      "            raise NotImplementedError\n",
      "\n",
      "    def visit_Var(self, node):\n",
      "        ty = self.env[node.id]\n",
      "        node.type = ty\n",
      "        return ty\n",
      "\n",
      "    def visit_Return(self, node):\n",
      "        ty = self.visit(node.val)\n",
      "        self.constraints += [(ty, self.retty)]\n",
      "\n",
      "    def visit_Loop(self, node):\n",
      "        self.env[node.var.id] = int32\n",
      "        varty = self.visit(node.var)\n",
      "        begin = self.visit(node.begin)\n",
      "        end = self.visit(node.end)\n",
      "        self.constraints += [(varty, int32), (\n",
      "            begin, int64), (end, int32)]\n",
      "        map(self.visit, node.body)\n",
      "\n",
      "    def generic_visit(self, node):\n",
      "        raise NotImplementedError"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 21
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "When the traversal is finished we'll have a set of constraints to solve:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def addup(n):\n",
      "    x = 1\n",
      "    for i in range(n):\n",
      "        n += 1 + x\n",
      "    return n\n",
      "\n",
      "transformer = PythonVisitor()\n",
      "core = transformer(addup)\n",
      "infer = TypeInfer()\n",
      "sig = infer.visit(core)\n",
      "\n",
      "print('Signature:%s \\n' % sig)\n",
      "\n",
      "print('Constraints:')\n",
      "for (a,b) in infer.constraints:\n",
      "    print(a, '~', b)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Signature:[$a] -> $retty \n",
        "\n",
        "Constraints:\n",
        "Int32 ~ Int32\n",
        "$c ~ Int64\n",
        "$a ~ Int32\n",
        "$d ~ $b\n",
        "$a ~ $b\n",
        "$b ~ $a\n",
        "$b ~ $retty\n"
       ]
      }
     ],
     "prompt_number": 22
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "So now we're left with a little riddle to reduce the number variables in the expression by equating like terms. We also notice that the inference has annotated our AST with explicit type terms for all the free variables."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print(pformat_ast(core))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "('Fun',\n",
        " {'args': [('Var', {'id': \"'n'\", 'type': $a})],\n",
        "  'body': [('Assign',\n",
        "            {'ref': \"'x'\", 'type': $b, 'val': ('LitInt', {'n': 1})}),\n",
        "           ('Loop',\n",
        "            {'begin': ('LitInt', {'n': 0}),\n",
        "             'body': [('Assign',\n",
        "                       {'ref': \"'n'\",\n",
        "                        'type': $b,\n",
        "                        'val': ('Prim',\n",
        "                                {'args': [('Var',\n",
        "                                           {'id': \"'n'\",\n",
        "                                            'type': $a}),\n",
        "                                          ('Prim',\n",
        "                                           {'args': [('LitInt',\n",
        "                                                      {'n': 1}),\n",
        "                                                     ('Var',\n",
        "                                                      {'id': \"'x'\",\n",
        "                                                       'type': $b})],\n",
        "                                            'fn': \"'add#'\"})],\n",
        "                                 'fn': \"'add#'\"})})],\n",
        "             'end': ('Var', {'id': \"'n'\", 'type': $a}),\n",
        "             'var': ('Var', {'id': \"'i'\", 'type': Int32})}),\n",
        "           ('Return', {'val': ('Var', {'id': \"'n'\", 'type': $b})})],\n",
        "  'fname': \"'addup'\"})\n"
       ]
      }
     ],
     "prompt_number": 23
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "So now we'll solve the system of equations using the very traditional unification solver via Robinson's algorithm. The solver will recursively build up the *most general unifier* (mgu) which is a substitution which when applied to the term yields the minimal singleton solution set."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def empty():\n",
      "    return {}\n",
      "\n",
      "def apply(s, t):\n",
      "    if isinstance(t, TCon):\n",
      "        return t\n",
      "    elif isinstance(t, TApp):\n",
      "        return TApp(apply(s, t.a), apply(s, t.b))\n",
      "    elif isinstance(t, TFun):\n",
      "        argtys = [apply(s, a) for a in t.argtys]\n",
      "        retty = apply(s, t.retty)\n",
      "        return TFun(argtys, retty)\n",
      "    elif isinstance(t, TVar):\n",
      "        return s.get(t.s, t)\n",
      "\n",
      "def applyList(s, xs):\n",
      "    return [(apply(s, x), apply(s, y)) for (x, y) in xs]\n",
      "\n",
      "def unify(x, y):\n",
      "    if isinstance(x, TApp) and isinstance(y, TApp):\n",
      "        s1 = unify(x.a, y.a)\n",
      "        s2 = unify(apply(s1, x.b), apply(s1, y.b))\n",
      "        return compose(s2, s1)\n",
      "    elif isinstance(x, TCon) and isinstance(y, TCon) and (x == y):\n",
      "        return empty()\n",
      "    elif isinstance(x, TFun) and isinstance(y, TFun):\n",
      "        if len(x.argtys) != len(y.argtys):\n",
      "            return Exception(\"Wrong number of arguments\")\n",
      "        s1 = solve(zip(x.argtys, y.argtys))\n",
      "        s2 = unify(apply(s1, x.retty), apply(s1, y.retty))\n",
      "        return compose(s2, s1)\n",
      "    elif isinstance(x, TVar):\n",
      "        return bind(x.s, y)\n",
      "    elif isinstance(y, TVar):\n",
      "        return bind(y.s, x)\n",
      "    else:\n",
      "        raise InferError(x, y)\n",
      "\n",
      "def solve(xs):\n",
      "    mgu = empty()\n",
      "    cs = deque(xs)\n",
      "    while len(cs):\n",
      "        (a, b) = cs.pop()\n",
      "        s = unify(a, b)\n",
      "        mgu = compose(s, mgu)\n",
      "        cs = deque(applyList(s, cs))\n",
      "    return mgu\n",
      "\n",
      "def bind(n, x):\n",
      "    if x == n:\n",
      "        return empty()\n",
      "    elif occurs_check(n, x):\n",
      "        raise InfiniteType(n, x)\n",
      "    else:\n",
      "        return dict([(n, x)])\n",
      "\n",
      "def occurs_check(n, x):\n",
      "    return n in ftv(x)\n",
      "\n",
      "def union(s1, s2):\n",
      "    nenv = s1.copy()\n",
      "    nenv.update(s2)\n",
      "    return nenv\n",
      "\n",
      "def compose(s1, s2):\n",
      "    s3 = dict((t, apply(s1, u)) for t, u in s2.items())\n",
      "    return union(s1, s3)\n",
      "\n",
      "class UnderDetermined(Exception):\n",
      "    def __str__(self):\n",
      "        return \"The types in the function are not fully determined by the \\\n",
      "                input types. Add annotations.\"\n",
      "\n",
      "class InferError(Exception):\n",
      "    def __init__(self, ty1, ty2):\n",
      "        self.ty1 = ty1\n",
      "        self.ty2 = ty2\n",
      "\n",
      "    def __str__(self):\n",
      "        return '\\n'.join([\n",
      "            \"Type mismatch: \",\n",
      "            \"Given: \", \"\\t\" + str(self.ty1),\n",
      "            \"Expected: \", \"\\t\" + str(self.ty2)\n",
      "        ])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 24
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def dot2(a, b):\n",
      "    c = 0\n",
      "    n = a.shape[0]\n",
      "    for i in range(n):\n",
      "       c += a[i]*b[i]\n",
      "    return c\n",
      "\n",
      "def test_infer(fn):\n",
      "    transformer = PythonVisitor()\n",
      "    ast = transformer(fn)\n",
      "    infer = TypeInfer()\n",
      "    ty = infer.visit(ast)\n",
      "    mgu = solve(infer.constraints)\n",
      "    infer_ty = apply(mgu, ty)\n",
      "    \n",
      "    print('Unifier: ')\n",
      "    for (a,b) in mgu.iteritems():\n",
      "        print(a + ' ~ ' + str(b))\n",
      "\n",
      "    print('Solution: ', infer_ty)\n",
      "\n",
      "test_infer(dot2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Unifier: \n",
        "$h ~ $c\n",
        "$f ~ Int64\n",
        "$g ~ $c\n",
        "$d ~ Int32\n",
        "$e ~ Int32\n",
        "$b ~ Array $c\n",
        "$c ~ $c\n",
        "$a ~ Array $c\n",
        "$retty ~ $c\n",
        "Solution:  [Array $c, Array $c] -> $c\n"
       ]
      }
     ],
     "prompt_number": 25
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "So in this case we have solution \n",
      "\n",
      "``[Array $c, Array $c] -> $c``\n",
      "\n",
      "indicating that our dot product function is polymorphic in both of it's arguments and return type. It works for any array."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def addup(n):\n",
      "    x = 1\n",
      "    for i in range(n):\n",
      "        n += i + x\n",
      "    return n\n",
      "\n",
      "test_infer(addup)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Unifier: \n",
        "$retty ~ Int32\n",
        "$b ~ Int32\n",
        "$c ~ Int64\n",
        "$a ~ Int32\n",
        "Solution:  [Int32] -> Int32\n"
       ]
      }
     ],
     "prompt_number": 26
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Where as for the addup function our inferred type is simply entirely determiend by the type of iteration variable, which we for range we defined to default to ``Int32`` which determines both the type of the input and the type of the output and the intermediate type of ``x``. \n",
      "\n",
      "Consider now a case where the system is underdetermined. If we ignore one of the arguments then our system doesn't have any constraints to solve for and it's simply left as a free variable."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def const(a,b):\n",
      "    return a\n",
      "\n",
      "test_infer(addup)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Unifier: \n",
        "$retty ~ Int32\n",
        "$b ~ Int32\n",
        "$c ~ Int64\n",
        "$a ~ Int32\n",
        "Solution:  [Int32] -> Int32\n"
       ]
      }
     ],
     "prompt_number": 27
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# LLVM Code Generator\n",
      "\n",
      "Now we set up another type system, the LLVM type system which map directly onto machine types for our platform.\n",
      "\n",
      "The only nonobvious thing going on here is that our NumPy arrays will be passed around as a structure object that holds metadata from the originally NumPy ndarray. The ``data`` pointer is simply the pointer to data buffer that NumPy allocated for it's values. In C we would write:\n",
      "\n",
      "```\n",
      "struct ndarray_double {\n",
      "    data *double;\n",
      "    dims int;\n",
      "    shape *int;\n",
      "}\n",
      "```"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "pointer     = Type.pointer\n",
      "int_type    = Type.int()\n",
      "float_type  = Type.float()\n",
      "double_type = Type.double()\n",
      "bool_type   = Type.int(1)\n",
      "void_type   = Type.void()\n",
      "void_ptr    = pointer(Type.int(8))\n",
      "\n",
      "def array_type(elt_type):\n",
      "    return Type.struct([\n",
      "        pointer(elt_type),  # data\n",
      "        int_type,           # dimensions\n",
      "        pointer(int_type),  # shape\n",
      "    ], name='ndarray_' + str(elt_type))\n",
      "\n",
      "int32_array = pointer(array_type(int_type))\n",
      "int64_array = pointer(array_type(Type.int(64)))\n",
      "double_array = pointer(array_type(double_type))\n",
      "\n",
      "lltypes_map = {\n",
      "    int32          : int_type,\n",
      "    int64          : int_type,\n",
      "    float32        : float_type,\n",
      "    double64       : double_type,\n",
      "    array_int32    : int32_array,\n",
      "    array_int64    : int64_array,\n",
      "    array_double64 : double_array\n",
      "}\n",
      "\n",
      "def to_lltype(ptype):\n",
      "    return lltypes_map[ptype]\n",
      "\n",
      "def determined(ty):\n",
      "    return len(ftv(ty)) == 0"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 28
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now the meat of the whole system is the LLVMEmitter class, which is a few hundred lines. Effectively we create a LLVM builder upon initialization and then traverse through our core AST. The important functions are:\n",
      "\n",
      "* **start_function**: Creates the initial basic block structure. Shifts the instruction \"cursor\" to the first block and then starts running through each of the statements in the body of the function to add logic.\n",
      "* **add_block**: Creates a new basic block.\n",
      "* **set_block**: Sets the active basic block that we are adding instructions to.\n",
      "* **specialize**: Extracts the type of the subexpression from the AST and maps our custom type into a LLVM type.\n",
      "\n",
      "The metadata for all array arguments is automatically stack allocated in the entry block so that subsequent accesses just have to look at the constant ``load``'d values. These are stored in the **arrays** dictionary which holds all NumPy array arguments and their metadata.\n",
      "\n",
      "The special **retval** reference holds the return value that the function will yield when the **exit_block**.\n",
      "in\n",
      "Whenever a name binder occurs we will look the AST, which is likely a type variable given to us from the inference engine. Since our type signature is fully determiend at this point we then need only look in the **spec_types** dictionary for what concrete type this subexpression has."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class LLVMEmitter(object):\n",
      "    def __init__(self, spec_types, retty, argtys):\n",
      "        self.function = None             # LLVM Function\n",
      "        self.builder = None              # LLVM Builder\n",
      "        self.locals = {}                 # Local variables\n",
      "        self.arrays = defaultdict(dict)  # Array metadata\n",
      "        self.exit_block = None           # Exit block\n",
      "        self.spec_types = spec_types     # Type specialization\n",
      "        self.retty = retty               # Return type\n",
      "        self.argtys = argtys             # Argument types\n",
      "\n",
      "    def start_function(self, name, module, rettype, argtypes):\n",
      "        func_type = lc.Type.function(rettype, argtypes, False)\n",
      "        function = lc.Function.new(module, func_type, name)\n",
      "        entry_block = function.append_basic_block(\"entry\")\n",
      "        builder = lc.Builder.new(entry_block)\n",
      "        self.exit_block = function.append_basic_block(\"exit\")\n",
      "        self.function = function\n",
      "        self.builder = builder\n",
      "\n",
      "    def end_function(self):\n",
      "        self.builder.position_at_end(self.exit_block)\n",
      "\n",
      "        if 'retval' in self.locals:\n",
      "            retval = self.builder.load(self.locals['retval'])\n",
      "            self.builder.ret(retval)\n",
      "        else:\n",
      "            self.builder.ret_void()\n",
      "\n",
      "    def add_block(self, name):\n",
      "        return self.function.append_basic_block(name)\n",
      "\n",
      "    def set_block(self, block):\n",
      "        self.block = block\n",
      "        self.builder.position_at_end(block)\n",
      "\n",
      "    def cbranch(self, cond, true_block, false_block):\n",
      "        self.builder.cbranch(cond, true_block, false_block)\n",
      "\n",
      "    def branch(self, next_block):\n",
      "        self.builder.branch(next_block)\n",
      "\n",
      "    def specialize(self, val):\n",
      "        if isinstance(val.type, TVar):\n",
      "            return to_lltype(self.spec_types[val.type.s])\n",
      "        else:\n",
      "            return val.type\n",
      "\n",
      "    def const(self, val):\n",
      "        if isinstance(val, (int, long)):\n",
      "            return Constant.int(int_type, val)\n",
      "        elif isinstance(val, float):\n",
      "            return Constant.real(double_type, val)\n",
      "        elif isinstance(val, bool):\n",
      "            return Constant.int(bool_type, int(val))\n",
      "        elif isinstance(val, str):\n",
      "            return Constant.stringz(val)\n",
      "        else:\n",
      "            raise NotImplementedError\n",
      "\n",
      "    def visit_LitInt(self, node):\n",
      "        ty = self.specialize(node)\n",
      "        if ty is double_type:\n",
      "            return Constant.real(double_type, node.n)\n",
      "        elif ty == int_type:\n",
      "            return Constant.int(int_type, node.n)\n",
      "\n",
      "    def visit_LitFloat(self, node):\n",
      "        ty = self.specialize(node)\n",
      "        if ty is double_type:\n",
      "            return Constant.real(double_type, node.n)\n",
      "        elif ty == int_type:\n",
      "            return Constant.int(int_type, node.n)\n",
      "\n",
      "    def visit_Noop(self, node):\n",
      "        pass\n",
      "\n",
      "    def visit_Fun(self, node):\n",
      "        rettype = to_lltype(self.retty)\n",
      "        argtypes = map(to_lltype, self.argtys)\n",
      "        # Create a unique specialized name\n",
      "        func_name = mangler(node.fname, self.argtys)\n",
      "        self.start_function(func_name, module, rettype, argtypes)\n",
      "\n",
      "        for (ar, llarg, argty) in zip(node.args, self.function.args, self.argtys):\n",
      "            name = ar.id\n",
      "            llarg.name = name\n",
      "\n",
      "            if is_array(argty):\n",
      "                zero = self.const(0)\n",
      "                one = self.const(1)\n",
      "                two = self.const(2)\n",
      "\n",
      "                data = self.builder.gep(llarg, [\n",
      "                                        zero, zero], name=(name + '_data'))\n",
      "                dims = self.builder.gep(llarg, [\n",
      "                                        zero, one], name=(name + '_dims'))\n",
      "                shape = self.builder.gep(llarg, [\n",
      "                                         zero, two], name=(name + '_strides'))\n",
      "\n",
      "                self.arrays[name]['data'] = self.builder.load(data)\n",
      "                self.arrays[name]['dims'] = self.builder.load(dims)\n",
      "                self.arrays[name]['shape'] = self.builder.load(shape)\n",
      "                self.locals[name] = llarg\n",
      "            else:\n",
      "                argref = self.builder.alloca(to_lltype(argty))\n",
      "                self.builder.store(llarg, argref)\n",
      "                self.locals[name] = argref\n",
      "\n",
      "        # Setup the register for return type.\n",
      "        if rettype is not void_type:\n",
      "            self.locals['retval'] = self.builder.alloca(rettype, \"retval\")\n",
      "\n",
      "        map(self.visit, node.body)\n",
      "        self.end_function()\n",
      "\n",
      "    def visit_Index(self, node):\n",
      "        if isinstance(node.val, Var) and node.val.id in self.arrays:\n",
      "            val = self.visit(node.val)\n",
      "            ix = self.visit(node.ix)\n",
      "            dataptr = self.arrays[node.val.id]['data']\n",
      "            ret = self.builder.gep(dataptr, [ix])\n",
      "            return self.builder.load(ret)\n",
      "        else:\n",
      "            val = self.visit(node.val)\n",
      "            ix = self.visit(node.ix)\n",
      "            ret = self.builder.gep(val, [ix])\n",
      "            return self.builder.load(ret)\n",
      "\n",
      "    def visit_Var(self, node):\n",
      "        return self.builder.load(self.locals[node.id])\n",
      "\n",
      "    def visit_Return(self, node):\n",
      "        val = self.visit(node.val)\n",
      "        if val.type != void_type:\n",
      "            self.builder.store(val, self.locals['retval'])\n",
      "        self.builder.branch(self.exit_block)\n",
      "\n",
      "    def visit_Loop(self, node):\n",
      "        init_block = self.function.append_basic_block('for.init')\n",
      "        test_block = self.function.append_basic_block('for.cond')\n",
      "        body_block = self.function.append_basic_block('for.body')\n",
      "        end_block = self.function.append_basic_block(\"for.end\")\n",
      "\n",
      "        self.branch(init_block)\n",
      "        self.set_block(init_block)\n",
      "\n",
      "        start = self.visit(node.begin)\n",
      "        stop = self.visit(node.end)\n",
      "        step = 1\n",
      "\n",
      "        # Setup the increment variable\n",
      "        varname = node.var.id\n",
      "        inc = self.builder.alloca(int_type, varname)\n",
      "        self.builder.store(start, inc)\n",
      "        self.locals[varname] = inc\n",
      "\n",
      "        # Setup the loop condition\n",
      "        self.branch(test_block)\n",
      "        self.set_block(test_block)\n",
      "        cond = self.builder.icmp(lc.ICMP_SLT, self.builder.load(inc), stop)\n",
      "        self.builder.cbranch(cond, body_block, end_block)\n",
      "\n",
      "        # Generate the loop body\n",
      "        self.set_block(body_block)\n",
      "        map(self.visit, node.body)\n",
      "\n",
      "        # Increment the counter\n",
      "        succ = self.builder.add(self.const(step), self.builder.load(inc))\n",
      "        self.builder.store(succ, inc)\n",
      "\n",
      "        # Exit the loop\n",
      "        self.builder.branch(test_block)\n",
      "        self.set_block(end_block)\n",
      "\n",
      "    def visit_Prim(self, node):\n",
      "        if node.fn == \"shape#\":\n",
      "            ref = node.args[0]\n",
      "            shape = self.arrays[ref.id]['shape']\n",
      "            return shape\n",
      "        elif node.fn == \"mult#\":\n",
      "            a = self.visit(node.args[0])\n",
      "            b = self.visit(node.args[1])\n",
      "            if a.type == double_type:\n",
      "                return self.builder.fmul(a, b)\n",
      "            else:\n",
      "                return self.builder.mul(a, b)\n",
      "        elif node.fn == \"add#\":\n",
      "            a = self.visit(node.args[0])\n",
      "            b = self.visit(node.args[1])\n",
      "            if a.type == double_type:\n",
      "                return self.builder.fadd(a, b)\n",
      "            else:\n",
      "                return self.builder.add(a, b)\n",
      "        else:\n",
      "            raise NotImplementedError\n",
      "\n",
      "    def visit_Assign(self, node):\n",
      "        # Subsequent assignment\n",
      "        if node.ref in self.locals:\n",
      "            name = node.ref\n",
      "            var = self.locals[name]\n",
      "            val = self.visit(node.val)\n",
      "            self.builder.store(val, var)\n",
      "            self.locals[name] = var\n",
      "            return var\n",
      "\n",
      "        # First assignment\n",
      "        else:\n",
      "            name = node.ref\n",
      "            val = self.visit(node.val)\n",
      "            ty = self.specialize(node)\n",
      "            var = self.builder.alloca(ty, name)\n",
      "            self.builder.store(val, var)\n",
      "            self.locals[name] = var\n",
      "            return var\n",
      "\n",
      "    def visit(self, node):\n",
      "        name = \"visit_%s\" % type(node).__name__\n",
      "        if hasattr(self, name):\n",
      "            return getattr(self, name)(node)\n",
      "        else:\n",
      "            return self.generic_visit(node)\n",
      "\n",
      "    def generic_visit(self, node):\n",
      "        raise NotImplementedError"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 29
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "This class may look big, but a lot of it is actually just the same logic over and over. The only non-trivial bit is the loop which is really just simple four basic blocks that jump between each other based on a loop condition just like the simple ``count`` example from the first section. If we graph the control flow for our loop constuctor it looks like:\n",
      "\n",
      "\n",
      "<img src=\"files/images/dot.png\">"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "So as not to duplicate work we'll create a unique mangled name for each function that is defined in terms of the hash of it's argument types. Every autojit'd function can map onto several mangled LLVM functions in the current module. This guarantees that names don't clash. It also gives us a way to cache on the argument types so that functions will not get recompiled and reJIT'd if the arguments given are identical to a function that has previously run."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def mangler(fname, sig):\n",
      "    return fname + str(hash(tuple(sig)))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 30
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now to actually invoke our function we'll use the ExecutionEngine as before, but we'd like to able to seamlessly go back and forth between Python/NumPy types without having to manually convert. To do this we'll use the ctypes/libffi wrapper to automatically lower the Python types into their C equivelants. Hat tip to Dave Beazley for documenting this technique in the Python Cookbook."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "_nptypemap = {\n",
      "    'i': ctypes.c_int,\n",
      "    'f': ctypes.c_float,\n",
      "    'd': ctypes.c_double,\n",
      "}\n",
      "\n",
      "def wrap_module(sig, llfunc):\n",
      "    pfunc = wrap_function(llfunc, engine)\n",
      "    dispatch = dispatcher(pfunc)\n",
      "    return dispatch\n",
      "\n",
      "def wrap_function(func, engine):\n",
      "    args = func.type.pointee.args\n",
      "    ret_type = func.type.pointee.return_type\n",
      "    ret_ctype = wrap_type(ret_type)\n",
      "    args_ctypes = map(wrap_type, args)\n",
      "\n",
      "    functype = ctypes.CFUNCTYPE(ret_ctype, *args_ctypes)\n",
      "    fptr = engine.get_pointer_to_function(func)\n",
      "\n",
      "    cfunc = functype(fptr)\n",
      "    cfunc.__name__ = func.name\n",
      "    return cfunc\n",
      "\n",
      "def wrap_type(llvm_type):\n",
      "    kind = llvm_type.kind\n",
      "    if kind == lc.TYPE_INTEGER:\n",
      "        ctype = getattr(ctypes, \"c_int\"+str(llvm_type.width))\n",
      "    elif kind == lc.TYPE_DOUBLE:\n",
      "        ctype = ctypes.c_double\n",
      "    elif kind == lc.TYPE_FLOAT:\n",
      "        ctype = ctypes.c_float\n",
      "    elif kind == lc.TYPE_VOID:\n",
      "        ctype = None\n",
      "    elif kind == lc.TYPE_POINTER:\n",
      "        pointee = llvm_type.pointee\n",
      "        p_kind = pointee.kind\n",
      "        if p_kind == lc.TYPE_INTEGER:\n",
      "            width = pointee.width\n",
      "            if width == 8:\n",
      "                ctype = ctypes.c_char_p\n",
      "            else:\n",
      "                ctype = ctypes.POINTER(wrap_type(pointee))\n",
      "        elif p_kind == lc.TYPE_VOID:\n",
      "            ctype = ctypes.c_void_p\n",
      "        else:\n",
      "            ctype = ctypes.POINTER(wrap_type(pointee))\n",
      "    elif kind == lc.TYPE_STRUCT:\n",
      "        struct_name = llvm_type.name.split('.')[-1]\n",
      "        struct_name = struct_name.encode('ascii')\n",
      "        struct_type = None\n",
      "\n",
      "        if struct_type and issubclass(struct_type, ctypes.Structure):\n",
      "            return struct_type\n",
      "\n",
      "        if hasattr(struct_type, '_fields_'):\n",
      "            names = struct_type._fields_\n",
      "        else:\n",
      "            names = [\"field\"+str(n) for n in range(llvm_type.element_count)]\n",
      "\n",
      "        ctype = type(ctypes.Structure)(struct_name, (ctypes.Structure,),\n",
      "                                       {'__module__': \"numpile\"})\n",
      "\n",
      "        fields = [(name, wrap_type(elem))\n",
      "                  for name, elem in zip(names, llvm_type.elements)]\n",
      "        setattr(ctype, '_fields_', fields)\n",
      "    else:\n",
      "        raise Exception(\"Unknown LLVM type %s\" % kind)\n",
      "    return ctype\n",
      "\n",
      "def wrap_ndarray(na):\n",
      "    # For NumPy arrays grab the underlying data pointer. Doesn't copy.\n",
      "    ctype = _nptypemap[na.dtype.char]\n",
      "    _shape = list(na.shape)\n",
      "    data = na.ctypes.data_as(ctypes.POINTER(ctype))\n",
      "    dims = len(na.strides)\n",
      "    shape = (ctypes.c_int*dims)(*_shape)\n",
      "    return (data, dims, shape)\n",
      "\n",
      "def wrap_arg(arg, val):\n",
      "    if isinstance(val, np.ndarray):\n",
      "        ndarray = arg._type_\n",
      "        data, dims, shape = wrap_ndarray(val)\n",
      "        return ndarray(data, dims, shape)\n",
      "    else:\n",
      "        return val\n",
      "\n",
      "def dispatcher(fn):\n",
      "    def _call_closure(*args):\n",
      "        cargs = list(fn._argtypes_)\n",
      "        pargs = list(args)\n",
      "        rargs = map(wrap_arg, cargs, pargs)\n",
      "        return fn(*rargs)\n",
      "    _call_closure.__name__ = fn.__name__\n",
      "    return _call_closure"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 31
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Toplevel\n",
      "\n",
      "The toplevel will consists of the ``autojit`` decorator which maps the function through translator, does type inference, and the creates a closure which when called will automatically specialize the function to the given argument types and compile a new version if needed. We will cache based on the arguments ( which entirely define the function ) and whenever a similar typed argument set is passed we just lookup the preJIT'd function and invoke it with no overhead."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "module = lc.Module.new('numpile.module')\n",
      "engine = None\n",
      "function_cache = {}\n",
      "\n",
      "tm = le.TargetMachine.new(features='', cm=le.CM_JITDEFAULT)\n",
      "eb = le.EngineBuilder.new(module)\n",
      "engine = eb.create(tm)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 32
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def autojit(fn):\n",
      "    transformer = PythonVisitor()\n",
      "    ast = transformer(fn)\n",
      "    (ty, mgu) = typeinfer(ast)\n",
      "    return specialize(ast, ty, mgu)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 33
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def typeinfer(ast):\n",
      "    infer = TypeInfer()\n",
      "    ty = infer.visit(ast)\n",
      "    mgu = solve(infer.constraints)\n",
      "    infer_ty = apply(mgu, ty)\n",
      "    return (infer_ty, mgu)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 34
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def codegen(ast, specializer, retty, argtys):\n",
      "    cgen = LLVMEmitter(specializer, retty, argtys)\n",
      "    mod = cgen.visit(ast)\n",
      "    cgen.function.verify()\n",
      "    print(cgen.function)\n",
      "    print(module.to_native_assembly())\n",
      "    return cgen.function"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 35
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "And finally the argument specializer logic."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def arg_pytype(arg):\n",
      "    if isinstance(arg, np.ndarray):\n",
      "        if arg.dtype == np.dtype('int32'):\n",
      "            return array(int32)\n",
      "        elif arg.dtype == np.dtype('int64'):\n",
      "            return array(int64)\n",
      "        elif arg.dtype == np.dtype('double'):\n",
      "            return array(double64)\n",
      "        elif arg.dtype == np.dtype('float'):\n",
      "            return array(float32)\n",
      "    elif isinstance(arg, int) & (arg < sys.maxint):\n",
      "        return int64\n",
      "    elif isinstance(arg, float):\n",
      "        return double64\n",
      "    else:\n",
      "        raise Exception(\"Type not supported: %s\" % type(arg))\n",
      "\n",
      "def specialize(ast, infer_ty, mgu):\n",
      "    def _wrapper(*args):\n",
      "        types = map(arg_pytype, list(args))\n",
      "        spec_ty = TFun(argtys=types, retty=TVar(\"$retty\"))\n",
      "        unifier = unify(infer_ty, spec_ty)\n",
      "        specializer = compose(unifier, mgu)\n",
      "\n",
      "        retty = apply(specializer, TVar(\"$retty\"))\n",
      "        argtys = [apply(specializer, ty) for ty in types]\n",
      "        print('Specialized Function:', TFun(argtys, retty))\n",
      "\n",
      "        if determined(retty) and all(map(determined, argtys)):\n",
      "            key = mangler(ast.fname, argtys)\n",
      "            # Don't recompile after we've specialized.\n",
      "            if key in function_cache:\n",
      "                return function_cache[key](*args)\n",
      "            else:\n",
      "                llfunc = codegen(ast, specializer, retty, argtys)\n",
      "                pyfunc = wrap_module(argtys, llfunc)\n",
      "                function_cache[key] = pyfunc\n",
      "                return pyfunc(*args)\n",
      "        else:\n",
      "            raise UnderDetermined()\n",
      "    return _wrapper"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 36
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**OK, so basically we're done**, we built the thing top to bottom so let's try it out. Keep in mind that this IR is without optimizations so it will do several naive things that the optimizer will clean up later."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "@autojit\n",
      "def add(a, b):\n",
      "    return a+b\n",
      "\n",
      "a = 3.1415926\n",
      "b = 2.7182818\n",
      "print('Result:', add(a,b))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Specialized Function: [Double, Double] -> Double\n",
        "\n",
        "define double @add4531207233431041901(double %a, double %b) {\n",
        "entry:\n",
        "  %0 = alloca double\n",
        "  store double %a, double* %0\n",
        "  %1 = alloca double\n",
        "  store double %b, double* %1\n",
        "  %retval = alloca double\n",
        "  %2 = load double* %0\n",
        "  %3 = load double* %1\n",
        "  %4 = fadd double %2, %3\n",
        "  store double %4, double* %retval\n",
        "  br label %exit\n",
        "\n",
        "exit:                                             ; preds = %entry\n",
        "  %5 = load double* %retval\n",
        "  ret double %5\n",
        "}\n",
        "\n",
        "\t.file\t\"numpile.module\"\n",
        "\t.text\n",
        "\t.globl\tadd4531207233431041901\n",
        "\t.align\t16, 0x90\n",
        "\t.type\tadd4531207233431041901,@function\n",
        "add4531207233431041901:\n",
        "\t.cfi_startproc\n",
        "\tvmovsd\t%xmm0, -8(%rsp)\n",
        "\tvmovsd\t%xmm1, -16(%rsp)\n",
        "\tvaddsd\t-8(%rsp), %xmm1, %xmm0\n",
        "\tvmovsd\t%xmm0, -24(%rsp)\n",
        "\tvmovsd\t-24(%rsp), %xmm0\n",
        "\tret\n",
        ".Ltmp0:\n",
        "\t.size\tadd4531207233431041901, .Ltmp0-add4531207233431041901\n",
        "\t.cfi_endproc\n",
        "\n",
        "\n",
        "\t.section\t\".note.GNU-stack\",\"\",@progbits\n",
        "\n",
        "Result: 5.8598744\n"
       ]
      }
     ],
     "prompt_number": 37
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "And how about for our dot product function."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "@autojit\n",
      "def dot(a, b):\n",
      "    c = 0\n",
      "    n = a.shape[0]\n",
      "    for i in range(n):\n",
      "       c += a[i]*b[i]\n",
      "    return c"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 38
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We'll get a lot of debug output for this one."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "a = np.array(range(1000,2000), dtype='int32')\n",
      "b = np.array(range(3000,4000), dtype='int32')\n",
      "\n",
      "print('Result:', dot(a,b))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Specialized Function: [Array Int32, Array Int32] -> Int32\n",
        "\n",
        "define i32 @dot-7244935599725600953(%ndarray_i32* %a, %ndarray_i32* %b) {\n",
        "entry:\n",
        "  %a_data = getelementptr %ndarray_i32* %a, i32 0, i32 0\n",
        "  %a_dims = getelementptr %ndarray_i32* %a, i32 0, i32 1\n",
        "  %a_strides = getelementptr %ndarray_i32* %a, i32 0, i32 2\n",
        "  %0 = load i32** %a_data\n",
        "  %1 = load i32* %a_dims\n",
        "  %2 = load i32** %a_strides\n",
        "  %b_data = getelementptr %ndarray_i32* %b, i32 0, i32 0\n",
        "  %b_dims = getelementptr %ndarray_i32* %b, i32 0, i32 1\n",
        "  %b_strides = getelementptr %ndarray_i32* %b, i32 0, i32 2\n",
        "  %3 = load i32** %b_data\n",
        "  %4 = load i32* %b_dims\n",
        "  %5 = load i32** %b_strides\n",
        "  %retval = alloca i32\n",
        "  %c = alloca i32\n",
        "  store i32 0, i32* %c\n",
        "  %6 = getelementptr i32* %2, i32 0\n",
        "  %7 = load i32* %6\n",
        "  %n = alloca i32\n",
        "  store i32 %7, i32* %n\n",
        "  br label %for.init\n",
        "\n",
        "exit:                                             ; preds = %for.end\n",
        "  %8 = load i32* %retval\n",
        "  ret i32 %8\n",
        "\n",
        "for.init:                                         ; preds = %entry\n",
        "  %9 = load i32* %n\n",
        "  %i = alloca i32\n",
        "  store i32 0, i32* %i\n",
        "  br label %for.cond\n",
        "\n",
        "for.cond:                                         ; preds = %for.body, %for.init\n",
        "  %10 = load i32* %i\n",
        "  %11 = icmp slt i32 %10, %9\n",
        "  br i1 %11, label %for.body, label %for.end\n",
        "\n",
        "for.body:                                         ; preds = %for.cond\n",
        "  %12 = load i32* %c\n",
        "  %13 = load %ndarray_i32* %a\n",
        "  %14 = load i32* %i\n",
        "  %15 = getelementptr i32* %0, i32 %14\n",
        "  %16 = load i32* %15\n",
        "  %17 = load %ndarray_i32* %b\n",
        "  %18 = load i32* %i\n",
        "  %19 = getelementptr i32* %3, i32 %18\n",
        "  %20 = load i32* %19\n",
        "  %21 = mul i32 %16, %20\n",
        "  %22 = add i32 %12, %21\n",
        "  store i32 %22, i32* %c\n",
        "  %23 = load i32* %i\n",
        "  %24 = add i32 1, %23\n",
        "  store i32 %24, i32* %i\n",
        "  br label %for.cond\n",
        "\n",
        "for.end:                                          ; preds = %for.cond\n",
        "  %25 = load i32* %c\n",
        "  store i32 %25, i32* %retval\n",
        "  br label %exit\n",
        "}\n",
        "\n",
        "\t.file\t\"numpile.module\"\n",
        "\t.text\n",
        "\t.globl\tadd4531207233431041901\n",
        "\t.align\t16, 0x90\n",
        "\t.type\tadd4531207233431041901,@function\n",
        "add4531207233431041901:\n",
        "\t.cfi_startproc\n",
        "\tvmovsd\t%xmm0, -8(%rsp)\n",
        "\tvmovsd\t%xmm1, -16(%rsp)\n",
        "\tvaddsd\t-8(%rsp), %xmm1, %xmm0\n",
        "\tvmovsd\t%xmm0, -24(%rsp)\n",
        "\tvmovsd\t-24(%rsp), %xmm0\n",
        "\tret\n",
        ".Ltmp0:\n",
        "\t.size\tadd4531207233431041901, .Ltmp0-add4531207233431041901\n",
        "\t.cfi_endproc\n",
        "\n",
        "\t.globl\tdot_2D_7244935599725600953\n",
        "\t.align\t16, 0x90\n",
        "\t.type\tdot_2D_7244935599725600953,@function\n",
        "dot_2D_7244935599725600953:\n",
        "\t.cfi_startproc\n",
        "\tmovq\t(%rsi), %rax\n",
        "\tmovq\t(%rdi), %rcx\n",
        "\tmovq\t16(%rdi), %rdx\n",
        "\tmovl\t$0, -8(%rsp)\n",
        "\tmovl\t(%rdx), %edx\n",
        "\tmovl\t%edx, -12(%rsp)\n",
        "\tmovl\t$0, -16(%rsp)\n",
        "\tjmp\t.LBB1_1\n",
        "\t.align\t16, 0x90\n",
        ".LBB1_2:\n",
        "\tmovslq\t-16(%rsp), %rsi\n",
        "\tmovl\t(%rcx,%rsi,4), %edi\n",
        "\timull\t(%rax,%rsi,4), %edi\n",
        "\taddl\t%edi, -8(%rsp)\n",
        "\tincl\t-16(%rsp)\n",
        ".LBB1_1:\n",
        "\tcmpl\t%edx, -16(%rsp)\n",
        "\tjl\t.LBB1_2\n",
        "\tmovl\t-8(%rsp), %eax\n",
        "\tmovl\t%eax, -4(%rsp)\n",
        "\tret\n",
        ".Ltmp1:\n",
        "\t.size\tdot_2D_7244935599725600953, .Ltmp1-dot_2D_7244935599725600953\n",
        "\t.cfi_endproc\n",
        "\n",
        "\n",
        "\t.section\t\".note.GNU-stack\",\"\",@progbits\n",
        "\n",
        "Result: 1035866204\n"
       ]
      }
     ],
     "prompt_number": 39
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Ok, now let's turn the optimizer on and have it have it automatically transform not only our naive code, but replace most of our inner loop with more optimial instructions."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def codegen(ast, specializer, retty, argtys):\n",
      "    cgen = LLVMEmitter(specializer, retty, argtys)\n",
      "    mod = cgen.visit(ast)\n",
      "    cgen.function.verify()\n",
      "\n",
      "    tm = le.TargetMachine.new(opt=3, cm=le.CM_JITDEFAULT, features='')\n",
      "    pms = lp.build_pass_managers(tm=tm,\n",
      "                                 fpm=False,\n",
      "                                 mod=module,\n",
      "                                 opt=3,\n",
      "                                 vectorize=False,\n",
      "                                 loop_vectorize=True)\n",
      "    pms.pm.run(module)\n",
      "\n",
      "    print(cgen.function)\n",
      "    return cgen.function"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 40
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "@autojit\n",
      "def dot_vectorize(a, b):\n",
      "    c = 0\n",
      "    n = a.shape[0]\n",
      "    for i in range(n):\n",
      "       c += a[i]*b[i]\n",
      "    return c"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 41
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "With the optimizer in full force LLVM has replaced most of our loop with SIMD and vector instructions and partially unrolled the loops for the dot product, as well as doing the usual dead code elimination, control flow simplification."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "a = np.array(range(1000,2000), dtype='int32')\n",
      "b = np.array(range(3000,4000), dtype='int32')\n",
      "\n",
      "print('Result:', dot_vectorize(a,b))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Specialized Function: [Array Int32, Array Int32] -> Int32\n",
        "\n",
        "define i32 @dot_vectorize-7244935599725600953(%ndarray_i32* nocapture %a, %ndarray_i32* nocapture %b) nounwind readonly {\n",
        "entry:\n",
        "  %a_data = getelementptr %ndarray_i32* %a, i64 0, i32 0\n",
        "  %a_strides = getelementptr %ndarray_i32* %a, i64 0, i32 2\n",
        "  %0 = load i32** %a_data, align 8\n",
        "  %1 = load i32** %a_strides, align 8\n",
        "  %b_data = getelementptr %ndarray_i32* %b, i64 0, i32 0\n",
        "  %2 = load i32** %b_data, align 8\n",
        "  %3 = load i32* %1, align 4\n",
        "  %4 = icmp sgt i32 %3, 0\n",
        "  br i1 %4, label %for.body.lr.ph, label %for.end\n",
        "\n",
        "for.body.lr.ph:                                   ; preds = %entry\n",
        "  %5 = zext i32 %3 to i64\n",
        "  %n.vec = and i64 %5, 4294967288\n",
        "  %cmp.zero = icmp eq i64 %n.vec, 0\n",
        "  br i1 %cmp.zero, label %middle.block, label %vector.body\n",
        "\n",
        "vector.body:                                      ; preds = %for.body.lr.ph, %vector.body\n",
        "  %index = phi i64 [ %index.next, %vector.body ], [ 0, %for.body.lr.ph ]\n",
        "  %vec.phi = phi <8 x i32> [ %13, %vector.body ], [ zeroinitializer, %for.body.lr.ph ]\n",
        "  %6 = getelementptr i32* %0, i64 %index\n",
        "  %7 = bitcast i32* %6 to <8 x i32>*\n",
        "  %8 = load <8 x i32>* %7, align 4\n",
        "  %9 = getelementptr i32* %2, i64 %index\n",
        "  %10 = bitcast i32* %9 to <8 x i32>*\n",
        "  %11 = load <8 x i32>* %10, align 4\n",
        "  %12 = mul <8 x i32> %11, %8\n",
        "  %13 = add <8 x i32> %12, %vec.phi\n",
        "  %index.next = add i64 %index, 8\n",
        "  %14 = icmp eq i64 %index.next, %n.vec\n",
        "  br i1 %14, label %middle.block, label %vector.body\n",
        "\n",
        "middle.block:                                     ; preds = %vector.body, %for.body.lr.ph\n",
        "  %resume.idx = phi i64 [ 0, %for.body.lr.ph ], [ %n.vec, %vector.body ]\n",
        "  %rdx.vec.exit.phi = phi <8 x i32> [ zeroinitializer, %for.body.lr.ph ], [ %13, %vector.body ]\n",
        "  %15 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 0\n",
        "  %16 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 1\n",
        "  %17 = add i32 %15, %16\n",
        "  %18 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 2\n",
        "  %19 = add i32 %17, %18\n",
        "  %20 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 3\n",
        "  %21 = add i32 %19, %20\n",
        "  %22 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 4\n",
        "  %23 = add i32 %21, %22\n",
        "  %24 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 5\n",
        "  %25 = add i32 %23, %24\n",
        "  %26 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 6\n",
        "  %27 = add i32 %25, %26\n",
        "  %28 = extractelement <8 x i32> %rdx.vec.exit.phi, i32 7\n",
        "  %29 = add i32 %27, %28\n",
        "  %cmp.n = icmp eq i64 %5, %resume.idx\n",
        "  br i1 %cmp.n, label %for.end, label %for.body\n",
        "\n",
        "for.body:                                         ; preds = %middle.block, %for.body\n",
        "  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ %resume.idx, %middle.block ]\n",
        "  %30 = phi i32 [ %36, %for.body ], [ %29, %middle.block ]\n",
        "  %31 = getelementptr i32* %0, i64 %indvars.iv\n",
        "  %32 = load i32* %31, align 4\n",
        "  %33 = getelementptr i32* %2, i64 %indvars.iv\n",
        "  %34 = load i32* %33, align 4\n",
        "  %35 = mul i32 %34, %32\n",
        "  %36 = add i32 %35, %30\n",
        "  %indvars.iv.next = add i64 %indvars.iv, 1\n",
        "  %lftr.wideiv = trunc i64 %indvars.iv.next to i32\n",
        "  %exitcond = icmp eq i32 %lftr.wideiv, %3\n",
        "  br i1 %exitcond, label %for.end, label %for.body\n",
        "\n",
        "for.end:                                          ; preds = %middle.block, %for.body, %entry\n",
        "  %.lcssa = phi i32 [ 0, %entry ], [ %29, %middle.block ], [ %36, %for.body ]\n",
        "  ret i32 %.lcssa\n",
        "}\n",
        "\n",
        "Result: 1035866204\n"
       ]
      }
     ],
     "prompt_number": 42
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Further Work\n",
      "\n",
      "While this example is kind of simplified (we only have addition and multiplication after all!), in principle all the ideas and machinary you would need to build out a full system are basically sketched here. Some further fruitful areas:\n",
      "\n",
      "* Translate all of the Python AST\n",
      "* Use [subpy](https://github.com/sdiehl/subpy/tree/master/subpy) to do feature\n",
      "  detection before lowering the function into LLVM. Gives better error reporting\n",
      "  when an invalid high-level feature is used instead of failing somewhere in the\n",
      "  middle of the compiler pipeline.\n",
      "* Use the lineno and column information on the Python AST to add better error handling.\n",
      "* Add more types and explicit casts or bidirectional type inference.\n",
      "* For untranslatable calls, use the Object Layer in the C-API to reach back into the interpreter.\n",
      "* Map a subset of numpy calls into LLVM.\n",
      "* Use the LLVM nvptx backend for targeting Python logic into Nvidia CUDA\n",
      "  kernels.\n",
      "* Use pthreads inside of LLVM to logic write multicore programs withou the usual GIL constraints.\n",
      "* Use MPI or ZeroMQ's zero-copy array transfer efficiently push distributed\n",
      "  numerical kernels across multiple machine.\n",
      "* Add composite numerical types, complex numbers, quaternions.\n"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Resources\n",
      "\n",
      "I've written about LLVM quite a bit lately. \n",
      "\n",
      "* [Implementing a JIT Compiled Language with Haskell and LLVM](http://www.stephendiehl.com/llvm)\n",
      "* [Write You a Haskell](http://dev.stephendiehl.com/fun/)"
     ]
    }
   ],
   "metadata": {}
  }
 ]
}