Introduction

Hi everyone !

In this post, I will share with you all the steps to write an optimized FP32 matrix multiplication on AMD RDNA3 GPU outperforming rocBLAS by 60%. I will cover some basics and explain all the optimizations I have implemented. This will be done in a iterative way in 8 differents Kernels.

Figure 1: sneak peek of the performance results

I primary intended to work on this to deepen my understanding of RDNA3 and try out HIP and I felt like I needed to share what I learned doing this :).

Few things I like to say before we start :

That being said, let’s start !

Problem statement

There is a lot of research happening on the way to improve the performance of matrix multiplication nowadays. Being a core algorithm in ML applications, any FLOPS we can exploit is golden.

Before proceeding, let’s recall the basics of matrix multiplication. Given two matrices:

Their product, \(C\), is computed as follows:

where \(C\) is the resulting matrix of size \(M,N\).

For each output value of matrix C, we compute the dot product between the rows of matrix A and the columns of matrix B.

Figure 2: example for the first element of C

In terms of complexity, we have \(\large O(n^3)\) computational complexity and \(\large O(n^2)\) memory accesses.
If we don’t think about architectural details, this is clearly a compute bound problem and our goal will be to be compute bound on the GPU.

Let’s say we manage to write the best implementation possible for the 7900 XTX. How fast could it run ? To answer this questions we need to look a bit at RDNA3 architecture.

RDNA3 GPUs are made of arrays of WorkGroup Processors (WGP). Every WGP are split into 2 Compute Units (CUs), themself split into 2 SIMDs. A SIMD handles the work of multiple threads organized in waves (or warps for CUDA folks) and has a set of components to do some work (like arithmetic operations). For Floating point operations, there are two 32 way VALU units.

Figure 3: simplified representation of WGPs

Figure 4: simplified representation of a single SIMD

We can compute our theoritical floating point operation per second with this formula:

Every SIMD can issue 2 Floating points intructions per cycle (one on each vALU unit). If we use FMA instructions (Fused Multiply Add), each SIMD can issue \(32*2*2=128\) floating point operations per cycle.
The 7900 XTX has 48 WGPs, that’s \(48*2*2=192\) SIMDs.

Our theoritical VRAM bandwidth is given by :

The 7900 XTX uses GDDR6 with a 384-bit bus running at 20 Gbps.

If we go back to our 4096x4096 matrix multiplication, we essentially need to do \(\large 2*4096*4096*4096\) operations.
With a 61 TFLops implementation, it would take roughly 2.23 ms to do the work and the bandwidth required to sustain this rate would be \(\large {4096*4096*4*3}/{2.23}*10^{-3} = 90.2 \text {GB/s}\).

Of course, these are oversimplified calculations as they totally ignore memory hierarchy but we see that the available bandwidth is sufficiently high so that we can increase the amount of data we read to be closer to compute bound.

Kernel 1: naive implementation

Let’s start with a naive implementation like this :

__global__ void kernel1_naive(const float *A, const float *B, float *C, int M, int K, int N, float alpha, float beta)
{
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < M && col < N)
    {
        float acc_c = 0.0f;
        for (int k = 0; k < K; ++k)
        {
            acc_c += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = alpha * acc_c + beta * C[row * N + col];
    }
}

You will notice I am doing  \(\large C={alpha}*A*B+beta*C\) instead of \(\large C=A*B\) here. This is because it makes easier to compare with libraries like rocBLAS where matrix multiplications is provided by SGEMM functions (Single-Precision General Matrix Multiply).

We launch 4096x4096 threads with a blocksize of 16x16 and each thread compute the inner dot product described before.

The performance for this kernel is 136 ms (1010.60 GFlops/s). I know, that’s pretty bad and far off our 61 TFLops target.

Kernel 0: rocBLAS reference implementation

Now that we have seen possibly the worst implementation in terms of performance, let’s look at the official rocBLAS implementation.

const int M = N;
    const int K = N;
    CHECK_ROCBLAS_STATUS(rocblas_sgemm(
        handle,
        rocblas_operation_none, // Transpose option for A
        rocblas_operation_none, // Transpose option for B
        M,                      // Number of rows in A and C
        N,                      // Number of columns in B and C
        K,                      // Number of columns in A and rows in B
        &alpha,                 // alpha
        d_a,                    // Matrix A on the device
        M,                      // Leading dimension of A
        d_b,                    // Matrix B on the device
        K,                      // Leading dimension of B
        &beta,                  // beta
        d_c,                    // Matrix C on the device
        M                       // Leading dimension of C
        ));

As discussed before, I used rocblas_sgemm function with alpha and beta set to 1.02

The performance for this kernel is 4.49 ms (30547 GFLOPs/s). This is clearly much better than our kernel 1 but still far from our theoritical 61.4 TFlops/s.

By inspecting the ISA in RGP3, I couldn’t find any dual issue instructions in the kernel (only v_fmac_f32_e32)4

Figure 5: extract of rocBLAS ISA code

This is very surprising as this essentially means one of the VALU unit is sitting there doing nothing.

Considering this, the VALU utilization of this kernel is pretty impressive and almost 100 %. However, it’s really surprising we can’t exploit these dual issue instructions properly. I’ll come to that later.

Kernel 2: LDS Tiling

The main issue with our naive kernel is that our inner loop directly accesses global memory. This is inefficient because fetching data from global memory has a high latency, typically on the order of hundreds of cycles. Since each memory read is followed by minimal computation (just one multiplication and one addition), the GPU struggles to hide this latency, even with a large number of concurrent threads. Moreover, the algorithm repeatedly reads the same rows and columns from global memory across different threads, leading to redundant memory accesses and further exacerbating the performance bottleneck.

A solution to this problem is to load the data once into faster local memory and then iterate efficiently over it with all the threads. On RDNA3, we have the Local Data Store (LDS), a high-speed, low-latency memory accessible by all threads within a workgroup.

Figure 6: simplified representation of the memory hierarchy

Since the LDS has a much smaller capacity than global memory, we need to use tiling to divide our problem into smaller sub-matrix multiplications. One way to facilitate this is to restructure the computation by moving the inner loop’s dot product to the outer loop. The key idea is to cache a column of matrix A and a row of matrix B, then perform the computation across the entire tile. This approach is more cache-efficient and significantly reduces memory access latency.

The pseudo code for our kernel 1 is :

for i from 0 to M - 1:                  # Loop over rows of A
    for j from 0 to N - 1:              # Loop over columns of B
        sum = 0
        for k from 0 to K - 1:          # Loop over columns of A / rows of B
            sum += A[i][k] * B[k][j]
        end for
        C[i][j] = sum
    end for
end for

If we move the dot product to the outer loop, we have this :

for k from 0 to K - 1:                  # Outer loop over the shared dimension
    for i from 0 to M - 1:              # Loop over rows of A
        for j from 0 to N - 1:          # Loop over columns of B
            C[i][j] += A[i][k] * B[k][j]
        end for
    end for
end for

Tiling in this form is straightforward: each workgroup operates on a tile and follows these steps: (\(BK\) is the batch size, ie number of rows/columns we load to the LDS)

Init c to 0
While kId is less than N:
  # Load A and B to Tile As and Bs
  Load BK columns A to As
  Load BK rows to Bs
  Syncthreads
  # Accumulate results using LDS
  for k from 0 to BK
    c += As[threadIdx.y][k] * Bs[k][threadIdx.x]
  Syncthreads
  Increment kId by BK
end for
c[row][col]=c

If we choose a tile size of 32x32 and \(BK=32\), our new kernel looks like this:

#define TILE_SIZE 32
__global__ void kernel2_lds(const float *A, const float *B, float *C, int N)
{
    __shared__ float As[TILE_SIZE][TILE_SIZE];
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];

    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;

    float sum = 0.0f;

    for (int t = 0; t < N; t += TILE_SIZE)
    {
        Bs[threadIdx.y][threadIdx.x] = B[N * (threadIdx.y + t) + col];
        As[threadIdx.y][threadIdx.x] = A[N * row + t + threadIdx.x];

        __syncthreads();

        for (int k = 0; k < TILE_SIZE; k++)
        {
            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        }

        __syncthreads();
    }

    if (row < N && col < N)
    {
        C[row * N + col] = sum;
    }
}

__syncthreads(); is required here to ensure that all threads in the workgroup can see the data loaded into the LDS and to synchronize before any updates are made to the data.

We also ensure that the contents of both matrices A and B are loaded into the LDS by rows rather than columns to avoid uncoalesced memory accesses.
Indeed, if we were to read by columns, each thread in a wave would access a non-contiguous memory region, resulting in multiple separate transactions and reduced efficiency as shown in the 2 diagrams below.

Figure 7: coalesced loads for matrix A. A single 128 bytes memory transaction for all threads

Figure 8: non coalesced loads for matrix A. Multiple 32 bytes memory transactions for a single wave

From the ISA guide, device memory is accessed through 32-, 64-, or 128-byte transactions, which must be naturally aligned. Maximizing memory throughput requires coalescing memory accesses across threads within a wave to minimize the number of transactions 5.

The performance for this kernel is 34.2 ms (4017 GFlops/s). 4 times faster than our naive kernel !

Let’s use RGP to understand what is going on.
Our occupancy is pretty good (100 %) but our VALU utilization is only 15%.

Figure 9 : stats taken from the instruction tab in RGP

If we look at the ISA in the instruction timing tab, we see a couple of interesting things :

Figure 10 : instruction timing

To understand what is happening, we need to quickly explain how SIMD scheduling works. During each clock cycle, the SIMD selects an instruction from a pool of wave to issue. A SIMD can manage up to 16 wavefronts in parallel. When we refer to occupancy, we are actually talking about the ratio of active wave to the theoretical maximum number of wave that a SIMD can support.
The more active wavefronts there are, the greater the likelihood that the SIMD can switch between wave, increasing the chances of hiding latency within individual wavefronts.6

If we go back to our case, we are likely having something like this :

Figure 11 : wavefront scheduling within a SIMD

Here, we have a high-occupancy kernel launching many waves in parallel, all contending for access to the LDS. Since the time taken by our VALU operations is shorter than the LDS latency, the latency cannot be hidden, even with additional threads. This results in both LDS bandwidth congestion and resource waste due to latency.

One way to address this issue is by increasing the arithmetic intensity of our kernel, ensuring that the VALU operations per wave take longer than the LDS memory reads.

Kernel 3 : Register tiling

Now, we want to increase the arithmetic complexity of our kernel. This means having each thread perform more computations. Essentially, we aim to increase the ratio of computation to data read.
One way to achieve this is to compute a small output tile per thread—for example, an 8x8 tile. To do this, we introduce an additional level of tiling.

Each thread will be responsible for producing a small tile of the output matrix. We can cache the contents of matrices \(A\) and \(B\) into registers to enable very low-latency access. However, registers are limited on the GPU, with 1536 VGPRs (Vector General-Purpose Registers) available per SIMD and a maximum of 256 registers per kernel. Increasing register usage means we won’t be able to launch as many waves per SIMD, effectively reducing occupancy. However, this shouldn’t be an issue if we can maximize utilization of the SIMD’s VALUs (Vector Arithmetic Logic Units) with just a few waves.

Now, let’s look at the different levels of tiling:

Figure 12 : tiling levels

Essentially, it means each thread will now be responsible to compute a 8x8 output tile.

Our kernel parameters looks like this:

