This post is intended to be an (informal) guide to the Taco1 compiler’s offspring, and a resource for a lot of the terminology used throughout various papers. Fred’s thesis2 is a very thorough and useful document that describes all the Taco lines of work up to 2020 (when it was published); it has become a little out-dated, though is still an excellent resource.

We’ll describe roughly the figure below, which is (approximately) how Burrito3 and Nacho4 produce kernels. There are some minor differences with Taco (including skipping CFIR, and just going straight to C++ from CIN), but the structure is largely the same. Burrito and Nacho also have a different iteration lattice construction algorithm: this is strictly necessary for supporting the operators that Burrito supports; and simply convenient for Nacho, because I find this algorithm more straight-forward than Taco’s original iteration lattice constructor.

digraph {
    splines=ortho; rankdir=LR;
    node [shape=box style=filled fillcolor="#f7f7f7" color="#555"
          fontname="Helvetica" fontsize=10 margin="0.14,0.07" height=0.4];
    edge [fontname="Helvetica" fontsize=9 color="#555" fontcolor="#555" arrowsize=0.7];

    TIN  [label="Tensor Index Notation" fillcolor="#f3e8fb" color="#6d09ba" group=main];
    CZ   [label="concretize" shape=plaintext style="" group=main];
    CIN  [label="Concrete Index Notation" group=main];
    IL   [label="iterate-locate" shape=plaintext style="" group=main];
    OCIN [label="oCIN" group=main];
    CFIR [label="Control Flow IR" group=main];
    CG   [label="code generation" shape=plaintext style="" group=main];
    CPP  [label="C++" fillcolor="#e8f1fb" color="#2c6fb5" group=main];
    FMT  [label="Format Specification" fillcolor="#f3e8fb" color="#6d09ba" group=fmt];
    PART [label="Partitioner"];

    // invisible branch points along the Format Specification line, so the
    // dashed offshoots leave that line instead of radiating from the box
    B1 [shape=point width=0.01 style=invis group=fmt];
    B2 [shape=point width=0.01 style=invis group=fmt];
    B3 [shape=point width=0.01 style=invis group=fmt];

    // TIN sits above FMT; each branch point shares a column with what it feeds.
    // The dashed edges below are declared downwards (with dir=back) because a
    // flat edge's direction is what orders the two nodes within their column.
    { rank=same; TIN; FMT; TIN -> FMT [style=invis]; }
    { rank=same; CZ;  B1; }
    { rank=same; IL;  B2; }
    { rank=same; CFIR; PART; B3; CFIR -> PART [style=invis]; }

    TIN  -> CZ   [arrowhead=none];
    CZ   -> CIN;
    CIN  -> IL   [arrowhead=none];
    IL   -> OCIN;
    OCIN -> CFIR [label="iteration lattices"];
    OCIN -> PART;
    CFIR -> CG   [arrowhead=none];
    CG   -> CPP;
    PART -> CG;
    FMT  -> B1   [arrowhead=none];
    B1   -> B2   [arrowhead=none];
    B2   -> B3   [arrowhead=none];
    B3   -> CG;
    CZ   -> B1   [xlabel="mode ordering" style=dashed dir=back];
    IL   -> B2   [xlabel="properties"    style=dashed dir=back];
    PART -> B3   [xlabel="nnz counters"  style=dashed dir=back];
}

Tensor Index Notation

We’ll start with some basics: specifying tensor algebra operations with Tensor Index Notation (also called einsum notation). TIN supports element-wise addition and multiplication, broadcasting, and summation (reduction along a tensor dimension). For example, the element-wise addition of two 2D tensors (AKA matrices), \(A\) and \(B\), is denoted as

\[Z_{ij} = A_{ij} + B_{ij}\]

Where \(i\) and \(j\) label the two tensor dimensions. Similarly, matrix multiplication can be specified as:

\[Z_{ik} = \sum_j A_{ij} \times B_{jk}\]

