# Estimation-Energy Tradeoff for Scalar Gauss-Markov Signals with Kalman
Filtering

Ioannis Krikidis, and Constantinos Psomas I. Krikidis and C. Psomas are with
the Department of Electrical and Computer Engineering, University of Cyprus,
Cyprus (email: {krikidis, psomas}@ucy.ac.cy). This work has received funding
from the European Research Council (ERC) under the European Union’s Horizon
2020 research and innovation programme (Grant agreement No. 819819).

###### Abstract

In this letter, we investigate a receiver architecture, which uses the
received signal in order to simultaneously harvest energy and estimate a
Gauss-Markov linear process. We study three communication scenarios: i) static
channel, ii) Rayleigh block-fading channel, and iii) high power amplifier
(HPA) nonlinearities at the transmitter side. Theoretical results for the
minimum mean square error as well as the average harvested energy are given
for all cases and the fundamental tradeoff between estimation quality and
harvested energy is characterized. We show that channel fading improves the
estimation performance while HPA requires an extended Kalman filter at the
receiver and significantly affects both the estimation and the harvesting
efficiency.

###### Index Terms:

Wireless power transfer, power-splitting, Kalman filter, extended Kalman
filter, fading channel, estimation.

##  I Introduction

Wireless power transfer (WPT) is an enabling technology for future cyber-
physical systems, where a massive number of devices (e.g., sensors, RFID tags,
actuators, etc) with low-rate and low-power requirements are interconnected in
order to monitor and control critical infrastructures. It can provide energy
sustainability, battery-less implementation and therefore it seems as an
essential technology for the next generation of cyber-physical systems that
require continuous estimation/monitoring of a physical process. The
integration of WPT in wireless communication systems has been studied
extensively over the last years and the fundamental tradeoff between
information and energy transfer has been studied from information theoretic
[1, 2], signal processing and networking perspectives [3]. However, the
interplay of WPT with conventional parameter/signal estimation, which is the
key task in the above systems, is still an unexplored area.

Recent works study the integration of WPT in signal estimation problems. The
authors in [4] investigate a distributed estimation problem, where sensors
harvest energy from radio-frequency signals and use it, in order to transmit
their observations to a fusion center; the parameters of WPT are designed to
minimize the mean-square error (MSE) of the final estimate. This work has been
extended in [5] for scenarios with multiple-antenna sensor nodes, while the
work in [6] studies the problem from a macroscopic point-of-view by taking
into account the spatial randomness of the sensor network. All the
aforementioned works decouple the harvesting operation from the signal
estimation process and energy harvesting is used in order to allow sensors to
communicate their observations. Another topology with practical interest is
the simultaneous parameter estimation and energy harvesting from the same
signal; in accordance with the simultaneous information and power transfer
concept, various practical architectures such as the power-splitting (PS)
technique can be used for this purpose [3]. To the best of the authors’
knowledge, this dual use of the received signal (signal/state estimation,
energy harvesting) is still an open problem in the literature.

In this paper, we investigate a PS-based architecture that uses the received
signal in order to simultaneously estimate a scalar Gauss-Markov signal (i.e.,
sequence of states evolving as a dynamic linear model) [7, 8] and harvest
energy. Specifically, we study the use of first-order Gauss-Markov signals for
WPT. We consider three scenarios: i) transmission over a static wireless
channel, ii) transmission over a Rayleigh block-fading channel, and iii) the
case where the transmitter is affected by high power amplifier (HPA)
nonlinearities [2]. Closed-form expressions for the average harvested energy
(linear model) are derived for both finite and asymptotic regimes. A receiver
architecture is also introduced, which enables a simultaneous estimation of
the transmitted signal and energy harvesting, by employing a PS technique. For
the estimation process, the receiver applies a linear Kalman filter on the
corresponding fraction of the signal by taking into account the memory of the
scalar Gauss-Markov signal; for the scenario with HPA, we employ an extended
Kalman filter. Theoretical results are given for the minimum MSE (MMSE) for
all the considered scenarios. The considered architecture is characterized by
a fundamental tradeoff between quality of the estimation (MMSE) and average
harvested energy. Notation: $\mathbb{E}\\{\cdot\\}$ denotes statistical
expectation, $\exp(\lambda)$ denotes the exponential distribution with rate
parameter $\lambda$, $\mathcal{CN}(\mu,\sigma^{2})$ represents the complex
Gaussian distribution with mean $\mu$ and variance $\sigma^{2}$,
$\Re\\{\cdot\\}$ and $\Im\\{\cdot\\}$ indicate real and imaginary operators,
respectively, and $j=\sqrt{-1}$.

##  II Estimation-energy tradeoff

In this section, we introduce the considered receiver architecture and study
the fundamental estimation-energy tradeoff for three scenarios: i) static
channel, ii) Rayleigh block-fading, and iii) transmission with HPA non-
linearities.

###  II-A Gauss-Markov over a static channel

We assume a scalar complex Gauss-Markov signal model[10], given by111The
coefficient $a$ depends on the specific application (e.g., navigation,
robotics, etc.) and determines the correlation between consecutive
measurements; in practical dynamic systems where small fluctuations occur, we
assume $|a|<1$, which results in stable systems. The case $a>1$ corresponds to
unstable open-loop processes where, in most cases, they are associated with a
control system to ensure stability.

State: $\displaystyle x(n)=ax(n-1)+u(n),\;\;n\geq 0,$ (1) Observations:
$\displaystyle y(n)=hx(n)+v(n),$ (2)

where $n$ is the time index $x(-1)\sim\mathcal{CN}(\mu_{0},\sigma_{0}^{2})$,
$u(n)\sim\mathcal{CN}(0,\sigma_{u}^{2})$,
$v(n)\sim\mathcal{CN}(0,\sigma_{v}^{2})$, $|a|<1$, and $h\in\mathbb{C}$
denotes the complex channel coefficient which is assumed constant and known at
the receiver. According to the principles of PS [3], the receiver splits the
received signal $y(n)$ into two components by using a splitting factor $\rho$
i.e., $\sqrt{\rho}y(n)$ is converted to the baseband and is used to estimate
$x(n)$ and $\sqrt{1-\rho}y(n)$ is driven towards the energy harvesting branch.
During the baseband conversion, an additional circuit noise, $q(n)$, is
present which is modeled as additive white Gaussian noise (AWGN) with zero
mean and variance $\sigma_{q}^{2}$. The system model considered is depicted in
Fig. 1 and could refer to wireless sensor networks that monitor the state of a
dynamic physical phenomenon (e.g., air quality monitoring, object tracking,
etc.). The obtained measurements are sent over a fading channel by using
analog forwarding to avoid processing delays [8]. This setup seems interesting
for future cyberphysical systems, where a massive number of sensors/devices
exchange measurements and are characterized by strict computation/delay/energy
constraints.

For data estimation, the equivalent (baseband) observation can be written as

$\displaystyle y^{\prime}(n)=\sqrt{\rho}hx(n)+e(n),$ (3)

where $e(n)=\sqrt{\rho}v(n)+q(n)$ is circularly symmetric complex Gaussian
with variance $\sigma_{e}^{2}=\rho\sigma_{v}^{2}+\sigma_{q}^{2}$. Due to the
data correlation, the estimation of $x(n)$ takes into account the whole
history of observations and is based on the Kalman filter equations222 The
Kalman filter is a time-varying linear filter and achieves linear MMSE in a
sequential way by introducing one new sample per time and by applying trivial
(low-complexity) algebraic computations [10]. [10, 7] i.e.,

$\displaystyle\hat{x}(n|n-1)=a\hat{x}(n-1|n-1),$ (4) $\displaystyle
M(n|n-1)=a^{2}M(n-1|n-1)+\sigma_{u}^{2},$ (5) $\displaystyle
K(n)=\frac{M(n|n-1)\sqrt{\rho}h^{*}}{\sigma_{e}^{2}+\rho|h|^{2}M(n|n-1)},$ (6)
$\displaystyle\hat{x}(n|n)=\hat{x}(n|n-1)+K(n)[y^{\prime}(n)-\sqrt{\rho}h\hat{x}(n|n-1)],$
(7) $\displaystyle M(n|n)=(1-K(n)\sqrt{\rho}h)M(n|n-1)$
$\displaystyle\;\;\;\;\;\;\;\;\;\;\;\;=\frac{a^{2}\sigma_{e}^{2}M(n-1|n-1)+\sigma_{e}^{2}\sigma_{u}^{2}}{\sigma_{e}^{2}+\rho|h|^{2}(a^{2}M(n-1|n-1)+\sigma_{u}^{2})},$
(8)

where $\hat{x}(n|n-1)=\mathbb{E}\\{x(n)|y(0),y(1),\ldots,y(n-1)\\}$ is the
MMSE prediction based on the previous estimated state, (5) represents the
minimum prediction MSE, (6) gives the Kalman gain, (7) denotes the correction
equations and (8) is the updated MMSE. The initial conditions are given by
$\hat{x}(-1|-1)=\mu_{0}$ and $M(-1|-1)=\sigma_{0}^{2}$.

For the asymptotic case ($n\rightarrow\infty$), the MMSE converges to a
deterministic value which can be calculated by solving the following
quadrature equation [10]

$\displaystyle
M(\infty)=\frac{\sigma_{e}^{2}(a^{2}M(\infty)+\sigma_{u}^{2})}{\sigma_{e}^{2}+\rho|h|^{2}(a^{2}M(\infty)+\sigma_{u}^{2})},$
(9)

with a solution equal to

$\displaystyle M(\infty)=\frac{-Q_{2}+\sqrt{Q_{2}^{2}-4Q_{1}Q_{3}}}{2Q_{3}},$
(10)

where $Q_{1}=\rho|h|^{2}a^{2}$,
$Q_{2}=\rho|h|^{2}\sigma_{u}^{2}+\sigma_{e}^{2}(1-a^{2})$, and
$Q_{3}=-\sigma_{u}^{2}\sigma_{e}^{2}$. For the special case with $\rho=0$
i.e., the received signal is used only for energy harvesting, the MMSE becomes
equal to

$\displaystyle M(\infty)=\sigma_{u}^{2}/(1-a^{2})=\;\sigma^{2}.$ (11)

Then, by using a linear WPT model333A linear WPT model refers to the linear
operation regime of the rectenna and serves as a useful guideline for
practical implementations; more sophisticated non-linear models can be also
used [3, 2]. and by neglecting energy harvesting from the AWGN component, the
average harvested energy for the time instant $n$ becomes equal to

$\displaystyle\mathcal{E}(n)=\zeta\mathbb{E}\\{|\sqrt{1-\rho}hx(n)|^{2}\\}=\zeta(1-\rho)|h|^{2}\mathbb{E}\\{|x(n)|^{2}\\},$
(12)

where $\zeta$ is the conversion efficiency of the energy harvester and
$|x(n)|$ is a Rician distributed random variable with parameters
$a^{n+1}\mu_{0}$ and
$\sqrt{(a^{2n+2}\sigma_{0}^{2}+\sigma^{2}(1-a^{2n+2}))/2}$, since
$x(n)\sim\mathcal{CN}(a^{n+1}\mu_{0},a^{2n+2}\sigma_{0}^{2}+\sigma^{2}(1-a^{2n+2}))$.
Therefore, we have