#define BLOCK_SIZE 256
    // Block Tile size
    constexpr int BN = 128;
    constexpr int BM = 128;

    // Number of Row or column we read per batch
    constexpr int BK = 8;

    // Thread Tile size . 4x4
    constexpr int TN = 4;
    constexpr int TM = 4;

    // A wave is a block of 8x4 of the output matrix
    constexpr int nbThreadXPerWave = 8;
    constexpr int nbThreadYPerWave = 4;

    // Number of waves in a block
    constexpr int nbWavesPerBlock = BLOCK_SIZE / 32;

    constexpr int WN = 64;
    constexpr int WM = BN * BM / nbWavesPerBlock / WN;

    constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
    constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);

    // LDS Tile
    __shared__ float As[BK][BM];
    __shared__ float Bs[BK][BN];

    // Column and row from A and B, stored into registers
    float A_col[nbIterWaveM * TM];
    float B_row[nbIterWaveN * TN];

    //Wave Tile (registers)
    float C_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};

The pseudo code for our new kernel:

Initialize kId to 0
    While kId is less than N:
        # Loading Tile to LDS
        Load BK columns from A to As
        Load BK rows from B to Bs
        Syncthreads

        For k from 0 to BK - 1 do:
            Load col k of As to A_col
            Load row k of Bs to B_row

            # Wave Tile
            For idY from 0 to nbIterWaveM:
              For idX from 0 to nbIterWaveN:

                # Thread Tile
                For i from 0 to TM:
                  For j from 0 to TN:
                     x = idX * TN + j;
                     y = idY * TM + i;
                     C_regs[y][x] = A_col[y] * B_row[x]

        Syncthreads
        Increment kId by BK
    Write C_regs to C

Full kernel source code can be found here

The performance for this kernel is 6.03 ms (22777 GFlops/s), 5 times faster than our previous kernel !

We have a lower occupancy but VALU utilization has significantly increased.

Figure 13 : kernel 3 stats

The ISA looks good. We now have a lot of v_dual_fmac instructions—exactly what we wanted, even though some are still single fma.

Figure 14 : kernel 3 instruction timing

Even though this is a significant improvement over Kernel 2, we can still see that we are waiting for the LDS. This is especially true for the first batch of ds_load instructions, where we observe more than 100 clock cycles of cumulative non-hidden latency as seen below :

Figure 15 : ds_load instructions latencies

Before diving into this, we need to first improve the way we read from global memory. According to RGP, this is now the biggest bottleneck in terms of performance.

Figure 16 : gmem wait latency

Our cumulative latency for global memory waits exceeds 12 million clock cycles, which is four times more than the LDS load wait in the inner loop.

To further optimize performance, we will focus on better hiding the Global memory read latency.

Kernel 4 : GMEM double buffering

With our current implementation, every wave must wait for global memory and then LDS write latency before doing any work. In a high-occupancy scenario, this shouldn’t be an issue if the GPU can find other waves to hide this latency. However, in practice, we often have multiple waves in the same state running simultaneously because we use a sync thread before and after reading from global memory.

Figure 17 : several waves waiting for GMEM loads

One way to mitigate this is by using double buffering. We could allocate twice the memory and perform reads and writes to the LDS in parallel.

Alternatively, we could use intermediate registers to load data from global memory while working on the LDS, only writing to LDS just before it is needed. This ensures no waiting on global memory.

I prefer this approach for now, as I don’t want to introduce additional LDS pressure in the inner loop just yet.

Figure 18 : double buffering on GMEM loads

If we update our pseudo code, we now have :

Initialize kId to 0
    # Load first batch before loop
    Load BK columns from A to As
    Load BK rows from B to Bs
    Syncthreads

    While kId is less than N:
        # Loading Tile to LDS
        Load BK columns from A to A_TMP (no wait)
        Load BK rows from B to B_TMP (no wait)

        For k from 0 to BK - 1 do:
            Load col k of As to A_col
            Load row k of Bs to B_row

            # Wave Tile
            For idY from 0 to nbIterWaveM:
              For idX from 0 to nbIterWaveN:

                # Thread Tile
                For i from 0 to TM:
                  For j from 0 to TN:
                     x = idX * TN + j;
                     y = idY * TM + i;
                     C_regs[y][x] = A_col[y] * B_row[x]

        Syncthreads
        Save A_TMP and B_TMP to As and Bs
        Syncthreads
        Increment kId by BK
    Write C_regs to C

To my surprise, the performance for this kernel decreased to 14.3032 ms (9612.48 GFLOPS), more than 2 times slower than kernel 3 !

Our double buffering algorithm utilizes more registers and reduces occupancy.
After inspecting the ISA in RGP, we see that the HIP compiler attempts to keep register usage low by using scratch memory instead—which is detrimental to performance7.

Figure 19 : scratch_load instructions introduced to reduce register usage

Unfortunately, we cannot directly set the maximum number of registers per kernel in HIP (which is theoretically 256). However, we can use the launch_bounds extension to provide hints to the compiler.

With this change, the performance is back to normal :  5.37 ms (25559.6 GFLOP/s).

Full kernel source code can be found here

VALU utilization has increased from 43 % to 52 %.

Figure 20 : kernel 4 stats

We can now go back to our LDS loads in the inner loop which have become the new bottleneck, as shown below.

Figure 21 : latency on LDS loads

Kernel 5 : Optimize LDS usage

One thing I didn’t look at in previous kernels is whether or not we had bank conflicts on the LDS. This information is actually not easily accessible in RGP. If we look at ISA section where we write to the LDS, we see that the latency are unexpectly high.

Figure 22 : latencies on LDS writes

According the RDNA3 programming guide, the LDS memory is split into 64 banks of DWORD-wide RAMS. These 64 banks are further sub-divided into two sets of 32-banks each where 32 of the banks are affiliated with a pair of SIMD32’s, and the other 32 banks are affiliated with the other pair of SIMD32’s within the WGP. Each bank is a 512x32 two-port RAM (1R/1W per clock cycle). DWORDs are placed in the banks serially, but all banks can execute a store or load simultaneously.1

So, if threads within a wave access the same bank, the memory transactions will be serialized, which is exactly what happens when we write a column of matrix A to As.

Figure 23 : Matrix A bank conflicts and how to remove them

Our current kernel reads the content of matrix A rows by rows to avoid uncoalesced memory loads. Given we then operates on columns of matrix A, we transpose matrix A into matrix As so that each line of As correspond to a tile column of A.

Now, if we look at how this work is mapped to waves, we see that we essentially write 8 times to 4 consecutive banks within each wave. One way to fix this is to add a padding of 4 elements to our LDS matrix, As.

__shared__ float As[BK][BM+4]; // 4 padding to avoid bank conflicts

Doing another RGP capture with this change:

Figure 24 : updated latency with padding

LDS latency has decreased a lot and our VALU utilization is now 62.3%.

However are kernel is still bound by these LDS loads. Let’s do some napkin math and check whether or not we are not reaching the limit of the LDS bandwidth.

As said before, each pair of SIMD has a 32-banks memory capable of reading DWORD. Our theoritical bandwidth should something like this:

Now, let’s analyze what our current algorithm does:

That’s for reading. We also write to the LDS : 4096x128x32x32x4x2 = 4294967296 bytes.

With our current execution time of 5.37 ms, the required LDS bandwidth is roughly 13.56 TBytes/s.
This is less than 46% of the maximum capacity, but it is highly likely that our kernel experiences congestion in the LDS when multiple waves attempt to read or write simultaneously.

To mitigate this, we can try these 2 things :

According the RDNA3 programming guide, the LDS can operate on 2 distincts mode : WGP Mode and CU mode. HIP usse by default WGP mode.
In WGP mode, the LDS is one large contiguous memory that all waves on the WGP can access meaning we are more likely to get congestion on the LDS.
In CU mode,  the LDS is effectively split into a separate upper and lower LDS, each serving two SIMD32’s. Waves are allocated LDS space within the half of LDS which is associated with the SIMD the wave is running on.
By enabling the CU mode, we should reduce the probability of wave contending for the LDS8

Second thing we can try is to increase our Thread tile to 16x8 instead of 8x8. This will improve the computation-to-data-read ratio. It should still fit within the 256 VGPR budget we have for the kernel and reduce our bandwidth requirements to 10.3 TBytes/s

With all these changes, the performance for this kernel is now 4.09 ms (33526 GFLOP/s). That’s better than rocBLAS !

Full kernel source code can be found here

We continue increasing VALU utilization, and now our kernel has doubled in register usage (which makes sense, as we have doubled our register space requirements). Even though occupancy is low, overall performance is better because we are making better use of the VALU units.

Figure 25 : kernel 5 stats

If we look at the ISA, we now have a small LDS latency of less than 30 cycles and most of it is hidden.

Figure 26 : kernel 5 instruction timing

OK, so our kernel outperforms rocBLAS, but since we are using dual_fmac instructions, the performance is still not as high as we would expect.

At this point, I tried several optimizations, but I struggled to get the HIP compiler to generate the code I wanted. Even small changes to the C++ code drastically altered the generated ISA, making optimization work very difficult. This was especially problematic with inline assembly, where the compiler would move instructions to incorrect locations due to a lack of explicit dependencies. Additionally, there is no way to manually assign specific VGPRs to particular instructions.

Because of these challenges, I decided to optimize directly at the ISA level, which is what we will focus on in the next steps.

Looking at RGP, one thing still puzzles me in the inner loop: the HIP compiler does not use dual_fmac instructions exclusively—we always see a few single FMA instructions mixed in. Another issue is that all the v_dual_fmac instructions have a minimum latency of 2–3 cycles. While this may seem insignificant, it adds up across all instructions and impacts overall performance at our current execution speed.

Kernel 6 : VALU optimization

Before we go into the next optimizations, I need to be able to directly modify the ISA. To do so, I will now use the Module Management API so that we can load pre-compiled kernel code. Of course the idea is that we generate the ISA of our kernel from C++ once and then iterate on the ISA for any further version.

To do so, I need to extract the ISA source file from my current C++ kernel and ask hip to  build hsaco binary format:

hipcc --genco --offload-arch=gfx1100 kernel5_lds_optim.cpp -mcumode --save-temps -o tmp.hsaco

The --save-tempsparameter will allow us to have access to the intermediate .s file containing the ISA.

HIP should produce these files:

kernel5_lds_optim-hip-amdgcn-amd-amdhsa-gfx1100.bc
kernel5_lds_optim-hip-amdgcn-amd-amdhsa-gfx1100.hipi
kernel5_lds_optim-hip-amdgcn-amd-amdhsa-gfx1100.o
kernel5_lds_optim-hip-amdgcn-amd-amdhsa-gfx1100.out
kernel5_lds_optim-hip-amdgcn-amd-amdhsa-gfx1100.out.resolution.txt
kernel5_lds_optim-hip-amdgcn-amd-amdhsa-gfx1100.s

kernel5_lds_optim-hip-amdgcn-amd-amdhsa-gfx1100.s is our guy.

Now we can take this file as a basis for our modifications and assemble it with the commands :

hipcc -target amdgcn-amd-amdhsa -mcpu=gfx1100 -mcumode -c kernel_modified.s -o kernel.o
ld.lld -shared kernel.o -o kernel.hsaco

The kernel.hsacofile can then be loaded at runtime using the Module Management API of HIP.

Direct control over the ISA is great for micro-benchmarking and makes it easier to instrument the code for performance assessment without worrying about unexpected compiler optimizations.

For example, I tried duplicating our dual_fmac instructions 32 times in the inner loop to see if we could artificially become VALU-bound. However, it turns out that our VALU utilization cannot exceed 75%!

Figure 27 : Fake kernel with 32x more VALU operations

Next thing I tried is to launch a single workgroup and have a single wave running. It turns out these 2-3 clocks latency are still there meaning it must comes from the VGPR distribution of these dual_fmac instructions.

Ok, so let’s take a closer look at these dual instructions and see if we can do something about it.
Dual instructions are of that form :