Note that there are different indexes into \(B\) and \(Z\): the output takes rows from \(A\) and columns from \(B\). In many of the Taco papers, multiplication isn’t written, e.g.: \(Z_{ik} = \sum_j A_{ij}B_{jk}\), and in some, summations are implicit: \(Z_{ik} = A_{ij} \times B_{jk}\) (the summation is over \(j\) because \(j\) does not appear in the output). Another implicit notation: \(A\) does not have a \(k\) dimension, and \(B\) does not have an \(i\) dimension: that means that \(A\) is broadcasted over dimension \(k\) and \(B\) is broadcasted over dimension \(i\). A simpler example of broadcasting is matrix-plus-vector:

\[Z_{ij} = A_{ij} + b_{i}\]

Where the vector \(b\) is broadcasted over dimension \(j\). Vectors are frequently denoted using lowercase letters, with higher-order tensors denoted using uppercase.

Two other example computations that I will use in this post are element-wise multiplication (AKA the Hadamard product) and matrix-vector multiplication:

\[Z_{ik} = A_{ij} \times B_{ij} \qquad \qquad z_{i} = A_{ij} \times b_{j}\]

Now, you may have noticed that tensor index notation says nothing about how the tensors are stored, and which tensors are sparse versus which are dense. That is part of the point of Taco: specify computation (sometimes called “the algorithm”) using pure, high-level mathematics, and separately specify the data structures, relying on an intelligent compiler to produce the correct (fast) code. This is a similar design to Halide5, which popularized the notion of decoupling a high-level algorithm specification (a functional image processing pipeline) from the execution schedule (loop tilings, parallelization details, etc.). Taco eventually added a scheduling language6 itself, but more on that later.

Other Operations

Another brief aside: follow-on works extended tensor index notation with support for a broader class of semiring operations7, shape operators3 (e.g., reshape, concat, etc.), and convolutions8; for the purpose of this blog post, I will focus on sparse tensor algebra, as other components of the pipeline (e.g., scheduling) have not been built in conjunction with support for these other iteration patterns. Hopefully it will all be unified someday ;-)

Format Specification

The format language9 specifies two important components of a tensor: the mode ordering and the storage format (AKA a list of level formats).

Mode Ordering

A mode ordering is the order that dimensions are stored in: for a 2D tensor (matrix), there are two options: row-major, or column-major; a 3D tensor has 6 options (RCD, RDC, CRD, CDR, DCR, DRC); and so on. For dense tensors, mode ordering simply changes stride information for access patters, but the mode ordering for sparse tensors generally describes the only way to efficiently iterate a sparse tensor. For example, the classic Compressed Sparse Row (CSR) format is a row-major format, which means that it is incredibly inefficient to iterate over its columns before its rows. But to understand why, we need to look at the second component of a format, the storage type.

Storage (Level Formats)

A storage type (or level format) describes how a dimension is iterated (or randomly accessed into). The original Taco paper1 supported a Dense and a Compressed level format, where Dense is the classic dense array supporting in-order iteration and O(1) random access, and Compressed is a sorted, compressed, key-value pairing that only stores non-zero coordinates. Thus, Compressed does not support O(1) random access (later10 works3 use the fact that it does support O(\(log(n)\)) random access.)

Consider the vector \(b = [5, 1, 0, 0, 2, 0, 9, 0]\). Stored as a dense array, every slot is kept, zeros included:

vals 5 0 0 0 2 0 8 0

Stored as a sparse (Compressed) vector, only the coordinates (crd) and values (vals) of the non-zeros are kept:

crd 0 4 6
vals 5 2 8

The highlighted entry (value 2, at coordinate 4) is the same element in both. For a vector with 3 non-zeros out of 8 entries, this saves a modest 8 bytes (assuming 32-bit values and 32-bit indices).

Chou et al.9 extended the original Taco’s level formats to support Singleton, Range, Offset, and Hash level formats; and Compressed and Singleton levels also support unique (or non-unique) and ordered (or unordered) modifiers. All in all, these formats describe how to iterate over, or index into, a dimension of a tensor. But the really interesting part is the composability.

Combining Level Formats

To go beyond specifying simple vector formats, a mode ordering and a list of level formats combine to fully specify a sparse tensor format. For example, the CSR format is a row-major [Dense, Compressed] format.

