MinMax Networks

Winfried Lohmiller, Philipp Gassert, and Jean-Jacques Slotine

Nonlinear Systems Laboratory

Massachusetts Institute of Technology

Cambridge, Massachusetts, 02139, USA

<EMAIL_ADDRESS>
###### Abstract

While much progress has been achieved over the last decades in neuro-inspired
machine learning, there are still fundamental theoretical problems in
gradient-based learning using combinations of neurons. These problems, such as
saddle points and suboptimal plateaus of the cost function, can lead in theory
and practice to failures of learning. In addition, the discrete step size
selection of the gradient is problematic since too large steps can lead to
instability and too small steps slow down the learning.

This paper describes an alternative discrete MinMax learning approach for
continuous piece-wise linear functions. Global exponential convergence of the
algorithm is established using Contraction Theory with Inequality Constraints
[6], which is extended from the continuous to the discrete case in this paper:

  * •

The parametrization of each linear function piece is, in contrast to deep
learning, linear in the proposed MinMax network. This allows a linear
regression stability proof as long as measurements do not transit from one
linear region to its neighbouring linear region.

  * •

The step size of the discrete gradient descent is Lagrangian limited
orthogonal to the edge of two neighbouring linear functions. It will be shown
that this Lagrangian step limitation does not decrease the convergence of the
unconstrained system dynamics in contrast to a step size limitation in the
direction of the gradient.

We show that the convergence rate of a constrained piece-wise linear function
learning is equivalent to the exponential convergence rates of the individual
local linear regions.

##  1 Introduction

In this paper, we revisit standard convergence difficulties of gradient
descent on a quadratic error cost, such as the possible presence of saddle
points, sub-optimal plateaus, non-Lipschitz edges, time-varying measurements
and the time discretization of the gradient. Initial results of the MinMax
learning we discuss were derived in [7][8].

The classical Rectified Linear Unit (ReLU) approach as e.g. in [4] implies a
piece-wise linear approximation function. The edges of the linear surfaces are
not Lipschitz continuous. Hence it is difficult to prove stability and even
uniqueness of the solution.

The stability problem of the non-Lipschitz edges is overcome in this paper
with a Lagrangian constraint step (2) to the edge. Since the edge belongs to
both Lipschitz continuous regions it is possible at the next iteration to
transit to the neighbouring region with a Lipschitz continuous step. Hence as
a first step Contraction Theory of Constrained Continuous Systems [6] is
generalized in section 2 to discrete dynamics of the form

${\bf x}^{i+1}={\bf f}({\bf x}^{i},i)$ (1)

with $n$-dimensional discrete state $x^{i}$, time index $i$ and
$l=1,...,L$-dimensional linear inequality constraints

$g_{l}={\bf g}_{l}^{T}(i+1){\bf x}^{i+1}+h_{l}(i+1)\leq 0.$ (2)

Note that many non-linear constraints can be brought into a linear form with a
coordinate transformation, such that the results of this paper apply again.

Another problem in the stability analysis of deep learning is that e.g. a
network of depth $100$ actually multiplies $100$ linear parameters with each
other. Hence the smooth surface pieces are piece-wise polynomials of order
$100$ w.r.t. the chosen parametrization, although the approximation function
is actually piece-wise linear. To overcome this problem this paper suggests to
use the sum of several piece-wise linear functions

$\hat{y}({\bf x})=\sum_{j=1}^{J_{\min}}\hat{y}_{j\min}({\bf
x})+\sum_{j=J_{\min}}^{J_{\max}}\hat{y}_{j\max}({\bf x})$ (3)

which uses the convex and concave neurons $j$

$\displaystyle\hat{y}_{j\min}({\bf x})$ $\displaystyle=$
$\displaystyle\min(\hat{z}_{j1},...,\hat{z}_{jK_{j}})$
$\displaystyle\hat{y}_{j\max}({\bf x})$ $\displaystyle=$
$\displaystyle\max(\hat{z}_{j1},...,\hat{z}_{jK_{j}})$

which consist of $k=1,...,K_{j}$ linear basic neurons $\hat{z}_{jk}={\bf
x}^{T}\hat{\bf w}_{jk}$ of estimated $N+1$-dimensional weight vector $\hat{\bf
w}_{jk}$.

The concave $\max$ and convex $\min$ functions are a direct generalization of
the ReLU to achieve piece-wise linear functions. Multiple local convex and
concave functions can be approximated with multiple $\min$ and $\max$
operators. The key advantage to deep networks is that the parametrization is
still linear in $\hat{\bf w}_{jk}$ between the edges, i.e. linear stability
proofs can be used with the mentioned step size limitation to the edge.

The following example shows the main modeling difference of both approaches:

  * ###### Example 1.1

:  Let us consider the unit pyramid in subfigure 1a

$\displaystyle y({\bf x})$ $\displaystyle=$
$\displaystyle\max(0,\min(x_{1}+1,x_{2}+1,-x_{1}+1,-x_{2}+1)).$

In a deep ReLU network each ReLU adds an edge to the piece-wise linear
function. A ReLU network estimation of depth 2 exactly models the pyramid
above with

$\displaystyle\hat{y}({\bf x})$ $\displaystyle=$
$\displaystyle\max(0,\hat{z}_{11})$ $\displaystyle\hat{z}_{11}$
$\displaystyle=$ $\displaystyle
1-0.5\left(\hat{y}_{21}+\hat{y}_{22}+\hat{y}_{23}+\hat{y}_{24}\right))$
$\displaystyle\hat{y}_{2j}$ $\displaystyle=$
$\displaystyle\max(0,\hat{z}_{2j})$ $\displaystyle\hat{z}_{21}$
$\displaystyle=$ $\displaystyle-x_{1}-x_{2},\ \hat{z}_{22}=x_{1}-x_{2},\
\hat{z}_{23}=x_{1}+x_{2},\ \hat{z}_{24}=-x_{1}+x_{2}$

The lower ReLUs $\hat{y}_{2j}=\max(0,\hat{z}_{2j})$ define the the 4 edges of
the pyramid without ground leading to the linear input $\hat{z}_{11}$ to the
last layer in subfigure 1b. The upper ReLU $\hat{y}=\max(0,\hat{z}_{11})$ then
adds the remaining 4 edges to the ground. In general a high over-
parametrization is needed to approximate a piece-wise linear functions with
deep ReLUs [4]. Also it is not very easy to understand which edge is defined
by which ReLU in which layer.

In contrast to the above the proposed MinMax approach (3) systematically
defines

    * –

all convex edges of the pyramid in $\hat{y}_{1\max}$ and

    * –

all concave edges of the pyramid in $\hat{y}_{2\min}$

with

$\displaystyle\hat{y}({\bf x})$ $\displaystyle=$
$\displaystyle\hat{y}_{1\max}+\hat{y}_{2\min}$ $\displaystyle\hat{y}_{1\max}$
$\displaystyle=$
$\displaystyle\max(\hat{z}_{11},\hat{z}_{12},\hat{z}_{13},\hat{z}_{14},\hat{z}_{15}),\
\hat{y}_{2\min}=\min(\hat{z}_{21},\hat{z}_{22},\hat{z}_{23},\hat{z}_{24})$
$\displaystyle\hat{z}_{11}$ $\displaystyle=$ $\displaystyle 0,\
\hat{z}_{12}=x_{1}+1,\ \hat{z}_{13}=x_{2}+1,\ \hat{z}_{14}=-x_{1}+1,\
\hat{z}_{15}=-x_{2}+1$ $\displaystyle\hat{z}_{21}$ $\displaystyle=$
$\displaystyle x_{1}+1,\ \hat{z}_{22}=x_{2}+1,\ \hat{z}_{23}=-x_{1}+1,\
\hat{z}_{24}=-x_{2}+1$

where subfigures 1c and 1d show the convex and concave neurons of the MinMax
approach. The legend indicates which basic neurons produce a nonzero output
for the colored surfaces.

(a) Unit pyramid

(b) Intermediate neuron $\hat{z}_{1}$ of ReLU network

(c) Concave neuron $\hat{y}_{1\max}$ of MinMax

(d) Convex neuron $\hat{y}_{2\min}$ of MinMax

Figure 1: Unit pyramid, intermediate ReLU neuron and a MinMax network with two
neurons

$\Box$

Motivated by the above section 3 introduces a piece-wise linear MinMax
discrete function learning (3) for the $N$-dimensional case. The approach
still uses gradient descent on a quadratic cost in discrete time. Exact
exponential stability guarantees are provided using the results of section 2.
Saddle points or sub-optimal plateaus are avoided with a linear
parametrization. Possible instabilities of the non-Lipschitz edges are avoided
with intermediate Lagrange constraints (2). Note that time-varying
measurements and the time discretization of the gradient are part of the
exponential stability proof. Finally the result is summarised in the Summary.

##  2 Discrete-time constrained systems

This section extends Contraction Analysis of Continuous Constrained Systems
[6] to the discrete case:

A constraint $l$ in (2) only has an impact on the unconstrained dynamics (1)
if the inequality turns into an equality $g_{l}=0$ , leading to the following
definition.

###### Definition 1

