A Framework for Fine-Grained Synchronization of Dependent GPU Kernels

Extracted source text · arXiv 2305.13450

\documentclass[conference]{IEEEtran}

\IEEEoverridecommandlockouts
% The preceding line is only needed to identify funding in the first footnote. If that is unneeded, please comment it out.

\usepackage{pervasives}

\def\BibTeX{{\rm B\kern-.05em{\sc i\kern-.025em b}\kern-.08em
    T\kern-.1667em\lower.7ex\hbox{E}\kern-.125emX}}
\begin{document}

\title{A Framework for Fine-Grained Synchronization of Dependent GPU Kernels}

\author{\IEEEauthorblockN{Abhinav Jangda}
\IEEEauthorblockA{Microsoft Research}
\and
\IEEEauthorblockN{Saeed Maleki}
\IEEEauthorblockA{Microsoft Research}
\and
\IEEEauthorblockN{Maryam Mehri Dehnavi}
\IEEEauthorblockA{University of Toronto}
\and
\IEEEauthorblockN{Madan Musuvathi}
\IEEEauthorblockA{Microsoft Research}
\and
\IEEEauthorblockN{Olli Saarikivi}
\IEEEauthorblockA{Microsoft Research}
}

\maketitle
\thispagestyle{plain}
\pagestyle{plain}

\begin{abstract}
Machine Learning (ML) models contain highly-parallel computations, such as, Matrix Multiplication, Convolutions, Dropout, etc.
These computations are commonly executed on Graphics Processing Units (GPUs), by dividing the computation in independent processing blocks, known as \emph{tiles}.
Since the number of tiles are usually higher than the execution units of a GPU, tiles are executed on all execution units in \emph{waves}.
However, the tiles executed in the last wave can under-utilize the execution units because tiles are not always a multiple of execution units. 
This under-utilization can be reduced by executing multiple \emph{independent} kernels concurrently on a GPU, but is not currently possible for \emph{dependent} kernels.

In this paper, we present \libname{}, a framework to write custom fine-grained synchronization policies for dependent kernels to improve GPU utilization.
\libname{}'s synchronizes tiles instead of kernels, which allows executing tiles of multiple dependent kernels.
Using \libname{} we expressed several synchronization policies in a few lines of code and reduced the inference times of GPT-3 and ResNet-38 by up to 1.19$\times$ and 1.16$\times$ respectively.  
\end{abstract}

\begin{IEEEkeywords}
CUDA, GPU, Synchronization, Machine Learning
\end{IEEEkeywords}

\section{Introduction}
The trend of large Machine Learning (ML) models has delivered remarkable results in multiple domains.
These results have suddenly expanded the demand of ML models in innumerable applications.
As a result, the infrastructure required to serve such large models has significantly increased.
Therefore, even small improvements in the execution of these models can get us huge savings in terms of cost and energy.

ML models typically consist of embarrassingly parallel common operations, such as Generalized Matrix Multiplication (GeMM), 2-D Convolution (Conv2D) etc.
Thus, massively parallel processors such as Graphics Processing Units (GPUs) are ideal for running ML models.
Utilizing the parallelism of GPUs requires breaking down the computations into multiple independent blocks, known as \emph{tiles}.
Each tile is executed by a \emph{thread block}, which runs on an execution unit of the GPU known as \emph{Simultaneous Multiprocessor} (SM).
Each SM can execute one or more thread blocks parallel depending on the resource requirements of the thread blocks.
Often times, not all thread blocks can be executed simultaneously on the limited number of parallel SMs. 
Therefore, the thread blocks are executed in one or more \emph{waves} where each wave utilizes all the SMs.
When thread blocks are not a multiple of SMs, the last wave utilizes less SMs leading to a lower utilization of the GPU.
This under-utilization is prevalent in large ML models.

\begin{table}[t]
  \centering
  \begin{tabular}{|r|r|r|r|r|r|r|}
    \hline
    \textbf{Batch} & \textbf{GeMM} & \textbf{TBs} & \textbf{TBs per SM} &\textbf{Waves} &\textbf{Utilization}\\
                                    % &  &\textbf{}  & \textbf{} &&\\
    \hline
    \multirow{2}{*}{256}& Producer & [1, 96, 2] & 2 & 1.2  & 60\%\\ %1$^{\text{st}}$
                        & Consumer & [1, 96, 1] & 2 & 0.6  & 60\%\\
    % & 1$\times$36$\times$2 & 0.45  & 45\%\\
    \hline
    \multirow{2}{*}{512}& Producer & [2, 48, 2] & 2 & 1.2  & 60\%\\
                        & Consumer & [2, 96, 1] & 2 & 1.2  & 60\%\\
    % & 2$\times$36$\times$2 & 0.9  & 90\%\\
    \hline
    \multirow{2}{*}{1024}& Producer & [4, 48, 1] & 2 & 1.2  & 60\%\\
                         & Consumer & [4, 96, 1] & 2 & 2.4  & 80\%\\
    % & 4$\times$36$\times$2 & 1.8  & 90\%\\
    \hline
  \end{tabular}
  \caption{Number of thread blocks, thread blocks per wave, waves and GPU utilization of two dependent GeMMs in MegatronLM GPT-3~\cite{megatronlm} on several batch sizes when executing on an NVIDIA Tesla V100 containing 80 SMs.
  Each SM can execute 2 thread blocks, i.e., maximum of 160 thread blocks can execute per wave.
  Both GeMM kernels achieve low utilization because number of thread blocks are not a multiple of SMs. \label{tab:intro}}
\end{table}

As an example, consider the utilization of two dependent GeMMs in MegatronLM~\cite{megatronlm} GPT-3 145 Billion parameter model inference.
Table~\ref{tab:intro} shows that two dependent GeMM CUDA kernels of GPT-3 achieves 60--80\% utilization on NVIDIA Tesla V100 because the number of tiles are not the multiple of SMs.
% Moreover, new GPUs released by the hardware manufacturers have an even higher number of SMs, which decreases the utilization even more.
Due to the dependence between the two GeMMs, the consumer kernel is synchronized with the producer kernel across all thread blocks and therefore,
%  two dependent GeMM kernels is to wait for the producer-kernel to finish before invoking the consumer-kernel
making it impossible to overlap the tile executions of the two kernels.
% Thus, in the traditional approach, the thread blocks of the consumer-kernel cannot be executed before all thread blocks of producer-kernel are finished, which inhibits full utilization of GPUs.
Alternatively, a common approach is to find independent kernels and execute them simultaneously~\cite{258921}.
However, inferencing large ML models do not contain enough such kernels~\cite{megatronlm, resnet}.

% There are two existing approaches to achieve high GPU utilization.
% First approach executes independent CUDA kernels simultaneously~\cite{258921}.
% However, existing large ML models do not contain enough independent computations.
A recent method by Stream-K~\cite{streamk} for GeMM mitigates the lower utilization of the last wave by further partitioning 
the tiles in the last wave in order to utilize more SMs.
However, this partitioning requires extra memory accesses by all thread blocks computing the same tile
and not extendable to other tile based computations such as Matrix Dot Product, Dropout, and Softmax.
As we will discuss in Section~\ref{sec:eval}, this method does not improve the utilization of inferencing large ML models.

In this paper, we show that the fine-grained synchronization of dependent tiles of dependent kernels allows the simultaneous execution of all thread blocks, hence, improving GPU utilization.
Furthermore, we show that diverse synchronization techniques provide best performance on different computation kernels and sizes.
To this end, we present, \libname{}, a framework to design efficient synchronization techniques for dependent tile based CUDA kernels.
In \libname{} users can quickly write and experiment with various synchronization techniques, such as synchronizing dependent tiles of $n$-D computations and row synchronization of 2-D computations.
\libname{} contains novel mechanisms to:
(i) ensure that all thread blocks of a producer kernel are executed before the consumer kernel (Section~\ref{sec:invoke-kernel}),
(ii) ensure that the dependence between tiles of producer and consumer-kernels is ensured using semaphores and memory fences (Section~\ref{sec:dependent-tiles}), 
and (iii) ensure that a processing order of producer and consumer tiles that minimizes the wait time(Section~\ref{sec:tile-order}).
% Additionally, we present optimizations that allows optimizations 
We used \libname{} to write novel synchronization techniques for diverse CUDA kernels, such as, GeMM, 2-D Convolutions, and Dot Product, within few lines of code.
\libname{} integrated CUDA kernels reduced the inference time of GPT-3~\cite{megatronlm} by up to 1.19$\times$ and of ResNet-38~\cite{resnet} by up to 1.16$\times$ (Section~\ref{sec:eval}).

\section{Background}
In this section, we discuss NVIDIA GPUs, computations done by Machine Learning models, and the process of executing computations on an NVIDIA GPU.

\subsection{NVIDIA Graphics Processing Units and CUDA}
A CUDA kernel executes a multiple concurrent \emph{threads} organized in a 3-dimensional \emph{grid}, where each thread of the grid has a 3-dimensional unique identifier.
All threads of the kernel are grouped into equally sized \emph{thread blocks}, with each thread block having a 3-dimensional index in the grid.
CUDA defines 3-dimensional indices and sizes with values for $x$, $y$, and $z$-dimension. 
Thus, the grid has thread blocks in $x$, $y$, and $z$-dimensions and the number of thread blocks in a grid is product of thread blocks in all dimensions.

A Graphics Processing Unit (GPU) contains multiple Streaming Multiprocessors (SMs), each of which executes one or more thread blocks.
Since each SM contains a fixed amount of registers and shared memory, there are a fixed number of thread blocks that can execute in parallel on an SM.
This number of thread blocks per SM, known as \emph{occupancy}, depends on the CUDA kernel's register and shared memory usage, and the number of threads in thread block and number of thread blocks in the CUDA kernel's grid. 

\paragraph*{Thread block Wave Execution}
A GPU executes all thread blocks of a 3-dimensional grid in a \emph{row major order}, i.e., first the $x$-dimension, then $y$-dimension, and finally $z$-dimension.
A GPU executes thread blocks on its SMs in one or more \emph{waves}, where each wave executes at maximum $occupancy \times Number\_of\_SMs$ thread blocks to all SMs.
Therefore, the number of waves are $\lceil \frac{Number\_of\_TBs\_in\_Grid}{occupancy \times Number\_of\_SMs} \rceil$.
For instance, consider a CUDA kernel with occupancy of two thread blocks per SM and the kernel is invoked with 240 thread blocks, then executing the kernel on a NVIDIA Tesla V100 containing 80 SMs requires $\lceil \frac{240}{2\times 80} \rceil = 2$ waves.
The first wave executes 160 thread blocks and the second wave executes the remaining 80 thread blocks.

