{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Python-moa (mathematics of arrays) is an approach to a high level tensor\n",
"compiler that is based on the work of [Lenore\n",
"Mullin](https://www.albany.edu/ceas/lenore-mullin.php) and her\n",
"[dissertation](https://www.researchgate.net/publication/308893116_A_Mathematics_of_Arrays).\n",
"A high level compiler is necessary because there are many optimizations\n",
"that a low level compiler such as `gcc` will miss. It is trying to solve\n",
"many of the same problems as other technologies such as the [taco\n",
"compiler](http://tensor-compiler.org/) and the [xla\n",
"compiler](https://www.tensorflow.org/xla). However, it takes a much\n",
"different approach than others guided by the following principles.\n",
"\n",
"1. What is the shape? Everything has a shape. scalars, vectors, arrays,\n",
" operations, and functions.\n",
"2. What are the given indicies and operations required to produce a\n",
" given index in the result?\n",
"\n",
"Having a compiler that is guided upon these principles allows for high\n",
"level reductions that other compilers will miss and allows for\n",
"optimization of algorithms as a whole. Keep in mind that MOA is **NOT**\n",
"a compiler. It is a theory that guides compiler development. Since\n",
"[python-moa](https://github.com/Quansight-Labs/python-moa) is based on\n",
"theory we get unique properties that other compilers cannot guarantee:\n",
"\n",
"\n",
"\n",
" - No out of bounds array accesses\n",
" - A computation is determined to be either valid or invalid at\n",
" compile time\n",
" - The computation will always reduce to a deterministic minimal form\n",
" (dnf) (see\n",
" [church-rosser](https://en.wikipedia.org/wiki/Church%E2%80%93Rosser_theorem)\n",
" property)\n",
" - All MOA operations are composable (including black box functions\n",
" and\n",
" [gufuncs](https://docs.scipy.org/doc/numpy-1.13.0/reference/c-api.generalized-ufuncs.html))\n",
" - Arbitrary high level operations will compile down to a minimal\n",
" backend instruction set. If the shape and indexing of a given\n",
" operation is known it can be added to python-moa.\n",
"\n",
"## Frontend"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Lenore Mullin originally developed a [moa\n",
"compiler](https://github.com/saulshanabrook/psi-compiler/) in the 90s\n",
"with programs that used a symbolic syntax heavily inspired by\n",
"[APL](https://en.wikipedia.org/wiki/APL_(programming_language))\n",
"([example\n",
"program](https://github.com/saulshanabrook/psi-compiler/blob/master/examples/ex1.m)).\n",
"This work was carried into python-moa initially with a lex/yacc compiler\n",
"with an example program below."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [],
"source": [
"from moa.frontend import parse\n",
"\n",
"context = parse('<0> psi (tran (A ^ + B ^ ))')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Upon pursuing this approach it became apparent that MOA should not\n",
"require that a new syntax be developed since it is only a theory! So a\n",
"pythonic interface to MOA was developed that expressed the same ideas\n",
"which look much like the current numeric python libraries. Ideally MOA\n",
"is hidden from the user. The python-moa compiler is broken into several\n",
"pieces each which their own responsibilities: shape, DNF, and ONF."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from moa.frontend import LazyArray\n",
"\n",
"A = LazyArray(shape=('n', 'm'), name='A')\n",
"B = LazyArray(shape=('k', 'l'), name='B')\n",
"\n",
"expression = (A + B).T.reduce('+')[0]\n",
"expression"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Shape Calculation\n",
"\n",
"The shape calculation is responsible for calculating the shape at every\n",
"step of the computation. This means that operations have a shape. Note\n",
"that the compiler handles symbolic shapes thus the exact shape does not\n",
"need to be known, only the dimension. After the shape calculation step\n",
"we can guarantee that an algorithm is a valid program and there will be\n",
"no out of bound memory accesses. Making MOA extremely compelling for\n",
"[FPGAs](https://en.wikipedia.org/wiki/Field-programmable_gate_array) and\n",
"compute units with a minimal OS. If an algorithm makes it past this\n",
"stage and fails then it is an issue with the compiler not the algorithm."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"expression.visualize(stage='shape')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Denotational Normal Form (DNF)\n",
"\n",
"The DNF\\'s responsibility is to reduce the high level MOA expression to\n",
"the minimal and optimal machine independent computation. This graph has\n",
"all of the indexing patterns of the computation and resulting shapes.\n",
"Notice that several operations disappear in this stage such a transpose\n",
"because transpose is simply index manipulation."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"expression.visualize(stage='dnf')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Operational Normal Form (ONF)\n",
"\n",
"The ONF is the stage of the compiler that will have to be the most\n",
"flexible. At its current stage the ONF is a naive compiler that does not\n",
"perform many important optimizations such as [PSI\n",
"reduction](https://www.researchgate.net/publication/264758384_Effective_data_parallel_computation_using_the_Psi_calculus)\n",
"which reduces the number of loops in the calculation, loop ordering, and\n",
"minimize the number of accumulators. MOA has ideas of dimension lifting\n",
"which make parallization and optimizing for cache sizes much easier.\n",
"Additionally algorithms must be implemented differently for sparse,\n",
"column major, row major. The ONF stage is responsible for any\n",
"\\\"optimal\\\" machine dependent implementation. \\\"optimal\\\" will vary from\n",
"user to user and thus will have to allow for multiple programs: optimal\n",
"single core, optimal parallel, optimal gpu, optimal low memory, etc."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"\n",
"@numba.jit\n",
"def f(A, B):\n",
" \n",
" \n",
" n = A.shape[0]\n",
" \n",
" m = A.shape[1]\n",
" \n",
" k = B.shape[0]\n",
" \n",
" l = B.shape[1]\n",
" \n",
" _a21 = numpy.zeros(())\n",
" \n",
" _a19 = numpy.zeros(())\n",
" \n",
" _a21 = 0\n",
" \n",
" for _i10 in range(0, m, 1):\n",
" \n",
" _a21 = (_a21 + (A[(0, _i10)] + B[(0, _i10)]))\n",
" \n",
" _a19[()] = _a21\n",
" return _a19\n"
]
}
],
"source": [
"print(expression.compile(use_numba=True, include_conditions=False))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Performance\n",
"\n",
"MOA excels at performing reductions and reducing the amount of actual\n",
"work done. You will see that the following algorithm only requires the\n",
"first index of the computation. Making the naive implementation `1000x`\n",
"more expensive for `1000x1000` shaped array. The following benchmarks\n",
"have been performed on my laptop with an intel i5-4200U. However, more\n",
"benchmarks are always available on the [Travis\n",
"CI](https://travis-ci.org/Quansight-Labs/python-moa) (these benchmarks\n",
"test python-moa\\'s weaknesses). You will see with the benchmarks that if\n",
"**any** indexing is required MOA will be significantly faster unless you\n",
"hand optimize the numerical computations."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [],
"source": [
"import numpy\n",
"import numba\n",
"\n",
"n, m = 1000, 1000\n",
"\n",
"exec(expression.compile(backend='python', use_numba=True, include_conditions=False))\n",
"\n",
"A = numpy.random.random((n, m))\n",
"B = numpy.random.random((n, m))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here we execute the MOA optimized code with the help of\n",
"[numba](https://github.com/numba/numba) which is a JIT LLVM compiler for\n",
"python."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2.14 µs ± 6.76 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
]
}
],
"source": [
"%%timeit\n",
"\n",
"f(A=A, B=B)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The following numpy computation is obviously the worst case expression\n",
"that you could write but this brings up the point that often times the\n",
"algorithm is expressed differently than the implementation. This is one\n",
"of the problems that MOA hopes to solve."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2.74 ms ± 29.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
"source": [
"%%timeit\n",
"\n",
"(A + B).T.sum(axis=0)[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We notice that even with the optimized version MOA is faster. This is\n",
"mostly due to the transpose operation the numpy performs that we have no\n",
"way around."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"name": "{.sourceCode"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"6.67 µs ± 91.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
]
}
],
"source": [
"%%timeit\n",
"\n",
"(A[0] + B[0]).T.sum(axis=0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Conclusions\n",
"-----------\n",
"\n",
"I hope that this walk through has shown the promising results that the\n",
"MOA theory can bring to tensor computations and the python ecosystem as\n",
"a whole. Please feel free to try out the project at [Quansight\n",
"Labs/python-moa](https://github.com/Quansight-Labs/python-moa). I hope\n",
"that this work can allow for the analysis and optimization of algorithms\n",
"in a mathematically rigorous way which allows users to express their\n",
"algorithms in an implementation independent manner."
]
}
],
"metadata": {
"jupytext": {
"main_language": "python",
"text_representation": {
"extension": ".md",
"format_name": "markdown"
}
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.3"
},
"nikola": {
"author": "Chris Ostrouchov",
"date": "2019-04-17",
"slug": "python-moa-tensor-compiler",
"title": "MOA: a theory for composable and verifiable tensor computations"
}
},
"nbformat": 4,
"nbformat_minor": 2
}