The set of active constraints $\mathcal{A}(({\bf
x}^{i+1},{i+1})\subseteq\\{1,...,L\\}$ contains the elements $l$ (2) which are
on the boundary of the original constraint

$\displaystyle g_{l}={\bf g}_{l}^{T}(i+1){\bf x}^{i+1}+h_{l}(i+1)$
$\displaystyle=$ $\displaystyle 0.$

The constrained dynamic equations (1, 2) are then of the form

${\bf x}^{i+1}={\bf f}({\bf x}^{i},i)+\sum_{all\ l\in\mathcal{A}}{\bf
g}_{l}{\lambda}_{l}.$ (4)

All constraints $g_{l},l\in\mathcal{A}$ are not violated at $i+1$ if

${g}_{l}={\bf g}^{T}_{l}(i+1)\left({\bf f}({\bf x}^{i},i)+\sum_{all\
l\in\mathcal{A}}{\bf g}_{l}(i+1){\lambda}_{l}\right)+h_{l}(i+1)\leq 0$

leading to the following definition:

###### Definition 2

The set of Lagrange multipliers $\lambda_{l},l\in\mathcal{A}$ of (2) is the
solution of the linear programming problem

$\displaystyle maximize\sum_{all\ l\in\mathcal{A}}\lambda_{l}\ subject\ to$
$\displaystyle{\bf g}^{T}_{l}\left({\bf f}({\bf x}^{i},i)+\sum_{all\
l\in\mathcal{A}}{\bf g}_{l}{\lambda}_{l}\right)+h_{l}$ $\displaystyle\leq$
$\displaystyle 0$ $\displaystyle\lambda_{l}$ $\displaystyle\leq$
$\displaystyle 0$

where we requested similar to [6] that the constraint term points in the
interior of the constraint with $\lambda_{l}\leq 0$ and the negative
$\lambda_{l}$ are maximized [3] to minimize their joint usage. The equation
above corresponds to the initial solution problem of Linear Programming (LP),
see e.g. section LP in [2]. Diverse solutions exist for this problem, such as
the simplex algorithm in section LP in [2].

Similar to [6] we introduce a virtual displacement between two neighbouring
trajectories, which is constrained by $g_{l}=0$ at the active
$l\in\mathcal{A}$. This virtual displacement has to be parallel to $g_{l}=0$,
i.e. orthogonal to the normals ${\bf g}_{l}$, which implies

$\displaystyle{\bf\delta x}^{i+1}$ $\displaystyle=$ $\displaystyle{\bf
G}_{\parallel}(i+1){\bf\delta x}^{i+1\ast}$ (5) $\displaystyle{\bf
G}_{\parallel}^{T}{\bf G}_{\parallel}$ $\displaystyle=$ $\displaystyle{\bf I}$
$\displaystyle{\bf G}_{\parallel}^{T}{\bf g}_{l}(i+1)$ $\displaystyle=$
$\displaystyle{\bf 0}\ \ \ \ \ \ \forall\ l\in\mathcal{A}$

where ${\bf\delta x}^{i+1\ast}$ is the reduced virtual displacement of
dimension $n$ minus by the number of active constraints in $\mathcal{A}$.

The constrained dynamic equations (4) can be rewritten as:

${\bf x}^{i+1}={\bf f}({\bf x}^{i},i)+\sum_{all\ l\in\mathcal{A}}{\bf
g}_{l}\lambda_{l}={\bf f}({\bf x}^{i},i)+\sum_{l=1}^{L}{\bf
g}_{l}step(g_{l})\lambda_{l}$

with the step function

$step(g_{l})=\left\\{\begin{array}[]{l}0\text{\ \ for \ \ }g_{l}<0\\\ 1\text{\
\ for \ \ }g_{l}=0\\\ undefined\text{\ \ for \ \ }g_{l}>0\end{array}\right.$

whose variation is

$\displaystyle\delta{\bf x}^{i+1}$ $\displaystyle=$
$\displaystyle\frac{\partial{\bf f}}{\partial{\bf x}^{i}}{\bf\delta
x}^{i}+\sum_{l=1}^{L}\left({\bf g}_{l}\frac{\partial\lambda_{l}}{\partial{\bf
x}^{i+1}}+{\bf g}_{l}\frac{\partial step(g_{l})}{\partial g_{l}}{\bf
g}_{l}^{T}\lambda_{l}\right){\bf\delta x}^{i+1}$

Now the squared virtual length dynamics with a metric ${\bf M}^{i}({\bf
x}^{i},i)$ can be bounded as

$\displaystyle\frac{1}{2}\ {\bf\delta x}^{i+1T}{\bf M}^{i+1}{\bf\delta
x}^{i+1}$ $\displaystyle=$ $\displaystyle{\bf\delta x}^{iT}\frac{\partial{\bf
f}}{\partial{\bf x}^{i}}^{T}{\bf M}^{i+1}\frac{\partial{\bf f}}{\partial{\bf
x}^{i}}{\bf\delta x}^{i}+\frac{1}{2}\ {\bf\delta x}^{i+1T}({\bf M}^{i+1}{\bf
g}_{l}\frac{\partial step(g_{l})}{\partial g_{l}}{\bf
g}_{l}^{T})_{H}\lambda_{l}{\bf\delta x}^{i+1}$

where we used ${\bf G}_{\parallel}^{T}{\bf
g}_{l}\frac{\partial\lambda_{l}}{\partial{\bf x}^{i+1}}={\bf 0}$ since on the
constraint the first two terms vanish and outside the constraint the last term
vanishes. The Dirac impulse $\frac{\partial step(g_{l})}{\partial g_{l}}$
discontinuously sets the virtual displacement ${\bf\delta x}^{i+1}$ to ${\bf
G}_{\parallel}{\bf\delta x}^{i+1\ast}$ when a constraint is activated.

Thus the dynamics of ${\bf\delta x}^{i+1T}{\bf\delta x}^{i+1}$ is composed of
exponentially convergent continuous segments and an enforcement of ${\bf\delta
x}^{i+1}$ to ${\bf G}_{\parallel}{\bf\delta x}^{i+1\ast}$ at the activation of
a constraint. Let us summarize this result.

###### Theorem 1

Consider the discrete dynamics

${\bf x}^{i+1}={\bf f}({\bf x}^{i},i)+\sum_{all\ l\in\mathcal{A}}{\bf
g}_{l}\lambda_{l}$ (6)

within the metric ${\bf M}({\bf x}^{i},i)$ constrained by a
$l=1,...,L$-dimensional inequality constraint

$g_{l}={\bf g}_{l}^{T}(i+1){\bf x}^{i+1}+h_{l}(i+1)\leq 0$ (7)

The set of active constraints $\mathcal{A}$ and Lagrange multipliers
$\lambda_{l}$ are given in Definition 1 and 2

The distance $s=\min\int_{{\bf x}^{i}(s)={\bf x}^{i}_{1}}^{{\bf
x}^{i}_{2}}\sqrt{\delta{\bf x}^{iT}{\bf M}^{i}\delta{\bf x}^{i}}$ within
$\mathbb{G}^{n}$ from any trajectory ${\bf x}^{i}_{1}(t)$ to any other
trajectory ${\bf x}^{i}_{2}(t)$ converges exponentially to $0$ with an
exponential convergence rate $\leq\max_{along\ s}(\sigma_{\max}({\bf
x}^{i},i)),\sigma_{\max}>0$ ($\geq\min_{along\ s}(\sigma_{\min}({\bf
x}^{i},i)),\sigma_{\min}>0$) with

$\displaystyle\sigma_{\min}^{2}{\bf G}^{T}_{\parallel}{\bf M}^{i}{\bf
G}^{T}_{\parallel}\leq{\bf G}_{\parallel}\frac{\partial{\bf f}}{\partial{\bf
x}^{i}}^{T}{\bf M}^{i+1}\frac{\partial{\bf f}}{\partial{\bf x}^{i}}{\bf
G}^{T}_{\parallel}$ $\displaystyle\leq$ $\displaystyle\sigma_{\max}^{2}{\bf
G}^{T}_{\parallel}{\bf M}^{i}{\bf G}^{T}_{\parallel}$ (8)

with the constrained tangential space ${\bf G}_{\parallel}$ from equation (5).

In addition the activation of a constraint discontinuously sets the virtual
displacement ${\bf\delta x}^{i+1}$ to ${\bf G}_{\parallel}{\bf\delta
x}^{i+1\ast}$.

Note that section 3.3. of [1] provides a stability condition on the
constrained step ${\bf f}_{i+1}-{\bf f}_{i}$ for the original contraction
mapping theorem. Theorem 1 makes this condition more concrete by giving
explicit conditions on the derivatives of ${\bf f}$. It also extends the
result to a metric. In addition all notes of the continuous theorem in [6]
apply here as well.

##  3 Discrete exponential stable learning

Let us us assume a piece-wise linear $N$-dimensional measurement function

$y=y({\bf x})$

where the $N$-dimensional input vector ${\bf x}^{\prime}$ is augmented to the
$N+1$-dimensional input vector ${\bf x}=({\bf x}^{{}^{\prime}T},1)^{T}$. We
have $m=1,...,M$ measurements

$y_{m}^{i}=y({\bf x}_{m}^{i})$

with the measured input vector ${\bf x}^{i}_{m}$ at time index $i$. The goal
is to approximate $y_{m}^{i}$ with $\hat{y}_{m}^{i}({\bf x}_{m}^{i})$. We
achieve the approximation of the true function by minimizing the weighted cost

$V=\frac{1}{2}\sum_{m=1}^{M}\alpha_{m}^{2}\tilde{y}_{m}^{i2}$

with $\tilde{y}_{m}^{i}=\hat{y}_{m}^{i}({\bf x}_{m}^{i})-y^{i}_{m}({\bf
x}_{m}^{i})$ and measurement weight $\alpha_{m}^{2}({\bf x}_{m}^{i})\geq 0$.
The unconstrained parameter learning law is the classical gradient descent of
$V$

$\displaystyle{\hat{\bf W}}_{j}^{i+1}$ $\displaystyle=$
$\displaystyle{\hat{\bf W}}_{j}^{i}-\frac{\partial V}{\partial\hat{\bf
W}_{j}}$ (9) $\displaystyle=$ $\displaystyle{\hat{\bf
W}}_{j}^{i}-\sum_{m=1}^{M}{\bf A}_{j}({\bf
x}_{m}^{i})\alpha_{m}^{2}\tilde{y}_{m}$ $\displaystyle=$
$\displaystyle{\hat{\bf W}}_{j}^{i}-\sum_{m=1}^{M}{\bf A}_{j}({\bf
x}_{m}^{i})\alpha_{m}^{2}\left(\sum_{l=1}^{J}{\bf A}_{l}^{T}({\bf
x}_{m}^{i}){\hat{\bf W}}_{l}^{i}-y^{i}_{m}\right)$

with $\hat{\bf W}_{j}^{iT}=(\hat{\bf w}_{j1}^{iT},...,\hat{\bf
w}_{jK_{j}}^{iT})$ and the activation function

${\bf A}_{j}({\bf x})=\left\\{\begin{array}[]{l}{\bf x}\text{\ \ if
}\hat{z}_{jk{\ast}}=\hat{y}_{j\min}(\hat{y}_{j\max})\text{ for a
}\min(\max)\text{ neuron}\\\ {\bf 0}\text{\ \ for all other \ }k\neq
k^{\ast}\end{array}\right.$ (10)

where for multiple solutions
$\hat{z}_{jk{\ast}}=\hat{y}_{j\min}(\hat{y}_{j\max})$ only one activation is
set to ${\bf x}$ and all others to ${\bf 0}$. Note that the gradient in (9)
could be multiplied on top with a gain $\alpha_{j}(i)$ which we have skipped
for simplicity here.

The weight dynamics (9) is equivalent to the measurement estimation dynamics

$\displaystyle\hat{y}^{i+1}_{n}$ $\displaystyle=$
$\displaystyle\sum_{j=1}^{J}{\bf A}_{j}^{T}({\bf x}_{n}^{i+1}){\hat{\bf
W}}_{j}^{i+1}$ (11) $\displaystyle=$
$\displaystyle\hat{y}^{i}_{n}-\sum_{j=1}^{J}{\bf A}_{j}^{T}({\bf
x}_{n}^{i})\sum_{m=1}^{M}{\bf A}_{j}({\bf
x}_{m}^{i})\alpha_{m}^{2}\tilde{y}_{m}.$

In this last formulation we assume static measurements ${\bf x}_{n}^{i+1}={\bf
x}_{n}^{i}$.

To assure Lipschitz continuity of the active basic neurons
$\hat{z}_{jk^{\ast}}$ at the learning step $i+1$ we have to exclude with
Theorem 1 a transition beyond the edges of the $\min$ and $\max$ operator with
the constraints

$\displaystyle g_{l}={\bf g}_{l}^{T}\hat{\bf W}_{j}^{i+1}={\bf
x}_{m}^{i+1T}\hat{\bf w}_{jk^{\ast}}^{i+1}-{\bf x}_{m}^{i+1T}\hat{\bf
w}_{jk}^{i+1}$ $\displaystyle\leq$ $\displaystyle 0\text{ for a }\min\text{
neuron j and all }k\neq k^{\ast}$ $\displaystyle g_{l}={\bf g}_{l}^{T}\hat{\bf
W}_{j}^{i+1}={\bf x}_{m}^{i+1T}\hat{\bf w}_{jk}^{i+1}-{\bf
x}_{m}^{i+1T}\hat{\bf w}_{jk^{\ast}}^{i+1}$ $\displaystyle\leq$ $\displaystyle
0\text{ for a }\max\text{ neuron j and all }k\neq k^{\ast}$

where the related Lagrange parameters are given in Definition 2.

Theorem 1 then implies global contraction behaviour with the largest and
smallest singular value of the variation of (9) or alternatively (11) within
the constraints above.

Summarizing the above leads to:

###### Theorem 2

Consider a piece-wise linear $N$-dimensional measurement function

$y=y({\bf x})$ (12)

where the $N$-dimensional input vector ${\bf x}^{\prime}$ is augmented to the
$N+1$-dimensional input vector ${\bf x}=(1,{\bf x}^{{}^{\prime}T})^{T}$. We
have $m=1,...,M$ measurements

$y_{m}^{i}=y({\bf x}_{m}^{i})$ (13)

with the measured input vector ${\bf x}_{m}^{i}$ at time index $i$. We
approximate (12) with

$\hat{y}({\bf x}_{m}^{i})=\sum_{j=1}^{J_{\min}}\hat{y}_{j\min}({\bf
x}_{m}^{i})+\sum_{j=J_{\min}}^{J_{\max}}\hat{y}_{j\max}({\bf x}_{m}^{i})$ (14)

which uses the convex and concave neurons $j$

$\displaystyle\hat{y}_{j\min}({\bf x}_{m}^{i})$ $\displaystyle=$
$\displaystyle\min(\hat{z}_{j1},...,\hat{z}_{jK_{j}})$
$\displaystyle\hat{y}_{j\max}({\bf x}_{m}^{i})$ $\displaystyle=$
$\displaystyle\max(\hat{z}_{j1},...,\hat{z}_{jK_{j}})$ (15)

which consist of $k=1,...,K_{j}$ linear basic neurons $\hat{z}_{jk}={\bf
x}_{m}^{iT}\hat{\bf w}_{jk}^{i}$ of estimated $N+1$-dimensional weight vector
$\hat{\bf w}_{jk}^{i}$ at time $i$ and the activation function (10). For a
neuron $j$ we constrain the dynamics with

$\displaystyle g_{l}={\bf g}_{l}^{T}\hat{\bf W}_{j}^{i+1}={\bf
x}_{m}^{i+1T}\hat{\bf w}_{jk^{\ast}}^{i+1}-{\bf x}_{m}^{i+1T}\hat{\bf
w}_{jk}^{i+1}$ $\displaystyle\leq$ $\displaystyle 0\text{ for a }\min\text{
neuron j, all }k\neq k^{\ast}$ $\displaystyle g_{l}={\bf g}_{l}^{T}\hat{\bf
W}_{j}^{i+1}={\bf x}_{m}^{i+1T}\hat{\bf w}_{jk}^{i+1}-{\bf
x}_{m}^{i+1T}\hat{\bf w}_{jk^{\ast}}^{i+1}$ $\displaystyle\leq$ $\displaystyle
0\text{ for a }\max\text{ neuron j, all }k\neq k^{\ast}$ (16)

with $\hat{\bf W}_{j}^{iT}=(\hat{\bf w}_{j1}^{iT},...,\hat{\bf
w}_{jK_{j}}^{iT})$. The constrained learning of the cost

$V=\frac{1}{2}\sum_{m=1}^{M}\alpha_{m}^{2}\tilde{y}_{m}^{i2}$ (17)

with $\tilde{y}_{m}^{i}=\hat{y}_{m}^{i}({\bf x}_{m}^{i})-y^{i}_{m}({\bf
x}_{m}^{i})$ and measurement weight $\alpha_{m}^{2}({\bf x}_{m}^{i})\geq 0$

$\displaystyle\hat{\bf W}_{j}^{i+1}$ $\displaystyle=$ $\displaystyle{\hat{\bf
W}}_{j}^{i}-\frac{\partial V}{\partial\hat{\bf W}_{j}}+\sum_{all\
l\in\mathcal{A}}{\bf g}_{l}\lambda_{l}$ (18) $\displaystyle=$
$\displaystyle\hat{\bf W}_{j}^{i}-\sum_{m=1}^{M}{\bf A}_{j}({\bf
x}_{m})\alpha_{m}^{2}(i)\tilde{y}_{m}^{i}+\sum_{all\ l\in\mathcal{A}}{\bf
g}_{l}\lambda_{l}$

where the set of active constraints $\mathcal{A}$ and Lagrange multipliers
$\lambda_{l}$ are given in Definition 1 and 2 is globally exponentially
converging to $\frac{\partial V}{\partial\hat{\bf W}_{j}}={\bf 0}$ with the
largest and smallest singular value of the matrix

$\displaystyle{\bf I}_{jk}-\sum_{m=1}^{M}{\bf A}_{j}({\bf
x}_{m}^{i})\alpha_{m}^{2}{\bf A}_{k}^{T}({\bf x}_{m}^{i})$ (19)

or alternatively in the metric $\alpha_{n}^{2}$ with the block diagonal matrix

$\displaystyle{\bf I}_{nm}-\sum_{j=1}^{J}\alpha_{n}{\bf A}_{j}^{T}({\bf
x}_{n}^{i}){\bf A}_{j}({\bf x}_{m}^{i})\alpha_{m}$ (20)

In this last formulation we assume static measurements ${\bf x}_{n}^{i+1}={\bf
x}_{n}^{i}$ e.g. for batch processing.

The equilibrium point $\frac{\partial V}{\partial\hat{\bf W}_{j}}={\bf 0}$ is
unique or global if the largest singular value is strictly less then $1$.

Note that e.g. $0<\alpha_{m}\leq\frac{1}{J\lvert{\bf x}_{m}^{i}\rvert}$
assures contraction behaviour in (20).

Also note that at the minimum of $V$ the cost $V$ will not be $0$ if an
incomplete active topology (14) was used. Hence the following pruning and
creation principles have to be applied to find a right active topology (14).
For a right topology $V$ will go to $0$ exponentially since a contracting
solution $V=0$ then exists.

  * •

Inactive neurons (15) close to $0$ can be pruned.

  * •

Basic neurons of a neuron, which never become active or which are similar to
another basic neuron, can be pruned.

  * •

New convex or concave neurons (15) $j$ can be initialized with all
$\hat{z}_{jk}=0$ and one activated basic neuron if persistent relevant errors
$\tilde{y}_{m}$ remain.

In principle one convex and one concave neuron are sufficient to activate
convex or concave edges everywhere. However, the numerical complexity of the
simplex algorithm in Definition 2 is exponential in the worst case. Hence it
is computationally more efficient to take several convex or concave neurons of
low dimension rather then one convex or concave neuron of very high dimension.

  * •

Similarly, new basic neurons can be created by duplicating existing basic
neurons $jk$ with persistent relevant errors $\tilde{y}_{m}$

$\hat{z}_{jnew}=\hat{z}_{jk}$

Note that when the Lagrangian constraint is active at the boundary between two
basic neurons the following cases exist for the trajectory which starts at the
next iteration on the boundary between both neurons:

  * •

Both trajectories move away from the boundary. Here one trajectory can be
selected. Several solutions may exist here since the boundary is not Lipschitz
continuous. Note that especially in this case $\delta{\bf x}^{i+1}$ does not
diverge since it was parallelized to the constraint the iteration before.

  * •

One trajectory moves to the boundary the other away from the boundary. Here
the trajectory which moves away from the boundary is selected.

  * •

Both trajectories again move to the boundary. Here the Lagrangian constraint
has to be maintained active at the next iteration, such that the trajectory
moves on the edge until it eventually leaves the edge at a later iteration.

The following example illustrates the effect above:

  * ###### Example 3.1

:  Figure 2 shows the learning of a single measurement at the edge of 2 basic
neurons of Theorem 2:

    * –

The left side shows a measurement on the convex side of 2 basic neurons.
Learning of the left or right neuron alone does not work since the measurement
cost (17) at $i$

$V=\frac{1}{2}\sum_{m=1}^{M}\tilde{y}_{m}^{i2}$

has different measurements $m$ at $i+1$. Learning works here only with an
active Lagrangian constraint (16) of Theorem 2

    * –

On the right side the measurement is on the concave side of 2 basic neurons.
Here both neurons can learn without an active Lagrangian constraint (16) since
the measurement stays at $i+1$ at the same basic neuron where it was at $i$.

Figure 2: Learning with measurements on convex or concave side

$\Box$

  * ###### Example 3.2

:  Figure 3 shows how a MinMax network of Theorem 2 evolves for a polygon. We
approximate the target polygon $y$ (12) in subfigure 3a with the MinMax
learning (18):

    * –

We start with one linear neuron $\hat{y}=y_{1,min}=min(\hat{z}_{11})$ (15) in
subfigure 3b.

    * –

We then insert a new basic neuron in
$\hat{y}=y_{1,min}=min(\hat{z}_{11},\hat{z}_{12})$, where $\hat{z}_{11}$ and
$\hat{z}_{12}$ have initially the same parameters. After training, the network
converges to the approximation in subplot 3c.

    * –

A new concave neuron with two basic neurons is then inserted leading to
$\hat{y}=min(\hat{z}_{11},\hat{z}_{12})+max(\hat{z}_{21},\hat{z}_{22})$ in
subfigure 3d.

    * –

After another insertion of a basic neuron, we get a perfect approximation in
subfigure 3e with
$\hat{y}=min(\hat{z}_{11},\hat{z}_{12})+max(\hat{z}_{21},\hat{z}_{22},\hat{z}_{23})$.
The final two neurons (15) are depicted in subfigures 3f and 3g.

The above example shows how important the (basic) neuron creation is to learn
the topology of the MinMax network.

The approach of this paper uses 5 basic neurons, whereas the benchmark in [10]
had 100 to 2000 neurons in several layers. The benchmark in [10] took up to
50000 iterations to converge to remaining persistent errors, where the MinMax
network only needs a few hundred. The approach of this paper needs 8
measurement points, i.e. the minimum number of points to define the polygon.
[10] used 100 measurement points for learning.

(a) Target function $y$

(b) $\hat{y}$ with one linear neuron

(c) $\hat{y}$ with one convex neuron

(d) $\hat{y}$ with one concave and one convex neuron

(e) $\hat{y}$ with with one concave and one convex neuron

(f) Concave neuron $\hat{y}_{1\min}$

(g) Convex neuron $\hat{y}_{2\max}$

Figure 3: Approximation of 1-dimensional target function by MinMax network

$\Box$

  * ###### Example 3.3

:  Figure 4 shows how a MinMax network of Theorem 2 evolves. We approximate
the target function $y$ (12) in subfigure 4a with the MinMax learning (18):

    * –

We start with one linear neuron in subfigure 4b.

    * –

The network then continuously inserts new basic neurons. Subfigure 4c depicts
the approximated function with one neuron consisting of 3 linear basic neurons
(15).

    * –

Figure 4d shows the final result, which perfectly matches the target function.
The final approximation $\hat{y}$ is the sum of the concave neuron
$\hat{y}_{1\max}$ and convex neuron $\hat{y}_{2\max}$ (15) shown in subfigure
4e and 4f respectively.

(a) Target function $y$

(b) $\hat{y}$ after initialization with one linear neuron

(c) $\hat{y}$ after two insertions of two basic neurons

(d) $\hat{y}$ after full training with two neurons

(e) Concave neuron $\hat{y}_{1\max}$ of MinMax

(f) Convex neuron $\hat{y}_{2\min}$ of MinMax

Figure 4: Approximation of 2-dimensional target function by MinMax network

$\Box$

  * ###### Example 3.4

:  Figure 5 shows the convergence of a MinMax Network of Theorem 2 to a
$8$-dimensional function

$y({\bf x})=\max_{n=1,...,8}(\lvert x_{n}\rvert)$

The network is initialized with a single neuron with random parameters
$w_{n}\in[-0.5,0.5]$ and further neurons are created after 100 iterations
until the total cost $V$ (17) converges to zero. Steep drops in the error
value indicate such insertions. Neurons are pruned if they become inactive or
too similar to other basic neurons, leading to a network with the minimal
required number of 16 basic neurons in one $\max$ neuron for the 16 linear
surfaces.

Figure 5: Convergence of the cost $V$ over $i$ for the approximation of 16
surfaces in an 8-dimensional space by a MinMax network.

$\Box$

##  4 Summary

This paper first extends discrete contraction theory to non-linear constrained
systems in Theorem 1. This allows to propose a new class of MinMax networks in
Theorem 2 for the learning of piece-wise linear functions:

  * •

Possible instabilities or even non-unique solutions at the discontinuity
between the linear regions are avoided by limiting ${\bf x}^{i+1}$ to its
linear subspace with a Lagrangian constraint of Theorem 1. This linear
subspace may change at the next time instance since ${\bf x}^{i+1}$ then
belongs to both neighbouring linear regions.

From a Contraction Theory perspective this constrained step parallelizes the
virtual displacement $\delta{\bf x}^{i+1}$ to the edge, which is then
orthogonal to the Dirac instability at the edge. Hence, it has no impact on
the contraction rate under this constraint.

  * •

Saddles points or sub-optimal plateaus are avoided with a linear
parametrization of the MinMax network. This is not possible with a deep
network since the parametrization is highly polynomial.

As a result exponential convergence guarantees are given with Theorem 2 for
the discrete time learning of piece-wise linear functions.

Although the learning was shown to be contracting in Theorem 2 the remaining
errors will not necessarily go to $0$ if a wrong topology was used. Hence the
current research focus is to define finite neuron creation principles to find
a correct topology of the MinMax network, where one error free solution
exists.

Also, since each basic neuron is linear in Theorem 2, all linear estimation
techniques can be exploited as e.g. computation of covariance matrices.

## References

  * [1] Bertesekas D., Tsitsiklis N., Parallel and Distributed Computation: Numerical Methods, Athena Scientific, 2014.
  * [2] Bronstein, Semendjajew, Taschenbuch der Mathematik, Teubner, 1991.
  * [3] Bryson A., Yu-Chi H., Applied Optimal Control, Taylor and Francis, 1975.
  * [4] LeCun Y., Bengio Y., and Hinton G., Deep learning, Nature, 2015.
  * [5] Lohmiller, W., and Slotine, J.J.E., On Contraction Analysis for Nonlinear Systems, Automatica, 34(6), 1998.
  * [6] Lohmiller, W., and Slotine, J.J.E., Contraction Theory with Inequality Constraints, arXiv:1804.10085, 2023.
  * [7] Lohmiller, W., Gassert P. and Slotine, J.J.E., Notes on stable learning with piecewise-linear basis functions, arXiv:1804.10085, 2018.
  * [8] Lohmiller, W., Gassert P. and Slotine, J.J.E., Deep MinMax Networks, 60th IEEE Conference on Decision and Control (CDC), IEEE, 2021.
  * [9] https://en.wikipedia.org/wiki/Minimax
  * [10] Shalev-Shwartz S., Shamir O. and Shammah S., Failures of Gradient-Based Deep Learning, arXiv:1703.07950v2, Apr 2017.


# Improving Locality in Sparse and Dense Matrix Multiplications

Mohammad Mahdi Salehi Dezfuli  McMaster UniversityHamiltonCanada
<EMAIL_ADDRESS>and  Kazem Cheshmi  0000-0002-2968-5176 McMaster
UniversityHamiltonCanada<EMAIL_ADDRESS>
###### Abstract.

Consecutive matrix multiplications are commonly used in graph neural networks
and sparse linear solvers. These operations frequently access the same
matrices for both reading and writing. While reusing these matrices improves
data locality, it presents a challenge due to the irregular dependencies
between iterations across the two multiplication operations. Existing fusion
methods often introduce excessive synchronization overhead or overlapped
computations with limited benefits. This paper proposes tile fusion, a runtime
approach that fuses tiles of the two matrix-matrix multiplications, where at
least one of the involved matrices is sparse. Tile fusion aims to improve data
locality while providing sufficient workload for cores in shared-memory multi-
core processors. For a pair of matrix-matrix multiplications, tile fusion
outperforms unfused baseline and MKL implementations with a geometric mean
speedup of 1.97$\times$ 1.64$\times$, respectively, on multi-core CPUs.

##  1\. Introduction

Consecutive calls to matrix multiplications are the computational bottleneck
in many scientific (O’Leary, 1980) and machine learning (Wang et al., 2019;
Fey and Lenssen, 2019) applications. Particularly this paper focuses on
accelerating a pair of matrix multiplications, represented as an equation:

(1) $D=A(BC)$

where matrix $\underset{n\times n}{A}$ is sparse, $\underset{n\times bCol}{B}$
is either sparse or dense, and $\underset{bCol\times cCol}{C}$ is dense. For
example, in a layer of graph convolution network (Kipf and Welling, 2016),
either cases happen. Existing frameworks such as PyTorch Geometric (PyG) (Fey
and Lenssen, 2019) and Deep Graph Library (DGL) (Wang et al., 2019) break the
expression into two matrix multiplication operations, $D_{1}=BC$ and
$D=AD_{1}$. The two operations are commonly mapped to a pair of General Matrix
Multiplication (GeMM)-Sparse Matrix-Matrix Multiplication (SpMM) or SpMM-SpMM
routines when $B$ is dense and sparse, respectively. These routines benefit
from efficient tiling and load balancing techniques (Hong et al., 2019; Wang
et al., 2014) that enable using memory and computing resources efficiently.
However, $D_{1}$ is shared between the two routines and often a large matrix
that can be reused but it is not used when the operation is mapped to GeMM or
SpMM separately.

Fusing operations or loops are commonly used to remove intermediate matrices
between the two operations. Tensor compilers (Kjolstad et al., 2017; Dias et
al., 2022; Mutlu et al., 2022) generate a fused code for Equation 1 when $A$
is sparse and $B$ and $C$ are dense. The generated code iterates over $A$ and
performs a general matrix-vector multiplication (GeMV) for each nonzero of
$A$. While this removes the need for storing intermediate results, i.e.
$D_{1}$, it causes random access to $B$ and thus inefficient use of memory
hierarchy. Additionally, this methodology does not apply when $A$ and $B$ are
sparse because memory accesses are unknown at compile time.

Figure 1. The ratio of computations in coarse fused tiles for all matrices
from SuiteSparse for GEMM-SpMM operation. Figure 2. Three different iteration
fusion schedules (Figure 2d–f) for the GeMM-SpMM in Figure 2b and the matrix
in Figure 2a. Figure 2c shows the dependence DAG between iterations of the
outermost loop of GeMM and SpMM, where colored and white vertices correspond
to GeMM and SpMM iterations, respectively. Dark solid lines show
synchronization barriers, the dotted red line shows a potential race
condition, and vertical dashed lines show per thread workload.

Prior approaches such as sparse tiling (Krieger et al., 2013) and
communication-avoiding (CA) (Demmel et al., 2008) methods have used sparsity
information at runtime to fuse sparse matrix-vector multiplications (SpMV) and
enable reuse between the two operations. They model SpMV operations as an
iteration data acyclic graph (DAG) where vertices are iterations of the
outermost loop of SpMV and edges represent dependencies between iterations.
Then a scheduler tiles iterations of operations by grouping vertices of DAG at
runtime. Then, sparse tiling uses barrier and atomic operations to ensure
dependence between tiles are not violated during parallel execution. Some CA
methods (Demmel et al., 2008) replicate dependent iterations within a tile to
make all tiles independent so they run in parallel without synchronization.
Since GeMM, SpMM, and SpMV have parallel iterations in their outermost loops,
the same techniques can be adopted for fusing GeMM-SpMM and SpMM-SpMM.
However, the computation of each fused iteration in the two operations, is
proportional with $bCol$ and $cCol$, increasing race conditions in sparse
tiling and redundant computation in CA methods.

Coarse-grain tiles provide opportunities for fusion in sparse matrices and
graphs without redundant computation or excessive synchronization. A coarse
tile contains large enough iterations of the first operation such that it
allows running some iterations of the second operation that solely depend on
iterations inside the tile. This allows tiles to execute in parallel without
synchronization. Figure 1 shows the percentage of GeMM-SpMM computations that
share data across the operations if coarse-grain tiles with the size of 2048
are selected for all 2893 matrices from SuiteSparse matrix collection (Davis
and Hu, 2011). As shown, an average of 34% of GeMM-SpMM computation reuse data
in fused coarse tiles. However, growing the tiles reduces the number of
parallel workloads, affecting load balance. Also, picking coarse grain tiles
groups a larger number of iterations from the two operations. This grouping
improves locality if the memory accesses of the tile fit within the size of
the fast memory.

We propose sparsity-oriented tile fusion, in short, tile fusion, that creates
fused tiles based on the opportunities shown in Figure 1 to improve locality
in GeMM-SpMM and SpMM-SpMM for shared memory multicore processors. This paper
makes the following contributions:

  * •

Tile fusion scheduler and fused code that turn data reuse between and across
iterations of GeMM and SpMM into locality. The tile fusion scheduler uses the
sparsity pattern of $A$ and selects tile sizes and a number of tiles to ensure
locality and load balance.

  * •

An implementation that is tested for a wide range of graphs and matrices and
provides a speedup of 1.97$\times$ and 3.52$\times$ compared to existing
unfused and best-fused codes. Also an analysis and adoption of prior tiling
approaches and comparison with tile fusion.

##  2\. Motivating Example

We use the matrix in Figure 2a to discuss how different fusion strategies
improve locality for computing Equation 1. The corresponding code to the
computation is shown in Figure 2b where lines 1–4 perform GeMM, $D_{1}=BC$,
and lines 5–8 perform SpMM, $D=AD_{1}$. Iterations of loops i1 and j1 are
independent so they execute in parallel. Fusing loops i1 and j1 can
potentially enable reusing $D_{1}$ but each iteration in j1 depends on a
variant number of i1 iterations. This irregular dependence is due to
D1[A.i[j2]][j3] in line 8 in Figure 2b, stemming from sparsity pattern of $A$.
The DAG shown in Figure 2c shows the dependence between i1 and j1. Colored and
white vertices in Figure 2c represent iterations of i1 and j1 loops,
respectively. Edges show dependence between iterations. While grouping
vertices with common edges as a tile improves locality, dependence between
tiles can prevent keeping all cores busy. Three different fused schedules of
iterations for the DAG shown in Figure 2c are shown in Figure 2d–f for a
processor with three cores.

Figure 2d shows five tiles composed of vertices of both computations with
common edges. Dependent tiles are separated by synchronization barriers to
ensure partial order. Tiles are atomic to prevent race conditions. For
example, tile $\mathcal{T}_{1,0}$ and $\mathcal{T}_{1,1}$ depend on tile
$\mathcal{T}_{0,0}$ and $\mathcal{T}_{0,1}$ thus a synchronization is needed
between them. Iteration j=4 is split among tiles $\mathcal{T}_{1,0}$ and
$\mathcal{T}_{1,1}$, writing to the same location of C, thus an atomic
operation is needed. The race condition is shown with the dotted red line in
Figure 2. This schedule is inspired by sparse tiling (Krieger et al., 2013)
and named atomic tiling due to atomic operations used in tiles. The chance of
race condition on writing to C increases as the number of columns in $B$ nad
$C$ increases.

Figure 2e shows overlapped tiles that create independent tiles by replicating
dependent iterations. Replicated iterations are shown with red vertices in two
tiles in Figure 2e. Therefore all fused tiles execute in parallel with no
synchronization. Each replicated vertex in the tile corresponds to an
iteration i1 which multiplies a row of $B$ with $C$. Therefore redundant
computations increase with the number of columns in $B$ and $C$. Due to
replicated iterations, this method is called overlapped tiling, inspired by CA
(Demmel et al., 2008) methods.

The tile fusion schedule is shown in Figure 2f where two groups of tiles are
created, fused tiles and tiles of the SpMM iterations separated by one
synchronization barrier. As shown, tiles in the schedule can be large, such as
tile $\mathcal{T}_{0,0}$, to enable fusing more SpMM iterations, benefiting
from coarse tile fusion shown in Figure 1. The tiles contain a variable number
of iterations to ensure the memory accesses of the tile remain local to the
fast memory. Also, both levels have three independent workloads for all three
cores. As a result of tile fusion, the performance of GeMM-SpMM for a subset
of SuiteSparse matrices on a 20-core processor is faster than atomic tiling,
overlapped tiling, and unfused code with a geometric mean of 13.6$\times$,
3.5$\times$, and 1.64$\times$, respectively.

##  3\. Tile Fusion

Tile fusion looks into the sparsity pattern of the input matrix $A$ in GeMM-
SpMM or SpMM-SpMM and creates a fused schedule. The tile fusion approach has
an iteration scheduler and fused code. Both pairs of operations are commonly
used in applications where the sparsity pattern of matrix $A$ and $B$ (when
sparse) remains static during the execution. Therefore the created schedule
will be computed once based on their sparsity and reused for the rest of the
computation. The rest of this section explains how the scheduler computes the
fused schedule and how codes are fused.

1

Input : $G$, $bCol$, $cCol$, $p$, $cacheSize$, $ctSize$

Output : $\mathcal{T}$

/* Step 1: Coarse Tile Fusion */

2 $I\leftarrow range(rows(G))$

3 $J\leftarrow range(cols(G))$

4 if  _$\lceil|I|/ctSize\rceil\geq p$_ then $t\leftarrow ctSize$ else
$t\leftarrow\lceil|I|/p\rceil$

5 $\mathcal{F}\leftarrow(\\{\\},\\{\\})$

6 for  _ $i\in I$_ do

7 $v\leftarrow i/t$

8 $\mathcal{F}_{0,v}\leftarrow\mathcal{F}_{0,v}\cup range(i,i+t)$

9 for  _$j\leftarrow i\quad to\quad i+t\quad\land\quad j\in J$_ do

10 if  _( $i<inEdges(G,j)<i+t$)_ then
$\mathcal{F}_{0,v}\leftarrow\mathcal{F}_{0,v}\cup j$

11 else $\mathcal{F}_{1,v}\leftarrow\mathcal{F}_{1,v}\cup j$

12 $j\leftarrow j+1$

13

14 end for

15 $i\leftarrow i+t$

16

17 end for

18$\mathcal{F}_{1,v}\leftarrow balance(\mathcal{F}_{1,v},t)$

/* Step 2: Fused Tile Splitting */

19 for  _$w\leftarrow 0\;to\;2$_ do

20 for  _$v\leftarrow 0\;to\;|\mathcal{F}_{w}|$_ do

21 if  _$cost(\mathcal{F}_{w,v},bCol,cCol) >cacheSize$ _ then
$\mathcal{T}_{w}\leftarrow\mathcal{T}_{w}\cup
split(\mathcal{F}_{w,v},bCol,cCol,cacheSize)$

22 else $\mathcal{T}_{w}\leftarrow\mathcal{T}_{w}\cup\mathcal{F}_{w,v}$)

23 $v\leftarrow v+1$

24

25 end for

26 $w\leftarrow w+1$

27

28 end for

Algorithm 1 Tile Fusion Scheduler

###  3.1. Scheduler

The tile fusion scheduler is shown in Algorithm 1 where it creates a schedule
of fused tiles based on sparsity pattern at runtime. This subsection explains
inputs, output, objective, and the two steps of the algorithm.

#### Inputs and output

The tile fusion scheduler takes the DAG $G$, number of columns of $B$ and $C$
as $bCol$ and $cCol$, architecture-specific information $p$ and $cacheSize$,
and heuristic parameter $ctSize$ as inputs and creates a set of fused tiles
$\mathcal{T}$. DAG $G$ represents dependence between $n$ iterations of the two
fused loops where $G_{i,j}=1$ shows iteration $j$ of the second loop depends
on the iteration $i$ of the first loop. $p$ represents the number of physical
cores and $cacheSize$ shows the total size of caches per core in the target
architecture. Each tile in the output fused tile $\mathcal{T}$ is shown with
$\mathcal{T}_{w,v}$ where $w$ and $v$ are wavefront number and tile number,
respectively. Wavefront is a set of iterations that can execute in any order
without violating correctness.

#### Objective and Constraints

The objective of the tile fusion scheduler is to maximize the fused ratio
across all tiles $\mathcal{T}$ while tiles fit into fast memory and only two
wavefronts are allowed. The fused ratio is computed as total iterations of the
second computation in the first wavefront over total iterations as shown in
Equation 2:

(2) $fused\\_ratio=\frac{\sum_{v=0}^{|\mathcal{T}_{0}|}|J_{0v}|}{|I|+|J|}$

where $J_{wv}$ represents the list of iterations from the second operation in
tile $\mathcal{T}_{wv}$ and $I$ and $J$ shows the list of all iterations (or
iteration space) of the first and second operations, respectively.
$\mathcal{T}_{0}$ is the set of tiles in the first wavefront. Operator $|.|$
shows the cardinal number for a set or size of a list. The tile fusion
scheduler maximizes fused ratio under two constraints, load balance constraint
and locality constraint. The load balance constraint ensures the schedule has
a maximum of two synchronization barriers (or two wavefronts) and the number
of tiles in each wavefront is larger than the number of cores, i.e., $\forall
0\leq w<2;\quad|\mathcal{T}_{w}|\geq p$. The locality constraint ensures the
data movement cost for a tile is smaller than cache sizes for each core
($cacheSize$). In other words the $\forall w,v;\quad
cost(\mathcal{T}_{w,v})<cacheSize$ where $cost(\mathcal{T}_{w,v})$ is data
movement cost of $T_{wv}$.

Figure 2c shows an example DAG $G$. There is an edge from the first iteration
of GeMM to the second iteration of SpMM thus $G_{1,2}=1$. The output schedule
in Figure 2f has two tiles. In tile $\mathcal{T}_{0,1}=\\{5,6,6\\}$,
$\\{5,6\\}\in I_{0,1}$ and $\\{6\\}\in J_{0,1}$.

####  3.1.1. Step 1

The first step of the tile fusion scheduler creates an intermediate fused
schedule $\mathcal{F}$, composed of uniform coarse fused tiles to maximize the
fused ratio while ensuring the load balance constraint. The scheduler first
finds fused iterations from tiles of consecutive iterations to improve spatial
locality and reduce the scheduler overhead. The scheduler also ensures
iterations in different tiles of a wavefront are independent, no
synchronization is needed.

Lines 1–1 in Algorithm 1 shows how the intermediate fused tiling $\mathcal{F}$
is created. The scheduler first computes the uniform tile size of $t$ using
the given coarse tile size $ctSize$ in line 1. As shown, the tile size is
chosen to be $ctSize$ if the number of tiles, i.e., $\lceil|I|/ctSize\rceil$
is larger than or equal to $p$ otherwise, it defines $t=|I|/p$. This ensures
the number of tiles in each wavefront is larger than $p$, i.e., the load
balance constraint. Each fused tile $\mathcal{F}_{0,k}$ is created from $t$
consecutive iterations of $I$ as shown in line 1 and some of $t$ consecutive
iterations of $J$ as shown in line 1–1. An iteration of $J$ is added to tile
$\mathcal{F}_{0,k}$ if all of its incoming edges are already in the tile as
shown in line 1. Iterations that do not satisfy the criteria in line 1 are
added to tile $\mathcal{F}_{1,k}$ in the second wavefront as shown in line 1.
The iterations in the second wavefront, $\mathcal{F}_{1}$ is evenly
distributed into $t$ tiles using the $balance$ routine in line 1 to ensure
load balance in the second wavefront.

Figure 3. Tile fusion schedule after step 1 Figure 4. Variation of fused ratio
versus tile size.

The coarse tile size parameter, $ctSize$ used for specifying $t$ in line 1 in
Algorithm 1, is determined heuristically. To select the best value for
$ctSize$, we compute how the fused ratio changes when tile size increases.
Figure 4 shows tile size changes on the x-axis and the average of fused ratio
changes for all matrices of the SuiteSparse repository on the y-axis. The
value of $ctSize$ should be selected to maximize the tile fusion objective.
Since after $ctSize=2048$ in Figure 4, the rate of fused ratio improvement is
slowed down, we use $ctSize=2048$. While going beyond this value can slightly
increase the fused ratio, it reduces the number of tiles in a wavefront,
potentially leading to load imbalance.

Figure 3 shows the output of step 1 for the example shown in Figure 2. For
this example, we assume $ctSize=4$ and $p=3$ which makes a tile size $t=4$.
Two coarse tile size is shown in Figure 3 $\mathcal{F}_{0,0}={1,2,3,4}$ and
then since iterations ${1,2,3}\in J$ depend on iterations ${1,2,3}\in I$ and
they already exist in $\mathcal{F}_{0,0}$, then the three iterations are
added.

####  3.1.2. Step 2

The second step of the tile fusion scheduler splits coarse tiles created in
the first step, $\mathcal{F}$, to fit them into fast memory. As a result,
tiles with different sizes are created to ensure the locality constraint of
the tile fusion. The scheduler iterates over $\mathcal{F}$ and measures their
data movement with a data movement cost model. Fused tiles whose data movement
is larger than the size of fast memory are split to fit into the fast memory.

The second step of the tile fusion scheduler is shown in Lines 1–1 in
Algorithm 1. The algorithm iterates over all tiles in the two wavefronts
$\mathcal{F}_{0}$ and $\mathcal{F}_{1}$ and computes the data movement cost of
each tile using a $cost$ function as shown in line 1. If the data movement
cost of a tile $\mathcal{F}_{i,j}$ is larger than the size of fast memory
$cacheSize$, the scheduler splits the tile recursively to create a set of
tiles, each fitting into the fast memory using the $split$ routine as shown in
line 1. The resulting split tiles are added to the same wavefront of the final
schedule $\mathcal{T}$.

#### Data movement cost

The tile fusion scheduler relies on a cost model to approximate data movement
for a coarse tile. As shown in line 1 in Algorithm 1, the data movement cost
model computes the cost of a tile $\mathcal{T}_{i,j}$ for a given $bCol$ and
$cCol$. The cost model is computed as shown in Equation 3:

(3) $cost(\mathcal{T}_{i,j},bCol,cCol)=\\\
(nz(\mathcal{T}_{i,j})+uc(\mathcal{T}_{i,j})+t+|J_{i,j}|)*cCol+idx$

where, $nz(\mathcal{T}_{i,j})$ is the number of unique nonzeros in the tile
from $A$ and $B$. When $B$ is dense, all $n\times bCol$ is added.
$uc(\mathcal{T}_{i,j})$ is the number of nonzeros with unique columns in the
tile, $|J_{i,j}|$ is the number of fused iterations from the second operation
and $idx$ is the indexing cost for when $A$ or $B$ are sparse.

The final fused schedule $\mathcal{T}$ in Figure 2f is computed from the
coarse fused tile schedule in $\mathcal{F}$ in Figure 3. The schedule has two
coarse fused tiles. Each tile is labeled with its data communication cost.
Assuming $bCol=cCol=1,t=4$, the cost of $\mathcal{F}_{0,0}$ is computed based
on Equation 3. Since the cost of $\mathcal{F}_{0,0}$ is less than
$cacheSize=30$, the tile will be directly added to $\mathcal{T}$. However,
$\mathcal{F}_{0,1}$ is larger than $cacheSize$ and is split in two smaller
tiles that has a cost less than $cacheSize$.

#### Computational Complexity

The first step of the algorithm, for every tile with size $t$, it checks
$inEdges$ for only $t$ columns of $G$ in the same range. Since tiles are
disjoint, $inEdges$ is called once per iteration, makes the first step of the
algorithm $O(nnz)$. Accessing incoming edges for an iteration is possible in
linear time. The second step only iterates over fused iterations in $J$ and
their dependent iterations from $I$ for splitting. Since the set of iterations
is split by a factor of 2, and each split function can visit up to $nnz$
edges, therefore, its complexity is $O(nnz*log(ctSize)$. The second wavefront
only operates on unfused iterations which take $O(|J|)$. Therefore the
complexity of the second step is $O(|J|+nnz*log(ctSize)$.

###  3.2. Fused Code

The fused code is created by fusing the outermost loop of the two operations.
The fused code is then replaced with a doubly nested loop that iterates over
the tiles in parallel using the OpenMP scheduler. The fused code uses the
fused tiling schedule $\mathcal{T}$ and maps iterations of fused tiles to
their associated code version to ensure correctness. Listing 1 and Listing 3
show the fused code for GEMM-SpMM (Figure 2) and SpMM-SpMM (Listing 2). As
shown the fused code is composed of a version of the two operations. Lines 4–7
and 8–11 in Listing 1 correspond to the innermost loops of GEMM and SpMM
respectively. Lines 3–7 and 8–11 in Listing 3 correspond to the innermost
loops of the first SpMM ($D_{1}=AC$) and the second SpMM ($D=AD_{1}$)
respectively. Loop-bounds in line 4 in Listing 1 and line 3 in Listing 3 are
determined by the schedule $\mathcal{T}$, playing as a runtime check to switch
between the two versions. Therefore, the fused code follows the schedule order
and ensures all dependence between fused iterations. The outermost loop of the
two computations is fused and replaced with a pair of loops that go over the
tile fusion schedule.

Fused code also turns data reuse between fused iterations into temporal
locality while also preserving the locality of each operation. The code
versions are next to each other in the code and when they execute right after
each other, the data reuse between them turns into temporal locality. For the
first tile in the schedule shown in Figure 2, the arrays corresponding to rows
1–4 stay in the cache when the fused code in Listing 1 switches to the SpMM
loop in line 8, thus improving temporal locality. Both fused codes in Listings
1 and 3 benefit from the fact that consecutive iterations are grouped in the
scheduler. The consecutive choice will eliminate the need for conditional
checking for every iteration of computations. It also improves spatial
locality and temporal locality exist within iterations of each operation, e.g.
spatial and temporal locality in GeMM will be in place in fused tiles.

Listing 1: Fused code of GeMM-SpMM, $D=A(BC)$.

⬇

1for (w in T){ // for each wavefront

2#pragma omp parallel for

3 for(t in T[w]){ // for each tile

4 for(i1 in t.first)

5 for(int i2=0; i2<cCol; ++i2)

6 for(int i3=0; i3<bCol; ++i3)

7 D1[i3][i2] += B[i1][i3] * C[i2][i3];

8 for(j1 in t.second)

9 for(int j2=A.p[j1]; j2<A.p[j1+1]; j2++)

10 for(int j3=0; j3<cCol; j3++)

11 D[j1][l2] += A.x[j2] * D1[A.i[j2]][j3];

12 }}

Listing 2:  $D=A(AC)$ where $A$ is CSR.

⬇

1fuse: for(int i1=0; i1<n; ++i1) // SpMM 1

2 for(int i2=A.p[i1]; i2<A.p[i1+1]; ++j1)

3 for(int i3=0; i3<cCol; ++i3)

4 D1[i1][i3] += A.x[i2] * C[A.i[i2]][i3];

5fuse: for(int j1=0; j1<n; ++j1) // SpMM 2

6 for(int j2=A.p[j1]; j2<A.p[j1+1]; ++j2)

7 for(int j3=0; j3<cCol; ++j3)

8 D[j1][j3] += A.x[j2] * D1[A.i[j2]][j3];

Listing 3:  Fused code for SpMM-SpMM, $D=A(AC)$.

⬇

1for (w in T){

2#pragma omp parallel for

3 for(t in T[w]){

4 for(i1 in t.first)

5 for(int i2=A.p[i1]; i2<A.p[i1+1]; ++i2)

6 for (int i3 = 0; i3<cCol; ++i3)

7 D1[i1][i3] += A.x[i2] * C[A.i[i2]][i3];

8 for(j1 in t.second)

9 for(int j2 = A.p[j1]; j2<A.p[j1+1]; j2++)

10 for(int j3 = 0; j3<cCol; j3++)

11 D[j1][j3] += A.x[j2] * D1[A.i[j2]][j3];

12 }}

The fused code enables thread-level parallelism by mapping fused tiles to
parallel threads. Mapping tiles to threads is done using omp scheduler as
shown in line 2 in Listing 1 and line 2 in Listing 3. The fused code will keep
all fine-grain parallelism opportunities such as vectorization that exist in
the unfused code. For example, lines 4–7 in Listing 1 is mapped to a highly
optimized GEMM BLAS to benefit from vector processors. Similarly, inner loop
vectorization is performed in lines 6 and 10 in Listing 3.

##  4\. Experimental Results

This section discusses the performance of tile fusion with existing fused and
unfused implementations for sparse matrices across two different shared memory
processors. Overall tile fusion is faster than unfused and best-fused code
with a geomean speedup of 1.98$\times$ and 3.52$\times$ and is scalable to 40
and 64 cores.

###  4.1. Setup

Platform | CascadeLake | EPYC
---|---|---
# of sockets $\times$ cores |  2 $\times$ 20 cores |  2 $\times$ 32 cores
L1/L2/L3 Cache Sizes | 32K/1M/28M | 32K/512K/256M
Compiler | ICC 2022 | GCC v.11
BLAS |  MKL BLAS 2022 (Wang et al., 2014) |  BLIS (Van Zee and Van De Geijn, 2015)
Table 1. Platform details

####  4.1.1. Environment

All experiments are executed on multi-core processors shown in Table 1 to show
the performance of tile fusion cross-platform. The experiments are done in
single-node except stated otherwise. Since both computations are used in GNN
(Zhou et al., 2020) and sparse iterative linear solvers with multiple right-
hand side (Aggarwal et al., 2021), we test tile fusion for double-precision
(DP) and single-precision (SP) data types, commonly used in machine learning
and scientific computing respectively. All experiments are tested for three
different numbers of columns, 32, 64, and 128 in $B$ to enable testing
different matrix sizes $B$ and $C$. Each reported time in the paper is the
median of 7 runs. For each matrix, the theoretical FLOPs for the unfused code
is computed and used for all implementations. Parameter $cacheSize$ in
Algorithm 1 is set to the sum of L1+L2+(L3/cores) cache sizes in row 2 in
Table 1. A close thread binding is selected and each thread is pinned to a
physical core.

####  4.1.2. Matrix Dataset

We select 233 matrices from SuiteSparse (Davis and Hu, 2011) matrix repository
to evaluate GeMM-SpMM and SpMM-SpMM computations. We select two groups of
matrices to represent scientific computing and machine learning applications.
I. all 132 symmetric positive definite matrices larger than $10^{5}$ nonzero
values. II. all 111 square matrices related to graph applications larger than
$10^{5}$ nonzeros with either integer or real types. We define a matrix as
graph-related if a“graph” keyword is included in its metadata.

Figure 5. GeMM-SpMM performance for all matrices on CascadeLAke (top) EPYC
(bottom)

####  4.1.3. Fused and Unfused Implementations

To compare with unfused implementations, we use OneAPI MKL (Wang et al., 2014)
library v2022. Since fused versions of GEMM or SpMM are not supported in MKL,
we called routines cblas_?gemm and mkl_sparse_?_mm for double/single precision
of GeMM and SpMM, respectively. We set the number of threads in MKL with
mkl_set_num_threads() to the number of physical cores. We also develop an
unfused parallel implementation for both GeMM-SpMM and SpMM-SpMM with the same
set of optimizations to show the effect of tile fusion.

Tile fusion is compared with existing fused implementations and some in-house
implementations of prior fusion techniques. For GeMM-SpMM, we use the
generated C++ code from TACO (Kjolstad et al., 2017) and SparseLNR (Dias et
al., 2022) tensor compilers for the expression D(i,l) = A(i,j) * B(j,k) *
C(k,l) where $A$ is sparse and other matrices are dense. We added both
generated codes and reported their best timing as Best of Tensor Compilers. We
also additionally vectorize the generated tensor compiler code by using MKL
GeMV BLAS (Wang et al., 2014) to ensure the effect of vectorization is
incorporated.

For SpMM-SpMM, tensor compilers do not support fusing SpMM-SpMM and, therefore
excluded from the benchmark. Since the code for communication-avoiding (Demmel
et al., 2008) and sparse tiling (Krieger et al., 2013) are not publicly
available, per the authors’ recommendation, we adopted the idea and applied
them to SpMM-SpMM. For communication-avoiding methods, we first equally
partition iterations of the first SpMV, and then we add all dependent
iterations to the same partitions. This implementation is called overlapped
tiling. For sparse tiling, we partition iterations of the first SpMM equally.
Then we add dependent iterations of the second SpMM to each partition to
create balanced partitions. Finally, dependence is resolved with atomic and
barriers. This implementation is named atomic tiling.

We only report the fused code execution time for all experiments operating on
matrices. The scheduler overhead is separately evaluated and illustrated.

|  |  | 32 | 64 | 128
---|---|---|---|---|---
CascadeLake | Single Precision | MKL | 1.64 | 1.41 | 1.36
UnFused | 1.36 | 1.24 | 1.14
Double Precision | MKL | 1.37 | 1.33 | 1.23
UnFused | 1.45 | 1.34 | 1.24
EPYC | Single Precision | UnFused | 1.67 | 1.73 | 1.84
Double Precision | UnFused | 1.81 | 1.93 | 1.97
Table 2. The summary of gmean of speedups for GeMM-SpMM. Figure 6. GeMM-SpMM
performance of fused implementations on graph matrices.

###  4.2. GEMM-SpMM Evaluation

####  4.2.1. Performance in FLOPs

Figure 5 shows the overall performance of GeMM-SpMM using tile fusion with
unfused MKL for the two architectures and three bcols. As shown tile fusion is
faster than MKL for 90% of matrices across bclos. Table 2 shows speedup
details for GeMM-SpMM and for single and double precision for the target
architectures shown in Table 1.

The performance of tile fusion increases as bCols increase due to increasing
arithmetic intensity. The tile fusion performance increases from a mean of 152
GFLOP/s when bCol=32 to 328 GFLOP/s when bCol=128. While MKL implementation
changes from 92 GLOP/s to 241 GFLOP/s when bCols changes from 32 to 128. As
bCols increase, the arithmetic intensity of fused tiles increases and tile
fusion can take advantage.

All implementations have a better performance for SPD matrices than graph
matrices. The reason is that the fused ratio in SPD matrices is on average 2
times higher than graph matrices. The performance of Tile Fusion for single
precision is 2$\times$ better than double precision. When operating on double,
the data movement increases, making computation more memory-bound than single,
thus reducing GFLOP/s. Also, since the EPYC processor has a larger L3 cache,
the performance gap between tile fusion and unfused baseline for large
matrices is higher than the CascadeLake processor.

Tile fusion also supports fusing Equation 1 when the transpose of $C$ should
be multiplied. Tile fusion provides a geometric mean of 1.49, 1.24, and 1.26
over unfused MKL on CascadeLake for bCol=cCol=32, 64, 128, respectively.

Figure 6 shows the performance of tile fusion compared to other fused
implementations. Tile fusion is faster than tensor compilers, atomic tiling,
and overlapped tiling with an average speedup of 9.4$\times$, 13.6$\times$,
and 3.5$\times$, respectively. Tensor compilers perform redundant computations
and also do not use memory hierarchy due to vector operations.

####  4.2.2. Ablation Study

This section analyzes the effect of tile fusion on locality and load balance
and the effect of the two steps of the tile fusion scheduler on the
performance. We selected all 111 graph matrices, a subset of the matrix
dataset for profiling and analysis. All analysis is also done on the
CascadeLake target architecture.

Figure 7. The effect of average memory access time on performance of tile
fusion for GeMM-SpMM

We measure an average memory cycle to analyze the effect of tile fusion on
improving locality in GeMM-SpMM. We measure average memory access time (AMT)
as AMT = hit time + miss ratio * miss penalty for all three levels of caches
in the target architecture. We use PAPI (Terpstra et al., 2010) performance
counters, PAPI_L1_TCM, PAPI_L2_TCM, PAPI_L3_TCM to measure L1 accesses, L2
accesses, L3 accesses, and main memory accesses, respectively to compute hit
and miss ratio for each level. Average memory access times for the selected
subset of matrices are shown in Figure 7. As shown, tile fusion improves AMT
for 92% of graph matrices between 1.1-1.3$\times$ compared to the unfused
implementation which is the main cause for improving the performance.

Figure 8. Potential Gain of GeMM-SpMM (lower is better)

We measure the potential gain of both fused and unfused code to show the
effect of tile fusion on the load balance of GeMM-SpMM. Potential gain (PG) is
defined as the maximum time that can be saved if all threads are balanced. We
measure the average difference between the maximum time of threads and other
threads’ time. We use PAPI counter PAPI_TOT_CYC to measure the number of
cycles for each thread. Figure 8 shows the PG compared to unfused. As shown,
the tile fusion load balance is close to unfused. The unfused code has a
larger number of fine-grain tasks, enabling it to be more balanced.

Figure 9. Effect of different steps of Tile Fusion in GeMM-SpMM for graph
matrices

Figure 9 shows the performance breakdown of the two steps of the tile fusion
inspector. As shown, the first step of tile fusion improves the performance of
sequential baseline code with a gmean speedup of 6.7$\times$. The second step
of tile fusion contributes to the performance of 90% of matrices shown in
Figure. This first step contributes more because it adds threading and
improves locality. The second step further balances the loads and improves the
parallel workloads of step 1. The second step selects tile sizes based on the
cost model provided in Equation 3. For the selected graph matrices, the tile
sizes selected by the second step vary between 64-2048.

Figure 10. Number of runs needed to amortize the scheduling cost for GeMM-SpMM
(lower and positive values are better)
---
Figure 11. SpMM-SpMM performance on CascadeLake (top) and EPYC (bottom) Figure 12. SpMM-SpMM performance of fused implementations |  |  | 32 | 64 | 128
---|---|---|---|---|---
CascadeLake | Single Precision | MKL | 1.2 | 1.02 | 1.11
UnFused | 1.17 | 1.15 | 1.14
Double Precision | MKL | 1.09 | 1.16 | 1.11
UnFused | 1.14 | 1.15 | 1.13
EPYC | Single Precision | UnFused | 1.14 | 1.17 | 1.19
Double Precision | UnFused | 1.14 | 1.20 | 1.22
Table 3. G-mean speedups for SpMM-SpMM

####  4.2.3. Scheduler Overhead analysis

The tile fusion performs scheduling once per sparsity pattern and can be
reused as long as sparsity remain static. Figure 10 shows the number of
iterations that fused code should run to amortize the scheduler overhead with
respect to the fastest baselines. The number of fused code runs is computed as
$\frac{scheduler\,time}{baseline-fusedCode\,time}$. As shown, tile fusion
needs less than 100 iterations to amortize the cost of the scheduler. In many
applications such as GNN training, GeMM-SpMM is called over hundreds or
thousands of times.

###  4.3. SpMM-SpMM Evaluation

The performance of tile fusion is compared with unfused implementations for
SpMM-SpMM as shown in Figure 11. Tile fusion is faster than unfused baseline
and MKL implementations in 100% and 70% of all matrices in any bCol that we
experimented on and for SP/DP. The detailed speedup for both CascadeLake and
EPYC and SP and DP are illustrated in Table 3. The performance of SpMM-SpMM is
overall lower than GeMM-SpMM for the same set of matrices due to the memory-
bound nature of SpMM.

Tile fusion provides a gmean speedup of 9.3$\times$, 13.2$\times$, and
13.7$\times$ over atomic tiling for bCol = 32, 64, and 128 respectively. A
similar trend exists for overlapped tiling where tile fusion provides a gmean
speedup of 5, 6.5, and 7.2 for bcols=32, 64, and 128. The main reason is the
amount of redundant computation that increases for overlapped tiles. For
example, matrix G2_circuit and inline_1 have redundant iterations of 126487
and 2844351 respectively while they only have 150102 and 503712 rows.

##  5\. Related Work

#### Tiling and Fusion for Sparse Codes

Loop tiling and fusion are common techniques to improve locality. A large body
of work applies these techniques to SpMM (Natesh et al., 2023; Hong et al.,
2019; Kurt et al., 2020; Ahmad et al., 2024; Cheshmi et al., 2022a; Wilkinson
et al., 2023) or GeMM (Wang et al., 2014; Kung et al., 2021) routines. Tile
fusion preserve locality between iterations of GeMM and SpMM thus enables
benefiting from existing developed efficient GeMM and SpMM. Using data
communication cost to define tile size is used in matrix signature (Kurt et
al., 2020) for SpMM CSR. Tile fusion however works on two operations that
needs to take into account common elements between the two operations. Fusing
based on compile-time information is commonly used in dense matrices. Indirect
memory accesses reduce fusion opportunities however, there are still
opportunities. Tensor expression compilers (Mutlu et al., 2022; Dias et al.,
2022; Kjolstad et al., 2017) generates code for a tensor expression. Sparse
tensor compilers (Mutlu et al., 2022; Dias et al., 2022; Kjolstad et al.,
2017) support fusing Equation 1. However, the generated fused code from
SparseLNR (Dias et al., 2022) and TACO (Kjolstad et al., 2017) make matrix
operations in Equation 1 to matrix-vector operations which do not efficiently
use the fast memory.

Modeling parallel loops such consecutive matrix multiplication as a graph
(Krieger et al., 2013; Strout et al., 2004; Demmel et al., 2008) or hypergraph
(Pawłowski et al., 2020) are commonly used for improve cache reuse. These
methods for shared memory processors either rely on synchronization (Krieger
et al., 2013) or overlapped computation (Demmel et al., 2008). The fused ratio
and its cost model can enhance locality and reduce overlapped computation in
these methods. Sympiler (Cheshmi, 2022; Cheshmi et al., 2017) uses DAG
schedulers (Cheshmi et al., 2018; Zarebavani et al., 2022) to build an initial
schedule of iteration then fuse the schedule with another loop using sparse
fusion (Cheshmi et al., 2023, 2022b). Sparse fusion scheduler is driven by
loop-carried dependency commonly occurs in scientific solvers (Cheshmi et al.,
2020) to ensure locality which does not exist in matrix multiplication
operations.

##  6\. Conclusion

This paper presents tile fusion to enable fusing GeMM-SpMM and SpMM-SpMM for
$D=A\times B\times C$ where $A$ is sparse and $B$, $C$, and $D$ are dense.
Tile fusion has a scheduler that builds a schedule of fused tiles based on the
sparsity pattern of the $A$ and the size of dense matrices. The created
schedule does not use redundant computation and its synchronizations are
always 2. Tile fusion outperforms existing unfused and fused implementations.

## References

  * (1)
  * Aggarwal et al. (2021) Isha Aggarwal, Aditya Kashi, Pratik Nayak, Cody J Balos, Carol S Woodward, and Hartwig Anzt. 2021.  Batched sparse iterative solvers for computational chemistry simulations on GPUs. In _2021 12th Workshop on Latest Advances in Scalable Algorithms for Large-Scale Systems (ScalA)_. IEEE, 35–43.
  * Ahmad et al. (2024) Khalid Ahmad, Cris Cecka, Michael Garland, and Mary Hall. 2024.  Exploring data layout for sparse tensor times dense matrix on GPUs.  _ACM Transactions on Architecture and Code Optimization_ 21, 1 (2024), 1–20.
  * Cheshmi (2022) Kazem Cheshmi. 2022.  _Transforming Sparse Matrix Computations_.  Ph. D. Dissertation. University of Toronto (Canada).
  * Cheshmi et al. (2022a) Kazem Cheshmi, Zachary Cetinic, and Maryam Mehri Dehnavi. 2022a.  Vectorizing sparse matrix computations with partially-strided codelets. In _SC22: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 1–15.
  * 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_. 1–13.
  * Cheshmi et al. (2018) Kazem Cheshmi, Shoaib Kamil, Michelle Mills Strout, and Maryam Mehri Dehnavi. 2018.  ParSy: inspection and transformation of sparse matrix computations for parallelism. In _SC18: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 779–793.
  * Cheshmi et al. (2020) Kazem Cheshmi, Danny M Kaufman, Shoaib Kamil, and Maryam Mehri Dehnavi. 2020.  NASOQ: numerically accurate sparsity-oriented QP solver.  _ACM Transactions on Graphics (TOG)_ 39, 4 (2020), 96–1.
  * Cheshmi et al. (2023) Kazem Cheshmi, Michelle Strout, and Maryam Mehri Dehnavi. 2023.  Runtime composition of iterations for fusing loop-carried sparse dependence. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_. 1–15.
  * Cheshmi et al. (2022b) Kazem Cheshmi, Michelle Mills Strout, and Maryam Mehri Dehnavi. 2022b.  Optimizing sparse computations jointly. In _Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming_. 459–460.
  * Davis and Hu (2011) Timothy A Davis and Yifan Hu. 2011.  The University of Florida sparse matrix collection.  _ACM Transactions on Mathematical Software (TOMS)_ 38, 1 (2011), 1–25.
  * Demmel et al. (2008) James Demmel, Mark Hoemmen, Marghoob Mohiyuddin, and Katherine Yelick. 2008.  Avoiding communication in sparse matrix computations. In _2008 IEEE International Symposium on Parallel and Distributed Processing_. IEEE, 1–12.
  * Dias et al. (2022) Adhitha Dias, Kirshanthan Sundararajah, Charitha Saumya, and Milind Kulkarni. 2022.  SparseLNR: accelerating sparse tensor computations using loop nest restructuring. In _Proceedings of the 36th ACM International Conference on Supercomputing_. 1–14.
  * Fey and Lenssen (2019) Matthias Fey and Jan Eric Lenssen. 2019.  Fast graph representation learning with PyTorch Geometric.  _arXiv preprint arXiv:1903.02428_ (2019).
  * 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_. 300–314.
  * Kipf and Welling (2016) Thomas N Kipf and Max Welling. 2016.  Semi-supervised classification with graph convolutional networks.  _arXiv preprint arXiv:1609.02907_ (2016).
  * Kjolstad et al. (2017) Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. 2017.  The tensor algebra compiler.  _Proceedings of the ACM on Programming Languages_ 1, OOPSLA (2017), 1–29.
  * Krieger et al. (2013) Christopher D Krieger, Michelle Mills Strout, Catherine Olschanowsky, Andrew Stone, Stephen Guzik, Xinfeng Gao, Carlo Bertolli, Paul HJ Kelly, Gihan Mudalige, Brian Van Straalen, et al. 2013\.  Loop chaining: A programming abstraction for balancing locality and parallelism. In _2013 IEEE International Symposium on Parallel & Distributed Processing, Workshops and Phd Forum_. IEEE, 375–384.
  * Kung et al. (2021) HT Kung, Vikas Natesh, and Andrew Sabot. 2021.  Cake: matrix multiplication using constant-bandwidth blocks. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_. 1–14.
  * Kurt et al. (2020) Süreyya Emre Kurt, Aravind Sukumaran-Rajam, Fabrice Rastello, and Ponnuswamy Sadayyapan. 2020.  Efficient tiled sparse matrix multiplication through matrix signatures. In _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 1–14.
  * Mutlu et al. (2022) Erdal Mutlu, Ruiqin Tian, Bin Ren, Sriram Krishnamoorthy, Roberto Gioiosa, Jacques Pienaar, and Gokcen” Kestor. 2022.  COMET: A Domain-Specific Compilation of High-Performance Computational Chemistry. In _Languages and Compilers for Parallel Computing_ , Barbara Chapman and José Moreira (Eds.). Springer International Publishing, Cham, 87–103.
  * Natesh et al. (2023) Vikas Natesh, Andrew Sabot, HT Kung, and Mark Ting. 2023.  Rosko: Row Skipping Outer Products for Sparse Matrix Multiplication Kernels.  _arXiv preprint arXiv:2307.03930_ (2023).
  * O’Leary (1980) Dianne P O’Leary. 1980.  The block conjugate gradient algorithm and related methods.  _Linear algebra and its applications_ 29 (1980), 293–322.
  * Pawłowski et al. (2020) Filip Pawłowski, Rob H Bisseling, Bora Uçar, and AN Yzelman. 2020.  Combinatorial Tiling for Sparse Neural Networks. In _2020 IEEE High Performance Extreme Computing Conference (HPEC)_. IEEE, 1–7.
  * Strout et al. (2004) Michelle Mills Strout, Larry Carter, Jeanne Ferrante, and Barbara Kreaseck. 2004.  Sparse tiling for stationary iterative methods.  _The International Journal of High Performance Computing Applications_ 18, 1 (2004), 95–113.
  * Terpstra et al. (2010) Dan Terpstra, Heike Jagode, Haihang You, and Jack Dongarra. 2010.  Collecting performance data with PAPI-C. In _Tools for High Performance Computing 2009: Proceedings of the 3rd International Workshop on Parallel Tools for High Performance Computing, September 2009, ZIH, Dresden_. Springer, 157–173.
  * Van Zee and Van De Geijn (2015) Field G Van Zee and Robert A Van De Geijn. 2015.  BLIS: A framework for rapidly instantiating BLAS functionality.  _ACM Transactions on Mathematical Software (TOMS)_ 41, 3 (2015), 1–33.
  * 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. (2019) Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, et al. 2019\.  Deep graph library: A graph-centric, highly-performant package for graph neural networks.  _arXiv preprint arXiv:1909.01315_ (2019).
  * Wilkinson et al. (2023) Lucas Wilkinson, Kazem Cheshmi, and Maryam Mehri Dehnavi. 2023.  Register Tiling for Unstructured Sparsity in Neural Network Inference.  _Proceedings of the ACM on Programming Languages_ 7, PLDI (2023), 1995–2020.
  * Zarebavani et al. (2022) Behrooz Zarebavani, Kazem Cheshmi, Bangtian Liu, Michelle Mills Strout, and Maryam Mehri Dehnavi. 2022.  HDagg: hybrid aggregation of loop-carried dependence iterations in sparse matrix computations. In _2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS)_. IEEE, 1217–1227.
  * Zhou et al. (2020) Jie Zhou, Ganqu Cui, Shengding Hu, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. 2020.  Graph neural networks: A review of methods and applications.  _AI open_ 1 (2020), 57–81.