\paragraph*{Stream Synchronization}
A CUDA \emph{stream} contains a sequence of CUDA operations that execute in the order they were issued.
When two CUDA kernels with a producer-consumer relationship are invoked on the same stream, the consumer-kernel is not started before all thread blocks of the producer-kernel has finished their execution. 
We call this synchronization as \emph{stream synchronization}.
However, independent CUDA kernels can be invoked on different streams to execute the kernels concurrently.
A stream also have an associated priority value.
The CUDA driver ensures that operations on a higher priority stream are issued first before the operations on a lower priority stream.

\subsection{Computations in Large ML Models}
Contemporary ML models contains embarrassingly parallel computations, such as Dot product, Generalized Matrix Multiplication (GeMM), 2-D Convolution (Conv2D), Dropout, and Softmax.
These parallel computations are executed on NVIDIA GPUs as CUDA kernels. 
In this paper, we consider two widely used machine learning models: GPT and ResNet.
Below we briefly explain the computations involved in these models and their execution on large clusters of GPUs.
 
\begin{figure}[t]
\begin{subfigure}[b]{\columnwidth}  
\begin{lstlisting}[language=CUDA, xleftmargin=1.0em]
//X: [B, H]; W|\color{commentgreen}{$_1$}|: [H, 4H/8]; W|\color{commentgreen}{$_2$}|: [4H/8, H]
//XW|\color{commentgreen}{$_1$}|: [B, 4*H/8]
XW|$_1$| = GeLU(X |$\times$| W|$_1$|) //1st GeMM fused with GeLU|\label{line:mlp:first-matmul}|
//XW|\color{commentgreen}{$_2$}|: [B, H]
XW|$_{12}$| = XW|$_1$| |$\times$| W|$_2$| //2nd GeMM
\end{lstlisting}
\caption{Multi-Layer Perceptron (MLP) contains two weight matrices: \texttt{W$_1$} of shape $\left[H, \frac{4H}{8}\right]$ and \texttt{W$_2$} of shape $\left[\frac{4H}{8}, H\right]$.
MLP takes an input matrix \texttt{X} of shape $\left[B, H\right]$ to obtain the result \texttt{XW$_{12}$} of same shape.
\label{fig:gpt-3:mlp}}
\end{subfigure}
\begin{subfigure}[b]{\columnwidth}
\begin{lstlisting}[language=CUDA, xleftmargin=1.0em]
//X: [B, H]; QKV: [H, 3H/8]; W|\color{commentgreen}{$_2$}|: [H/8, H]
//XQKV: [B, 3H/8] //1st GeMM
XQKV = X |$\times$| QKV
//XQ: [B, H/8]; XK: [B, H/8]; XV: [B, H/8]
XQ = XQKV[:,0:H/8]     //First matrix slice
XV = XQKV[:,H/8:2*H/8] //Second matrix slice
XK = XQKV[:,2*H/8:]    //Third matrix slice
//XDot: [B, H/8]
XDot = Dropout(Softmax(XQ|$\cdot$|XV))|$\cdot$|XK//Fused Kernel |\label{line:attn:pointwise}|
//XW|\color{commentgreen}{$_2$}|: [B, H]
XW|$_2$| = XDot |$\times$| W|$_2$| //2nd GeMM
\end{lstlisting}
\caption{Self-Attention contains two weight matrices: \texttt{QKV} of shape $\left[\frac{3H}{8}, H\right]$, and \texttt{W$_2$} of shape $\left[\frac{H}{8}, H\right]$.
A Self-Attention layer takes an input matrix \texttt{X} of shape $\left[B, H\right]$ and obtain the result \texttt{XW$_{12}$} of the same shape.
\label{fig:gpt-3:attention}}
\end{subfigure}
\caption{Architecture of two building blocks of GPT models: Multi-Layer Perceptron (MLP) and Self-Attention.
With model parallelism on 8 GPUs, the weight matrices of both layers are equally divided among all 8 GPUs.
B is the Batch size and H is the hidden dimension.
In GPT-3, H is 12288.\label{fig:gpt-3}}
\end{figure}
% \begin{figure*}
%   \begin{subfigure}[b]{0.43\textwidth}
%     \includegraphics[scale=0.39]{figures/MLP.png}
%   \caption{}
%   \end{subfigure}
%   \hspace{1em}
%   \begin{subfigure}[b]{0.56\textwidth}
%     \includegraphics[scale=0.39]{figures/attention.png}
%   \caption{}
%   \end{subfigure}
%   \caption{}
% \end{figure*}
\subsubsection{GPT Models} 
Generative Pre-trained Transformers (GPT) is a class of natural language models that consists of multiple Multi-Layer Perceptron (MLP) and Self-Attention blocks.
Both MLP and Self-Attention contains two weight matrices and takes an input matrix, \texttt{X}.
With model parallelism these weight matrices are divided among all GPUs~\cite{megatronlm}.
Figure~\ref{fig:gpt-3} shows the computations of GPT with model parallelism of 8 GPUs, which is the commonly used setting.
Both MLP and Self-Attention perform the first GeMM of \texttt{X} with their first weight matrix.
Then, they perform several pointwise operations, such as GeLU, Softmax, and Dropout on the result of first GeMM.
State-of-the-art implementations for MLP fuses GeLU activation with the first GeMM (line~\ref{line:mlp:first-matmul} in Figure~\ref{fig:gpt-3:mlp}) and for Self-Attention the dot product and activations are fused in a single CUDA kernel (line~\ref{line:attn:pointwise} in Figure~\ref{fig:gpt-3:attention}).
Finally, they perform the second GeMM with the second weight matrix.

\subsubsection{ResNet} Residual Network (ResNet) is a computer vision model consisting of multiple convolution layers.
Each convolution layer perform two Conv2D operations with same kernel size and channels.
Table~\ref{tab:resnet-sizes} shows the details of each layer.

\begin{table}[t]
  \centering
  \begin{tabular}{|r|r|r|r|r|}
    \hline
    \textbf{Image Size} & \textbf{Kernel} & \textbf{Channels} & \textbf{Convs per Layer} & \textbf{Layers}\\
    \hline
    $\threeDShape{56}{56}{64}$ & $\twoDShape{3}{3}$ & 64  & 2 & 3 \\
    $\threeDShape{28}{28}{128}$ & $\twoDShape{3}{3}$ & 128 & 2 & 4 \\
    $\threeDShape{14}{14}{256}$ & $\twoDShape{3}{3}$ & 256 & 2 & 6 \\
    $\threeDShape{7}{7}{512}$   & $\twoDShape{3}{3}$ & 512 & 2 & 3\\
    \hline
  \end{tabular}
  \caption{Input/output image size, kernel size, and input/output channels for each pair of Conv2D in ResNet-38.
  The last column shows the number of layers for each Conv2D pair. \label{tab:resnet-sizes}}
\end{table}

\subsection{Tile based Computations}
Efficient CUDA kernels of common computations, including GeMM and Conv2D, divides the computation into multiple \emph{tiles}.
Each tile is computed by one or more thread blocks and all threads of a thread block compute one or more elements of the tile. 
This decomposition enables computing tiles independent of each other, which leads to high-parallelism.

Figure~\ref{fig:motivation:tiled-matmul} shows an example workflow of a computing GeMM, $C_{[M, N]} = A_{[M, K]} \times B_{[K, N]}$, with tile size $[T_M, T_N]$ on a GPU.
A tile, $C^{xy}$, represents a 2-D sub-matrix of $C$ starting at row $x$ and column $y$ of $C$, and is of the same size as the tile size, i.e., $\left[T_M, T_N\right]$.
Therefore, the number of tiles in the row dimension are $\frac{M}{T_M}$ and in column dimension are $\frac{N}{T_N}$.
% The 3-D index of thread blocks in grid enables direct assignment of 2-D tiles to thread blocks.
Each tile is computed by one thread block and the CUDA kernel is invoked with a 3-D grid $\left[ \frac{M}{T_M}, \frac{N}{T_N}\right]$, which is same as the number of tiles in row and column dimensions.
To compute $C^{xy}$, the thread block multiplies $T_M$ rows starting at $x$ row of $A$ and $T_N$ columns starting at $y$ column of $B$.

High-performance CUDA libraries, such as NVIDIA CUTLASS and CUBLAS, uses the above tile based methodology for GeMM.
Additionally, these libraries computes Convolutions using GeMM to utilize the optimized GeMM kernels.
In this paper, we focus on these tile based computations and their implementation in NVIDIA CUTLASS.

\section{Motivation}
In this section, we show how the traditional stream synchronization can lead to under-utilization and how our \emph{fine-grained synchronization} can ensure maximum utilization of GPU resources.

Let us consider two dependent GeMMs on three matrices: $A$ of shape $\twoDShape{12}{8}$, $B$ of shape $\twoDShape{8}{8}$, and $D$ of shape $\twoDShape{8}{8}$:
\begin{gather*}
\text{producer-gemm:  } C_{[12, 8]} = A_{[12, 8]} \times B_{[8, 8]}\\
\text{consumer-gemm:  } E_{[12, 8]} = C_{[12, 8]} \times D_{[8, 8]}
\end{gather*}
The producer-gemm computes $C$ by multiplying $A$ and $B$, and the consumer-gemm computes $E$ by multiplying $C$ and $D$.
Let us also assume that the CUDA kernel for both GeMMs has a tile size of $[4, 4]$ and each tile is computed by one thread block.
Thus, both gemm kernels are invoked with a 3-D grid of size $\twoDShape{\frac{12}{4}}{\frac{8}{8}}$, i.e., six thread blocks.

Figure~\ref{fig:motivation:current-scenario} shows the execution of thread blocks of both kernels on a GPU with four SMs.
The GPU executes all six thread blocks of both kernels in $\lceil\frac{6}{4}\rceil = 2$ waves.
The producer-gemm kernel executes four thread blocks in the first wave and after these thread blocks are finished it executes the remaining two thread blocks in the second wave.
Similarly, the consumer-gemm kernel executes its thread blocks in two waves.
With stream synchronization between kernels, thread blocks of the consumer-gemm kernel are executed only after the producer-gemm kernel is finished.
Hence, in the second wave of both kernels only \texttt{SM-0} and \texttt{SM-1} are utilized while \texttt{SM-2} and \texttt{SM-3} are idle.
In summary, the stream synchronization between producer- and consumer-kernels can lead to under-utilization of GPU resources. 

In this paper, we present a novel \emph{fine-grained synchronization} of thread blocks, thus synchronizing dependent tiles, of producer- and consumer-kernels to improve the utilization of GPU resources.
Figure~\ref{fig:motivation:fine-grained synchronization} shows how we obtain full utilization in our synthetic example.
We invoke both kernels on separate streams to execute thread blocks of both kernels simultaneously, and synchronize only the dependent thread blocks using a semaphore stored in GPU's memory.
In our example, computing an $i^{\text{th}}$ row of $E$ requires computing the $i^{\text{th}}$ row of $C$, i.e., $E^{i, 0}$ and $E^{i,1}$ depends on $C^{i, 0}$ and $C^{i, 1}$. 
To maintain this dependency, each thread block of the consumer-gemm waits until all thread blocks of the corresponding row of producer-gemm has been computed.
In our example, both $E^{i, 0}$ and $E^{i,1}$ waits until $C^{i, 0}$ and $C^{i, 1}$ has been computed.
Since both kernels are running on different CUDA streams, thread blocks of both kernels are executed in only three waves as compared to four waves in stream synchronization.
% Hence, in only three waves all SMs of the GPU are utilized.
Hence, with fine-grained synchronization between thread blocks of producer and consumer-kernels can improve the utilization of all SMs on a GPU.