$\displaystyle\mathbb{E}\\{|x(n)|^{2}\\}=a^{2n+2}(\mu_{0}^{2}+\sigma_{0}^{2})+\sigma^{2}(1-a^{2n+2}).$
(13)

Then, for the asymptotic case with $n\rightarrow\infty$, we have
$x(n)\sim\mathcal{CN}(0,\sigma^{2})$ and thus $|x(n)|$ follows a Rayleigh
distribution with

$\displaystyle\mathbb{E}\\{|x(n)|^{2}\\}\rightarrow\sigma^{2},$ (14)

and so

$\displaystyle\mathcal{E}(\infty)=\lim_{n\rightarrow\infty}\mathcal{E}(n)=\zeta\frac{(1-\rho)|h|^{2}\sigma_{u}^{2}}{1-a^{2}}.$
(15)

It is worth noting that $\mathcal{E}(\infty)=\zeta|h|^{2}\mathcal{M}(\infty)$
for $\rho=0$.

Figure 1: System model; A receiver architecture that simultaneously estimates
a Gauss-Markov signal and harvests energy by using PS.

###  II-B Gauss-Markov over a Rayleigh fading channel

In this case, the state $x(n)$ is transmitted through a Rayleigh fading
channel with $|h(n)|^{2}\sim\exp(-\lambda)$ [7, 8]. The Kalman filter
equations are similar to (4)-(8) by replacing the constant $|h|^{2}$ with the
random variable $|h(n)|^{2}$. The MMSE becomes a nonstationary random variable
with a cumulative density function (CDF) upper/lower bounded (for any $n$)
according to Theorem 1 given below.

######  Theorem 1.

For a Rayleigh fading channel, the MMSE is a nonstationary random variable
with a CDF bounded as follows

$\displaystyle F\left(\sigma_{u}^{2}(1+a^{2}),x\right)\leq F_{M_{n}}(x)\leq
F\left(\frac{\sigma_{u}^{2}}{1-a^{2}},x\right),$ (16)

where

$\displaystyle
F(c,x)=\left\\{\begin{array}[]{ll}\exp\left(-\frac{\lambda\sigma_{e}^{2}}{\rho}\left(\frac{1}{x}-\frac{1}{c}\right)\right),\;\;\textrm{if}\;x<c,\\\
1,\;\;\textrm{elsewhere}.\end{array}\right.$ (19)

###### Proof.

The proof is given in Appendix A. ∎

It is obvious that for small values of $a$, the two bounds become tight and
both approximate efficiently the actual CDF i.e., $F_{M_{n}}(x)\approx
F(\sigma_{u}^{2},x)$.

Regarding the bounding of the mean value of the MMSE, we derive the expected
values associated with the two considered bounds. Specifically, the expected
value corresponding to the CDF in (19) is given by

$\displaystyle\Theta(c)$ $\displaystyle=\int_{0}^{c}(1-F(c,x))dx$
$\displaystyle=\frac{\lambda\sigma_{e}^{2}}{\rho}\exp\left(\frac{\lambda\sigma_{e}^{2}}{\rho
c}\right)\Gamma\left(0,\frac{\lambda\sigma_{e}^{2}}{\rho c}\right),$ (20)

which follows through integration by parts and from the definition of the
upper incomplete gamma function
$\Gamma(s,x)=\int_{x}^{\infty}t^{s-1}\exp(-t)dt$. Therefore, the mean MMSE is
bounded as follows i.e.,

$\displaystyle\Theta(\sigma_{u}^{2}(1+a^{2}))\leq\mathbb{E}\\{M_{n}\\}\leq\Theta\left(\frac{\sigma_{u}^{2}}{1-a^{2}}\right),$
(21)

and thus for low values of $a$, we have
$\mathbb{E}\\{M_{n}\\}\approx\Theta(\sigma_{u}^{2})$. By using the expression
in (12), the mean harvested energy becomes

$\displaystyle\mathcal{E}(n)=\zeta\mathbb{E}\\{|\sqrt{1-\rho}h(n)x(n)|^{2}\\}$
$\displaystyle=\zeta(1-\rho)\mathbb{E}\\{|h(n)|^{2}\\}\mathbb{E}\\{|x(n)|^{2}\\}=\frac{\zeta(1-\rho)}{\lambda}\mathbb{E}\\{|x(n)|^{2}\\},$
(22)

where $\mathbb{E}\\{|x(n)|^{2}\\}$ is given by (13) or (14). It is worth
noting that the above mathematical framework is general and can be applied to
other channel fading models e.g., Rice/Nakagami.

###  II-C High power amplifier nonlinearities

We study the impact of high power amplifier nonlinearities on the estimation-
energy tradeoff. Specifically, we assume that the transmitter suffers from HPA
nonlinearities [2]; in this case the transformation of the states to
observations becomes nonlinear. The output of the power amplifier for a
complex input $z$ is given by

$\displaystyle F_{\text{HPA}}(z)$
$\displaystyle=f_{A}(|z|)\exp(j(f_{P}(|z|)+\arg{z})),$ (23)

where $f_{A}(|z|)$ and $f_{P}(|z|)$ denote the AM/AM and AM/PM distortion,
respectively. For the solid state power amplifier (SSPA) model [2], the AM/AM
and AM/PM characteristic functions are given by

$\displaystyle
f_{A}(|z|)=\frac{|z|}{\left[\left(\frac{|z|}{A_{\text{sat}}}\right)^{2\beta}+1\right]^{\frac{1}{2\beta}}},\;\;\;f_{P}(|z|)=0,$
(24)

where $A_{\text{sat}}$ is the output saturation voltage and $\beta$ represents
the smoothness of the transition from the linear regime to the saturation. In
this case, the Gauss-Markov nonlinear model is

State: $\displaystyle x(n)=ax(n-1))+u(n),\;\;n\geq 0,$ (25) Observations:
$\displaystyle y(n)=hF_{\text{HPA}}(x(n))+v(n).$ (26)

By linearizing the function $F_{\text{HPA}}(x(n))$ through the first-order
Taylor expansion i.e., $F_{\text{HPA}}(x(n))\approx
F_{\text{HPA}}(\hat{x}(n|n-1))+f(n)(x(n)-\hat{x}(n|n-1))$ and by applying the
conventional linear Kalman filter, we have the extended Kalman filter given by
[10, Ch. 13.7]

$\displaystyle\hat{x}(n|n-1)=a\hat{x}(n-1|n-1),$ (27) $\displaystyle
M(n|n-1)=a^{2}M(n-1|n-1)+\sigma_{u}^{2},$ (28) $\displaystyle
K(n)=\frac{M(n|n-1)\sqrt{\rho}h^{*}f(n)^{*}}{\sigma_{e}^{2}+\rho|h|^{2}|f(n)|^{2}M(n|n-1)},$
(29) $\displaystyle\hat{x}(n|n)=\hat{x}(n|n-1)+K(n)[y^{\prime}(n)$
$\displaystyle\qquad\qquad\qquad\qquad\qquad-\sqrt{\rho}hF_{\text{HPA}}(\hat{x}(n|n-1))],$
(30) $\displaystyle M(n|n)=(1-K(n)\sqrt{\rho}hf(n))M(n|n-1)$
$\displaystyle=\frac{a^{2}\sigma_{e}^{2}M(n-1|n-1)+\sigma_{e}^{2}\sigma_{u}^{2}}{\sigma_{e}^{2}+\rho|h|^{2}|f(n)|^{2}(a^{2}M(n-1|n-1)+\sigma_{u}^{2})},$
(31)

where $f(n)=\left.\frac{dF_{\text{HPA}}}{dx(n)}\right|_{x(n)=\hat{x}(n|n-1)}$
is given in Appendix B. The average harvested energy is given by

$\displaystyle\mathcal{E}(n)$
$\displaystyle=\zeta\mathbb{E}\\{|y^{\prime}(n)|^{2}\\}=\zeta\mathbb{E}\\{|\sqrt{1-\rho}hF_{\text{HPA}}(x(n))|^{2}\\}$
$\displaystyle=\zeta(1-\rho)|h|^{2}\mathbb{E}\\{|F_{\text{HPA}}(x(n))|^{2}\\}\\},$
(32)

where

$\displaystyle\mathbb{E}\\{|F_{\text{HPA}}(x(n))|^{2}\\}$
$\displaystyle=\int_{0}^{\infty}f_{A}(|x|)^{2}\frac{2x}{\nu^{2}}\exp\\!\left(\\!-\frac{x^{2}\\!+\\!a^{2n+2}\mu_{0}^{2}}{\nu^{2}}\right)$
$\displaystyle\qquad\qquad\times
I_{0}\left(\frac{2xa^{n+1}\mu_{0}}{\nu^{2}}\right)dx,$ (33)

where $f_{A}(|x|)$ is given by (24),
$\nu^{2}=a^{2n+2}\sigma_{0}^{2}+\sigma^{2}(1-a^{2n+2})$, and $I_{0}(\cdot)$ is
the modified Bessel function of the first kind. The above integral can be
evaluated numerically but we can provide a closed-form for high values of $n$.
In this case, the envelope of $x(n)$ follows a Rayleigh distribution with
scale parameter $\sigma/\sqrt{2}$. Therefore, we have

$\displaystyle\mathbb{E}\\{|F_{\text{HPA}}(x(n))|^{2}\\}\\!=\\!\int_{0}^{\infty}\\!\\!\frac{x^{2}}{\left[\left(\frac{x}{A_{\text{sat}}}\right)^{2\beta}\\!+\\!1\right]^{\frac{1}{\beta}}}\frac{2x}{\sigma^{2}}\exp\left(-\frac{x^{2}}{\sigma^{2}}\right)dx$
$\displaystyle=\frac{A_{\text{sat}}^{2}}{\beta\sigma^{2}}\int_{0}^{\infty}\\!\\!\\!\\!\\!\frac{z^{2/\beta-1}}{\left[z\\!+\\!A_{\text{sat}}^{2\beta}\right]^{\frac{1}{\beta}}}\exp\left(-\frac{z^{1/\beta}}{\sigma^{2}}\right)dz,$
(34)

which follows from the transformation $x^{2\beta}\to z$. Then, by using the
identity
$\exp(x)=G^{1,0}_{0,1}\left(-x\big{|}\genfrac{}{}{0.0pt}{}{-}{0}\right)$ and
the result from [11, 2.24.2.4], we have for $\beta\in\mathbb{Z}^{+}$

$\displaystyle\mathbb{E}\\{|F_{\text{HPA}}(x(n))$
$\displaystyle|^{2}\\}=\frac{A_{\text{sat}}^{4}(2\pi)^{(1-\beta)/2}}{\sqrt{\beta}\sigma^{2}\Gamma(1/\beta)}$
$\displaystyle\times
G_{1,\beta+1}^{\beta+1,1}\left(\left[\frac{A_{\text{sat}}^{2}}{\beta\sigma^{2}}\right]^{\beta}\bigg{|}\genfrac{}{}{0.0pt}{}{1-2/\beta,-}{-1/\beta,\Delta(\beta)}\right),$
(35)

