1

# Programming Matrices as Staged Sparse Rows to Generate Efficient Matrix-free
Differential Equation Solver

Huanqi Cao  Department of Computer Science and Technology & BNRistTsinghua
UniversityBeijingChina<EMAIL_ADDRESS>,  Shizhi Tang
Department of Computer Science and Technology & BNRistTsinghua UniversityChina
<EMAIL_ADDRESS>,  Bowen Yu  Department of Computer Science and
Technology & BNRistTsinghua UniversityChina<EMAIL_ADDRESS>and
Wenguang Chen  Department of Computer Science and Technology & BNRistTsinghua
UniversityChina<EMAIL_ADDRESS>
###### Abstract.

Solving differential equations is a critical task in scientific computing.
Domain-specific languages (DSLs) have been a promising direction in achieving
performance and productivity, but the current state of the art only supports
stencil computation, leaving solvers requiring loop-carried dependencies
aside. Alternatively, sparse matrices can represent such equation solvers and
are more general than existing DSLs, but the performance is sacrificed.

This paper points out that sparse matrices can be represented as programs
instead of data, having both the generality from the matrix-based
representation and the performance from program optimizations. Based on the
idea, we propose the Staged Sparse Row (SSR) sparse matrix representation that
can efficiently cover applications on structured grids. With SSR
representation, users can intuitively define SSR matrices using generator
functions and use SSR matrices through a concise object-oriented interface.
SSR matrices can then be chained and applied to construct the algorithm,
including those with loop-carried dependences. We then apply a set of
dedicated optimizations, and ultimately simplify the SSR matrix-based codes
into straightforward matrix-free ones, which are efficient and friendly for
further analysis.

Implementing BT pseudo application in the NAS Parallel Benchmark, with less
than $10\%$ lines of code compared with the matrix-free reference FORTRAN
implementation, we achieved up to $92.8\%$ performance. Implementing a matrix-
free variant for the High-Performance Conjugate Gradient benchmark, we achieve
$3.29\times$ performance compared with the reference implementation, while our
implementation shares the same algorithm on the same programming abstraction,
which is sparse matrices.

††journal: PACMPL††journalvolume: 1††journalnumber: CONF††article:
1††journalyear: 2018††publicationmonth: 1††copyright: none††ccs: Software and
its engineering General programming languages††ccs: Social and professional
topics History of programming languages

##  1\. Introduction

Numerically solving differential equations is an essential class of scientific
computing programs. In many such applications, including weather prediction
(Skamarock et al., 2019) and seismic imaging (Mcmechan, 2006), the space is
discretized with structured grids, especially Cartesian grids. The data is
then a 2-dimensional or 3-dimensional array, and the computation walks over
the array according to the algorithm. Historically, science researchers write
differential equation solvers on the structured grid by hand with general-
purpose languages, usually FORTRAN and C/C++, with “for” loops and arithmetic
operations. It usually requires thousands of lines of code to implement a
solver for just one problem; a new equation will require a new solver with
another thousands-of-line FORTRAN module. Things become even worse nowadays:
with multi-threading and single instruction multiple data (SIMD), it is more
unacceptable to write a solver by hand for each equation for each hardware
backend.

To address the issue, science researchers have built various domain-specific
languages (DSLs) (Tang et al., 2011; Louboutin et al., 2019; Lengauer et al.,
2020), targeting stencil computation, which can be used to express typical
classes of numerical solvers on structured grids: the explicit solving method
and a particular class of the implicit solving methods restricted to Jacobian
iterations. There can be multiple stencil operators in stencil computation,
where each stencil operator walks

Figure 1. An example for 5-point stencil on a 2D grid.

on the Cartesian grid and computes values for each grid point as a linear
combination of input arrays’ values on nearby points. A typical example is
shown in Fig. 1, which is a 5-point stencil operator on a 2D grid. Given that
stencil operators are linear operators, they can be chained, summed up, or
subtracted, further simplifying the programming. We refer to these operations
as compositions. These DSLs support such compositions of stencil operators.
The composition of stencils is isomorphic to arithmetic between sparse
matrices and allows the code to be closer to the math formula.

But stencil computation only covers limited classes of solvers. Many important
kernels for implicit solvers, including Gauss-Seidel iteration and variants of
Incomplete LU, are not supported by those DSLs due to their loop-carried
dependences. Moreover, these DSLs also struggle with some complicated stencil
operators, e.g., higher-order dissipation that requires special boundary
handling (Olsson and Johnsson, 1989) and restriction/prolongation that relies
on the parity of the grid point coordinates (Trottenberg and Schuller, 2000);
such stencils are both hard to express and hard to participate in the
composition. The DSL designers had to build complicated symbolic rules to
accomplish simpler user programming.

On the other hand, sparse linear algebra libraries serve structured grids as
well, although they primarily target unstructured grids. Applying a stencil
operator is equivalent to multiplying a sparse matrix by a dense vector, in
which the sparse matrix encodes the stencil. Composing the stencil operators
is equivalent to arithmetics between sparse matrices. Also, algorithms with
loop-carried dependences are possible to be expressed as general sparse
matrices as well, and can be implemented as libraries, e.g. SciPy (Virtanen et
al., 2020) includes some classes of solvers as sparse matrix computation. With
such a sparse-matrix representation, programming solvers are much easier.
However, implementing general-purpose sparse routines is hard and only expert
programmers can work on it through extensive human efforts, leaving the
algorithms set incomplete. Besides, the structured grid generates fixed
sparsity patterns, while such libraries are not aware of the concrete sparsity
pattern of the matrix and can only guess for optimization, leading to poor
performance compared with manually implemented routines. People still seek
implementations that are not through the sparse matrices, which they refer to
as “matrix-free”, when performance is under concern. The pros and cons of the
above approach are listed in Table 1.

| Stencil Computations | Other Computations | Performance
---|---|---|---
| Expressivness | Ease of Use | Expressivness | Ease of Use
Hand-written | ✓ | $\times$ | ✓ | $\times$ | ✓
Stencil DSLs | ✓ | ✓ | $-$ | $-$ | ✓
Sparse Matrices | ✓ | ✓ | $\times$ | ✓ | $\times$
SSR | ✓ | ✓ | ✓ | ✓ | ✓
Table 1. Pros (✓) and Cons ($\times$) of different approaches to numerically
solving differential equations.

In this paper, we present a new representation for the sparse matrices, and
the corresponded dedicated optimizations, to accomplish both sparse-linear-
algebra-based easy programming, and high-performance matrix-free
implementation supporting parallelization through multi-threading. We list our
contributions below.

  * •

We present Staged Sparse Row, abbreviated as SSR, which represents sparse
matrices on structured grids as programs instead of data. Borrowing the idea
of multi-stage programming, SSR represents the matrix as a first stage
program, enabling the second stage program to be matrix-free. It encodes the
sparse rows as lazy sequences of indices and values of non-zeros, allowing
programs to traverse the rows without materializing the matrix. By providing a
set of object-oriented interfaces, SSR supports compositions of operators
through arithmetics between matrices and implicit solving with loop-carried
dependences. Furthermore, we present a set of programs based on SSR and four
basic SSR matrices, which cover common types of algorithms and basic operators
on structured grid differential equation solving. The SSR DSL is embedded into
Python for ease of use.

  * •

We further develop a set of dedicated optimizations that optimize the
generated program into a more regular form. The optimizations aim at
simplifying the state machine codes introduced by the lazy sequences, and
ultimately produces optimal serial program with regular nested loops. They
thus enable later parallelization and acceleration by our backend tensor
compiler, FreeTensor (Tang et al., 2022) in our case.

  * •

We evaluate our methodology by implementing benchmarking applications. In NPB-
BT (Bailey et al., 1994), we achieve up to $0.928\times$ speedup compared with
the reference FORTRAN implementation, with less than $10\%$ of codes. The
original FORTRAN implementation is matrix-free and well-optimized, proving the
efficiency of our generated code. In the matrix-free solver for an HPCG-
defined problem (Dongarra et al., 2016), we achieve $3.29\times$ performance
than the reference implementation of HPCG with the programs both written
against sparse matrices. Our matrix-free HPCG implementation is
straightforward and involves only 209 lines of code, while the original HPCG
implementation includes 553 lines of algorithm code.

The next sections will be organized as below. We first introduce the
background and our motivation in detail in Section 2. Then we present our
programming model with example programs in Section 3. The transformations for
dedicated optimizations are presented in Section 4. Finally we present
application cases in Section 5 and performance evaluations in Section 6. As
closing sections, Section 7 will discuss our related work, Section 8 will
discuss future potential of our design and Section 9 will present our
conclusion.

##  2\. Background and Motivation

Many scientific computing workloads involving 2-dimensional or 3-dimensional
physical simulation are solving differential equations. Decades ago, domain
researchers have to write solver code by hand. Today, things are getting
better. In explicit solvers, efficient stencil computation can already be
accomplished via various domain-specific languages (DSLs) and their specific
compilers. In implicit solvers, they can generate sparse matrices and leverage
efficient general-purpose solver libraries to solve sparse linear equations.
But matrix-free sparse linear solving, which has more potential for
optimization, still requires human efforts. In this section, we discuss the
current status and problems in differential equation solving, and then the
status of tensor programming approaches.

###  2.1. Solving Differential Equations on Structured Grids

####  2.1.1. Matrix-free Approaches: Stencil and More

Explicitly solving differential equations on structured grids involves stencil
computations. The stencil computations are usually derived from spatial
differentiation (e.g. $\frac{\delta u}{\delta x}$, in which $x$ is a spatial
dimension) and various point-wise operations. They read neighbor points in the
input 2D or 3D tensors with a certain pattern, go through a certain function
as its computation, and yield the corresponding value in output tensor(s). Due
to its nature of independent computation, simply parallelizing stencils is
straightforward. Further optimizations (Habich et al., 2009; Williams et al.,
2007; Nguyen et al., 2010) include different spatial blocking (or tiling)
strategies (wavefront, 3D, 2.5D, …), temporal blocking, etc.

Since the computation is highly customized across applications and
optimization by hand requires huge redundant work, many domain-specific
languages (DSLs), including Devito (Louboutin et al., 2019) and ExaStencil
(Lengauer et al., 2020), have been designed to address the optimization of
stencil computation. They express stencil computation in different ways, as is
shown in Fig. 2 While in ExaStencil the multiplication operator is used,
instead of algebraic multiplication, it indeed represents “applying the
stencil operator $D$ to the data $x$”, which is the only way to consume a
stencil operator in ExaStencil. Upon their similar design, these works focus
on automatic stencil optimization primarily and have not explored the language
extensively.

$\forall i,y_{i}=x_{i+1}-x_{i}$ (a) Math formula