In summary, stream synchronization of CUDA kernels can suffer from under-utilization of GPU resources, while our fine-grained synchronization of tile based CUDA kernels improves the utilization of GPU resources.
We now show how that executing ML models using stream synchronization leads to under utilization of GPU resources. 

\begin{figure*}[ht]
  \centering
  \begin{subfigure}[b]{0.31\textwidth}
    \includegraphics[scale=0.6]{figures/tile-sync.png}
    \caption{\textit{Tiled GeMM kernels}: Tiles $C_{i,j}$ are computed by multiplying sub-matrices $A_{i}$ and $B_{j}$.
    Similarly, tiles $E_{i,j}$ are computed by multiplying $C_{i,0}$ and $C_{i, 1}$ with $D_{j}$.
    Since each tile is computed by one thread block, the tile size of $4\times 4$ gives the grid size of $\{\frac{12}{4}, \frac{8}{4}\} = \{3, 2\}$ for both kernels.
    Both kernels have the occupancy of 1 thread block per SM.
      \label{fig:motivation:tiled-matmul}} 
  \end{subfigure}
  \hfill{}
  \begin{subfigure}[b]{0.31\textwidth}
  \includegraphics{figures/current-scenario.png}
  \caption{
    \textit{Stream Synchronization}:
  Thread blocks of both producer ($C_{i,j}$) and consumer ($E_{i,j}$) are executed in two waves.
  Four thread blocks are executed in the first wave and remaining two in the second wave.
  Stream synchronization is used to synchronize between producer- and consumer-kernel.
  In the second wave of both kernels, SM-3 and SM-4 are idle, leading to under-utilization.
  \label{fig:motivation:current-scenario}}
  \end{subfigure}
  \hfill{}
  \begin{subfigure}[b]{0.32\textwidth}
  \includegraphics[scale=0.97]{figures/orchestration.png}
  \caption{
    \textit{Fine-grained Synchronization}:
  We execute both kernels on different streams to allow simultaneous execution of all thread blocks in only three waves.
  Arrows shows the dependence between thread blocks.
  Thread blocks of consumer-kernel waits using a semaphore until its producer-kernel's thread block has computed dependent tile. 
  Since in every wave all SMs are utilized, we achieve full utilization.
  \label{fig:motivation:fine-grained synchronization}}
  \end{subfigure}
  \caption{Thread block execution with existing stream synchronization and fine-grained fine-grained synchronization on 4 SMs for a GeMM kernel to compute two dependent GeMMs: $C_{12\times 8} = A_{12\times 8} \times B_{8\times 8}$ and $E_{12\times 8} = C_{12\times 8} \times B_{8\times 8}$.\label{fig:motivation}}
\end{figure*}

\subsection{Under-utilization in Training of GPT-3}
The training and inference of large ML models suffer from under-utilization due to stream synchronization among dependent CUDA kernels.
As a motivation, consider both GeMMs of MLP in GPT-3 with \texttt{float16} weights, hidden dimension \texttt{H}$ = 12288$, and batch size \texttt{B}$ = 1024$.
Figure~\ref{fig:gpt-3:mlp} shows the GeMM sizes in MLP for GPT-3 are: \texttt{XW$_1$} of shape $\left[1024, 6144\right]$ and \texttt{XW$_{12}$} of shape $\left[1024, 12288\right]$.
Hand-written GeMM CUDA kernels in NVIDIA CUTLASS use the tile size for both GeMMs to $\left[256, 128\right]$.
The GeMM kernel has an occupancy of 2 thread blocks per SM and when executed on an NVIDIA Tesla V100 GPU with 80 SMs, each wave can execute up to 160 thread blocks.

Based on the tile sizes, the producer-gemm is invoked with $\twoDShape{4}{48} = 192$ thread blocks and requires $\lceil\frac{192}{160} \rceil = 2$ waves to execute all thread blocks.
While the first wave executes 160 thread blocks, the second wave can only execute the remaining $192 \% 160 = 32$ thread blocks.
Thus, the second wave under-utilizes GPU resources.
Similarly, the consumer-gemm is invoked with $\twoDShape{4}{96} = 384$ thread blocks and the last wave executes $384 \% 160 = 64$ thread blocks.
Therefore, thread block execution of both kernels is done in $2+3 = 5$ waves using stream synchronization.

With our fine-grained synchronization, thread blocks of both kernels are executed in only $\lceil\frac{192+384}{160} \rceil = 4$ waves, which is one less than with stream synchronization.
Moreover, with fine-grained synchronization the last wave improves utilization of GPU by executing $(192+384)\% 160 = 96$ thread blocks, which is far more than the thread blocks executed in the last wave of both kernels.
That is why, our fine-grained synchronization obtains higher utilization of the GPU and provides 1.31$\times$ speedup over the stream synchronization. 

In the next section, we present, \libname{}, a framework to write synchronization techniques among diverse kernels.
Using \libname{} we show that diverse fine-grained synchronization techniques can improve the GPU utilization of ML models.

\section{Fine-Grained Synchronization using \libname{}}
% \begin{figure}[ht]
%   \begin{lstlisting}[language=CUDA]
% class OrkConsumer{|\label{line:api:consumer-start}|
%   device dim3 tile() {
%     return OrkHandle::tile(consTileOrderBuff,
%                            consTileCounter);}
%   device void wait(dim3 prodTile) {|\label{line:api:cons:wait}|
%     if(threadIdx == {0,0,0})
%       syncVal = SyncMap(prodTile);
%       while(tileStatusBuff[syncVal] != |\label{line:api:cons:while}|
%             expectedVal);
%     __syncthreads();|\label{line:api:cons:syncthreads}|
%   }
%   cudaStream stream();
% };|\label{line:api:consumer-end}|
% class OrkProducer{|\label{line:api:producer-start}|
%   device setWaitSemaphore();
%   device dim3 tile(){
%     return OrkHandle::tile(prodTileOrderBuff, 
%       prodTileCounter);}
%   device void post(dim3 prodTile) {|\label{line:api:prod:post}|
%     __syncthreads();  |\label{line:api:prod:post:syncthreads}|
%     if(threadIdx == {0,0,0})
%       __threadfence_system();|\label{line:api:prod:post:threadfence}|
%       syncVal = SyncMap(prodTile);
%       atomicAdd(&tileStatusBuff[syncVal],1);|\label{line:api:prod:setval}|
%   }
%   cudaStream stream();
% };|\label{line:api:producer-end}|
% class OrkHandle{|\label{line:api:handle-start}|
%   OrkHandle(dim3 prodGrid, dim3 prodTiles,
%             dim3 consGrid, dim3 consTiles,
%             SyncMap& sync, SchedMap& sched);
%   void setTileWaitVal(unsigned int val);|\label{line:api:handle:setTileWaitVal}|
%   void waitKernel();
%   dim3 tile(dim3* tileOrderBuff, 
%             unsigned int* tbCounter) {|\label{line:api:tile-start}|
%     shared dim3 _tile;
%     if (threadIdx == {0, 0, 0})
%       _tile = tileOrderBuff
%         [atomicAdd(tbCounter)];
%     __syncthreads();
%   }|\label{line:api:tile-end}|
%   OrkProducer<SyncMap>& prod();
%   OrkProducer<SyncMap>& emptyP();
%   OrkConsumer<SyncMap>& cons();
%   OrkConsumer<SyncMap>& emptyC();
% };|\label{line:api:handle-end}|
% \end{lstlisting}
% \caption{The \libname{} API: \texttt{OrkHandle} handles the fine-grained synchronization for each producer-consumer kernel pair.
% The producer kernel can set the status of each of one or more tiles using \texttt{OrkProducer}'s \texttt{post} method. 
% The consumer kernel can wait for tiles of the producer kernel to be processed using \texttt{OrkConsumer}'s \texttt{wait}.
% \texttt{dim3} is a CUDA struct to represent 3-dimensional values.\aj{Change Ork} 
% \label{fig:api}}
% \end{figure}

Our fine-grained synchronization of dependent CUDA kernels consists of three mechanisms: (i) executing thread blocks of all kernels on SMs without coarse-grained stream synchronizations, (ii) adding fine-grained synchronization among the dependent tiles of every producer-consumer kernel pair, and (iii) controlling the order of tile processing in each kernel to minimize the wait time of synchronization. 
We have implemented these mechanisms in a header-only standalone CUDA library, we call as \libname{}. 
% Figure~\ref{fig:api} shows the functions of \libname{} that abstracts the implementation details of our approach and can be used in diverse tile based computations.
% \texttt{OrkHandle} defines functions for fine-grained synchronization each pair of producer-consumer kernels.
% \texttt{OrkProducer} contains methods for the producer kernel and \texttt{OrkConsumer} contains methods for the consumer kernel.
%We now explain all three mechanisms in detail.

We explain these mechanisms through the MLP~(Figure~\ref{fig:gpt-3:mlp}) running example with two dependent GeMM kernels.
Figure~\ref{fig:ex:matmuls} shows an implementation of MLP that uses \libname{} for fine-grained synchronization.
The {\tt gemm} function is a standard generalized matrix multiplication kernel (we use CUTLASS for our experiments) with additional code (shown as \underline{underlined}) to call into \libname{}. Specifically, \libname{} associates each kernel with a {\tt CuStage} object that supports fine-grained synchronization facilities among kernels. The {\tt MLP} function creates these stage objects, declares dependencies between them, and invokes the kernels.  

\begin{figure*}[ht]
  \begin{subfigure}[b]{0.5\textwidth}
  \begin{lstlisting}[language=CUDA, xleftmargin=3.0ex]
