# GPU Load Balancing

Muhammad Osama

(2022)

###### Abstract

Fine-grained workload and resource balancing is the key to high performance
for regular and irregular computations on the GPUs. In this dissertation, we
conduct an extensive survey of existing load-balancing techniques to build an
abstraction that addresses the difficulty of scheduling computations on the
GPU.

We propose a GPU fine-grained load-balancing abstraction that decouples load
balancing from work processing and aims to support both static and dynamic
schedules with a programmable interface to implement new load-balancing
schedules. Prior to our work, the only way to unleash the GPU’s potential on
irregular problems has been to workload-balance through application-specific,
tightly coupled load-balancing techniques. With our open-source framework for
load-balancing, we hope to improve programmers’ productivity when developing
irregular-parallel algorithms on the GPU, and also improve the overall
performance characteristics for such applications by allowing a quick path to
experimentation with a variety of existing load-balancing techniques.
Consequently, we also hope that by separating the concerns of load-balancing
from work processing within our abstraction, managing and extending existing
code to future architectures becomes easier.

Using our insights from load-balancing irregular workloads, we build _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.

[inconsolata] it [inconsolatait] Electrical and Computer Engineering
Electrical and Computer Engineering 2022 December John D. Owens Venkatesh
Akella Jason Lowe-Power To my parents, Shahid and Fozia, my siblings, Umair
and Sahar,
and my fiancée Melissa.

###### Acknowledgements.

I owe my deepest gratitude to my advisor John D. Owens. He gave me the
opportunity and the freedom to seek exciting research and become who I am
today. His enthusiasm for research and admiration of his students made the
last seven years of my life truly exciting and endearing. He has helped me
find my own style of research, writing and mentoring, and these are the best
gifts anyone can ever hope to receive. Thank you John, you taught me how to
look for and attempt to solve challenging problems, and I am honored to have
been a part of your exceptional research family. Over the years, I also had
the pleasure of learning and discussing research with remarkable individuals.
Through early part of my Ph.D., Yuechao Pan and Yangzihao Wang mentored me and
taught me much of what I know about parallel programming, graphical processing
units (GPUs) and graph analytics. During the later half I had the pleasure of
learning from Serban D. Porumbescu, a terrific post graduate of our research
group. Since joining our group, Serban has been a key resource and has helped
me not just with the academic but also personal and professional matters. He
is also responsible for helping me edit many of my research works, and I will
forever be grateful for the time that he has spent making my graduate school
experience truly extraordinary. Among the folks in the industry, I was lucky
enough to learn CUDA from the CUDA architect himself, Stephen Jones! Stephen
has always been a great friend to the group, but has been especially kind to
me and taught me the value of good software design and engineering. From Duane
Merrill and Cris Cecka I learned everything about low-level optimizations and
writing high-performance GEMM kernels. I also had the chance to work with
Michael Garland, with whom I enjoyed fruitful discussions that led to many
significant design decisions within my research. I truly enjoyed working with
my co-authors and collaborators: Jonathan Wapman, Carl Yang, Minh Truong,
Chenshan Yuan, Leyuan Wang, Weitang Liu, Andrew Davidson, Yuduo Wu, Andy T.
Riffel, Aydin Buluç, Toluwanimi Odemuyiwa, Cameron Shinn, Collin Michael
McCarthy, and Saurav Muralidharan. I also would like to thank my students Minh
Truong, Jonathan Wapman, Daniel Loran, Cameron Shinn and Annie Robison, who I
had the pleasure of mentoring as they found their own place within the
research group. It was a privilege. Among the terrific people already
mentioned, I would like to add the rest of my outstanding colleagues, with
whom I have spent large parts of the last seven years: Muhammad Awad, Kerry
Seitz, Afton Geil, Agnieszka Łupińska, Saman Ashkiani, Vehbi Eşref Bayrktar,
Ahmed H. Mahmoud, Yuxin Chen, Zhongyi Lin, Radoyeh Shojaei, Chuck Rozhon,
Matthew Drescher, Shalini Venkataraman, Nima Johari and Teja Aluru. I am
fortunate to have experienced the influence of so many talented and
hardworking people in my life. I would also like to extend my gratitude to
Prof. Venkatesh Akella, Prof. Soheil Ghiasi, Jason Lowe-Power and Sean
Treichler for dedicating their time and attention to be a part of my
qualification exam committee, and also Prof. Akella and Jason for further
providing me with feedback and guidance on my dissertation. Additionally, I
would like to thank Sean and Aamer Jaleel for dedicating an hour out of their
busy schedule every week for months to give feedback on mine and my
colleagues’ research works. I also wish to thank the past and current UC Davis
ECE staff, who always work hard to remove frictions from the path of students
and professors, and enable us to focus on our research. Among the staff, I
especially thank Sacksith Ekkaphanh, Phil Young, Fredrick P. Singh, Jennifer
Y. Torres, Carole Bustamante, and Michelle T. Walker. Graduate school would be
way too stressful without friends, and I have been lucky enough to have a lot
of wonderful friends that helped me get through some of the most stressful
times with lots (and lots) of distractions. The RPG, anime and manga theory
crafting with Aaron Feaster and Thea Tesoro, the Dungeon & Dragons sessions
and the countless hours111Please don’t tell John. spent playing Destiny and
Overwatch with Nate Bryant, Jesse Navarro, Christian Perry, William “Billy”
Lucas, Krestine Whaley, Jace Locario and Tiffany Simonet, and the League of
Legends games with Megan Vieira. My late friend Pierre “Link” Schreurs, dearly
missed, who’s absurdity and jokes always made us laugh till we had tears in
our eyes. Thank you all for being the best of friends. I extend my love and
gratitude to my fiancée, Melissa Vieira. Despite the distance that separated
us during these years, you have supported me with your love and joy. I am
lucky to have you and truly excited to spend the rest of my life with you. My
research would not have been possible without the financial support from the
National Science Foundation (SI2-SSE), the Defense Advanced Research Projects
Agency (XDATA, HIVE and Symphony), the United States Geological Survey Agency
(USG23), equipment donations from NVIDIA and their funding of an NVIDIA AI Lab
at UC Davis, and a 2022 UC Davis ECE Dissertation Fellowship. And finally, to
my family—I am what I am today because of you. I thank my parents, Shahid and
Fozia, who have worked tirelessly throughout their lives, foregoing personal
pleasure to raise me and my sibilings with all the love, care and support in
the world, and who always taught us to learn and unravel the secrets of this
world. I thank my sister, Sahar, who continues to be the symbol of patience
and courage, which I needed to chase after my dreams. I thank my brother,
Umair, who ever since I was a kid, has been my role model—your hardwork,
passion and persevarence has led me down a path of excellence. I thank my
sister-in-law and brother-in-law, Maria and Rameez, for their support and
love. I thank my niece and nephews, Hadiya, Hashir, Musab, and Mahad, who are
my sweet little bundles of joy and my source of happiness. Whatever I did to
deserve you, it could not have been enough.

Thank you.

##  Chapter 0 Introduction

Graphics Processing Units (GPUs), born out of the needs of graphics for games,
dealing with vertices that make up triangles and subsequent fragments
generated by the rasterizer, have evolved to be more general-purpose and
programmable. GPUs now excel in many other domains beyond graphics, such as
machine learning applications, graph analytics, sparse linear algebra, real-
time interactive media, physical simulations, and bioinformatics. GPUs are
used as a staple for accelerating data-parallel workloads within these
domains, where parallelism is exploited by operating on multiple elements of
data simultaneously.

Data parallelism within these application domains can be classified based on
the granularity of work, where granularity (or grain size) is a measure of the
amount of work performed by each task [44]. _Fine-grained_ data parallelism is
present when the granularity of work is much smaller than the entire problem.
The problem is broken down into smaller components with small amounts of
computations that map well to less sophisticated processing elements.
Computations between fine-grained components have high communication overhead,
and the impact of workload imbalance among processing elements is fairly
large. _Coarse-grained_ data parallelism is present when a program is broken
down into a small number of large, computationally intensive tasks, and
between the tasks, there is low communication overhead [44]. Due to their
large number of cores and their throughput-oriented design, modern GPUs excel
at processing fine-grained data-parallel tasks. NVIDIA’s GA100 GPU, for
example, boasts 128 Streaming Multiprocessors (SMs) with 8192 total cores and
High Bandwidth Memory (HBM) to keep the cores fed with data, both promoting
fine-grained parallelism [71].

In this dissertation we take a closer look at two categories of fine-grained
parallel work, _regular_ and _irregular_ computations, and explore the
fundamental underlying software technology that allows these computations to
run efficiently on the GPU architecture: load balancing. In this context, a
computation is defined to be regular when the neighboring processing elements
have similar or identical workloads. In contrast, a computation is irregular
when neighboring processing elements have varying amounts of work to process.
Figure 1(a) shows an example of a dense matrix used in a General-Matrix
Multiplication (GEMM) computation, a regular workload where the partition of
the work can be statically and trivially determined, such that each processing
element performs identical amounts of reads, writes, and multiply-accumulate
instructions. In contrast, Figure 1(b) visualizes a graph represented as a
sparse matrix used in sparse-linear algebra operations such as Sparse-Matrix
Matrix Multiplication (SpMM), an irregular workload where each processing
element has varying amounts of reads, writes, and nonzero values to process,
and thus an inherent load imbalance within the dataset.

(a) An example dense matrix of size 2048$\times$2048 used as an input to
linear layer of a ConvNeXt architecture [59]. In a dense matrix-matrix
multiply (GEMM), all values of this matrix will be consumed, with neighboring
threads each processing an identical number of elements.

(b) Xyce circuit simulation matrix (graph) of size 321,671$\times$321,671 with
1,316,085 nonzero values irregularly distributed [29]. If this data is stored
in a typical sparse-matrix format (e.g., CSR), only nonzero values will be
consumed, and the varying amount of work per matrix row means that a
straightforward mapping to a GPU results in neighboring threads consuming
different amounts of work.

Figure 1: Visualization of regular and irregular matrices.

To leverage fine-grained parallelism within regular and irregular computations
on a GPU, we attempt to maximize the available memory bandwidth and compute
throughput. For compute-bound, regular problems such as GEMM, the predominant
approach is to find the optimal size of chunks of the complete workload, which
efficiently utilizes the GPU memory hierarchy, and map that onto the GPU cores
in oversubscribed “waves” of work. These waves get scheduled onto the GPUs as
resources become available until all waves/work is completed. In practice, the
last wave of scheduled work often underutilizes the GPU, causing millions of
instructions to parallelize over only a small subset of the total available
cores. In contrast to prior approaches, which built complicated heuristics
selecting different chunk sizes for the waves of work to better utilize the
device, this dissertation addresses the _resource imbalance_ problem for GEMM
by instead balancing the total available work to a fixed, device-filling
number of processing elements. Our approach is a generalization of previous
parallel decompositions and removes the need for complicated heuristics for
kernel selection for GEMM (and other regular problems alike).

For _irregular_ computations, however, the irregularity within the program
loop-structure or the data itself varies tremendously, making it difficult to
find _one_ parallel decomposition that balances all problems onto available
GPU resources. Instead, there is an inherent tradeoff between the quality of
the balance and the amount of extra work needed to compute the workload
balance. This dissertation addresses the _workload balancing_ problem for
irregular workloads with a load-balancing abstraction, which separates
workload mapping from work execution. We argue and demonstrate that this is
the superior approach to building high-performance kernels for irregular
applications as it promotes reuse of existing load-balancing techniques and
facilitates improved portability and programmability.

###  1 Contributions

This dissertation makes several contributions to the areas of load balancing
regular and irregular computations on the GPU.

  * •

We conduct a survey of current load-balancing techniques used for sparse-
irregular applications (Chapter 2).

  * •

We present our GPU load-balancing abstraction that promotes the separation of
concerns between workload mapping and work execution. Our GPU load-balancing
framework reconstructs existing application-dependent techniques that address
irregularity to be more general, portable, and programmable (Chapter 3).

  * •

We provide a solution to the _resource balancing_ problem for dense-regular
computations such as GEMM. Our approach provides improved performance compared
to highly-optimized vendor library, cuBLAS, generalizes existing solutions and
removes the need for complicated heuristics to select the right set of tile
sizes and shapes when scheduling (Chapter 4).

##  Chapter 1 Background

###  1 Graphics Processing Unit (GPU) Architecture

Modern GPU architectures feature two separate hierarchies, one targeted at the
individual processing cores organized into a parallel compute hierarchy
(Figure 1(b)) and the second targeted at the memory scopes such as L1 and L2
cache, programmable shared memory, and global memory, all organized into a
memory hierarchy (Figure 1(a)). Using NVIDIA’s popular Compute Unified Device
Architecture (CUDA) programming model for parallel computing, the compute and
memory hierarchies are exposed for programmers to build general-purpose
computing applications for GPUs. In this chapter, we provide a brief overview
of CUDA’s compute and memory hierarchies and highlight their significance for
load balancing.

(a) CUDA’s Memory System NVIDIA A100 sizes for global memory, L1- and
L2-caches, registers used as an example.

(b) CUDA’s Compute Hierarchy Shows how a kernel is mapped onto the individual
threads of a CUDA grid.

Figure 1: GPU architecture illustrated as CUDA’s memory system and compute
hierarchy.

####  1 CUDA’s Compute Hierarchy