$y=x.dx$ (b) Devito

$\displaystyle y=D*x$ $\displaystyle\text{Stencil}\ D$
$\displaystyle\quad[1]\Rightarrow 1.0$ $\displaystyle\quad[0]\Rightarrow-1.0$
(c) ExaStencil

Figure 2. Differentiation operators, in math formula, Devito and ExaStencil.

Due to the lack of expressiveness, complicated boundary conditions and
staggered grids, which has spatial patterns and computing functions depending
on the coordinates, are rarely supported in those DSLs. ExaStencil partially
supports such cases by skipping invalid (with indices out-of-bound or non-
integer) entries in the non-zero list, but not yet covering all cases.

While stencil computations are widely researched, many other cases exist to be
non-stencils. Such kernels, while still operating on 2D or 3D structured
grids, have loop-carried dependencies on one or more dimensions. They occur in
implicit differential equation solving as linear system solvers, either
iterative or direct, including Gauss-Seidel iteration, multi-diagonal solver,
incomplete lower-upper (ILU) solver, fine-grained parallel ILU (Chow and
Patel, 2015), and so on. Those operations with loop-carried dependencies are
usually much more complicated than stencils and require even more human
effort. For example, the block tri-diagonal pseudo-application in NAS Parallel
Benchmarks (NPB-BT) writes over 400 lines of Fortran code for solving the
block tri-diagonal matrix in one direction, and it has three directions to
solve, involving over 1300 LOC. However, such routines are not yet supported
by DSLs. ExaStencil is capable of implicitly solving differential equations,
but only through Jacobian iterations in a multi-grid (MG) solver, which is
just a special case limiting its computation pattern to only stencils.

####  2.1.2. Generic Approach: Sparse Linear Algebra

Another method to deal with structured grids involves sparse matrices. Take
the 5-point stencil in Fig. 1 as an example, it computes the following:

$B_{x,y}=\alpha A_{x,y}+\beta(A_{x-1,y}+A_{x+1,y}+A_{x,y-1}+A_{x,y+1})$

If the $A$ and $B$ fields are flattened into vectors $A^{(f)},B^{(f)}$, the
transform above can be seen as a matrix-vector product:

Row domain
---
($(n-2)\times(n-2)$)
Column domain ($n\times n$)$\begin{aligned} B^{(f)}_{xN+y}&=\beta
A^{(f)}_{(x-1)N+y}+\beta A^{(f)}_{xN+(y-1)}+\alpha A^{(f)}_{xN+y}+\beta
A^{(f)}_{xN+(y+1)}+\beta A^{(f)}_{(x+1)N+y}\\\ \\\ (B_{2..n-1,\
2..n-1})^{(f)}&=\setcounter{MaxMatrixCols}{12}\begin{bmatrix}\beta&&&\beta&\alpha&\beta&&&\beta\\\
&\beta&&&\beta&\alpha&\beta&&&\beta\\\
&&\ddots&&&\ddots&\ddots&\ddots&&&\ddots\\\
&&&\beta&&&\beta&\alpha&\beta&&&\beta\\\ \end{bmatrix}\;(A_{1..n,\
1..n})^{(f)}\end{aligned}$

In which the matrix is sparse with a certain structure. Thus, stencils turn
into sparse matrix-vector products (SpMV). The other algorithms with loop
carried dependencies mentioned above, including Gauss-Seidel, ILU, etc., are
also possible to be written as sparse linear algebra routines. Thus, by
implementing those algorithms against sparse matrices, library providers
liberate domain experts from writing those complicated codes, only leaving the
generation of sparse matrices.

The High Performance Conjugate Gradient (HPCG) Benchmark (Dongarra et al.,
2016), which is a widely recognized effort for ranking HPC systems including
supercomputers, generates a synthetic partial differential equation and solves
it through preconditioned conjugate-gradient iterations. As a benchmark
adopted by nearly all top-level supercomputers, it represents a wide range of
real-world applications needs. It explicitly prohibits exploiting the sparsity
pattern, thus forbidding matrix-free approaches, indicating their concern that
real applications require generic optimization methods instead of one ad-hoc
optimization after another. It turns out such prohibition limits the
performance drastically. It is reported that implementing the HPCG computation
in a matrix-free form significantly improves the performance, by $4.67\times$
on the New Sunway supercomputer (Zhu et al., 2021). When possible, HPC
researchers still seek matrix-free approaches even for implicit approaches,
e.g., the 2016 Gordon Bell Prize winner (Yang et al., 2016) designed and
manually implemented a geometry-based pipelined ILU method that maps the data
dependency to hardware-supported inter-core communication, which is a case for
further optimizing with the sparsity pattern in hand.

###  2.2. Multi-stage Programming

Multi-stage programming has been a promising method for building embedded
domain-specific languages. In such practices of multi-stage programming, a
program usually is split into two stages: the first stage generates and
compiles the domain-specific program, and the second executes the compiled
artifact. Lightweight Modular Staging (LMS) (Rompf and Odersky, 2012) is one
of such frameworks and has been adopted in multiple DSL works (Sujeeth et al.,
2014; Ofenbeck et al., 2013; Shaikhha et al., 2018). LMS separates the two
stages through only types: Rep[T] refers to a T-typed value at the second
stage. Interacting with Rep[T] instances, including taking it for control
flows, generates code in the second stage instead of operating them
immediately. Since that LMS is a Scala library, we implement our own and adopt
a similar method to build our embedded DSL in Python. Var[T] in our
implementation is equivalent to the Rep[T] in LMS.

###  2.3. Tensor Programs

Through multi-stage programming, we apply first dedicated optimizations for
differential-equation solvers, and then general optimizations for tensor
programs. Here tensor programs refer to programs manipulating multi-
dimensional arrays, or tensors. There have been many studies on optimizing
tensor programs since the invention of FORTRAN. Program analysis and
transformations for multithreading, vectorization and better cache management
are summarized in (Padua and Wolfe, 1986) and (Allen and Kennedy, 2001). In
order to apply optimizations, languages and their compilers for tensor
programs require programs to be multiple levels of loops densely iterating
along tensors. The purpose of our first stage of compiling is to meet the
requirements of a tensor program. Among recent works, FreeTensor (Tang et al.,
2022) is a language and compiler to optimize such a tensor program. FreeTensor
provides a flexible Python interface to construct the tensor program, which is
ideal as the target language of our multi-stage programming. It further
provides a set of code transforming schedules, as performance tuning
primitives.

Recent years also witness an emergence of tensor program frameworks and
compilers for the development of deep learning, including PyTorch (Paszke et
al., 2017), TensorFlow (Abadi et al., 2015), JAX (Frostig et al., 2018), and
TVM (Chen et al., 2018). These frameworks also provide flexible interfaces for
a variety of tensor programs, by decomposing them as the combination of a set
of basic tensor operations. However, the staged code by our method contains
loop-carried dependences implicitly derived from the loop ordering, the
inference of which is not possible without fine-grained analysis and is not
supported by these frameworks.

##  3\. Matrices as Staged Sparse Row

In this section, we introduce the new sparse matrix representation, namely
Staged Sparse Row (SSR). SSR is a program-level representation instead of a
runtime storage format, unlocking the potential that programs using it can be
lowered to matrix-free code. Through multi-staged programming, the matrix has
its rows being lazy sequences in the first compilation stage, and gets
unrolled into matrix-free codes in the second tensor program stage.

As indicated by the name, SSR organizes sparse matrices by rows, represented
as programs, in the compilation stage. Each row, being a lazy sequence, is
implemented as a generator function, which yields non-zeros one by one once
given the row index. In Fig. 3 we demonstrate an end-to-end example comparing
a sparse matrix as data, as the code to generate its adjacency list, and as
our staged representation.

Figure 3. An end-to-end example for SSR.

(a) The sparse matrix. It has $n$ columns and $n-1$ rows. (b) The Python code
generating the matrix in adjacency lists. (c) The SSR implementation of the
matrix. We see that it is similar to (b).

###  3.1. Defining SSR Matrix with Generator Function

Figure 4. An example of implementing the prolongation operator in SSR.

(a) A coarser grid (orange lines) and its prolonged finer grid (additional
black lines). The prolongation operator interpolates the points in the coarser
grid (orange circles) to the points in the finer grid (all markers). (b)
Prolongation operator written in an SSR generator. It first distinguishes the
types of grid points according to the parity of coordinates, and then selects
the corresponding coarse grid points for averaging. Neighbors’ coordinates and
weights are yielded accordingly. (c) Iterator class translated from the
generator function. The next function moves between states for one step
according to the control flow of the input generator function. The value and
terminated functions retrieve the current value and whether the iterator
terminates, correspondingly.

Implementing lazy sequences by hand can be painful in an imperative language,
and generator functions have been a common way to implement lazy sequences, or
say iterators. Thus, we employ generator functions to express the sparse rows.
The generator function takes two arguments, with one being the CartesianDomain
for retrieving metadata, and the other one being the requested row index.
Considering the nature of spatially-structured sparse matrices, we limit the
generator function to support only if control flow without harming
expressiveness. It allows us to infer the maximum non-zeros count, which is
useful in many routines; we will see the usages in Section 3.3.

Unlike conventional matrices mapping between two 1D vectors, the grids that
occur in differential equations usually have multiple dimensions. To avoid
manually calculating the coordinates from a flattened index, we instead retain
the original shape of the grid, giving the rows and columns their own
Cartesian space. For example, in Fig. 3(c), we give the coordinates to the
generator function as a list of indices, instead of one flattened index; the
generator function yields the non-zero columns with their coordinates as well.
Thus, to finish the definition of a SSR matrix, we further request mapping
functions from the column domain to the row domain: one for mapping the grid
shape, the other for mapping the grid step length (also known as $\delta$).
The $\delta$ mapping function defaults to the identity function. Different
sparse matrices usually have their own domain mapping functions, for either
adjusting the boundary or working with finer or coarser grids in multi-grid
methods. The corresponded column and row domains represent multi-dimensional
grids, i.e. the spatial structure before flattening, which has been introduced
in Section 2.1.2.

Altogether, it is possible to define complex stencil operators through a
generator function, with the power of control flows in it. We demonstrate it
with a simple 2-dimensional prolongation operator in Fig. 4(a) which
interpolates from a $3\times 3$ coarse grid (orange lines) to a $5\times 5$
finer grid (orange and black lines). On even points in the finer grid, it
matches an exact point in the coarser grid, and thus takes the input value. On
odd points, it instead interpolates according to the nearest points.
Implemented as SSR matrix, it simply branches according to the index and
yields non-zeros, as is shown in Fig. 4(b). Then the generator function is
transformed into a finite state machine implementing an iterator, more
precisely an instance of SpRow interface which will be introduced in Section
3.2. The iterator is responsible for iterating over non-zeros in the sparse
row. The domain mapping function is provided through parameters of @spmat: the
first calculates the size of the coarser grid, and the other doubles the grid
step length accordingly. Similar prolongation operators are common in an
important and widely used class of implicit solvers, namely multi-grid, which
generally first scales down (restrict) from a finer grid, then scales up
(prolong) from the restricted coarser grid.

