lir]textfontsize=, escapeinside=——, mathescape lirinline]text
lirln]textfontsize=, escapeinside=——, mathescape,linenos

# Compressing Structured Tensor Algebra

Mahdi Ghorbani∗ 0000-0003-3686-2158<EMAIL_ADDRESS>University of
EdinburghEdinburghUK ,  Emilien Bauer∗<EMAIL_ADDRESS>University of
EdinburghEdinburghUK ,  Tobias Grosser<EMAIL_ADDRESS>University of CambridgeCambridgeUK and  Amir Shaikhha  0000-0002-9062-759X
<EMAIL_ADDRESS>University of EdinburghEdinburghUK

###### Abstract.

Tensor algebra is a crucial component for data-intensive workloads such as
machine learning and scientific computing. As the complexity of data grows,
scientists often encounter a dilemma between the highly specialized dense
tensor algebra and efficient structure-aware algorithms provided by sparse
tensor algebra. In this paper, we introduce DASTAC, a framework to propagate
the tensors’s captured high-level structure down to low-level code generation
by incorporating techniques such as automatic data layout compression,
polyhedral analysis, and affine code generation. Our methodology reduces
memory footprint by automatically detecting the best data layout, heavily
benefits from polyhedral optimizations, leverages further optimizations, and
enables parallelization through MLIR. Through extensive experimentation, we
show that DASTAC achieves 1 to 2 orders of magnitude speedup over TACO, a
state-of-the-art sparse tensor compiler, and StructTensor, a state-of-the-art
structured tensor algebra compiler, with a significantly lower memory
footprint.

Sparse Tensors, Compiler Optimization, Barvinok Algorithm, MLIR

**footnotetext: These authors contributed equally to this work

##  1\. Introduction

Figure 1. Comparison of tensor processing frameworks. DASTAC is the first code
generation framework that combines the algorithmic optimizations known from
sparse tensor algebras with performance-optimized low-level code known from
decades of tuning dense tensors.

Tensor algebra is a fundamental component in several data-intensive workloads,
such as signal processing (De Lathauwer and De Moor, 1998; Nion and
Sidiropoulos, 2010; Muti and Bourennane, 2007), machine learning (Sidiropoulos
et al., 2017), computer vision (Panagakis et al., 2021), quantum chemistery
(Khoromskaia and Khoromskij, 2018; Mutlu et al., 2019), and bioinformatics
(Terán et al., 2019). Numerous lines of research have attempted to specialize
and enhance the performance of the tensor algebra computations either through
hardware level (e.g., tensor accelerators (Hegde et al., 2019; Jia et al.,
2021; Gondimalla et al., 2019) and TPUs (Jouppi et al., 2023)), or software
level (e.g., tuned tensor kernels (Wang et al., 2014; Dongarra et al., 1990)
and compilation based optimizations (Gareev et al., 2018)).

As the complexity of the data grows, scientists often face a trade-off between
highly tuned and specialized frameworks provided for dense tensor algebra
computation and efficient and flexible algorithms provided for sparse tensor
algebra that leverage the sparsity of input tensors. Dense tensor algebra
frameworks (Baghdadi et al., 2019; Vasilache et al., 2018; Grosser et al.,
2012; Mullapudi et al., 2015) support a rich set of compile-time optimizations
such as vectorization, tiling, and parallelization. This is thanks to well-
known memory access patterns and efficient utilization of contiguous memory,
despite not knowing the actual data.

Computations over structured tensor algebras are at the core of several
important workloads. The structure is an expression of static information on
the sparsity and redundancy of the tensors. Sparsity information defines the
symbolic set of indices of non-zero values (e.g., diagonal) while redundancy
defines a symbolic mapping from the set of indices of repetitive elements to
their unique equivalent (e.g., symmetric). Structured matrices have many
applications in machine learning (Khamis et al., 2020) and deep learning
(Kissel and Diepold, 2023). Diagonal structure used in the ACDC layer
(Moczulski et al., 2016), tridiagonal weight matrices mentioned by (Dumitras
and Kossentini, 2000), and structured neural network pruning for large
language models (Wang et al., 2020) are examples of applications of matrices
with simple structures. Furthermore, more intricate structures such as hyper
triangular for covariance matrices (Ghorbani et al., 2023) and the butterfly
structure used for neural network training (Chen et al., 2022) are examples of
structured tensor applications.

Sparse tensor algebra frameworks (Ghorbani et al., 2023; Chou and Amarasinghe,
2022; Virtanen et al., 2020; Spampinato and Püschel, 2016) provide
asymptotically improved algorithms by leveraging the sparsity structure of the
data with efficient memory storage requirements and better scalability.
However, unlike dense counterparts, sparse frameworks cannot leverage the full
computation power of the hardware. This is due to irregular data structures,
leading to less predictable memory access patterns and making it challenging
to fully leverage the same level of optimizations as their dense counterparts.

In this paper, we introduce DASTAC, a framework that propagates the high-level
structural information of tensor computations down to the low-level highly
tuned code generation (cf. Figure 1). DASTAC leverages a novel symbolic
indexing method that compresses an input structured tensor into a densely
packed vector representation. In contrast to well-established sparse data
layouts like CSR, CSC, or COO, which involve indirect accesses and incur index
storage overhead, our method employs efficient direct symbolic index
computations. This direct indexing not only diminishes the overall memory
footprint but also opens up opportunities for vectorization. These advantages
become more important in scenarios where memory bandwidth is the main
bottleneck in program efficiency, as is the case on many widespread
contemporary hardware, such as CPUs and GPUs.

DASTAC achieves the densely assembled data layout by relying on a well-
established mathematical foundation: the polyhedral model. First, DASTAC uses
StructTensor to capture and propagate the structure throughout the tensor
algebra computation and then proceeds to apply polyhedral set and map
optimizations over the captured structure. Second, DASTAC proceeds to find a
densely packed data layout for the tensors involved in the computation and
compress them automatically by a novel symbolic indexing algorithm, leading to
better cache locality and less memory footprint. Finally, DASTAC generates
structure-aware low-level code for CPU using the affine dialect of MLIR
(Lattner et al., 2021), enabling effective optimizations at different levels
of abstraction. Consequently, DASTAC brings out the best of both dense and
sparse tensor algebra worlds together by relying on polyhedral analysis,
memory, and complexity efficacy of sparse tensor algebra alongside
specialization and compile-time optimization of dense tensor algebra based on
polyhedral and compiler optimizations.

Specifically, we make the following contributions:

  * •

We introduce DASTAC, the first framework that integrates the best of both
sparse and dense tensor algebra worlds (Section 3). The algorithmic
improvements are achieved by leveraging the structure of tensors and this
structure is progressively lowered down to the machine-level code.

  * •

DASTAC introduces a densely assembled data layout that enables low-level
specialization opportunities (Section 4). The central technique is a novel
static indexing algorithm that maps the indices of the input tensor to a
contiguous memory buffer.

  * •

We provide a progressive code generation algorithm (Section 5). First, the
tensor computations are translated to affine operations provided out-of-the-
box by the MLIR framework. Subsequently, at the next stages, the MLIR
framework allows DASTAC to benefit from optimizations such as code motion,
common subexpression elimination and parallelization.

  * •

We experimentally evaluate DASTAC against state-of-the-art tensor algebra
frameworks and show that leveraging structure while benefiting from our
proposed compressed data layout and polyhedral optimizations leads to better
performance both sequentially and on multi-threaded scenarios (Section 6). We
show that DASTAC achieves 1 to 2 orders of magnitude speedup over the TACO
sparse tensor compiler and StructTensor, a state-of-the-art structured tensor
algebra compiler, with a significantly lower memory footprint. We also show
that frameworks such as Polygeist cannot recover the optimizations offered by
DASTAC.

##  2\. Background on Structured Tensor Algebra

Matrices and tensors can have several structures, such as diagonal,
tridiagonal, symmetric, and triangular. Leveraging such structures can improve
the computational cost of tensor operations by orders of magnitude. For
example, leveraging the matrix structure in sparse matrix-vector
multiplication (SpMV), where the matrix is diagonal, can lower the
computational complexity from $O(n^{2})$ to $O(n)$.

Tensor structures can be classified into two categories: 1) sparsity patterns
and 2) redundancy patterns. The sparsity patterns refer to structures such as
diagonals where zero and non-zero values are distinguished. Structures with
redundancy pattern refers to structures such as symmetry, with unique and
repetitive elements. The tensor can be reconstructed using the unique values
and a mapping from the indices of redundant elements to the unique ones.

STUR (Ghorbani et al., 2023) is the StructTensor’s unified intermediate
representation used for representing both arithmetic and structural
information. StructTensor relies on structure inference for tensor
computations in the form of unique sets, redundancy maps, and compressed
tensors. We explain each component of STUR and how it works using a running
example.

Unique set. It represents the non-zero and unique (non-repetitive) values of
the input tensor and corresponds to capturing the sparsity pattern. For
example, the unique set of an $n\times n$ diagonal matrix $M$ represented as
$M_{U}$ is as follows:

$M_{U}$ | $:=$ |  $\\{(x,y)\mid(0\leq x<n)\wedge(0\leq y<n)\wedge(x=y)$}
---|---|---

Any set representation can be shown in the form of tensor computation over a
boolean domain by replacing $\wedge$ with $*$ and $\vee$ with $+$. Therefore,
the unique set can be represented as follows:

$M_{U}(x,y)$ | $:=$ | $(0\leq x<n)*(0\leq y<n)*(x=y)$
---|---|---

Redundancy map. It provides the mapping from redundant elements to their
corresponding unique elements and corresponds to capturing the redundancy
pattern. For example, the unique set and redundancy map for an $n\times n$
symmetric matrix $V$ are as follows:

$V_{U}(x,y)$ | $:=$ | $(0\leq x<n)*(0\leq y<n)*(x\leq y)$
---|---|---
$V_{R}(x,y,x^{\prime},y^{\prime})$ | $:=$ | $(0\leq x<n)*(0\leq y<n)*(x>y)*(x^{\prime}=y)*(y^{\prime}=x)$

Here we are assuming that the lower triangular part of the matrix is the
unique part, and the upper triangular part of it is repetitive. Variables
$x,y$ represent the iteration space, and the mapping from the upper triangular
to the lower triangular part is represented through variables
$x^{\prime},y^{\prime}$. So this representation means $x$ and $y$ values must
be swapped to have access to the unique elements.

Structure inference. Now imagine the example of sparse matrix-matrix element-
wise multiplication where the operands are the aforementioned matrices, $M$
and $V$. To perform this multiplication while leveraging the structure, STUR
applies the program reasoning rules (Ghorbani et al., 2023) to infer this
computation’s output structure and compressed iteration space. For the
mentioned example, the computation and the inferred output structure after
simplification are:

$T(x,y)$ | $:=$ | $M(x,y)*V(x,y)$
---|---|---
$T_{U}(x,y)$ | $:=$ | $(0\leq x<n)*(0\leq y<n)*(x=y)$
$T_{R}(x,y,x^{\prime},y^{\prime})$ | $:=$ | $\emptyset$

Compressed tensor. It is used for code generation and contains the structural
information (the unique set and redundancy map) and the arithmetic operation
over tensors using STUR’s syntax. The computation over sets and arithmetic is
written using the same intermediate representation (STUR) for structure
inference purposes but is treated separately in the code generation process.
StructTensor internally separates the structural information and arithmetic
parts. The structural information is used to generate loops while the
arithmetic part generates the computational code inside the loop nests. For
example, matrix $T$ will have the following compressed tensor:

$T_{C}(x,y)$ | $:=$ | $(0\leq x<n)*(0\leq y<n)*(x=y)*M(x,y)*V(x,y)$
---|---|---

This means that in this computation, the iterators $x,y$ should iterate
between $0$ and $n$ and must be equal to each other. The arithmetic operation
that they perform inside the loop nests is $M(x,y)*V(x,y)$. Note that
compressed tensor here refers to a tensor that has the structural information,
hence, the iteration space can be compressed using this information. The
tensor data layout itself is not compressed here. StructTensor generates the
following code:

⬇

for (int x = 0; x < n; ++x) {

int y = x;

T(x, y) = M(x, y) * V(x, y);

}

The first 2 lines are generated by the structural information and the third
line is generated because of the arithmetic part of the compressed tensor.

State-of-the-Art. StructTensor (Ghorbani et al., 2023) introduces the STUR
intermediate language with the program reasoning rules to infer structures. As
a final step, StructTensor provides C++ code that performs the computation
over the compressed iteration space. StructTensor outperforms sparse and dense
tensor algebra competitors by leveraging and propagating the compile-time-
known structures throughout the computation.

Limitations. Even though StructTensor infers the structure, it relies on an
uncompressed dense representation of tensors that contains zero and repetitive
elements. Alternatively, StructTensor requires the user to specify a compact
data layout manually, but it is limited to only data layout for inputs. The
outputs are still stored in an uncompressed dense format. However,
StructTensor cannot capture this compressed output data layout. Furthermore,
the generated C++ code relies on the low-level compiler (e.g., Clang or GCC)
for further optimizations. Thus, one cannot benefit from the more advanced
optimizations not supported directly by the compiler (e.g., parallelization).

DASTAC solves both limitations by relying on the mathematical foundation of
the polyhedral model. First, it uses a novel symbolic indexing algorithm to
densely pack the input and output tensors. Second, benefitting from the
progressive lowering between intermediate languages (dialects) provided by the
MLIR framework, it applies additional compiler optimizations, including
parallelization and vectorization.

Figure 2. DASTAC architecture overview.

##  3\. Overview

In this section, we describe the overall architecture of DASTAC (cf. Figure
2).

(1) $\begin{split}A(i,j,k)&:=B(i,j,l)*C(k,l)\\\ C_{U}(k,l)&:=(0\leq
k)*(k<P)*(0\leq l)*(l<Q)\\\ B_{U}(i,j,l)&:=(i\leq j)*(0\leq i)*(i<M)*(0\leq
j)*(j<N)*(0\leq l)*(l<Q)\\\ A_{C}(i,j,k)&:=C(k,l)*B(i,j,l)*(0\leq l)*(Q>l)\\\
&*(i\leq j)*(N>j)*(0\leq i)*(M>i)*(0\leq k)*(P>k)\end{split}$ Figure 3. The
TTM running example where the first tensor has an upper half cube structure.
$A_{C}$ represents its optimized structured computation.

Input. DASTAC uses STUR as its intermediate language, as it relies on
StructTensor (Ghorbani et al., 2023) for the structure inference as the first
step. STUR can express generalized Einsum expressions as a sum of products
(SoP). As an example, Figure 3 shows the STUR representation of the tensor
times matrix (TTM) computation with an upper half cube structure.

Structure inference. DASTAC proceeds to propagate the input structure
throughout the computation by applying program reasoning rules provided by
StructTensor. The output structure and compressed computation are inferred
during this stage. DASTAC uses the inferred compressed domain by StructTensor
as input. For the running example, the output structure is as follows:

$A_{U}(i,j,k)$ | $:=$ | $B_{U}(i,j,l)*C_{U}(k,l)$
---|---|---

Structure Optimizations. DASTAC improves the structured tensor computation by
applying structure optimizations. StructTensor provides a set of optimizations
such as inlining and logical simplifications (e.g., set idempotence and
addition/multiplication identity optimizations) on the output structure and
compressed computation. A polyhedral tool such as isl (Verdoolaege, 2010) can
easily support and apply the aforementioned optimizations alongside further
polyhedral set optimizations. Therefore, relying on polyhedral tools can
reduce the development cost while increasing the robustness. For these
reasons, DASTAC relies on isl for further structure optimizations. Figure 3
shows the structured computation for the running example after applying
structure optimizations.

Data Layout Compression. DASTAC leverages the statically known sparsity
structure to assemble a densely packed buffer, aiming for cache-friendliness
and vectorization opportunities. Our compression technique consists of a
symbolic indexing function, mapping each original tensor index to a linear
index of unique and non-zero values. For example, the inferred compressed
index of $B$ in our running example Figure 3 is:

$P_{B}(i,j,l)=\left(\left(\left(-\frac{1}{2}+N\right)Q\cdot i-\frac{1}{2}\cdot
Q\cdot i^{2}\right)+Q\cdot j\right)+l$

In this expression, the division operator is rational, $Q$ and $N$ correspond
to the dimensions in Figure 3. Compression and decompression can then be
achieved by iterating over the accessed indices and copying the data in and
out of the compressed buffer.

Progressive Code Generation. DASTAC then proceeds to code generation. It
leverages isl’s Abstract Syntax Tree (AST) generation feature, taking a union
of iteration domains and optional validity dependencies and returns an AST
enforcing those conditions. This AST is then used to generate code in the
Affine dialect, a polyhedral-based intermediate representation defined in the
MLIR framework. We drive the compilation from this IR down to executable code
through MLIR’s provided passes to generate low-level code implementing the
used affine operations, simplify the polynomials computation and hoist each
part as high as possible in the loop nest, and parallelize the computation by
mapping to OpenMP runtime calls. Finally, the resulting LLVM IR is passed to
Clang, which provides auto-vectorization and profits from the new
opportunities to do so offered by this densely packed data layout
representation.