OpCodeX DSTX, SRCX0, SRCX1 :: OpCodeY DSTY, SRCY0, SRCY1

In our case :

v_dual_fmac_f32 DSTX, SRCX0, SRCX1 :: v_dual_fmac_f32 DSTY, SRCY0, SRCY1

The two instructions are executed at the same time, so there are no races between them if one reads a VGPR and the other writes the same VGPR. The ‘read’ gets the old value.

There are a number of constrains in order to use these instructions. Namely:

On top of that, the RDNA 3 programming guide states :

Figure 28 : my visualization of register banks and dual instructions. Example of an instruction using bank 0,1,2,3

Figure 29 : SRCX0 and SRCY0 must use different banks

Bank number of register \(\large X\) is given by \(\large X\%4\)

Taking a example:

v_dual_fmac_f32 v10, v189, v207 :: v_dual_fmac_f32 v9, v190, v20

FMAC Bank2, Bank1, Bank3 :: FMAC Bank1, Bank2, Bank0

With this instruction, we are reading the 4 different banks in parallel and writing to bank 1 and 2 the next cycle.
In practice we could read from the same bank in both OPX and OPY if it’s not using the same operand. For example this is valid given that SRCX0 and SRCY0 use different banks :

v_dual_fmac_f32 v123, v139, v144 :: v_dual_fmac_f32 v114, v140, v143

FMAC Bank3, Bank3, Bank0 :: FMAC Bank2, Bank0, Bank3

Both instructions are reading the same banks (0 & 3). The way I see it (that’s not covered in the ISA guide AFAIK), 2 things could happen here  :

On top of that, we also need to take into account the VGPRs we are writing to even though we write on the next cycle.

So, even though we may have a valid VGPR distribution that successfully compiles, we could still encounter register bank conflicts, impacting performance.

Let’s take a look at what the HIP compiler has generated for us.

v_dual_fmac_f32 v127, v138, v144 :: v_dual_fmac_f32	v122, v139, v143
v_dual_fmac_f32 v128, v138, v145 :: v_dual_fmac_f32	v121, v139, v142
v_dual_fmac_f32 v123, v139, v144 :: v_dual_fmac_f32	v114, v140, v143
v_dual_fmac_f32 v124, v139, v145 :: v_dual_fmac_f32	v113, v140, v142
v_dual_fmac_f32 v115, v140, v144 :: v_dual_fmac_f32	v110, v141, v143
v_dual_fmac_f32 v116, v140, v145 :: v_dual_fmac_f32	v109, v141, v142
v_dual_fmac_f32 v111, v141, v144 :: v_dual_fmac_f32	v90, v138, v147
v_dual_fmac_f32 v112, v141, v145 :: v_dual_fmac_f32	v89, v138, v146
v_dual_fmac_f32 v91, v138, v148 :: v_dual_fmac_f32 v94, v139, v147
v_dual_fmac_f32 v92, v138, v149 :: v_dual_fmac_f32 v93, v139, v146
v_dual_fmac_f32 v95, v139, v148 :: v_dual_fmac_f32 v98, v140, v147
v_dual_fmac_f32 v96, v139, v149 :: v_dual_fmac_f32 v97, v140, v146
v_dual_fmac_f32 v99, v140, v148 :: v_dual_fmac_f32 v118, v141, v147
v_dual_fmac_f32 v100, v140, v149 :: v_dual_fmac_f32	v117, v141, v146
v_dual_fmac_f32 v119, v141, v148 :: v_dual_fmac_f32	v70, v138, v151
v_dual_fmac_f32 v120, v141, v149 :: v_dual_fmac_f32	v69, v138, v150
;...

If we analyse both the banks and the cache state for the first instructions, we get something like this:

I assumed that writes are performed with a 1-clock delay, which is why the first row of DSTX and DSTY is empty.

We can see that the compiler does a great job of reusing the cache, as we only read a small amount of data. However, the access pattern is not consistent over time, and we often use the same bank more than twice.

I started creating microbenchmarks to understand the latencies displayed in RGP based on different access patterns in terms of VGPR banks. However, this turned out to be quite complex, likely due to the underlying architecture’s complexity.

Instead of spending too much time on this, I tried to design an implementation following these principles:

The good news is that we can ignore alignment constraints on the output matrix C, given that the number of iterations in the inner loop is quite high. In other words, we can freely shuffle register allocations during the accumulation phase and reorder them only once before writing to memory. This effectively removes one constraint, as we no longer need to maintain a direct mapping between contiguous memory locations and contiguous registers. This might be the reason the HIP compiler was struggling to only use dual_fmac instructions (write of matrix C_reg to C by global_store_b128 requires 4 consecutive VGPRs)

Since kernel 4, our inner loop consist of doing the multiplication between the 8 elements of a column of A and the 16 elements of a row of B. We can assume both A and B is contiguously distributed on the 4 different VGPR banks. Something like this :

Figure 30 : inner loop - product between A_col and B_row

For the sake of simplicity, I will only represent the algorithm on a 8x4 tile from now on. A naive approach is to create a dual instructions by shifting a small diagonal like this making sure both SRC0s & SRC1 use different banks.

Figure 31 : naive distribution

The cell number represents the instruction index.

We see it is perfectly doable to only use dual issue instructions however some of them are using multiple times the same banks. This is something we wanted to avoid. One way to get rid of this is to store A and B on a set of non-overlapping banks. For example B only on bank 0-1 and A on bank 2-3. The issue with this is that we won’t be able to use ds_load_b128 instruction anymore as they target a 4 consecutive VGPRs. So instead of having 6 ds_load_b128 instructions like we have now, we will have 12 ds_load_b64 instead. If the performance uplift from our change is good enough, it shouldn’t matter.

Figure 32 : separate banks between A and B

All green ! However if we look at the cache use and the read pattern, we have this :

Figure 33 : Cache usage per instruction

We have good reuse of the values from A, although instruction 8 performs two reads. However, if we look at the read pattern in the detailed table below, we can see that we mostly read from bank 0 and bank 1, and from instruction Y (which is not as symmetrical as we would like it to be)

Instead of iterating over the values of A alone, we could iterate over both A and B, swapping registers between instructions X and Y to maximize cache usage

Figure 34 : Optimal solution

Looking at the detailed view, we now have a nice and symetrical access pattern. Both instruction X and Y read the same amount of data from the register file and we iterate over the 4 banks in sequential way (not just bank 0 and 1)

Right, now we are happy with our new access pattern how can we apply this change to the code ?
Here are the steps:

List the VGPRs used

Let’s start with the ds_load_b128instructions:

ds_load_b128 v[184:187], v183
ds_load_b128 v[188:191], v183 offset:64
ds_load_b128 v[192:195], v204
ds_load_b128 v[196:199], v204 offset:128
ds_load_b128 v[200:203], v204 offset:256
ds_load_b128 v[204:207], v204 offset:384

The first 2 instructions are responsible for loading A.

If we look at the fma instructions now :

v_fmac_f32_e32 v124, v184, v192
v_fmac_f32_e32 v133, v184, v193
v_dual_fmac_f32 v132, v184, v194 :: v_dual_fmac_f32 v129, v185, v192
v_dual_fmac_f32 v131, v184, v195 :: v_dual_fmac_f32 v128, v185, v193
v_dual_fmac_f32 v127, v185, v194 :: v_dual_fmac_f32 v122, v186, v193
v_dual_fmac_f32 v126, v185, v195 :: v_dual_fmac_f32 v123, v186, v192
v_dual_fmac_f32 v121, v186, v194 :: v_dual_fmac_f32 v116, v187, v193
v_dual_fmac_f32 v120, v186, v195 :: v_dual_fmac_f32 v117, v187, v192
v_dual_fmac_f32 v115, v187, v194 :: v_dual_fmac_f32 v112, v184, v197
v_dual_fmac_f32 v114, v187, v195 :: v_dual_fmac_f32 v113, v184, v196
v_dual_fmac_f32 v111, v184, v198 :: v_dual_fmac_f32 v108, v185, v197
v_dual_fmac_f32 v110, v184, v199 :: v_dual_fmac_f32 v109, v185, v196
v_dual_fmac_f32 v107, v185, v198 :: v_dual_fmac_f32 v104, v186, v197
v_dual_fmac_f32 v106, v185, v199 :: v_dual_fmac_f32 v105, v186, v196
v_dual_fmac_f32 v103, v186, v198 :: v_dual_fmac_f32 v100, v187, v197
v_dual_fmac_f32 v102, v186, v199 :: v_dual_fmac_f32 v101, v187, v196
v_dual_fmac_f32 v99, v187, v198 :: v_dual_fmac_f32 v96, v184, v201
v_dual_fmac_f32 v98, v187, v199 :: v_dual_fmac_f32 v97, v184, v200
v_dual_fmac_f32 v95, v184, v202 :: v_dual_fmac_f32 v92, v185, v201
v_dual_fmac_f32 v94, v184, v203 :: v_dual_fmac_f32 v93, v185, v200
v_dual_fmac_f32 v91, v185, v202 :: v_dual_fmac_f32 v88, v186, v201
v_dual_fmac_f32 v90, v185, v203 :: v_dual_fmac_f32 v89, v186, v200
v_dual_fmac_f32 v87, v186, v202 :: v_dual_fmac_f32 v84, v187, v201
v_dual_fmac_f32 v86, v186, v203 :: v_dual_fmac_f32 v85, v187, v200
v_dual_fmac_f32 v83, v187, v202 :: v_dual_fmac_f32 v80, v184, v205
v_dual_fmac_f32 v82, v187, v203 :: v_dual_fmac_f32 v81, v184, v204
v_dual_fmac_f32 v79, v184, v206 :: v_dual_fmac_f32 v76, v185, v205
v_dual_fmac_f32 v78, v184, v207 :: v_dual_fmac_f32 v77, v185, v204
v_dual_fmac_f32 v75, v185, v206 :: v_dual_fmac_f32 v72, v186, v205
v_dual_fmac_f32 v74, v185, v207 :: v_dual_fmac_f32 v73, v186, v204
v_dual_fmac_f32 v71, v186, v206 :: v_dual_fmac_f32 v68, v187, v205
v_dual_fmac_f32 v70, v186, v207 :: v_dual_fmac_f32 v69, v187, v204
v_dual_fmac_f32 v67, v187, v206 :: v_dual_fmac_f32 v64, v188, v193
v_dual_fmac_f32 v66, v187, v207 :: v_dual_fmac_f32 v65, v188, v192
v_dual_fmac_f32 v63, v188, v194 :: v_dual_fmac_f32 v60, v189, v193
v_dual_fmac_f32 v62, v188, v195 :: v_dual_fmac_f32 v61, v189, v192
v_dual_fmac_f32 v59, v189, v194 :: v_dual_fmac_f32 v56, v190, v193
v_dual_fmac_f32 v58, v189, v195 :: v_dual_fmac_f32 v57, v190, v192
v_dual_fmac_f32 v55, v190, v194 :: v_dual_fmac_f32 v52, v191, v193
v_dual_fmac_f32 v54, v190, v195 :: v_dual_fmac_f32 v53, v191, v192
v_dual_fmac_f32 v51, v191, v194 :: v_dual_fmac_f32 v48, v188, v197
v_dual_fmac_f32 v50, v191, v195 :: v_dual_fmac_f32 v49, v188, v196
v_dual_fmac_f32 v47, v188, v198 :: v_dual_fmac_f32 v44, v189, v197
v_dual_fmac_f32 v46, v188, v199 :: v_dual_fmac_f32 v45, v189, v196
v_dual_fmac_f32 v43, v189, v198 :: v_dual_fmac_f32 v40, v190, v197
v_dual_fmac_f32 v42, v189, v199 :: v_dual_fmac_f32 v41, v190, v196
v_dual_fmac_f32 v39, v190, v198 :: v_dual_fmac_f32 v36, v191, v197
v_dual_fmac_f32 v38, v190, v199 :: v_dual_fmac_f32 v37, v191, v196
v_dual_fmac_f32 v35, v191, v198 :: v_dual_fmac_f32 v32, v188, v201
v_dual_fmac_f32 v34, v191, v199 :: v_dual_fmac_f32 v33, v188, v200
v_dual_fmac_f32 v31, v188, v202 :: v_dual_fmac_f32 v28, v189, v201
v_dual_fmac_f32 v30, v188, v203 :: v_dual_fmac_f32 v29, v189, v200
v_dual_fmac_f32 v27, v189, v202 :: v_dual_fmac_f32 v24, v190, v201
v_dual_fmac_f32 v26, v189, v203 :: v_dual_fmac_f32 v25, v190, v200
v_dual_fmac_f32 v23, v190, v202 :: v_dual_fmac_f32 v20, v191, v201
v_dual_fmac_f32 v22, v190, v203 :: v_dual_fmac_f32 v21, v191, v200
v_dual_fmac_f32 v19, v191, v202 :: v_dual_fmac_f32 v16, v188, v205
v_dual_fmac_f32 v18, v191, v203 :: v_dual_fmac_f32 v17, v188, v204
v_dual_fmac_f32 v15, v188, v206 :: v_dual_fmac_f32 v12, v189, v205
v_dual_fmac_f32 v14, v188, v207 :: v_dual_fmac_f32 v13, v189, v204
v_dual_fmac_f32 v11, v189, v206 :: v_dual_fmac_f32 v8, v190, v205
v_dual_fmac_f32 v10, v189, v207 :: v_dual_fmac_f32 v9, v190, v204
v_dual_fmac_f32 v7, v190, v206 :: v_dual_fmac_f32 v4, v191, v205
v_dual_fmac_f32 v6, v190, v207 :: v_dual_fmac_f32 v5, v191, v204
v_fmac_f32_e32 v3, v191, v206
v_fmac_f32_e32 v2, v191, v207