###  3.2. Using SSR Matrix through Object-Oriented Interface

The above generator function defines a SSR matrix, the type of which is SpMat.
Unlike data formats that can be read by user programs directly, to cooperate
with such matrices as programs, a set of programming interfaces are provided.
We introduce what interfaces they provide in this subsection.

Type Definitions |  Description
---|---
type Idx = List[Var[Int]] |  Staged index over a Cartesian domain. Elements are staged integer scalars.
| type CartesianDomain =
---
List[(Int, Float)]
Cartesian domain that fields (Vectors) may lie in. Each dimension is attached
with a grid size and step length.
Members of Vector |
domain: CartesianDomain |  Domain of this Vector. Each element of this Vector is a value in the domain, forming a field in space.
inner_shape: List[Int] |  Shape of each element in the Vector.
Interface provided by Vector |
__getitem__(i: Idx) -> Var[Float] |  Retrieve element at the specified index.
Members of SpMat |
col_domain: CartesianDomain |  Domain of the column index. When multiplied to a Vector, it should match the domain of the field. Initialized as None.
row_domain: CartesianDomain |  Domain of the row index. When multiplied by a Vector, it becomes the domain of the output Vector. Initialized as None.
inner_shape: List[Int] |  Shape of each element in the SpMat.
max_nnz: Int |  Upper bound of non-zeros count in the rows of this matrix. Useful for unrolling loops manually (extensively used in SpGEMM) and preparing space for materialization of matrix values (used in ILU).
Interfaces provided by SpMat |
set_domain(domain: CartesianDomain) |  Set col_domain and infer row_domain accordingly.
row(i: Idx) -> SpRow |  Retrieve a staged row with a given index. Should only construct the tensors used in row iterator without initializing them in order to support placeholder rows.
Interfaces provided by SpRow |
delayed_init() |  Delayed initialization for safe placeholder row creation.
next() |  Move to the next non-zero in the current row.
value() -> (Idx, Var[Float]) |  Retrieve column index and value of the current non-zero. Value is represented as a floating-point tensor of shape inner_shape.
terminated() -> Var[Bool] |  Test if the current row has ended.
Table 2. Core interface of SSR.

The core interfaces are shown in Table 2. SpMat is the base class of any
matrix in SSR. By calling row() on a SpMat, a SpRow instance is retrieved.
SpRow represents a staged lazy sequence of non-zeros, which in turn are tuples
of column indices and values. It provides three interfaces, next(), value(),
and terminated(), to form a typical iterator. Translated from the generator
function, it yields non-zeros in ascending order of the column indices,
similar to the CSR data format. For example in Fig. 4(c), the state-transition
function, next(), is generated from the row generator function, and the other
two interfaces return their results according to the internal state. Unlike
non-staged iterators, these interfaces of SpRow produce staged code instead of
operating on their own. Through these interfaces, SpRow supports to be
iterated through a staged for-each loop.

Corresponding to the multi-dimensional rows and columns, CartesianDomain is
introduced to represent a Cartesian space. Instead of flattening the multi-
dimensional data into flat vectors which is conventional in using sparse
linear algebra libraries, we use CartesianDomain to encode the multi-
dimensional grid of Vectors and SpMats. It consists of the grid shape in
integers, and the grid steps (distance between points on each axis) as
floating-point numbers.

A Vector thus becomes a tensor on the Cartesian domain, index-able with a
Cartesian Idx, which is a list of staged integer expressions. Similarly, an
SpMat has its rows and columns indices being Cartesian Idxs in its row_domain
and col_domain, accordingly. The row’s domain decides the domain of the result
Vector when used as a stencil operator.

As we have mentioned in the previous section, a SSR matrix is suitable for an
arbitrary domain. Its col_domain and row_domain is not decided on
construction, but is set through set_domain. The set_domain also calls the
domain mapping function passed to @spmat to compute the row_domain.

In SpMat and Vector we further allow each element to be tensor of any shape,
nesting inside the grid. We describe it as inner_shape. It is common to have
multiple scalars in the same grid, representing different components of a
vector (velocity, force, etc.; not our Vector) or with different physical
meanings. In our implementation, we further support conventional operations on
the inner tensor, including element-wise add and subtraction, matrix-vector
multiplication, matrix-matrix multiplication, and matrix inversion.

User program using SSR matrices and sparse rows will generate matrix-free yet
complex code, containing all those state machine codes representing rows,
which lazily yield non-zeros. It’s possible to simplify such staged codes into
neat, handwritten-like ones. We will cover the generated code and its
simplification in Section 4.

###  3.3. Writing Programs on SSR

Figure 5. Implementing SpMV on (a) adjacency lists, (b) SSR through manual
staging (code emission), and (c) SSR through type-based automatic staging. (c)
is what we write in practice.

Writing programs on SSR matrix is simple and similar to Compressed Sparse Row.
Only some issues around variable-length lists need to be avoided through
utilizing SpMat.max_nnz and SpRow.delayed_init. We demonstrate four important
examples in this subsection. Among the four examples, SpMV and SpGEMM are
included in our prelude library as common building blocks.

####  3.3.1. Stencil Through SpMV

Stencil operators are equivalent to sparse matrix-vector multiplications
(SpMV) with a sparse matrix representing the connectivity on the structural
grid. To calculate an element in a result vector, we simply walk through the
corresponding row and accumulate the products of matrix non-zeros and source
vector elements. The code is shown in Fig. 5(c). It is worth noticing that the
implementation is nearly identical to the SpMV code on adjacency lists in Fig.
5(a); also, it is much more straightforward compared with manually emitting
the staged code, as is shown in Fig. 5(b).

####  3.3.2. Compose Operators with SpGEMM

Multiple stencil operators applied to a Vector can be composed into one. Speak
in sparse matrices, it is a sparse matrix to sparse matrix multiplication,
abbreviated as SpGEMM. With both operands being SSR matrices, it is possible
to calculate their product still in SSR through row-wise SpGEMM as is shown in
Fig. 6(a).

(a) Algorithm via row iterators.

(b) Implementation on SSR.

Figure 6. Algorithm and implementation of SpGEMM.

As is shown in delayed_init in Fig. 6(b), we walk through the row in left-hand
side (LHS) during initializing a row, retrieving its non-zeros. The non-zeros’
indices are used to index the right-hand side (RHS) matrix, creating row
generators accordingly. With SpRow being a compile-stage type, we cannot
maintain a runtime list of rows, resulting in the list of row generators being
a padded compile-time list of size lhs_mat.max_nnz.

During iteration as is shown in the next method, these ordered row generators
are merged together. The first iteration over RHS rows locates the rows that
have non-zeros with the smallest column indices and accumulates their value.
The second walk advances those rows, moving to the next state. It’s worth
noticing that static_range is used: with rhs_rows being compile-stage list, we
have to use compile-time index on it. Thus the loops are unrolled at compile
stage, yielding lhs_mat.max_nnz distinct branches handling different RHS rows.
The branches will be eliminated by the optimizations we will discuss in
Section 4.

Figure 7. Symmetric Gauss-Seidel: (a) Python code on adjacency list, and (b)
implementation in SSR.

Composing operators is important in implementing an implicit solver: unlike
the explicit parts which can be written as chains of SpMVs, the matrix in
implicit parts has to be first simplified from the original math formula to
enable further solving. Other types of operators composition include summation
and subtraction between two SSR matrices. Such compositions are simpler than
SpGEMM since they only require merging two row sequences. We will see concrete
usages in Section 5.2.

Figure 8. SSR implementation for ILU.

####  3.3.3. Loop-Carried Dependences: SymGS

Symmetric Gauss-Seidel is a commonly-used routine in iterative solving
methods. The algorithm and implementation in SSR is shown in Fig. 7. The
forward sweep iterates over the sparse matrix similar to in SpMV, but with two
differences: it takes branch on each non-zero checking for diagonal elements,
demonstrating the highly customizable control flow; more importantly, the $x$
Vector is modified in-place, introducing loop-carried dependences. The
backward sweep does the same thing, but in reversed dictionary order of the
domain coordinates. SSR again provides a straightforward programming interface
to the user, with the codes on adjacency lists nearly identical to codes on
SSR matrices, just like in SpMV.

####  3.3.4. A Complex Example: Incomplete LU Solver

Incomplete Lower-Upper factorization (abbreviated as ILU) is a sparse
approximation of LU factorization on matrices. It is based on the insight that
LU factorization might generate lower and upper matrices with much more non-
zeros than the input sparse matrix. ILU instead seeks an approximation that is
accurate on non-zero positions in the input matrix. Thus, ILU is a
decomposition with form $A=LU-R$ in which $R$ is the residual, and by solving
$LUx=b$, ILU solver gives an $x$ that $Ax\approx b$. ILU is useful as a
preconditioner in solving sparse linear equations iteratively. It also
generates accurate results for banded sparse matrices, e.g. tri-/penta-
diagonal matrices.

Implementing with SSR matrices, the code for an ILU solver is shown in Fig. 8.
It is required to materialize the matrices value of the decomposed $L$ and $U$
since we need two passes to calculate the result. With the $max\\_nnz$ in
hand, we can materialize the matrices with ELL format, but without writing
down the indices. Only non-zero values of the input matrix are subjected to
materialization, and we always know the sparsity pattern through the SSR
format. When we need to find common non-zero columns in the two rows in the
forward elimination, we simply apply an $O(n^{2})$ linear scan. It will be
simplified to direct accesses through later optimization, as we will see in
Section 5.1.

Figure 9. The basic SSR matrices, presented with code and the corresponded
matrix: $D$, $A$, $P$, and $S$.

###  3.4. Basic SSR Matrices

Based on the SSR matrix representation, we introduce several basic SSR
matrices in our prelude library. Inspired by OpenArray (Huang et al., 2019),
we include spatial differentiation ($\mathbf{D}$) and averaging ($\mathbf{A}$)
operators in our prelude library to allow simple implementation of complicated
sparse matrices. Besides, we further introduce padding ($\mathbf{P}$) and
slicing ($\mathbf{S}$), to allow simple boundary handling: $\mathbf{P}x$ pads
1-dimensional $x$ with two extra zeros at both ends, and $\mathbf{S}x$
inversely removes the two ends of $x$. These matrices usually exist as stencil
operators in the stencil frameworks and DSLs; as we have shown before, stencil
operators are naturally sparse matrices by themselves. The matrices are shown
in Fig. 9.

