Staged compilation of tensor expressions
Abstract
We present our current progress towards a metaprogramming framework for tensor expressions embedded in Haskell; the system o ﬀ ers a high-level syntax for dimension-annotated linear algebra, and generates specialized source code corresponding to the input expression.
Introduction
The design of domain-speciﬁc languages (DSL) for numerical computing is characterized by a tension between performance and expressiveness [1].
Traditional numerical libraries expose a number of highly optimized “kernels” (for instance, the linear combination Ax + b ); a user is expected to program against this interface in order to achieve high performance.
This approach is only limited by the availability of such kernels for the operations and data formats of interest.
Recently, generative programming techniques have been applied to synthesize specialized code from high-level speciﬁcations by leveraging statically known data (e.g. vector dimensions and types of the input expression). Multi-stage compilation (or staging ) is a form of metaprogramming in which two or more code synthesis phases take place, each evaluating at least in part the output of the previous one and emitting source code and metadata that will be consumed downstream [2–6].
Contribution In this work we describe a two-stage compiler for a high-level tensor DSL embedded in Haskell, which aims to strike a balance between user-friendliness and high performance.
In particular, the multi-stage approach makes it possible to synthesize code and corresponding specialized type signatures, circumventing certain limitations of the host language.
Compiler pipeline
AST analysis Users deﬁne the expression to be compiled using the grammar shown in Figure 1; for instance, contract (1) (0) x y denotes a contraction of variables x and y over the second and ﬁrst index respectively, and it may correspond to the matrix product X ij Y jk if both operands x and y have order 2.
The concrete syntax looks very similar to the abstract one, since it is encoded as a Haskell sum type. [Figure]
After checking its dimensional consistency, a number of structural operations are performed on the expression tree. First, subexpressions are recovered from the abstract syntax tree (AST) by hash-consing, which reconstructs the dataﬂow graph[7].
Afterwards, ontraction nodes are decorated with stride data structures which map tensor indices to the underlying memory layout of the variable (e.g. specifying that the ﬁbers 3 × 4 × 5 tensor to be reduced over its second index are stored everyelements). Given the expression DAG, we order the internal nodes into a min-heap according to their distance to the leaves, which then induces the expression reduction order (i.e. from the bound variables to the result).
Memory layout All operands are stored in memory as 1D unboxed arrays, as implemented by the vector library.
This library provides implicit fusion of intermediate mapping and reduction operations over contiguous ranges, which can decrease memory pressure and improve performance. The tensor storage type T is decorated with dimension annotations, implemented with a type-level list of naturals, as shown in Figure 2.
[Figure]
[Equation] [Figure]
[Equation] [Figure] Code and type signature generation The Q monad (for quotation , implemented in the template-haskell package) provides an e ﬀ ect system for creating unique names (analogous to Lisp’s “gensym”), as well as the syntactic marker separating our compiler stages.
In the code generation stage we employ a simple monad transformer type over Q which adds safe mutation of code fragments and read-only access to the expression DAG.
The expression heap from the AST analysis phase is popped until empty, and code corresponding to a let declaration is created and inserted into a hashmap; the last node represents the expression value, and will be returned by reference in the body of a let (whereas the internal nodes are bound as clauses). Template Haskell also permits the generation of types, including those having existentially quantiﬁed type variables and type-level literals such as strings and natural numbers.
In our case, this lets us specialize the signature of the generated bindings by adding mentions of the T tensor type as well as dimension annotations computed from the AST itself.
This can be seen as a convenient way of circumventing part of the typesystem that don’t support fully dependent-typing techniques (e.g. computing speciﬁc functions of types).  Example session
Here we show a typical interaction with the system. The expression in this example is the action of a 2 × 3 matrix on a 3-vector. The user declares their expression in a separate module, using the provided combinators, and compile s it elsewhere (the GHC “stage restriction” prevents code from di ﬀ erent stages to run in the same module) :
[Figure]
let y = contract
[Figure]
aa v pure y
[Figure]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
Related work
Array compilers and metaprogramming Generative programming has a history of successes in self-tuning numerical libraries such as ATLAS [14] and FFTW [15].
Later research produced whole-program compilers that optimize parallelism and arithmetic intensity [16], adapt deep learning workloads to various hardware backends or generic compiler IR [17–20] and o ﬀ er a mathematicallyintuitive API for tensor expressions while retaining the performance of hand-tuned kernels [8].
Similarly to other polyhedral compilers [21], our DSL is limited to modeling the static control part (“SCoP”) of a numerical program, and it overlaps in scope with TACO [8], with the major di ﬀ erence of being embedded in a declarative language.
Discussion
Some authors (e.g. [20, 23]) have shown that transforming memory layout can improve overall arithmetic e ﬃ ciency of tensor reduction sequences, even though picking the optimal such schedule is NP-hard [24], which motivates the use of approximations and heuristics at the expense of compiler complexity.
Others (e.g. [19] ) have demonstrated user-facing schedule combinators in their compiler API, but this is beyond the scope of the present work.
Hardware support is another important aspect; since the one we present is a Haskell metaprogramming environment, it only targets the hardware backends that are natively supported by the GHC compiler (i.e. no GPUs or FPGAs at present).  Acknowledgements The author is grateful to Jamie H. Willis and Oleg Kiselyov for helpful discussions, and to the anonymous reviewers whose feedback greatly improved this manuscript. A Typed Template Haskell
An early version of this work used typed template Haskell for code generation, [9], a type-checked extension to template Haskell (TH) [10] which provides syntactic support for quoting and splicing (“anti-quoting”).
In practice this means that we can compose target programs using higher-order combinators (e.g. those shown in Figure 5), resting assured that the resulting code will be well-formed.
This approach to compiling embedded DSLs has been demonstrated e.g. in [11] and recently applied to parser combinators [12] and extensible algebraic datatypes [13].
In typed TH, values of type Code a denote program fragments producing values of type a , whereas quoting and splicing are here denoted with (cid:126) · (cid:127) and σ ( . ) respectively.
[Figure]
The initial idea was to provide a type-safe user-facing DSL based on Code combinators, which would then be compiled into e ﬃ cient array programs.
Unfortunately this approach fell short because GHC doesn’t currently support impredicative polymorphism [22], i.e. type variables cannot be instantiated with polymorphic types.
More speciﬁcally, precise, dependently-typed signatures such as Code ( ∀ n .
KnownNat n ⇒ Vec n a → Vec n a → a) indicating binary operations on vectors of statically-known identical length will be rejected.