Consider the \(3 \times 4\) matrix \(A\), whose rows are labelled on the left:

0 5 0 1 0
1 0 0 0 0
2 0 2 7 8

Because the outer (row) level is Dense, it stores no coordinates at all: but it does need to store information for the level below it; in this case, the ranges of key-value pairs that correspond to a given row, which is the pos array below. Because the inner (column) level is Compressed it stores the column coordinates (crd) and values (vals) of the non-zeros:

row: 0 1 2 end
pos 0 2 2 5
crd 0 2 1 2 3
vals 5 1 2 7 8

The highlighted entry is \(A_{02} = 1\): it lives at position 1 of vals, with its column coordinate 2 at the same position of crd. Row \(i\) owns positions pos[i] through pos[i+1], so row 0 owns positions 0–2, row 1 owns 2–2 (it is empty, which is why 2 appears twice), and row 2 owns 2–5. This is exactly why CSR is inefficient to iterate column-first: recovering a single column means scanning every row’s segment of crd, because only the row level supports O(1) access.

There are many different sparse tensor formats that can be produced by compositions of level formats and mode orders: Compressed Sparse Column is a column-major [Dense, Compressed] matrix; while Doubly Compressed Sparse Row is a row-major [Compressed, Compressed] matrix; and the 3D Compressed Sparse Forest is [Compressed, Compressed, Compressed] (with some choice of mode ordering). Read Chou et al.9 for more about formats, or Fred’s thesis2.

Concretization

The first step of lowering a sparse tensor algebra kernel is concretizing (in the language of Kjolstad et al.11) sparse tensor algebra into a loop nest. The natural question is, given a tensor algebra computation with multiple labeled dimensions, how do I choose an order of dimensions to iterate? For most sparse computations, it is relatively easy: you choose the order that the mode ordering allows. This is called concordant traversal; if a computation cannot be traversed in mode order (e.g., adding a CSR matrix to a CSC matrix), generally at least one operand will need to be transposed into the mode ordering of the other, instead of performing a discordant traversal. The Scorch paper12 chooses a loop ordering via exactly this: construct a graph, where the nodes represent each (unique) dimension, and directed edges indicate the mode ordering. For example, adding two CSR matrices: \(Z_{ij} = A_{ij} + B_{ij}\) produces:

digraph {
    rankdir=LR;
    node [shape=circle fixedsize=true width=0.35 color="#555"
          fontname="Helvetica" fontsize=10];
    edge [fontname="Helvetica" fontsize=9 color="#555" fontcolor="#555" arrowsize=0.7];
    i -> j [label="A, B"];
}

Where the labels on the edge simply correspond to the tensor(s) whose mode orderings create that edge. The matrix multiplication of two CSR matrices (AKA SpGEMM) \(Z_{ik} = \sum_j A_{ij} \times B_{jk}\) produces:

digraph {
    rankdir=LR;
    node [shape=circle fixedsize=true width=0.35 color="#555"
          fontname="Helvetica" fontsize=10];
    edge [fontname="Helvetica" fontsize=9 color="#555" fontcolor="#555" arrowsize=0.7];
    i -> j [label="A"];
    j -> k [label="B"];
}

Both of these graphs are chains, so there is exactly one topological ordering, and the loop order is forced. To get a real choice, the graph has to branch, which happens as soon as one dimension constrains two others. The Galerkin product \(Z = P^\top A P\) (the workhorse of algebraic multigrid, which coarsens a fine-grid operator \(A\) through a prolongation operator \(P\)) is a smallest realistic example of branching. In TIN, the Galerkin product is:

\[Z_{kl} = \sum_i \sum_j P_{ik} \times A_{ij} \times P_{jl}\]

Note that \(P\) appears twice, under different index variables, so it contributes two edges. Assuming all matrices are stored in a row-major format, the index graph looks like:

digraph {
    rankdir=LR;
    node [shape=circle fixedsize=true width=0.35 color="#555"
          fontname="Helvetica" fontsize=10];
    edge [fontname="Helvetica" fontsize=9 color="#555" fontcolor="#555" arrowsize=0.7];
    i -> j [label="A"];
    i -> k [label="P"];
    j -> l [label="P"];
}