##  4\. Dedicated Optimization for SSR

Figure 10. An example of staged tensor program.

Through multi-staged programming, SSR generates staged codes with vast usages
of state machines. These usages contain lengthy state-transition codes,
leading to unacceptable redundant computation and loss of optimization
(especially parallelization) chance. In this section, we analyze the staged
code and discuss how to optimize it into a more regular form which is similar
to hand-written matrix-free code.

$\displaystyle p\quad\Coloneqq$ $\displaystyle\ \quad\texttt{fn}\
(id_{1},id_{2},\dots,id_{n})\ \\{\ s;\ \\}$ (program) $\displaystyle
s\quad\Coloneqq$ $\displaystyle\ \quad s;s$ (statement sequence)
$\displaystyle|\quad\textbf{{vardef}}\ id\
[e_{\textit{size}_{1}},e_{\textit{size}_{2}},\dots,e_{\textit{size}_{n}}]\
\\{\ s;\ \\}$ (define tensor variable) $\displaystyle|\quad\textbf{{for}}\
iter\ \text{in}\ e_{\textit{begin}}:e_{\textit{end}}:e_{\textit{step}}\ \\{\
s;\ \\}$ (for loop) $\displaystyle|\quad\textbf{{while}}\ e\ \\{\ s;\ \\}$
(while loop) $\displaystyle|\quad\textbf{{if}}\ e\ \\{\ s_{\textit{then}};\
\\}\ \textbf{{else}}\ \\{\ s_{\textit{else}};\ \\}$ (if-then-else conditional)
$\displaystyle|\quad id[e_{1},e_{2},\dots,e_{n}]\leftarrow e$ (store)
$\displaystyle|\quad\textit{nop}$ $\displaystyle e\quad\Coloneqq$
$\displaystyle\ \quad e+e\ |\ e*e\ |\ \dots$ (arithmetic operations)
$\displaystyle|\quad id[e_{1},e_{2},\dots,e_{n}]$ (load) $\displaystyle|\quad
iter$ (iteration index) $\displaystyle|\quad literal$ (literal value) Figure
11. The intermediate language for staged code.

To understand what we need to do in optimizing the staged code, we demonstrate
a simple example in Fig. 10. The SSR matrix `d2` is written as a generator
function, and later transformed to the state-transition function next. Then it
is consumed by a lazy SpMV instance, which is materialized into a
1-dimensional tensor spmv, finishing the computation. We observe that the
staged code on the right-hand side is still far from optimal, which is shown
in the top-right corner. But it is possible to optimize the staged code into
the optimal one: constant propagation will do most of it, except the while
loop must be unrolled to the end. More complicated code will require more than
that. Based on the domain knowledge, we introduce the dedicated optimizations
below to optimize the staged code. The rules in this subsection are based on
the language described in Fig. 11.

###  4.1. Unroll While Loops

In programs targeting SSR matrices, while loops occur when iterating over a
row. Considering that the row is of a certain shape determined by its row
index, it is possible to unroll the while loop into a plain body sequence and
thus enable further simplification. One unrolling step takes the following
rule:

$\displaystyle\textbf{{while}}\ e\ \\{\ s;\ \\}\quad\Rightarrow\quad$
$\displaystyle\textbf{{if}}\ e\ \\{\ s;\ \\}$ $\displaystyle\textbf{{while}}\
e\ \\{\ s;\ \\}$

We expect the first if to be expanded by constant propagation. With the while
loop unrolling and constant propagation taken alternately, a while loop
iterating over a state machine generated by SpRow.foreach will be flattened
into a plain code sequence.

$\displaystyle spmv[i]\leftarrow 0.0;$ $\displaystyle state\leftarrow 0;$
$\displaystyle\textbf{{while}}\ state\neq-1\ \\{$
$\displaystyle\quad\textbf{{if}}\ state=0\ \\{\ state\leftarrow
1;nnz\\_idx\leftarrow i;nnz\\_val\leftarrow-1.0;\ \\};$
$\displaystyle\quad\textbf{{if}}\ state=1\ \\{\
state\leftarrow-1;nnz\\_idx\leftarrow i+1;nnz\\_val\leftarrow 1.0;\ \\};$
$\displaystyle\quad spmv[i]\leftarrow spmv[i]+nnz\\_val\times x[nnz\\_idx];$
$\displaystyle\\}$ $\displaystyle\xRightarrow{\text{unroll}\times 2}$
$\displaystyle\xRightarrow{\text{prop}}$ $\displaystyle spmv[i]\leftarrow
0.0;$ $\displaystyle state\leftarrow 0;$ $\displaystyle state\leftarrow
1;nnz\\_idx\leftarrow i;nnz\\_val\leftarrow-1.0;$ $\displaystyle
spmv[i]\leftarrow spmv[i]+nnz\\_val\times x[nnz\\_idx];$ $\displaystyle
state\leftarrow-1;nnz\\_idx\leftarrow i+1;nnz\\_val\leftarrow 1.0;$
$\displaystyle spmv[i]\leftarrow spmv[i]+nnz\\_val\times x[nnz\\_idx];$
$\displaystyle\xRightarrow{\text{general opt.}}$ $\displaystyle
spmv[i]\leftarrow-x[i]+x[i+1];$ Figure 12. An example of while loop unrolling.

We demonstrate an example in Fig. 12. The input code is a part of the staged
code from an SpMV with the $D$ differentiation matrix mentioned in Section
3.4. With the first step of while loop unrolling and constant propagation, one
copy of the loop body is unrolled out of the while loop, and simplified to
straight-line codes. The next unrolling and propagation pair in turn unrolls
out the next step, eliminating the rest while loop as well since it sees
$state=-1$. As we enforce acyclic control flow in the generator function
producing sparse rows, the unrolling is guaranteed to halt since each unroll
at least moves some iterator by one step. With later dead-code elimination,
combining store operations, and other simplification passes, the code finally
gets simplified to as simple as one statement.

$\displaystyle\mathcal{P}_{s}\llbracket(x[e_{\textit{idx}}]\leftarrow
e),E^{\mathcal{P}}\rrbracket\coloneqq$ $\displaystyle\qquad\textbf{{let}}\
e_{\textit{idx}}^{\prime}=\mathcal{P}_{e}\llbracket
e_{\textit{idx}},E^{\mathcal{P}}\rrbracket\ \text{in}$
$\displaystyle\qquad\textbf{{let}}\ e^{\prime}=\mathcal{P}_{e}\llbracket
e,E^{\mathcal{P}}\rrbracket\ \text{in}$
$\displaystyle\qquad(x[e_{\textit{idx}}^{\prime}]\leftarrow e^{\prime}),\\{$
$\displaystyle\qquad\quad\textbf{{if}}\ \text{no {(load)} occurs in
$e^{\prime}$ and $e_{\textit{idx}}^{\prime}$ is lexical integer}\ \\{\
E^{\mathcal{P}}.gen(x[e_{\textit{idx}}^{\prime}]\rightarrow e^{\prime})\ \\}$
$\displaystyle\qquad\quad\textbf{{elif}}\ \text{$e_{\textit{idx}}^{\prime}$ is
lexical integer}\ \\{\ E^{\mathcal{P}}.kill(x[e_{\textit{idx}}^{\prime}])\
\\}$ $\displaystyle\qquad\quad\textbf{{else}}\ \\{\
E^{\mathcal{P}}.kill(x[any])\ \\}$ $\displaystyle\qquad\\}$
$\displaystyle\mathcal{P}_{s}\llbracket(\textbf{{for}}\ i\ \text{in}\
e_{1}:e_{2}:e_{3}\ \\{\ s;\ \\}),E^{\mathcal{P}}\rrbracket\coloneqq$
$\displaystyle\qquad\textbf{{let}}\
(s_{new},E^{\mathcal{P}}_{new})=\mathcal{P}_{s}\llbracket
s,E^{\mathcal{P}}\rrbracket\ \text{in}$ $\displaystyle\qquad\textbf{{let}}\
E^{\mathcal{P}}_{next}=E^{\mathcal{P}}\cap\\{(x[e]\rightarrow e_{x})\in
E^{\mathcal{P}}_{new}\ |\ i\text{ not used in }e_{x}\\}\ \text{in}$
$\displaystyle\qquad\quad\textbf{{if}}\
E^{\mathcal{P}}=E^{\mathcal{P}}_{next}\ \\{$
$\displaystyle\qquad\quad\quad(\textbf{{for}}\ i\ \text{in}\
\mathcal{P}_{e}\llbracket
e_{1},E^{\mathcal{P}}\rrbracket:\mathcal{P}_{e}\llbracket
e_{2},E^{\mathcal{P}}\rrbracket:\mathcal{P}_{e}\llbracket
e_{3},E^{\mathcal{P}}\rrbracket\ \\{\ s_{new};\ \\}),E^{\mathcal{P}}_{next}$
$\displaystyle\qquad\quad\\}\ \textbf{{else}}\ \\{\
\mathcal{P}_{s}\llbracket(\textbf{{for}}\ i\ \text{in}\ e_{1}:e_{2}:e_{3}\
\\{\ s;\ \\}),E^{\mathcal{P}}_{next}\rrbracket\ \\}$
$\displaystyle\mathcal{P}_{e}\llbracket
x[e_{\textit{idx}}],E^{\mathcal{P}}\rrbracket\coloneqq$
$\displaystyle\qquad\textbf{{let}}\
e_{\textit{idx}}^{\prime}=\mathcal{P}_{e}\llbracket
e_{\textit{idx}},E^{\mathcal{P}}\rrbracket\ \text{in}$
$\displaystyle\qquad\textbf{{if}}\ \text{$e_{i}$ are all lexical integers and
$E^{\mathcal{P}}.has(x[e_{\textit{idx}}^{\prime}])$}\ \\{\
E^{\mathcal{P}}.get(x[e_{\textit{idx}}^{\prime}])\ \\}$
$\displaystyle\qquad\textbf{{else}}\ \\{\ x[e_{\textit{idx}}^{\prime}]\ \\}$
Figure 13. Core transformation rules for iteration indices propagation.

###  4.2. Propagate Iteration Indices

The iteration indices created by for statements are extensively used to index
SSR matrices and obtain lazy rows. These lazy rows’ generator state machines
need to store the iteration indices into closure variables. Then the stored
indices are used to compute non-zero indices in the row, which will in turn
again be stored into future variables. As such, we need not only propagate
constants but also iteration indices to enable further simplification.