The compute hierarchy within the CUDA programming model leverages multiple
compute “perspectives” that get mapped onto physical streaming multiprocessor
cores (SMs) on the GPU. The SMs feature a _Single Instruction Multiple Thread
(SIMT)_ programming model to achieve multithreaded data-parallel execution
[24]. CUDA further expands the compute hierarchy within the SMs using a
software abstraction of threads, the smallest unit, grouped to create larger
compute perspectives like a warp, a block, and more. Figure 1(b) illustrates
how a CUDA grid, mapped onto the physical SM, is broken down into a number of
compute perspectives. The finer-grained perspectives within the hierarchy
allow us to do more flexible, higher-performance communication when compared
to coarser-grained perspectives. The following list summarizes each of these
perspectives and their purpose:

  * •

Thread The smallest processing unit. Each thread runs an instance of a GPU
program called a _kernel_.

  * •

Warp 32 CUDA threads that run in lockstep. Threads in a warp are divergent-
free, and run in a Single Instruction Multiple Data (SIMD) fashion.

  * •

Block A block or Cooperative Thread Array (CTA) in a GPU is a group of threads
that cooperate together to process a slice of data in an algorithm and map
onto the same Streaming Multiprocessor (SM). Multiple blocks may run on the
same SM concurrently.

  * •

Clusters A thread block cluster is a collective of blocks, co-scheduled on
adjacent SMs.

  * •

Grid The collection of all blocks that run on all the SMs in a single device.

####  2 CUDA’s Memory Hierarchy

The CUDA memory model is similar to modern CPU architectures. It features _L1-
and L2-cache_ , with varying caching policies and traits depending on the GPU
architecture. It also features the following list of programmable memory
scopes:

  * •

Global Memory An off-chip memory shared by all SMs.

  * •

Shared Memory Fast on-chip memory shared among different threads in a single
block, but not between blocks. Shared memory is limited per SM and is often
used to store data that is accessed across an entire block. On some
architectures, shared memory shares the same pool of memory as the GPU’s L1
cache and the programmer may specify how much of this pool to assign to cached
vs. explicit data storage.

  * •

Registers Any scalar variables declared within the scope of a kernel are by
default stored within registers. Register data is only visible locally at the
thread-level scope. It is the fastest memory in a GPU.

####  3 Impact on Load-Balancing

High-performance implementations of fine-grained parallel workloads require
programmers to design their kernels with the previously described compute and
memory hierarchies in mind. Warps’ SIMD execution model is hardware-efficient
and is a great match for many workloads. However, it forces programmers to
avoid any discrepancies in instructions between the neighboring threads within
the same warp. Load-balancing techniques must ensure that threads within the
same warp have approximately the same amounts of work to process as well as
the same set of instructions. At a block level, NVIDIA’s hardware scheduler
also plays an important role in scheduling blocks to the underlying SM as
resources become available. Work scheduled to the GPU is often oversubscribed
to its resources, therefore, the blocks get scheduled out as they finish their
set of work with new blocks assigned to keep the GPU fully utilized.
Programmers often spend time balancing the workload for a given block (or
various compute hierarchies) and then rely on the hardware scheduler or the
data workload distribution itself to balance the resources of an SM to a given
problem. We discuss this and other load-balancing techniques in detail in
Chapter 3.

Memory hierarchy also plays a significant role in optimizing kernels, where
the most important consideration is minimizing the memory accesses required to
access data at a thread-level granularity. This includes coalesced memory
accesses, where multiple memory accesses are grouped into a single
transaction. However, many irregular problems inherently result in uncoalesced
loads and stores where the accesses become serialized; this is due to the
sparse or random access nature of these workloads. For both regular and
irregular workloads, memory hierarchy is utilized to reduce the number of off-
chip (global memory) accesses needed by caching the accesses to L2/L1,
registers, and the shared memory.

##  Chapter 2 A Survey of GPU Load-Balancing Techniques for Irregular
Applications

Graphical Processing Units (GPUs) excel at and are generally designed for
regular fine-grained parallel problems, such as General Matrix Multiplication
(GEMM). In regular problems like GEMM, neighboring threads have similar or
identical workloads and often achieve nearly 100% of peak GPU theoretical
performance. What is much more challenging is an application with ample _fine-
grained_ parallelism but _irregular_ parallelism. In such applications,
neighboring threads running in a lockstep fashion will have different
workloads—perhaps different amounts of work—making an efficient implementation
on a highly parallel machine like a GPU a significant challenge.

Consider Sparse-Matrix Vector Multiplication (SpMV), a critical kernel within
scientific computing for sparse linear algebra and sparse eigenvalue solvers
[31]. A sparse matrix A and a dense vector $x$ defined as inputs, SpMV
computes the output vector $y=\textbf{A}x$ and is an example of irregular
fine-grained parallelism. Unlike in GEMM, the sparse matrix in SpMV can
contain irregularity within the rows of the matrix: the rows of the matrix can
have different numbers of non-zero entries. A simple mapping of one row to
each GPU thread can expose this irregularity, where the neighboring threads
may get different amounts of non-zeros to process, causing threads within the
same warp111Reminder: A CUDA warp is a collection of 32 threads that execute
instructions in lockstep. Threads in a warp are divergent-free, and run in a
Single Instruction Multiple Data (SIMD) fashion. to wait on threads with large
amounts of non-zeros. The imbalance created due to this
irregularity—specifically, when the work is not equally distributed among the
parallel actors, and consequently, some actors are idle while others do more
work—is defined as the load imbalance problem. Current approaches to
addressing the load-imbalance problem is by distributing the work evenly among
threads using sophisticated load-balancing techniques.

In this survey, we look at these fine-grained GPU load-balancing techniques
that are used to date to implement high-performance kernels for irregular
applications such as graph algorithms (Breadth-First Search), sparse-linear
algebra (Sparse-Matrix Multiplication), and many others. But this problem is
not specific to GPUs alone. In the past several decades, researchers have
proposed many solutions to load imbalance for SIMD and SIMT programming models
alike [12, 10, 58]. However, given the prevalence of GPUs, their unique
compute and memory model (see Chapter 1), this work addresses the need for a
broad understanding of the literature space. Our contributions are the
following:

  1. 1.

A taxonomy of GPU load-balancing techniques to categorize the load-balancing
algorithms (Section 2).

  2. 2.

A deeper dive on key techniques that appear in the literature (Section 3).

  3. 3.

An understanding of implementation details and building-block algorithms
required to implement these techniques (Section 4).

The irregularly-parallel workloads we encountered in our survey predominantly
focus on sparse linear algebra (with problems such as sparse-matrix dense-
vector multiplication, sparse-matrix sparse-matrix multiplication), graph
analytics (with algorithms such as breadth-first search, single-source
shortest path, triangle counting, and general graph analytics programming
abstractions), and computer graphics (ray traversal).

###  1 Related Works

To the best of our knowledge, only one other paper from 2017 conducts a survey
of GPU load-balancing techniques [16]. Busato and Bombeiri explain some of the
key techniques used at that time to implement load balancing on older GPU
architectures, and provide and performance, power and energy analysis on a
small subset of sparse datasets. Contemporary to Busato and Bobeiri’s
approach, our survey not only updates the load-balancing survey literature
with new and old techniques, but also provides a deeper dive into their
implementations and the underlying low-level primitives. We also build a fine-
grained GPU load balancing taxonomy to characterize any current or future
load-balancing algorithms.

####  1 Compressed Sparse Data Structures

Irregularly-parallel applications often operate on _sparse_ datasets, or
datasets in which a significant portion of the data consists of “zero” values
that make up the bulk of the data in the data structure, but that would be
wasteful to store. Sparse data formats allow the programmer to both store the
data more efficiently in memory by omitting explicit “zero” values as well as
access the data more efficiently once it is in a sparse format without
performing useless work. Since many sparse workloads are bandwidth-bound
rather than compute-bound, the cost of moving data into the GPU’s memory
hierarchy is a key factor in determining the performance of a sparse workload.
In fact, load balancing is often a function of making sure that all processors
are _reading or writing_ equal amounts of data rather than performing
arithmetic on equal amounts of data. Performant sparse data formats commonly
group logically-adjacent work items in memory so that they can be read or
written in a single memory transaction, so that elements within a group share
data (such as row or column metadata), or so that the size of all groups in
the data structure are equal. In this section we describe a few common
compressed data structures that are used to represent these datasets.

First, consider a matrix where only a handful of the elements are non-zero. We
can represent this matrix as a _COO_ (coordinate) format, which is a simple
list of nonzero elements. Each element is stored as a tuple of row, column,
and value. These tuples are often sorted first by row and then by column
within each row, but this step is optional. The COO format works well when the
programmer simply wants to split nonzeros into multiple groups.

The primary downsides of the COO format are that it uses more space than
necessary due to explicitly-stored row and column indices for each nonzero,
and it is computationally expensive to perform queries on the matrix,
especially if it is unsorted by row and column. For example, determining the
number of nonzeros in a row requires an iteration over the list, as do random
accesses into the matrix. To solve this problem, we can use the Compressed
Sparse Row (CSR) format, which is a more space-efficient format that stores
nonzero elements in a compressed format. CSR uses a simple row-major list of
nonzero column indices and values, as well as a prefix-sum that stores a
running total of the sum of the nonzeros in all previous rows. Compared to
COO, this eliminates the need for row index storage, allows kernels to quickly
determine the number of nonzeros in each row with a single subtraction
operation, and enables faster random or iterative accesses into a given row.

There are many other alternative data structures that typically require more
preprocessing, as described in Table II by Fillipone et al. [31]. In general,
preprocessed formats are a form of static load balancing, where the programmer
analyzes how the data will be accessed and uses this information to optimize
the data structure, with an assumption that the cost of preprocessing the data
format will pay off on net, either by amortizing the preprocessing cost over
many computations or by significantly speeding up the computation. In this
work we primarily focus on CSR and COO formats in part because they are the
most common formats used in sparse linear algebra on GPUs, and in part because
they represent two different ways of organizing data with opposing trade-offs.
COO is a list of coordinate/value tuples, where the programmer can easily
subdivide the data structure based on an equal split of nonzeros, but where
the cost to access information about the row the nonzero is a part of is more
expensive. In contrast, CSR is a compressed format allows the programmer to
easily subdivide the data structure into an equal number of rows, but requires
the programmer to pay an additional cost to load balance _across_ rows or
access a random nonzero element.

###  2 Taxonomy of GPU Load Balancing

This section focuses on building a taxonomy of GPU load balancing. Our goal
with this taxonomy is to (1) characterize existing load-balancing techniques
found in GPU literature (see Section 3 for a deeper-dive on implementations,
Section 4 for a study on common low-level algorithmic primitives used for
load-balancing, and 5 for a complete summary), and (2) provide an abstraction
to classify future works searchability.

Configuration | Accuracy | Granularity | Communication | Topology
---|---|---|---|---
Static | Exact | Hierarchical | Cooperative | Centralized
Dynamic | Approximate | Flat | Non-Cooperative | Distributed

Table 1: Taxonomy of Load-Balancing techniques for GPUs, a layout elaborating
how load-balancing algorithms can be characterized.

####  1 Configuration: Static vs. Dynamic

Load-balancing algorithms fall into two distinct approaches; static and
dynamic. Static algorithms use only the information known before the algorithm
launches. This may include the number of work items or tiles in the dataset,
the number of threads or SMs in a given GPU, or any other characteristics of a
given dataset, algorithm, or GPU architecture that can be determined without
actually running the algorithm. The work distribution does not change as the
algorithm proceeds, and static load-balancing strategies do not account for
any runtime changes in the state of the processors or the problem. Therefore,
static load-balancing algorithms have no notion of the load on each processor
during the execution period. Conversely, dynamic load-balancing algorithms
account for the changing incoming or outgoing load on each execution unit by
measuring and (re-)distributing that load at runtime. For example, some
kernels may dynamically create new work during runtime, and a dynamic load-
balancing strategy allows the kernel to immediately process this new work with
idle processors rather than waiting for a later kernel launch to do so. Other
load balancing strategies may rely on processors greedily requesting or
donating work during runtime. Because dynamic approaches can redistribute work
at runtime, they can potentially have better overall performance than static
approaches, which do not. But this performance comes at the cost of additional
runtime work. Thus dynamic approaches win only if the performance gains from
better load balance exceed the additional cost of the runtime work to compute
it. This is not always a win; Yang et al.’s sparse-matrix multiplication
analysis [88] and concludes that static approaches for this problem deliver
better overall performance.

For load-balancing techniques that use both static and dynamic approaches, we
use the classification of Hybrid [22]. Hybrid algorithms may use a static
technique for one segment of the problem (or one part of the compute
hierarchy) and a dynamic technique for another.

####  2 Accuracy: Exact vs. Approximate

The ideal load-balancing strategy aims to minimize runtime. In practice, such
a strategy aims to balance work evenly across all compute elements (threads,
warps, or blocks) so that each compute element completes its work at the same
time and no units are ever idle waiting for others to finish. Optimal
algorithms use the properties of the problem and the number of available
execution units to determine the exact amount of work to assign to each unit
to ensure all execution units are fully utilized at all times. These
algorithms may achieve optimal load balancing for a given cost function. In
GPUs, the cost function often results in modeling the variance of work between
neighbouring threads, and the goal of the load-balancing is to attempt to
assign equal amount of work items per thread (an “even-share” schedule [35,
26, 64, 63, 8, 78], another cost function may model execution time, and a
load-balancing schedule optimizes this cost function to ensure that each
thread runs for approximately the same time). However there is often an
overhead associated with load balancing techniques that attempt to achieve
exact workload balance across all processing units. Approximate load-balancing
algorithms, on the other hand, focus on metrics _within_ the problem or the
dataset, or the attributes of the system (such as the number of threads) to
attempt to distribute the work equally across processors. Approximate
algorithms may achieve inconsistent load balancing depending on the metrics
used, but can be tuned to reduce the overhead required to achieve a balanced
workload. The resulting trade-off between the cost of load balancing and the
quality of the workload balanced achieved is an important consideration when
opting for an exact or approximate load-balancing algorithm.