scf.parallel (step (scf.for scf.for scf.for store

Figure 4. The compressed kernel code only consists of cheap arithmetic
operations, loops, and loads and stores on data. In particular, no index and
offset values are loaded from memory such that all memory bandwidth is
available for the actual compute workload

##  4\. Data Layout Compression

In this section, we present our method for data layout compression using a
novel symbolic indexing algorithm, enabling a densely packed representation of
the structured sparse tensor elements. This method relies on a symbolic and
linear indexing of the values used in a computation. As opposed to well-
established sparse data layouts, such as CSR, CSC, or COO, that incur access
indirections and index storage overhead, our method uses direct symbolic index
computations. This direct indexing simultaneously reduces the global memory
footprint and opens vectorization opportunities, two excellent properties when
memory bandwidth is a significant constraint on the program efficiency, as is
the case on many widespread contemporary hardware, such as CPUs and GPUs.

DASTAC starts the compression process from a tensor computation expressed in
STUR (e.g., Figure 3). DASTAC restricts itself to the STUR rules expressed
with quasi-affine expressions (affine expressions augmented with modulos).
This enables the implementation to leverage mature tools to infer this
compressed indexing efficiently.

To sidestep combinatorial complexity, DASTAC only works with individual
summands of a rule. Take for example the following rule:

$O(j):=A(i,j)*(i=1)*(0\leq j<N)+A(i,j)*(i=3)*(0\leq j<N)$

DASTAC will first infer linear indices for $O(j)$ and $A(i,j)$ over the first
region, where $i=1$. Only then does it move on to the $i=3$ region. In the
second region, with $i=3$, a disjoint region of $A$ is read from, so a second
compressed buffer is inferred for this region. On the other hand, the same
region of $O$ is being written to; DASTAC thus ”deduplicates” the compressed
buffer by simply reusing the one infered for the first region.

This design decision avoids compile-time costly generalization of the process
and is found to work effectively on many motivated cases. However, in all
generality, this yields two limitations. First, some values might be
duplicated in different compressed buffers, potentially leading to a
suboptimal memory footprint reduction in some specific cases. Second, the
output compression is only used when output regions are equal (as in the above
example) or disjoint, as any other case would potentially lead to partial
results being spread across multiple buffers, requiring further computation
overhead to have the right values in each buffer.

1:function SymbolicIndexing($\mathcal{D}_{\mathbf{p}}$, $T$)

2: $\mathcal{A}_{\mathbf{p}}\leftarrow$
AccessMapping($\mathcal{D}_{\mathbf{p}},T$)

3: $\mathcal{B}_{\mathbf{p},\mathbf{i}}\leftarrow$
LocalOrdering($\mathcal{A}_{\mathbf{p}},T$)

4: $T_{index}(\mathbf{i},\mathbf{p})\leftarrow$
SymbolicCount($\mathcal{B}_{\mathbf{p},\mathbf{i}}$)

5: return PiecewiseFusion($T_{index}$)

6:end function

Algorithm 1 Symbolic Indexing Algorithm

Symbolic Indexing Algorithm. At this stage, DASTAC uses our symbolic indexing
algorithm (Algorithm 1) to infer a compressed indexing function. The algorithm
takes two inputs. The first input is $\mathcal{D}_{\mathbf{p}}$ which is an
ordered iteration space, where the parameter $\mathbf{p}$ defines symbols such
as dimension information. In other words, $\mathcal{D}_{\mathbf{p}}$
represents the set of iterator values with which the tensor access $T$ is
used. By ordered we mean a sequential execution of the computation. The second
input denoted by $T$, is the considered tensor access in the form of a
function from an iteration to the indices used in this access. The described
process in Algorithm 1 is illustrated using buffer A of our example in Figure
3.

Iteration space extraction. To apply Algorithm 1, we define the iteration
space of the summand at hand, taking all comparisons of the considered STUR
summand. For our running example in Figure 3, the iteration space is:

$\mathcal{D}_{Q,N,M,P}:=\\{(i,j,k,l)\in\mathbb{Z}^{4}\mid 0\leq i<M,i\leq
j<N,0\leq k<P,0\leq l<Q\\}$

Here, symbols correspond to the boundary information of dimensions
($Q,N,M,P$). By restricting STUR to affine expressions, one can observe that
this set is a polyhedron by construction.

Because we define the iteration space from a high-level representation, we
control the order of the dimensions; we simply order them based on their
appearance in the rule, ensuring the first iterators are used to index the
output, with the last one of them indexing contiguous elements.

Access mapping. For each buffer, we define its access map using the access
variables used in the STUR expression. By construction, those are affine maps.
The corresponding map for buffer $B$ in Figure 3 is as follows:

$a_{B}(i,j,k,l)=(i,j,l)$
$\mathcal{A}_{B,\overrightarrow{p}}:=a_{B}(\mathcal{D}_{\overrightarrow{p}})=\\{\overrightarrow{j}\in\mathbb{Z}^{m}\mid\exists
i\in\mathbb{Z}^{n}:a_{B}\cdot\overrightarrow{i}=\overrightarrow{j},\mathcal{C}\cdot\overrightarrow{i}+\overrightarrow{c}\geq
0\\}$

We apply this mapping on the above iteration space, yielding the set of
accessed indices:

(2)
$\begin{split}\mathcal{A}^{B}_{Q,N,M,P}=a_{B}(\mathbf{\mathcal{D}_{Q,N,M,P}})&=\\{(i,j,l)\in\mathbb{Z}^{3}\mid
0\leq i<M,i\leq j<N,0\leq l<Q\\}\end{split}$

Preceding Accesses. We proceed to define an order on this accessed domain,
using the access map to respect the iteration order. This order is similar to
the lexicographic order, only using indices in their iteration domain order.
For the map $a(i,j,k,l)=(k,j,l)$, this order will be:

$a(i^{\prime},j^{\prime},k^{\prime},l^{\prime})<a(i,j,k,l)\iff
j^{\prime}<j\vee j^{\prime}=j\wedge k^{\prime}<k\vee j^{\prime}=j\wedge
k^{\prime}=k\wedge l^{\prime}<l$

We use the set of iterators used in the access map in their initial order:

Applying this order to the accessed domain, we define the preceding access
domain as the set of all accessed indices preceding a symbolic index. For our
running examples, this yields:

$\begin{split}\mathcal{P}^{B}_{Q,N,M,P,i,j,l}&=\\{(i^{\prime},j^{\prime},l^{\prime})\in\mathcal{A}^{B}_{Q,N,M,P}\mid
i^{\prime}<i\\}\\\
&\cup\\{(i^{\prime},j^{\prime},l^{\prime})\in\mathcal{A}^{B}_{Q,N,M,P}\mid
i^{\prime}=i\wedge j^{\prime}<j\\}\\\
&\cup\\{(i^{\prime},j^{\prime},l^{\prime})\in\mathcal{A}^{B}_{Q,N,M,P}\mid
i^{\prime}=i\wedge j^{\prime}=j\wedge l^{\prime}<l\\}\end{split}$

Access domain linearization. Finally, we can synthesize a mapping from a
current iteration to the desired compressed linear index by counting the
number of effective accesses preceding this iteration, that is, counting the
number of elements of $\mathcal{A}_{\mathbf{p},\mathbf{i}}$. The Barvinok
algorithm (Verdoolaege et al., 5 01) is a method for counting the number of
integer points within a parametric polyhedron, returning a piecewise quasi-
polynomial (a generalization of polynomials having periodic coefficients) of
the polyhedron’s parameters.

We can thus use this algorithm, implemented in isl, on the preceding access
domain and obtain a linear indexing of the accessed domains element, following
our chosen order:

$P_{B}=\\#\mathcal{P}^{B}_{\mathbf{p},\mathbf{i}}$

0
---
1
---
2
---
$\frac{i(i+1)}{2}$
---
$\frac{i(i+1)}{2}+j$
---
$j$
---
⋮
---
$\frac{n(n+1)}{2}$
---
?
---
$j$
---
$i$
---

Figure 5. Illustration of the affine symbolic indexing method, ensuring a
dense packing and memory compression.

1:function PiecewiseFusion($T_{index}$)

2: for all $(D,P)$ in $T_{index}$ do

3: for all $(D^{\prime},P^{\prime})$ in $T_{index}-(D,P)$ do

4: if $(P-P^{\prime})|_{D^{\prime}}=0$ then

5: $D\leftarrow D\cup D^{\prime}$

6: $T_{index}\leftarrow T_{index}-(D^{\prime},P^{\prime})$

7: end if

8: end for

9: end for

10: return $T_{index}$

11:end function

Algorithm 2 Piecewise QuasiPolynomial Fusion Algorithm

Piecewise Polynomial Simplification. In general, the counting algorithm yields
a _piece-wise_ quasi-polynomial, which could have a different expression on
distinct subsets of the counted set. In practice, it often yields special
cases at boundaries, that are not desirable for simple code generation. For
example, when counting the preceding accesses for a triangle as illustrated in
Figure 5, the resulting piecewise polynomial is:

(3) $P=\begin{cases}\frac{1}{2}i+\frac{1}{2}i^{2}+j&\text{if }j>0\\\
\frac{1}{2}i+\frac{1}{2}i^{2}&\text{if }j=0\end{cases}$

We apply a piecewise polynomial fusion algorithm presented in Algorithm 2 to
simplify the expression of this polynomial. Let us apply this algorithm to our
example. We name the two regions as $D$ and $D^{\prime}$, and their
corresponding polynomials as $P$ and $P^{\prime}$:
$P=\frac{1}{2}i+\frac{1}{2}i^{2}+j$,
$P^{\prime}=\frac{1}{2}i+\frac{1}{2}i^{2}$. Following line 4 of Algorithm 2,
we have $P-P^{\prime}=j$. Restricting this difference to $D^{\prime}:j=0$, we
have $(P-P^{\prime})|_{D^{\prime}}=0$; in other words, $P$ is equal to
$P^{\prime}$ on the latter’s domain. We thus remove the special case on this
domain (line 6), only keeping $P$ on the union of $D$ and $D^{\prime}$,
yielding a single polynomial expressive enough for the whole buffer:

$P=\frac{1}{2}i+\frac{1}{2}i^{2}+j$

The computed compressed index for buffer $A$ in Figure 3 is as follows:

$P_{B}(i,j,l)=\left(\left(\left(-\frac{1}{2}Q+Q\cdot N\right)\cdot
i-\frac{1}{2}Q\cdot i^{2}\right)+Q\cdot j\right)+k$

In this example, one is expressive enough for the whole buffer; we can thus
safely ignore the conditions, knowing by construction we will not index any
value out of those bounds.

An interesting side-effect of our design is that the generated piece-wise
quasi-polynomials are readily expressed in a form where the slowest varying
indices are evaluated the deepest in the expression, making it amenable to
effective loop-invariant code hoisting. For example, in $P_{B}$, the most
product-intensive term ($-\frac{1}{2}Q+Q\cdot N$) can be hoisted of all loops
in the computation, as it only depends on the sizes of the tensors.

func.func @plus10(func.return

Figure 6. Example function in MLIR IR.

As our experimental results show, our symbolic indexing algorithm
significantly reduces the memory footprint. Furthermore, we observed
performance improvements thanks to improved cache locality and enabling
vectorization opportunities. Finally, we have not observed significant
performance overhead caused by the index computation time on the selected
motivated cases.

##  5\. Progressive Code Generation

Finally, STUR starts generating code using the MLIR framework. It has been
chosen as a compiler for this work for its ability to express high-level
Intermediate Representation (IR) and lower it progressively to LLVM IR, a low-
level IR that the CLang compiler can compile down to binary. This enables MLIR
to provide many optimizations at a high level, where the information needed to
apply them safely is still present.

###  5.1. MLIR Background

MLIR (Lattner et al., 2021) is a compiler framework successfully applied in
the development of compilers in deep learning (Paszke et al., 2019; Abadi et
al., 2015), fully-homomorphic encryption (Govindarajan and Moses, 2020), and
hardware design (Eldridge et al., 2021; Majumder and Bondhugula, 2021). It
also opens opportunities, such as compiling our high-level IR down to GPU
kernels or other targets rather than generating specialized code and
reinventing the wheel, but this has been left as future work.

MLIR IR elementary objects are Operations, Values, Regions, and Attributes.
Values represent runtime values and are constrained by the Static Single-
Assignment (SSA) form: they can only be assigned once, at definition time.
Operations may take values as input and may define new values as output. They
can be augmented through attributes, representing various compile-time
information. Finally, Operations may contain Regions, which include lists of
Operations and possibly entry arguments: Values defined in the region and
taking their runtime value from the operation.

An example is given in Figure 6. func.func is an Operation representing a
function definition. It has a Region representing its body and a function
type. In its region, arith.constant is an operation only defining a value from
a compile-time known value and type. arith.addi takes this value and the
function’s argument and returns its sum in a new value. Finally, func.return
only consumes this value to represent returning it from the function.

MLIR implements this idea of mixing levels of abstraction through dialects:
collections of operations and types that represent some level of abstraction
but can be combined in a program to allow progressive lowering of different
parts in a controllable manner. Operations names are prefixed by their
dialect’s name: in Figure 6, arith and func were used.

mark: parallel iter: —i—, init: —0—, cond: —i ¡ min(N,M)—, inc: 1 iter: —j—,
init: —i—, cond: —j ¡ N—, inc: 1 iter: —k—, init: —0—, cond: —k ¡ P—, inc: 1
iter: —l—, init: —0—, cond: —l ¡ Q—, inc: 1 S0[i, j, k, l] .

$\longrightarrow$

affine.parallel —affine.for —affine.for —affine.for —S0(i, j, k, l);

Figure 7. The translation of ISL AST to MLIR’s Affine loops.

###  5.2. Affine code generation

Parallelization. DASTAC drives ISL to generate an AST expressing parallelism
over the outermost index if it is not a reduction. We believe that parallel
code generation is a benefit of our densely packed data layout; while similar
compaction could be achieved by simple linear indices increments judiciously
placed in the loop nest, those would effectively result in loop-carried
dependencies, and necessitate further analysis or treatment to parallelize
correctly. DASTAC’s symbolic indexing, on the other hand, only depends on
current iterators and symbols of the computation; once inferred, it stays a
correct mapping however loops are parallelised or further transformed. Thus,
parallelism, loop transformations, and compression usage in the generated code
become completely orthogonal. Our experiments demonstrate the effectiveness of
our compression scheme with simple parallelization (cf. Section 6.2).

Control Flow Code Generation. The first step is obtaining an Abstract Syntax
Tree (AST) scanning the iteration domain from ISL; an example is in Figure 7.
By construction, this AST is guaranteed to have affine bounds and increments.
DASTAC uses MLIR’s affine dialect, which expresses affine loops, affine
conditions, and affine accesses, to generate IR at the closest abstraction
level from its internal one. The main benefit of this approach is that DASTAC
does not need to re-implement the basics: it simply translates affine
expressions from ISL’s to MLIR’s representation and defines loops using those.
MLIR is equipped to lower those expressions to arithmetic computations and has
advanced infrastructure to simplify those computations at a lower level.
Finally, DASTAC generates code corresponding to each statement. For the
computation expressed in the STUR language, that means loading each designated
value from the right-hand side, multiplying them, and adding that product to
the loaded left-hand side value. Memory accesses are translated in two steps
to simplify the usage of the compressed indexing polynomial, as explained
below.

Uncompressed Compute Code Generation. The first step is to apply ISL’s
_schedule_ , an affine mapping from the current iteration to the values to
pass to the statement. In the example given in Figure 4, this schedule is the
identity $\theta(i,j,k,l)=(i,j,k,l)$.

S0[—i—, —j—, —k—, —l—]

$\longrightarrow$

—%si— = affine.apply affine_map¡—(i,j,k,l)-¿(i)—¿(%i, %j, %k, %l) —%sj— =
affine.apply affine_map¡—(i,j,k,l)-¿(j)—¿(%i, %j, %k, %l) —%sk— = affine.apply
affine_map¡—(i,j,k,l)-¿(k)—¿(%i, %j, %k, %l) —%sl— = affine.apply
affine_map¡—(i,j,k,l)-¿(l)—¿(%i, %j, %k, %l)

$\begin{split}&\smash{{\hbox{\pagecolor{color5}A(i, j, k)}}}:=\\\
&\smash{{\hbox{\pagecolor{color6!80}C(k,
l)}}}*\smash{{\hbox{\pagecolor{colorm1}B(i, j, l)}}}\end{split}$

$\longrightarrow$

—%—lastval = affine.load —%A[%si, %sj, %sl]— : memref¡?x?x?xf64¿ —%—factor1 =
affine.load —%C[%sk, %sl]— : memref¡?x?xf64¿ —%—factor2 = affine.load —%B[%si,
%sj, %sk]— : memref¡?x?x?xf64¿ —%—prod = arith.mulf —%—factor1, —%—factor2 :
f64 —%—newval = arith.addf —%—lastval, —%—prod : f64 affine.store —%A[%si,
%sj, %sl]— : memref¡?x?x?xf64¿

Figure 8. ISL’s schedule is naturally translated to affine maps application
and the input computation to accesses and arithmetic.

MLIR’s affine dialect provides an operation to apply a given affine mapping to
a list of indices, affine.apply, so we use it to define new values
corresponding to the scheduled access variables. It then remains to generate
the product and accumulation code from the corresponding accesses in the input
STUR rule, using MLIR’s basic arithmetic operations. THis is all illustrated
in Figure 8

$\begin{split}&\smash{{\hbox{\pagecolor{color5}A(i, j, k)}}}:=\\\
&\smash{{\hbox{\pagecolor{color6!80}C(k,
l)}}}*\smash{{\hbox{\pagecolor{colorm1}B(i, j, l)}}}\end{split}$

$\longrightarrow$

———memref.store — Compressed Compute Code Generation. Because the inferred
compression polynomial is indexing a linear contiguous buffer by directly
mapping indices of the original tensor to the corresponding compressed index,
generating compressed compute code can be done independently of any control
flow; it suffices to generate accesses to the compressed buffers using the
inferred polynomial. This is illustrated in Figure 5.2, where those
polynomials are represented as function calls.  The attentive reader might
notice that we switched from affine.load to memref.load. Both representing a
load operation, the affine variant restricts the access indices to affine
expressions of iterators and symbols. Our usage of quasipolynomials breaks
this constraint. Here, MLIR’s capacity to mix abstractions is helping because
this doesn’t forbid the usage of affine loops and maps where still applicable.

###  5.3. Progressive Lowering

The final step is progressively translating and optimizing this high-level
representation to LLVM IR through MLIR. The result of the optimizations is
illustrated in Figure 4. The critical steps are as follows (cf. Section 6 for
the exact pipeline).  Function inling. The first optimization DASTAC applies
is MLIR’s function inlining, inline. This is done to expose the quasi-
polynomials arithmetic directly in the loop nest, allowing to share
computations and hoisting of different parts as high as possible, amortizing
their theoretical computational cost. Lowering Affine. The first conversion to
execute is from affine operations to arithmetic operations, implementing them
through the lower-affine pass exposed by MLIR. This will expand all affine map
applications, loop bounds, conditions, and accesses to as many basic
arithmetic operations as necessary to correctly implement them. Generic
Optimizations. CSE is directly applied to the resulting IR. Amongst the
different affine and quasipolynomial expressions used in the computation, all
expressed on the same iterators and symbols, there are potentially a lot of
common sub-expressions. By applying it directly at this level, we allow more
opportunities than expecting another lower-level compiler to analyze
simplification safety again. The subsequent critical optimization is
canonicalization, a collection of individual standard optimizations. In this
case, we are primarily interested in its constant folding in case any
opportunity is open. Loop Invariant Code Motion. Finally, DASTAC uses loop-
invariant-code-motion to hoist all arithmetic operations as high as possible
in the loop nest, leaving only minimal iteration arithmetic in the innermost
loop to maximize auto-vectorization opportunities.

Figure 9. Using our symbolic indexing, compressed indexing is natural to
switch to from dense indexing.

##  6\. Experimental Results

In this section, we experimentally evaluate our system on a set of
application-driven selected kernels. We study the following questions:

  * •

How does DASTAC perform compared to the state-of-the-art tensor algebra
frameworks in single- and multi-threaded environments?

  * •

How does the automatic memory compression impact the allocated memory
footprint? How do different compression levels improve performance?

  * •

Can DASTAC support complicated structures (e.g., block butterfly factor matrix
(Chen et al., 2022))? How does the structure’s complexity impact the
compression, run time, and compilation time?

  * •

How does DASTAC scale on various numbers of threads in comparison to state-of-
the-art tensor algebra frameworks?

  * •

What is the impact of applying various optimizations compared to the
StructTensor? Can state-of-the-art backends (e.g., Polygeist) recover DASTAC’s
optimizations?

###  6.1. Experimental Setup

Table 1. Information on the kernels used for the end-to-end benchmarks. Kernel | Structure of $B$ | Name | Compile Time (s)
---|---|---|---
TTM | Diagonal (plane) | TTM_DP | 0.31
|  Fixed $j$ | TTM_J | 0.31
$A(i,j,k):=B(i,j,l)*C(k,l)$ | Upper half cube (UHC) | TTM_UT | 0.35
THP | Diagonal (plane) | THP_DP | 0.32
|  Fixed $i$ | THP_I | 0.30
$A(i,j,k):=$ $B(i,j,k)*C(i,j,k)$ |  Fixed $j$ | THP_J | 0.30
MTTKRP | Diagonal | MTT_D | 0.30
|  Fixed $j$ & UHC | MTT_JUT | 0.33
$A(i,j):=B(i,k,l)$ $*$ $C(k,j)*D(l,j)$ |  Fixed $j$ | MTT_J | 0.33
SpMV | Leslie | SpMV_L | 0.30
| Upper triangular | SpMV_UT | 0.30
$A(i):=B(i,j)*C(j)$ | Diagonal | SpMV_D | 0.28
Polynomial Regression Degree-D | Dense, First Dimension Size = 100000 | PR2 | 0.47
$A^{2},\ldots,A^{2D}$ where | Dense, First Dimension Size = 10000 | PR3 | 1.76
$A^{k}(i_{1},\ldots,i_{k}):=$ $B(x,i_{1})*\ldots*B(x,i_{k})$ | Dense, First Dimension Size = 2000 | PR4 | 43.18

We use a server running Ubuntu 22.04 LTS equipped with a 10-core 2.2 GHz Intel
Xeon Silver 4210 CPU with one thread per core and 220 GB of main memory. The
server has 32 KB, 1 MB, and 13.8 MB of level 1, 2, and 3 cache, respectively.
All C++ code is compiled with Clang 18.0.0 using the following flags:

⬇

-std=c++17 -O3 -fopenmp -ffast-math -ftree-vectorize -march=native -mtune=native

MLIR code is lowered with the pass pipeline:

⬇

builtin.module(canonicalize, cse, affine-expand-index-ops, lower-affine,
finalize-memref-to-llvm,

canonicalize, cse, loop-invariant-code-motion, convert-arith-to-llvm, convert-
scf-to-openmp, canonicalize,

convert-scf-to-cf,convert-func-to-llvm,convert-cf-to-llvm,convert-openmp-to-
llvm,reconcile-unrealized-casts)

DASTAC does not use tiling, as we did not observe a major benefit (cf. Section
6.6).

As a competitor, we use Polygeist (Moses et al., 2021) to activate out-of-the-
box MLIR code generation, polyhedral optimizations (e.g., tiling), and
parallelization on top of StructTensor, since StructTensor lacks such
optimizations as it is. Polygeist is shown to have a better single- and multi-
threaded performance (Moses et al., 2021) compared to other tools such as
Pluto (Bondhugula et al., 2008) and Polly (Grosser et al., 2012). Polygeist
transforms a C code generated based on the inferred structure by StructTensor
to MLIR with the following lowering pipeline:

⬇

–canonicalize –cse –affine-expand-index-ops –affine-parallelize –affine-loop-
tile –affine-loop-

invariant-code-motion –canonicalize –cse –lower-affine –loop-invariant-code-
motion –convert-scf-to-

openmp –convert-polygeist-to-llvm –canonicalize –cse –loop-invariant-code-
motion –canonicalize –cse

The lowered code goes through Clang with identical flags to generate
executable binary files. For all experiments, the average run time of three
runs is reported, as the numbers were stable. All experiments are evaluated
using a warm cache. Cache counters are measured using PAPI 7.1 (Laboratory,
2023). As the competitors, we use the latest version of
TACO111https://github.com/tensor-compiler/taco/tree/2b8ece4c230a5f (Kjolstad
et al., 2017), NumPy 1.26.0 (Harris et al., 2020), PyTorch 2.0 (Paszke et al.,
2019), and TensorFlow 2.14.0 (Abadi et al., 2015) with the XLA backend. All
the Python frameworks are run using Python 3.10.13. A dense implementation of
the kernels does not leverage any structure and has been shown to perform
worse than both sparse and dense tensor algebra frameworks (Ghorbani et al.,
2023); hence, we do not evaluate DASTAC against it.

Table 1 shows kernel information (including the ones used in StructTensor
(Ghorbani et al., 2023)) with their structure, compilation time (time required
for generating MLIR code using DASTAC), and the name we used for them on the
plots. This experimental setup has been used previously for StructTensor
(Ghorbani et al., 2023), a state-of-the-art structured tensor compiler. All
the experiments except for the three kernels SpMV_D, SpMV_L, and MTT_D fit in
the cache. Table 2 lists the kernels used in our experiments and their input
structure representation. For TACO, we picked the best data layout according
to the input structure for each kernel (cf. Table 2).

Table 2. Tensor kernels for DASTAC evaluation. In MTT_J, $j$ is fixed for $D$. In MTT_D, $l=j$ too. Kernel |  $B_{U}$ in STUR | TACO Data Layout
---|---|---
TTM_DP | $(0\leq i<n_{i})*(i=j)*(0\leq l<n_{l})$ | $(D,S,D)*(D,D)\rightarrow(D,S,D)$
TTM_J | $(0\leq i<n_{i})*(j=J)*(0\leq l<n_{l})$ | $(D,S,D)*(D,D)\rightarrow(D,S,D)$
TTM_UT | $(0\leq i<n_{i})*(i\leq j<n_{j})*(0\leq l<n_{l})$ | $(D,S,D)*(D,D)\rightarrow(D,S,D)$
THP_DP | $(0\leq i<n_{i})*(i=j)*(0\leq l<n_{l})$ | $(D,S,D)*(D,D,D)\rightarrow(D,S,D)$
THP_I | $(i=I)*(0\leq j<n_{j})*(0\leq l<n_{l})$ | $(S,D,D)*(D,D,D)\rightarrow(S,D,D)$
THP_J | $(0\leq i<n_{i})*(j=J)*(0\leq l<n_{l})$ | $(D,S,D)*(D,D,D)\rightarrow(D,S,D)$
MTT_D | $(i=k=l)*(0\leq i<n_{i})$ | $(D,S,S)*(D,D)*(D,S)\rightarrow(D,S)$
MTT_JUT | $(0\leq i<k)*(0\leq k<n_{k})*(0\leq l<n_{l})$ | $(D,S,D)*(D,D)*(D,S)\rightarrow(D,S)$
MTT_J | $(0\leq i<n_{i})*(0\leq k<n_{k})*(0\leq l<n_{l})$ | $(D,D,D)*(D,D)*(D,S)\rightarrow(D,S)$
SpMV_L | $(i=0)*(0\leq j<n_{j})+(1\leq i<n_{i})*(j=i-1)$ | $(D,S)*(D)\rightarrow(D)$
SpMV_UT | $(0\leq i<n_{i})*(i\leq j<n_{j})$ | $(D,S)*(D)\rightarrow(D)$
SpMV_D | $(0\leq i<n_{i})*(i=j)$ | $(D,S)*(D)\rightarrow(D)$

###  6.2. End-to-End Experiments

In this subsection, we evaluated the performance of DASTAC against the best
alternative of StructTensor (w/ and w/o manual input data layout compression),
TACO with the best structure for each kernel, and the best run time of
Polygeist applied on the output of StructTensor (w/ and w/o manual input and
output data compression, w/ and w/o tiling) on all the kernels (cf. Table 2)
in single- and multi-threaded environments.

---
|  |
---|---|---
|  |
|  |
|  |
|  |

Figure 10. Performance comparison of various frameworks on ten threads on
several structured tensor kernels. In almost all cases, DASTAC shows
significantly better performance in comparison with all other competitors.

The number of floating point operations (FLOPs) is computed by multiplying the
number of structure-aware iterations for each kernel by the summation of the
number of additions and multiplications in that kernel. We used the structure-
aware FLOPs since all the systems are leveraging the structure. We set a
timeout of ten seconds for all frameworks. All the dense frameworks timed out
on all kernels except SpMV with an upper triangular structure and half of TTM
with an upper triangular structure. Therefore, we do not include them in any
of the figures. The missing points for TACO are due to timeout (for polynomial
regression kernels) or segmentation fault because of excessive memory
allocation (other kernels). On MTT_D, all competitors had a segmentation fault
(excessive memory allocation) for sizes bigger than 120000.

---
|  |
---|---|---
|  |
|  |
|  |
|  |

Figure 11. Performance comparison of various frameworks on one thread on
several structured tensor kernels. In all cases, DASTAC performs on par or
better than StructTensor (w/ and w/o Polygeist).

Multi-Thread. Figure 10 compares the performance of each framework’s best
version. Static schedule is used for all the experiments. We only show the
single-threaded performance results for systems that do not support
parallelism (TACO222TACO parallelization problem is mentioned in this GitHub
issue: The link is removed due to anonymization. for all kernels except SpMV
and StructTensor). As shown in Figure 10, DASTAC performs significantly better
than all the competitors while running on multi threads. Overall, compared to
the competitors, DASTAC can get up to 50$\times$ performance enhancement.

In all scenarios, StructTensor stands in second place from the performance
perspective despite lacking parallelization. This shows the importance of
capturing data structure and having an efficient algorithm that leverages this
information. TACO’s code cannot be parallelized on any of the kernels, except
SpMV, which leads to slower performance. Polygeist automatic parallelization
fails to parallelize the code in most cases. The reason behind it is discussed
in Section 6.5. In the polynomial regression kernels, TACO times out,
Polygeist parallelization descales, and StructTensor is running on a single
thread; hence, their performance is significantly diminished. DASTAC
outperforms them by leveraging the flexibility and efficacy of sparse tensor
algebra algorithms and merging them with optimizations and specializations
provided by dense tensor algebra through polyhedral techniques. Combined with
parallelization, these techniques lead to up to 1 to 2 orders of magnitude
performance boost compared to structure-aware competitors.

Single-Thread. Figure 11 represents the performance of all frameworks running
on a single thread. DASTAC performs on par or outperforms all competitors in
almost all cases while running on a single thread. TACO and StructTensor
achieve closely to DASTAC since they also leverage the structure. However,
DASTAC manages to outperform them in most cases (by up to 1 order of
magnitude) by having a better memory layout, leading to better cache locality
and vectorization.

The downward slope on some kernels (e.g., TTM_DP and TTM_J) is due to
increased cache miss rate (from 0.2% to 12.6% for TTM_DP and from 0.2% to
13.2% for TTM_J) by increasing the size (cache miss rate is measured through
PAPI for all experiments, but not shown due to space constraints). The
generated code by Polygeist fails to outperform DASTAC despite leveraging the
structure, manual data layout compression, and having out-of-the-box
polyhedral and MLIR optimizations. This shows that using Polygeist to recover
affine structure from a lower-level code and applying such optimizations out-
of-the-box without a higher-level analysis does not benefit the performance.

###  6.3. Memory Compression

---
Figure 12. Evaluation of compression impact on memory footprint of DASTAC and
StructTensor against an implementation with no compression. The dimension size
for all kernels is $10000$. DASTAC reaches the maximum compression on all
tensors for each kernel, while StructTensor often does a sub-optimal job.

Impact of Compression on Memory Footprint. We evaluate the impact of the
memory compression technique used by DASTAC in comparison to an implementation
without compression and StructTensor’s manual input compression on six
different structures: 1) TTM_UT, 2) SpMV_UT, 3) SpMV_L, 4) TTM_J, 5) THP_DP,
and 6) MTT_D (cf. Table 2). We calculate the allocated memory for each
structured kernel based on the number of elements allocated for all input and
output tensors. The implementation without compression allocates memory for
all the elements of the tensors. StructTensor reduces the allocated memory to
unique elements of the input based on the structure and all output elements.
DASTAC has a better understanding of compression based on the compressed
tensor computation and keeps only the elements required for tensor computation
in both inputs and output. Therefore, even if a tensor is dense, but some of
its elements are not used in the computation, DASTAC compresses the tensor so
that unused elements are not stored.