The immutability of iteration indices in the loop body makes the propagation
really easy. Everything goes as is in constant propagation, except that when
exiting a for loop, the recorded values depending on its iteration index must
be killed. To demonstrate the transforming rules, we notate the recorded
values environment as $E^{\mathcal{P}}$, containing the mapping between tensor
entries and their value, which is an expression consisting of constants and
iteration indices. $\mathcal{P}_{s}$ transforms statement $s$ under
environment $E^{\mathcal{P}}$ into $\mathcal{P}_{s}\llbracket
s,E^{\mathcal{P}}\rrbracket=(s_{1},E^{\mathcal{P}}_{1})$, in which $s_{1}$ is
the transformed statement and $E^{\mathcal{P}}_{1}$ is the updated
environment. $\mathcal{P}_{e}\llbracket e,E^{\mathcal{P}}\rrbracket$ is what
$e$ propagates to within environment $E^{\mathcal{P}}$. Then we list the core
rules (only on one-dimensional tensors for brevity) in Fig. 13, as an
extension to constant propagation; note that we iterate for a fixed point in
propagating for loops, just to be conservative.

###  4.3. Simplify Predicates

In programs using SSR matrices, conditional branches based on iteration
indices are common. For example, in SpGEMM, it requires comparing non-zero
indices of different rows to merge multiple rows. After unrolling and
iteration indices propagation, the comparisons would base on the iteration
indices. Many of these comparisons can be resolved statically, e.g. $i<i+1$ is
always true. Besides, comparisons for boundary check, which is common in
stencil operators, can also be resolved statically according to for loop range
or predicates of outer if statements.

To achieve the above targets, we present the transformation as follows. To
maintain the logical theorems’ context, on entering if branches and for loops,
we add the corresponded information to the environment: for for loops, the
loop range is added; for if branches, $cond=true$ is added to its “then”
case’s environment and $cond=false$ into its “else” case’s environment. Based
on the information, we check the condition of each if statement with an SAT
solver, in our case, Z3 (De Moura and Bjørner, 2008). If either $cond$ or
$!cond$ can be proved, we simplify it into its corresponding case, removing
the condition check. Thus we list the rules in Fig. 14, in which
$\mathcal{S}\llbracket s,E^{\mathcal{S}}\rrbracket$ represents simplified code
of $s$ in environment $E^{\mathcal{S}}$, and all transformation only takes
place when the expressions have no dependency to (the mutable) variables.

$\displaystyle\mathcal{S}\llbracket(\textbf{{for}}\ iter\ \text{in}\
e_{1}:e_{2}:e_{3}\ \\{\ s;\ \\}),E^{\mathcal{S}}\rrbracket\coloneqq$
$\displaystyle\qquad(\textbf{{for}}\ i\ \text{in}\ e_{1}:e_{2}:e_{3}\ \\{$
$\displaystyle\qquad\quad\mathcal{S}\llbracket s,E^{\mathcal{S}}.add($
$\displaystyle\qquad\quad\quad(e_{3}>0\ ?\ (i\geq e_{1}\land i<e_{2}):(i\leq
e_{1}\land i>e_{2})),$ $\displaystyle\qquad\quad\quad((i-e_{1})\ mod\
e_{3}=0)$ $\displaystyle\qquad\quad)\rrbracket;$ $\displaystyle\qquad\\})$
$\displaystyle\mathcal{S}\llbracket(\textbf{{if}}\ e\ \\{\ s_{1};\ \\}\
\textbf{{else}}\ \\{\ s_{2};\ \\}),E^{\mathcal{S}}\rrbracket\coloneqq$
$\displaystyle\qquad\textbf{{if}}\ E^{\mathcal{S}}.unsat(!e)\ \\{\
\mathcal{S}\llbracket s_{1},E^{\mathcal{S}}\rrbracket\ \\}$
$\displaystyle\qquad\textbf{{elif}}\ E^{\mathcal{S}}.unsat(e)\ \\{\
\mathcal{S}\llbracket s_{2},E^{\mathcal{S}}\rrbracket\ \\}$
$\displaystyle\qquad\textbf{{else}}\ \\{\ (\textbf{{if}}\ e\ \\{\
\mathcal{S}\llbracket s_{1},E^{\mathcal{S}}.add(e)\rrbracket;\ \\}\
\textbf{{else}}\ \\{\ \mathcal{S}\llbracket
s_{2},E^{\mathcal{S}}.add(!e)\rrbracket;\ \\})\ \\}$ Figure 14. Core
transformation rules for predicate simplification.

###  4.4. Lifting Branches Depending on Iteration Indices

Apparently, not all branches are decidable before runtime, e.g., stencils with
special boundary conditions and prolongation operations like in Fig. 4 both
require runtime condition check. After taking a branch, the two cases apply
different modifications to the variables, especially the state variable of
SpRow. During propagation, we have to join the two cases, losing information
about the state variable. It then hinders later while loop unrolling and
iteration indices propagation, resulting in a failure of optimizing the state-
machine-based code into a straightforward one.

The transformation is as below.

$\displaystyle\textbf{{for}}\ i\ \text{in}\ e_{1}:e_{2}:e_{3}\ \\{$
$\displaystyle\Rightarrow\qquad\qquad$ $\displaystyle\textbf{{for}}\ i\
\text{in}\ e_{1}:e_{2}:e_{3}\ \\{$ $\displaystyle\quad s_{1};$
$\displaystyle\quad\textbf{{if}}\ e\ \\{\ s_{1};\ s_{2};\ s_{4};\ \\}$
$\displaystyle\quad\textbf{{if}}\ e\ \\{\ s_{2};\ \\}\ \textbf{{else}}\ \\{\
s_{3};\ \\};$ $\displaystyle\quad\textbf{{else}}\ \\{\ s_{1};\ s_{3};\ s_{4};\
\\}$ $\displaystyle\quad s_{4};$ $\displaystyle\\}$ $\displaystyle\\}$

With $e$ undecidable at compiling, we have to join the environment after
propagating on the two branches, losing information about what $s_{2}$ and
$s_{3}$ will do and leading to failure in optimizing $s_{4}$. By lifting the
branch, we can separately propagate on $s_{2}\rightarrow s_{4}$ and
$s_{3}\rightarrow s_{4}$, fixing the problem.

Though this transformation causes duplication of the codes before the
condition check and seems to be sub-optimal, such lifting will not create many
cases to be created in practice. Because many of the conditions are overlapped
or even repeated with each other, programs generate only a handful number of
cases in total. Moreover, lifting to just inside the dimensions loop enables
loop separation, which generates individual regular loops for the interior of
the domain, providing more potential for acceleration.

###  4.5. Passes Fusion

Among the transformations presented above, while loop unrolling and branch
lifting both causes the size of code to increase. Especially, if processed
alone, each while loop will be unrolled for its corresponded $max\\_nnz$
times. Deep nested loops are common in SSR programs, e.g. SpGEMM and ILU.
Optimizing staged code from NPB-BT shown in Section 5.2, it generates over 1
million lines of intermediate code and still does not halt. The optimizations
need to be carefully applied to avoid code explosion and non-halting
compilation.

We observe that the passes can be fused together, applying the transform rules
incrementally and alternately. The transforms only operate on a part of a
program at each move. While unrolling and lifting introduces duplicated code,
the immediately following propagation and simplification will separate the
continuing code. The duplication keeps living shortly, and never bloats up the
code size. Through the whole transformation, we essentially divide the domain
coordinates, which are represented as loop indices in the staged code, into
cases and simplify the body to the end for each case.

##  5\. Application Cases

In this section, we introduce three application cases to demonstrate how to
use SSR matrices, what the generated code is like, and how the generated code
will perform. The three examples will be: numerical integration with solving
differentiation operator, which is a minimal example for ILU; a pseudo
application in the NAS Parallel Benchmarks, NPB-BT; and a matrix-free non-
distributed version of HPCG, the High Performance Conjugate Gradient
Benchmarks.

###  5.1. A Minimal Example for Incomplete LU

Discretized spatial differentiation can be represented as a sparse matrix.
Taking $v=(\delta/\delta x)u$ as an example, it is equivalent to

$\mathbf{v}=\begin{bmatrix}-1/\delta x&1/\delta x\\\ &-1/\delta x&1/\delta
x\\\ &&\ddots&\ddots\\\ &&&-1/\delta x&1/\delta x\\\
\end{bmatrix}\mathbf{u}=\mathbf{D}_{x}\mathbf{u}.$

As such, a numerical integration with base $u_{0}=0$, which calculates
$\mathbf{u}_{1\ldots N}$ by $v$ as input, will follow the equation that

(1) $\mathbf{v}=\begin{bmatrix}1/\delta x\\\ -1/\delta x&1/\delta x\\\
&\ddots&\ddots\\\ &&-1/\delta x&1/\delta x\\\ \end{bmatrix}\mathbf{u}_{1\ldots
N}$

Noticing that the matrix is a banded sparse matrix, an Incomplete LU solver
will yield an accurate result for this equation. In such a case, it will be
equivalent to direct solving, if written in hand. The resulting solver, when
derived by hand, will be

(2) $\displaystyle u_{1}=$ $\displaystyle v_{1}\delta x$ (initialize)
$\displaystyle u_{i}=$ $\displaystyle u_{i-1}+v_{i}\delta x$ (integrate add)

This is exactly one will do for numerical integration: a prefix sum with
constant factor $\delta x$.

We show the corresponded code through solving SSR matrix in Fig. 15(a). The
matrix representation is straightforward as a generator. Then using the ILU
routine ilu_solve shown in Section 3.3.4, the integration linear equations
Equation 1 are solved. The code generated after the dedicated optimizations is
shown in Fig. 15(b). We observe that the generated code does exactly what a
human will do, being identical to Equation 2.

While overkill at this point, the ability for abstracting over different
matrices gives us unprecedented expressiveness, allowing us even to write
solvers in code very similar to the original discretized equation. We will see
an example in the next section.

(a) SSR implementation.

(b) Generated C code.

Figure 15. Implementation of the ILU case and generated code.

###  5.2. A Pseudo-Application in the NAS Parallel Benchmarks: Block Tri-
diagonal