Matrix C_reg is spread accross these ranges :
\(\large[2, 117], [120, 124], [126, 129], [131,133]\)

VGPR redistribution

It turns out that the VGPR allocation for C_reg is already close to what we need. We just need to add an extra bank 2 VGPR to ensure that all 128 VGPRs are allocated sequentially across banks 0-3.

This is good news, as it allows us to maintain compatibility with the initialization code for C_reg (setting all values to 0.0).

New allocation for C_reg : \([2, 117], [120, 124], [126, 129], [131,133], [214]\)

For A_col and B_row, we also need to allocate extra registers given that B_row will only use bank 0-1.

New allocation for A_col and B_row :

A_col : \([186,187],[190,191],[194,195],[198,199]\) (banks 2-3)

B_row : \([184,185] , [188,189] , [192,193] , [196,197] , [200,201], [204,205] , [208,209] , [212,213]\) (banks 0-1)

Re-write LDS loads

Our new code for loading A_col from As:

;A on bank 2-3
ds_load_b64 v[186:187], v183
ds_load_b64 v[190:191], v183 offset: 8
ds_load_b64 v[194:195], v183 offset: 64
ds_load_b64 v[198:199], v183 offset: 72

Loading B_row from Bs:

;B on bank 0-1
ds_load_b64 v[184:185], v202
ds_load_b64 v[188:189], v202 offset: 8
ds_load_b64 v[192:193], v202 offset: 128
ds_load_b64 v[196:197], v202 offset: 136
ds_load_b64 v[200:201], v202 offset: 256
ds_load_b64 v[204:205], v202 offset: 264
ds_load_b64 v[208:209], v202 offset: 384
ds_load_b64 v[212:213], v202 offset: 392

v183 and v202 are the new VGPRs holding the addresses of A and B in the LDS memory.

Re-write dual_fmas

We can then write our inner loop only as v_dual_fmac :

v_dual_fmac_f32 v5, v186, v184 :: v_dual_fmac_f32 v2, v187, v185
v_dual_fmac_f32 v3, v186, v185 :: v_dual_fmac_f32 v4, v187, v184
v_dual_fmac_f32 v9, v186, v188 :: v_dual_fmac_f32 v6, v187, v189
v_dual_fmac_f32 v7, v187, v188 :: v_dual_fmac_f32 v8, v186, v189
v_dual_fmac_f32 v13, v190, v188 :: v_dual_fmac_f32 v10, v191, v189
v_dual_fmac_f32 v11, v190, v189 :: v_dual_fmac_f32 v12, v191, v188
v_dual_fmac_f32 v17, v190, v184 :: v_dual_fmac_f32 v14, v191, v185
v_dual_fmac_f32 v15, v191, v184 :: v_dual_fmac_f32 v16, v190, v185
v_dual_fmac_f32 v21, v194, v184 :: v_dual_fmac_f32 v18, v195, v185
v_dual_fmac_f32 v19, v194, v185 :: v_dual_fmac_f32 v20, v195, v184
v_dual_fmac_f32 v25, v194, v188 :: v_dual_fmac_f32 v22, v195, v189
v_dual_fmac_f32 v23, v195, v188 :: v_dual_fmac_f32 v24, v194, v189
v_dual_fmac_f32 v29, v198, v188 :: v_dual_fmac_f32 v26, v199, v189
v_dual_fmac_f32 v27, v198, v189 :: v_dual_fmac_f32 v28, v199, v188
v_dual_fmac_f32 v33, v198, v192 :: v_dual_fmac_f32 v30, v199, v193
v_dual_fmac_f32 v31, v199, v192 :: v_dual_fmac_f32 v32, v198, v193
v_dual_fmac_f32 v37, v186, v192 :: v_dual_fmac_f32 v34, v187, v193
v_dual_fmac_f32 v35, v186, v193 :: v_dual_fmac_f32 v36, v187, v192
v_dual_fmac_f32 v41, v186, v196 :: v_dual_fmac_f32 v38, v187, v197
v_dual_fmac_f32 v39, v187, v196 :: v_dual_fmac_f32 v40, v186, v197
v_dual_fmac_f32 v45, v190, v196 :: v_dual_fmac_f32 v42, v191, v197
v_dual_fmac_f32 v43, v190, v197 :: v_dual_fmac_f32 v44, v191, v196
v_dual_fmac_f32 v49, v190, v192 :: v_dual_fmac_f32 v46, v191, v193
v_dual_fmac_f32 v47, v191, v192 :: v_dual_fmac_f32 v48, v190, v193
v_dual_fmac_f32 v53, v194, v192 :: v_dual_fmac_f32 v50, v195, v193
v_dual_fmac_f32 v51, v194, v193 :: v_dual_fmac_f32 v52, v195, v192
v_dual_fmac_f32 v57, v194, v196 :: v_dual_fmac_f32 v54, v195, v197
v_dual_fmac_f32 v55, v195, v196 :: v_dual_fmac_f32 v56, v194, v197
v_dual_fmac_f32 v61, v198, v196 :: v_dual_fmac_f32 v58, v199, v197
v_dual_fmac_f32 v59, v198, v197 :: v_dual_fmac_f32 v60, v199, v196
v_dual_fmac_f32 v65, v198, v200 :: v_dual_fmac_f32 v62, v199, v201
v_dual_fmac_f32 v63, v199, v200 :: v_dual_fmac_f32 v64, v198, v201
v_dual_fmac_f32 v69, v186, v200 :: v_dual_fmac_f32 v66, v187, v201
v_dual_fmac_f32 v67, v186, v201 :: v_dual_fmac_f32 v68, v187, v200
v_dual_fmac_f32 v73, v186, v204 :: v_dual_fmac_f32 v70, v187, v205
v_dual_fmac_f32 v71, v187, v204 :: v_dual_fmac_f32 v72, v186, v205
v_dual_fmac_f32 v77, v190, v204 :: v_dual_fmac_f32 v74, v191, v205
v_dual_fmac_f32 v75, v190, v205 :: v_dual_fmac_f32 v76, v191, v204
v_dual_fmac_f32 v81, v190, v200 :: v_dual_fmac_f32 v78, v191, v201
v_dual_fmac_f32 v79, v191, v200 :: v_dual_fmac_f32 v80, v190, v201
v_dual_fmac_f32 v85, v194, v200 :: v_dual_fmac_f32 v82, v195, v201
v_dual_fmac_f32 v83, v194, v201 :: v_dual_fmac_f32 v84, v195, v200
v_dual_fmac_f32 v89, v194, v204 :: v_dual_fmac_f32 v86, v195, v205
v_dual_fmac_f32 v87, v195, v204 :: v_dual_fmac_f32 v88, v194, v205
v_dual_fmac_f32 v93, v198, v204 :: v_dual_fmac_f32 v90, v199, v205
v_dual_fmac_f32 v91, v198, v205 :: v_dual_fmac_f32 v92, v199, v204
v_dual_fmac_f32 v97, v198, v208 :: v_dual_fmac_f32 v94, v199, v209
v_dual_fmac_f32 v95, v199, v208 :: v_dual_fmac_f32 v96, v198, v209
v_dual_fmac_f32 v101, v186, v208 :: v_dual_fmac_f32 v98, v187, v209
v_dual_fmac_f32 v99, v186, v209 :: v_dual_fmac_f32 v100, v187, v208
v_dual_fmac_f32 v105, v186, v212 :: v_dual_fmac_f32 v102, v187, v213
v_dual_fmac_f32 v103, v187, v212 :: v_dual_fmac_f32 v104, v186, v213
v_dual_fmac_f32 v109, v190, v212 :: v_dual_fmac_f32 v106, v191, v213
v_dual_fmac_f32 v107, v190, v213 :: v_dual_fmac_f32 v108, v191, v212
v_dual_fmac_f32 v113, v190, v208 :: v_dual_fmac_f32 v110, v191, v209
v_dual_fmac_f32 v111, v191, v208 :: v_dual_fmac_f32 v112, v190, v209
v_dual_fmac_f32 v117, v194, v208 :: v_dual_fmac_f32 v114, v195, v209
v_dual_fmac_f32 v115, v194, v209 :: v_dual_fmac_f32 v116, v195, v208
v_dual_fmac_f32 v121, v194, v212 :: v_dual_fmac_f32 v122, v195, v213
v_dual_fmac_f32 v123, v195, v212 :: v_dual_fmac_f32 v120, v194, v213
v_dual_fmac_f32 v129, v198, v212 :: v_dual_fmac_f32 v126, v199, v213
v_dual_fmac_f32 v127, v198, v213 :: v_dual_fmac_f32 v124, v199, v212
v_dual_fmac_f32 v133, v198, v184 :: v_dual_fmac_f32 v214, v199, v185
v_dual_fmac_f32 v131, v199, v184 :: v_dual_fmac_f32 v128, v198, v185

Restore VGPR mapping

We use a temporary VGPR to restore the full mapping like this:

; v2 -> v128 & v128 ->  v2
v_mov_b32 v200, v128
v_mov_b32 v128, v2
v_mov_b32 v2, v200
; v128 -> v56 & v56 ->  v128
v_mov_b32 v200, v56
v_mov_b32 v56, v2
v_mov_b32 v2, v200
; v56 -> v46 & v46 ->  v56
v_mov_b32 v200, v46
v_mov_b32 v46, v2
v_mov_b32 v2, v200
; v46 -> v100 & v100 ->  v46
v_mov_b32 v200, v100
v_mov_b32 v100, v2
v_mov_b32 v2, v200
 ...

To facilitate these changes, I wrote a small C++ program to parse the ISA, extract the mapping between the old and new VGPR distribution, and automatically generate all the necessary instructions.