####  3 Granularity: Hierarchical vs. Flat

Hierarchical schedules leverage the compute hierarchy available within NVIDIA
GPUs by breaking the workload into coarse-grained chunks that map to the
highest level of the compute hierarchy (for example the blocks) and then
further break down the work into smaller granularities that proportionally map
to smaller computation constructs such as warps and/or threads [64]. Section 1
explores this concept in more detail. Flat scheduling algorithms are designed
to target the smallest computational unit in the hierarchy (threads) and are
not mapped to the compute hierarchy available within GPU architectures [65,
28, 84]. Hierarchical schedules often improve load balance by ensuring that a
given work chunk maps to a group of computational units instead of a single
(smallest) unit with proportional computation capabilities, and process the
mapped work in parallel using the compute hierarchy within the group (CUDA’s
compute hierarchy explained in Chapter 1). For example, a row in a CSR matrix
with few nonzeros may be most efficiently processed with a single thread,
since processing this row with an entire block would leave many threads
sitting idle when they could otherwise be processing other rows. As another
example, to avoid significant amounts of atomic contention, a queue-based
algorithm may assign a single thread to fetch a large chunk work on behalf of
all threads in a block, but then further subdivide this work amongst the
threads in the block. Beyond their increased complexity, the primary downside
of hierarchical schedules is that the performance gains from leveraging
coarser-granularity workers may not be enough to outweigh imperfect load
balance within the coarser-granularity work chunks.

####  4 Topology: Centralized vs. Distributed

Centralized schedules focus on a distribution unit responsible for assigning
work items to multiple processors. As an example, the GPU’s (hardware) block
scheduler distributes work to all SMs, assigning a new block to an SM when the
SM is ready for new work. In contrast, distributed schedules rely on
individual processors making independent, local decisions about when they need
more work.

Consider a shared queue of work, accessible by all blocks. A centralized
strategy for managing this queue has a single entity (either hardware or
software) that is responsible for distributing work to processors. A
distributed strategy instead delegates the responsibility for fetching new
work to each processor: processors must actively fetch new work for themselves
as needed, while competing with other processors to claim work from the queue.
Centralized and Distributed schedules only apply to _dynamic_ schedules. In
static schedules, each processor already knows exactly which work items it
must process, so there is no concept of “shared” work that could go to one of
many processors depending on runtime decisions.

####  5 Communication: Cooperative vs. Non-Cooperative

Cooperative schedules allow processing entities to communicate with each
other, for example in work-stealing load-balancing algorithms where threads
are able to steal work from other threads within a warp or another block [11].
Non-cooperative algorithms require processors to operate autonomously within
their own pool of work and with no ability to access pools of work that may be
simultaneously accessed by _other_ processors. In a throughput-oriented
architecture, such as a GPU, cooperation often requires expensive
synchronization, such as locks or atomics on queues. However, cooperation may
result in a better workload balance [22]. Similar to Centralized and
Distributed schedules, the concept of “cooperation” only applies to _dynamic_
schedules. Cooperation requires that processors are able to make decisions at
runtime about whether they need more work, or whether they have additional
work to give to another processor.

###  3 Implementations of Load Balancing on GPUs

In this section we conduct a study of how load-balancing techniques
(summarized in Table 2) actually get implemented. We characterize these
implementations using the taxonomy built in the previous section, and group
each implementation under common patterns of scheduling work.

###### Terminology

For consistency, we will use the following terms to describe the partitions of
_work_ within the different load-balancing techniques.

  * •

Work Item: A single unit of work that is to be scheduled onto the GPU; for
example, a non-zero element of a sparse matrix.

  * •

Work Tile: A collection of work items, for example, a single row of a sparse
matrix.

  * •

Tile Set: A collection of work tiles that together comprise the entire
problem, for example, an entire sparse matrix.

The terminologies summarized above are later built into a complete abstraction
for load-balancing on the GPU, and are explained in detail in Section 3. As an
example to illustrate the terminology, consider Sparse-Matrix Dense-Vector
Multiplication (SpMV). The irregular work in SpMV is within the sparse matrix,
where the work items (the smallest units of work) are the non-zero elements, a
work tile is a row of the sparse matrix, and a tile set is the entire sparse
matrix. The work items are scheduled onto the GPU using a load-balancing
technique.

The rest of this section describes different strategies for load-balancing
work. We can choose to assign an equal number of work tiles to each thread
(Section 1) or to a group of threads (Section 2), or instead assign an equal
number of work items to each thread (Section 3). We use several low-level
algorithmic building blocks to construct these load-balancing techniques, such
as _parallel prefix-sum (scan)_ , _binary search_ , _sorting_ , and more. We
defer a detailed explanation of these building blocks to Section 4.

####  1 Thread-Mapped

Definition
Assign a fixed, constant number of work tiles to each thread. Resultant work
items from each work tile are processed sequentially within the thread.

Characterization
Static, Approximate, Flat.

Evaluation
Works well for small, balanced work tiles; not recommended for tile sets where
the variance between work tile size is large (e.g., scale-free graphs).

Perhaps the most natural way to parallelize across work tiles is to assign
each work tile to a GPU thread, and sequentially process the work items within
the work tile. For problems where the workload is inherently balanced, the
additional overhead of a sophisticated load-balancing scheduling of work can
often be detrimental to performance. In such cases, this simple thread-mapped
assignment can prove to be a viable substitute. However, the resultant
workload balance with this method is often at a coarser granularity. As
discussed in detail in Merrill et al.’s “Merge-Based Parallel Sparse Matrix-
Vector Multiplication” [64], the performance for a thread-mapped assignment is
typically not as good as other schedules if the problem is highly irregular or
if there are large amounts of fine-grained parallelism left to explore.
Because each thread contains a fixed-size set of tiles and in a highly
irregular dataset, each tile may contain a different amount of work, such
datasets cause threads within the same warp to be waiting on threads with
significantly more work, causing very low warp or block utilization and poor
performance. Furthermore, since a thread is assigned an entire work tile, the
work items within a tile—for example, the nonzeros within a sparse-matrix’s
row—are sequentially processed, and the fine-grained parallelism within the
work tile is effectively serialized, leaving the device underutilized.
Algorithm 1 shows a simple example of thread-mapped assignment as a scheduling
strategy where each thread is assigned a fixed, statically determined input
and output elements to process [65, 84, 14].

Algorithm 1 Simple pseudocode to illustrate thread-mapped scheduling where
each thread is assigned to work on $N$ input and output elements of the saxpy
operation.

inconsolata

$N\leftarrow 4$ $\rhd$ Four items per thread.

procedure Thread-Mapped($\alpha,x,y$)

$t_{id}\leftarrow N\times\textbf{ThreadIndex}()$

$\rhd$ Loop can be perfectly unrolled.

for $k\leftarrow 0$, $k{+}{+}$, while $k<N$ do

$y[t_{id}+k]\leftarrow\alpha\cdot x[t_{id}+k]+y[t_{id}+k]$

end for

end procedureinconsolatait

####  2 Group-Mapped

Definition
Assign an equal amount of work tiles to a group of threads (warp or block).
Threads within each group process individual work items in parallel.

Characterization
Static, Approximate, Hierarchical.

Evaluation
Exposes fine-grained parallelism within a work tile, and is highly tunable to
different problems (allowing the group size to be configured to fit a
problem). Works well for fairly regular workloads where no one group takes too
much GPU resources.

In contrast to the coarse-grained thread-mapped approach, for irregular
problems with ample amounts of fine-grained work items, consider assigning
groups of threads to process a single work tile in parallel. For this
approach, we can leverage the existing compute hierarchy available within
CUDA’s programming model: threads, warps, and blocks. A larger compute unit
such as a block or a warp could potentially prove to be a better mapping for a
work tile that has a large number of work items to process. This is because
instead of serializing the work within a work tile, we can assign all threads
within a block or a warp to process that work in parallel, promoting better
utilization and load balancing within the thread group.

This strategy of assigning an entire tile of work to a group of threads
(either a warp, a block, or a “Cooperative Group” [69]) of a GPU is called
group-mapped load balancing. Group-mapped is typically implemented by first
assigning an equal number of work tiles to each group, $\frac{\text{Total Work
Tiles}}{\text{Total Groups}}$. Then to effectively map the assigned tiles onto
the threads within a group, each thread individually needs two sets of
information: (1) The total number of work items it is processing (a range from
work item $i$ to work item $j$) and (2) what work tile each work item belongs
to. A parallel prefix sum gives us both pieces of information. Discussed in
detail in Section 4, a prefix sum is a sequence of numbers that is the running
total of an input sequence. In group-mapped’s case, a prefix sum of work items
per assigned work tiles is constructed, where the last element of the prefix
sum sequence is the total number of work items the group needs to process. To
solve (1), we simply divide the total number of work items assigned to the
group by the number of threads within the group to get the “range” of work
items any given thread is going to process. And for (2), when processing the
work items within a loop, each thread performs a binary search within the
prefix sum sequence to determine the work tile index the current work item
corresponds to (see Algorithm 2, lines 8–10).

Using the running SpMV example, if a row of a sparse matrix is assigned to an
entire warp, all 32 threads within a warp will collectively process each
nonzero element within the row. Similarly, if a row is assigned to an entire
block, all threads within the block will process the individual nonzero
elements of the row [28].

Group-mapped load balancing exposes the fine-grained parallelism available in
each work tile to a group, allowing the work items to be processed in
parallel, with the cost being the overhead of a parallel prefix sum per group
(warp or block) and the binary search to find the work tile within the
parallel prefix sum array. A group-mapped assignment also heavily relies on
the underlying hardware scheduler’s ability to schedule new thread groups as a
group finishes processing a work tile. This method is also commonly known as
warp-mapped or block-mapped scheduling, and is often used in conjunction with
dynamic scheduling techniques such as a dynamic task queue [22] to create a
hybrid load-balancing strategy discussed later (see Section 5). Generally,
irregular problems with structured regular blocks within them map well to this
load-balancing schedule.

Algorithm 2 An example group-mapped scheduled kernel for sparse-matrix dense-
vector multiplication (spmv), where each group is assigned an entire row and
threads within the group process individual nonzero elements.

inconsolata

1:input: $A$, CSR Matrix. $x$, Dense Vector.

2:output: $y$, Dense Vector.

3:procedure Group-Mapped($A,x,y$)

4: $items\\_per\\_tile\leftarrow$ Zeros($GROUP\\_SIZE$)

5: $t_{id}\leftarrow$ ThreadIndex()

6: $gt_{id}\leftarrow$ ThreadIndexInGroup()

7: $\rhd$ Each group populates an array with the number

8: $\rhd$ of work items per tile.

9: $items\\_per\\_tile[gt_{id}]\leftarrow
A.offsets[t_{id}+1]-A.offsets[t_{id}]$

10: $prefix\\_sum\\_array\leftarrow$
ParallelPrefixSum($items\\_per\\_tile,t_{id}$)

11: $\rhd$ Last element of prefix-sum array corresponds

12: $\rhd$ to total work items per group.

13: $total\\_items\leftarrow prefix\\_sum\\_array[GROUP\\_SIZE-1]$

14: $\rhd$ Loop over total work, each thread processing individual work items.

15: for $k\leftarrow t_{id}$, $k+GROUP\\_SIZE$, while $k<total\\_items$ do

16: $\rhd$ Perform a binary-search to find the tile index.

17: $row\leftarrow$ BinarySearch($prefix\\_sum\\_array,k$)

18: $y[row]\leftarrow A.values[k]\times x[A.indices[k]]$

19: end for

20:end procedureinconsolatait

####  3 Work-Oriented

Definition
All threads are assigned $\frac{\textit{Total\ Work}}{\textit{Number\ of\
Processors}}$ work items. Each thread then sequentially processes assigned
work items in a loop.

Characterization
Static, Exact, Flat (+ GPU hierarchy).

Evaluation
Works well when the upfront cost of performing a prefix sum or binary search
is minor compared to the cost of processing the work items; very effective at
balancing highly irregular work (such as in scale-free graphs).

So far we have considered two load-balancing techniques that map _work tiles_
, the input, onto compute entities (thread- or group-mapped), however, neither
are guaranteed to achieve perfect workload balancing at the device level.
Instead, we consider the _work items_ as the granularity targetted for load-
balancing, where additional computation is done to assign equal amounts of
work items to each thread in a GPU, achieving exact workload balancing. This
approach is known as work-oriented scheduling and can be implemented via two
predominant methods, which largely differ in the definition of a “work unit”.

Since the goal is to assign a constant number of work items per thread, the
definition of a work item is an important design decision. The first method,
known as _non-zero splitting_ (used for sparse-linear algebra kernels such as
SpMV), considers the total number of nonzero elements within a sparse-matrix
or total multiply-accumulate instructions ($y_{i}\mathrel{+}=A_{ij}\times
x_{j}$) as the total work [8, 26, 88]. The second method, _merge-path_ ,
considers a work item as either a nonzero element or an output, effectively
associating an equal cost to a nonzero and to outputting a value to the GPU’s
global memory [64, 35]. Fundamentally, the work-oriented assignment is a
mapping of the $\frac{Total\ Work}{Number\ of\ Threads}$ to each thread, where
each thread is assigned a contiguous range of work items to process in a
sequential loop. A binary search per thread (a key low-level algorithm
discussed in Section 2) then finds the work tile for each range of work items
thread must process (see Algorithm 3).