$\displaystyle\begin{aligned}
\mathbf{E},\mathbf{F},\mathbf{G}:&\quad\text{elementwise functions with $(5)$
{inner\\_shape} of }\mathbf{U}\\\
\mathbf{T},\mathbf{V},\mathbf{W}:&\quad\\{D_{\xi},D_{\eta},D_{\zeta}\\}\times\text{elementwise
functions with $(5)$ {inner\\_shape} of }\mathbf{U}\\\
\mathbf{A},\mathbf{B},\mathbf{C}:&\quad\text{elementwise functions with
$(5,5)$ {inner\\_shape} of }\mathbf{U}\\\
\mathbf{N},\mathbf{Q},\mathbf{S}:&\quad\text{elementwise functions with
$(5,5)$ {inner\\_shape} of }\mathbf{U}\end{aligned}$
$\displaystyle\begin{aligned} \mathbf{RHS}\ =&\ [D_{\xi}(T+E)-\epsilon
h_{\xi}^{4}D_{\xi}^{4}\mathbf{U}]&\\\ +&\ [D_{\eta}(V+F)-\epsilon
h_{\eta}^{4}D_{\eta}^{4}\mathbf{U}]&\\\ +&\ [D_{\zeta}(W+G)-\epsilon
h_{\zeta}^{4}D_{\zeta}^{4}\mathbf{U}]&\\\ +&\ Forcing&\textit{(RHS
compute)}\\\ \Delta\mathbf{U}_{1}\ =&\
\\{\mathbf{I}-\Delta\tau[D_{\xi}\mathbf{A}+D_{\xi}^{2}\mathbf{N}]\\}^{-1}(\Delta\tau\
\mathbf{RHS})&\textit{($\xi$ sweep)}\\\ \Delta\mathbf{U}_{2}\ =&\
\\{\mathbf{I}-\Delta\tau[D_{\eta}\mathbf{B}+D_{\eta}^{2}\mathbf{Q}]\\}^{-1}\Delta\mathbf{U}_{1}&\textit{($\eta$
sweep)}\\\ \Delta\mathbf{U}\ =&\
\\{\mathbf{I}-\Delta\tau[D_{\zeta}\mathbf{C}+D_{\zeta}^{2}\mathbf{S}]\\}^{-1}\Delta\mathbf{U}_{2}&\textit{($\zeta$
sweep)}\\\ \end{aligned}$ (a) The discretized and decomposed problems to solve
in NPB-BT, for one time step.

(b) Implementation of NPB-BT in SSR, showing code for one time step.

We use t, x, y, z to denote $\tau,\xi,\eta,\zeta$. int_ stands for
“integrate”.

Figure 16. Demonstration of problems and SSR codes for NPB-BT, the BT pseudo-
application in the NAS Parallel Benchmarks.

The NAS Parallel Benchmarks (Bailey et al., 1994) include three pseudo-
applications solving the same differential equation. We implement the first,
BT (Block Tri-diagonal), as an example.

####  5.2.1. Problem Setup

The equation it solves is derived from Navier-Stokes equations, anonymized due
to confidential real application codes. Discretized spatially and temporally,
the equations required to solve are listed in Fig. 16(a). It first defines 12
functions on the input field $\mathbf{U}$, in which 3 of them
($\mathbf{T},\mathbf{V},\mathbf{W}$) are spatial differentiation of
elementwise functions while all others are elementwise functions of
$\mathbf{U}$. Then it computes the residual, $\mathbf{RHS}$, through
complicated computation including spatial correlations. The spatial
correlations are expressed with the $D_{*}$ operators, which can be
interpreted as sparse matrices as is mentioned in Section 3.4. Then it goes
through 3 different linear systems solving to compute the final
$\Delta\mathbf{U}$. The linear systems solved here are block tri-diagonal
matrices when reordered by the corresponding axis as the inner-most one,
giving the pseudo-application its name.

####  5.2.2. Implementation

To solve the problem, codes in SSR is presented in Fig. 16(b). Elementwise
functions are implemented as vfuncs straightforwardly. The vfunc transforms a
elementwise function to a vector-wise function. Further, the
$\mathbf{T},\mathbf{V},\mathbf{W}$ are equivalent to
$D_{\xi},D_{\eta},D_{\zeta}$ applied to an elementwise functions. Thus we
instead implement int_t with vfunc, and similarly int_v and int_w. Then we
have $\mathbf{T}=D_{\xi}\times\texttt{int\\_t}$, etc.

Upon these, $\mathbf{RHS}$ is computed through a linear algebra expression.
The prelude matrices, like Px padding along $\xi$ axis and Ay averaging two
contiguous grid points along $\eta$ axis, help implementing the $\mathbf{RHS}$
computation easily. Notice that we keep the number of applied Ax and Dx
consistent, to ensure that added vectors are well aligned spatially. We
implement 4th-order dissipation transformation $\epsilon h^{4}D^{4}$ specified
by NPB as a SpMat instance operating on a 1-D domain. It is then broadcasted
into $\epsilon h_{\xi}^{4}D_{\xi}^{4},\epsilon
h_{\eta}^{4}D_{\eta}^{4},\epsilon h_{\zeta}^{4}D_{\zeta}^{4}$ matrices, which
operates on a 3D domain. The results, dsspx, dsspy, and dsspz, are used in the
$\mathbf{RHS}$ computation.

Based on the computed residual, we solve the three linear problems, also known
as $\xi$, $\eta$ and $\zeta$ sweeps, with three ilu_solve calls. The matrices
to be solved, which is composed of identity matrix $\mathbf{I}$ and spatial
correlations like $D_{\xi}$, are banded sparse matrices, and thus ILU produces
an accurate result in solving it. The final computed $\Delta\mathbf{U}$,
written as du3, is returned, finishing the iteration step.

With about 200 lines of code following the math definitions of
$h^{4}D^{4},\mathbf{E},\mathbf{F}$, etc., and 29 lines describing the
computation that are also nearly identical to the math equations, we implement
the solver which requires 4200 lines of code in the original hand-written
implementation by the NAS. Another twenties lines of code for human-directed
optimization schedules on the generated tensor program will boost the
performance, which will be presented in Section 6. Together with ILU routine,
initialization, and driver code, we implement the complete NPB-BT end-to-end
with about 400 lines of code, $<10\%$ compared with the original FORTRAN
implementation.

$\displaystyle\textbf{MG}_{i}:r\rightarrow x$
$\displaystyle\quad\begin{aligned} &x\coloneq 0\text{ in the same shape of
}r\\\ &\textbf{if }i>0\\\ &\quad\textbf{SymGS}(A,r,x)\\\ &\quad
xc\leftarrow\textbf{MG}_{i-1}(\textbf{Restrict}(r-Ax))\\\ &\quad x\leftarrow
x+\textbf{Prolong}(xc)\\\ &\quad\textbf{SymGS}(A,r,x)\\\ &\textbf{else}\\\
&\quad\textbf{SymGS}(A,r,x)\end{aligned}$ $\displaystyle\textbf{CG-
Iter}:(x,r,p,rtz)_{t-1}\rightarrow(x,r,p,rtz)_{t}$
$\displaystyle\quad\begin{aligned} &z_{t}\leftarrow\textbf{MG}_{4}(r_{t-1})\\\
&{rtz}_{t}\leftarrow\textbf{dot}(r_{t-1},z_{t})\\\
&p_{t}\leftarrow({rtz}_{t}/{rtz}_{t-1})p_{t-1}+z_{t}\\\
&\alpha_{t}\leftarrow{rtz}_{t}/\textbf{dot}(p_{t},Ap_{t})\\\ &x_{t}\leftarrow
x_{t-1}+\alpha_{t}p_{t}\\\ &r_{t}\leftarrow
r_{t-1}+\alpha_{t}Ap_{t}\end{aligned}$ (a) The CG iteration with MG
preconditioner.

(b) Implementation through SSR matrices.

Figure 17. The backbone of the CG algorithm in HPCG, with its implementation
on SSR.

###  5.3. Matrix-free High Performance Conjugate Gradient

We then demonstrate how we achieve the same computation with the HPCG
benchmark (Dongarra et al., 2016) in matrix-free, with SSR. The HPCG benchmark
forbids matrix-free approaches, so our implementation should not be regarded
as a valid HPCG test; instead, it is a demonstration of how easy the
programming and how performant the implementation would be with our SSR.

####  5.3.1. Problem Setup

The HPCG benchmark, as is specified in (Heroux et al., 2013), involves solving
a linear system defined by a 27-point stencil. It deploys an iterative solving
method with a multi-grid (MG) preconditioner, in which symmetric Gauss-Seidel
(SymGS) presmoother is used. The algorithm of SymGS has been mentioned in
Section 3.3.3. Based on the SymGS iterations, the MG preconditioner
recursively restricts a grid to a coarser one, iterates the coarser grid, and
prolongs it back. The procedure is called a “V-Cycle” and HPCG requires a
recursion depth of 4. The MG preconditioner is used in each conjugate gradient
(CG) iteration. The algorithm is shown in Fig. 17(a).

####  5.3.2. Implementation

Based on the SymGS iteration mentioned in Section 3.3.3, we implement mg, the
preconditioner, and then cg_iter, the complete conjugate gradient iteration,
as is shown in Fig. 17(b). The dot operation is implemented via FreeTensor
tensor operators, since both operands (Vectors) are dense tensors. The
matrices for the 27-point stencil with boundary condition, the restriction,
and the prolongation are straightforward generator functions, omitted for
space.

##  6\. Performance Evaluation

In this section, we evaluate the performance of our implementation of the two
applications, NPB-BT (Bailey et al., 1994) and matrix-free HPCG (Dongarra et
al., 2016). The simple ILU case is omitted as ILU is covered extensively in
NPB-BT.

We setup our experiments on a dual-socket server, each socket being one AMD
EPYC 7742 CPU, with 64 physical cores and 128 logical cores per CPU. The
server has 576GB of DDR4-2933 main memory. The server runs Debian 11.1 with
Linux kernel 5.10. We use GCC 10.2.1 to compile the programs, using g++ for
C++ and gfortran for FORTRAN, including our generated codes. All parallelism
is through OpenMP, including reference codes and our implementation.

We use NPB 3.4.2 and HPCG 3.1 as the baselines accordingly. Our SSR
implementation takes 1181 lines of Python code implementing the frontend and
433 lines of C++ code as a FreeTensor plugin implementing the dedicated
optimizations. Implementing NPB-BT on SSR matrices, including the ILU routine
and performance-optimizing schedules, takes 411 lines of Python code. The
matrix-free HPCG, including the SymGS routine, takes 209 lines of Python code.

(a) Performance of different NPB-BT implementations, time in seconds.

Ref is the reference implementation; SSR Naive is our implementation with
optimizing schedules turned off; and SSR Opt is with optimizing schedules on.

(b) Performance of different HPCG implementation, time in seconds.

Ref is the matrix-based reference implementation; SSR is our matrix-free
implementation.

Figure 18. Performance results of SSR implemented applications.

###  6.1. The Performance on NPB-BT