//CUDA Kernel to compute C = A * B
global void gemm(f16* A, f16* B, f16* C, |\label{line:ex:kernel-start}| 
                    int K, |\underline{CuStage stage}|) {
  |\underline{stage.start();}| |\label{line:ex:init-prod-cons}|
  |\underline{row, col = stage.tile();}| |\label{line:ex:get-tile}|
  for (tk = 0; tk < K; tk += TileK) {
    |\underline{stage.wait({A, row, tk});}| |\label{line:ex:tile-sync}|
    LoadTileToShMem(Ash, A, row, tk); |\label{line:ex:tileA}|
    |\underline{stage.wait({B, col, tk});}| |\label{line:ex:tile-sync-B}|
    LoadTileToShMem(Bsh, B, col, tk); |\label{line:ex:tileB}|
    MultiplyAccumulate(C, Ash, Bsh, 
                        row, col, tk);}|\label{line:ex:mac}|
  |\underline{stage.post({row, col});}| |\label{line:ex:post}||\label{line:ex:kernel-end}|
}
void MLP(int BS, int H, f16* X, f16* W|$_1$|, |\label{line:ex:mlp-start}|
          f16* XW|$_1$|, f16* W|$_2$|, f16* XW|$_{12}$|) {
  grid1 = {B, 4*H/8}/tile1; |\label{line:ex:grid1}|
  grid2 = {B, H}/tile2; |\label{line:ex:grid2}|
  CuSync cs;
  CuStage prod = cs.stage(grid1,tile1,RowMajor);|\label{line:ex:alloc-prod}|
  CuStage cons = cs.stage(grid2,tile2,RowMajor);|\label{line:ex:alloc-cons}|
  // declare prod to cons[XW1] dependency
  cs.dependency<RowSync>(prod, cons, XW1); |\label{line:ex:declare-dep}|
  // invoke the producer gemm
  gemm<<<grid1, tb1, prod.stream()>>>
    (X, W|$_1$|, XW|$_1$|, H, prod);|\label{line:ex:mlp-kernel-1}|
  // invoke waitKernel and then consumer 
  cons.waitKernel(); |\label{line:ex:wait-kernel}| 
  gemm<<<grid2, tb2, cons.stream()>>>
    (XW|$_1$|, W|$_2$|, XW|$_{12}$|, 4*H/8, cons);} |\label{line:ex:mlp-kernel-2}|
  \end{lstlisting}
  \caption{
  Synchronize two GeMM kernels using \libname{}.
  The kernels are invoked on different streams. 
  The wait kernel ensures the order of kernel invocation. The {\tt post} and {\tt wait} methods ensure tile dependency.
  Changes to the GeMM kernel are \underline{underlined}.
  \label{fig:ex:matmuls}}
  \end{subfigure}
  \hfill{}
  \begin{subfigure}[b]{0.48\textwidth}
  \begin{lstlisting}[language=CUDA, xleftmargin=3.0ex]
void wait_till(sem, expected) {|\label{line:api:cons:wait}|
  if(threadIdx == {0,0,0})
    while(*sem != expected);|\label{line:api:cons:while}|
  __syncthreads();}|\label{line:api:cons:syncthreads}|

void post(sem) {|\label{line:api:prod:post}|
  __syncthreads();  |\label{line:api:prod:post:syncthreads}|
  if(threadIdx == {0,0,0})
    __threadfence_system();|\label{line:api:prod:post:threadfence}|
    atomicAdd(sem,1);}|\label{line:api:prod:setval}|


class TileSync
  void init(int* sem) {
  // Create semaphore for each tile
  }
  void wait(row, col){|\label{line:tilesync:wait}|
    wait_till(
      &sem[row*numCols + col], 1);}

  void post(row, col){|\label{line:tilesync:post}|
    post(&sem[row*numCols + col])}

class RowSync
  void init(int* sem){
  // Create semaphore for each row    
  }
  void wait(row, col){|\label{line:rowsync:wait}|
    if(col == 0)
      wait_till(sem[row], numCols);}
  
  void post(row, col) {|\label{line:rowsync:post}|
    post(&sem[row]);}
  \end{lstlisting}
  % tile.z*grid.y*grid.x+tile.x*grid.y+
  \caption{TileSync creates a semaphore for each tile. 
  RowSync trades concurrency for synchronization by creating a single semaphore per row.\label{fig:ex:tilesync}\label{fig:ex:rowsync}}
  \end{subfigure}
  \caption{Fine-grained synchronization of two GeMMs of MLP using TileSync and RowSync policies.
  \label{fig:ex}}
  \end{figure*}

\subsection{Invoke Dependent Kernels}
\label{sec:invoke-kernel}
The first part of fine-grained synchronization is to eliminate the stream synchronization between kernels. 
\libname{} achieves this by invoking the kernels on different CUDA streams. \libname{} associates a {\tt CuStage} object with each kernel. Lines~\ref{line:ex:alloc-prod} and~\ref{line:ex:alloc-cons} in Figure~\ref{fig:ex:matmuls} creates a producer {\tt prod} and consumer {\tt cons} stage for the two GeMM kernels. These kernels are invoked on different streams associated with respective stages. Before doing so, the code declares the dependency between the two stages in line~\ref{line:ex:declare-dep} specifying that input {\tt A} of the consumer depends on the output of the producer. 
Section~\ref{sec:dependent-tiles} describes how \libname{} enforces this dependency. 

\subsection{Stage Processing Order}
\label{sec:invoke-order}
CUDA runtime lacks a straightforward mechanism to enforce a scheduling order among kernels belonging to different streams. Therefore, there is a possibility that the consumer kernel is scheduled to execute on the GPU before the producer kernel. This can lead to poor performance as thread blocks of the consumer kernel occupy SMs without any useful work to do. In the worst case, this can lead to deadlocks if there are no SMs available for thread blocks of the producer kernel. 

To avoid this, \libname{} provides a mechanism for enforcing the scheduling order of kernels. In this example, the consumer stage has a \texttt{wait}-kernel invoked in line~\ref{line:ex:wait-kernel} in Figure~\ref{fig:ex:matmuls}. Wait kernels are invoked with a single thread associated with stream of the consumer stage. This thread waits on semaphores stored in global memory for each dependent kernel using a busy-wait while loop. \libname{} requires each kernel to call the {\tt stage.start()} method (line~\ref{line:ex:init-prod-cons}) that sets this semaphore. Once the wait kernel exits, the CUDA runtime will schedule the kernel associated with the stage. In our example, this mechanisms ensures none of the thread blocks of the consumer kernel are scheduled before at least one of the thread blocks of the producer kernel. 


\subsection{Tile Processing Order}
\label{sec:tile-order}
The CUDA runtime can schedule thread blocks to SMs in an arbitrary order, which can lead to unpredictable wait times in dependent kernels. Ideally, we want the thread blocks of the consumer kernel to be scheduled in the order producer kernels generate the tiles. \libname{} enforces this order as follows. Each kernel calls into {\tt stage.tile()}~(line~\ref{line:ex:get-tile}) to compute the tile it needs to compute next. Internally, \libname{} maintains a global counter per stage that determines the tile computation order. In the example, the template parameter {\tt RowMajor} in line~\ref{line:ex:alloc-prod} and line~\ref{line:ex:alloc-cons} ensures that the two kernels produce tiles in row major order independent of how the CUDA runtime schedules thread blocks. \libname{} supports other tile orders such as column major and strided row major.     

% We enforce an ordering on thread blocks, consequently on tiles, as follows.
% (1) \texttt{OrkHandle} takes a schedule function \texttt{SchedFn} that maps a linear index to a 3-dimensional index in the space of kernel grid.
% During its initialization, the \texttt{OrkHandle} creates a \texttt{tileOrderBuff} buffer based on the schedule function that defines the order of thread blocks.
% (2) \texttt{OrkProducer} has access to a global counter of producer tiles, \texttt{prodTileCounter}.
% Similarly, \texttt{OrkConsumer} has access to counter of consumer tiles, \texttt{consTileCounter}.
% (3) The \texttt{tile} method of \texttt{OrkProducer} and \texttt{OrkConsumer} returns the next tile index at the start of the thread block.
% The method \texttt{tile} atomically increments their respective tile counters and returns the thread block index stored at the current value of counter (lines~\ref{line:api:tile-start}--\ref{line:api:tile-end}).

% In summary, ... \aj{a transition?}

\subsection{Synchronizing Dependent Tiles}
\label{sec:dependent-tiles}
The second mechanism required for fine-grained synchronization is to ensure the dependence between tiles of producer- and consumer-kernels because we have now eliminated the stream synchronization between the kernels. This is accomplished by invoking the {\tt wait} (line~\ref{line:ex:tile-sync} and line~\ref{line:ex:tile-sync-B}) and {\tt post} (line~\ref{line:ex:kernel-end}) methods of the {\tt CuStage} object associated with the consumer kernel. To compute a tile of GeMM, the GeMM kernel depends on a row of tiles from input {\tt A} and a column of tiles from input {\tt B}. For the consumer kernel, the input row tiles are generated by the producer kernel. This dependency is declared in line~\ref{line:ex:declare-dep}. Accordingly, the {\tt wait} before loading a tile of {\tt A} waits for the corresponding {\tt post} of the producer kernel. Since there are no depenences for loading the tiles of {\tt B}, the corresponding {\tt wait} becomes a no-op. Similarly, both {\tt wait}s are no-ops for the producer kernel. Figure~\ref{fig:ex:tilesync} describes the implementation of the {\tt post} and {\tt wait} methods, which will be described in detail below. 

% Figure~\ref{fig:api} shows the implementation of three functions to maintain the dependence between tiles: (i) the \texttt{post} method of \texttt{OrkProducer} sets the completion status of a tile of producer-kernel (line~\ref{line:api:prod:post}),
% (ii) the \texttt{wait} method of \texttt{OrkConsumer} stop the further execution until the status of producer-kernel's tile reaches an expected value (line~\ref{line:api:cons:wait}), and 
% (iii) the \texttt{setTileWaitVal} to set the expected value of semaphore to wait on (line~\ref{line:api:handle:setTileWaitVal}).
% The \texttt{wait} and \texttt{post} methods has access to \texttt{tileStatusBuff}, which is an array of semaphores for each tile and is stored in the global memory.
% Since tiles are represented as a multi-dimensional index, the API takes a device function parameter, \texttt{SyncFn} that maps 3-dimensional tile to a linear index in \texttt{tileStatusBuff}.

\subsection{Synchronization Policies}
\label{sec:synchronization-policies}
\libname{} allows implementation of diverse synchronization policies.
Each synchronization policy is implemented as a CUDA class with three functions: (i) \texttt{init} allocates an array of semaphore, (ii) \texttt{post} increments the status of semaphore for given tile, and (iii) \texttt{wait} waits until the value of semaphore for the tile reaches an expected value. 
Below we discuss three policies we implemented in \libname{} in few lines of code based on our workloads.

\paragraph*{TileSync}
We refer to the policy of synchronizing on the semaphore of each tile of the producer-kernel as \emph{TileSync} policy.
To minimize the wait time of each tile of the consumer-kernel, both kernels computes their tiles in a row major order.
Figure~\ref{fig:ex:tilesync} shows the implementation of this policy in \libname{} in a few lines of code.
Both \texttt{wait} and \texttt{post} methods works on a single semaphore for each tile (line~\ref{line:tilesync:wait} and \ref{line:tilesync:post}).
For example, in Figure~\ref{fig:ex:matmuls} to compute a tile $E^{xy}$, the TileSync policy requires waiting first on $C^{x0}$ and $C^{x1}$.