All dimensions for all kernels are considered to be $10000$. The compression
rate on the y-axis of Figure 12 is the memory allocated by the implementation
without compression divided by the memory allocated by other systems. The
higher compression rate leads to less memory complexity. Most elements are
unique in kernels such as TTM_UT and SpMV_UT; therefore, the compression rate
is low. Kernels such as SpMV_UT and SpMV_L only need the input tensors to be
compressed since all the output elements are available and unique; hence,
StructTensor and DASTAC have the same compression rate. In all other cases,
input and output tensors require compression in an ideal scenario. In several
kernels, there are unique elements that are not used in the tensor computation
(e.g., MTT_D). Since DASTAC performs computationally structure-aware
compression on these kernels, DASTAC’s compression rate is significantly
higher than StructTensor.

Figure 13. Evaluation of the impact of different levels of compression on the
performance of DASTAC.

Impact of Different Levels of Compression. Figure 13 shows that each level of
compression improves the performance for the majority of kernels. Compressing
the input has no impact on polynomial regression cases since the input is
fully dense. In TTM cases, compression cannot improve the cache locality and
the cache miss rate stays the same even after compression, hence, the
performance is similar to uncompressed code.

In the tensor computations where the input compression significantly reduces
the memory footprint, which leads to a significantly better cache miss rate,
output compression has a minimal impact. On the other hand, in tensor
expressions where the input is dense or compression over only the input cannot
improve the cache accesses significantly, output compression has a larger
impact on enhancing the performance.