\(i\) clearly needs to be the first loop index, and \(l\) must come after \(j\) but \(k\) is free to be reordered anywhere around \(j\) and \(l\). The choice has a real performance impact: they differ in the patterns of reading from the inputs, and the write pattern of the output. The Scorch algorithm relies on a simple heuristic: it is better to “cull” work earlier in the loop nest, and the \(j\) loop is an intersection between two indexes, \(j_A\) and \(j_P\), so it is the second choice for a loop nest. \(k\) and \(l\) have no structural differences (and are, in fact, equal, in this example!) and can be ordered in any order. This choice does have an interesting performance impact: writes into the output are completely random scatters. But more on that later.

Loop Bounds

After a loop ordering is chosen, the iteration bounds of a loop needs to be computed: Burrito3 refers to sparse iteration bounds as a “set expression” over the non-zero coordinates. For example, a simple vector addition \(z_i = a_i + b_i\) can be computed as a loop over the union of non-zero coordinates of \(a\) and \(b\), denoted:

forall \(i \in i_a \cup i_b\): z[\(i\)] = a[\(i_a\)] + b[\(i_b\)]

This is based on the fact that the sum of two zeros is zero, meaning the output is only written to when one of the vectors has a non-zero. Likewise, an element-wise multiplication, \(z_i = a_i * b_i\) is computed over the intersection of non-zero coordinates:

forall \(i \in i_a \cap i_b\): z[\(i\)] = a[\(i_a\)] * b[\(i_b\)]

This is because the multiplication of any value with zero is zero, so both vectors need a non-zero value for the output to be non-zero. This idea composes very cleanly: a fused multiply-accumulate, \(z_i = a_i + (b_i * c_i)\), iterates over:

forall \(i \in i_a \cup (i_b \cap i_c)\): z[\(i\)] = a[\(i_a\)] + b[\(i_b\)] * c[\(i_c\)]

While sparse tensor algebra can be expressed solely in terms of intersections and unions, three works extend this model: Henry and Hsu et al.7 show that adding set complement to the language massively expands the iteration patterns to support arbitrary (element-wise or reduction) functions applied to sparse tensors; Burrito3 adds other set operations (projection, product, disjoint union) to extend this programming model to support shape operators like reshape and concatenation; and Liu et al.8 shows that a type of dependent slicing can be used to model sparse convolutions in this model, with asymptotic performance improvements over prior techniques. This blog post is going to focus on just sparse tensor algebra (for now), but no system has been built that supports all of these iteration patterns (yet).

Deriving Loop Bounds

The derivation of loop bounds is relatively straight-forward for tensor algebra: for each loop index, recursively traverse the AST to build the set expression that represents the non-zeros of the expression in that dimension.

func derive(expr, idx): match expr with | array(I) \(\mapsto\) \(idx_{array}\) | a + b \(\mapsto\) derive(a, idx) \(\cup\) derive(b, idx) | a * b \(\mapsto\) derive(a, idx) \(\cap\) derive(b, idx) | sum(_, a) \(\mapsto\) derive(a, idx) | broadcast(idx, a) \(\mapsto\) \(\mathbb{U}_{idx}\) | broadcast(j, a) \(\mapsto\) derive(a, idx) .

This algorithm produces the loop bounds we saw earlier for the sparse vector examples, as well as sparse matrix addition:

forall \(i \in i_A \cup i_B\): forall \(j \in j_A \cup j_B\): z[\(i, j\)] = A[\(i_A, j_A\)] + B[\(i_B, j_B\)]

and sparse matrix multiply:

forall \(i \in i_A \cap \mathbb{U}_i\): forall \(j \in j_A \cap j_B\): forall \(k \in \mathbb{U}_k \cap k_B\): Z[\(i, k\)] += A[\(i_A, j_A\)] * B[\(j_B, k_B\)]