Our kernel now uses 214 VGPRs instead of 208. We need to modify this in the .s file in the amdhsa.kernels section:

.vgpr_count:     214

Full kernel source code can be found here

Performance for this kernel is 3.63 ms (37791.2 GFLOP/s).

Our VALU utilization has gone up again to 76.2 % (more than our 75 % with 32x the inner loop).

Figure 35 : kernel 6 stats

If we look at this ISA,  our inner loop consists solely of v_dual_fmac instructions, each with a 1-cycle latency. Beautiful !

Figure 36 : kernel 7 stats

We can also see that many cycles are wasted at the end of the loop on branching. Let’s try to optimize that in the next kernel.

Kernel 7 : Loop unrolling

I previously tried unrolling the inner loop in the C++ HIP implementation, but it didn’t work out. The kernel became too large as the compiler pre-fetched more values from the LDS, and performance remained unchanged.

Now that we have a highly efficient loop and full control over the ISA, we might have better luck. For this step, I will simply duplicate the added code from Kernel 6 eight times and remove the loop mechanism.

s_cmpk_lg_i32 s14, 0x1000 ; Remove this line at the beginning of the loop
s_waitcnt lgkmcnt(0)
v_dual_fmac_f32  ...
v_dual_fmac_f32  ...
s_cbranch_scc1 .LBB0_9 ; Remove this line at the end of the loop

Duplicate 8 times our load and multiplication and make sure we increment the addresses :

v_add_nc_u32_e32 v183, 0x210, v183 ; B : 0x210 = (128+4)*4
v_add_nc_u32_e32 v202, 0x200, v202 ; A : 0x200 = (128)*4

Full kernel source code can be found here

The performance for this kernel is 3.33 ms (41255.6 GFLOPS/s).

VALU utilization is now above 80 %.

Figure 37 : kernel 7 stats

The instruction timing starts to look very good as well :

Figure 38 : instruction timing

So why aren’t we faster ?

If we look at the total latency clk in RGP, our biggest offender is the wait on the barrier. The s_waitcnt just before is the wait on the global memory loads.

Figure 39 : Remaining non hidden latencies

We can’t eliminate the barrier since we need to synchronize the threads before writing to LDS. However, if we examine the generated code for global memory loads, we notice a large code segment dedicated to it (128 lines)

Figure 40 : GMEM loads

I didn’t notice it before but even though latency is partly hidden, cumulated latency for a single load is around 1.3 Millions clks. Given that we do 16 different loads (8 for each matrix), that’s 20 millions clks latency here !

Let’s see how we can improve this in the next kernel.

Kernel 8 : Batched GMEM loads

Ok, let’s start looking at what HIP has generated for us (I have removed s_delay_alu instructions for better readability)

v_add_nc_u32_e32 v169, s4, v168
v_ashrrev_i32_e32 v170, 31, v169
v_lshlrev_b64 v[170:171], 2, v[169:170]
v_add_co_u32 v170, vcc_lo, s10, v170
v_add_co_ci_u32_e32 v171, vcc_lo, s11, v171, vcc_lo
global_load_b32 v168, v[170:171], off

v_add_nc_u32_e32 v170, s4, v169
v_ashrrev_i32_e32 v171, 31, v170
v_lshlrev_b64 v[171:172], 2, v[170:171]
v_add_co_u32 v171, vcc_lo, s10, v171
v_add_co_ci_u32_e32 v172, vcc_lo, s11, v172, vcc_lo
global_load_b32 v169, v[171:172], off

Here s[10:11] hold the address of matrix B. For each each global_load_b32, the compiler computes the read offset using VGPRs from the previous iteration (v170 and v171 here). This is not ideal for a couple of reasons :

So ideally, we would like this 128 line section of code to be just 16 lines:

global_load_b32 v169, v[171:172], off
global_load_b32 v170, v[173:174], off
global_load_b32 v171, v[175:176], off
....

However, this would require us to maintain additional VGPRs and potentially use VALU instructions to update the memory addresses as well. Given that we are already using 214 VGPRs, this is clearly not feasible.

That said, we still have a fairly good SGPR budget, and according to the RDNA3 programming guide, global_load instructions can use SGPR for base addressing.

global_load_b32 v171, v214, s[10:11]

v214 is now a offset in bytes. s[10:11] a 64bit address in memory.

So we could pre-compute once all the needed addresses for the 16 loads and just increment once the offset in the loop. That would require an additional 16*2 SGPRs and 2 VGPRs to handle the offset.

After inspecting the ISA, it looks like :

That’s all we need to compute the needed base addresses.

First, we load the first 128 bytes to get matrix A and B addresses into s[20:21] and s[22:23]:

s_load_b128 s[20:23], s[0:1], 0x0 ; Matrix A and B
s_waitcnt lgkmcnt(0)

For matrix B, we will save the  base addresses with the pre-computed offsets in s[24:39]. If we go back to our C++ code, each offset is separated by strideReadB*N = BLOCK_SIZE / BN * N, that’s 4096x4 = 0x4000 bytes

s_add_u32 s24, s22, 0x0000
s_addc_u32 s25, s23, 0
s_add_u32 s26, s22, 0x4000
s_addc_u32 s27, s23, 0
s_add_u32 s28, s22, 0x8000
s_addc_u32 s29, s23, 0
s_add_u32 s30, s22, 0xc000
s_addc_u32 s31, s23, 0
s_add_u32 s32, s22, 0x10000
s_addc_u32 s33, s23, 0
s_add_u32 s34, s22, 0x14000
s_addc_u32 s35, s23, 0
s_add_u32 s36, s22, 0x18000
s_addc_u32 s37, s23, 0
s_add_u32 s38, s22, 0x1c000
s_addc_u32 s39, s23, 0

And to compute the index in bytes, we can do :

; compute Matrix B offset
s_lshl_b32 s19, s14, 7         		; BN * blockIdx.x
v_add_nc_u32_e32 v203, s19, v0 		; index = BN * blockIdx.x + threadIdx.x
v_lshlrev_b32_e32  v203,2, v203   ; offset = 4*index (to bytes offset)

We apply the same logic for matrix A using s[40:55] for the base addresses and v215 for the offset.

s_add_u32 s40, s20, 0x0000
s_addc_u32 s41, s21, 0
s_add_u32 s42, s20, 0x40000
s_addc_u32 s43, s21, 0
s_add_u32 s44, s20, 0x80000
s_addc_u32 s45, s21, 0
s_add_u32 s46, s20, 0xc0000
s_addc_u32 s47, s21, 0
s_add_u32 s48, s20, 0x100000
s_addc_u32 s49, s21, 0
s_add_u32 s50, s20, 0x140000
s_addc_u32 s51, s21, 0
s_add_u32 s52, s20, 0x180000
s_addc_u32 s53, s21, 0
s_add_u32 s54, s20, 0x1c0000
s_addc_u32 s55, s21, 0

; compute Matrix A offset
s_lshl_b32 s19, s15, 19          ; 4096 * 128 * blockIdx.y
v_lshrrev_b32_e32 v1, 3, v0      ; threadIdx.x / 8
v_lshlrev_b32_e32 v1, 12, v1     ; 4096 * (threadIdx.x/8)
v_and_b32_e32 v215, 7, v0        ; threadIdx.x % 8
v_add_u32_e32 v215, v1, v215     ; index = 4096*(threadIdx.x/8) + threadIdx.x % 8
v_add_nc_u32_e32 v215, s19, v215 ; index += 4096*128*blockIdx.y
v_lshlrev_b32_e32  v215,2, v215  ; offset = 4*index

Now, in our main loop we can replace the 128 lines of code with this :

v_add_nc_u32_e32 v203, 0x20000, v203
v_add_nc_u32_e32 v215, 0x20, v215

global_load_b32	 v167, v203, s[24:25]
global_load_b32	 v168, v203, s[26:27]
global_load_b32	 v169, v203, s[28:29]
global_load_b32	 v170, v203, s[30:31]
global_load_b32	 v171, v203, s[32:33]
global_load_b32	 v172, v203, s[34:35]
global_load_b32	 v173, v203, s[36:37]
global_load_b32	 v174, v203, s[38:39]
global_load_b32	 v175, v215, s[40:41]
global_load_b32	 v176, v215, s[42:43]
global_load_b32	 v177, v215, s[44:45]
global_load_b32	 v178, v215, s[46:47]
global_load_b32	 v179, v215, s[48:49]
global_load_b32	 v180, v215, s[50:51]
global_load_b32	 v181, v215, s[52:53]
global_load_b32	 v182, v215, s[54:55]

Our modified kernel now uses 55 SGPRs instead of 18 and 216 VGPRs instead of 214.
If we take another RGP capture, we can see that this is much better, with less than 2 million clock cycles of latency for the entire process now

Figure 41 : Simplified GMEMs

After some experimentation, I found that spreading these 16 loads across the inner loop was more efficient.
Our kernel currently executes six wavefronts per SIMD. Since our workgroup consists of 128 threads (4 waves), every time we execute a syncthreads, at least 2 of the 6 waves on the SIMD will compete for GMEM access. Additionally, if any of the remaining 4 waves happen to be in the same state, even more waves could be contending for memory access

Figure 42 : at least 2 waves requesting GMEM at the same time

By splitting these loads into chunks of 2, we reduce the likelihood of overlap between waves, as shown in the following diagram:

Figure 43 : splitting GMEM instructions in chunks of 2

Performance for this kernel is 2.80 ms (49047 GFLOPS/s). That’s now 60% faster than our reference rocBLAS version and almost 50 times faster than our naive approach !

Full kernel source code can be found here

Figure 44 : performance results in GFLOPS/s

Conclusion

This has been an exciting journey. What started as a simple experiment to try out HIP on Windows turned into a deep dive into the hardware details of RDNA3. My biggest inspiration for this blog was Simon Boehm’s technical post9 on matrix multiplication in CUDA—an incredibly well-written piece that clearly influenced Kernel 3.

HIP tooling on Windows is quite limited. For instance, RGP does not display bank conflicts by default. However, with enough practice, it becomes possible to analyze most performance bottlenecks using the instruction timing view.

Even though the performance results are impressive—outperforming rocBLAS by 60%—this code is clearly not scalable in its current state. Furthermore, performing custom ISA optimizations makes these changes RDNA3-specific, limiting portability. As the codebase grows, modifications become increasingly difficult to implement.

That being said, the goal of this personal project was to push performance to the limit without worrying about maintainability or flexibility. While matrix multiplication can be implemented in just a few lines of code, writing an optimized implementation is incredibly challenging. We achieved a 50x speedup between the naive kernel and our best kernel, which, in my experience, would not have been possible using only HIP C++. This highlights the value of projects like OpenAI’s Triton, which I find particularly interesting and worth exploring in the future.

Although reaching almost 50 TFLOP/s is a solid achievement, we are still not fully VALU-bound, meaning there’s likely more performance left on the table. One technique I haven’t tested yet is LDS double buffering, which could eliminate one of the barriers and potentially improve the distribution of LDS instructions across the SIMD.

Finally, I want to thank Francois Guthmann for our brainstorming session on LDS optimization, which inspired the approach used in Kernel 4.

This project has been both fun and insightful, and I look forward to investigating further optimizations in the future

All the code for the 8 kernels can be found on this github here

Changelog

18/02/2024 : Updated Figures 28 and 29. Added missing SRC2 used by destination operand for FMAC_F32. Thanks to Aditya Atluri for pointing this out.

References

RDNA3 Instruction Set Architecture ↩ ↩2

I used ROCm 6.2.4 which was the latest version available on Windows 11 at the time I wrote this link ↩