\paragraph*{RowSync}
Computing tile $E^{xy}$ in Figure~\ref{fig:ex:matmuls} using TileSync requires waiting for 2 tiles of $C$.
In general, for two dependent GeMM, TileSync requires $\frac{N}{T_n}$ number of synchronizations, where $N$ is columns of $C$ and $T_n$ is the column tile.
For sufficiently high value of $N$, the high number of synchronizations can become a bottleneck.
We can reduce the number of synchronizations for tile $E^{ij}$ by synchronizing over dependent rows, i.e., share the semaphore of all tiles computing the $i^{th}$ row of $C$.
Figure~\ref{fig:ex:rowsync} shows the implementation of this policy in CUDA.
The \texttt{post} function increments the semaphore for the row of given tile (line~\ref{line:rowsync:post}) and \texttt{wait} function only waits for the row of tile (line~\ref{line:rowsync:wait}).
To minimize the wait time of each tile of the consumer-kernel, both kernels computes their tiles in the row major order.

\paragraph*{StridedSync}
Figure~\ref{fig:use-case:self-attention} shows the dependence of tiles between the three kernels of Self-Attention of GPT-3.
The result of first GeMM, \texttt{XQKV}, is equally divided along the columns into three matrices: \texttt{XQ}, \texttt{XK}, and \texttt{XV}.
The dot product kernel performs the dot product of \texttt{XQ}, \texttt{XK}, and \texttt{XV}, and softmax on the output of dot product, to return matrix \texttt{XDot}.
Hence, the $i^{th}$ column tile of \texttt{XDot} depends on the $i^{th}$ column tile of \texttt{XQ}, \texttt{XK}, and \texttt{XV}.
Moreover, each row tile of \texttt{XW$_{12}$} depends on all tiles of \texttt{XDot} computing the same row.

Figure~\ref{fig:use-case:self-attention:strided} shows the implementation of \emph{StridedSync} policy for synchronizing the first GeMM and the dot product kernel.
In StridedSync, all tiles of the first GeMM that are \texttt{stride} apart shares the same semaphore.
Moreover, the \texttt{prodOrder} function orders tiles of the first GeMM that are \texttt{stride} apart first and then next group of tiles.
For Self-Attention, the value of stride is $\frac{\text{Column Tiles of }\texttt{XQKV}}{3} = \frac{H}{8\times T_y}$, where $T_y$ is the tile size.

\paragraph*{Conv2DTileSync}
The implicit GeMM kernel performs the convolution by converting the input image and convolution matrix to a matrix representation in shared memory.
Our \emph{Conv2DTileSync} synchronizes a tile $\left[i,j\right]$ of the output matrix $\left[i,\frac{j}{K\times K}\right]$ tile of input image, where $K\times K$ is the size of convolution matrix.

\subsubsection{Synchronization Implementation}
We now discuss the implementation of semaphores.

\textbf{Post}: 
The \texttt{post} method of a policy obtains a semaphore and then call \texttt{post} on the semaphore.
In Figure~\ref{fig:ex:tilesync}, the \texttt{post} method first performs a \texttt{\_\_syncthreads} to ensure that all threads of the thread block has computed the tile and has issued store instructions to write computed tile elements to the global memory (line~\ref{line:api:prod:post:syncthreads}).
Then the method executes a memory fence to ensure that all global memory writes will be visible to threads of the other kernel (line~\ref{line:api:prod:post:threadfence}).
Finally, the method increments the semaphore value by 1.\\

\textbf{Wait}:
The \texttt{wait} method of a policy obtains a semaphore and then call \texttt{wait\_till} on the semaphore.
In Figure~\ref{fig:ex:tilesync}, the \texttt{wait\_till} function reads the value of semaphore in a while loop using only the first thread of the thread block (line~\ref{line:api:cons:while}).
While the first thread is waiting on the semaphore, all other threads of thread block waits on the \texttt{\_\_syncthreads} (line~\ref{line:api:cons:syncthreads}).
When the value changes to the expected value, the first thread reaches the \texttt{\_\_syncthreads} and all threads of the thread-block reads from the producer tile.

% This is how \libname{} creates custom tile processing order.
% In section~\ref{sec:}, we show different implementations of the schedule function. 

% We will present more policies for Self-Attention and ResNet in Section~\ref{sec:use-cases}.

% \subsection{Fine-Grained Synchronization of GPT-3 MLP}
% In this section, we show an example \aj{...}

% Figure~\ref{fig:ex} shows fine-grained synchronization of two dependent matmuls of MLP shown in Figure~\ref{fig:gpt-3:mlp}.
% The function \texttt{MLP} takes 8 arguments that contains the usual six arguments: batch size (\texttt{B}), hidden dimension (\texttt{H}), and global addresses of the dataset (\texttt{X}), weights (\texttt{W}$_1$ and \texttt{W}$_2$), and result matrices (\texttt{XW}$_1$ and \texttt{XW}$_{12}$), with two additional arguments for the synchronization and scheduling policy functions
% (line~\ref{line:ex:mlp-start}--\ref{line:ex:mlp-end}).
% The \texttt{MLP} function first determines the grid size for both matmul kernels based on the tile sizes (line~\ref{line:ex:grid1}--\ref{line:ex:grid2}) and constructs the \texttt{OrkHandle}.
% Then, the function invokes the producer matmul kernel on the producer-stream to compute \texttt{XW}$_1$ and the wait-kernel on the consumer stream (line~\ref{line:ex:mlp-kernel-1}--\ref{line:ex:mlp-wait}).
% Finally, the function calls the consumer matmul kernel on the consumer-stream to compute \texttt{XW}$_{12}$ (line~\ref{line:ex:mlp-kernel-2}).

% The matmul CUDA kernel is modified to perform the fine-grained synchronization (lines~\ref{line:ex:kernel-start}--\ref{line:ex:kernel-end}).
% In addition to the four arguments of memory address of input and output matrices, the kernel takes two extra arguments: \texttt{prod} to set the tile status in the producer-matmul, and \texttt{cons} to wait on a producer-tile status in the consumer-matmul.
% The producer-matmul works as follows.
% First, the producer-matmul set the wait-kernel's semaphore (line~\ref{line:ex:init-prod-cons}), so that, the wait-kernel can exit and the CUDA runtime can execute thread blocks of the consumer-matmul.
% Then, each thread block of the producer-matmul obtains the tile of first output matrix, \texttt{XW}$_1$, to compute using the \texttt{tile()} function (line~\ref{line:ex:get-tile}).
% Finally, when the producer-matmul has computed a tile (lines~\ref{line:ex:mac}) and it sets the completion status of the tile (line~\ref{line:ex:post}).
% The consumer-matmul works similar to the producer-matmul except instead of setting the status of each tile, either
% the consumer-matmul waits for each tile of \texttt{XW}$_1$ using the TileSync policy (line~\ref{line:ex:tile-sync}), or waits for completion of the row of \texttt{XW}$_1$ using the RowSync policy (line~\ref{line:ex:row-sync}).

\begin{figure}[t]
  \begin{subfigure}[b]{0.5\textwidth}
  \centering
  \includegraphics[scale=0.8]{figures/self-attention.png}
  \caption{Dependency between tiles of three kernels of Self-Attention. 
  Tile \texttt{P1} of \texttt{XDot} depends on first tiles of \texttt{XQ}, \texttt{XK}, and \texttt{XV}, i.e., \texttt{T1}, \texttt{T3}, and \texttt{T5}.
  Similarly, \texttt{P2} depends on \texttt{T2}, \texttt{T4}, and \texttt{T6}.
  Each tile of second GeMM, i.e., \texttt{X1}--\texttt{X3}, depends on all tiles of row of \texttt{XDot}, i.e, \texttt{P1}--\texttt{P2}.
  \label{fig:use-case:self-attention:dependencies}}
  \end{subfigure}

  \begin{subfigure}[b]{0.48\textwidth}
    \begin{lstlisting}[language=CUDA, xleftmargin=3.0ex]
template<int stride> class StrideSync
  void init(){
    // numRows * numCols/stride semaphores
  }
  void wait(row, col){
    if(col%stride == 0)
      wait_till(&sem[row*numCols/stride 
      + col%stride], stride);
  }
  void post(row, col) {
    increment(&sem[row*numCols/stride 
      + col%stride]);
  }

  int prodOrder(grid, tile){
    tile.y = (tile.y%stride)*stride + 
      tile.y/stride;
    return RowMajor(grid, {tile.x, tile.y});
  }
    \end{lstlisting}
    \caption{StridedSync policy creates a semaphore for tiles that are stride apart.
    \label{fig:use-case:self-attention:strided}}
  \end{subfigure}
  \caption{Tile dependence between all three CUDA kernels of Self-Attention and the StridedTileSync policy in \libname{} for synchronizing the first GeMM and dot product.
  The dot product and second GeMM are synchronized using TileSync.\label{fig:use-case:self-attention}}
\end{figure}

\section{Optimizations}
\label{sec:opt}
We present two optimizations that the user can do to improve the performance of fine-grained synchronization.

\subsubsection{Avoid Wait Kernel} \label{sec:opt:wait}
The wait-kernel is required to ensure that all thread blocks of the producer-kernel are scheduled on the GPU before the consumer-kernel.
However, we can avoid invoking the wait-kernel if thread blocks of both producer- and consumer-kernel can be scheduled in a single wave.
This condition is satisfied when the sum of thread blocks of both kernels is less than or equal to the minimum occupancy of both kernels.
% Table~\ref{tab:impact-opt} shows that avoiding wait kernel can improve the performance for small computation kernels.

\subsubsection{Reorder Tile Loads and Synchronization} \label{sec:opt:reorder}
The general workflow of tile based CUDA kernels is to load tile of each input and then perform operations on the tile.
We can re-order the waiting of tile of one input with the loading of other tile, to overlap the waiting of one tile with the loading of the other input's tile.

For example, in Figure~\ref{fig:ex:matmuls} the second GeMM kernel loads a tile of both inputs (\texttt{A} and \texttt{B}) and compute the tile of output matrix (\texttt{C}) (line~\ref{line:ex:tile-sync}--\ref{line:ex:tileB}).
We can reorder the loading of \texttt{B} tile with the waiting on \texttt{A} tile, i.e., swap lines~\ref{line:ex:tile-sync}--\ref{line:ex:tileA} with lines~\ref{line:ex:tile-sync-B}--\ref{line:ex:tileB}.
Since there is no waiting for tile of \texttt{B}, loading \texttt{B} tile can overlap with waiting of \texttt{A} tile, thus, improving performance.

\iffalse
\section{Use Cases}
\label{sec:use-cases}
In this section, we use \libname{} to synchronize Self-Attention kernels of GPT-3 and Conv2Ds in ResNet.

\subsubsection{Self-Attention}
\label{sec:use-cases:attn}
\aj{Madan: Simplify this as the description of StridedTileSync is above.}