Note that the “universes” (e.g., \(\mathbb{U}_i\)) come from the fact that \(B\) is broadcast over \(i\) and \(A\) is broadcast over \(k\). This represents the set of all possible coordinates of that dimension. We can similarly produce the following \(i, j, l, k\) loop nest for the Galerkin product:

forall \(i \in i_P \cap i_A \cap \mathbb{U}_i\): forall \(j \in \mathbb{U}_j \cap j_A \cap j_P\): forall \(l \in \mathbb{U}_l \cap l_P\): forall \(k \in k_P \cap \mathbb{U}_k\): Z[\(k, l\)] += P[\(i_P, k_P\)] * A[\(i_A, j_A\)] * P[\(j_P, l_P\)]

Because \(P\) appears twice, it contributes two independent pairs of coordinate sets: \(i_P, k_P\) from \(P_{ik}\), and \(j_P, l_P\) from \(P_{jl}\).

Optimizing Loop Bounds (iterate-locate)

The naive lowerings have a few (possibly obvious) structural inefficiencies: for example, iterating over a set intersected with the universe is equivalent to iterating over just the set. This is true for simplifying broadcasts, but is also true for simplifying kernels that mix dense and sparse dimensions. Reconsider the SpGEMM CIN:

forall \(i \in i_A \cap \mathbb{U}_i\): forall \(j \in j_A \cap j_B\): forall \(k \in \mathbb{U}_k \cap k_B\): Z[\(i, k\)] += A[\(i_A, j_A\)] * B[\(j_B, k_B\)]

Both of the universes are intersected with sets, and can therefore be optimized away:

forall \(i \in i_A\): forall \(j \in j_A \cap j_B\): forall \(k \in k_B\): Z[\(i, k\)] += A[\(i_A, j_A\)] * B[\(j_B, k_B\)]

But going further, if \(B\) is stored in a CSR matrix, then \(j_B\) is a dense dimension, which is equivalent to iterating over the universe! If \(j_A\) is sparse (e.g., because \(A\) is also stored as a CSR), then it is asymptotically efficient (but still correct!) to just iterate \(j_A\):

forall \(i \in i_A\): forall \(j \in j_A\) with \(j_B = j\): forall \(k \in k_B\): Z[\(i, k\)] += A[\(i_A, j_A\)] * B[\(j_B, k_B\)]

The with syntax denotes computing \(j_B\) from \(j\) – for sparse tensor algebra, this is only ever an equality. Extending to shape operators allows the equation to get a little bit more complicated. This pattern is known as iterate-locate, because \(j_A\) is iterated, and \(j_B\) is located into. Taco (and the follow-ons) use the properties of sparse dimensions to correctly perform this look-up: any dimension with O(1) random access (dense and hashed) will be located into if part of a set intersection.

A similar optimization can be done if all index variables are dense, for example, sparse matrix addition can be optimized to:

forall \(i \in \mathbb{U}_i\) with \(i_A = i\), \(i_B = i\): forall \(j \in j_A \cup j_B\): z[\(i, j\)] = A[\(i_A, j_A\)] + B[\(i_B, j_B\)]

The various Taco compiler extensions choose to do this optimization in different places: Taco does it as part of iteration lattice construction (the next section) but lowers directly to C/C++ from there, Burrito defines an intermediate representation (IR) called CFIR that stands between CIN and C++, and performs this optimization when constructing CFIR via a generalized iteration lattice; Nacho chooses to simplify CIN into an iterate-locate optimized form (labeled oCIN in the overview figure of this blog, but not discussed in the Nacho paper directly), because this optimization needs to be applied before performing Nacho’s load balancing algorithm (again, more on that below).

Iteration Lattices

Now, while the iterate-locate pattern does produce asymptotic benefits, but there can still be (perceived) inefficiencies in CIN. For example, consider a sparse matrix addition over two DCSR matrices (row-major, [Compressed, Compressed]):

forall \(i \in i_A \cup i_B\): forall \(j \in j_A \cup j_B\): z[\(i, j\)] = A[\(i_A, j_A\)] + B[\(i_B, j_B\)]