Radeon Graphic Profiler is the recommended profiler on Windows ↩

I didn’t spend much time analysing how rocBLAS works but after exploring the ROCBlas repo, it looks like rocBLAS uses a project called Tensile to generate highly optimized GEMM codes for AMD GPU. ↩

ROCm performance guidelines ↩

There is a great blogpost by Francois Guthmann here that goes into these details. ↩

According to the RDNA3 ISA programming guide, the Scratch memory is similar to the global memory but instructions access a private (per-thread) memory space ↩

To enable cumode, just add -mcumode option when building with hipcc. ↩

How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog ↩

seb-v

Insights into GPU Performance, Optimization, and Beyond


Akshay Subramaniam
asubramaniam@nvidia.com
WHAT THE PROFILER IS TELLING YOU:
OPTIMIZING GPU KERNELS
2
BEFORE YOU START
1.
Know your hardware
•
What are the target machines, how many nodes? Machine-specific optimizations okay?
2.
Know your tools
•
Strengths and weaknesses of each tool? Learn how to use them (and learn one well!)
3.
Know your application
•
What does it compute? How is it parallelized? What final performance is expected?
4.
Know your process
•
Performance optimization is a constant learning process
5.
Make it so!
The five steps to enlightenment
3
THE APOD CYCLE
1. Assess
•
Identify Performance Limiter
•
Analyze Profile
•
Find Indicators
2. Parallelize
3. Optimize
3b. Build Knowledge
4. Deploy
and Test
4
Scope
GUIDING OPTIMIZATION EFFORT
• Challenge: How to know where to start?
• Top-down Approach:
•
Find Hotspot Kernel
•
Identify Performance Limiter of the Hotspot
•
Find performance bottleneck indicators related to the limiter
•
Identify associated regions in the source code
•
Come up with strategy to fix and change the code
•
Start again
“Drilling Down into the Metrics”
5
KNOW YOUR HARDWARE:
VOLTA ARCHITECTURE
6
VOLTA V100 FEATURES
Volta Architecture
Most Productive GPU
Tensor Core
120 Programmable
TFLOPS Deep Learning
Improved SIMT Model
New Algorithms
Volta MPS
Inference Utilization
Improved NVLink &
HBM2
Efficient Bandwidth
7
GPU COMPARISON
P100 (SXM2)
V100 (SXM2)
Double/Single/Half TFlop/s
5.3/10.6/21.2
7.8/15.7/125 (TensorCores)
Memory Bandwidth (GB/s)
732
900
Memory Size
16GB
16GB
L2 Cache Size
4096 KB
6144 KB
Base/Boost Clock (Mhz)
1328/1480
1312/1530
TDP (Watts)
300
300
8
VOLTA SM
GV100
GP100
FP32 Cores
64
64
INT32 Cores
64
0
FP64 Cores
32
32
Register File
256 KB
256 KB
Active Threads
2048
2048
Active Blocks
32
32
Same active threads/warps/blocks on SM
Same amount of registers
Expect similar occupancy, if not limited by shared mem.
10
Shared
Memory
64 KB
L1$
24 KB
L2$
4 MB
Load/Store Units
Pascal SM
L2$
6 MB
Load/Store Units
Volta SM
L1$ and Shared Memory
128 KB
Low Latency
Streaming
IMPROVED L1 CACHE
13
INSTRUCTION LATENCY
Dependent instruction issue latency for core FMA operations:
Volta:   4 clock cycles
Pascal: 6 clock cycles
14
TENSOR CORE
Each Tensor Core performs 64 FMA mixed-precision operations per clock
Mixed precision multiplication and accumulation
15
TENSOR CORE
cuBLAS/cuDNN: set TENSOR_OP_MATH
CUDA: nvcuda::wmma API
Example to use tensor core
#include <mma.h>
using namespace nvcuda;
__global__ void wmma_ker(half *a, half *b, float *c) {
wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::col_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;
wmma::fill_fragment(c_frag, 0.0f);
wmma::load_matrix_sync(a_frag, a, 16);
wmma::load_matrix_sync(b_frag, b, 16);
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
wmma::store_matrix_sync(c, c_frag, 16, wmma::mem_row_major);
}
16
KNOW YOUR TOOLS:
PROFILERS
18
PROFILING TOOLS
From NVIDIA
Volta, Turing, Ampere and future:
• NVIDIA Nsight Systems
• NVIDIA Nsight Compute
Older generations
• nvprof
• NVIDIA Visual Profiler (nvvp)
• Nsight Visual Studio Edition
Third Party
• TAU Performance System
• VampirTrace
• PAPI CUDA component
• HPC Toolkit
• (Tools using CUPTI)
Many Options!
Without loss of generality, in this talk we will be showing Nsight systems and
compute screenshots
20
Nsight Systems
●
Nsight Systems:
nsys profile -o profile_v4_2O ./build/bin/hpgmg-fv 7 8
System level analysis tool (think timeline)
GPU info
CPU info
21
Nsight Compute
Kernel analysis tool (think metrics)
●
Nsight Compute:
nv-nsight-cu-cli -o profile_v4_2O \
--launch-count 1 ./build/bin/hpgmg-fv 7 8
30
KNOW YOUR APPLICATION:
HPGMG
31
HPGMG
High-Performance Geometric Multi-Grid, Hybrid Implementation
Fine levels are executed on throughput-optimized processors (GPU)
Coarse levels are executed on latency-optimized processors (CPU)
5/20/2019
GPU
CPU
THRESHOLD
F-CYCLE
V-CYCLE
DIRECT SOLVE
SMOOTHER
& RESIDUAL
SMOOTHER
& RESIDUAL
SMOOTHER
SMOOTHER
RESTRICTION
INTERPOLATION
http://crd.lbl.gov/departments/computer-science/PAR/research/hpgmg/
32
MULTI-GRID BOTTLENECK
Cost of operations
5/20/2019
level
kernel time / total time
MOST TIME SPENT
ON STENCILS
level
kernel time / level time
VOLUME
SURFACE
33
MAKE IT SO:
ITERATION 1
2ND ORDER 7-POINT STENCIL
36
Identify the hotspot: smooth_kernel()
IDENTIFY HOTSPOT
Hotspot
Kernel
Time
Speedup
Original Version
2.079ms
1.00x
38
IDENTIFY PERFORMANCE LIMITER
Memory utilization
Compute utilization
39
Memory Utilization vs Compute Utilization
Four possible combinations:
PERFORMANCE LIMITER CATEGORIES
Comp
Mem
Compute
Bound
Comp
Mem
Bandwidth
Bound
Comp
Mem
Latency
Bound
Comp
Mem
Compute and
Bandwidth
Bound
42
BANDWIDTH BOUND ON V100
45
DRILLING DOWN: LATENCY ANALYSIS (V100)
The profiler
warns about
low occupancy
Limited by block size of
only 8x4=32 threads
46
OCCUPANCY
Each SM has limited resources:
•
max. 64K Registers (32 bit) distributed between threads
•
max. 48KB of shared memory per block (96KB per SMM)
•
max. 32 Active Blocks per SMM
•
Full occupancy: 2048 threads per SM (64 warps)
When a resource is used up, occupancy is reduced
GPU Utilization
(*) Values vary with Compute Capability
47
LATENCY
GPUs cover latencies by having a lot of work in flight
warp 0
warp 1
warp 2
warp 3
warp 4
warp 5
warp 6
warp 7
warp 8
warp 9
The warp issues
The warp waits (latency)
Fully covered latency
warp 0
warp 1
warp 2
warp 3
No warp issues
Exposed latency, not enough warps
48
LATENCY AT HIGH OCCUPANCY
Many active warps but with high latency instructions
Exposed latency at high occupancy
No warp issuing
warp 0
warp 1
warp 2
warp 3
warp 4
warp 5
warp 6
warp 7
warp 8
warp 9
49
GLOBAL MEMORY
Basic optimization is the same: Coalescing, Alignment, SOA pattern.
Granularity is 32 bytes, i.e. 8 threads are accessing a continuous 32 byte space.
Latency: what is the occupancy we need to saturate global load/store?
One V100:
BW = 4096 bit * 877Mhz * 2 / 8 = 898 GB/s   ~ 1.23x of P100 (theoretical)
SM ratio: 80/56 = 1.43x of P100
HBM2 increase bandwidth from 732 GB/s to 900 GB/s
52
LOOKING FOR MORE INDICATORS
12 Global Load
Transactions per 1 Request
For line numbers use:
nvcc -lineinfo
Source Code
Association
53
MEMORY TRANSACTIONS: BEST CASE
A warp issues 32x4B aligned and consecutive load/store request
Threads read different elements of the same 128B segment
1x L1 transaction: 128B needed / 128B transferred
4x L2 transactions: 128B needed / 128B transferred
1x 128B L1 transaction per warp
4x 32B L2 transactions per warp
1x 128B load/store request per warp
54
MEMORY TRANSACTIONS: WORST CASE
Threads in a warp read/write 4B words, 128B between words
Each thread reads the first 4B of a 128B segment
32x L1 transactions: 128B needed / 32x 128B transferred
32x L2 transactions: 128B needed / 32x 32B transferred
1x 128B L1 transaction per thread
1x 32B L2 transaction per thread
1x 128B load/store request per warp
Stride: 32x4B
thread 2
55
TRANSACTIONS AND REPLAYS
With replays, requests take more time and use more resources
More instructions issued
More memory traffic
Increased execution time
Inst. 0
Issued
Inst. 1
Issued
Inst. 2
Issued
Execution time
Threads
0-7/24-31
Threads
8-15
Threads
16-23
Inst. 0
Completed
Inst. 1
Completed
Inst. 2
Completed
Threads
0-7/24-31
Threads
8-15
Threads
16-23
Transfer data for inst. 0
Transfer data for inst. 1
Transfer data for inst. 2
Extra latency
Extra work (SM)
Extra memory traffic
58
FIX: BETTER GPU TILING
Before
After
Memory
Utilization Up
Transactions Per
Access Down to 9
Kernel
Time
Speedup
Original Version
2.079ms
1.00x
Better Memory Accesses
1.756ms
1.18x
+10%
Block Size Up from (8,4,1) to (32,4,1)
59
Category:
Latency Bound – Occupancy
Problem:
Latency is exposed due to low occupancy
Goal:
Hide latency behind more parallel work
Indicators:
Occupancy low (< 60%)
Execution Dependency High
Strategy:
Increase occupancy by:
• Varying block size
• Varying shared memory usage
• Varying register count (use __launch_bounds)
PERF-OPT QUICK REFERENCE CARD
60
Category:
Latency Bound – Coalescing
Problem:
Memory is accessed inefficiently => high latency
Goal:
Reduce #transactions/request to reduce latency
Indicators:
Low global load/store efficiency,
High #transactions/#request compared to ideal
Strategy:
Improve memory coalescing by:
•
Cooperative loading inside a block
•
Change block layout
•
Aligning data
•
Changing data layout to improve locality
PERF-OPT QUICK REFERENCE CARD
61
Category:
Bandwidth Bound - Coalescing
Problem:
Too much unused data clogging memory system
Goal:
Reduce traffic, move more useful data per request
Indicators:
Low global load/store efficiency,
High #transactions/#request compared to ideal
Strategy:
Improve memory coalescing by:
•
Cooperative loading inside a block
•
Change block layout
•
Aligning data
•
Changing data layout to improve locality
PERF-OPT QUICK REFERENCE CARD
62
ITERATION 2: DATA MIGRATION
68
PAGE FAULTS
Details
69
MEMORY MANAGEMENT
Using Unified Memory
No changes to data structures
No explicit data movements
Single pointer for CPU and GPU data
Use cudaMallocManaged for allocations
5/20/2
019
Developer View With
Unified Memory
Unified Memory
70
Solution: allocate the first CPU level with cudaMallocHost (zero-copy memory)
UNIFIED MEMORY
Eliminating page migrations and faults
5/20/2
019
GPU
CPU
THRESHOLD
F-CYCLE
Page faults
72
PAGE FAULTS
Almost gone
74
PAGE FAULTS
Significant speedup for affected kernel
75
MEM ADVICE API
Not used here
cudaMemPrefetchAsync(ptr, length, destDevice, stream)
    Migrate data to destDevice: overlap with compute
    Update page table: much lower overhead than page fault in kernel
    Async operation that follows CUDA stream semantics