Merge-path was originally proposed for SpMV kernels. In a merge-path SpMV, the
sparse-matrix is stored in a compressed sparse row format (CSR), where the row
offsets of the CSR matrix are already known (the row offset array is
effectively a prefix-sum array). The schedule aims to perform a 2-D split of
the grid created using the row-offsets and the nonzero indices, and the search
along the diagonal of the grid allows each thread to find its starting and
ending rows and nonzeros. The merge-path schedule ensures that there is
constant number of fix-up steps (i.e., the number of threads) required to
accumulate the partial rows that were split between two threads. The result of
a merge-path schedule is that work is mapped to the smallest possible compute
unit in the hierarchy, but merge-path takes advantage of the compute hierarchy
by first splitting the work across blocks, and then to threads within a block
to reduce the search space each thread has to search through (each thread now
only searches through its block’s share of the work). If the data is already
stored in a CSR representation (as proposed by the original paper by Merrill
and Garland [64]), the cost of this method is the 2-dimensional binary search
along the row offsets and the nonzero indices of the CSR matrix, and the
required fix-up step for the partial tiles of work. Beyond load-balancing a
CSR sparse matrix, this method can also be generalized to other sparse
representations and domains such as graph analytics, with an added cost of a
parallel prefix sum to create the required offsets array [84, 14]. Since there
is an additional overhead attached to computing this prefix-sum, this method
tends to be less performant than other low-overhead methods on an well-
structured sparse data [88].

Algorithm 3 An example _work-oriented_ scheduled kernel for sparse-matrix
dense-vector multiplication (spmv), where each thread is assigned an even
share of work [64].

inconsolata

1:input: $A$, CSR Matrix. $x$, Dense Vector.

2:output: $y$, Dense Vector.

3:procedure Work-Oriented($A,x,y$)

4: $\rhd$ Calculate work division and bounds.

5: $t_{id}\leftarrow$ ThreadIndex()

6: $total\\_work\leftarrow A.rows+A.nonzeros$

7: $items\\_per\\_thread\leftarrow\frac{total\\_work}{num\\_threads}$

8: $diag\leftarrow$ min($items\\_per\\_thread\times t_{id},total\\_work$)

9: $diag\\_end\leftarrow$ min($diag+items\\_per\\_thread,total\\_work$)

10: $(row_{start},nz_{start})\leftarrow$ 2DSearch($diag,A$)

11: $(row_{end},nz_{end})\leftarrow$ 2DSearch($diag\\_end,A$)

12: $\rhd$ Perform work on full tiles.

13: $running\\_total\leftarrow 0$

14: for $m\leftarrow row_{start}$, $m{+}{+}$, while $m<row_{end}$ do

15: for $k\leftarrow nz_{start}$, $k{+}{+}$, while $k<A.offsets[m]$ do

16: $running\\_total\mathrel{+}=x[A.indices[k]]\times A.values[k]$

17: end for

18: $y[row]\leftarrow running\\_total$

19: $running\\_total\leftarrow 0$

20: end for

21: $\rhd$ Perform work on partial tiles.

22: for $k\leftarrow k$, $k{+}{+}$, while $k<nz_{end}$ do

23: $running\\_total\mathrel{+}=x[A.indices[k]]\times A.values[k]$

24: end for

25: $row\\_carry\\_out[t_{id}]\leftarrow row_{end}$

26: $value\\_carry\\_out[t_{id}]\leftarrow running\\_total$

27:end procedure

28:$\rhd$ Fix-up step to accumulate partial tiles.

29:procedure Fix-up($A,y,row\\_carry\\_out,value\\_carry\\_out$)

30: for $t_{id}\leftarrow 0$, $t_{id}{+}{+}$, while $t_{id}<num\\_threads$ do

31: $y[row\\_carry\\_out[t_{id}]]\mathrel{+}=value\\_carry\\_out[t_{id}]$

32: end for

33:end procedureinconsolatait

####  4 Binning and Reordering

Definition
Separate work tiles into a fixed number of bins based on heuristics, such that
each bin holds work tiles with approximately the same amount of work items.

Characterization
Dynamic, Approximate, Hierarchical, Non-Cooperative, Centralized or
Distributed.

Evaluation
Effective at balancing workloads with contiguous work tiles with approximately
the same number of work items, and with imbalance among these contiguous
ranges of work tiles. This method maps well to CUDA’s compute hierarchy, where
bins can be assigned to a level in the hierarchy depending on the granularity
of work within them.

The static load-balancing techniques discussed above aim for an even work
distribution for a given problem; these techniques are inefficient when either
the overhead to load-balance is greater than processing the entire
computation, or when there is a huge mismatch between the worker size and the
work unit. Instead, when a given problem has a known bound on the amount of
work, we can implement a load-balancing schedule that is reasonably efficient
for a given bounded range, and particularly so if we choose the correct
amounts of compute resources that match the given input range (e.g., assign a
work tile with exactly 32 work items to be processed by a warp composed of 32
threads.) This is where binning-based load-balancing algorithms, which attempt
to dynamically balance a given workload by categorizing each work tile into a
fixed number of bins with the goal that each bin will hold work tiles that
have approximately the same amount of work items, perform reasonably well.
Binning-based techniques can largely be subdivided into three different
phases: (1) choosing bin sizes (compile time); (2) dynamically distributing
work tiles into bins based on a programmer-specified criteria (runtime); then
(3) computing the contents of each bin (runtime).

One approach to binning that specifically targets the GPU compute hierarchy
chooses three bins to partition the work tiles: (1) a bin with work tiles that
have a number of work items equal to or greater than the size of the GPU block
(number of threads per block); (2) a bin with work tiles that have a number of
work items less than the size of the GPU block but larger than a warp (32
threads); and (3) a bin with work tiles that have a number of work items less
than the size of the warp. For brevity, we will use the terms block-sized,
warp-sized, and thread-sized bins. After the initial assignment of work tiles
to their respective bins, the processing stage that perform the desired
computation on the contents of each bins can be implemented in three different
ways:

  1. 1.

Launch three kernels, where each kernel is specialized to process one of the
three categories of bins. For the block-sized bin, one kernel uses all threads
with the block to cooperatively process the tiles. Similarly, in the warp-
sized bin, the kernel uses all threads within a warp to cooperatively process
the bin, And finally, the last kernel processes the thread-sized bin, where
each work tile has only a few work items, my assigning one thread to each work
tile [65, 84].

  2. 2.

Launch one kernel, where each thread is initially assigned a work tile,
determines the amount of work, and assumes control of the block or the warp if
the size of work is enough to saturate the compute unit. All threads within
the compute unit then cooperatively work together to process each work item
within the “winning” thread’s tile. This process continues until the work
tiles are small enough for each thread to individually process them (Algorithm
4) [28, 84].

  3. 3.

Brahmakshatriya et al. proposed a variation of the one-kernel approach, where
instead of threads competing to assume control of a block, each block first
processes a multiple of its size worth of work items of a tile (if any), and
then a warp processes multiple of its size worth of work items, and finally
the threads complete the work tile by processing the remaining work items
[14].

The benefit of the three-kernel specialization is that each kernel is designed
to process one bin-type with no added communication or synchronization for
threads within the kernel to assume control of a given compute unit or
hierarchy. However, the total work now is now split among three kernel
launches, making it challenging to fully utilize the device per kernel for a
given work distribution. Furthermore, for problems that require a large number
of iterations to successfully converge, the kernel launch overhead, which is
commonly an insignificant cost and largely ignored, is now 3$\times$ larger.

Algorithm 4 An example _binning_ scheduled kernel for sparse-matrix dense-
vector multiplication (spmv), where work tiles are placed into bins based on
the number of work items within them, and processed with the required compute
resource [65, 28].

inconsolata

1:input: $A$, CSR Matrix.

2:procedure Binning($A$)

3: $row\leftarrow$ ThreadIndex()

4: $num\\_nonzeros\leftarrow A.offsets[row+1]-A.offsets[row]$

5: if $num\\_nonzeros\geq block\\_size$ then

6: $cta\\_bin\\_ids[cta\\_bin\\_size]\leftarrow row$

7: $cta\\_bin\\_size{+}{+}$

8: else if $num\\_nonzeros\geq warp\\_size$ then

9: $warp\\_bin\\_ids[warp\\_bin\\_size]\leftarrow row$

10: $warp\\_bin\\_size{+}{+}$

11: else

12: $thread\\_bin\\_ids[thread\\_bin\\_size]\leftarrow row$

13: $thread\\_bin\\_size{+}{+}$

14: end if

15:end procedure

16:input: $A$, CSR Matrix. $x$, Dense Vector.

17:output: $y$, Dense Vector.

18:procedure CTA_BIN($A,x,y$)

19: $row\leftarrow cta\\_bin\\_ids[idx]$

20: $\rhd$ Use an entire CTA to process the row.

21:end procedure

22:procedure WARP_BIN($A,x,y$)

23: $row\leftarrow warp\\_bin\\_ids[idx]$

24: $\rhd$ Use an entire warp to process the row.

25:end procedure

26:procedure THREAD_BIN($A,x,y$)

27: $row\leftarrow thread\\_bin\\_ids[idx]$

28: $\rhd$ Sequentially process a row in the thread.

29: for $k\leftarrow A.offsets[row]$, $k{+}{+}$, while $k<A.offsets[row+1]$ do

30: $y[A.indices[k]]\mathrel{+}=x[A.indices[k]]\times A.values[k]$

31: end for

32:end procedureinconsolatait

Another proposed approach, Logarithmic Radix Binning (LRB), first introduced
by Green et al. within a triangle counting implementation [36], allows the
load balancer to schedule work items with similar amounts of work within the
same spatial and temporal region [32]. LRB assigns tasks to bins based on the
_logarithm_ of the amount of work required for a given task. Thus, the range
of possible amounts of work across all items in a bin varies by no more than a
factor of two. When a work tile is encountered with work items greater than
equal to $2^{b}$ and less than $2^{b+1}$, a counter at location $b$ is
incremented in an array of size $B\in{\\{32,64\\}}$. Using these bins instead
of the ones associated with CUDA’s compute hierarchy (threads, warps, blocks)
in the previous approach, the total workload is binned at a finer granularity
such that work tiles containing a similar number of work items are grouped
together for processing. To assign work tiles to threads, a modulo operation
is performed on the reordered task index and a constant number of work items,
$P$. Each thread then processes $(i,P+i,2*P+i,\ldots)$ work items [32, 36].

Another widely used technique for reordering or grouping like-sized work
together is a simple (but often costly) sort to arrange tiles from most to
fewest work items [13]. The sort operation reduces the work variance between
adjacent tiles and therefore improves load balance. In iterative problems such
as sparse-matrix dense-matrix multiplication in a deep learning workload, the
overhead of the sort is amortized over the number of runs of an algorithm
[33].

####  5 Task-Oriented Scheduling

Definition
Independent workers fetch work from a work queue or queues, process the work,
and optionally add new work to a queue.

Characterization
Dynamic Approximate, Flat, Cooperative, Centralized or Distributed.

Evaluation
Works well when it is too expensive to compute a static assignment of work
items to processors at the start of the kernel, if the kernel dynamically
generates new work items, or if the cost to process each work item is unknown.
Performs best if one thread can fetch chunks of work on behalf of an entire
block, to reduce synchronization contention. Powerful when paired with
hierarchical load balancing within a block, where a block fetches a chunk of
work on behalf of all threads, and then further subdivides this work among its
threads.

The previously-described load balancing methods divide up work at compile time
or runtime and associate each chunk of work with a virtual processor, which
the hardware scheduler then dynamically maps to a physical processor based on
resource utilization. In contrast, task parallelism shifts the job of
assigning work to processors to the programmer. Task parallelism methods are
more flexible at runtime, and benefit workloads where it would be expensive to
compute the optimal work distribution and assign it to individual processors
in advance, or where the total amount of work to be generated in a kernel is
unknown. Programmers can customize the task generation and consumption
behavior to their specific use case, at the cost of some software overhead
compared with the GPU’s hardware block scheduler. A primary benefit of a task-
parallel schedule is that it allows processors that finish their currently
assigned task to immediately start work on another task without waiting for
synchronization at the end of a kernel, or to generate new irregular tasks
that other processors may consume without waiting for a new kernel iteration.

A classic example of task parallelism is a queue-based implementation of
Breadth-First Search. After a block removes a vertex from the FIFO queue, it
sets the vertex’s depth then adds neighbors of that vertex to the back of the
queue. In a parallel GPU system, these new work items may be popped and
processed by other processors in the system. This process continues until the
queue is empty and the processors no longer have new work to add to the queue.

The primary downsides of queue-based schedules are that they suffer from
overheads necessary to ensure that each processing element adds or retrieves
work at the correct location, and that the queue must be large enough to
accomodate the worst-case number of work items in the queue. This is because
memory allocations are expensive and happen outside the kernel boundary.
Additionally, there is a chance that some processors may have long-running
work that leaves other processors sitting idle when the queue is empty.
Finally, queue-based schedules replace the hardware work distribution done by
the GPU’s block scheduler with user-controlled software work distribution,
which adds additional overheads.