|  |  |  |  |  |  |  | …
---|---|---|---|---|---|---|---|---

|  |
---|---|---

Figure 14. Evaluation of connection between compression impact and the number
of summands using DASTAC. A scaled version of the used structure for this
experiment is shown on top.

###  6.4. Impact of Structure’s Complexity on Compression

In this section, we introduce four new complex structures to study the impact
of structure complexity on DASTAC’s compression technique: 1) strided diagonal
to study the impact of the number of summands in a tensor operation, 2) block
butterfly factor matrix to study the impact of the number of the polyhedrals
in the structure, 3) strided band to study the impact of distance of elements,
and 4) sub-triangular to study the impact of density.

Number of Summands in Tensor Computation and Compression. To measure the
connection between these two factors (cf. Figure 14), we consider the
following tensor computation:

$A(i,k)$ | $:=$ | $\sum_{w=1}^{p}B^{w}(i,j)*C^{w}(i,k)$
---|---|---

where $B$ and $C$ have a strided diagonal structure with a random stride. In
other words:

$B^{w}_{U}(i,j)$ | $:=$ | $(0\leq i<N)*(0\leq j<N)*((j-i)\%N=s_{B}^{w})$
---|---|---

The size $N$ is fixed to $2^{14}$ here, and the stride $s_{B}^{w}$ is selected
randomly for each tensor. The output of each of these multiplications is a
strided diagonal with stride $s_{B}^{w}+s_{C}^{w}$. We vary the number of
summands from $1$ to $2^{10}$ in powers of $2$. More summands result in longer
compilation time. Generating code for the compressed version takes longer than
without compression since the compression indexing should be computed for more
summands. For the compressed version, as the size increases, the number of
elements fitting in the cache decreases since the compression cannot
compensate anymore; therefore, the cache miss rate goes from almost $0$ to
almost $100\%$. The performance starts dropping on the same size that the
cache miss rate starts increasing. The generated code without compression has
a consistent cache miss rate around $60\%$, and its performance is also
consistently low. This is due to not caching because of high sparsity.