We run the Class D problem size specified by the NAS Parallel Benchmark, a
$408\times 408\times 408$ grid with $\Delta\tau=0.00002$ and runs for $250$
time steps. The performance of the NAS-provided OpenMP-based implementation,
our basic SSR implementation without human-directed optimizing schedules, and
SSR implementation with those schedules are listed in Fig. 18(a). The
reference implementation of the NPB benchmark is a well-optimized human-
written one, with 4200 lines of FORTRAN code, not taking comments and blank
lines into account. Instead, our basic implementation is close to the math
formulas and takes less than 400 lines of code. It achieves $32.7\%$
performance compared with the reference.

Inspecting the generated code, we found a critical performance problem being
that the sweeping axis is not always the inner-most one, resulting in the
entire $A,B,C,N,Q,S$ 3D arrays to be materialized, which is enormous in memory
consumption and hurts cache performance. With just 26 lines of schedules in
FreeTensor, we can rotate the axes and make the axis with dependence the
inner-most one, allowing the arrays to be only materialized on that axis. This
optimization can also be observed in the reference implementation. After the
scheduling, optimized SSR implementation of NPB-BT achieves $92.8\%$
performance compared to the reference, closing the performance gap.

###  6.2. The Performance of Matrix-free HPCG

For HPCG, we run on a $256\times 256\times 256$ grid and compute the GFLOP/s
(giga- floating-point operations per second) according to the HPCG official
rules. We keep the parallel algorithm of our SSR-based implementation and the
reference HPCG codebase aligned to demonstrate the performance advantage of
matrix-free methods. The reference code and the SSR matrix-free version both
perform SymGS sequentially, strictly following the data dependence. Other
parts are parallelized with multi-threading. Further, we tweak the HPCG code
to output its accuracy requirement (“tolerance” as written in its code) and
iterate our implementation until it reaches the same accuracy, aligning with
the HPCG specification. The results are shown in Fig. 18(b).

It can be seen that the performance of our naive matrix-free implementation
can already outperform the reference with a $3.29\times$ speedup. Recall that
the implementation on SSR is at the same abstraction level as the original
HPCG; the speedup shows the power of the SSR design. Thanks to the brief
syntax of our embedded DSL, the code is even shorter than the HPCG reference
implementation (553 lines in *_ref.cpp as of HPCG v3.1).

##  7\. Related Work

Historically optimization of stencil computation has been well studied (Habich
et al., 2009; Nguyen et al., 2010; Williams et al., 2007; Kamil et al., 2006).
Based on the optimization techniques, automatic optimization techniques
(Krishnamoorthy et al., 2007; Datta et al., 2008) have been proposed, but
working on the nested loops without reducing the programming burden.

Addressing the issue, researchers have come up with many stencil libraries,
and DSLs, including (Huang et al., 2019; Tang et al., 2011; Louboutin et al.,
2019; Lengauer et al., 2020; Pieper et al., 2021; Zhang et al., 2017),
providing easier interfaces for programming stencils. Among those, Devito
(Louboutin et al., 2019) and ExaStencil (Lengauer et al., 2020) presents a set
of program representations that is pretty much close to the math formula of
the differential equation, being the easiest ones in programming. However,
none of these work enables the programming of complicated implicit solving
algorithms, including direct solving of banded matrices, Gauss-Seidel
iterating, Incomplete Lower-Upper (Saad, 2003), Fine-Grained Parallel
Incomplete Lower-Upper (Chow and Patel, 2015), etc. The great diversity of
such algorithms makes it hard for DSLs to cover all of them.

Contrary to those performance-centric work, numerically solving differential
equations has also been widely treated as sparse linear algebra computation.
Sparse matrix-vector multiplication (SpMV) (Pinar and Heath, 1999; Williams et
al., 2010; Merrill and Garland, 2016) and sparse matrix-sparse matrix
multiplication (SpGEMM) (Yuster and Zwick, 2005; Hong et al., 2019) are two
well studied scopes, naturally covering stencil computation and operator
composition. Other routines are also available in sparse linear algebra, e.g.
SuperLU contains a general-purpose efficient ILU implementation (Li and Shao,
2010). SciPy (Virtanen et al., 2020) is a commonly used Python library for
scientific computing and covers many such routines. While those approaches
decouple the equation and solvers, they are too general to be highly optimized
against structured grids. Borrowing their idea of representing differential
equation solving in sparse linear algebra, we propose our SSR matrix
representation, enabling easy algorithm programming while persisting the
performance benefits of matrix-free codes.

Optimization methods based on general program analysis (Allen and Kennedy,
2001; Padua and Wolfe, 1986) automatically apply to both explicit and implicit
solving routines once presented as regular nested loops, but the optimization
is not. In the contrary, recent emerging tensor programming frameworks (Paszke
et al., 2017; Abadi et al., 2015; Frostig et al., 2018; Chen et al., 2018)
designed mainly for deep learning widely adopts a coarse-grained programming
interface, dealing with vectorized operators at high level. Such an interface
can hardly implement stencils without redundant work or merely writing
implicit routines. Thus we use FreeTensor (Tang et al., 2022) as our backend
for its capability of inferring loop-carried dependences from serial nested
loops, flexible scheduling primitives for code transformation,
parallelization, and other optimization. The dedicated optimizations are
implemented as a FreeTensor plugin.

Our embedded DSL implementation is inspired by Lightweight Modular Staging
(LMS) (Rompf and Odersky, 2012), which represents multi-staged programming
through types in Scala.

##  8\. Discussion

Our core design, SSR, is inspired by the spirit of abstract data type (ADT) in
object-oriented programming. Nevertheless, simply applying ADT to sparse
matrices yields suboptimal performance: without separating the two stages, the
result matrix-free program cannot be efficient due to the complex control flow
it will introduce. The staging design and its dedicated optimizations are
critical to building an efficient matrix-free solver.

While this paper primarily targets the representation of solving differential
equations on structured grids, the generated code from SSR programs still has
large performance potential. More automatic optimizations are possible by
applying more schedules but are out of this paper’s scope and will require
future work. Automatic distributed code generation is also possible through
polyhedral methods analyzing loop-carried dependences, which has been explored
previously (Baghdadi et al., 2019).

Though our current abstraction focuses on Cartesian grids, other types, e.g.,
star-shaped or hexagonal ones, can be supported with SSR as long as they can
be expressed through acyclic control flows in the generator function. However,
the basic matrices choice and definition should be reconsidered with such
grids. Speaking even further, semi-structured grids such as block-structured
grids might also benefit from our method by combining the unstructured matrix-
based representation at the coarse level and our SSR representation for the
finer grids, requiring more programming optimization designs.

##  9\. Conclusion

This paper proposes Staged Sparse Row, a novel representation for sparse
matrices in programs instead of data, covering the differential equation
solving on structured grids. By lifting the matrices by one stage in multi-
stage programming and applying a set of dedicated optimizations, we manage to
do efficient matrix-free solving in the generated code with a user-friendly
matrix-based interface. Evaluated on NPB-BT, we achieve up to $92.8\%$
performance compared with the hand-rolled matrix-free reference implementation
within $10\%$ lines of code; on the problem specified by HPCG, a
straightforward SSR implementation achieves $3.29\times$ performance than the
reference code, thanks to the matrix-free feature.

###### Acknowledgements.

This work was partially supported by National Key Research & Development Plan
of China under grant 2017YFA0604500 and NSFC U20B2044. The corresponding
author of this work is Wenguang Chen.