We use \libname{} to synchronize all three dependent kernels of Self-Attention in Figure~\ref{fig:gpt-3:attention}, i.e., the first GeMM, the dot product, and the second GeMM.
Figure~\ref{fig:use-case:self-attention} shows the tile dependence of all three kernels.
The result of first GeMM, \texttt{XQKV}, is equally divided along the columns into three matrices: \texttt{XQ}, \texttt{XK}, and \texttt{XV}.
The dot product kernel performs the dot product of \texttt{XQ}, \texttt{XK}, and \texttt{XV}, and softmax on the output of dot product, to return matrix \texttt{XDot}.
Hence, the $i^{th}$ column tile of \texttt{XDot} depends on the $i^{th}$ column tile of \texttt{XQ}, \texttt{XK}, and \texttt{XV}.
Moreover, each row tile of \texttt{XW$_{12}$} depends on all tiles of \texttt{XDot} computing the same row.

We can use RowSync and TileSync to synchronize all three kernels.
However, to decrease wait time of the product kernel, we developed \emph{StridedTileSync} policy, for synchronizing the first GeMM and the dot product kernel.
Figure~\ref{fig:use-case:self-attention:strided} shows the implementation of StridedTileSync policy in \libname{}, where the $i^{th}$ column tile of \texttt{XDot} waits on $i^{th}$ column tiles of \texttt{XQ}, \texttt{XK}, and \texttt{XV} by sharing the same semaphore.
Since there are $\frac{H}{8}$ columns in \texttt{XDot}, the \texttt{stridedSync} function returns the same semaphore for all tiles that are $\frac{H}{8\times Ty}$ apart (line~\ref{line:ex:tilesyncmap}).
Moreover, the policy orders the tile such that each pair of strided tiles are processed consecutively by the GeMM kernel while the row major order is followed for the dot product kernel.

\subsubsection{Consecutive Conv2Ds}
The implicit GeMM kernel performs the convolution by converting the input image and convolution matrix to a matrix representation in shared memory.
Hence, we can straightforwardly use the same RowSync and TileSync to synchronize dependent Conv2D kernels.
\fi

\section{Evaluation}
\label{sec:eval}
In this section, we evaluate the performance of \libname{} against state-of-the-art baselines on diverse computations.

\paragraph{Baseline and Dataset}
We consider computations in two most widely used machine learning models:
(i) MegatronLM's GPT-3 145 Billion~\cite{megatronlm} parameter model that contains GeMM, Dot Products, Softmax, and Dropout (Figure~\ref{fig:gpt-3}), and 
(ii) ResNet-38 model that contains different Conv2Ds (Table~\ref{tab:resnet-sizes}).
We consider multiple batch sizes from 1 to 2048 for GPT-3 and 1 to 32 for ResNet-38~\cite{resnet}.
We use the CUDA kernels of GeMM and Conv2D in NVIDIA CUTLASS 2.13 and we implemented a fused pointwise CUDA kernel that performs dot product, softmax, and dropout in Self-Attention.
We modified these kernels to do fine-grained synchronization using \libname{} and compare \libname{} against two baselines:
\begin{itemize}
  \item \textbf{StreamSync} the traditional uses stream synchronization between dependent kernels.
  \item \textbf{Stream-K}~\cite{streamk} divides GeMM workload of the last thread block wave equally among all SMs to improve utilization.
  Stream-K divides the GeMM computation into two kernels.
  The first kernel computes tiles for thread blocks executed in the full waves using blocked GeMM.
  The second kernel computes remaining tiles by partitioning the summation dimension across all SMs.
  Each thread block of the second kernel computes it tile-partition and update the partial tile stored in the global memory.
  Therefore, Stream-K performs extra memory accesses for last remaining tiles than the traditional GeMM.
\end{itemize}
\paragraph{Sync. Policies} We consider following synchronization policies and optimizations cases:
\begin{itemize}
  \item \textbf{RowSync} synchronizes dependent rows (Figure~\ref{fig:ex:rowsync}).
  \item \textbf{TileSync} synchronizes dependent tiles (Figure~\ref{fig:ex:tilesync}).
  \item \textbf{TileSync+W} extends TileSync by avoiding the wait-kernel for kernels with low grid sizes (Section~\ref{sec:opt:wait}).
  \item \textbf{TileSync+WR} extends TileSync+W by reordering the tile
  synchronization on first input with loads of other inputs (Section~\ref{sec:opt:reorder}). 
  \item \textbf{Strided+TileSync+WR}, only for Self-Attention, synchronizes the first GeMM with dot product using StridedSync, and the dot product with second GeMM using TileSync (Figure~\ref{fig:use-case:self-attention}).
  The policy also avoids wait-kernel and reorder tile loads.
  \item \textbf{Conv2DTileSync} synchronizes dependent tiles of Conv2Ds. 
  \item \textbf{Conv2DTileSync+W} and \textbf{Conv2DTileSync+WR} apply avoid wait-kernel and reordering tile loads optimizations.
\end{itemize} 

\begin{table}[t]
  \centering
  \begin{tabular}{|r|r|r|r|r|}
    \hline
    \textbf{Kernel} & \textbf{Implementation} & \multicolumn{2}{c|}{\textbf{Lines Changed}}\\
    \cline{3-4}   
          &         & \textbf{Number} & \textbf{Fraction}\\\hline %Lines of Code & Lines Changed &
    GeMM & CUTLASS & 25& 0.5\%\\
    Dot-Product & Ours & 5& 1\%\\
    Conv2D & CUTLASS & 22& 0.6\%\\
    \hline
  \end{tabular}
  \caption{Fraction of lines of code changed in GeMM, Dot-Product, and Conv2D kernels to enable fine-grained synchronization using \libname{}.\label{tab:loc}}
\end{table}

\begin{figure*}[th]
  \begin{subfigure}[b]{0.30\textwidth}
    \includegraphics[scale=0.65]{results/mlp.pdf}
    \caption{Synchronizing both GeMMs in MLP\label{fig:results:gpt-3:mlp}}
  \end{subfigure}
  \hfill{}
  \begin{subfigure}[b]{0.45\textwidth}
    \includegraphics[scale=0.65]{results/attention.pdf}
    \caption{Synchronizing GeMMs and dot-product in Self-Attention.\label{fig:results:gpt-3:attention}}
  \end{subfigure}
  \begin{subfigure}[b]{0.20\textwidth}
    \includegraphics[scale=0.65]{results/gptend2End.pdf}
    \caption{GPT-3 Inference\label{fig:results:gpt-3:e2e}}
  \end{subfigure}
  \caption{Speedup of \libname{}'s policies for synchronizing CUDA kernels in MLP and Self-Attention over StreamSync, and speedup of GPT-3 inference for batch sizes 1--2048.
  Numbers shows the maximum speedup of all policies.\label{fig:results:gpt-3}}
\end{figure*}
\begin{figure}
  \centering
  \includegraphics[scale=0.7]{results/gptstreamk.pdf}
  \caption{Speedup of best \libname{} policy of Self-Attention and MLP over StreamK for different batch sizes.\label{fig:results:streamk}}
\end{figure}

\begin{table*}
  \centering
%  \begin{tabular}{|r|r|r|r|r|r|r|r|r|}
  % \hline
    % Batch  & \multicolumn{2}{c|}{First Matmul} & \multicolumn{2}{c|}{Second Matmul} & Seq. & \multicolumn{2}{c|}{Ork Waves} & Best\\
    % \cline{2-3} \cline{4-5} \cline{7-8}
    %       & \multicolumn{1}{c|}{Grid.x}&Waves & \multicolumn{1}{c|}{Grid.x} & Waves & Waves & ? & ? & Sync\\ \hline
    % 1--64 & 1&0.3&1 &0.2& 2 & 0.5 & 1 & Tile\\
    % 128   & 1&0.4&1 &0.4& 2 & 0.8 & 1 & Tile\\
    % 256   & 1&1.2&1 &0.6& 3 & 1.8 & 2 & Tile\\
    % 512   & 2&1.2&2 &1.2& 4 & 2.4 & 3 & Row\\
    % 1024  & 4&1.2&4 &2.4& 5 & 3.6 & 4 & Row\\
    % 2048  & 8&2.4&8 &4.8& 8 & 7.2 & 8 & Row\\
    % \hline
  % \end{tabular}
  \begin{tabular}{|r|r|r|r|r|r|r|r|r|r|r|}
  \hline
  \textbf{Batch}  & \multicolumn{2}{c|}{\textbf{First GeMM}} & \multicolumn{2}{c|}{\textbf{Second GeMM}} & \multicolumn{2}{c|}{\textbf{StreamSync}} & \multicolumn{3}{c|}{\textbf{\libname{}}} & \textbf{Speedup}\\
    \cline{2-3} \cline{4-5} \cline{6-7} \cline{8-10}
          & \multicolumn{1}{c|}{\textbf{Grid}} & \textbf{Waves} & \multicolumn{1}{c|}{\textbf{Grid}} & \textbf{Waves} & \textbf{Waves} & \textbf{Time($\mu$s)} &\textbf{Waves} & \textbf{Policy} & \textbf{Time($\mu$s)} & \\ \hline
    1--64 & 1$\times$24$\times$3&0.3&1$\times$48$\times$1 &0.2& 2 & 445& 0.5 & Tile & 411&1.04--1.07$\times$\\
    128   & 1$\times$48$\times$2&0.4&1$\times$96$\times$1 &0.4& 2 & 780& 0.8 & Tile & 730&1.07$\times$\\
    256   & 1$\times$96$\times$2&1.2&1$\times$96$\times$1 &0.6& 3 & 951& 1.8 & Tile & 840&1.12$\times$\\
    512   & 2$\times$48$\times$2&1.2&2$\times$96$\times$1 &1.2& 4 & 1651& 2.4 & Row  & 1450&1.13$\times$\\
    1024  & 4$\times$48$\times$1&1.2&4$\times$96$\times$1 &2.4& 5 & 2832& 3.6 & Row  & 2451&1.31$\times$\\
    2048  & 8$\times$48$\times$1&2.4&8$\times$96$\times$1 &4.8& 8 & 4838& 7.2 & Row  & 4262&1.12$\times$\\
    \hline
  \end{tabular}
\caption{Grid size and the number of waves for each GeMM, total waves and execution time in stream synchronization and in \libname{} for both GeMMs of MLP for batch size from 1 to 2048.
The grid $x$ and $y$-dims are obtained by dividing the size of GeMM with the tile size and the $z$-dim is the number of thread blocks used for split-k.\label{tab:waves:mlp}}
\end{table*}



\paragraph*{Experimental Setup} We run our experiments on a machine containing a 2.60GHz 12-core Intel Xeon CPU E5-2690 v4 with 448GB RAM and an NVIDIA Tesla V100 GPU with 32 GB memory.
All kernels are compiled with CUDA 11.2.
We report the average time of 20 executions after a warmup of 5 executions.
The execution time only contains running time of CUDA kernels and do not include any host and device memory transfer time.

  % \begin{subtable}[t]{\columnwidth} 
  % \begin{tabular}{|r|r|r|r|r|r|r|r|r|}
  %   \hline
  %   % Computation  & MLP  & Attention\\
  %   % Matmul & XW$_1$ = [BS, 12288] $\times$ [12288, 6144] & XQKV = [BS, 12288] $\times$ [12288, 4608]\\
  %   % Matmul & [BS, 6144] $\times$ [6144, 12288] & \\ 
  %   % [BS, 1536] $\times$ [1536, 12288]\\
  % \end{tabular}
  % \end{subtable}