|  |  |  |
---|---|---|---|---

|  |
---|---|---

Figure 15. Evaluation of connection between compression impact and the number
of polyhedrals using DASTAC. A scaled version of the used structure for this
experiment is shown on top.

Number of Polyhedrals and Compression. To measure the connection between these
two factors (cf. Figure 15), we considered a SpMV computation where the matrix
dimension is $2^{14}$ and has a block butterfly factor matrix (Chen et al.,
2022). The structure is as follows:

$B_{U}(x,y)$ $:=$ $(0\leq i<b)*(0\leq j<b)*(0\leq s<k/2)*(0\leq c<2)*(0\leq
d<2)*$
---
$(0\leq
r<\frac{n}{k})*(x=r*k*b+(\frac{c*b*k}{2})+b*s+i)*(y=r*k*b+(\frac{d*b*k}{2})+b*s+j)$

Matrix $B$ is an $nb\times nb$ matrix. We set $n=2^{10}$ and $b=2^{4}$. To
vary the number of polyhedrals, we unroll this unique set over the variables
$c$, $d$, and $s$ and write it as a summation/union format. We also change the
value of $k$ from $2$ to $2^{10}$ in powers of two. This way, the matrix size
and density stay exactly the same. However, the underlying structure and the
number of polyhedrals change. Therefore, we can measure the impact of changing
the number of polyhedrals alone without other factors interfering. The number
of polyhedrals is $2k$. Increasing the number of polyhedrals increases the
compilation time for compressed and uncompressed code generation with a
similar slope. The cache miss rate for the compressed version increases by
increasing the number of polyhedrals since the data is compressed in several
buffers rather than fewer contiguous buffers; hence the performance drops. The
cache miss rate and performance are consistent for the uncompressed version.
The cache miss rate is lower than the compressed version because the iteration
space already traverses $b\times b$ blocks, and compression disrupts the
contiguous nature of the buffer.

|  |  |  |  |
---|---|---|---|---|---
|
---|---

Figure 16. Evaluation of connection between compression impact and stride of
the matrix using DASTAC. A scaled version of the used structure for this
experiment is shown on top.

Distance of Elements and Compression. To measure the connection between these
two factors (cf. Figure 16), we considered a SpMV computation where the matrix
dimension is $2^{16}$ and has a strided band matrix structure. The structure
is as follows:

$B_{U}(i,j)$ | $:=$ | $(0\leq i<N)*(j=i)+(0\leq i<N)*(j-i=s)+(0\leq i<N)*(i-j=s)$
---|---|---

We vary the stride $s$ from $1$ (tridiagonal) to $2^{16}$ (diagonal). By
increasing $s$, the number of elements is reduced but their distance is
increased. The compressed version mostly fits in the cache and is only
impacted by the number of elements. The fewer the elements, the higher the
performance and the lower the cache miss rate for the compressed version. The
uncompressed version’s cache miss rate increases as the distance of the
elements increases. However, when the distance reaches beyond $2^{12}$, the
impact of the decrease in the number of elements helps the cache miss rate to
drop. Both versions reach maximum performance when the matrix is diagonal.

|  |  |  |  |
---|---|---|---|---|---
|
---|---

Figure 17. Evaluation of connection between compression impact and the density
of the matrix using DASTAC. A scaled version of the used structure for this
experiment is shown on top.

---
|  |
---|---|---
|  |
|  |
|  |
|  |

Figure 18. Comparison of the speed-up scaling of various frameworks on
different numbers of threads on several structured tensor operation kernels.
In almost all cases, DASTAC scales on par or better than the state-of-the-art
tensor algebra frameworks.

Density and Compression. To measure the connection between these two factors
(cf. Figure 17), we considered a SpMV computation where the matrix dimension
is $2^{14}$ and has a sub-triangular structure, meaning that the shape is
triangular but only contains a subset of a upper triangular matrix. The
structure is as follows:

$B_{U}(i,j)$ | $:=$ | $(0\leq i<N)*(0\leq j<N)*(j-i\geq N-k)$
---|---|---

Here, $k$ is the size of the triangle. We vary $k$ from $2^{6}$ to $2^{14}$ in
powers of two. Up to $k=2^{9}$, the density is low enough to keep all the
elements in the cache for the compressed version. Both cases have their best
performance on $k=2^{9}$. After that point, the density and number of elements
become larger, hence the compression technique starts losing its benefit
compared to uncompressed code. They finally merge into the same cache miss
rate and the same performance when the matrix is upper triangular
($k=2^{14}$).

###  6.5. Parallelization Scalability

We evaluate the DASTAC’s ability to scale on all kernels mentioned in Table 2
by varying the number of threads from one to ten (cf. Figure 18). A
considerably large size is selected for all kernels to ensure that the
frameworks do not timeout or generate errors (e.g., segmentation fault). The
scaling factor on the y-axis shows the speed-up of each framework compared to
running them sequentially. As shown in Figure 18, DASTAC can scale almost
linearly in most cases by increasing the number of threads. In nearly all
cases, DASTAC scales on par or better than state-of-the-art tensor algebra
frameworks. In several THP, MTTKRP, and SpMV kernels, DASTAC reaches an almost
linear speed-up compared to running it on one thread. TACO only parallelizes
on SpMV kernels. StructTensor and All other kernels of TACO can not be run in
parallel; therefore, they do not scale by varying the number of threads. Only
DASTAC can consistently scale up on all kernels. PyTorch fails to scale on
polynomial regression kernels, NumPy fails to scale on THP kernels, and
TensorFlow fails to scale on SpMV_UT and SpMV_D. Polygeist does not scale or
even de-scales on the majority of kernels. The automatic parallelization
provided by Polygeist leads to parallelizing inner loops and not the outer
loops on a majority of kernels. This increases overhead while running on more
threads, which leads to a decrease in performance. Furthermore, the code
generated by Polygeist for SpMV_D kernel shows an abnormal behavior on the
lower number of threads, which appeared on a range of sizes.

###  6.6. Polygeist Performance Using the Structure

Figure 19. Evaluation of the impact of different systems’ compression
techniques on the performance.
---
Figure 20. Evaluation of the impact of Polygeist automatic tiling on the
performance.

Figure 19 shows that manually compressing both input and output almost always
improves the performance, except for the cases with a hyper triangular
structure, where the compression damages locality. Converting the manually
compressed code to Polygeist is usually not helpful and in several cases, it
causes performance degradation by removing the existing optimization
opportunities for Clang. DASTAC automatically compresses both input and output
based on the iteration space and the compression algorithm almost always
enhances the performance over all other competitors.

Impact of Automatic Tiling Using Polygeist. Figure 20 shows that out-of-the-
box tiling provided by Polygeist does not enhance the performance. Polygeist
relies on MLIR compile pass to tile the generated code. However, the out-of-
the-box tiling has no impact on the majority of the cases. In cases such as
SpMV_D and SpMV_L, the iteration space is one-dimensional, and tiling
decreases the performance.

##  7\. Related Work

In this section, we elaborate on the research that has been done on dense and
sparse (unstructured and structured) tensor algebra, as well as data layout
compression.

Dense Tensor Algebra. Tensor Comprehension (Vasilache et al., 2018) provides
optimized code for CUDA kernels by designing a polyhedral-based JIT compiler.
TensorFlow (Abadi et al., 2015) and PyTorch (Paszke et al., 2019) are deep
learning frameworks providing efficient CPU and GPU code for tensor operation
kernels. NumPy (Harris et al., 2020) is a fundamental tool for numerical
computing, providing optimized code for dense tensor algebra computations. TVM
(Chen et al., 2018) proposes a compiler to address machine learning-specific
challenges, making it portable to various hardware devices. All of these
systems generate significantly efficient code for dense tensor algebra
computations. However, they lack leveraging the structure of the underlying
data. isl (Verdoolaege, 2010) and CLooG (Bastoul, 2004) are generic
polyhedral-based frameworks that provide advanced loop optimizations and
scheduling and are capable of efficient affine code generation, but they are
not full-fledged compiler frameworks and thus lack the ability to densely pack
the data layout.

Unstructured Sparse Tensor Algebra. The sparse polyhedral framework (Strout et
al., 2018) combines the support for sparse tensor algebra with polyhedral
techniques. TACO (Kjolstad et al., 2017; Chou et al., 2018) proposes a generic
way to handle computation over sparse tensor algebra. SDQL (Shaikhha et al.,
2022; Schleich et al., 2023; Shaikhha et al., 2024) uses a functional language
to compile sparse tensor programs. JAX (Bradbury et al., 2018) is a high-
performance computing tool relying on the MLIR (Lattner et al., 2021)
infrastructure with preliminary support for sparse tensor algebra. SciPy
(Virtanen et al., 2020) is an algebraic library supporting a variety of
kernels supporting both sparse and dense tensor algebra. A DSL to capture
sparsity structure for recursive, pointer-based data structures and perform
optimized computation on them is provided in (Chou and Amarasinghe, 2022).
Recent work on Register Tiling (Wilkinson et al., 2023) relies on low-level
optimizations to better exploit compute resources in those cases, but focuses
on SpMM only. Those frameworks capture dynamic sparsity nicely. However, they
spend extra memory storage to store indexing and go through indirect accesses
for computation. DASTAC automatically compresses the tensor data and uses the
most efficient data layout. Moreover, DASTAC has direct access to tensor
elements by relying on the symbolic polynomial index rather than reading it
from an auxiliary array.

Given the interdependencies between the optimal data layout and the target
architecture details, recent work has focused on composability (Ye et al.,
2023) and interoperability (Bansal et al., 2023) of such transformations to
autotune an optimal scheme. These techniques are complementary to our proposed
compression technique; exposing symbolic indexing as a composable
transformation and augmenting search spaces with its usage is worth exploring
in further work.

