For-Loops Are Not a Frontend: Index Notation And The Future Of Array Compilers
The task of writing a complete library for linear algebra can be summarized as:
- For all linear algebra problems:
- For all matrix and problem structures:
- For all data types:
- For all architectures and networks:
- …
- Produce the best algorithms with respect to performance and accuracy UC Berkeley CS267 Class Notes
In Julia, the explosion in the number of array types and operations (both in terms of representation and function) has added even more dimensions of complexity to this problem. Automation is necessary to handle the software engineering complexity of a performant linear algebra compiler, and defining the inputs with a unified generic computational framework for array operations is the first step.
Index notation (Einstein summation) is a powerful representation of array operations that is general enough to encapsulate most nested loops on array accesses, but simple enough for humans and machines to understand and manipulate. With some small generalizations of the datatypes and operations, this notation can elegantly express important computations in a staggering range of applications. A few examples include machine learning, quantum chemistry, graph algorithms, data mining, and image processing.
Index notation is catching on as an API for array metaprogramming. Several emerging computational frameworks for linear algebra and tensor operations have have emerged that use index notation (or subsets thereof) to describe input operations. TACO (a compiler for operations on sparse tensors), Cyclops (A library for performing distributed tensor contractions), Facebook’s Tensor Comprehensions (A compiler for tensor expressions seen in deep learning applications) and the Tensor Contraction Engine (A library for dense tensor operations arising in computational chemistry) are just a few prominent examples of libraries which use index notation to describe input operations. Libraries for stencil computation (such as Halide) also often use syntax which can be thought of as a generalization of index notation.
By using this simple representation of array computations, we can separate the description of an operation from how the operation is actually implemented. When the operation is described in index notation, it becomes easy to optimize the implementation or specialize the whole operation for new or mixed array or element types. With small modifications, index notation can generalize broadcast, transpose, and reduce operations, providing more specific control of which axes should be manipulated. Index notation can also generalize Julia’s more complicated getindex and setindex! operations. This was originally recognized by the ArrayMeta.jl project, with the goal of unifying the implementation of an entire Julia Array type into a single generated function. Several Julia projects have used index notation as an interface to linear-algebraic metaprogramming, among them ArrayMeta.jl, TensorOperations.jl, and Tensors.jl.
Since index notation was created for use in scientific domains to describe tensor operations, few descriptions are targeted towards programmers. In this talk, I will give a programmer-centric description of index notation and it’s various flavors, showing how it can be used to describe some example operations. I will summarize how index notation separates mechanism and policy in the ecosystem of tensor implementations and compilers. Finally, I will explain how to transform and specialize index notation from the operation to the array types to the elementwise operations.
Speaker's bio
I write algorithms to improve interfaces and automate performance engineering tasks in scientific computing. When I get distracted, I enjoy woodworking, glassblowing, burritos, and climbing. I am a second year computer science graduate student at MIT, and Department of Energy Computational Science Graduate Fellow.