% \begin{table}
%   \begin{subtable}[t]{\columnwidth}
%     \begin{tabular}{l}
% X: [BS, 12288]; W$_1$: [12288, 6144]; W$_2$: [6144, 12288]\\
% XW$_1$ = X $\times$ W$_1$ \\
% XW$_2$ = XW$_1$ $\times$ W$_2$\\
% \\
% X: [BS, 12288]; Q: [12288, 1536]; K: [12288, 1536]; V: [12288, 1536]; W$_2$: [1536, 12288]\\
% XQKV = X $\times$ QKV\\
% X = Softmax(XQ $\dot{}$ XK)\\
% X = Dropout(X) $\cdot{}$ XV\\
% X = X $\times$ W$_2$
%     \end{tabular}
%     \caption{MLP}
%   \end{subtable}
%   \caption{Dataset sizes of GPT-3 175B and ResNet-38 Models \label{tab:dataset}}
% \end{table}

\subsection{Ease of Programming}

Table~\ref{tab:loc} shows that the number of lines added and changed to support fine-grained synchronization of GeMM, Conv2D, and Dot Product kernels using \libname{} are negligible compared to the lines of code of these kernels.
Thus, \libname{} approach enables diverse synchronizations of tile based computation kernels through few modifications.

\subsection{GPT-3 Inference Results}
In this section, we evaluate reduction in inference times of GPT-3 using \libname{}.
We first show the speedup of \libname{}'s synchronization of computation kernels of MLP and Self-Attention, and then discuss the reduction in end-to-end inference time of GPT-3.
Figure~\ref{fig:results:gpt-3} shows our inference results of GPT-3 over StreamSync.

\subsubsection{MLP Results} 
Figure~\ref{fig:results:gpt-3:mlp} shows that synchronizing two dependent GeMMs using \libname{} in MLP decreases the combined execution time of both GeMMs by 1.04$\times$--1.31$\times$ for different batch sizes.
We discuss these improvements using Table~\ref{tab:waves:mlp} that shows the number of waves for each batch size for both StreamSync and \libname{}.
TileSync+WR performs best for 1 to 256 batch sizes because there is a single thread block in the $x$-dimension of grid as shown in Table~\ref{tab:waves:mlp}.
The speedup at batch size of 256 is higher than small batch sizes because the GPU can execute both GeMMs in 2 waves using TileSync+WR while StreamSync requires 3 waves.
On small batch sizes, even though the number of waves are not decreased, TileSync+WR still provides a speedup of up to 1.07$\times$ because the second GeMM can overlap the loading of \texttt{W$_2$} tile in shared memory with the computation of first GeMM.

RowSync performs best for sizes greater than 256 because synchronizing over a row once than multiple tiles reduces memory accesses, and multiple rows provide more opportunities for row synchronization.
Moreover, increasing the number of rows also increases the speedup of RowSync from 1.13$\times$ at 512 to 1.31$\times$ at 1024.
However, the speedup decreases to 1.12$\times$ at 2048 because the fraction of waves reduced by \libname{} decreases with more thread blocks in the grid.

We measured the time spent in invocation of a kernel is around 6$\mu$s.
The difference in execution time of StreamSync and \libname{} in Table~\ref{tab:waves:mlp} is significantly higher than the kernel invocation.
Therefore, the speedups are significantly higher than what would be only achieved by only overlapping the invocation of consumer-kernel with producer-kernel's execution.

\subsubsection{Self-Attention Results}
Figure~\ref{fig:results:gpt-3:attention} shows that synchronizing both GeMMs and the dot-product kernel of Self-Attention using \libname{} provides up to 1.12$\times$ speedup over StreamSync for different batch sizes.
Similar to the MLP, both tile synchronization policies, Strided+TileSync+WR and TileSync+WR, performs better than RowSync for batch sizes up to 256
and RowSync performs better than other policies for batch sizes more than 256.
For batch sizes up to 256, Strided+TileSync+WR obtains 1.15$\times$ speedup over StreamSync and performs better than TileSync+WR because the former policy combines synchronizations of three tiles into a single synchronization.
RowSync achieves the maximum speedup of 1.12$\times$ at 1024 batch size.

\subsubsection{End-to-End GPT-3 Inference Results} 
We integrated \libname{} synchronized MLP and Self-Attention in Megatron-LM~\cite{megatronlm} to evaluate the reduction in inference times of GPT-3.
Figure~\ref{fig:results:gpt-3:e2e} shows that our fine-grained synchronization decreases the inference times by 1.05$\times$--1.19$\times$.
 
\subsubsection{Comparison with Stream-K}
Since Stream-K improves the utilization of individual GeMMs, we compare the performance of \libname{} against Stream-K's GeMM kernels in MLP and Self-Attention.
Figure~\ref{fig:results:streamk} shows that best policy of \libname{} obtains up to 1.31$\times$ speedup over Stream-K.

The speedup of \libname{} over Stream-K is due to three reasons.
First, computing a tile using multiple thread blocks in Stream-K requires extra memory accesses than computing a tile using single thread block.
Whereas \libname{} perform a single write to post the status of computed tile and single read to wait on the status of producer-tile.
Second, dividing the work of last wave among all SMs require smaller tile sizes that decreases the locality and hence, the performance of individual thread blocks, while with \libname{} the tile sizes are not needed to be changed.
Third, Stream-K is not straightforward to apply to other tile based kernels other than GeMM, such as the dot-product kernel in Self-Attention, while \libname{} is valid for any tile based kernels.

In summary, \libname{} provides significant speedups over StreamSync and Stream-K in GPT-3, and there is no one policy that provides best speedup for all cases.

\begin{table}
  \begin{tabular}{|r|r|r|r|r|r|r|r|r|}
   \hline
   \textbf{B} & \textbf{grid.x} & \textbf{Conv} & \multicolumn{2}{c|}{\textbf{StreamSync}} & \multicolumn{3}{c|}{\textbf{\libname{}}}\\% & Speedup\\
     \cline{4-5} \cline{6-8}
     &  &\textbf{Waves}& \textbf{Waves} & \textbf{Time} & \textbf{Waves} & \textbf{Policy}& \textbf{Time}\\
     \hline
     1  & 13  & 0.24 & 2& 56& 0.48 & Tile & 48\\%&1.11$\times$\\
     4  & 49  & 0.61 & 2& 93& 1.22 & Tile & 77\\%&1.20$\times$\\
     8  & 98  & 0.61 & 2& 151& 1.22 & Row &  132\\%&1.15$\times$\\
     12 & 147 & 0.91 & 2& 160&1.84 & Row &  150\\%&1.07$\times$\\
     16 & 196 & 1.22 & 4& 210&2.45 & Row &  196\\%&1.12$\times$\\
     20 & 245 & 1.53 & 4& 262&3.06 & Row &  234\\%&1.07$\times$\\
     24 & 294 & 1.83 & 4& 290&3.68 & Row & 267\\%&1.09$\times$\\
     28 & 343 & 2.14 & 6& 332&4.3 & Row & 316\\%&1.06$\times$\\
     32 & 392 & 2.45 & 6& 369&4.9 & Row & 350\\%&1.05$\times$\\
     \hline
  \end{tabular}
\caption{Number of $x$-dim thread blocks and waves of individual Conv2D kernels, waves and execution time in $\mu$s of StreamSync and of \libname{} of two Conv2D with 128 channels for different batch sizes. \label{tab:waves:conv-128}}
\end{table}
\begin{figure*}[th]
  \begin{subfigure}[b]{0.78\textwidth}  
    \includegraphics[scale=0.75]{results/resnet.pdf}
    \caption{\libname{} on both Conv2Ds of each layer \label{fig:results:conv}}
  \end{subfigure}
  \hfill{}
  \begin{subfigure}[b]{0.21\textwidth}  
    \includegraphics[scale=0.75]{results/resnetEnd2End.pdf}
    \caption{ResNet-38 Inference \label{fig:results:e2e-resnet}}
  \end{subfigure}
  \caption{Speedup of \libname{} policies for both Conv2D kernels of each layer over StreamSync in ResNet-38 and the inference speedup of ResNet-38 for batch sizes from 1 to 32.\label{fig:results:resnet}}
\end{figure*}
\begin{table}[H]
  \centering
  \begin{subtable}[t]{\columnwidth}
    \centering
    \begin{tabular}{|r|r|r|r|r|r|r|r|r|}
      \hline
      \textbf{B} & \textbf{TBs} & \multicolumn{3}{c|}{\textbf{TileSync}}\\
      \cline{3-5}
        & & \textbf{Vanilla} & \textbf{TileSync+R} & \textbf{TileSync+WR} \\\hline
       1--64 & 120 & 433 &  418 & 411\\
       128   & 192 & 780 &  736 & 730\\ \hline
     \end{tabular}
     \caption{Execution times of TileSync of GeMM kernels in MLP with and without optimizations for different batch sizes. 
     Maximum thread blocks per wave for both kernels are 240.\label{tab:impact-opt:mlp}}
  \end{subtable}

  \begin{subtable}[t]{\columnwidth}
    \centering
  \begin{tabular}{|r|r|r|r|r|r|r|r|r|}
    \hline
    \textbf{Channels} & \textbf{B} & \textbf{TBs} & \multicolumn{3}{c|}{\textbf{Conv2DTileSync}} \\
    \cline{4-6}
    & & & \textbf{Vanilla} & \textbf{+W} & \textbf{+WR} \\\hline
    64       & 1  & 98 & 45 & 41   & 37\\
    128      & 1  & 156& 56 & 50   & 45\\
    256      & 1  & 128&61  & 56   & 51\\ \hline
\multirow{2}{*}{512} & 1 & 96 &94 &89    & 85\\ 
                     & 4 & 256&128&120   & 115\\\hline
  \end{tabular}
  \caption{Execution times of Conv2DTileSync of ResNet's Conv2D with and without optimizations for channels and batch sizes.
  Maximum thread blocks per wave for both kernels are 320.\label{tab:impact-opt:resnet}}
\end{subtable}
\caption{Execution times in $\mu$s of adding optimizations over TileSync policy in MLP and Conv2DTileSync in ResNet for smaller grid sizes. \emph{TileSync+W} avoids the wait-kernel and \emph{TileSync+WR} also reorders the tile loading with tile synchronization. Similarly, \emph{Conv2DTileSync+W} avoids the wait-kernel and \emph{Conv2DTileSync+WR} also reorders the tile loading with tile synchronization.\label{tab:opt-impact}}
\end{table}
\subsection{ResNet-38 Inference Results}
We now evaluate \libname{} on reducing the inference times of ResNet-38 by synchronizing both Conv2D kernels of each layer of ResNet-38.
Figure~\ref{fig:results:resnet} shows the speedup of \libname{} over StreamSync for synchronizing both Conv2D kernels of each layer in ResNet.
We first discuss the performance of each pair of Conv2Ds using RowSync and TileSync policies, and then show the end-to-end ResNet-38 inference results.