where $\Delta(\beta)=0,1/\beta,\dots,(\beta-1)/\beta$ and
$G_{pq}^{mn}\left(\cdot~{}\Big{|}~{}\genfrac{}{}{0.0pt}{}{a_{1},\dots,a_{p}}{b_{1},\dots,b_{q}}\right)$
denotes the Meijer G-function [9, 9.301]. Clearly, the HPA nonlinearities have
a negative effect on the harvested energy, that is, (33) and (35) are always
less or equal to (13) and $\sigma_{u}^{2}/(1-a^{2})$, respectively. Note that
equality is achieved for large values of $A_{\text{sat}}$ since then we have
$f_{A}(|z|)\to|z|$.

Figure 2: MMSE versus $n$; $a=0.8$, $\sigma_{0}^{2}=0.1$, $\mu_{0}=0$,
$\sigma_{u}^{2}=\\{0.0001,0.001,0.02\\}$, $\sigma_{v}^{2}=1$,
$\sigma_{q}^{2}=0.5$, $\rho=\\{0.1,0.9\\}$ and $|h|^{2}=1$. Figure 3:
Estimation-energy tradeoff for various $n$; $\mu_{0}=0$, $\sigma_{0}^{2}=0.1$,
$a=0.8$, $\sigma_{u}^{2}=0.001$, $\sigma_{v}^{2}=1$, $\sigma_{q}^{2}=0.5$,
$|h|^{2}=1$, $\zeta=1$, $\rho\in[0,1]$.

##  III Numerical results

Computer simulations are carried-out to validate the performance of the
proposed architecture. Fig. 2 refers to the static channel scenario and plots
the achieved MMSE $M(n|n)$ versus the symbol time $n$ for different parameters
$\rho$ and $\sigma_{u}^{2}$. Since $\sigma_{0}^{2}>\sigma_{u}^{2}$, it can be
seen that the MMSE performance is improved as $n$ increases. Asymptotically
and according to the expression in (10), the MMSE performance converges to the
floor $M(\infty)$. It can be also seen that the MMSE performance is improved
as the parameter $\rho$ increases since a higher fraction of the received
signal is used for state estimation. A similar trend is observed when the
parameter $\sigma_{u}^{2}$ decreases; the excitation noise $\sigma_{u}^{2}$ is
more critical for the estimation process since it determines the dynamics of
the states and therefore significantly affects the performance of the linear
Kalman filter.

Fig. 3 shows the MMSE-harvested energy tradeoff for various values of $n$ and
for $0\leq\rho\leq 1$ ($n$ affects the tradeoff when the system operates in
the non asymptotic regime); each point of the curves corresponds to a
different power splitting parameter $\rho$. It can be seen that as $n$
increases both estimation and harvesting performance are improved; however,
the performance becomes independent of $n$ for $n>30$ (asymptotic regime). In
addition, we observe that $\rho$ significantly affects the tradeoff since it
defines the fraction of the received signal that is used for estimation and
harvesting, respectively. The parameter $\rho$ is designed based on the system
requirements to achieve a specific estimation-harvesting tradeoff.

Fig. 4 deals with the fading case and presents the proposed CDF upper/lower
bounds for different values of $a$; the empirical/true CDF is also plotted
based on Monte Carlo simulations. As it can be seen, the proposed bounds
efficiently approximate the true CDF for all the values of $a$; in addition,
it can be seen that as $a$ decreases the CDF curves are shifted to the left (
i.e., the MMSE performance degrades since the correlation of consecutive
states decreases) and their gap is squeezed. Fig. 5 plots the average MMSE
versus the harvested energy for a fading scenario with small values of $a$ and
large values of $n$ (asymptotic case); for different values of $\rho$, the
proposed architecture achieves a different estimation-energy tradeoff. A
comparison with the static case, shows that channel fading boosts the MMSE
performance due to the associated time diversity. Finally, in Fig. 6, we study
the impact of HPA and plot the MMSE performance versus the symbol time $n$ (we
assume $\sigma_{0}^{2}=\sigma_{u}^{2}<\sigma^{2}$ and thus the MMSE increases
as $n$ increases). It can be seen that HPA significantly affects the MMSE
performance since it makes the observation equation nonlinear and degrades the
performance of the linear Kalman filter (extended Kalman).

Figure 4: CDF of MMSE for the fading scenario- empirical and upper/lower
bounds; $n=10$, $\rho=0.9$, $\mu_{0}=0$, $\sigma_{0}^{2}=0.1$,
$\sigma_{u}^{2}=0.002$, $\sigma_{v}^{2}=1$, $\sigma_{q}^{2}=0.5$, $\lambda=1$;
(a) $a=0.9$, (b) $a=0.6$, (c) $a=0.4$, and (d) $a=0.1$. Figure 5: MMSE versus
energy harvesting for a fading scenario (asymptotic case) with $\lambda=1$,
$a=0.3$, $\zeta=1$, $\sigma_{u}^{2}=0.003$, $\sigma_{v}^{2}=1$,
$\sigma_{q}^{2}=0.5$. Figure 6: MMSE versus the symbol time $n$ for the
scenario with HPA; $\rho=0.9$, $a=0.8$, $\mu_{0}=0$, $\sigma_{0}^{2}=0.01$,
$\sigma_{u}^{2}=0.01$, $\sigma_{v}^{2}=1$, $\sigma_{q}^{2}=0.5$, and
$|h|^{2}=1$.

##  IV Conclusion

This paper has studied a receiver architecture which enables simultaneous
estimation of a Gauss-Markov signal and WPT. We have considered three
operation scenarios: static channel, Rayleigh block-fading, and transmission
with HPA. The considered architecture employs a Kalman filter for estimation
and is characterized by a fundamental tradeoff between estimation quality
(MMSE) and average harvested energy. Closed-form expressions are given for
both MSE performance and average harvested energy. It has been shown that
channel fading enhances MSE while HPA significantly affects both
estimation/harvesting performance.

##  Appendix A Proof of Theorem 1

An upper bound corresponds to the case where, due to the channel fading, no
observation is available at the receiver. In this case, the MMSE is upper
bounded by
$M(n|n-1)\leq\text{variance}(x(n))\rightarrow\frac{\sigma_{u}^{2}}{1-a^{2}}$
(for high $n$) which is similar to the initialization phase of the Kalman
filter [10]. Regarding the lower bound, from the expressions in (5), we have
$M(n|n-1)\geq\sigma_{u}^{2}$ since $M(n-1|n-1)\geq 0$.

Since the MMSE $M(n|n-1)$ is bounded in both cases by a constant, the random
variable $M(n|n)$ is simplified as

$\displaystyle M(n|n)=\frac{\sigma_{e}^{2}}{\sigma_{e}^{2}/c+\rho|h(n)|^{2}},$
(36)

where $c=a^{2}c_{b}+\sigma_{u}^{2}$ is a constant and $c_{b}$ denotes the
upper/lower bound. The CDF is given by