Structured Tensor Algebra. LGen (Spampinato and Püschel, 2016) proposes a
polyhedral-based algorithm to capture structure on small-scale fixed-size
linear algebra computations and generate faster code. EGGS (Tang et al., 2020)
unrolls the computation tree to specialize the computation over the sparse
data. Sympiler (Cheshmi et al., 2017) performs compile-time symbolic analysis
to produce optimized code for a set of linear algebra kernels. A method to
capture the unstructured sparsity of matrices in a sparse way to generate
polyhedral-based efficient code for the SpMV kernel is proposed in (Augustine
et al., 2019). All the mentioned works are specialized in structured linear
algebra and leverage a limited set of operations. However, they are not usable
for higher-order tensors and lack a proper automatic data layout compression,
and thus are not suitable competitors. StructTensor (Ghorbani et al., 2023)
captures the structure of tensor algebra computation and propagates it
throughout the computation, leading to efficient C++ code generation. However,
StructTensor’s support for data layout is limited to inputs and requires the
user to provide it manually. StructTensor does not support an automatic data
layout compression despite inferring the structure. Moreover, StructTensor
does not leverage specialized dense tensor algebra computation and relies on
the underlying compiler for optimizations, even though it lowers the problem
of structured tensor algebra computation to dense tensor algebra computation.
Finch (Ahrens et al., 2024, 2023) supports unstructured and structured
sparsity but shares similar limitations. DASTAC supports structured
computation over higher-order tensors. Furthermore, DASTAC packs the
structured data densely using a symbolic indexing algorithm and enables
leveraging polyhedral- and affine-based optimizations.

Data Layout Compression. Tiramisu (Baghdadi et al., 2019) proposes a
polyhedral compiler producing efficient and portable sparse tensor algebra
code by introducing several optimizations, including a hardware-specific data
layout transformation on both CPU and GPU. Polyhedral-based techniques to
infer a data layout transformation with a better locality suitable for stencil
computation over dense tensor algebra are provided in (Henretty et al., 2013).
PolyMage (Mullapudi et al., 2015) utilizes polyhedral techniques to transform
the data layout of intermediate computation results and generates affine code
for it. All the aforementioned frameworks propose and utilize advanced data
layout transformation techniques for enhancing the data locality most suitably
for the underlying hardware or based on the computation graph. However, they
are only permuting the data without compressing based on the structure of the
data. DASTAC automatically infers a data transformation to densely pack the
sparse data and provides a compressed data layout for the input, intermediate,
and output data. A method using the Barvinok counting algorithm (Verdoolaege
et al., 5 01) to index iterations in a parallel loop nest to collapse them for
coarser parallelism through OpenMP is proposed in (Clauss et al., 2017).
However, this approach does not compress the data and is much more costly
compared to DASTAC since it computes the inverse of a quasi-polynomial.

##  8\. Conclusion

In this paper, we introduce DASTAC, the first code generation framework
combining algorithmic optimizations known from sparse tensor algebra
frameworks with performance-optimized low-level code known from decades of
tuning dense tensors. The central design behind DASTAC is a novel symbolic
indexing algorithm to compress the structured tensor data into a contiguous
memory buffer. This algorithm relies on the well-established mathematical
foundation of polyhedral models, enabling additional optimizations such as
optimal conditional placements and parallelization. Our experimental results
show that the compressed data layout by DASTAC achieves state-of-the-art
memory consumption and sequential/parallel execution run time. In the future,
we plan to target architectures such as GPU with applications such as sparse
neural networks that can benefit from our proposed algorithm for densely
packing tensor data. We also plan to consider other structured data including
graphs.

## Acknowledgement

This work was supported by the Engineering and Physical Sciences Research
Council (EPSRC) grant EP/W007940/1 and partly funded by a gift from
RelationalAI. The authors thank Huawei for their support of the distributed
data management and processing laboratory at the University of Edinburgh.