## 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.
  * Allen and Kennedy (2001) Randy Allen and Ken Kennedy. 2001.  _Optimizing Compilers for Modern Architectures: A Dependence-based Approach_.  Morgan Kaufmann.
  * Baghdadi et al. (2019) Riyadh Baghdadi, Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yunming Zhang, Patricia Suriana, Shoaib Kamil, and Saman P. Amarasinghe. 2019.  Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code. In _IEEE/ACM International Symposium on Code Generation and Optimization, CGO 2019, Washington, DC, USA, February 16-20, 2019_ , Mahmut Taylan Kandemir, Alexandra Jimborean, and Tipp Moseley (Eds.). IEEE, 193–205.  https://doi.org/10.1109/CGO.2019.8661197
  * Bailey et al. (1994) D. H. Bailey, E. Barszcz, J. T. Barton, D. S. Browning, R. L. Carter, L. Dagum, R. A. Fatoohi, P. O. Frederickson, T. A. Lasinski, R. S. Schreiber, H. D. Simon, and V. Venkatakrishnan. 1994\.  The NAS Parallel Benchmarks.
  * Chen et al. (2018) Tianqi Chen, Thierry Moreau, Ziheng Jiang, Haichen Shen, Eddie Q. Yan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018.  TVM: End-to-End Optimization Stack for Deep Learning.  _CoRR_ abs/1802.04799 (2018).  arXiv:1802.04799 http://arxiv.org/abs/1802.04799
  * Chow and Patel (2015) Edmond Chow and Aftab Patel. 2015.  Fine-Grained Parallel Incomplete LU Factorization.  _SIAM Journal on Scientific Computing_ 37, 2 (2015), C169–C193.  https://doi.org/10.1137/140968896 arXiv:https://doi.org/10.1137/140968896
  * Datta et al. (2008) K Datta, M Murphy, V Volkov, S Williams, J Carter, L Oliker, D Patterson, J Shalf, and K Yelick. 2008\.  Stencil computation optimization and auto-tuning on state-of-the-art multicore architectures. In _In Proc, of the ACM/IEEE Conference on Supercomputing_. IEEE Press.
  * De Moura and Bjørner (2008) Leonardo De Moura and Nikolaj Bjørner. 2008.  Z3: An Efficient SMT Solver. In _Proceedings of the Theory and Practice of Software, 14th International Conference on Tools and Algorithms for the Construction and Analysis of Systems_ (Budapest, Hungary) _(TACAS’08/ETAPS’08)_. Springer-Verlag, Berlin, Heidelberg, 337–340.
  * Dongarra et al. (2016) Jack Dongarra, Michael A. Heroux, and Piotr Luszczek. 2016\.  A new metric for ranking high-performance computing systems.  _National Science Review_ 3, 1 (01 2016), 30–35.  https://doi.org/10.1093/nsr/nwv084 arXiv:https://academic.oup.com/nsr/article-pdf/3/1/30/31565532/nwv084.pdf
  * Frostig et al. (2018) Roy Frostig, Matthew James Johnson, and Chris Leary. 2018\.  Compiling machine learning programs via high-level tracing.  _Systems for Machine Learning_ (2018).
  * Habich et al. (2009) Johannes Habich, T. Zeiser, Georg Hager, and Gerhard Wellein. 2009.  Enabling temporal blocking for a lattice Boltzmann flow solver through multicore-aware wavefront parallelization.  (01 2009).
  * Heroux et al. (2013) Michael Allen Heroux, Jack Dongarra, and Piotr Luszczek. 2013\.  HPCG Benchmark Technical Specification.  (10 2013).  https://doi.org/10.2172/1113870
  * Hong et al. (2019) Changwan Hong, Aravind Sukumaran-Rajam, Israt Nisa, Kunal Singh, and P. Sadayappan. 2019\.  Adaptive Sparse Tiling for Sparse Matrix Multiplication. In _Proceedings of the 24th Symposium on Principles and Practice of Parallel Programming_ (Washington, District of Columbia) _(PPoPP ’19)_. Association for Computing Machinery, New York, NY, USA, 300–314.  https://doi.org/10.1145/3293883.3295712
  * Huang et al. (2019) X. Huang, X. Huang, D. Wang, Q. Wu, Y. Li, S. Zhang, Y. Chen, M. Wang, Y. Gao, Q. Tang, Y. Chen, Z. Fang, Z. Song, and G. Yang. 2019.  OpenArray v1.0: a simple operator library for the decoupling of ocean modeling and parallel computing.  _Geoscientific Model Development_ 12, 11 (2019), 4729–4749.  https://doi.org/10.5194/gmd-12-4729-2019
  * Kamil et al. (2006) S Kamil, K Datta, S Williams, L Oliker, J Shalf, and K Yelick. 2006.  Implicit and explicit optimizations for stencil computations.. In _In Workshop on Memory System Performance and Correctness,_. ACM,, 51–60.
  * Krishnamoorthy et al. (2007) S Krishnamoorthy, M Baskaran, U Bondhugula, J Ramanujam, A Rountev, and P Sadayappan. 2007\.  Effective automatic parallelization of stencil computations. In _In Proc. of the 2007 ACM SIGPLAN Conference on Programming Language Design and Implementation_. ACM, 235–244.
  * Lengauer et al. (2020) Christian Lengauer, Sven Apel, Matthias Bolten, Shigeru Chiba, Ulrich Rüde, Jürgen Teich, Armin Größlinger, Frank Hannig, Harald Köstler, Lisa Claus, Alexander Grebhahn, Stefan Groth, Stefan Kronawitter, Sebastian Kuckuk, Hannah Rittich, Christian Schmitt, and Jonas Schmitt. 2020.  _ExaStencils: Advanced Multigrid Solver Generation_.  405–452.  https://doi.org/10.1007/978-3-030-47956-5_14
  * Li and Shao (2010) Xiaoye S. Li and Meiyue Shao. 2010.  A Supernodal approach to incomplete LU factorization with partial pivoting.  _ACM Trans. Mathematical Software_ 37, 4 (2010).
  * Louboutin et al. (2019) M. Louboutin, M. Lange, F. Luporini, N. Kukreja, P. A. Witte, F. J. Herrmann, P. Velesko, and G. J. Gorman. 2019\.  Devito (v3.1.0): an embedded domain-specific language for finite differences and geophysical exploration.  _Geoscientific Model Development_ 12, 3 (2019), 1165–1187.  https://doi.org/10.5194/gmd-12-1165-2019
  * Mcmechan (2006) George Mcmechan. 2006\.  Migration by extrapolation of time-dependent boundary values.  _Geophysical Prospecting_ 31 (04 2006), 413 – 420.  https://doi.org/10.1111/j.1365-2478.1983.tb01060.x
  * Merrill and Garland (2016) Duane Merrill and Michael Garland. 2016.  Merge-Based Parallel Sparse Matrix-Vector Multiplication. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ (Salt Lake City, Utah) _(SC ’16)_. IEEE Press, Article 58, 12 pages.
  * Nguyen et al. (2010) Anthony Nguyen, Nadathur Satish, Jatin Chhugani, Changkyu Kim, and Pradeep Dubey. 2010.  3.5-D Blocking Optimization for Stencil Computations on Modern CPUs and GPUs. In _SC ’10: Proceedings of the 2010 ACM/IEEE International Conference for High Performance Computing, Networking, Storage and Analysis_. 1–13.  https://doi.org/10.1109/SC.2010.2
  * Ofenbeck et al. (2013) Georg Ofenbeck, Tiark Rompf, Alen Stojanov, Martin Odersky, and Markus Püschel. 2013.  Spiral in Scala: Towards the Systematic Construction of Generators for Performance Libraries.  _SIGPLAN Not._ 49, 3 (oct 2013), 125–134.  https://doi.org/10.1145/2637365.2517228
  * Olsson and Johnsson (1989) Pelle Olsson and Lennart Johnsson. 1989.  Boundary modifications of the dissipation operators for the three-dimensional Euler equations.  _Journal of Scientific Computing_ 4 (06 1989), 159–195.  https://doi.org/10.1007/BF01061500
  * Padua and Wolfe (1986) David A. Padua and Michael Wolfe. 1986.  Advanced Compiler Optimizations for Supercomputers.  _Commun. ACM_ 29, 12 (1986), 1184–1201.  https://doi.org/10.1145/7902.7904
  * Paszke et al. (2017) Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. 2017\.  Automatic differentiation in PyTorch.  (2017).
  * Pieper et al. (2021) Andreas Pieper, Georg Hager, and Holger Fehske. 2021\.  A domain-specific language and matrix-free stencil code for investigating electronic properties of Dirac and topological materials.  _The International Journal of High Performance Computing Applications_ 35, 1 (2021), 60–77.  https://doi.org/10.1177/1094342020959423 arXiv:https://doi.org/10.1177/1094342020959423
  * Pinar and Heath (1999) Ali Pinar and Michael T. Heath. 1999.  Improving Performance of Sparse Matrix-Vector Multiplication. In _Proceedings of the 1999 ACM/IEEE Conference on Supercomputing_ (Portland, Oregon, USA) _(SC ’99)_. Association for Computing Machinery, New York, NY, USA, 30–es.  https://doi.org/10.1145/331532.331562
  * Rompf and Odersky (2012) Tiark Rompf and Martin Odersky. 2012.  Lightweight Modular Staging: A Pragmatic Approach to Runtime Code Generation and Compiled DSLs.  _Commun. ACM_ 55, 6 (jun 2012), 121–130.  https://doi.org/10.1145/2184319.2184345
  * Saad (2003) Yousef Saad. 2003\.  Iterative Methods for Sparse Linear Systems, 2nd edition.
  * Shaikhha et al. (2018) Amir Shaikhha, Yannis Klonatos, and Christoph Koch. 2018\.  Building Efficient Query Engines in a High-Level Language.  _ACM Trans. Database Syst._ 43, 1, Article 4 (apr 2018), 45 pages.  https://doi.org/10.1145/3183653
  * Skamarock et al. (2019) C. Skamarock, BogumiŁa Klemp, Jimy Dudhia, O. Gill, Zhiquan Liu, Judith Berner, Wei Wang, G. Powers, Greg Duda, Dale M. Barker, and Xiangyu Huang. 2019.  A Description of the Advanced Research WRF Model Version 4.
  * Sujeeth et al. (2014) Arvind K. Sujeeth, Kevin J. Brown, Hyoukjoong Lee, Tiark Rompf, Hassan Chafi, Martin Odersky, and Kunle Olukotun. 2014\.  Delite: A Compiler Architecture for Performance-Oriented Embedded Domain-Specific Languages.  _ACM Trans. Embed. Comput. Syst._ 13, 4s, Article 134 (apr 2014), 25 pages.  https://doi.org/10.1145/2584665
  * Tang et al. (2022) Shizhi Tang, Jidong Zhai, Haojie Wang, Lin Jiang, Liyan Zheng, Zhenhao Yuan, and Chen Zhang. 2022.  FreeTensor: A Free-Form DSL with Holistic Optimizations for Irregular Tensor Programs. In _Proceedings of the 43rd ACM SIGPLAN International Conference on Programming Language Design and Implementation (PLDI ’22)_ (San Diego, CA, USA) _(PLDI ’22)_. New York, NY, USA, 16 pages.  https://doi.org/10.1145/3519939.3523448
  * Tang et al. (2011) Yuan Tang, Rezaul Chowdhury, Bradley C. Kuszmaul, Chi keung Luk, and Charles E. Leiserson. 2011\.  The Pochoir stencil compiler. In _In SPAA_.
  * Trottenberg and Schuller (2000) Ulrich Trottenberg and Anton Schuller. 2000.  _Multigrid_.  Academic Press, Inc., USA.
  * 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, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, İlhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. 2020.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ 17 (2020), 261–272.  https://doi.org/10.1038/s41592-019-0686-2
  * Williams et al. (2007) Samuel Williams, John Shalf, Leonid Oliker, Shoaib Kamil, Parry Husbands, and Katherine Yelick. 2007\.  Scientific Computing Kernels on the Cell Processor.  _International Journal of Parallel Programming_ 35 (06 2007).  https://doi.org/10.1007/s10766-007-0034-5
  * Williams et al. (2010) Samuel W. Williams, Nathan Bell, Jee Whan Choi, Michael Garland, Leonid Oliker, and Richard Vuduc. 2010\.  Sparse Matrix-Vector Multiplication on Multicore and Accelerators.  20102756 (12 2010).  https://doi.org/10.1201/b10376
  * Yang et al. (2016) Chao Yang, Wei Xue, Haohuan Fu, Hongtao You, Xinliang Wang, Yulong Ao, Fangfang Liu, Lin Gan, Ping Xu, Lanning Wang, Guangwen Yang, and Weimin Zheng. 2016\.  10M-Core Scalable Fully-Implicit Solver for Nonhydrostatic Atmospheric Dynamics.  https://doi.org/10.1109/SC.2016.5
  * Yuster and Zwick (2005) Raphael Yuster and Uri Zwick. 2005.  Fast Sparse Matrix Multiplication.  _ACM Trans. Algorithms_ 1, 1 (jul 2005), 2–13.  https://doi.org/10.1145/1077464.1077466
  * Zhang et al. (2017) Nathan Zhang, Michael Driscoll, Charles Markley, Samuel Williams, Protonu Basu, and Armando Fox. 2017.  Snowflake: A Lightweight Portable Stencil DSL.  _Proceedings - 2017 IEEE 31st International Parallel and Distributed Processing Symposium Workshops, IPDPSW 2017_ (5 2017).  https://doi.org/10.1109/IPDPSW.2017.89
  * Zhu et al. (2021) Qianchao Zhu, Hao Luo, Chao Yang, Mingshuo Ding, Wanwang Yin, and Xinhui Yuan. 2021\.  Enabling and Scaling the HPCG Benchmark on the Newest Generation Sunway Supercomputer with 42 Million Heterogeneous Cores. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ (St. Louis, Missouri) _(SC ’21)_. Association for Computing Machinery, New York, NY, USA, Article 57, 13 pages.  https://doi.org/10.1145/3458817.3476158