No dimension is dense, so the iterate-locate optimization does not apply. However, it appears as if the \(j\) loop iterates over both \(j_A\) and \(j_B\) even when \(j_A\) is empty (e.g., \(i \notin i_A\)). Ideally, in the case that \(i \notin i_A\), the loop over \(j\) is just a loop opver \(j_B\). This reasoning is encoded in iteration lattices in Taco’s terminology. The desired loop nest looks like:

forall \(i \in i_A \cup i_B\): case \(i \in i_A \cap i_B\): forall \(j \in j_A \cup j_B\): z[\(i, j\)] = A[\(i_A, j_A\)] + B[\(i_B, j_B\)] case \(i \in i_A\): forall \(j \in j_A\): z[\(i, j\)] = A[\(i_A, j_A\)] case \(i \in i_B\): forall \(j \in j_B\): z[\(i, j\)] = B[\(i_B, j_B\)]

The cases are checked in order, so the later two only run when the intersection case does not: reaching case \(i \in i_A\) already implies \(i \notin i_B\), and there is no need to iterate (or even look at) \(j_B\). Note also that only the first case needs to merge two dimensions in the \(j\) loop; the other two degenerate into a single-tensor loop over a single sparse dimension, which is exactly the inefficiency we wanted to remove.

We arrive at this loop nest by constructing the following iteration lattice for the \(i\) loop:

digraph {
    rankdir=TB;
    node [shape=box style=filled fillcolor="#f7f7f7" color="#555"
          fontname="Helvetica" fontsize=10 margin="0.16,0.08"];
    edge [color="#555" arrowsize=0.7];
    AB [label="i_A ∪ i_B"];
    A  [label="i_A"];
    B  [label="i_B"];
    BOT [label="∅"];
    AB -> A;
    AB -> B;
    A -> BOT;
    B -> BOT;
}

Each lattice point names the dimensions that must be iterated together, and the expression that is computed when exactly those tensors have a non-zero at \(i\). The top point is the most constrained (both operands present), and edges point to the points that remain when an operand drops out. Lowering walks the lattice from the top down, emitting one case per point, which is where the ordered cases above come from. Because addition is a union, every point in the lattice is reachable; for a multiplication, only the intersection point produces a non-zero, so the lattice collapses to a single (non-empty) point:

digraph {
    rankdir=TB;
    node [shape=box style=filled fillcolor="#f7f7f7" color="#555"
          fontname="Helvetica" fontsize=10 margin="0.16,0.08"];
    edge [color="#555" arrowsize=0.7];
    AB [label="i_A ∩ i_B"];
    BOT [label="∅"];
    AB -> BOT;
}

TODO: finish here.

  1. Kjolstad et al. The Tensor Algebra Compiler, OOPSLA 2017. ↩ ↩2

  2. Fredrik Kjolstad. The Tensor Algebra Compiler. PhD thesis, MIT, 2020. ↩ ↩2

  3. Root et al. Compilation of Shape Operators on Sparse Arrays, OOPSLA 2024. ↩ ↩2 ↩3 ↩4 ↩5

  4. Chougule and Root et al. Partitioning Unstructured Sparse Tensor Algebra for Load-Balanced Parallel Execution, ASPLOS 2027. ↩

  5. Ragan-Kelley et al. Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines, PLDI 2013. ↩

  6. Senanayake et al. A Sparse Iteration Space Transformation Framework for Sparse Tensor Algebra, OOPSLA 2020. ↩

  7. Henry and Hsu et al. Compilation of Sparse Array Programming Models, OOPSLA 2021. ↩ ↩2

  8. Liu et al. Compiler Support for Sparse Tensor Convolutions, OOPSLA 2024. ↩ ↩2

  9. Chou et al. Format Abstraction for Sparse Tensor Algebra Compilers, OOPSLA 2018. ↩ ↩2 ↩3

  10. Kovach et al. Indexed Streams: A Formal Intermediate Representation for the Fused Execution of Contraction Operations, PLDI 2023. ↩

  11. Kjolstad et al. Tensor Algebra Compilation with Workspaces, CGO 2019. ↩

  12. Yan et al. Fast Autoscheduling for Sparse ML Frameworks, CGO 2026. ↩