## References

  * (1)
  * Abadi et al. (2015) Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. 2015.  TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems.  https://www.tensorflow.org/ Software available from tensorflow.org.
  * Ahrens et al. (2024) Willow Ahrens, Teodoro Fields Collin, Radha Patel, Kyle Deeds, Changwan Hong, and Saman P. Amarasinghe. 2024.  Finch: Sparse and Structured Array Programming with Control Flow.  _CoRR_ abs/2404.16730 (2024).  https://doi.org/10.48550/ARXIV.2404.16730 arXiv:2404.16730
  * Ahrens et al. (2023) Willow Ahrens, Daniel Donenfeld, Fredrik Kjolstad, and Saman P. Amarasinghe. 2023.  Looplets: A Language for Structured Coiteration. In _Proceedings of the 21st ACM/IEEE International Symposium on Code Generation and Optimization, CGO 2023, Montréal, QC, Canada, 25 February 2023- 1 March 2023_ , Christophe Dubach, Derek Bruening, and Ben Hardekopf (Eds.). ACM, 41–54.  https://doi.org/10.1145/3579990.3580020
  * Augustine et al. (2019) Travis Augustine, Janarthanan Sarma, Louis-Noël Pouchet, and Gabriel Rodríguez. 2019.  Generating piecewise-regular code from irregular structures. In _Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI 2019, Phoenix, AZ, USA, June 22-26, 2019_ , Kathryn S. McKinley and Kathleen Fisher (Eds.). ACM, 625–639.  https://doi.org/10.1145/3314221.3314615
  * Baghdadi et al. (2019) Riyadh Baghdadi, Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yunming Zhang, Patricia Suriana, Shoaib Kamil, and Saman Amarasinghe. 2019.  Tiramisu: A polyhedral compiler for expressing fast and portable code. In _2019 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)_. IEEE, 193–205.
  * Bansal et al. (2023) Manya Bansal, Olivia Hsu, Kunle Olukotun, and Fredrik Kjolstad. 2023.  Mosaic: An Interoperable Compiler for Tensor Algebra.  _Proc. ACM Program. Lang._ 7, PLDI, Article 122 (jun 2023), 26 pages.  https://doi.org/10.1145/3591236
  * Bastoul (2004) Cédric Bastoul. 2004.  Code Generation in the Polyhedral Model Is Easier Than You Think. In _13th International Conference on Parallel Architectures and Compilation Techniques (PACT 2004), 29 September - 3 October 2004, Antibes Juan-les-Pins, France_. IEEE Computer Society, 7–16.  https://doi.org/10.1109/PACT.2004.10018
  * Bondhugula et al. (2008) Uday Bondhugula, Albert Hartono, J Ramanujam, and P Sadayappan. 2008.  Pluto: A practical and fully automatic polyhedral program optimization system. In _Proceedings of the ACM SIGPLAN 2008 Conference on Programming Language Design and Implementation (PLDI 08), Tucson, AZ (June 2008)_. Citeseer.
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018.  _JAX: composable transformations of Python+NumPy programs_.  http://github.com/google/jax
  * Chen et al. (2022) Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Ré. 2022.  Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models. In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net.  https://openreview.net/forum?id=Nfl-iXa-y7R
  * Chen et al. (2018) Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. 2018\.  $\\{$TVM$\\}$: An automated $\\{$End-to-End$\\}$ optimizing compiler for deep learning. In _13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)_. 578–594.
  * Cheshmi et al. (2017) Kazem Cheshmi, Shoaib Kamil, Michelle Mills Strout, and Maryam Mehri Dehnavi. 2017.  Sympiler: Transforming Sparse Matrix Codes by Decoupling Symbolic Analysis. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ (Denver, Colorado) _(SC ’17)_. Association for Computing Machinery, New York, NY, USA, Article 13, 13 pages.  https://doi.org/10.1145/3126908.3126936
  * Chou and Amarasinghe (2022) Stephen Chou and Saman Amarasinghe. 2022.  Compilation of dynamic sparse tensor algebra.  _Proceedings of the ACM on Programming Languages_ 6, OOPSLA2 (2022), 1408–1437.
  * Chou et al. (2018) Stephen Chou, Fredrik Kjolstad, and Saman P. Amarasinghe. 2018.  Format abstraction for sparse tensor algebra compilers.  _Proc. ACM Program. Lang._ 2, OOPSLA (2018), 123:1–123:30.  https://doi.org/10.1145/3276493
  * Clauss et al. (2017) Philippe Clauss, Ervin Altintas, and Matthieu Kuhn. 2017.  Automatic collapsing of non-rectangular loops. In _2017 IEEE International Parallel and Distributed Processing Symposium (IPDPS)_. IEEE, 778–787.
  * De Lathauwer and De Moor (1998) Lieven De Lathauwer and Bart De Moor. 1998.  From matrix to tensor: Multilinear algebra and signal processing. In _Institute of mathematics and its applications conference series_ , Vol. 67. Citeseer, 1–16.
  * Dongarra et al. (1990) J. J. Dongarra, Jeremy Du Croz, Sven Hammarling, and I. S. Duff. 1990.  A Set of Level 3 Basic Linear Algebra Subprograms.  _ACM Trans. Math. Softw._ 16, 1 (mar 1990), 1–17.  https://doi.org/10.1145/77626.79170
  * Dumitras and Kossentini (2000) Adriana Dumitras and Faouzi Kossentini. 2000.  Feedforward neural network design with tridiagonal symmetry constraints.  _IEEE Trans. Signal Process._ 48, 5 (2000), 1446–1454.  https://doi.org/10.1109/78.839989
  * Eldridge et al. (2021) Schuyler Eldridge, Prithayan Barua, Aliaksei Chapyzhenka, Adam Izraelevitz, Jack Koenig, Chris Lattner, Andrew Lenharth, George Leontiev, Fabian Schuiki, Ram Sunder, et al. 2021\.  MLIR as hardware compiler infrastructure. In _Workshop on Open-Source EDA Technology (WOSET)_.
  * Gareev et al. (2018) Roman Gareev, Tobias Grosser, and Michael Kruse. 2018.  High-performance generalized tensor operations: A compiler-oriented approach.  _ACM Transactions on Architecture and Code Optimization (TACO)_ 15, 3 (2018), 1–27.
  * Ghorbani et al. (2023) Mahdi Ghorbani, Mathieu Huot, Shideh Hashemian, and Amir Shaikhha. 2023.  Compiling Structured Tensor Algebra.  _Proc. ACM Program. Lang._ 7, OOPSLA2, Article 229 (oct 2023), 30 pages.  https://doi.org/10.1145/3622804
  * Gondimalla et al. (2019) Ashish Gondimalla, Noah Chesnut, Mithuna Thottethodi, and TN Vijaykumar. 2019.  SparTen: A sparse tensor accelerator for convolutional neural networks. In _Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture_. 151–165.
  * Govindarajan and Moses (2020) Sanath Govindarajan and William S. Moses. 2020.  SyFER-MLIR: Integrating Fully Homomorphic Encryption Into the MLIR Compiler Framework.
  * Grosser et al. (2012) Tobias Grosser, Armin Größlinger, and Christian Lengauer. 2012.  Polly - Performing Polyhedral Optimizations on a Low-Level Intermediate Representation.  _Parallel Process. Lett._ 22 (2012).  https://api.semanticscholar.org/CorpusID:18533155
  * Harris et al. (2020) Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. 2020.  Array programming with NumPy.  _Nature_ 585, 7825 (Sept. 2020), 357–362.  https://doi.org/10.1038/s41586-020-2649-2
  * Hegde et al. (2019) Kartik Hegde, Hadi Asghari-Moghaddam, Michael Pellauer, Neal Crago, Aamer Jaleel, Edgar Solomonik, Joel Emer, and Christopher W Fletcher. 2019.  Extensor: An accelerator for sparse tensor algebra. In _Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture_. 319–333.
  * Henretty et al. (2013) Tom Henretty, Richard Veras, Franz Franchetti, Louis-Noël Pouchet, Jagannathan Ramanujam, and Ponnuswamy Sadayappan. 2013.  A stencil compiler for short-vector simd architectures. In _Proceedings of the 27th international ACM conference on International conference on supercomputing_. 13–24.
  * Jia et al. (2021) Liancheng Jia, Zizhang Luo, Liqiang Lu, and Yun Liang. 2021.  Tensorlib: A spatial accelerator generation framework for tensor algebra. In _2021 58th ACM/IEEE Design Automation Conference (DAC)_. IEEE, 865–870.
  * Jouppi et al. (2023) Norm Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, et al. 2023\.  Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In _Proceedings of the 50th Annual International Symposium on Computer Architecture_. 1–14.
  * Khamis et al. (2020) Mahmoud Abo Khamis, Hung Q. Ngo, XuanLong Nguyen, Dan Olteanu, and Maximilian Schleich. 2020.  Learning Models over Relational Data Using Sparse Tensors and Functional Dependencies.  _ACM Trans. Database Syst._ 45, 2 (2020), 7:1–7:66.  https://doi.org/10.1145/3375661
  * Khoromskaia and Khoromskij (2018) Venera Khoromskaia and Boris N Khoromskij. 2018.  _Tensor numerical methods in quantum chemistry_.  Walter de Gruyter GmbH & Co KG.
  * Kissel and Diepold (2023) Matthias Kissel and Klaus Diepold. 2023.  Structured Matrices and Their Application in Neural Networks: A Survey.  _New Gener. Comput._ 41, 3 (2023), 697–722.  https://doi.org/10.1007/S00354-023-00226-1
  * Kjolstad et al. (2017) Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. 2017.  The Tensor Algebra Compiler.  _Proc. ACM Program. Lang._ 1, OOPSLA, Article 77 (Oct. 2017), 29 pages.  https://doi.org/10.1145/3133901
  * Laboratory (2023) Innovative Computing Laboratory. 2023.  The Performance Application Programming Interface (PAPI).  https://icl.utk.edu/papi/
  * Lattner et al. (2021) Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. 2021.  MLIR: Scaling Compiler Infrastructure for Domain Specific Computation. In _2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)_. IEEE, 2–14.  https://doi.org/10.1109/CGO51591.2021.9370308
  * Majumder and Bondhugula (2021) Kingshuk Majumder and Uday Bondhugula. 2021.  HIR: An mlir-based intermediate representation for hardware accelerator description.  _arXiv preprint arXiv:2103.00194_ (2021).
  * Moczulski et al. (2016) Marcin Moczulski, Misha Denil, Jeremy Appleyard, and Nando de Freitas. 2016.  ACDC: A Structured Efficient Linear Layer. In _4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings_ , Yoshua Bengio and Yann LeCun (Eds.).  http://arxiv.org/abs/1511.05946
  * Moses et al. (2021) William S. Moses, Lorenzo Chelini, Ruizhe Zhao, and Oleksandr Zinenko. 2021.  Polygeist: Raising C to Polyhedral MLIR. In _Proceedings of the ACM International Conference on Parallel Architectures and Compilation Techniques_ (Virtual Event) _(PACT ’21)_. Association for Computing Machinery, New York, NY, USA, 12 pages.
  * Mullapudi et al. (2015) Ravi Teja Mullapudi, Vinay Vasista, and Uday Bondhugula. 2015.  Polymage: Automatic optimization for image processing pipelines.  _ACM SIGARCH Computer Architecture News_ 43, 1 (2015), 429–443.
  * Muti and Bourennane (2007) Damien Muti and Salah Bourennane. 2007.  Survey on tensor signal algebraic filtering.  _Signal Processing_ 87, 2 (2007), 237–249.
  * Mutlu et al. (2019) Erdal Mutlu, Karol Kowalski, and Sriram Krishnamoorthy. 2019.  Toward generalized tensor algebra for ab initio quantum chemistry methods. In _Proceedings of the 6th ACM SIGPLAN International Workshop on Libraries, Languages and Compilers for Array Programming_. 46–56.
  * Nion and Sidiropoulos (2010) Dimitri Nion and Nicholas D Sidiropoulos. 2010.  Tensor algebra and multidimensional harmonic retrieval in signal processing for MIMO radar.  _IEEE Transactions on Signal Processing_ 58, 11 (2010), 5693–5705.
  * Panagakis et al. (2021) Yannis Panagakis, Jean Kossaifi, Grigorios G Chrysos, James Oldfield, Mihalis A Nicolaou, Anima Anandkumar, and Stefanos Zafeiriou. 2021.  Tensor methods in computer vision and deep learning.  _Proc. IEEE_ 109, 5 (2021), 863–890.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019.  PyTorch: An Imperative Style, High-Performance Deep Learning Library. In _Advances in Neural Information Processing Systems 32_. Curran Associates, Inc., 8024–8035.  http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
  * Schleich et al. (2023) Maximilian Schleich, Amir Shaikhha, and Dan Suciu. 2023.  Optimizing tensor programs on flexible storage.  _Proceedings of the ACM on Management of Data_ 1, 1 (2023), 1–27.
  * Shaikhha et al. (2024) Amir Shaikhha, Mathieu Huot, and Shideh Hashemian. 2024.  A Tensor Algebra Compiler for Sparse Differentiation. In _IEEE/ACM International Symposium on Code Generation and Optimization, CGO 2024, Edinburgh, United Kingdom, March 2-6, 2024_ , Tobias Grosser, Christophe Dubach, Michel Steuwer, Jingling Xue, Guilherme Ottoni, and ernando Magno Quintão Pereira (Eds.). IEEE, 1–12.  https://doi.org/10.1109/CGO57630.2024.10444787
  * Shaikhha et al. (2022) Amir Shaikhha, Mathieu Huot, Jaclyn Smith, and Dan Olteanu. 2022.  Functional collection programming with semi-ring dictionaries.  _Proc. ACM Program. Lang._ 6, OOPSLA1 (2022), 1–33.  https://doi.org/10.1145/3527333
  * Sidiropoulos et al. (2017) Nicholas D Sidiropoulos, Lieven De Lathauwer, Xiao Fu, Kejun Huang, Evangelos E Papalexakis, and Christos Faloutsos. 2017.  Tensor decomposition for signal processing and machine learning.  _IEEE Transactions on signal processing_ 65, 13 (2017), 3551–3582.
  * Spampinato and Püschel (2016) Daniele G. Spampinato and Markus Püschel. 2016.  A basic linear algebra compiler for structured matrices. In _Proceedings of the 2016 International Symposium on Code Generation and Optimization, CGO 2016, Barcelona, Spain, March 12-18, 2016_ , Björn Franke, Youfeng Wu, and Fabrice Rastello (Eds.). ACM, 117–127.  https://doi.org/10.1145/2854038.2854060
  * Strout et al. (2018) Michelle Mills Strout, Mary W. Hall, and Catherine Olschanowsky. 2018.  The Sparse Polyhedral Framework: Composing Compiler-Generated Inspector-Executor Code.  _Proc. IEEE_ 106, 11 (2018), 1921–1934.  https://doi.org/10.1109/JPROC.2018.2857721
  * Tang et al. (2020) Xuan Tang, Teseo Schneider, Shoaib Kamil, Aurojit Panda, Jinyang Li, and Daniele Panozzo. 2020.  EGGS: Sparsity-Specific Code Generation.  _Comput. Graph. Forum_ 39, 5 (2020), 209–219.  https://doi.org/10.1111/cgf.14080
  * Terán et al. (2019) Julio E Terán, Yovani Marrero-Ponce, Ernesto Contreras-Torres, César R García-Jacas, Ricardo Vivas-Reyes, Enrique Terán, and F Javier Torres. 2019.  Tensor algebra-based geometrical (3D) biomacro-molecular descriptors for protein research: Theory, applications and comparison with other methods.  _Scientific Reports_ 9, 1 (2019), 11391.
  * Vasilache et al. (2018) Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. 2018.  Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions.  _CoRR_ abs/1802.04730 (2018).  arXiv:1802.04730 http://arxiv.org/abs/1802.04730
  * Verdoolaege (2010) Sven Verdoolaege. 2010.  isl: An Integer Set Library for the Polyhedral Model. In _Mathematical Software – ICMS 2010_ , Komei Fukuda, Joris van der Hoeven, Michael Joswig, and Nobuki Takayama (Eds.). Springer Berlin Heidelberg, Berlin, Heidelberg, 299–302.
  * Verdoolaege et al. (5 01) Sven Verdoolaege, Rachid Seghir, Kristof Beyls, Vincent Loechner, and Maurice Bruynooghe. 2007-05-01.  Counting integer points in parametric polytopes using Barvinok’s rational functions.  , 37 - 66 pages.
  * Virtanen et al. (2020) Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, et al. 2020\.  SciPy 1.0: fundamental algorithms for scientific computing in Python.  _Nature methods_ 17, 3 (2020), 261–272.
  * Wang et al. (2014) Endong Wang, Qing Zhang, Bo Shen, Guangyong Zhang, Xiaowei Lu, Qing Wu, Yajuan Wang, Endong Wang, Qing Zhang, Bo Shen, et al. 2014\.  Intel math kernel library.  _High-Performance Computing on the Intel® Xeon Phi™: How to Fully Exploit MIC Architectures_ (2014), 167–188.
  * Wang et al. (2020) Ziheng Wang, Jeremy Wohlwend, and Tao Lei. 2020.  Structured Pruning of Large Language Models. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020_ , Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu (Eds.). Association for Computational Linguistics, 6151–6162.  https://doi.org/10.18653/V1/2020.EMNLP-MAIN.496
  * Wilkinson et al. (2023) Lucas Wilkinson, Kazem Cheshmi, and Maryam Mehri Dehnavi. 2023.  Register Tiling for Unstructured Sparsity in Neural Network Inference.  _Proc. ACM Program. Lang._ 7, PLDI, Article 188 (jun 2023), 26 pages.  https://doi.org/10.1145/3591302
  * Ye et al. (2023) Zihao Ye, Ruihang Lai, Junru Shao, Tianqi Chen, and Luis Ceze. 2023.  SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning. In _Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3_ (¡conf-loc¿, ¡city¿Vancouver¡/city¿, ¡state¿BC¡/state¿, ¡country¿Canada¡/country¿, ¡/conf-loc¿) _(ASPLOS 2023)_. Association for Computing Machinery, New York, NY, USA, 660–678.  https://doi.org/10.1145/3582016.3582047