##### Centralized Queue

Evaluation
Useful when the algorithm dynamically creates work at runtime, or when it’s
undesirable to precompute an assignment of work items to processors.

Algorithm 5 An example queue-based kernel for breadth-first search, where each
worker is assigned an entire vertex.

inconsolata

1:input: $G$, CSR Graph. $Q$, Work Queue. $s$, Source Vertex.

2:output: $D$, Depth Array. $P$, Predecessor Array.

3:procedure Task-Oriented($G,Q,s$)

4: $t_{id}\leftarrow$ ThreadIndex()

5: $\rhd$ Initialization phase.

6: $D[t_{id}]\leftarrow\infty$

7: $P[t_{id}]\leftarrow None$

8: if $t_{id}=0$ then

9: $Q[0]\leftarrow s$

10: $Q.push\\_idx\leftarrow 1$

11: $Q.pop\\_idx\leftarrow 0$

12: end if

13: $\rhd$ Loop until the queue is empty.

14: while not $Q.empty()$ do

15: $my\\_pop\\_idx\leftarrow atomicAdd(Q.pop\\_idx,1)$

16: $my\\_pop\\_idx\leftarrow$mod$(my\\_pop\\_idx,Q.size())$

17: $v\leftarrow Q[my\\_pop\\_idx]$

18: for $i\leftarrow 0$, $i{+}{+}$, while $i<G.offsets[v+1]-G.offsets[v]$ do

19: $n\leftarrow G.indices[G.offsets[v]+i]$

20: $\rhd$ BFS condition: update the neighbor’s depth.

21: if $D[v]+1<D[n]$ then

22: $D[n]\leftarrow D[v]+1$

23: $P[n]\leftarrow v$

24: $my\\_push\\_idx\leftarrow atomicAdd(Q.push\\_idx,1)$

25: $my\\_push\\_idx\leftarrow$mod$(my\\_push\\_idx,Q.size())$

26: $Q[my\\_push\\_idx]\leftarrow n$

27: end if

28: end for

29: end while

30:end procedureinconsolatait

The static task list approach presented by Cederman et al. is the most basic
task-oriented schedule: the kernel consists of two arrays: an “in-array” and
an “out-array” [18]. The in-array contains tasks to be executed by the
processors in the current iteration. As the kernel proceeds, processors
iterate over their statically assigned tasks (where a block with index _i_
handles all tasks at integer multiples of _i_), perform the task, and
potentially add additional tasks to the “out array”. This process continues
until the in-array is empty, at which point the out-array becomes the new in-
array. The kernel iterates between the two arrays until there are no remaining
tasks in either array. Like all queue methods, this approach requires that the
arrays be sized to hold the worst-case number of concurrently-queued tasks,
and that the processors use some form of synchronization when adding tasks to
the out-array. Note that Cederman et al. do not enforce synchronization on
reads since blocks execute tasks at predefined positions. This removes the
overheads inherent to synchronization, but also loses the greedy consumption
property described in the Task-Oriented Scheduling introduction. However, more
complex versions of this load-balancing schedule could add synchronization to
the in-array to enable greedy task pops from the head of the in-array. Aila et
al. present a similar method that performs only reads from a single task list
with atomic synchronization, and generates no additional work [4].

At the next level of complexity, centralized queues allow any processors on
the GPU to dynamically consume tasks from the head of the queue when they need
more work, or add tasks to the tail of the same queue when they generate new
work, which will eventually be consumed by processors without an intermediate
synchronization step, such as the one needed for the static task list schedule
to swap the in- and out-arrays. Many queue-based load-balancing variants
exist, such as variants that use mutual exclusion to control access to the
queue, variants that use cheaper atomic increments for synchronization,
variants that get an entire block’s worth of work with a single thread or
warp, variants that issue new work to the GPU’s kernel queue as additional
kernel launches, variants that use multiple queues to store tasks from
different stages of a rendering pipeline, and variants that add tasks from the
CPU and consume from the GPU [4, 20, 18, 91, 77, 22].

The distributed queue schedule presented by Zhang et al. assigns per-processor
queues with task indices stored in shared memory [92]. This allows each
processor’s threads to add and consume tasks from the block’s private queue,
and avoids the atomic contention overheads seen in the monolithic queue
schedules if the thousands of threads on a GPU were to compete for a single
globally-shared atomic. However, as GPU atomics have improved in performance
over time, atomic contention in modern centralized queues is no longer a
significant bottleneck [22]. The schedule also receives latency benefits from
storing task indices in shared memory. The primary downsides of this method is
that the GPU cannot load-balance between blocks, and each block’s individual
queue must be sized for the worst-case queue size. This may cause the
aggregate amount of memory used to be larger than the centralized task queue
variant. As part of this work, Zhang et al. present the CUIRREE library, which
provides general-purpose APIs and implementations for queue-based load
balancing.

##### Task Stealing

Evaluation
Reduces contention compared to the global queue method since processors
usually utilize their own private queue. Solves load imbalance between block
queues with stealing, and improves locality.

Task stealing queue variants provide per-processor task queues [19, 82, 92].
Processors typically both consume work from and add work to the tail of their
own queue. If a processor has no tasks available to process in its queue, it
may “steal” a task from the head of another processor’s queue. Since stealing
is the less common case, such schedules minimize the effect of atomic
contention that occurs in the monolithic queue case. Setting the queue so that
task pops come from the head, while task additions and steals come from the
tail, also ensures that the owning processor almost never needs to synchronize
when acquiring a new task. They also improve in data locality, since
neighboring tasks in time likely access nearby data.

However, task stealing schedules still have the problem that each individual
queue must be sized so that it could hold the worst-case of work that could be
simultaneously in the processor’s queue, which is not ideal on GPUs with
limited memory (both on-chip and off-chip). If the queues are too small, the
processor must stall until another processor steals work from a full queue.

##### Task Donation

Evaluation
When considering Task Stealing, Task Donation is even better. It allows
processors to offload overflow work and shrinks queue sizes.

Task Donation extends the work stealing schedule to also allow processors to
add overflow work to another processor’s queue that still has space available
[82]. This improves memory utilization on the system since now only the
aggregate capacity of all queues needs to be large enough to hold the maximum
amount of work. This also improves latency if queues are small enough to
reside in on-chip cache.

In many ways, task donation is the “ideal” version of a queue-based load-
balancing schedule. It preserves locality between work items, reduces
synchronization contention since processors likely work out of their own
queues most of the time, uses the same amount of memory as a global queue, and
preserves the ability to load balance between blocks using stealing and
donation.

##### Hierarchical Task Scheduling

Evaluation
Reduces contention at the global queue since a single thread can fetch a chunk
of work for the block, and can load balance within the block. Potentially
suffers from some load imbalance between blocks.

The centralized queue, task stealing, and task donation variants previously
described use only a single level of load balancing, where threads or blocks
retrieve work from the queue, but do not do anything _within_ the block to
balance the work. Consider, as an example, a breadth-first-search kernel where
the first thread in each block acquires a _chunk_ of work (multiple vertices)
on behalf of all threads in the block, perhaps in an attempt to reduce the
overheads of atomic contention that were to occur if each thread in the GPU
retrieved its work directly from the queue. This type of kernel would benefit
from an addional level of load balancing within the block, such as work-
oriented scheduling or binning. Chen et al. use a hybrid task/work scheduling
approach in Atos [22].

Another method presented by Steinberger et al. uses task lists in both global
memory and shared memory hierarchies [77]. The shared-memory queues are
faster, but have limited memory and cause load imbalance between blocks. The
global queues have significantly more space and can load balance between
blocks, but suffer from high-latency accesses and require more expensive
contention. Using two levels of queue hierarchy, where newly-generated work
items can go in either level of the hierarchy based on some heuristic or the
available storage space, finds a useful balance point. Additionally,
Steinberger et al. use a binning strategy, with separate queues for block-
sized, warp-sized, or thread-sized items to reduce load imbalance within a
processor.

###  4 Common Low-Level Algorithmic Primitives

In this section we identify the key building blocks that are used to construct
the load-balancing implementations discussed in Section 3. These building
blocks are used to address the irregularity/sparsity of their inputs and would
likely not be necessary if the inputs were dense or regular. For instance,
consider load-balancing an operation on every element of a matrix across two
workers. If that matrix is dense, it is likely simple to split it in half
knowing only the matrix dimensions such that each worker is assigned half of
the matrix elements. With a sparse matrix, determining the split point is more
complex; a robust solution must first count all non-zero elements in the
matrix, then identify the middle element to be the split point, and finally
assign half of the elements to each worker.

We first identify the challenges faced by the load-balancing implementations,
then the building blocks needed to address that challenge. Most load-balancing
strategies face challenges in two general areas: (1) Discovering work items
and work tiles and (2) scheduling/assigning work onto worker units.

####  1 Challenge: Counting Non-zeros or Work Items

Goal:  Several load balancers must first scan a sparse matrix or graph to
determine the number of non-zeros or work-items within each work tile of the
input data. For example, SpGEMM implementations using Gustavson’s row-wise
implementation [38] first count the number of non-zeros within each row of the
left-hand matrix to provide an estimate on the amount of work needed [37].
Similarly, kernels such as SpMV, SpMM and higher-order tensors using work-
oriented, thread-mapped, or group-mapped approaches first divide the sparse
matrix/tensor into regions of equal-sized non-zero blocks [87]. Target kernels
include nearly all sparse kernels—SpMV, SpGEMM, SpMM, and graph algorithms.

##### Implementation Primitives

When counting regions of work items within a workload, the following
algorithms are commonly used:

  * •

Parallel prefix sum, i.e., parallel scan, is a widely used algorithm
popularized by Guy Blelloch in the parallel setting [9]. Later, Harris et al.
proposed a parallel GPU implementation [41]. The algorithm takes as input (1)
a binary associative operator—which is $+$ for prefix sum—and (2) an array. It
then produces a new array where the element at any position is a sum of all
previous elements. Prefix-sum when applied to the work items per work tile for
all tiles produces an array that not only has the _total number of work_ to be
processed, but also the indices of this array (i.e. the location of each
element) represents the work tile which any given work item belongs to.
Parallel scan is especially useful when creating a balanced workload within
the GPU’s hierarchy [41]. The load balancer performs a parallel prefix sum at
a device-, block- or warp-wide level on an input array of expected work items
per task. Using the resulting array, the work is effectively partitioned onto
GPU threads. This implementation is particularly common in non-zero splitting
load balancers [8, 26, 88, 78]. It is also common in matrix algorithms
splitting work by row length [37]. Moreover, schedulers needing the total
amount of work, such as work-oriented approaches (Section 3), simply access
the last element of the prefix sum array.

A search on the prefix sum array can be used to determine which task is being
processed for any given work item. Examples of this process include a Single-
Source Shortest Path algorithms [28], Breadth-First Search algorithms [65],
various merging algorithms [36], and SpMV [35, 26, 64, 63]. Likewise, the IrGL
[74] compiler and the CUDA-quicksort algorithm [60] explicitly use prefix sums
to allocate space for work items.

  * •

Segmented scan and segmented reduce, closely related to prefix sum, reduces
data in a work tile or _segments_. Given an array of arrays, the segmented
_scan_ performs a prefix sum on each array within the larger array. Instead of
storing the entire scan array, a segmented _reduce_ stores the final, single
sum per input sub-array. Binning approaches requiring row-length or vertex
output degree information typically employ segmented reduce [33, 6, 67, 42].
Both Yan et al. and Liu et al. use segmented scan when partitioning sparse 2D,
3D and 4D tensors across threads [87, 56].

####  2 Challenge: Searching for Work Tiles and Items

Goal:  For several load-balancing schedules, each thread has a pre-assigned
number of work items to process, however, the individual threads may not know
exactly where in the total work pool they should begin processing their
assigned work, such that each thread has unique work items to process.
Furthermore, threads may not know which work tile an assigned work item
belongs to. For example, for kernels such as SpMV, if threads are assigned to
process non-zero values, they _need_ to determine (1) a range of unique non-
zeros to multiply and accumulate, and (2) the row and column for the
corresponding non-zero. The primitives below help each thread search, in
parallel, a sorted list of number of work items to determine their starting
position and the corresponding work tiles for a given work item.

##### Implementation Primitives

The following algorithms are commonly used (searches are called with bulk
inputs in parallel):

Figure 1: Shows a simple problem where 8 work items are perfectly balanced
across 4 threads, and a running sum of work items per work tile (prefix-sum
array) is provided as an input. Each thread processes work items range of
$[\text{thread\\_idx}\times\frac{8}{4}]\dots[(\text{thread\\_idx}+1)\times\frac{8}{4}]$
and uses a lower bound search into the prefix-sum array to determine the index
of the work tile for each work item.

  * •

Binary Search, operating on a sorted input, repeatedly divides the input in
half until the desired item is found. In approaches using parallel prefix sums
to divide work items, each thread performs a lower bound search (using binary
search) on the prefix sum array with the work item index as the input. The
lower bound search specifically returns the index of the first position where
the work item can be inserted. This index corresponds to the work tile the
input work item belongs to. Figure 1 shows an illustration of this process
with an example. In graph algorithms, threads may use this method to search on
the edge offset array (created using a prefix sum), to determine the source
vertex ID of an edge [28, 50, 84, 15].

  * •