cudaMemAdvise(ptr, length, advice, device)
    Specifies allocation and usage policy for memory region
       User can set and unset at any time
5/20/2019
76
CONCURRENCY THROUGH PIPELINING
Serial
Concurrent– overlap kernel and D2H copy
Use CUDA streams to hide data transfers
K1
K2
K3
K4
cudaMemcpyAsync(H2D)
cudaMemcpyAsync(D2H)
Kernel<<<>>>
time
cudaMemcpyAsync(H2D)
DH1
DH2
DH3
DH4
time
performance
improvement
77
ITERATION 3:
REGISTER OPTIMIZATION AND CACHING
80
LIMITER: STILL MEMORY BANDWIDTH
81
SM
Functional Units
Register File
SM
Functional Units
Register File
GPU MEMORY HIERARCHY
V100
Global Memory (Framebuffer)
L2$
Bring reused
data closer to
the SMs
•
Registers (256 KB/SM): good for
intra-thread data reuse
•
Shared mem / L1$ (128 KB/SM):
good for explicit intra-block data
reuse
•
L2$ (6144 KB): implicit data reuse
Shared Memory /
 L1$
Shared Memory /
 L1$
82
CACHING IN REGISTERS
No data loaded initially
5/20/2019
83
CACHING IN REGISTERS
Load first set of data
5/20/2019
load
84
CACHING IN REGISTERS
Perform calculation
5/20/2019
Stencil
85
CACHING IN REGISTERS
Naively load next set of data?
5/20/2019
load
86
CACHING IN REGISTERS
Reusing already loaded data is better
5/20/2019
load
keep
keep
87
CACHING IN REGISTERS
Repeat
5/20/2019
Stencil
Higher register usage may
result in reduced
occupancy => trade off
(run experiments!)
91
THE EFFECT OF REGISTER CACHING
Transactions for cached
loads reduced by a
factor of 8
Memory utilization still
high, but transferring less
redundant data
Kernel
Time
Speedup
Original Version
2.079ms
1.00x
Better Memory Accesses
1.756ms
1.18x
Register Caching
1.486ms
1.40x
92
SHARED MEMORY
• Programmer-managed cache
• Great for caching data reused across threads in a CTA
• 128KB split between shared memory and L1 cache per SM
• Each block can use at most 96KB shared memory on GV100
• Search for cudaFuncAttributePreferredSharedMemoryCarveout in the docs
__global__ void sharedMemExample(int *d) {
   __shared__ float s[64];
   int t = threadIdx.x;
   s[t] = d[t];
   __syncthreads();
   if(t>0 && t<63)
      stencil[t] = -2.0f*s[t] + s[t-1] + s[t+1];
}
global
global
registers
registers
shared
93
Category:
Bandwidth Bound – Register Caching
Problem:
Data is reused within threads and memory bw
utilization is high
Goal:
Reduce amount of data traffic to/from global mem
Indicators:
High device memory usage, latency exposed
Data reuse within threads and small-ish working set
Low arithmetic intensity of the kernel
Strategy:
•
Assign registers to cache data
•
Avoid storing and reloading data (possibly by
assigning work to threads differently)
•
Avoid register spilling
PERF-OPT QUICK REFERENCE CARD
94
Category:
Latency Bound – Texture Cache
Problem:
Load/Store Unit becomes bottleneck
Goal:
Relieve Load/Store Unit from read-only data
Indicators:
High utilization of Load/Store Unit, pipe-busy stall
reason, significant amount of read-only data
Strategy:
Load read-only data through Texture Units:
•
Annotate read-only pointers with const __restrict__
•
Use __ldg() intrinsic
PERF-OPT QUICK REFERENCE CARD
95
Category:
Device Mem Bandwidth Bound – Shared Memory
Problem:
Too much data movement
Goal:
Reduce amount of data traffic to/from global mem
Indicators:
Higher than expected memory traffic to/from global
memory
Low arithmetic intensity of the kernel
Strategy:
(Cooperatively) move data closer to SM:
•
Shared Memory
•
(or Registers)
•
(or Constant Memory)
•
(or Texture Cache)
PERF-OPT QUICK REFERENCE CARD
96
Category:
Shared Mem Bandwidth Bound – Shared Memory
Problem:
Shared memory bandwidth bottleneck
Goal:
Reduce amount of data traffic to/from global mem
Indicators:
Shared memory loads or stores saturate
Strategy:
Reduce Bank Conflicts (insert padding)
Move data from shared memory into registers
Change data layout in shared memory
PERF-OPT QUICK REFERENCE CARD
97
ITERATION 4:
KERNELS WITH INCREASED
ARITHMETIC INTENSITY
98
OPERATIONAL INTENSITY
•
Operational intensity = arithmetic operations/bytes written and read
•
Our stencil kernels have very low operational intensity
•
It might be beneficial to use a different algorithm with higher operational intensity.
•
In this case this might be achieved by using higher order stencils
99
ILP VS OCCUPANCY
•
Earlier we looked at how occupancy helps hide latency by providing independent threads of
execution.
•
When our code requires many registers the occupancy will be limited but we can still get
instruction level parallelism inside the threads.
•
Occupancy is helpful to achieving performance but not always
required
•
Some algorithms such as matrix multiplications allow
increases in operational intensity by using more registers
for local storage while simultaneously offering decent ILP.
In these cases it might be beneficial to maximize ILP and
operational intensity at the cost of occupancy.
a = b + c;
d = e + f;
a = b + c;
d = a + f;
Independent instr.
Dependent instr.
104
STALL REASONS:
EXECUTION DEPENDENCY
Memory accesses may influence execution dependencies
Global accesses create longer dependencies than shared accesses
Read-only/texture dependencies are counted in Texture
Instruction level parallelism can reduce dependencies
a = b + c; // ADD
d = a + e; // ADD
a = b[i]; // LOAD
d = a + e; // ADD
a = b + c;  // Independent ADDs
d = e + f;
105
ILP AND MEMORY ACCESSES
#pragma unroll is useful to extract ILP
Manually rewrite code if not a simple loop
float a = 0.0f;
for( int i = 0 ; i < N ; ++i )
  a += logf(b[i]);
