2019-11-20 08:55:34 +08:00
[README ](/README.md#documentation ) > **Efficient GEMM in CUDA**
# Efficient GEMM in CUDA
2022-11-19 22:02:15 +08:00
CUTLASS implements the hierarchically blocked structure described in
2019-11-20 08:55:34 +08:00
[CUTLASS: Fast Linear Algebra in CUDA C++ ](https://devblogs.nvidia.com/cutlass-linear-algebra-cuda/ )
and the [CUTLASS GTC2018 talk ](http://on-demand.gputechconf.com/gtc/2018/presentation/s8854-cutlass-software-primitives-for-dense-linear-algebra-at-all-levels-and-scales-within-cuda.pdf ).
## Hierarchical Structure
The basic triple loop nest computing matrix multiply may be blocked and tiled to match
2022-11-19 22:02:15 +08:00
concurrency in hardware, memory locality, and parallel programming models. In CUTLASS,
2019-11-20 08:55:34 +08:00
GEMM is mapped to NVIDIA GPUs with the structure illustrated by the following loop nest.
```c++
for (int cta_n = 0; cta_n < GemmN ; cta_n + = CtaTileN ) { / / for each threadblock_y } threadblock-level concurrency
for (int cta_m = 0; cta_m < GemmM ; cta_m + = CtaTileM ) { / / for each threadblock_x }
2022-11-19 22:02:15 +08:00
for (int cta_k = 0; cta_k < GemmK ; cta_k + = CtaTileK ) { / / " GEMM mainloop " - no unrolling
2019-11-20 08:55:34 +08:00
// - one iteration of this loop is one "stage"
//
for (int warp_n = 0; warp_n < CtaTileN ; warp_n + = WarpTileN ) { / / for each warp_y } warp-level parallelism
for (int warp_m = 0; warp_m < CtaTileM ; warp_m + = WarpTileM ) { / / for each warp_x }
//
2021-07-23 12:40:53 +08:00
for (int warp_k = 0; warp_k < CtaTileK ; warp_k + = WarpTileK ) { / / fully unroll across CtaTileK
2019-11-20 08:55:34 +08:00
// - one iteration of this loop is one "k Group"
//
for (int mma_k = 0; mma_k < WarpTileK ; mma_k + = MmaK ) { / / for each mma instruction } instruction-level parallelism
for (int mma_n = 0; mma_n < WarpTileN ; mma_n + = MmaN ) { / / for each mma instruction }
for (int mma_m = 0; mma_m < WarpTileM ; mma_m + = MmaM ) { / / for each mma instruction }
2022-11-19 22:02:15 +08:00
//
2019-11-20 08:55:34 +08:00
mma_instruction(d, a, b, c); // TensorCore matrix computation
} // for mma_m
} // for mma_n
} // for mma_k
} // for warp_k
} // for warp_m
} // for warp_n
} // for cta_k
} // for cta_m
} // for cta_n
```
This tiled loop nest targets concurrency among
2022-11-19 22:02:15 +08:00
- threadblocks,
- warps, and
- CUDA and Tensor Cores.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
It takes advantage of memory locality within
- shared memory and
- registers.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
The figure below illustrates the flow of data within this structure.
This is the hierarchical GEMM computation embodied by CUTLASS. Each stage depicts a
nested level of tiling which corresponds to a layer of concurrency within the CUDA execution model and to a
2019-11-20 08:55:34 +08:00
level within the memory hierarchy, becoming increasingly finer moving left to right.
### Threadblock-level GEMM
Each threadblock computes its portion of the output GEMM by iteratively loading tiles of input
2022-11-19 22:02:15 +08:00
matrices and computing an accumulated matrix product. At the threadblock level, data are loaded from
global memory. The blocking strategy in general is key to achieving efficiency. However, the programmer
must balance multiple conflicting goals. A
2019-11-20 08:55:34 +08:00
larger threadblock means fewer fetches from global memory, thereby ensuring that DRAM bandwidth
2022-11-19 22:02:15 +08:00
does not become a bottleneck.
2019-11-20 08:55:34 +08:00
However, large threadblock tiles may not match the dimensions of the problem well. If either the
GEMM _M_ or _N_ dimension is small, some threads within the threadblock may not perform meaningful
work, as the threadblock may be partially outside the bounds of the problem. If both _M_ and _N_
are small while _K_ is large, this scheme may launch relatively few threadblocks and fail to
2022-11-19 22:02:15 +08:00
make full use of all multiprocessors within the GPU. Strategies to optimize performance for this case,
as described in the section [Parallelized Reductions ](efficient_gemm.md#parallelized-reductions ),
partition the GEMM K dimension across multiple threadblocks or multiple warps. These threadblocks
or warps compute matrix products in parallel; the products are then reduced to compute the result.
2019-11-20 08:55:34 +08:00
In CUTLASS, the dimensions of the threadblock tile are specified as `ThreadblockShape::{kM, kN, kK}`
and may be tuned to specialize the GEMM computation for the target processor and dimensions of
the GEMM problem.
### Warp-level GEMM
The warp-level GEMM maps to the warp-level parallelism within the CUDA execution model. Multiple
warps within a threadblock fetch data from shared memory into registers and perform computations.
2022-11-19 22:02:15 +08:00
Warp-level GEMMs may be implemented either by TensorCores issuing
[mma.sync ](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma )
or [wmma ](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-wmma-mma )
instructions, or by thread-level matrix computations issued to CUDA cores.
2019-11-20 08:55:34 +08:00
For maximum performance, access to shared memory should be bank conflict free. To maximize data
reuse within the warp, a large warp-level GEMM tile should be chosen.
### Thread-level GEMM
At the lowest level of blocking, each thread is responsible for processing a certain number of
2022-11-19 22:02:15 +08:00
elements. Threads cannot access each other's registers, so we choose an organization that enables
reuse of values held in registers for multiple math instructions. This results in a 2D tiled
2019-11-20 08:55:34 +08:00
structure within a thread, in which each thread issues a sequence of independent math instructions
to the CUDA cores and computes an accumulated outer product.
SGEMM, IGEMM, HGEMM, and DGEMM are computed by SIMT math instructions issued by thread-level matrix multiply
procedures.
## Epilogue
The above code focuses only on the matrix multiply computation **C = AB** whose result is
held in the registers of each thread within the threadblock. The mapping of logical elements
in the output tile to each thread is chosen to maximize performance of the matrix multiply
computation but does not result in efficient, coalesced loads and stores to global memory.
The epilogue is a separate phase in which threads exchange data through shared memory then
cooperatively access global memory using efficient striped access patterns. It is also
the phase in which linear scaling and other elementwise operations may be conveniently
computed using the matrix product results as inputs.
CUTLASS defines several typical epilogue operations such as linear scaling and clamping,
but other device-side function call operators may be used to perform custom operations.
## Optimizations
2022-11-19 22:02:15 +08:00
The hierarchical structure described above yields an efficient mapping to the CUDA execution model and
2019-11-20 08:55:34 +08:00
CUDA/TensorCores in NVIDIA GPUs. The following sections describe strategies for obtaining peak performance
for all corners of the design space, maximizing parallelism and exploiting data locality wherever possible.
### Pipelining
The blocked structure demands a large storage allocation within the registers of each CUDA thread. The
2022-11-19 22:02:15 +08:00
accumulator elements typically occupy at least half a thread's total register budget. Consequently,
2019-11-20 08:55:34 +08:00
occupancy -- the number of concurrent threads, warps, and threadblocks -- is relatively low compared
2022-11-19 22:02:15 +08:00
to other classes of GPU workloads. This limits the GPU's ability to hide memory latency and other stalls
2019-11-20 08:55:34 +08:00
by context switching to other concurrent threads within an SM.
2022-11-19 22:02:15 +08:00
To mitigate the effects of memory latency, CUTLASS uses *software pipelining* to overlap memory accesses
with other computation within a thread. CUTLASS accomplishes this by double buffering at the
following scopes.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
- **Threadblock-scoped shared memory tiles:** two tiles are allocated in shared memory.
One is used to load data for the current matrix operation,
while the other tile is used to buffer data loaded from global memory
for the next mainloop iteration.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
- **Warp-scoped matrix fragments:** two fragments are allocated within registers.
One fragment is passed to CUDA and TensorCores during the current matrix computation,
while the other is used to receive shared memory fetch returns
for the next warp-level matrix operation.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
The following diagram illustrates the efficient, pipelined mainloop body used in CUTLASS GEMMs.
2019-11-20 08:55:34 +08:00
### Threadblock Rasterization
To maximize reuse of data held in the last level cache, CUTLASS defines several functions to
affect the mapping of threadblocks to logical partitions of the GEMM problem. These map
consecutively launched threadblocks to packed two-dimensional regions of the partitioned GEMM
problem to increase the probability that these will access the same tiles of global memory at
approximately the same time.
2021-07-29 08:43:27 +08:00
Several functions are defined in [cutlass/gemm/threadblock_swizzle.h ](/include/cutlass/gemm/threadblock/threadblock_swizzle.h ).
2019-11-20 08:55:34 +08:00
### Parallelized Reductions
**Split K - reduction across threadblocks**
Matrix product computations expose parallelism among _O(MN)_ independent inner product
computations. For sufficiently large problem sizes, a GEMM kernel in CUTLASS may approach
the theoretical maximum computational throughput. For small problems, however, there are
too few threadblocks to efficiently occupy the entire GPU.
As a recourse, parallelizing the reduction performed during the inner product computation
enables more threadblocks to execute concurrently while still taking advantage of the throughput
benefits of large threadblock-level GEMM tiles.
CUTLASS implements parallel reductions across threadblocks by partitioning the GEMM _K_ dimension
and launching an additional set of threadblocks for each partition. Consequently, we refer to
2022-11-19 22:02:15 +08:00
this strategy within CUTLASS as "parallel reduction splitK." The "parallel reduction splitK" strategy
requires the execution of 2 kernels: partitionedK GEMM, and batched reduction.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
PartitionedK GEMM resembles one flavor of batched strided GEMM. Instead of requiring users
to specify the problem size of each batch, partitionedK GEMM asks for the overall problem size and the
number of partitions that will be applied along the K dimension for operands A and B. For example,
parameters of m=128, n=128, k=4096 and partition=16 will result in 16 batched strided GEMMs
with each batch of m=128, n=128, k=256. PartitionedK also allows scenario where k is not divisible
by the partition count.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
For example, parameters of m=128, n=128, k=4096 and partition=20
will result in 20 batched strided GEMMs.
The first 19 batches will have m=128, n=128, and k=4096/20=204,
and the last batch will have m=128, n=128, and k=220.
2019-11-20 08:55:34 +08:00
2022-11-19 22:02:15 +08:00
The batched reduction kernel takes as input the output (C) of partitionedK GEMM,
and performs a reduction along the K-dimension.
Users must manage workspace memory to store this intermediate result.
2019-11-20 08:55:34 +08:00
**Sliced K - reduction across warps**
2022-11-19 22:02:15 +08:00
Similar to the split-k scenario, sliced-k aims at improving the efficiency of kernels
with smaller M and N dimensions, but large K dimension.
At the thread-block level, the parameters CtaTileN and CtaTileM expose parallelism
by partitioning the work among warps.
Larger warpTiles expose better instruction-level parallelism (ILP) and reuse,
but also limit the number of warps running per threadblock, which reduces efficiency.
In order to improve efficiency in such scenarios, partitioning the warpTiles also along ctaTileK
helps use the hardware more efficiently by allowing more warps to run concurrently in a CTA.
Sliced-k kernels break down a threadblock's computation among participating warps
not just among the CtaTileN, CtaTileM dimension, but also the CtaTileK dimension.
Thus, sliced-k entails a small cost in form of a reduction
which has to happen at the end among the participating warps.
This is because each warp computes using only a "slice" of CtaTileK,
so each warp only has a partial sum before the reduction.
2019-11-20 08:55:34 +08:00
# Resources
The following additional resources describe design and implementation details of GEMMs
targeting NVIDIA GPUs.
2020-06-09 07:17:35 +08:00
- [Developing CUDA Kernels to Push Tensor Cores to the Absolute Limit on NVIDIA A100. ](https://www.nvidia.com/en-us/gtc ) (SR 21745)
2019-11-20 08:55:34 +08:00
- [CUTLASS: Fast Linear Algebra in CUDA C++ ](https://devblogs.nvidia.com/cutlass-linear-algebra-cuda/ )
- [CUTLASS: SOFTWARE PRIMITIVES FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES WITHIN CUDA ](https://on-demand-gtc.gputechconf.com/gtcnew/sessionview.php?sessionName=s8854-cutlass%3a+software+primitives+for+dense+linear+algebra+at+all+levels+and+scales+within+cuda )
- [Programming Tensor Cores: NATIVE VOLTA TENSOR CORES WITH CUTLASS ](https://developer.download.nvidia.com/video/gputechconf/gtc/2019/presentation/s9593-cutensor-high-performance-tensor-operations-in-cuda-v2.pdf )
- [CUDA Programming Guide: warp matrix functions ](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#wmma )
- [Matrix Multiply Accumulate Instructions ](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma )
# Copyright
2022-04-24 03:02:38 +08:00
Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
SPDX-License-Identifier: BSD-3-Clause
2019-11-20 08:55:34 +08:00
```
2022-04-24 03:02:38 +08:00
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
2019-11-20 08:55:34 +08:00
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
```