Vectorized Sorted Search or Load-balanced Search, proposed by Sean Baxter in
ModernGPU [8], first sorts the queries to avoid the thread divergence caused
by the traditional binary search algorithm. Given a sorted list of queries,
$A$, and the sorted database $B$, it recasts the search problem into a merge
problem, where the program now linearly searches for $A$ in $B$. This improves
work-efficiency over binary search, with complexity $\mathcal{O}(A+B)$ versus
$\mathcal{O}(A\log{}B)$. Overall, each thread now has more locality in the
search since the keys will be sorted. Within a load balancing context,
vectorized sorted search is applied to a prefix sum array of work items
offsets to determine the index of the work tile that generated the work item.
Within graph algorithms, and similar to strategies using a binary search,
threads cooperatively apply vectorized sorted search to determine source
vertex IDs for a given edge [86, 28, 84].

####  3 Challenge: Efficient Binning of Work Items

Goal:  As mentioned in Section 4, several load balancers use various
heuristics to place work items into bins, such that each bin has a similar
amount of work. Finding a heuristic with low overhead and near-perfect
distribution of work is a critical step in developing an efficient
implementation.

##### Implementation Primitives

Some strategies used in the implementation of binning strategies include:

  * •

Sorting Based on Work Variance: Parallel sorting algorithms are often used as
a means to balance the workload in a given irregular application. To minimize
the preprocessing overhead of sorting, load balancers often employ this
strategy when the cost can be amortized over multiple runs of the same
application for a given input [33]. For example, in Sparse-Matrix Dense-Matrix
Multiplication (SpMM) for sparse deep learning workloads, Gale et al. first
reorder the sparse matrix by sorting its rows by length before assigning row
bundles to each warp [33]. Several other implementations for SpGEMM, SpMM,
triangle counting, and graph pattern mining likewise rely on an initial
sorting or reordering by the amount of work, number of non-zeros, or vertex
degree before assigning to threads/warps/blocks [13, 62, 5, 30, 45, 43, 26,
28].

####  4 Challenge: Assigning Compute Resources

Goal:  In dynamic approaches to load balancing, workers compete for compute
resources and/or work items. Such implementations need a way to determine the
outcome of the competition.

##### Implementation Primitives

Common primitives in determining resource assignment include:

  * •

Voting: This strategy is often used in conjuction with binning approaches,
where work is assigned to be processed by individual threads, warps, or CTAs,
depending on the work tile size [65, 28]. Since there are limited resources,
each worker unit competes for the resources of the entire group (warps or
CTAs), by writing, atomically, to a shared memory location. The winner of the
atomic write takes over the resources and proceeds with compute. Originally
designed by Merrill et al. [65], several works employ this strategy within
their own load balancing implementations [27, 47, 46, 84].

  * •

Dynamic Parallelism: Starting with architectures of CUDA capability 3.5 and
higher, kernels can be launched using dynamic parallelism, which allows CUDA
kernels to recursively launch child kernels. Each child kernel then has
complete access over the compute resources. This strategy can be used when
processing work tiles containing a large amount of work items. The load
balancer dynamically launches one or more kernels sized according to the
amount of work within the work tile [6, 15, 91, 17].

###  5 Summary of Load-Balancing Techniques

To the best of our knowledge, Table 2 and 3 characterizes all of the
predominant approaches practiced for load balancing irregular-parallel
problems on the GPU to the load-balancing taxonomy described in the previous
section.

Work(s) | Application Space | Method(s) | Implementation Detail | Characterization
---|---|---|---|---
Green et al. [35] Dalton et al. [26] Merrill and Garland [64] Merrill [63] |  Merging Algorithm, Sparse-Matrix Dense-Vector Multiplication |  Merge-Path Based decomposition |  Work (number of outputs and nonzeros) divided equally to number of threads. 2-D binary-search on row-offsets and nonzeros to figure out the active rows and nonzero indices. |  Static, Exact, Flat or Hierarchical (to reduce the search space per thread)
Baxter [8] Dalton et al. [26] Yang et al. [88] Steinberger et al. [78] |  SpMV, General Graph Traversal |  Nonzero Splitting |  Work (number of nonzeros) split equally to number of threads. 1-D binary search on row-offsets to find the active row for any given nonzero/thread. |  Static, Exact, Flat
Merrill et al. [65] Davidson et al. [28] |  Breadth-First Search, Single-Source Shortest Path |  Sequential Gather |  Vertices asigned to threads, neighbor-list processed sequentially |  Static, Approximate, Flat
Merrill et al. [65] Davidson et al. [28] |  Breadth-First Search, Single-Source Shortest Path |  Warp Gather |  Vertices assigned to warps, paralle prefix sum and binary-search used to asign a neighbor to each thread within the warp |  Static, Approximate, Hierarchical
Merrill et al. [65] Davidson et al. [28] |  Breadth-First Search, Single-Source Shortest Path |  Binning (Scan+Warp+CTA Gather) |  Place vertices with similar amount of work in the same bin, process large bins using a block, smaller bins using a warp and smallest using a thread. |  Dynamic, Approximate, Hierarchical, Non-Cooperative, Centralized (warp- and block-level)
Davidson et al. [28] |  Single-Source Shortest Path |  Edge partitioning |  Using parallel sorted search on the scanned edge offset queue find an intersection of block’s edge-list within the work list. |  Static, Exact, Hierarchical
Brahmakshatriya et al. [14] |  General Graph Traversal |  Binning |  Work items multiple of block-size processed using a block, trickle the remaining work to be (multiple of warp-size) processed by a warp, and finally, if more work is remaining, trickle to a thread. |  Dynamic, Exact, Hierarchical, Cooperative, Centralized
Green et al. [36] Fox et al. [32] |  Triangle Counting, Load-Balancing Algorithm |  Logrithmic Radix Binning |  Assign work to bins based on the logarithmic work estimate. Approximate method of reordering without sort. |  Dynamic, Approximate, Hierarchical, Non-Cooperative, Distributed
Baxter [8] |  GPU Parallel Primitives |  Vectorized Sorted Search (Load-Balanced Search) |  Upper-bound search of work-item indices with exclusive scan of the counts (subtract one). |  Static, Exact, Flat
Gale et al. [33] |  Sparse-Matrix Dense-Matrix Multiplication, Sampled Dense-Dense Matrix Product |  Row Binning and Bundling |  Sort row-indices based on a heuristics over row-length, grouping like-sized rows and tiles together |  Static, Approximate, Hierarchical (Block and Warp-wide)
Lee et al. [54] |  Sparse-Matrix Dense-Matrix Multiplication |  Block Reorganizer |  Group blocks into overloaded, normal, underloaded blocks based on the size of their work. Split large blocks, merge underloaded blocks. |  Static, Approximate, Hierarchical (Block-wide)
Zhang et al. [91] Busato and Bombieri [15] Busato and Bombieri [17] |  Graph Algorithms |  Dynamic Kernels |  Launch a dynamic parallel kernel configured to a given workload. |  Dynamic, Approximate, Hierarchical, Non-Cooperative, Centralized
Cederman and Tsigas [18] Aila and Laine [4] |  Ray Traversal |  Task List |  Work items popped from a read-only queue. If the blocks generate new work, this gets added to a write-only queue. Queues swap when the read-only queue is empty. |  Dynamic, Approximate, Flat, Cooperative, Centralized
Cederman and Tsigas [18] Chen et al. [20] Zhang et al. [92] Chen et al. [22] |  Dynamic Load-Balancing (Octree Partitioning), Molecular Dynamics (Atoms-Decomposition), Ray-Tracing, N-queens solver, Potential distribution, Asynchronous Graph Algorithms. |  Task Queue |  [18] Blocking and lock-free (atomic) queue variants. [20] CPU pushes, GPU blocks pop. [92] Device-wide and per-block queue variants (no stealing or donation.) [22] A single thread acquires tasks on behalf of an entire block to reduce synchronization overheads. |  Dynamic, Approximate, Flat, Cooperative, Centralized
Tzeng et al. [82] Cederman and Tsigas [19] Zhang et al. [92] |  General Load-Balancing (Rendering Primitives and Synthetic Workloads) |  Task Stealing |  Threads consume from per-block queues, can steal from other blocks’ queues |  Dynamic, Approximate, Hierarchical, Cooperative, Distributed
Tzeng et al. [82] |  General Load-Balancing (Rendering Primitives and Synthetic Workloads) |  Task Donation |  Per-block queues, can steal from or donate to other queues. |  Dynamic, Approximate, Hierarchical, Cooperative, Distributed

Table 2: Annotated bibliography of load-balancing techniques. _Continued in
Table 3._

Work(s) | Application Space | Method(s) | Implementation Detail | Characterization
---|---|---|---|---
Chen et al. [22] |  Asynchronous Graph Algorithms (BFS and PageRank) |  Hybrid task and work, load-balancing schedule |  Task queue shared by all blocks, parallel scan within blocks. |  Hybrid, Approximate (device), Exact (block), Hierarchical
Steinberger et al. [77] |  General Load-Balancing |  Hierarchical task queues, task binning |  Uses queues at both the device and block level. Queue items are binned into separate block-sized, warp-sized, or thread-sized queues. |  Dynamic, Approximate, Hierarchical, Cooperative, Centralized (device-level), Distributed (block-, warp- and thread-level)
Nisa et al. [67] |  Sparse Matricized Tensor Times Khatri-Rao Product |  Binning (inter-block load imbalance), Pre-partitioning CSF tensor (inter-warp load imbalance) |  Two strategies: 1\. Preprocessing step where each fiber (equivalent to row/column), is split such that each sub-fiber now contains a predefined maximum number of nonzeros (occupancy-based partitioning). Thread-warps are assigned to each sub-fiber. 2\. Configuration step where slices are binned (see Ashari et al. [6]). Thread-blocks are assigned a certain number of nnzs within a slice. ($\frac{\textit{nnz}}{\text{threads per block}})~{}=~{}\text{thread blocks}$ per slice. |  Static, Exact, Hierarchical
Ashari et al. [6] |  Sparse-Matrix Dense-Vector Multiplication |  Row binning and dynamic parallelism |  Target: CSR format, Row binning: place rows into bins with roughly the same number of nonzeros. For rows smaller than a threshold, bin-specific SpMV kernels are launched, where constant number of threads are assigned to each row. For bins with rows larger than a threshold, a separate grid is launched (through CDP) per row. Grid size is based on a parameter that determines compute per thread. Multiple thread-warps are assigned per row, based on nnzs in that row. |  Dynamic, Approximate, Hierarchical, Cooperative, Centralized
Liu et al. [56] |  SpMTTKRP, Sparse-Tensor Tensor Multiplication |  Group-Mapped (nonzero splitting) |  Introduces new F-COO format (variant of COO) to more easily access nnzs. Launches 2-D grids with 1-D thread blocks. On the sparse tensor, each block (x-axis) is assigned a set of nnzs, where each thread has the same number of nnzs. For the dense matrices each block (y-axis) is assigned a partition of the column. Partition size (nnzs and columns) is found by tuning. |  Static, Exact, Hierarchical
Yan et al. [87] |  Sparse-Matrix Dense-Vector Multiplication |  Pre-partitioning COO matrix and a segmented scan |  Introduces new blocked COO format for better reuse, where sparse matrix partitioned into vertically-aligned tiles. Thread blocks assigned equal number of nonzero tiles. Each thread within a thread-block is assigned equal number of consecutive nonzero tiles. Threads perform sequential segmented scan/sum. Blocks perform parallel scan. |  Static, Exact, Hierarchical
Winter et al. [85] |  Sparse General Matrix Multiplication |  (Thread-block) nonzero splitting of A matrix, (Thread-level) dynamically assign work items to each thread |  Partition the A matrix such that each thread block gest the same number of nonzeros. Within a thread block: Every thread is assigned the same number of work items. ESC algorithm is run within each thread block, where each thread works on computing a certain number of multiplies. Threads communicate through scratchpad memory. Device-wide prefix scan is run to merge all output chunks. |  Dynamic, Exact, Hierarchical, Cooperative, Distributed
Gremse et al. [37] Liu et al. [57] |  Sparse General Matrix Multiplication |  Nonzero splitting within a row |  (RMerge) A sub-warp with N threads is assigned N$\times$2 rows of the left-hand operand, where each row has N nonzeros. If the operand has more than A nonzeros in a row, the matrix is iteratively split into multiple matrices where each sub-matrix has up to N nnzs per row. |  Dynamic, Approximate, Flat
Parger et al. [75] |  Sparse General Matrix Multiplication |  (Thread-block) binning based on lightweight row-analysis of A and B, (Thread-level) hybrid load balancing, work-oriented |  Thread-block (global load balancing) analyzes row patterns of inputs to select one of 5 kernels. Assigns blocks based on binning of rows. Thread-level (local load balancing) Threads are split into groups where each group works on a nonzero of A (several rows of B). Group sizes chosen heuristically based on potential length of B row. |  Hybrid, Approximate, Hierarchical
Niu et al. [68] |  Sparse General Matrix Multiplication |  Group-Mapped |  Divide the matrices into evenly-shaped tiles, and assigns a warp to each sparse output tile. |  Static, Approximate, Flat

Table 3: (Continued) Annotated Bibliography—Summary and characterization of
research work on load balancing parallel-irregular problems on the GPU.

###  6 Optimizations Orthogonal to Load Balancing

High-performance load-balancing algorithms often need to take the GPU’s unique
architectural features into account. Critical characteristics include choice
of kernel configuration, the use of synchronization between parallel workers,
and management of on-chip memory. In this section, we discuss GPU
optimizations that are orthogonal to load balancing, and are often needed to
achieve high performance.