c = b[0]
No ILP
2-way ILP (with loop unrolling)
float a, a0 = 0.0f, a1 = 0.0f;
for( int i = 0 ; i < N ; i += 2 )
{
  a0 += logf(b[i]);
  a1 += logf(b[i+1]);
}
a = a0 + a1
a += logf(c)
c = b[1]
a += logf(c)
c = b[2]
a += logf(c)
c = b[3]
a += logf(c)
c0 = b[0]
a0 += logf(c0)
c0 = b[2]
a0 += logf(c0)
c1 = b[1]
a1 += logf(c1)
c1 = b[3]
a1 += logf(c1)
a = a0 + a1
...
106
Category:
Latency Bound – Instruction Level Parallelism
Problem:
Not enough independent work per thread
Goal:
Do more parallel work inside single threads
Indicators:
High execution dependency, increasing occupancy has
no/little positive effect, still registers available
Strategy:
•
Unroll loops (#pragma unroll)
•
Refactor threads to compute n output values at the
same time (code duplication)
PERF-OPT QUICK REFERENCE CARD
107
Category:
Compute Bound – Algorithmic Changes
Problem:
GPU is computing as fast as possible
Goal:
Reduce computation if possible
Indicators:
Clearly compute bound problem, speedup only with
less computation
Strategy:
•
Pre-compute or store (intermediate) results
•
Trade memory for compute time
•
Use a computationally less expensive algorithm
•
Possibly: run with low occupancy and high ILP
PERF-OPT QUICK REFERENCE CARD
108
SUMMARY
109
SUMMARY
1.
Know your application
2.
Know your hardware
3.
Know your tools
4.
Know your process
•
Identify the Hotspot
•
Classify the Performance Limiter
•
Look for indicators
5.
Make it so!
Performance Optimization is a Constant Learning Process
112
nsys profile -o profile_v4_2O
\
./build/bin/hpgmg-fv 7 8
nv-nsight-cu-cli -o profile_v4_2O \
--kernel-regex ".*smooth_kernel*" \
--launch-count 1 ./build/bin/hpgmg-fv 7
8
113
GUIDING OPTIMIZATION EFFORT
• Challenge: How to know where to start?
• Top-down Approach:
•
Find Hotspot Kernel
•
Identify Performance Limiter of the Hotspot
•
Find performance bottleneck indicators related to the limiter
•
Identify associated regions in the source code
•
Come up with strategy to fix and change the code
•
Start again
“Drilling Down into the Metrics”
Nsight Systems
Nsight Compute
116
REFERENCES
CUDA Documentation
Best Practices: http://docs.nvidia.com/cuda/cuda-c-best-practices-guide/
Volta Tuning Guide: http://docs.nvidia.com/cuda/volta-tuning-guide/
Ampere Tuning Guide: https://docs.nvidia.com/cuda/ampere-tuning-guide/
NVIDIA Developer Blog on HPGMG
https://devblogs.nvidia.com/high-performance-geometric-multi-grid-gpu-acceleration/
Nsight Tools
https://devblogs.nvidia.com/migrating-nvidia-nsight-tools-nvvp-nvprof/
https://devblogs.nvidia.com/transitioning-nsight-systems-nvidia-visual-profiler-nvprof/
https://devblogs.nvidia.com/using-nsight-compute-to-inspect-your-kernels/


DeMoriarty

Beep Boop

Modifying Custom Matmul CUDA Kernels

I started to learn CUDA last year, and started writing matrix multiplication kernels as a learning project. After some struggles, I made them to work, but then got disappointed when I saw my kernels are 10 times slower than cuBLAS GEMM kernels. Maybe my expectations were a bit too high. I’ve tried lots of open sourced matmul kernels on github, but the best one I found was still about 5 times slower (some of them were optimized for older architectures). So I started the journey of optimizing my own matmul kernel. After few months of trial and error, my matmul kernel finally has comparable speed to cuBLAS GEMM.

Why would you need to write your own matmul kernel? the answer is, in most cases you don’t need to and also shouldn’t. It simply isn’t worth the effort when there are highly optimized cuBLAS kernels available.

There are two reasons why I did it:

In this post, I mainly want to talk about the second point: some modifications that can be done on matmul kernels and their applications.

Note: all of the kernel optimizations are targeted towards devices with compute capability 7.5 (Tesla T4, RTX 20 series), all experiments are done with a Tesla T4 GPU on Google Colab. Performance on other GPUs may not be optimal.

Note: all source code can be found in this repository.

Table of contents:

Batch Matrix Multiplication (BMM)

BMM is basically multiplying a batch of (M x K) matrices with a batch of (K x N) matrices, and get a batch of (M x N) matrices as a result. When batch size is equal to 1, it becomes a regular matrix multiplication. here is an example code in pytorch:

a = torch.randn(batch_size, M, K)
b = torch.randn(batch_size, K, N)
c = torch.bmm(a, b)
# c has shape [batch_size, M, N]

Here is the source code of my BMM kernel.

Here are some speed comparisons between my BMM kernel and torch.bmm which calls cuBLAS GEMM.

Square matrices: batch size = 1, varying M, N, K (M = N = K)

Runtime (ms):

Batch size = 128, K = 128, varying M, N (M = N)

Runtime (ms):

Fused Reduce Matmul

Sometimes we want to apply reduction right after a matrix multiplication:

c = torch.bmm(a, b)
max_c, argmax_c = torch.max(c, dim=1)
sum_c = torch.sum(c, dim=2)
argmin_c = torch.argmin(c, dim=1)

Normally there isn’t a problem with doing this, however, when I was implementing k-means clustering for TorchPQ, this became an issue. One of the major steps of k-means clustering algorithm is the computations of pairwise distance between all data points and all centroids (cluster centers), then getting the index of the closest cluster for each data point (argmin).

When the number of data points (n_data) and the number of clusters (n_clusters) are very large, this step will produce a huge (n_data x n_clusters) matrix that might not fit into the GPU memory (imagine a 1,000,000 x 10,000 fp32 matrix)

One workaround is to split this step into multiple tiny steps: in each step only compute the distance between a subset of data points and centroids (let’s say 10,000 x 10,000), then asign each data point in the subset to the closest cluster, and repeat.

A better solution could be to fuse argmin into the matmul kernel. Advantages of doing this are:

I will not go into details in this post, if you are interested, you can read implementation details, and check the source code.

I compared the new MinBMM kernel with 2 things:

values, indices = torch.min( torch.bmm(a, b), dim = dim)

and

values, indices = torch.min( custom_bmm(a, b), dim = dim)

And here are the results:

batch_size = 1, K = 16, varying M, N (M = N)

Runtime (ms)
 Peak Memory Usage (MB)

Runtime (ms)

Peak Memory Usage (MB)

batch_size = 1, K = 64, varying M, N (M = N)

Runtime (ms)
 Peak Memory Usage (MB)

Runtime (ms)

Peak Memory Usage (MB)

batch_size = 1, K = 256, varying M, N (M = N)

Runtime (ms)
 Peak Memory Usage (MB)

Runtime (ms)

Peak Memory Usage (MB)

batch_size = 64, K = 64, varying M, N (M = N)

Runtime (ms)
 Peak Memory Usage (MB)

Runtime (ms)

Peak Memory Usage (MB)

From the benchmark results we can see that fused MinBMM kernel is much faster than calling min and bmm kernels seperately, especially when k is smaller. The speedup seems to decay as K growth larger. Another advantage of MinBMM is the required memory is much smaller than other methods, the inputs can be (1,000,000 x K) matrices, and we still don’t need to be worried of going out of memory.

L2 distance

For L2 distance, all we need to change is a single line in  thread_matmul:

cCache[m].val[n] = fmaf(a, b, cCache[m].val[n]);

to

float dif = a - b;
cCache[m].val[n] = fmaf(dif, dif, cCache[m].val[n]);

Now the BMM kernel will calculate squared L2 distance instead of dot product. We don’t calculate square root since it doesn’t affect the order of nearest neighbors.

Topk Search

Nearest Neighbor Search (NNS) and Maximum Inner Product Search (MIPS) are some of the most important components of information retrieval and data science. Sometimes it’s required to search in million or billions of vectors with thousands of queries within seconds.

The naive way is to calculate pairwise distance matrix (or dot product) of dataset vectors and query vectors first, then apply topk search after. This is also called “bruteforce search”. Here is an example in pytorch:

n_data = 1_000_000
n_query = 1000
d_vector = 256
a = torch.randn(n_data, d_vector)
b = torch.randn(n_queries, d_vector)
c = 2 * a @ b.T - a.pow(2).sum(dim=-1)[:, None] - b.pow(2).sum(dim=-1)[None] # negative squared L2 distance
# c = a @ b.T # inner product

topk_val, topk_idx = torch.topk(c, dim=0, k=128)

Doing this is OK in smaller scale, but as n_data grows larger, bruteforce search will get more and more expensive. It also has the same problem as “matmul & reduce” that I talked about in previous section: the pairwise distance (or dot product) matrix may not fit into GPU RAM.

There are some approximate nearest neighbor search algorithms that try to solve the performance problem, for example Hierarchical Navigable Small World (HNSW), Product Quantization (PQ) and Locality Sensitve Hashing (LHS). I’ve implemented some variations of PQ algorithm earlier this year, you can take a look if interested.

In this post I will not be focusing on approximation methods, instead I will try to speed up the “bruteforce search” by fusing topk search into matmul kernel. Like in MinBMM kernel, all there needs to be changed is the part where we store calculated results to global memory (DRAM). In MinBMM, after matrix multiplication, instead of writing 128 x 128 block of C matrix to DRAM, we first did an in-block reduction (128 x 128) –> (128), then a global reduction using atomic operations. Similarly, in TopkBMM kernel, we will first perform an in-block bitonic sort along one of the two axis of C matrix, and a global sort after. We also need to have mutex (mutual exclusion) in order to avoid racing conditions (multiple threads trying to write to the same global memory location). The source code can be found here.

I compared it to

torch.topk(torch.bmm(a, b), k=128)

and here are the results:

A is a (n_data x d_vector) matrix, B is a (d_vector x n_query) matrix , n_data is the number of data points, n_query is the number of query vectors, d_vector is the vector dimensionality.

n_query = 1024, d_vector = 64, varying n_data:

Runtime (ms)
 Peak Memory Usage (MB)

Runtime (ms)

Peak Memory Usage (MB)

n_query = 1024, d_vector = 256, varying n_data:

Runtime (ms)
 Maximum Memory Footprint (MB)

Runtime (ms)

Maximum Memory Footprint (MB)

n_query = 1024, d_vector = 1024, varying n_data:

Runtime (ms)
 Peak Memory Usage (MB)

Runtime (ms)

Peak Memory Usage (MB)

We can see that TopkBMM kernel is extremely fast! it’s up to 10x faster than the naive “bruteforce search” when n_data is close to a million and d_vector = 64. To be honest I totally wasn’t expecting this. And this is not the end! We can further improve its speed by switching to half type (float16 or bfloat16), it’s also possible to use tensor cores of more recent nvidia GPUs to accelerate the matrix multiplication. But that’s for future, I still haven’t figured out how to use WMMA with CuPy.

One limitation is that TopkBMM can search up to 128 nearest neighbors at most. In order to search for more than 128 neighbors, we can do this iteratively. For example, first search the top 128 neighbors, then 128 ~ 256, then 256 ~ 384, and so on.

UPDATE

The TopkBMM kernel that I originally wrote was a prototype that only works under specific assumptions, such as:

I wrote a more complete version of TopkBMM later, that works for both C order and Fortan order input matrices, and supports any K value that is under 128, and supports topk on both first and second dimentions of c. However, when benchmarking the TopkBMM kernel when dim == 1, I found out that the torch.topk was performing much faster than before. It turns out that  torch.topk is faster on the contiguous dimention compared to other non-contiguous dimentions.

c = torch.randn(M, N)
topk_value, topk_index = c.topk(dim=1, k=k) # this is fast
topk_value, topk_index = c.topk(dim=0, k=k) # this is a lot slower

This is why when I benchmarked my original TopkBMM kernel, it was up to 10 times faster than torch.topk(torch.bmm(a, b)), it’s not because my kernel was that much faster, it’s because the torch.topk was not used correctly. There is a trick to speed up torch.topk when performed on a non-contiguous dimention:
\((A \cdot B)^T = B^T \cdot A^T\)

c = torch.bmm(b.transpose(-1, -2), a.transpose(-1, -2))
topk_value, topk_index = c.topk(dim=-1, k=k)

TODO: add new benchmark results.

Masked BMM

Masked BMM means “hiding” some portion of the output of BMM using a binary mask. One use case of this operation is the self-attention layer of GPT style language models. After computing dot product of queries and keys, we fill the upper triangular half of the output matrix with -∞. Here is how it’s usually done in pytorch:

# shape of keys: [batch_size, M, K]
# shape of values: [batch_size, K, N]
# shape of mask: [M, N]
dot = torch.bmm(keys, queries)
dot.masked_fill_(mask = mask, value = float("-inf") )

Mask pattern

If we could fuse masked_fill into the matmul kernel, we not only can get rid of computation of masked_fill, but the BMM kernel itself can take advantage of the sparcity given by the mask (by not computing the masked part), reducing computation even further.

The matmul kernel splits the output matrix into a grid of 128 x 128 submatrices, each submatrix is assigned to a thread block. Each thread block consists of 256 threads, and each thread computes an 8 x 8 block of the 128 x 128 submatrix.

First we need to do some preprocessing. We need 3 different levels of mask: block mask, thread mask and element mask. Block mask indicates which thread blocks should be ignored, and the same goes for thread mask and element mask.

M, N = mask.shape
assert M % 128 == 0
assert N % 128 == 0

element_mask = mask
thread_mask = mask.view( M//8, 8, N//8, 8)
thread_mask = thread_mask.sum(dim=1).sum(dim=3)
block_mask = mask.view(M//128, 128, N//128, 128)
block_mask = block_mask.sum(dim=1).sum(dim=3)

For simplicity, here we assume both M and N are divisible by 128.

We input all 3 masks to the MBMM kernel:

extern "C"
__global__ void mbmm_nn(
  const float* __restrict__ A,
  const float* __restrict__ B,
  float* __restrict__ C,
  const uint8_t* __restrict__ BlockMask,   //
  const uint8_t* __restrict__ ThreadMask,  //
  const uint8_t* __restrict__ ElementMask, //
  int M, int N, int K
){
  ...
}

At the very start of the kernel, all threads within a thread block will read the block mask value, if it is 0, all threads within that thread block will exit.  (No memory store/loads, no arithmetic ops.)

int bN = (N + 128 - 1) / 128;
uint8_t block_mask = BlockMask[(blockIdx.y)*bN + (blockIdx.x)];
if (block_mask == 0){
  return;
}

Each thread in the grid also reads a unique thread mask value:

int vx = threadIdx.x % 16;
int vy = threadIdx.x / 16;
uint8_t thread_mask = ThreadMask[(blockIdx.y*16 + vy)*tN + (blockIdx.x*16 + vx) ];

If it’s 0, then that thread skips thread_matmul.

// main loop
for (int i=0; i<nItr; i++)
{
  ...

  if (thread_mask != 0){
    thread_matmul(aSM, bSM, cCache, vx, vy);
  }
  __syncthreads();
}

However, all threads still participate in loading matrix A and B from global memory.

At the end of the kernel, before storing cached results to C, each thread will check its thread mask value once again, if it’s between 0 and 64, that thread will mask the cached result using element mask.

if (0 < thread_mask < 64){
  mask_cCache(cCache, ElementMask, gStartx, gStarty, vx, vy, bid, M, N);
}
write_c(cCache, C, gStartx, gStarty, vx, vy, bid, M, N);

This is kind of similar to the idea of Block Sparse. The difference is instead of designing a new kernel, we’re taking advantage of the grid-threadblock-thread hierarchy of BMM kernel to introduce sparcity. This works very well when there are lots of (128 x 128) masked blocks in our mask.

Now let’s see how the new MBMM kernel performs compared to the original BMM + masked_fill and torch.bmm (cuBLAS) + masked_fill:

1. batch size = 128, M = N = 1024, varying K

Runtime (ms):

2. batch size = 128, K = 64, varying M, N (M = N)

Runtime (ms):

3. M = N = 1024, K = 128, varying batch size

Runtime (ms):

We can see that, the new MBMM kernel has roughly 2 times the speed of the original BMM kernel + masked_fill, and 1.2 ~ 1.4 times the speed of torch.bmm + masked_fill.

Selective BMM

to be continued…