$\displaystyle F_{M_{n}}(x)$
$\displaystyle=\mathbb{P}\left\\{\frac{\sigma_{e}^{2}}{\sigma_{e}^{2}/c+\rho|h(n)|^{2}}\leq
x\right\\}$
$\displaystyle=\mathbb{P}\left\\{|h(n)|^{2}\geq\frac{\sigma_{e}^{2}}{\rho}\left[\frac{1}{x}-\frac{1}{c}\right]\right\\}$
$\displaystyle=\left\\{\begin{array}[]{ll}\exp\left(-\frac{\lambda\sigma_{e}^{2}}{\rho}\left(\frac{1}{x}-\frac{1}{c}\right)\right),\;\;\textrm{if}\;x<c,\\\
1,\;\;\textrm{elsewhere.}\end{array}\right.$ (39)

For the upper bound, we have $c_{b}=\sigma_{u}^{2}/(1-a^{2})$ and thus
$c=c_{b}=\sigma_{u}^{2}/(1-a^{2})$; for the lower bound $c_{b}=\sigma_{u}^{2}$
and thus $c=\sigma_{u}^{2}(1+a^{2})$. This completes the proof.

##  Appendix B Extended Kalman - Computation of $f(n)$

Let $z=x+jy$ be the complex input to the power amplifier. Then, the output is
given by

$\displaystyle F_{\text{HPA}}(z)$ $\displaystyle=f_{A}(|z|)\exp(j\arg{z})$
$\displaystyle=\frac{\sqrt{x^{2}+y^{2}}}{\left[\left(\frac{\sqrt{x^{2}+y^{2}}}{A_{\text{sat}}}\right)^{2\beta}+1\right]^{\frac{1}{2\beta}}}\cos\left(\tan^{-1}\frac{y}{x}\right)$
$\displaystyle+j\frac{\sqrt{x^{2}+y^{2}}}{\left[\left(\frac{\sqrt{x^{2}+y^{2}}}{A_{\text{sat}}}\right)^{2\beta}+1\right]^{\frac{1}{2\beta}}}\sin\left(\tan^{-1}\frac{y}{x}\right)$
$\displaystyle=U(x,y)+jV(x,y).$ (40)

By applying the Cauchy Riemann equations [12, 11.2],

$\displaystyle
f(n)=\left.\frac{dF_{\text{HPA}}}{dz}\right|_{z=x_{0}+jy_{0}}=\frac{dU}{dx}(x_{0},y_{0})+j\frac{dV}{dx}(x_{0},y_{0})$
$\displaystyle=\frac{\sqrt{x_{0}^{2}+y_{0}^{2}}\left(1+\left(\frac{x_{0}^{2}+y_{0}^{2}}{A_{\text{sat}}^{2}}\right)^{\beta}\right)^{-\frac{1+2\beta}{2\beta}}\left(x_{0}^{2}+y_{0}^{2}\left(1+\left(\frac{x_{0}^{2}+y_{0}^{2}}{A_{\text{sat}}^{2}}\right)^{\beta}\right)\right)}{\left(x_{0}^{2}+y_{0}^{2}\right)^{3/2}}$
$\displaystyle\qquad-j\frac{y_{0}\left(\frac{\sqrt{x_{0}^{2}+y_{0}^{2}}}{A_{\text{sat}}}\right)^{2\beta-1}\left(1+\left(\frac{x_{0}^{2}+y_{0}^{2}}{A^{2}_{\text{sat}}}\right)^{\beta}\right)^{-\frac{1+2\beta}{2\beta}}}{A_{\text{sat}}\sqrt{1+\frac{y_{0}^{2}}{x_{0}^{2}}}},$
(41)

where $z=\hat{x}(n|n-1)$, $x_{0}=\Re\\{\hat{x}(n|n-1)\\}$ and
$y_{0}=\Im\\{\hat{x}(n|n-1)\\}$.

## References

  * [1] L. R. Varshney, “Transporting information and energy simultaneously,” in Proc. IEEE Int. Symp. Inf. Theory, Toronto, Canada, July 2008, pp. 1612–1616.
  * [2] I. Krikidis, “Information-energy capacity region for SWIPT systems with power amplifier nonlinearity,” in Proc. IEEE Int. Symp. Inf. Theory, Los Angeles, USA, June 2020.
  * [3] B. Clerckx, R. Zhang, R. Schober, D. W. K. Ng, D. I. Kim, and H. V. Poor, “Fundamentals of wireless information and power transfer: From RF energy harvester models to signal and system designs,” IEEE J. Selec. Areas Commun., vol. 37, pp. 4–33, Jan. 2019.
  * [4] Y.-W. P. Hong, T.-C. Hsu, and P. Chennakesavula, “Wireless power transfer for distributed estimation in wireless passive sensor networks,” IEEE Trans. Sig. Proc., vol. 64, pp. 5382–5395, Oct. 2016.
  * [5] N. K. D. Venkategowda, H. Lee, and I. Lee, “Data precoding and energy transmission for parameter estimation in MIMO wireless powered sensor networks,” in Proc. Veh. Techn. Conf., Toronto, Canada, Sept. 2017.
  * [6] C. Psomas and I. Krikidis, “Energy beamforming in wireless powered mmWave sensor networks,” IEEE J. Selected Areas Commun., vol. 37, pp. 424–438, Feb. 2019.
  * [7] R. Parseh and K. Kansanen, “On estimation error outage for scalar Gauss-Markov signals sent over fading channels,” IEEE Trans. Sign. Proc., vol. 62, pp. 6225–6234, Dec. 2014.
  * [8] R. Parseh, K. Kansanen, and D. Slock,“Distortion outage analysis for joint space-time coding and Kalman filtering,” IEEE Trans. Sign. Proc., vol. 65, pp. 2291–2305, May 2017.
  * [9] I. S. Gradshteyn and I. M. Ryzhik, _Table of Integrals, Series, and Products_ , San Diego, CA: Academic Press, 2007.
  * [10] S. M. Kay, Fundamentals of Statistical Signal Processing: Estimation Theory, Prentice Hall, USA, 1993.
  * [11] A. P. Prudnikov, Y. A. Brychkov, and O. I. Marichev, _Integral and Series. Vol. 3: More Special Functions_ , Amsterdam: Gordon and Breach Science Publishers, 1986.
  * [12] W. Rudin, Real and Complex Analysis, McGraw-Hill Education, 1986.


1

# Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix
Multiplication on the GPU

Muhammad Osama<EMAIL_ADDRESS>0000-0003-1616-6817 University of
California, Davis1 Shields AveDavisCaliforniaUSA95616 ,  Duane Merrill
<EMAIL_ADDRESS>0000-0002-4869-9477 NVIDIA Corporation2788 San Tomas
ExpySanta ClaraCaliforniaUSA95051 ,  Cris Cecka<EMAIL_ADDRESS>0000-0002-3542-7675 NVIDIA Corporation2788 San Tomas ExpySanta
ClaraCaliforniaUSA95051 ,  Michael Garland<EMAIL_ADDRESS>0000-0001-6093-7602 NVIDIA Corporation2788 San Tomas ExpySanta
ClaraCaliforniaUSA95051 and  John D. Owens<EMAIL_ADDRESS>0000-0001-6582-8237 University of California, Davis1 Shields
AveDavisCaliforniaUSA95616

###### Abstract.

We introduce _Stream-K_ , a work-centric parallelization of matrix
multiplication (GEMM) and related computations in dense linear algebra.
Whereas contemporary decompositions are primarily tile-based, our method
operates by partitioning an even share of the aggregate inner loop iterations
among physical processing elements. This provides a near-perfect utilization
of computing resources, regardless of how efficiently the output tiling for
any given problem quantizes across the underlying processing elements.

On GPU processors, our _Stream-K_ parallelization of GEMM produces a peak
speedup of up to 14$\times$ and 6.7$\times$, and an average performance
response that is both higher and more consistent across 32,824 GEMM problem
geometries than state-of-the-art math libraries such as CUTLASS and cuBLAS.
Furthermore, we achieve this performance from a _single_ tile size
configuration per floating-point precision, whereas today’s math libraries
employ complex kernel-selection heuristics to select from a large ensemble of
kernel variants.

Matrix-Multiplication, GPU, Load-Balancing

Distribution Statement “A” (Approved for Public Release, Distribution
Unlimited).

††copyright: rightsretained

##  1\. Introduction

General matrix-matrix product (GEMM), convolution, and other similar
computations constitute the dominant workloads in many deep learning and
scientific computing applications. High-performance processors such as GPUs,
for example, are designed to achieve nearly 100% of their theoretical peak
math throughput when computing GEMM. Doing so, however, requires a work
decomposition that perfectly occupies the underlying physical cores. As we
show, attaining such high levels of processor utilization across a broad
landscape of problems shapes and sizes can be challenging.

Classically, GEMM implementations block their computation using a _data-
parallel_ tiling of the output matrix, assigning the independent production of
output tiles among concurrent threads (or thread groups) (Abdelfattah et al.,
2016; Kerr et al., 2017; Nath et al., 2010). The work per output tile is
regular, and tile production tends to dispatch across idle physical cores in
“waves”. The overall workload is well-balanced and processor utilization is
highest when there are many waves, i.e., the number of output tiles greatly
oversubscribes the number of cores.

However, such oversubscription has shrunk considerably as processors have
grown in size. An increased core count will require fewer waves to produce a
given tile count. Bigger cores will compel larger matrix blocking factors,
leading to fewer waves of larger tiles. In general, execution schedules with
fewer waves are much more likely to suffer from _quantization inefficiency_ ,
i.e., the processor underutilization that occurs when the number of output
tiles is not an even multiple of the number of processor cores. When the last
wave is partially full, the unused cores must wait for the remaining threads
to execute millions (if not billions) of multiply-accumulate (MAC)
instructions before they are able to execute any dependent work.

Figure 1(a) illustrates such a scenario on a hypothetical GPU with four
streaming multiprocessor cores (SMs). If we block a $384\times 384\times 128$
GEMM computation into nine $128\times 128$ output tiles, a _data-parallel_
decomposition cannot achieve more than 75% of the processor’s rated
throughput. This theoretical utilization ceiling can be improved to 90% by
halving the tile size as shown in Figure 1(b). However, the finer-grained
blocking factor will be less cache and scratchpad efficient, and may preclude
any practical performance improvement.

(a) Data parallel decomposition with grid size $g$=9 CTAs,
large $128\times 128\times 128$ CTA work volumes,
and 75% processor utilization ceiling

(b) _Data parallel_ decomposition with grid size $g$=18 CTAs,
smaller $128\times 64\times 128$ CTA work volumes,
and 90% processor utilization ceiling

Figure 1. _Data-parallel_ execution schedules for $384\times 384\times 128$
GEMM across a hypothetical four-SM GPU.

(a) _Fixed-split_ decomposition with splitting factor $s$=2,
grid size $g$=18 CTAs, smaller $128\times 128\times 64$ CTA work volumes,
and 90% quantization efficiency

(b) Basic _Stream-K_ decomposition with grid size $g$=4 CTAs,
larger $128\times 128\times 288$ CTA work volumes,
and nearly 100% quantization efficiency

Figure 2. Tile-splitting execution schedules for $384\times 384\times 128$
GEMM across a hypothetical four-SM GPU.

Quantization inefficiency is a concern for increasingly wide processors such
as GPUs, where ALUs-per-core and cores-per-processor both currently number in
the hundreds. Consequently, many common GEMM-like workloads now exhibit a
final, partially full wave that comprises a significant fraction of the total
computation time.

The current remedy employed by GPU-based math and deep learning libraries is
to deploy an ensemble of tiling configurations. When the ideal blocking factor
does not quantize well, the library chooses among tiling alternatives with
smaller concurrent work volumes, such as those illustrated in Figure 1(b) and
Figure 2(a).

Tile-based ensembles, however, present performance and logistical challenges
for math libraries seeking to deliver the best-achievable performance across
diverse problem sizes and shapes. Distributable code size can be problematic
for large ensembles. For example, NVIDIA’s cuBLAS library (NVIDIA Corporation,
2020) is hundreds of megabytes, often providing more than twenty pre-compiled
kernel specializations per architecture for a given API entry point. Large
ensembles also require sophisticated selection heuristics. In our evaluation,
we show these heuristics can struggle to consistently identify the optimal
configuration for arbitrary problems.

Unlike these tile-based methods, our _Stream-K_ decomposition always
distributes an even share (within one) of the aggregate multiply-accumulate
loop iterations required by the GEMM computation across SMs. Because the
instruction workload of a single MAC-loop iteration is far smaller than that
of an entire output tile, any variance in core workload is practically
negligible. _Stream-K_ uses the ideal blocking factor regardless of problem
shape, has communication overheads that scale with processor width (rather
than output tiles), and compiles to a single kernel.

We use an enormous corpus of 32,824 GEMM shapes and sizes to evaluate _Stream-
K_ , which we implemented within NVIDIA’s CUTLASS library (Kerr et al., 2017).
In comparison with CUTLASS’s _data-parallel_ implementation of the same
blocking factor, _Stream-K_ provides a substantially higher performance
response across our landscape of GEMM problems, demonstrating up to 14$\times$
speedup on NVIDIA A100 GPUs.

To highlight the practical challenges of ensemble-based solutions, we also
evaluate NVIDIA’s cuBLAS library as well as an oracle-driven ensemble of
_data-parallel_ CUTLASS tilings. Relative to both ensembles, we show that our
single-kernel _Stream-K_ achieves both (1) higher average performance, and (2)
higher performance consistency. Versus cuBLAS, _Stream-K_ demonstrates up to
6.7$\times$ speedup and virtually no instances of slowdown for compute-bound
problems.

##  2\. Background

General Matrix Multiplication (GEMM) is defined as the product
$\textbf{C}=\alpha\textbf{AB}+\beta\textbf{C}$ where $\alpha$ and $\beta$ are
scalar values and A, B, and C are matrices. (For simplicity, we assume
$\alpha=1$, $\beta=0$ throughout this paper.) We refer to the _shape_ of a
given GEMM problem by the volumetric extents of its computation. For example,
a $m\times n\times k$ GEMM consumes $m\times k$ and $k\times n$ input matrices
A and B, respectively, performs $m\times n\times k$ multiply-accumulate
operations, and produces an $m\times n$ output matrix C.

GEMM is a performance-critical subroutine in many large-scale engineering and
scientific applications. It plays an important role in matrix factorization
methods such as LU, QR, and Cholesky decomposition. High-performance modeling
and simulation applications in engineering, climate simulation, cosmology,
quantum chemistry, and other scientific domains rely on these factorization
methods.

Matrix multiplication is also the fundamental building block of modern deep
learning (DL) methods. The training of deep neural networks (DNNs) is often
performed on massive datasets across large distributed systems (Mattson et
al., 2020). Many DL training and inference operations are cast as matrix
multiplications. For example, image recognition and computer vision models
rely on convolution, which can be implemented directly as the product of
filter and image datasets (Chetlur et al., 2014). Transformer architectures,
which have come to dominate natural language processing and other
applications, are almost entirely limited by the performance of large matrix
products.

Early work on GPU matrix-matrix multiplication from Larsen and McAllister
framed the computation as a multi-texture multiplication and blending
operation (Larsen and McAllister, 2001). The user-programmable shared memory
provided by subsequent GPU architectures enabled higher-performing _data
parallel_ schemes with two levels of blocking (shared memory and registers)
with tile sizes informed via extensive micro-benchmarking analysis (Barrachina
et al., 2008; Nath et al., 2010; Tan et al., 2011; Tillet et al., 2019) and
auto-tuning (Cui et al., 2010; Jiang and Snir, 2005; Li et al., 2009).

The MAGMA GPU math library was perhaps the first to optimize for diverse GEMM
problem shapes (Kurzak et al., 2012). Their solution applied a constrained set
of tiling parameters to a templated CUDA C++ code stencil, generating several
hundred _data-parallel_ variants per API primitive (e.g., hgemm_tt() for half-
precision transpose-transpose GEMM). They evaluated these variants to distill
a small ensemble of typically three to five kernels that collectively perform
well across a diversity of problem shapes. Kernel selection and dispatch for a
given problem was governed by size thresholds expressed via simple handwritten
rules.

Subsequent GPU math libraries have employed more sophisticated code-generation
and kernel-selection components. For example, the ISAAC project uses machine
learning techniques to predict an optimal tiling and/or splitting
parameterization for a given GEMM shape, which can then be instantiated either
online or offline via a PTX-level code generator (Tillet and Cox, 2017).

NVIDIA’s cuBLAS (NVIDIA Corporation, 2020) library has provided an extended
`cublasGemmEx` interface that allows the caller to select from among 24
different GEMM “algorithms”. Carefully trained heuristics choose between this
large space of alternatives when using the default interface. These algorithms
implement a variety of different _data-parallel_ and _fixed-split_ variants,
and it is common for cuBLAS to have assembled each variant into its own
architecture-specific kernel program for code optimization purposes. The cross
product of GEMM API functionality, strategic variants, and microarchitecture
has resulted in distributions that are increasingly enormous, exceeding
hundreds of megabytes of executable code.

Given the fast-paced and rapidly changing nature of contemporary deep
learning, recent work has focused on programming models for simplifying the
expression and construction high performance kernels that alter or supplement
the GEMM computation. The CUTLASS C++ library provides data-movement and
multiply-accumulation classes for composing custom GEMM-like computations at
all levels of the GPU thread hierarchy (Kerr et al., 2017). Triton (Tillet et
al., 2019) is a domain-specific language for tensor programming centered on
the expression, transformation, and optimization of block/tile concepts. Other
domain-specific programming languages such as Halide (Ragan-Kelley et al.,
2013) and TVM (Chen et al., 2018) separate the expression of pointwise
operators from that of loop scheduling. Fireiron (Hagedorn et al., 2020)
further adds data movement constructs into the scheduling grammar.

##  3\. Existing Work Decomposition Strategies

Modern processors typically store A, B, and C in a large, slow, distant memory
and have access to a small, fast, scratchpad or cache memory. A primary goal
for any GEMM implementation is to leverage these local storage resources so
that the resulting implementation is computation-bound.

###  3.1. Sequential Cache-Blocked

The classic cache-blocked formulation of GEMM divides its computational volume
into blocks and chooses a traversal order that exposes memory locality.
Algorithm 1 presents a simplified implementation comprising six loops. The
innermost three loops iterate within the blocking factors BLK_M, BLK_N, and
BLK_K, while the outermost three iterate across them. If the cache can capture
one block from each of the three matrices, the resulting data reuse among
those elements will significantly reduce the number of last-level memory
accesses (Lam et al., 1991).

Algorithm 1 Sequential cache-blocked GEMM.

1:$\rhd$ tile-processing outer loops

2:for mm $\leftarrow$ 0 to m step BLK_M do

3: for nn $\leftarrow$ 0 to n step BLK_N do

4: $\rhd$ zero-initialize output tile

5: for mmm $\leftarrow$ mm to (mm + BLK_M) do

6: for nnn $\leftarrow$ nn to (nn + BLK_N) do

7: C[mmm,nnn] $\leftarrow$ 0

8: end for

9: end for

10: $\rhd$ perform the MAC iterations for this tile

11: for kk $\leftarrow$ 0 to k step BLK_K do

12: $\rhd$ MAC iteration (fully unrolled)

13: for mmm $\leftarrow$ mm to (mm + BLK_M) do

14: for nnn $\leftarrow$ nn to (nn + BLK_N) do

15: for kkk $\leftarrow$ kk to (kk + BLK_K) do

16: C[mmm,nnn] $\leftarrow$ C[mmm, nnn] +

17: (A[mmm,kkk] $\times$ B[kkk,nnn])

18: end for

19: end for

20: end for

21: end for

22: end for

23:end for

###  3.2. Data-parallel

As shown in Algorithm 2, the _data-parallel_ GPU formulation of GEMM is
decomposed across a grid of parallel thread blocks, or _cooperative thread
arrays_ (CTAs)111Blocks of GPU threads are coscheduled in CTAs, which
virtualize the hardware’s streaming multiprocessor cores (SMs).. The grid is
sized such that each CTA produces its own (BLK_M $\times$ BLK_N) output tile.

For exposition, the MacLoop() subroutine of Algorithm 3 encapsulates the
multiply-accumulate workloads that compute the values of the CTA’s output
tile. It performs a sequence of _MAC-loop_ iterations in the accumulation
domain, e.g., the _k_ -axis for GEMM. Each _MAC-loop_ iteration comprises a
per-thread volume of (BLK_M $\times$ BLK_N $\times$ BLK_K) $/$ CTA_THREADS MAC
operations. As the computation proceeds, fragments of the input matrices are
staged through the SM’s shared memory for local reuse among individual
threads.

Although this particular presentation of MacLoop() deploys one thread per
output tile element, the sophisticated implementations in CUTLASS (Kerr et
al., 2017) and cuBLAS (Kerr et al., 2017) will: (1) fully unroll the per-
thread MAC-loop iteration; (2) implement additional blocking at the warp
and/or thread levels; and (3) orchestrate a software pipeline of shared memory
data movement across MAC-loop iterations.

Unfortunately, this classic _data-parallel_ decomposition is liable to suffer
from quantization inefficiency on modern GPUs, as illustrated in Figure 1.
Although an ensemble of diverse blocking factors may uncover opportunities for
greater processor utilization, it is unlikely to facilitate perfect
quantizations for arbitrary problem sizes. Furthermore, smaller blocking
factors have two drawbacks: (1) fewer instructions per MAC-loop iteration for
covering the latencies of global and shared memory transfers in pipelined
implementations; and (2) a higher proportion of memory operations relative to
MAC instructions, which may prevent them from being computation-bound.

Algorithm 2 Data-parallel GPU GEMM.

1:_shared_ accum[BLK_M,BLK_N]

2:iters_per_tile $\leftarrow\lceil$k/BLK_K$\rceil$

3:$\rhd$ instantiate one CTA per output tile

4:fork CTA[x] in [ $\lceil$m/BLK_M$\rceil$ $\times$ $\lceil$n/BLK_N$\rceil$ ]
do

5: $\rhd$ perform the MAC iterations for this tile

6: accum $\leftarrow$ MacLoop(x, 0, iters_per_tile)

7: $\rhd$ store accumulators to output tile

8: StoreTile(C, x, accum)

9:join

Algorithm 3 CTA-wide MacLoop() subroutine for performing a sequence of MAC-
loop iterations.

1:procedure MacLoop(tile_idx, iter_begin, iter_end)

2: _shared_ accum[BLK_M,BLK_N]

3: _shared_ frag_a[BLK_M,BLK_K]

4: _shared_ frag_b[BLK_K,BLK_N]

5: $\rhd$ determine output tile coordinates

6: mm $\leftarrow$ BLK_M $\times$ (tile_idx / $\lceil$m/BLK_M$\rceil$)

7: nn $\leftarrow$ BLK_N $\times$ (tile_idx % $\lceil$m/BLK_M$\rceil$)

8: $\rhd$ zero-initialize local accumulator storage

9: accum $\leftarrow$ 0

10: $\rhd$ perform the specified range of MAC iters for this tile

11: for iter $\leftarrow$ iter_begin to iter_end do

12: kk $\leftarrow$ iter $\times$ BLK_K

13: $\rhd$ copy global matrix fragments to local storage

14: frag_a $\leftarrow$ LoadFragment(A, mm, kk)

15: frag_b $\leftarrow$ LoadFragment(B, kk, nn)

16: fork THREAD[mmm,nnn] in [BLK_M, BLK_N] do

17: $\rhd$ MAC iteration per thread (fully unrolled)

18: for kkk $\leftarrow$ 0 to BLK_K do

19: accum[mmm, nnn] $\leftarrow$ accum[mmm,nnn] +

20: (frag_a[mmm,kkk] $\times$ frag_b[kkk,nnn])

21: end for

22: join

23: end for

24: return accum

25:end procedure

###  3.3. Fixed-split

Alternatively, the granularity of work assigned to each CTA can be reduced via
parallelization across the accumulation dimension. For a given output tile,
the associativity of addition allows the iteration domain to be split among
multiple concurrent CTAs, followed by a dependent “fixup” step to reduce the
partial sums computed by each CTA. We highlight this _fixed-split_ approach in
Algorithm 4, where each output tile is cooperatively produced by $s$ CTAs.
Notably, it functions identically to the _data-parallel_ decomposition when
the splitting factor $s=1$.

The _fixed-split_ decomposition is also featured in CUTLASS and cuBLAS. The
splitting factor is implemented as a runtime parameter, allowing a single
kernel executable to support multiple work volumes while retaining the ideal
blocking factors for optimal data sharing and latency hiding. However, as
illustrated in Figure 2(a), the prospect of achieving a perfect quantization
from a uniform tile-splitting is unlikely. Furthermore, the extra overheads of
communication and synchronization scale with both the overall problem size as
well as the splitting factor.

Algorithm 4 Fixed-split GPU GEMM with splitting factor $s$.

1:_shared_ accum[BLK_M,BLK_N]

2:iters_per_tile $\leftarrow\lceil$k/BLK_K$\rceil$

3:iters_per_split $\leftarrow\lceil$iters_per_tile/$s$$\rceil$

4:$\rhd$ instantiate $s$ CTAs per output tile

5:fork CTA[x,y] in [ $\lceil$m/BLK_M$\rceil$ $\times$ $\lceil$n/BLK_N$\rceil$,
$s$] do

6: $\rhd$ perform the range of MAC iterations for this split

7: iter $\leftarrow$ y $\times$ iters_per_split

8: iter_end $\leftarrow$ min(iters_per_tile, iter + iters_per_split)

9: accum $\leftarrow$ MacLoop(x, iter, iter_end)

10: $\rhd$ consolidate partial-sums across CTAs

11: if y $\neq$ 0 then

12: $\rhd$ store accumulators to temporary global storage

13: StorePartials(partials[x,y], accum)

14: Signal(flags[x,y])

15: else

16: $\rhd$ accumulate partial sums from other CTAs contributing to this tile

17: for cta $\leftarrow$ 1 to $s$ do

18: Wait(flags[x,cta])

19: accum $\leftarrow$ accum + LoadPartials(partials[x,cta])

20: end for

21: $\rhd$ store accumulators to output tile

22: StoreTile(C, tile_id, accum)

23: end if

24:join

##  4\. Our _Stream-K_ Decomposition

Our _Stream-K_ decomposition is a tile-splitting parallelization in which the
splitting seams are completely dissociated from the tiling structure itself.
Although we employ familiar blocking and tiling strategies for data reuse, we
instead quantize the GEMM computation into MAC-loop iterations, i.e., small
volumes of CTA-wide BLK_M $\times$ BLK_N $\times$ BLK_K work. As presented in
Algorithm 5, _Stream-K_ evenly partitions the GEMM’s aggregate workload of
MAC-loop iterations across a constant-sized grid of $g$ CTAs. Each CTA’s range
of MAC-loop iterations is mapped contiguously into the $m\rightarrow
n\rightarrow k$ linearization of the GEMM shape, crossing output-tile
boundaries as it may.

Algorithm 5 Basic Stream-K GPU GEMM with grid size $g$.

1:_shared_ accum[BLK_M,BLK_N]

2:iters_per_tile $\leftarrow\lceil$k/BLK_K$\rceil$

3:total_iters $\leftarrow\lceil$m/BLK_M$\rceil\times\lceil$n/ BLK_N$\rceil$
$\times$ iters_per_tile

4:iters_per_cta $\leftarrow\lceil$total_iters / g$\rceil$

5:$\rhd$ instantiate g CTAs

6:fork CTA[x] in [g] do

7: iter $\leftarrow$ x $\times$ iters_per_cta

8: iter_end $\leftarrow$ iter + iters_per_cta

9: $\rhd$ iteration-processing outer loop

10: while iter $<$ iter_end do

11: tile_idx $\leftarrow$ iter / iters_per_tile

12: tile_iter $\leftarrow$ tile_idx $\times$ iters_per_tile

13: tile_iter_end $\leftarrow$ tile_iter + iters_per_tile

14: $\rhd$ perform the range of MAC iterations for this tile

15: local_iter $\leftarrow$ iter - tile_iter

16: local_iter_end $\leftarrow$

17: min(iter_end, tile_iter_end) - tile_iter

18: accum $\leftarrow$

19: MacLoop(tile_id, local_iter, local_iter_end)

20: $\rhd$ consolidate partial-sums across CTAs

21: tile_started $\leftarrow$ iter = tile_iter

22: tile_ended $\leftarrow$ (iter_end $\geq$ tile_iter_end)

23: if $\neg$tile_started then

24: $\rhd$ store accum to temporary global storage

25: StorePartials(partials[x], accum)

26: Signal(flags[x])

27: else

28: $\rhd$ store accumulators to output tile

29: if $\neg$tile_ended then

30: $\rhd$ accumulate partial sums from other CTA contributing to this tile

31: cta_end $\leftarrow$ tile_iter_end / iters_per_tile

32: for cta $\leftarrow$ (x+1) in cta_end do

33: Wait(flags[cta])

34: accum $\leftarrow$ accum

35: \+ LoadPartials(partials[cta])

36: end for

37: end if

38: StoreTile(C, tile_id, accum)

39: end if

40: iter $\leftarrow$ tile_iter_end

41: end while

42:join

Should a given CTA’s starting and/or ending iterations not coincide with tile
boundaries (as is expected to be the common case), it must consolidate its
partial results with those of the other CTA(s) also covering that tile. In
this basic implementation, each output tile in C is written by the CTA that
performed that tile’s $k=0$ MAC-loop iteration. Before it can do so, however,
it must accumulate any partial sums shared from other CTAs in temporary global
storage. Notably, _Stream-K’_ s communication, synchronization, and global
storage overheads are independent of problem size, scaling instead with the
number of CTAs $g$.

A secondary benefit of _Stream-K_ is that synchronization-waiting is likely
negligible when the number of output tiles is greater than the number of CTAs.
In this regime, each output tile is covered by at most two CTAs, and the tile-
processing skew ensures that the accumulating CTA will not need its peer
contributions until well after those collaborators have finished producing
them.

Continuing our earlier example, Figure 2(b) illustrates the basic _Stream-K_
execution schedule of the $384\times 384\times 128$ GEMM problem on a
hypothetical four-SM GPU. To fully occupy the GPU, we launch $g=4$ CTAs.
Assuming BLK_M $=128$, BLK_N $=128$, and BLK_K $=4$, each CTA is tasked with a
$128\times 128\times 288$ work volume comprising 72 MAC-loop iterations. This
results in a 100% quantization efficiency, as all four SMs will execute the
same number of MAC instructions.

Additionally, the work volume of a single MAC-loop iteration is 32$\times$
smaller than that of an entire output tile. Consequently, a 32-way _fixed-
split_ decomposition would also provide a 100% quantization efficiency, but at
the expense of an 8$\times$ larger “fixup” overhead. Furthermore, _Stream-K_
is better able to hide the latency of inter-CTA synchronization due to the
temporal skew between writers and readers when sharing partial sums.

_Stream-K_ also generalizes to both _fixed-split_ and _data-parallel_
decompositions. When the grid size $g$ is an even multiple of the number of
output tiles, _Stream-K_ functions exactly as the _fixed-split_ decomposition.
Similarly, when $g$ equals the number of output tiles, _Stream-K_ behaves
identically to the _data-parallel_ decomposition. We take advantage of this
generalization to create an optimized hybridization of the _Stream-K_
decomposition in following section (5.2).

##  5\. Implementation Details

(a) Basic Stream-K

(b) DP + one-tile SK

(c) Two-tile SK + DP

Figure 3. Basic Stream-K vs. hybrid execution schedules for $896\times
384\times 128$ GEMM across a hypothetical four-SM GPU.

The work decomposition we introduced in the last section can be instantiated
in a number of different ways to suit the needs of different hardware
architectures and software library designs. Our implementation targets NVIDIA
GPUs and is designed to be integrated into existing libraries like cuBLAS and
CUTLASS. In this section, we describe how we configure the kernels we launch
and introduce a hybridization scheme that helps ensure users achieve maximum
GEMM performance across the widest possible range of problem shapes.

We also emphasize that these are truly internal implementation details. They
are completely transparent to the user of a BLAS-like library and do not alter
the library’s interface. The only observable impact is the improved
performance characteristics that we analyze in Section 6.

###  5.1. Kernel Configuration

The tile size chosen for blocking the GEMM computation is, of course, a
critical parameter controlling the performance of the GEMM kernel. For modern
NVIDIA GPUs, appropriate tile sizes are determined by the shape of matrices
supported by the GPU’s Tensor Cores. Based on extensive empirical experience,
we selected the smallest CTA-wide tile size capable of achieving 99% of the
GPU’s peak TFLOP/s for very large GEMM volumes for each supported precision.
For the NVIDIA A100 GPU used in our experiments, these sizes are
64$\times$64$\times$16 for FP64 problems and 128$\times$128$\times$32 for
FP16$\rightarrow$32 problems.

Achieving maximal GEMM performance from _Stream-K_ parallelization also
requires some degree of dynamic problem-specific configuration. Before
launching a kernel we choose a grid size likely to yield the best performance
on the specific problem shape at hand. This is in contrast to ensemble-based
approaches which accommodate diverse problem shapes through the static
generation of many kernel variants based on workload decomposition and
blocking factor.

Our grid size selection heuristic is based on a simple analytical model that
minimizes the cost of reading, writing, and accumulating partial sums while
equally distributing the MAC-loop iterations per CTA. Details of this
analytical model are provided in the supplementary material (Appendix A.1).
Parameters to the model are trivially chosen with empirical measurements and
need only be done once per target architecture. The resulting parameters can
then be compiled statically into the library. Again, this is in contrast to
ensemble-based approaches that rely on potentially complex heuristics and
machine learning models for kernel selection at run time.

###  5.2. Data-parallel Hybridization

The basic _Stream-K_ decomposition can, in certain cases, exhibit tile-
processing skew that leads to potentially adverse effects on cache
performance. When the number of output tiles $t$ is not an even multiple of
the grid size $g$, the starting $k$-offset for the first MAC-loop iteration in
each CTA will be different. Depending on the sizes and shapes of the input
matrices and blocking factors, this skew may preclude these fragments from
seeing reuse across CTAs in the GPU’s cache structure. In Figure 3(a), for
example, the initial $k$-axis fragment offsets for each of the four CTAs will
be $k=0$, $k=32$, $k=64$, and $k=96$, respectively. Furthermore, this
32-element skew between CTAs will persist for the duration of the GEMM
computation.

Tile-processing skew is a direct consequence of _Stream-K_ ’s workload
balancing strategy. However, we can take measures to limit its duration by
applying _Stream-K_ ’s iteration balancing to a smaller, tile-aligned region
of the total iteration domain such that the remaining tiles can be produced in
full, temporally aligned waves.

The simplest hybrid scheme is the “ _data-parallel_ \+ one-tile _Stream-K_ ”
schedule illustrated in Figure 3(b). It applies iteration balancing only among
the tiles otherwise remaining for a final, partially full _data-parallel_
wave. The total number of full waves is $w=\lfloor t/p\rfloor$, where $t$ is
the number of output tiles and $p$ is the number of SM cores in the GPU.
Consequently, each _Stream-K_ CTA receives an even share of iterations that is
less than one tile’s worth. Unfortunately, this strategy has little ability to
hide the synchronization latency for the exchange of partial sums when three
or more CTAs cover the same tile. In these scenarios, the accumulating CTA may
be forced to wait for the contributions of other CTAs to become visible, as
all but the last will be completing their final iterations at roughly the same
time. Furthermore, the basic version of our scheme for aggregating partials is
serialized within a single CTA, and thus will likely cause SM workload
imbalance when the number of contributing CTAs per tile is large.

We address these problems with our “two-tile _Stream-K_ \+ _data-parallel_ ”
hybrid schedule, illustrated in Figure 3(c). It performs one fewer full data-
parallel wave in exchange for each _Stream-K_ CTA receiving more than one
tile’s worth of iterations (but fewer than two). This provides much better
latency hiding when $w\geq 2$, and each accumulating CTA will only need to
receive partials from one other contributing CTA. Otherwise, it behaves
identically to the “ _DP + one tile SK_ ” schedule. This hybrid approach
results in both improved memory access patterns and latency hiding. It also
shows the versatility of the generic _Stream-K_ looping structure to implement
different scheduling policies within the same kernel instance.

##  6\. Performance Evaluation

We have implemented our _Stream-K_ decomposition using NVIDIA’s CUTLASS
library of CUDA C++ template abstractions for authoring GEMM-like
computations. CUTLASS provides the optimized equivalent of the CTA-wide
MacLoop() subroutine in Algorithm 3, which performs blocking, tiling, and
software-pipelined data movement that is analogous to the closed-source cuBLAS
and cuDNN implementations. Our evaluation encompasses both (1) double-
precision FP64 GEMM, and (2) mixed-precision FP16$\rightarrow$32 GEMM. For the
latter, the input matrices A and B comprise half-precision FP16 values, yet
the internal accumulation and output matrix C values are single-precision
FP32.

#### Hardware environment.

Our test GPU is the NVIDIA A100, which contains 108 SM cores. For measurement
stability, we lock the power envelope at 400 W and SM clocks at 1005 MHz
($\sim$71% of their dynamic peak). This establishes FP64 tensor-core peak
throughput of 13.9 TFLOP/s, and mixed FP16$\rightarrow$32 tensor-core peak
throughput of 222.3 TFLOP/s.

Figure 4. The test domain of 32,824 GEMM problem shapes and sizes used for
performance evaluation.
$\\{m\\}=\\{128\ldots 8192\\}$, $\\{n\\}=\\{128\ldots 8192\\}$,
$\\{k\\}=\\{128\ldots 8192\\}$

#### Dataset.

Our test corpus intends to approximate the enormous breadth and scope of
device-wide GEMM problems that GPU math kernel libraries are designed to
accommodate. As shown in Figure 4, we evaluate 32,824 different problem sizes
and shapes, log-sampled at random within a domain of $m$, $n$, and $k$ matrix
dimensions whose volume spans six orders of magnitude.

#### Methodology.

For both GEMM precisions, we build a single _Stream-K_ kernel that has been
specialized per the guidelines in the Section 5. Furthermore, these kernels
implement our “two-tile _Stream-K_ \+ _data-parallel_ ” hybrid decomposition.
Our evaluation compares each _Stream-K_ kernel with:

  1. (1)

the default _data-parallel_ CUTLASS kernel of the same blocking factor;

  2. (2)

the cuBLAS ensemble for that precision (CUDA 11.6); and

  3. (3)

an idealized oracle that will always select the highest performing _data-
parallel_ CUTLASS blocking factor to execute for a given GEMM instance.

For FP64 problems, this oracle selects among the ensemble of
{(32$\times$32$\times$16), (32$\times$64$\times$16), (64$\times$64$\times$16),
(64$\times$128$\times$16), (128$\times$128$\times$16)} blocking factor
specializations. For FP16$\rightarrow$32, it selects among the ensemble of
{(64$\times$64$\times$64), (64$\times$128$\times$32),
(128$\times$128$\times$32), (128$\times$256$\times$32)} blocking factor
specializations. These specific specializations are an open-sourced strict
subsets alternative of the corresponding cuBLAS GEMM kernel ensembles.

(a) CUTLASS FP16$\rightarrow$32 _data-parallel_ “roofline”
performance ($\text{blocking factors}=$ 128$\times$128$\times$32).

(b) cuBLAS (ensemble)

(c) Idealized CUTLASS oracle (ensemble)

(d) _Stream-K_ ($\text{blocking factors}=$ 128$\times$128$\times$32)

Figure 5. FP16$\rightarrow$FP32 GEMM “roofline” performance utilization
landscapes on NVIDIA A100 across 32K GEMM problem shapes and sizes.

(a) CUTLASS _data-parallel_
($\text{blocking factors}=$ 64$\times$64$\times$16)

(b) cuBLAS (ensemble)

(c) Idealized CUTLASS oracle (ensemble)

(d) _Stream-K_ ($\text{blocking factors}=$ 64$\times$64$\times$16)

Figure 6. FP64 GEMM “roofline” performance utilization landscapes on NVIDIA
A100 across 32K problem shapes and sizes.

|  vs. CUTLASS $64\times 64\times 16$ |  vs. cuBLAS |  vs. cuBLAS $>150$ ops/B |  vs. CUTLASS oracle
---|---|---|---|---
Average |  $1.23\times$ |  $1.06\times$ |  $1.03\times$ |  $1.05\times$
StdDev |  $0.45$ |  $0.10$ |  $0.03$ |  $0.09$
Min |  $0.77\times$ |  $0.68\times$ |  $0.99\times$ |  $0.70\times$
Max |  $5.63\times$ |  $2.55\times$ |  $1.24\times$ |  $1.64\times$

Table 1. _Stream-K_ FP64 Relative Performance

|  vs. CUTLASS $128\times 128\times 32$ |  vs. cuBLAS |  vs. cuBLAS $>150$ ops/B |  vs. CUTLASS oracle
---|---|---|---|---
Average |  $1.63\times$ |  $1.13\times$ |  $1.15\times$ |  $1.12\times$
StdDev |  $1.46$ |  $0.45$ |  $0.12$ |  $0.37$
Min |  $0.80\times$ |  $0.64\times$ |  $0.98\times$ |  $0.61\times$
Max |  $14.7\times$ |  $6.74\times$ |  $1.85\times$ |  $4.63\times$

Table 2. _Stream-K_ FP16$\rightarrow$32 Relative Performance

The “roofline” plots of Figure 6(a) and Figure 5(a) highlight the spread of
performance produced by the singleton _data-parallel_ CUTLASS kernels. They
plot the percentage of FP64 and FP16$\rightarrow$32 processor utilization as a
function of computational intensity. Ideally, a GEMM implementation’s
performance response would manifest as a narrow band that adheres tightly to
the machine’s bandwidth- and compute-bound performance ceilings. Here, the
_data-parallel_ kernels exhibit a fairly large dynamic range for any given
regime of arithmetic intensity. In contrast, the performance responses from
the equivalent _Stream-K_ kernels in Figure 6(d) and Figure 5(d) are much
tighter. These observations are corroborated by Table 2 and Table 2, which
show the _Stream-K_ kernels outperforming their _data-parallel_ FP64 and
FP16$\rightarrow$32 equivalents by an average of 1.23$\times$ and
1.63$\times$, respectively. For extreme strong-scaling scenarios where
$m\times n$ is small and $k$ is large, our _Stream-K_ kernels demonstrate up
to 5.63$\times$ and 14.7 $\times$ speedup, respectively.

The second columns of Table 2 and Table 2 compare our _Stream-K_ performance
with that of cuBLAS. On average, our FP64 and FP16$\rightarrow$32 _Stream-K_
GEMM kernels respectively deliver 6% and 13% greater throughput than their
corresponding cuBLAS ensembles, with peak improvement of 2.55$\times$ and
6.74$\times$. This is a significant improvement over the breadth of 32K GEMM
problem shapes and sizes with 20$\times$ _less_ executable code (a single
kernel for each precision) than NVIDIA’s vendor GEMM library, cuBLAS.

Furthermore, the contrast between the FP64 and FP16$\rightarrow$32 cuBLAS
performance responses (Figure 6(b) and Figure 5(b)) versus those of our
hypothetical CUTLASS oracle ensembles (Figure 6(c) and Figure 5(c)) reveal the
difficulties of designing kernel selection heuristics that deliver
consistently good performance. Despite having access to the same blocking
factor specializations, cuBLAS exhibits substantially wider dynamic ranges
than the idealized _data-parallel_ CUTLASS oracle. The performance spreads of
our _Stream-K_ kernels are narrower still, achieving up to 4.6$\times$ the
idealized oracle performance and underscoring their ability to achieve
utilization levels that are simply not possible from tile-centric work
decompositions.

(a) FP64 _Stream-K_ speedup vs. cuBLAS.

(b) FP16$\rightarrow$32 _Stream-K_ speedup vs. cuBLAS.

Figure 7. _Stream-K_ speedup vs. cuBLAS.

Finally, we observe regimes of small, bandwidth-bound problem shapes where our
largish blocking factors do not compete well against cuBLAS. However, if we
restrict our scope to the domain of compute-bound problems (i.e., FP64
problems having compute intensity $>$ 150 ops/byte and FP16 $\rightarrow$ 32
problems $>$ 400 ops/byte), Figure 7(a) and Figure 7(b) demonstrate that our
singleton _Stream-K_ kernels achieve unilaterally higher performance than the
cuBLAS ensembles. The “noisy” relative performance in the regimes below these
thresholds is not surprising, as _Stream-K_ is attempting to make memory-bound
computations run faster by adding more memory workload. This suggests a few
avenues for future work, namely separate cost-modeling for the memory-bound
regime and/or the bundling of a second _Stream-K_ kernel having smaller tile
size into a two-kernel ensemble.

##  7\. Conclusion

We presented _Stream-K_ , a novel parallel workload decomposition technique
for scheduling general matrix multiplication (GEMM) and similar computations
on wide architectures such as GPUs. Unlike other tile-splitting techniques,
the MAC-loop iteration is our unit of workload quantization across processor
cores. This affords excellent strong scaling and workload balancing because
its cost is (1) a constant with respect to the problem shape, and (2)
substantially smaller than that of an entire output tile.

Furthermore, _Stream-K_ produces an $O(p)$ number of splitting seams that are
bound by the number of processor cores. Consequently, the overheads of strong
scaling and workload balancing scale with processor width rather than problem
size. This is a welcome feature for many applications that cannot afford to
allocate large amounts of temporary storage equivalent to the problem output.

Finally, we evaluated our _Stream-K_ approach across a broad spectrum of GEMM
shapes and sizes. We showed that a single blocking configuration of _Stream-K_
can (1) achieve levels of absolute performance that match and/or exceed that
of NVIDIA’s cuBLAS library, even when the latter is operating at near-peak
processor utilization, and (2) do so with much higher levels of performance
consistency. Additionally, _Stream-K_ is an attractive option for library
construction and maintenance, as it presents an opportunity to reduce
distribution sizes by an order of magnitude and removes the need for complex
handcoded heuristics or machine learning models for kernel selection without
compromising performance. _Stream-K_ is open-sourced within CUTLASS 2.11
(https://github.com/NVIDIA/cutlass) and the performance shown within this
paper can be reproduced when compiled using CUDA 11.8.

For future works, we identify cache-aware, tile-access patterns such as Morton
Order, an avenue for optimization. We also believe that _Stream-K_
decomposition could provide a similar improved performance response for other
GEMM-like workloads that struggle with the same quantization inefficiencies.

###### Acknowledgements.

This material is based upon work supported by Sponsor Defense Advanced
Research Projects Agency (DARPA)  under Contract No. Grant #HR0011-18-3-0007.
Any opinions, findings and conclusions or recommendations expressed in this
material are those of the author(s) and do not necessarily reflect the views
of the U.S. Government. Distribution Statement “A” (Approved for Public
Release, Distribution Unlimited). We would like to acknowledge Louis Feng,
Valentin Andrei, Zhongyi Lin and Serban D. Porumbescu for their feedback on
early drafts of the paper.

## References

  * (1)
  * Abdelfattah et al. (2016) Ahmad Abdelfattah, David Keyes, and Hatem Ltaief. 2016\.  KBLAS: An Optimized Library for Dense Matrix-Vector Multiplication on GPU Accelerators.  _ACM Trans. Math. Software_ 42, 3 (June 2016), 1–31.  https://doi.org/10.1145/2818311
  * Barrachina et al. (2008) Sergio Barrachina, Maribel Castillo, Francisco D. Igual, Rafael Mayo, and Enrique S. Quintana-Orti. 2008\.  Evaluation and tuning of the Level 3 CUBLAS for graphics processors. In _2008 IEEE International Symposium on Parallel and Distributed Processing_. IEEE.  https://doi.org/10.1109/ipdps.2008.4536485
  * 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_ 1802.04799 (Feb. 2018).  arXiv:1802.04799
  * Chetlur et al. (2014) Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. 2014.  cuDNN: Efficient Primitives for Deep Learning.  _CoRR_ 1410.0759 (Oct. 2014).  arXiv:cs.NE/1410.0759v1
  * Cui et al. (2010) Xiang Cui, Yifeng Chen, Changyou Zhang, and Hong Mei. 2010\.  Auto-tuning Dense Matrix Multiplication for GPGPU with Cache. In _Proceedings of the 16th International Conference on Parallel and Distributed Systems_ _(ICPADS 2010)_. 237–242.  https://doi.org/10.1109/icpads.2010.64
  * Hagedorn et al. (2020) Bastian Hagedorn, Archibald Samuel Elliott, Henrik Barthels, Rastislav Bodik, and Vinod Grover. 2020.  Fireiron: A Data-Movement-Aware Scheduling Language for GPUs. In _Proceedings of the ACM International Conference on Parallel Architectures and Compilation Techniques_ _(PACT ’20)_.  https://doi.org/10.1145/3410463.3414632
  * Jiang and Snir (2005) Changhao Jiang and Marc Snir. 2005.  Automatic Tuning Matrix Multiplication Performance on Graphics Hardware. In _Proceedings of the 14th International Conference on Parallel Architectures and Compilation Techniques_ _(PACT ’05)_. 185–194.  https://doi.org/10.1109/pact.2005.10
  * Kerr et al. (2017) Andrew Kerr, Duane Merrill, Julien Demouth, and John Tran. 2017.  CUTLASS: Fast Linear Algebra in CUDA C++.  (2017).  https://devblogs.nvidia.com/cutlass-linear-algebra-cuda/
  * Kurzak et al. (2012) Jakub Kurzak, Stanimire Tomov, and Jack Dongarra. 2012\.  Autotuning GEMM Kernels for the Fermi GPU.  _IEEE Transactions on Parallel and Distributed Systems_ 23, 11 (Nov. 2012), 2045–2057.  https://doi.org/10.1109/tpds.2011.311
  * Lam et al. (1991) Monica D. Lam, Edward E. Rothberg, and Michael E. Wolf. 1991\.  The Cache Performance and Optimizations of Blocked Algorithms. In _Proceedings of the Fourth International Conference on Architectural Support for Programming Languages and Operating Systems_ _(ASPLOS IV)_. 63–74.  https://doi.org/10.1145/106972.106981
  * Larsen and McAllister (2001) E. Scott Larsen and David McAllister. 2001.  Fast Matrix Multiplies using Graphics Hardware. In _Proceedings of the 2001 ACM/IEEE Conference on Supercomputing_ _(SC ’01)_. 55:1–55:6.  https://doi.org/10.1145/582034.582089
  * Li et al. (2009) Yinan Li, Jack Dongarra, and Stanimire Tomov. 2009.  A Note on Auto-tuning GEMM for GPUs. In _International Conference on Computational Science_ _(ICCS 2009)_. 884–892.  https://doi.org/10.1007/978-3-642-01970-8_89
  * Mattson et al. (2020) Peter Mattson, Christine Cheng, Gregory Diamos, Cody Coleman, Paulius Micikevicius, David Patterson, Hanlin Tang, Gu-Yeon Wei, Peter Bailis, Victor Bittorf, David Brooks, Dehao Chen, Debo Dutta, Udit Gupta, Kim Hazelwood, Andy Hock, Xinyuan Huang, Daniel Kang, David Kanter, Naveen Kumar, Jeffery Liao, Deepak Narayanan, Tayo Oguntebi, Gennady Pekhimenko, Lillian Pentecost, Vijay Janapa Reddi, Taylor Robie, Tom St John, Carole-Jean Wu, Lingjie Xu, Cliff Young, and Matei Zaharia. 2020\.  MLPerf Training Benchmark. In _Proceedings of Machine Learning and Systems_ , I. Dhillon, D. Papailiopoulos, and V. Sze (Eds.), Vol. 2. 336–349.
  * Nath et al. (2010) Rajib Nath, Stanimire Tomov, and Jack Dongarra. 2010\.  An Improved Magma Gemm For Fermi Graphics Processing Units.  _The International Journal of High Performance Computing Applications_ 24, 4 (Nov. 2010), 511–515.  https://doi.org/10.1177/1094342010385729
  * NVIDIA Corporation (2020) NVIDIA Corporation. 2020\.  CUDA cuBLAS Library (v9.2).  (2020).  http://developer.nvidia.com/cublas.
  * Ragan-Kelley et al. (2013) Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. 2013.  Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines. In _Proceedings of the 34th ACM SIGPLAN Conference on Programming Language Design and Implementation_ _(PLDI ’13)_. 519–530.  https://doi.org/10.1145/2491956.2462176
  * Tan et al. (2011) Guangming Tan, Linchuan Li, Sean Treichle, Everett Phillips, Yungang Bao, and Ninghui Sun. 2011\.  Fast Implementation of DGEMM on Fermi GPU. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ _(SC11)_. Seattle, Washington, 35:1–35:11.  https://doi.org/10.1145/2063384.2063431
  * Tillet and Cox (2017) Philippe Tillet and David Cox. 2017.  Input-Aware Auto-Tuning of Compute-Bound HPC Kernels. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ _(SC17)_. Article 43, 12 pages.  https://doi.org/10.1145/3126908.3126939
  * Tillet et al. (2019) Philippe Tillet, H. T. Kung, and David Cox. 2019.  Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. In _Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages_ _(MAPL 2019)_. 10–19.  https://doi.org/10.1145/3315508.3329973

##  Appendix A Supplementary Material

###  A.1. Analytical Modeling for Stream-K Configuration

In practice, it is not always advantageous to invoke the _Stream-K_
decomposition with as many CTAs as can be actively resident on the GPU.
Because it is a tile-splitting approach, it incurs fixup costs above and
beyond the simple _data-parallel_ decomposition. Consequently, the fundamental
proposition is one of strong scaling: how much additional parallelism can be
expressed before the extra overhead causes a negative return on investment.
Depending on the problem shape, the optimal splitting could be enough to fill
the entire processor (i.e., $g\leftarrow p$), no splitting at all (i.e.,
$g\leftarrow t$), or somewhere in between.

To predict this inflection point, we present a simple approach to model the
runtime of _Stream-K_ as a function of grid size $g$. In the absence of other
work on the GPU, the runtime of the entire _Stream-K_ schedule will be the
same as that of one of its tile-outputting CTAs, which we formulate as
follows:

$\displaystyle time_{CTA}(g)\leftarrow$
$\displaystyle\mathpzc{a}+\mathpzc{b}(FixupPeers(g)>1)$
$\displaystyle+\mathpzc{c}(ItersPerCta(g))+\mathpzc{d}(FixupPeers(g)-1)$

where:

$\displaystyle ItersPerCta(g)\leftarrow$
$\displaystyle\left\lceil\frac{\lceil\frac{m}{\text{BLK\\_M}}\rceil\times\lceil\frac{n}{\text{BLK\\_N}}\rceil\times\lceil\frac{k}{\text{BLK\\_K}}\rceil}{g}\right\rceil$
$\displaystyle FixupPeers(g)\leftarrow$
$\displaystyle\left\lceil\frac{\left\lceil\frac{k}{\text{BLK\\_K}}\right\rceil}{IterationsPerCta(g)}\right\rceil$

This CTA runtime model comprises four components. The $\mathpzc{a}$ workload
encompasses the one-time, fixed-size costs incurred by each CTA, e.g., the
grid launch latency, the initial compulsory cache misses, the cost of writing
the final output tile to C, etc. The second component $\mathpzc{b}$
incorporates the conditional costs of outputting temporary partial sums for
scenarios where the number of output tiles does not quantize perfectly across
the processor. The third—the per-iteration workload $\mathpzc{c}$—represents
the instruction and stall workload of each MAC-iteration. The final, per-
collaborator workload $\mathpzc{d}$ is the cost of reading and accumulating
the partial sums from another CTA covering the same tile. The set of workload
constants {$\mathpzc{a}$, $\mathpzc{b}$, $\mathpzc{c}$, $\mathpzc{d}$} will be
unique to each combination of blocking factors, matrix data type, and GPU
microarchitecture, and can be determined empirically via microbenchmarks.

(a) GEMM $256\times 3584\times 8192$
56 output tiles, 256 iterations per tile
$g_{best}\leftarrow 108$ CTAs, 132/133 iterations per CTA

(b) GEMM $1024\times 1024\times 1024$
64 output tiles, 32 iterations per tile
$g_{best}\leftarrow 64$ CTAs, 32 iterations per CTA

(c) GEMM $128\times 128\times 16384$
1 output tile, 512 iterations per tile
$g_{best}\leftarrow 8$ CTAs, 64 iterations per CTA

Figure 8. Modeled Stream-K performance on NVIDIA A100 (108 SMs) for BLK_M=128,
BLK_N=128, BLK_K=32

Figure 8 illustrates the behavior of our grid size selection model as
parameterized for fp16-precision GEMM on NVIDIA’s A100 GPU using blocking
factors BLK_M $=128$, BLK_N $=128$, and BLK_K $=32$. Specifically, we
highlight three strong-scaling GEMM scenarios where the number of output tiles
is insufficient to produce a single full wave across the processor’s 108 SM
cores.

The first GEMM shape accumulates through a large-sized $k$-dimension to
produce a short, wide output matrix. In this scenario, the reduction in MAC-
loop time relative to the increasing costs of seam fixup is monotonically
improving. Consequently, the optimal grid size coincides with maximal
parallelism at $g=108$ CTAs.

The second shape accumulates through a medium-sized $k$-dimension to produce a
square matrix with 64 output tiles. In this case, the fixup costs of
$\mathpzc{b}$ and $\mathpzc{d}$ outweigh any reduction in MAC-loop iteration
count, as seen by the global minima “dip” at $g=64$ CTAs.

Figure 9. Strong-scaling comparison of _data-parallel_ and _Stream-K_
execution schedules for $128\times 128\times 384$ GEMM across a hypothetical
four-SM GPU. _Data-parallel_ causes the enormous $k$-dimension to be
sequentially processed within single CTA, whereas _Stream-K_ is able to take
advantage of the parallelism available across the $k$-dimension.

The third shape produces a single output tile after accumulating through an
enormous $k$-dimension, analogous to the execution schedule in Figure 9.
Although the opportunity for strong scaling is quite large, the per-peer cost
of serial reduction is entirely incurred by a single CTA. These accumulation
costs begin to outweigh any further reductions in iteration count for grid
sizes $g>8$.