A Tensor Language
This is a working space for sketching a tensor language. I am not placing a high degree of stress on portability right now. Please schedule time with me to iron out various portability issues if you want to get this working on Linux or Windows, your machine etc.
Dependencies
These are mostly documented (in simple form) in requirements.txt. However, note that the pysmt library requires an SMT solver to be installed on the machine external to the python dependency management. I have been using Z3, though you're welcome to see if you can get another solver working. To debug this please see pySMT's own documentation; the simplest path seems to make use of a command line tool called pysmt-install.
Early Notebooks
Most of the work has been sketched in notebooks and then copied into the src/ files with some degree of editing. Perhaps a final version should edit the notebooks up to match the source exactly for posterity's sake.
The notebooks ought to be read chronologically. That order is recorded here:
- How-to Represent IRs
- Memoization of IRs
- A Tensor Language v0
- Wrapping Halide Part 1
- Wrapping Halide Part 2
- Compiling A Tensor Language v0
I will now briefly gloss each of these.
How-to Represent IRs and Memoization of IRs
These notebooks generated src/adt.py, which is a tiny DSL designed to help compiler writers create type-checked Python class hierarchies out of descriptions that look like BNF grammars. The grammar language syntax is called ASDL, from the Princeton Zephyr project. It is used by the standard Python compiler to document Python's own IR.
A Tensor Language v0
This is the main course of the notebooks so far. It shows how to bootstrap up a language sketch using the ADT tools and Python hackery very quickly. This is my overriding modus operandi with this project: figure out how to get the big ideas across without writing much code. (see src/atlv0.py for the final result)
Wrapping Halide, Part 1 & 2
This is an exercise in how to pare down a whole build-system and library wrapper into less than 500 lines of code using clever meta-programming and filesystem inspection/introspection tricks. It exists to support a careful attempt to compile the tensor language efficiently. (see src/halide.py for the result of this notebook, and test_halide.py for a very few test cases present)
Compiling a Tensor Language v0
An experiment (quite long) in forcing the tensor language to compile down to Halide. The results are collected in src/atlv0_compile.py
Brief Overview of Compiler Organization
Preliminary Files
adt.py defines the ASDL support language/DSL for defining IRs.
prelude.py defines a bunch of useful helper objects/functions that are shared amongst the whole rest of the compiler
atl_types.py is part of the language proper, but should precede in include order basically anything else.
Main Compiler Flow (Front-end)
parse_pyast.py does the AST hijacking and generates a UST.function that gets wrapped in a Function object for the python interface. That Function (defined in functions.py) acts as the Python-facing API and also directs the rest of the compilation flow.
frontend.py defines UST and AST. The UST is generated by the parser, and then typechecked to generate the AST; typechecking is defined in this file.
bounds.py uses an SMT encoding to hit an ILP-SAT solver and check that all array accesses are in-bounds
Main Compiler Flow (Middle)
norm_ast.py implements the "pre-norm" passes, which lift all Let-bindings to the top level, pushes all tuples out to the edges (i.e. I/O) and moves all Gen-operators to the outer position of expressions, occuring as the outermost part of the right-hand-side of a let-binding.
Then, there is norm_ir.py which defines NIR, which is a lateral/same-level IR as AST (i.e. you can convert both ways between the two). A bunch of different stuff for NIR is defined in files named nir_*.py. NOTE: NIR is a DAG-based IR, not a Tree-based one.
nir_convert.py defines the mappings to and from AST. nir_simplify.py defines the memo-ized simplification that occurs automatically in a lot of places.
nir_filterdown.py is an additional kind of optimization not accounted for by nir_simplify.py but potentially important; likewise nir_factor.py (unimplmented right now) does factorization and will be undone by simplification. So it needs to be done (carefully) right before converting back to AST.
Derivative passes (Middle)
On NIR, both forward and reverse mode AD are defined in nir_deriv.py as two modes of a common pass object.
Forward-mode derivatives can also be computed on AST but not reverse-mode. That is defined in deriv_ast.py.
Main Compiler Flow (Back-ends)
Interpreters
The default reference is interpreter.py. It is not designed to be fast at all.
perf_analysis.py mimics the interpreter and also counts how many of each scalar operator are executed.
Halide
halide.py defines a wrapper for the Halide library in order to load it directly into process and use it as a JIT compiler for ATL.
halide_ir.py puts a spacer in-between ATL and the Halide backend, by providing an IR-proxy for Halide that can be executed in different ways. Right now it's using halide.py, but we want to have it alternately dump out to a C++ file and then compile and load that file. This will make it easier to debug for everyone, but especially for Halide developers who don't know anything about ATL and expect to see a C++ program using Halide.
ast_to_halide.py translates ATL (pre-normed) to the Halide IR.
JAX
jax_backend.py is an incomplete sketch of executing ATL code via JAX.
C
c_lang.py is an IR mimicing some aspects of C (a particular C--). It defines conversion of that IR to strings of C-code.
lower_to_c.py translates ATL into the IR defined by c_lang, and also supports a JIT-compilation/execution model which dumps those C-code strings into cached files, compiles them (escape to command-line) and then loads in the generated .so files via dynamic linking to be wrapped with the ctypes python library and called from Python.
Other surrounding support
functions.py defines a great deal of the Python interface, which is more specifically detailed by __init__.py too.
builtins.py defines an extensible set of language built-in functions, such as sin or cos and is not equivalent to a user-definable library.