\subsubsection{Individual Layer Results}
Figure~\ref{fig:results:conv} shows that synchronizing two Conv2D kernels using \libname{} provides 1.02$\times$--1.20$\times$ speedup over StreamSync for different channels and batch sizes.
We notice following two interesting trends, which we explain using the number of grids and waves for Conv2D with 128 channels shown in Table~\ref{tab:waves:conv-128}.

First, for each channel, the speedup increases to a maximum possible value with larger batch size
because larger batch size invokes more number of thread blocks, which leads to higher fraction of waves reduced by \libname{}.
However, the speedup decreases after the maximum value because further increasing the number of thread blocks eventually decreases the fraction of waves reduced by \libname{}.
For example, for 128 channels, Table~\ref{tab:waves:conv-128} shows that RowSync obtains maximum reduction in the number of waves, i.e., 33\%, which leads to maximum speedup of 1.12$\times$ for batch size 16.

Second, similar to MLP and Self-Attention, Conv2DTileSync performs better than RowSync for smaller batch sizes while RowSync performs better for larger batch sizes.
For 128 channels, Table~\ref{tab:waves:conv-128} shows that batch size 8 and more invokes high number of thread blocks in $x$-dimension of grid leading to better speedup by RowSync.
However, for 64 channels, both Conv2DTileSync and RowSync performs similarly because there is a single tile in the $y$-dimension.

\subsubsection{End to End Inference}
We integrated \libname{} version of Conv2D kernels of each layer in ResNet-38 to evaluate the decrease in inference times with vanilla ResNet-38 on different batch sizes.
Figure~\ref{fig:results:e2e-resnet} shows that the integration of ResNet-38 obtains up to 1.16$\times$ speedup over the vanilla version.

\subsection{Impact of Optimizations}
We now discuss the performance improvements provided by the optimizations on the top of TileSync for ResNet and MLP.
Table~\ref{tab:impact-opt:mlp} shows that both TileSync+W and TileSync+WR has lower execution times than TileSync for kernels with low thread blocks.
While the reordering of tile loads and waiting provides larger improvement, avoiding wait kernel also reduces some execution time. 

% \subsection{Code changes}
% \subsection{Performance Analysis}
% Is SyncFn giving improvements?\\
% How to show the improvement? There is no way to see the improved utilization of combined kernels. So, how do we show that the improvement is because of the utilization?\\
% Show that everything else is the same: global memory accesses, register usage, and occupancy.
% Is adding register usage too high
% \cite{megatronlm}

\section{Related Work}
Several works have focussed on efficient software-based synchronization between threads of the same CUDA kernel for irregular GPU applications~\cite{10.1145/2751205.2751232, 10.1145/2903150.2903155, 10.1145/3297858.3304055}.
Li et. al.~\cite{10.1145/2751205.2751232} developed an approach for inter-thread synchronizations by reassembling the micro-instructions of shared memory atomic operations in an efficient way.
Kai et. al.~\cite{10.1145/3297858.3304055} presented a hierarchical approach to synchronization for irregular applications by synchronizing thread blocks using global memory and threads of a thread block using shared memory.
Xu et. al.~\cite{10.1145/2903150.2903155} present a lock design
that uses lock stealing to avoid deadlocks.
\textsc{CoCoNet}~\cite{10.1145/3503222.3507778} performs synchronization between computation and communication kernel to overlap the communication transfers with computation.
\libname{} targets synchronization between threads of multiple CUDA kernels and provides an abstraction to easily design several synchronization policies, both of these are missing from above works.

Moreover, some works have focussed on hardware-supported synchronization primitives for inter-kernel threads.
GLocks~\cite{10.1109/IPDPS.2011.87} is the first hardware supported implementation for highly-contented locks using message passing.
HQL~\cite{10.1109/IPDPS.2013.82} is a hardware-accelerated fine-grained lock scheme for GPUs, which adds support for queuing locks in L1 and L2 caches and uses a customized communication protocol for faster lock transfer and reduced lock retries.
ElTantway~\cite{8327023} et. al. propose a hardware warp scheduling policy that reduces lock retries by de-prioritizing warps whose threads are spin waiting.
They also propose a hardware mechanism for accurately detecting busy-wait synchronization on GPUs.
Dalmia~\cite{9933620} et. al. designed multi-level barrier and priority mechanisms for semaphores for GPU based synchronization primitives. 
\libname{} is a software solution for synchronization threads of multiple CUDA kernels and these hardware-supported mechanisms are complementary to \libname{}.
% to the approaches presented in this paper.

Lingqi et. al.~\cite{9139854} studied the performance and pitfalls of several CUDA synchronization methods for reduction operation.
Sinclair et. al.~\cite{8167781} presented a benchmark suite to measure the performance of synchronization primitives for different coherence protocols and consistency models.

Stream-K~\cite{streamk} is a GeMM implementation that improves the utilization of SMs of a GPU by dividing the workload among all SMs.
However, Stream-K is not straightforward to apply to computations other than GeMMs. 
In contrast, \libname{} fits thread blocks of multiple kernels in each wave and is applicable to any tile based computations.

\section{Conclusion}
Machine learning models contains several dependent computations.
Each computation cannot completely utilize the GPU but it is not possible to independently execute dependent computations.
In this paper, we show that fine-grained synchronization of dependent computations can improve the performance by higher utilization and more optimizations.
In future, we will apply our techniques on better computation and communication overlap techniques. 

\bibliographystyle{plain}
\bibliography{paper}
\end{document}
\begin{thebibliography}{10}

\bibitem{10.1109/IPDPS.2011.87}
Jose~L. Abell\'{a}n, Juan Fern\'{a}ndez, and Manuel~E. Acacio.
\newblock Glocks: Efficient support for highly-contended locks in many-core
  cmps.
\newblock In {\em Proceedings of the 2011 IEEE International Parallel \&
  Distributed Processing Symposium}, IPDPS '11, pages 893--905, USA, 2011. IEEE
  Computer Society.

\bibitem{9933620}
Preyesh Dalmia, Rohan Mahapatra, Jeremy Intan, Dan Negrut, and Matthew~D.
  Sinclair.
\newblock Improving the scalability of gpu synchronization primitives.
\newblock {\em IEEE Transactions on Parallel and Distributed Systems},
  34(1):275--290, 2023.

\bibitem{8327023}
Ahmed ElTantawy and Tor~M. Aamodt.
\newblock Warp scheduling for fine-grained synchronization.
\newblock In {\em 2018 IEEE International Symposium on High Performance
  Computer Architecture (HPCA)}, pages 375--388, 2018.

\bibitem{resnet}
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
\newblock Deep residual learning for image recognition.
\newblock {\em CoRR}, abs/1512.03385, 2015.

\bibitem{10.1145/3503222.3507778}
Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein~Nodehi Sabet, Saeed
  Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi.
\newblock Breaking the computation and communication abstraction barrier in
  distributed machine learning workloads.
\newblock In {\em Proceedings of the 27th ACM International Conference on
  Architectural Support for Programming Languages and Operating Systems},
  ASPLOS '22, pages 402--416, New York, NY, USA, 2022. Association for
  Computing Machinery.

\bibitem{10.1145/2751205.2751232}
Ang Li, Gert-Jan van~den Braak, Henk Corporaal, and Akash Kumar.
\newblock Fine-grained synchronizations and dataflow programming on gpus.
\newblock In {\em Proceedings of the 29th ACM on International Conference on
  Supercomputing}, ICS '15, pages 109--118, New York, NY, USA, 2015.
  Association for Computing Machinery.

\bibitem{258921}
Lingxiao Ma, Zhiqiang Xie, Zhi Yang, Jilong Xue, Youshan Miao, Wei Cui,
  Wenxiang Hu, Fan Yang, Lintao Zhang, and Lidong Zhou.
\newblock Rammer: Enabling holistic deep learning compiler optimizations with
  {rTasks}.
\newblock In {\em 14th USENIX Symposium on Operating Systems Design and
  Implementation (OSDI 20)}, pages 881--897. USENIX Association, November 2020.

\bibitem{streamk}
Muhammad Osama, Duane Merrill, Cris Cecka, Michael Garland, and John~D. Owens.
\newblock Stream-k: Work-centric parallel decomposition for dense matrix-matrix
  multiplication on the gpu.
\newblock In {\em Proceedings of the 28th ACM SIGPLAN Annual Symposium on
  Principles and Practice of Parallel Programming}, PPoPP '23, page 429–431,
  New York, NY, USA, 2023. Association for Computing Machinery.

\bibitem{megatronlm}
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper,
  and Bryan Catanzaro.
\newblock Megatron-lm: Training multi-billion parameter language models using
  model parallelism, 2020.

\bibitem{8167781}
Matthew~D. Sinclair, Johnathan Alsop, and Sarita~V. Adve.
\newblock Heterosync: A benchmark suite for fine-grained synchronization on
  tightly coupled gpus.
\newblock In {\em 2017 IEEE International Symposium on Workload
  Characterization (IISWC)}, pages 239--249, 2017.

\bibitem{10.1145/3297858.3304055}
Kai Wang, Don Fussell, and Calvin Lin.
\newblock Fast fine-grained global synchronization on gpus.
\newblock In {\em Proceedings of the Twenty-Fourth International Conference on
  Architectural Support for Programming Languages and Operating Systems},
  ASPLOS '19, pages 793--806, New York, NY, USA, 2019. Association for
  Computing Machinery.

\bibitem{10.1145/2903150.2903155}
Yunlong Xu, Lan Gao, Rui Wang, Zhongzhi Luan, Weiguo Wu, and Depei Qian.
\newblock Lock-based synchronization for gpu architectures.
\newblock In {\em Proceedings of the ACM International Conference on Computing
  Frontiers}, CF '16, page 205–213, New York, NY, USA, 2016. Association for
  Computing Machinery.

\bibitem{10.1109/IPDPS.2013.82}
Ayse Yilmazer and David Kaeli.
\newblock Hql: A scalable synchronization mechanism for gpus.
\newblock In {\em Proceedings of the 2013 IEEE 27th International Symposium on
  Parallel and Distributed Processing}, IPDPS '13, pages 475--486, USA, 2013.
  IEEE Computer Society.

\bibitem{9139854}
Lingqi Zhang, Mohamed Wahib, Haoyu Zhang, and Satoshi Matsuoka.
\newblock A study of single and multi-device synchronization methods in nvidia
  gpus.
\newblock In {\em 2020 IEEE International Parallel and Distributed Processing
  Symposium (IPDPS)}, pages 483--493, 2020.

\end{thebibliography}