####  1 Kernel Strategy for Load-Balancing Operations

Although selecting an optimal kernel configuration strategy is independent
from a high-level load-balancing algorithm, it nevertheless can have a large
impact on its resulting performance. One significant configuration parameter
that influences kernel performance is _Occupancy_ , defined as the number of
blocks concurrently active on each Streaming Multiprocessor (SM) Occupancy
depends on several factors including the amount of shared memory used by each
thread block, the number of threads per block, the number of registers used
per thread, and the hardware capabilities of the targeted GPU. In general, the
more of a SM’s resources an individual block uses, the fewer blocks may run
concurrently on a given SM due to resource constraints. If a programmer does
not tune the kernel parameters such as block size, shared memory usage, or
register consumption correctly, they may find that they are underutilizing
their hardware. Ultimately, even with the kernel tuned to provide maximum
occupancy, an optimal grid and block size may be difficult to determine
without experimentation. Several works perform a simple sweep over block and
thread sizes, without any explanation of why one configuration is better than
another [35, 4], while other works account for nonobvious architectural
factors such as the warp scheduler overheads, register constraints, and
synchronization overheads [18, 19], or algorithmic benefits from
oversubscribing the number of blocks [65, 33].

After the programmer tunes their kernel’s resource usage to achieve maximum
occupancy and performance, there are two primary ways to configure the kernel.
The first, most common approach is to launch a large number of blocks, often
on the order of tens of thousands. The GPU’s block scheduler assigns a block
to an available SM, where it executes its assigned portion of the parallel
work and then terminates, allowing the block scheduler to assign a new block
to the SM. Programmers often pick the _Many-Blocks_ kernel configuration
because it is both easy to use and scalable. With hundreds, thousands, or even
tens of thousands of blocks available to the SM scheduler, a given kernel is
agnostic to the number of SMs present on a GPU as long as the number of blocks
reaches the minimum threshold to fully occupy the GPU. The same kernel can
scale down to mobile GPUs with only a handful of SMs, up to server GPUs with
over a hundred SMs (such as the NVIDIA A100), and can continue to scale as
future GPU releases include ever more SMs. In this kernel configuration, the
programmer is responsible for creating an appropriate division of work among
the blocks so that all blocks have approximately equal amounts of work.
Additionally, the programmer should pick a number of blocks equal to a
multiple of the number of SMs on the GPU if the architecture is known in
advance. This ensures that all SMs are fully utilized throughout the lifetime
of the kernel, and avoids problems of _Wave Quantization_ , where the last
wave of blocks scheduled onto the GPU occupies only a fraction of the GPU [73,
70]. By extension, the programmer should also ensure that the size of a tile
of work mapped to a block is is a multiple of the number of threads, to avoid
_Tile Quantization_.

The other common approach is to use a _Persistent Kernel_ , in which the
programmer launches just enough blocks to fully occupy all SMs of the GPU and
which stay resident on the SMs throughout the lifetime of the kernel,
iterating over work items in a loop as needed [4, 82, 20, 77]. Using a
persistent kernel grants programmers significant benefits such as reduced
kernel launch overheads, improved data reuse, and better load balancing.
However, this method requires careful management of GPU resources such as
shared memory, register usage, and thread count to ensure that the kernel runs
at maximum occupancy. To reduce kernel launch overheads and device-wide
synchronization overheads at kernel termination, a programmer may use an
_uberkernel_ , where within the kernel, blocks run in a continuous loop and
can select one of multiple paths within the kernel depending on the desired
computation. For example, in a multi-stage processing pipeline such as the
Reyes renderer proposed by Tzeng et al. [82], processors may alternate between
different section of the processing pipeline without device-wide barriers
between stages. A persistent kernel also improves data reuse in two primary
ways. First, by never exiting the kernel, a program set up as a pipelined
uberkernel may keep processed work in shared memory between stages of the
pipeline, rather than performing expensive reads and writes between the GPU
and off-chip memory between kernel launches. Secondly, initialization or
clean-up tasks required when a block first launches or exits only need to run
once, rather than each time a short-lived block acquires residence on a SM.
The persistent kernel method also gives load-balancing benefits in some
situations. If a single warp within a block has significantly more work than
the other warps, it will prevent new blocks from being scheduled to that SM.
If instead the programmer uses a persistent kernel approach, the other warps
within the block can move on to consume other work items as long as the
algorithm does not require a block-wide synchronization [4].

####  2 Synchronization Avoidance

Many load-balancing kernels use synchronization to ensure that parallel
workers create, consume, and process work correctly. For instance, consider a
queue used to sequence work, with many workers that want to push to or pop
from the queue [82, 18, 19, 20, 4, 92]. This algorithm requires each worker to
perform an exclusive modification to a shared queue so that all workers on the
device know where to add or remove work. Some queue implementations use mutex
locks on the queue, which are expensive to acquire and release. Other
implementations use atomic operations to ensure that only one worker can
access the queue at a time. Such synchronization approaches become
increasingly costly as the number of workers accessing the queue increases.
Modern GPUs support hundreds of concurrent blocks, equivalent to tens of
thousands of concurrent threads. One practical method of minimizing the impact
of synchronization is to interact with global memory queues at warp or block
granularity, rather than at thread granularity. This reduces the number of
atomics necessary for work creation and consumption, since a single thread in
a block can fetch a chunk of work for the entire block [20, 77]. The primary
downside of this method is that there may be some load imbalance within a
block if the work obtained by the block-wide chunk is not sufficient to supply
all threads with even amounts of work. Aila and Laine solve this problem by
periodically replacing finished work items with new work items on some threads
while long-running work continues on the busy threads [4]. In other cases, it
is possible to perform a second round of load balancing _within_ the block
[22]. Alternatively, Cederman and Tsigas reduce the synchronization overhead
of their queue by using lazy updates [18]. This queue implementation only
updates the queue’s head and tail pointers using expensive CAS operations
every $n$ accesses, and instead has each requesting block check several
consecutive positions in the queue offset from the head or tail to find the
true location between CAS updates. One important point to note is that as the
performance of atomic operations on modern GPUs continues to improve,
globally-synchronized, queue-based methods may be promising areas for
continued research.

In some algorithms atomic operations are not needed, and the kernel authors
tune their kernel implementations accordingly. For example, rather than using
a work queue, Merge Path [35, 64] precomputes the maximum amount of output
work items each block may produce and partitions the output storage
accordingly. This often requires a larger necessary output size (such as for
the output frontier in BFS) and additional computation at the beginning of the
load-balancing algorithm to determine such split points, but this method
eliminates the synchronization required to add work items to the output work
queue. Additionally, Merrill et al.’s BFS implementation [65] uses
unsynchronized bitmask reads and writes when updating a vertex’s “visited
flag”. This causes the kernel to add redundant work to the output frontier if
the result of one thread’s write is not visible to another thread’s read, but
the authors claim that the speedup from eliminating atomics is greater than
the additional work their kernel must do to remove duplicates. Although this
example is algorithm-specific, a programmer can often take one of multiple
approaches to implementing a kernel, and it is worth considering whether an
atomic-free implementation may show overall performance improvements despite
creating redundant work. Even kernel tuning has overlaps with synchronization
avoidance. For example, many kernels rely on block-wide thread synchronization
methods to ensure that all threads in a block are working on the same work
items. However, if a small number of threads within the block are not finished
with their assigned work, all other threads must sit idle. Tzeng et al. [82]
solve this problem by using blocks sized to be exactly 1 warp (32 threads),
which eliminates the need for explicit synchronizations since warps
automatically advance in lockstep.

####  3 Shared Memory

GPU kernels need to maximize memory bandwidth to achieve best performance. As
a result, the kernel programmer must take advantage of spatial and temporal
data locality, and must be aware of the characteristics of each level of the
memory hierarchy. Off-chip DRAM has the slowest bandwidth of any GPU memory
pool and experiences significant slowdowns from random accesses. To achieve
maximum DRAM bandwidth, a programmer must read or write a large sequential
chunk of data. On-chip shared memory, however, does not suffer from these
issues to the same degree On modern NVIDIA GPUs, a programmer may configure a
per-block pool of shared memory as either an explicitly-managed scratchpad
accessible by all threads, or as a cache where the GPU manages the contents of
the storage. For example, consider the merge path load-balancing methods,
which require binary searches along the diagonals of the input arrays to
determine work decomposition [28, 35, 64, 26]. These would be prohibitively
slow in global memory, but by using each block to read sequential chunks of
the array into explicitly-managed shared memory, the threads can then perform
the random memory accesses required by the binary searches in shared memory
with significantly reduced overhead and latency. Meanwhile, the accesses to
global memory use a smaller number of more efficient coalesced reads.

Ideally, the programmer should use vectorized reads to achieve maximum
throughput. However, sparse data formats make vectorized reads difficult,
since a data structure such as a sparse CSR matrix will likely not align with
the vector width of the memory system due to irregular row lengths. To solve
this challenge, Gale et al. introduce _Reverse-Offset Memory Alignment_ ,
which allows the kernel to use vector loads without padding the rows of the
sparse matrix with zeros [33]. Each block decrements its row offset to align
it with with the vector width, increases the number of nonzeros it needs to
process by the amount it decremented previously, and maintains a bitmask to
avoid duplicate work processing for the first values in the row.

Although the queue-based load-balancing methods achieve high-performance load
balancing through the use of both work queues and synchronization variables
stored in DRAM, blocks must suffer long-latency memory accesses when
attempting to synchronize, add work, or consume work. An alternative to a
single DRAM queue is to use distributed per-block queues as presented in
CUIREE [92]. These queues could be stored in either global memory or shared
memory (depending on their size), which presents an interesting trade-off
where some kernels may benefit from device-wide load balancing, and others may
benefit from faster access to the queue in the block’s shared memory at the
expense of using load balancing only within a block. If the queue is too large
to fit in shared memory, distributed queue systems can still take advantage of
shared memory by keeping per-block synchronization and indexing variables in
the faster shared memory. Whippletree [77] extends this concept to use small,
fast, per-block queues stored in each block’s shared memory. Under normal
runtime conditions, blocks consume work from and add new work to their private
queues. If the block runs out of space in the shared memory queue, or if there
is no remaining work in the shared memory queue, Whippletree provides a global
memory queue used for overflow data.

Explicitly managing shared memory to store frequently-used data is critical
for many kernels, but others can instead benefit from configuring on-chip
storage as a GPU-managed cache. For example, Gale et al. [33] evaluate a
trade-off between using on-chip memory as explicitly-managed shared memory to
store a transpose of the input matrix, versus leaving the memory as a cache
while re-computing the matrix transpose in registers as needed. The authors
determine that increased L1 cache capacity is more important for their Sampled
Dense-Dense Matrix Multiplication (SDDMM) kernel than reducing compute
complexity. We hypothesize that this is because other parts of the kernel,
separate from the matrix transpose, benefit from the cache’s ability to expose
spatial and temporal reuse more than the kernel benefits from an explicitly
managed transpose scratchpad. When developing GPU kernels, the programmer
should consider whether explicitly managing a portion of the data is
beneficial, or if it hurts performance by eliminating caching opportunities in
other areas of the program’s data accesses.

###  7 Conclusion: A Look Ahead

The focus of this work is to design a method for characterization of fine-
grained GPU load-balancing techniques that are key to high-performance for
irregular sparse kernels. Our survey uniquely builds the taxonomy of fine-
grained load-balancing, a deeper dive on key techniques practiced in
literature today, and an understanding of implementation details and building-
block algorithms required to implement these load-balancing algorithms. Our
goal is to fasicilitate development of future work in this space that further
accelerates irregular problems on the GPUs, and encourages support for load-
balancing primitives as a first-class software and hardware citizen.

The aforementioned survey gives us not only a broad, but also a very deep
understanding of the load-balancing landscape. In the chaper that follows we
bring the knowledge of common load-balancing patterns to bear and build a
powerful load-balancing abstraction for GPUs. We strongly believe that this is
an important first step towards a future that facilitates ease of programming
and portability of code for sparse-irregular problems.

##  Chapter 3 A Programming Model for GPU Load Balancing

The survey in Chapter 2 shows that the predominant approach to solving the
load imbalance problem is through application-specific load-balancing
techniques that aim to evenly distribute the work such that each thread gets
the same number of work items to achieve maximum performance (for instance,
Merrill and Garland’s load-balanced SpMV implementation [64]). These load-
balancing techniques are often tightly coupled with the application itself.
The load-balancing components within these implementations are both complex
and are often collectively the most significant contributor to the performance
of an application. The work presented in this chapter generalizes today’s
application-specific load-balancing algorithms into a clean, modular, powerful
abstraction that can be applied to many complex irregular workloads.

In the process of building our abstraction, we identified common load-
balancing approaches deployed within sparse, irregular applications on GPUs:
application-specific frameworks such as GraphIt [14], Gunrock [84], and
GraphBLAST [89]; techniques from low-level CUDA libraries such as ModernGPU
[8] and CUB [63]; and other hand-coded implementations of load-balancing
algorithms within applications such as SpMV/SpMM [64, 28, 33], triangle
counting [32, 36] and breadth-first search [65, 15]. We show that with a
simple, intuitive, powerful abstraction, these load-balancing schedules can be
extended to support irregular workloads that are more general than the
specific problem for which they were designed. We demonstrate this by using
sparse-linear-algebra-based load balancing for data-centric graph traversal
kernels.

Writing high-performance load-balancing code is complex, in large part because
this code must perform many roles. Among other tasks, it must ingest data from
a specific data structure, perform user-defined computation on that data, and
schedule that computation in a load-balanced way. The key insight in our
abstraction is to separate the concerns between workload mapping (the load-
balance task) and work execution (the user-defined computation), where we
_map_ sparse formats (such as Compressed Sparse Row (CSR)) to simple
abstraction components called work atoms, tiles, and sets. These fundamental
components are expressed as composable C++ ranges and range-based for loops,
and are used to build load-balancing schedules. Programmers can then use these
APIs to build load-balanced, high-performance applications and primitives.
Expressed in this way, we can reconstruct existing application-dependent load-
balancing techniques that address irregularity to be more _general_ ,
_portable_ , and _programmable_. The contributions of our work are as follows:

  1. 1.

We present a novel abstraction for irregular-parallel workloads on GPUs. Our
abstraction at a high level allows programmers to develop sparse, irregular-
parallel algorithms with minimal code while delivering high performance.

  2. 2.

We design and implement a set of intuitive APIs, available in our open-source
GPU load-balancing framework, built on the proposed abstraction using CUDA-C++
ranges and range-based for loops.

  3. 3.

We show the ease of implementing new load-balancing schedules by implementing
a novel cooperative groups-based load-balancing schedule (Section 2), which is
a generalization of previous thread-, warp-, and block-level load-balancing
schedules [88].

  4. 4.

We provide best-in-class SpMV performance as a benchmark with a geomean of
speedup of 2.7$\times$ for the SuiteSparse Matrix Collection [29] over
cuSparse’s state-of-the-art implementation using simple heuristics and 3 GPU
load-balancing schedules.

###  1 Design Goals

The design goals of our load balancing abstraction are as follows:

###### Achieve high performance.

First and foremost, the goal of our work is to achieve the high performance of
existing load balancing algorithms for irregular applications. Our abstraction
cannot come at the cost of significant overhead or performance degradation. We
measure our success in achieving high performance by comparing the performance
of our abstraction against the performance of existing hardwired
implementations.

###### A composable and programmable interface.

Importantly, we do not want to restrict the user to a library interface that
takes control of the larger system. Programmers strongly prefer to adopt new
software components that fit into their control structures rather than require
them to adopt a new control structure. We want to allow the users to (1)
maintain control of GPU kernel boundaries (kernel launches), (2) be able to
add new load-balancing algorithms, and (3) compose new load-balanced
primitives from existing load-balancing APIs. We measure the programmability
of our work by comparing the Lines of Code (LOC) of our abstraction against
existing implementations and show composability by implementing a new load-
balancing algorithm in terms of our existing APIs.

###### Extensible to new applications.

We aim to decouple and extend application-specific load-balancing techniques
to new irregular-parallel domains. Our abstraction seeks to promote the reuse
of existing load-balancing techniques for new applications. We use SpMV as a
benchmark application implemented using three different load-balancing
techniques, some of which were previously used to implement parallel graph
analytics kernels [28, 14, 84, 15].

###### Facilitate the exploration of optimizations.

A key goal of our abstraction is to facilitate the exploration of
optimizations for a given application by switching the underlying load-
balancing algorithms used to balance the work. We want to encourage our users
to experiment with heuristics and new load-balancing techniques to discover
what works best for their application needs. We measure the success of this
goal by optimizing SpMV’s performance response for a large corpus of sparse
matrices across several different load-balancing techniques.

#### Non-Goals

In addition to the above design goals, we also define our non-goals:

###### Targeting other parallel architectures.

Although we believe the lessons learned should apply to other parallel
architectures, we explicitly target NVIDIA’s CUDA architecture and programming
model [69]. Many components of our abstraction leverage CUDA’s compute
hierarchy of threads, warps and blocks mapped onto the physical streaming
multiprocessors, the oversubscription model of assigning more work than the
number of processors to fully saturate the underlying hardware, and CUDA’s
Cooperative Groups programming model [40], described in Section 2, to achieve
high performance.

###### Multi-GPU support.

This work focuses on load-imbalance issues for a single GPU and does not
consider multi-GPU single-node or multi-node systems, although these are
interesting directions for future work.

###  2 Load Balancing Abstraction

Figure 1: Load balancing as a simple pipeline of the three key concepts of our
abstraction: (1) sparse data structures represented as iterators, (2) load-
balancing algorithm that partitions the work onto threads, and (3) user-
defined computation consuming the balanced work and executing on each thread.

The key insight behind our GPU load balancing abstraction is the _separation
of concerns_ between the mapping of the work items to processing units and
work execution. We divide our abstraction into three key concepts (illustrated
in Figure 1), each of which describes a different aspect of an implementation:
(1) defining the work; (2) defining the workload balance across GPU threads,
warps or blocks; and (3) defining the work execution and computation per
thread on the balanced work. This separation allows us to distinctly divide
the work between an application developer and a load-balanced-library
developer and facilitates the exploration of optimizations by mixing different
load-balancing techniques and sparse-irregular algorithms. Sidebar 1 presents
a practical example of the motivation for our load balancing abstraction.

####  1 Input from Sparse Data Structures

We begin with our input data expressed in some form of sparse data structure.
Examples of such data structures include, but are not limited to, Compressed
Sparse Row (CSR) and Coordinate (COO) formats. The goal of the first stage of
our abstraction is to map the input data format to a common data framework and
vocabulary that is the input to the next stage. This vocabulary has three
simple components that together express the input data:

  1. 1.

A work atom, a single unit of work that is to be scheduled onto the processors
(for example, a non-zero element of a sparse matrix). We assume that all work
atoms have an equal cost during execution.

  2. 2.

A work tile, a logical entity represented as a set of work atoms (for example,
a row of a sparse matrix). Work tiles may have different costs during
execution. As we highlighted in the introduction, work is most _logically_
parallelized over work tiles but is often most _efficiently_ parallelized over
work atoms, and mapping between work tiles and work atoms may be expensive and
complex.

  3. 3.

A tile set, a set of work tiles that together comprise the entire working
problem (for example, a sparse matrix). In our abstraction, the tiles within a
tile set must be independent (and thus can run in parallel across multiple
processors).

This mapping between sparse formats and atoms/tiles/tile sets is defined by
the user. Though we have not implemented all of them, we believe our mapping
abstraction here is flexible enough to express a wide variety of existing
sparse data formats in the literature [31] in such a way that they are
suitable for load balancing in our abstraction’s next stage. As well, we have
already included several common sparse formats (CSR, CSC, COO) in our load-
balancing library implementation so that users can simply select and use them
without having to implement them. Given a mapping to atoms/tiles/tile sets, we
can next implement a load-balancing algorithm that can parallelize over work
atoms or tiles transparently from the computation’s perspective.

Sidebar 3.1 A practical example of existing, predominant approach to load
balancing sparse-irregular workloads. Consider an SpMV implementation on the
GPU provided in the open-source CUDA CUB library [63]. CUB implements and
maintains the SpMV algorithm presented in the paper by Merrill and Garland
[64]. Merge-based SpMV, explained in detail in Section 2, is a CSR-based,
perfectly load-balanced SpMV, where each thread gets an even share of work,
and the amount of work is defined by the total number of matrix rows and the
total number of non-zeros, summed. In the reference, this highly efficient,
state-of-the-art implementation took 1,100 lines of code (LoC) (or 503 LoC of
kernel code) across 3 files (not including a 4th file required for a segmented
fixup step of an additional 234 LoC). In contrast, the actual computation of
SpMV within this reference implementation is expressed within a _single_ for-
loop and 4–5 LoC! This disparity between the LoC required to map the work
items to processing units in a load-balanced way and the LoC required to
express the desired computation is the key motivation behind our work.
Additionally, the CUB implementation is specifically dedicated to the SpMV
algorithm and would require a significant rewrite to apply it to other
algorithms, even within the same computing domain. One such example of this
exact rewrite is by Yang et al., where the authors extend merge-path load
balancing from SpMV to Sparse-Matrix Dense-Matrix Multiplication (SpMM)
algorithm [88]. The load balancing algorithm in both works is the same but
applied to different computations, which motivates the need for reuse.

####  2 Defining Load Balancing

By expressing workloads through an abstraction that captures work at differing
levels of granularity (i.e., tile set, atoms, and tiles), we can more easily
distribute computation evenly across the GPU’s available resources. Given a
user-defined input tile set and associated sequences of atoms and tiles, along
with a user-selected partitioning algorithm, our load balancing stage outputs
subsequences of atoms and tiles assigned to processor ids (i.e., where atoms
or tiles will be processed).

The resulting assignment of subsequences to processor ids is critical to
effectively balancing workloads across processing elements and is generally
problem- and dataset-specific. The user must specify the necessary sequences.
Ideally, an _oracle_ would take these sequences and select the most optimal
subsequences for every processing element. Finding such an oracle is an open
problem and as such, we provide the next best thing: the ability for a user to
choose and experiment from set of predefined schedules and the ability to
implement their own schedules. In general, load-balancing algorithm designers
must balance between the cost of scheduling and the benefits from better
scheduling. A schedule could be as straightforward as assigning processing
elements to tiles with arbitrary numbers of atoms (e.g., rows with an
arbitrary number of non-zeros in a sparse matrix) to something more
complicated/expensive that takes on a more holistic approach to work (e.g.,
considering work across multiple rows with a varying number of non-zeros in a
sparse matrix).

####  3 Defining Work Execution

The final component of our load-balancing abstraction expresses the irregular-
parallel computation itself. The previous stage inputs load-imbalanced work
and load-balances it; this stage then consumes that load-balanced work by
performing computation on it. The scope of what computation can be expressed
is extensive, and is only limited by how the load-balanced work, represented
as sequences, can be consumed within a CUDA kernel. Since the framework does
not assume control of the kernel, anything you can write in a CUDA kernel will
also work in our framework. For instance, programmers can express a
mathematical operation performed on each atom or each tile of the work, or
build cooperative algorithms that not only consume the work assigned to each
thread but also combine the results with neighboring threads to implement more
complex algorithms such as parallel reduce or scan. Practical examples that we
have implemented in our framework (see Section 3 and 3) using this abstraction
include, but are not limited to, sparse-linear algebra kernels, such as
Sparse-Matrix and Sparse-Tensor contractions, and data-centric parallel graph
algorithms, such as Single-Source Shortest Path (SSSP) and Breadth-First
Search (BFS) built on a neighborhood traversal kernel.

We expect typical users of our library will _only_ write their own code for
_this_ stage of the abstraction and use standard data structures and load-
balancing schedules that are already part of our library. However, those users
can also implement custom data formats and load-balancing schedules.

###  3 High-Level Framework Implementation

Our GPU load-balancing framework implements the abstraction described in
Section 2 using C++17 and CUDA. In our system, programmers use CUDA/C++ to
develop irregular-parallel algorithms and implement new load-balancing
schedules. Per our design goals of composable APIs, extensibility and reuse,
this and the following section introduce the implementation details of our
API, and how it is used to develop new applications that promote the reuse of
high-performance load-balancing techniques available within the framework. We
also explore a new load-balancing method (Section 2) built on CUDA’s
Cooperative Groups model. Furthermore, we identify how our work can be used to
facilitate the exploration of optimizations for a given application such as
SpMV.

####  1 Implementing Sparse Data Structures

Our framework translates sparse data structures (e.g., COO, CSR, CSC) into
work atoms, work tiles, and tile sets (Section 1) using simple C++ iterators.
C++ iterators are objects that point to some element in a range of elements
and enable iteration through the elements of that range using a set of
operators. For example, a counting_iterator is an iterator that represents a
pointer into a range of sequential values [1]. Our framework requires the user
to define three important iterators using C++: (1) an iterator over all work
atoms; (2) an iterator over the work tiles; and (3) an iterator over the
number of atoms in each work tile. (Our library already supports several
common sparse data structures.) Using these iterators, the load-balancing
schedule can then determine and distribute load-balanced work across the
underlying hardware. Listing 2 shows how our abstraction expresses the
commonly used CSR format as a tile set within our framework.

Figure 2: Compressed-Sparse Row (CSR) format expressed within our framework
using C++17. The CSR format describes a matrix using three arrays: (1) column
indices of nonzero values; (2) the extent of rows (row offsets); and (3) the
nonzero values. Since the CSR data structure does not contain arrays that
point to indices of atoms and tiles (nonzeros and rows), in the listing above
we define atom and tile iterators as simple counting iterators from 0 to the
total number of nonzeros (nnz) and from 0 to the total rows in the matrix
(rows), respectively (Lines 2 and 3). The iterator over the atoms-per-work-
tile is expressed using a transform iterator, which computes the expression
within a provided function for each tile id. For CSR, this is simply the row
offset of the current tile subtracted from the offset of the next tile (Lines
5–11).

⬇

// Simple iterators for atoms and tiles.

counting_iterator<int> atoms_iter(0, nnz);

counting_iterator<int> tile_iter(0, rows);

// Iterator over the atoms within tile i.

auto atoms_per_tile = make_transform_iterator(

tile_iter,

[tile_iter, row_offsets]

__host__ __device__(const int& i) {

return (row_offsets[tile_iter[i + 1]] -

row_offsets[tile_iter[i]]);

});

####  2 Implementing Load-Balancing Schedules

Perhaps the most straightforward schedule is scheduling each work tile onto
one GPU thread. This approach is common in the literature and practice [65,
28, 78, 8, 88]; although this strategy is ineffective in the presence of