Swizzling explained properly

"Please let this be a normal matrix layout..."
"With the Swiz? No way!"
"Awwww..."

Swizzling is an essential CUDA programming technique, but weeks of browsing turned up no explanation of it that satisfied me. Let's change that.

Why swizzle?

GPUs, like CPUs, have a memory hierarchy consisting of one large pool of main memory, and progressively smaller pools of progressively faster memory. In general, Nvidia GPUs contain:

not including non-programmer-managed caches. (Recent GPU generations introduce extra kinds of memory, such as Hopper's distributed shared memory (DSMEM) and Blackwell's tensor memory (TMEM). They won't be too important for this article.)

Swizzling concerns the memory region called shared memory or SMEM. SMEM is organized into 32 consecutive banks, each bank serving one word of 4 consecutive bytes. Each bank can only serve one word per cycle, and so accesses to the same bank are serialized. This means we should organize our data in SMEM such that no load or store accesses two words from the same bank—in other words, we must avoid bank conflicts. Swizzling is one way to achieve this.

What does swizzling do?

We may visualize the SMEM bank layout as follows:

In the above diagram, each 4-byte (1-word) bank is one column, and each consecutive group of 32 banks (= 128 bytes) is one row.

On Nvidia GPUs (and all other GPUs that matter), exactly 32 consecutive threads execute the same instruction in parallel on each SM. A group of 32 consecutive threads is called a warp. [2] If two threads in a warp access the same SMEM bank at once, one of them has to wait until the other is done. This is called a bank conflict, and it decreases SMEM throughput linearly: if 4 threads access the same bank, the last thread has to wait 4 cycles; if 32 threads access the same bank, we wait 32 cycles. Obviously, this is bad and we want to avoid it. (Conversely, if all 32 threads access different banks, all 32 accesses complete in the same 1 cycle.)

Suppose we're storing a 32-column matrix of 4-byte floats. If we wanted to access one row of 32 floats, we'd hit 32 different banks, and our 128 bytes would be served in a single cycle. However, if we were to access one column, every access would hit the same bank, and we'd have to wait for as many cycles as there are rows in our matrix.

The theoretical maximum SMEM bandwidth is 4 bytes/bank/cycle × 32 banks = 128B/cycle. A 32-way bank conflict like above would limit us to a measly 4B per cycle. But if we paid for this 128B per cycle, we should use it all! A number of solutions exist:

1. Use a matrix whose width isn't a multiple of 32. If we have a 33-column matrix instead, then $M_{00}$ goes in bank 0, $M_{10}$ goes in bank 1, etc. However, this introduces a bunch of complications.

• Cache lines are 128B wide; SMEM comes in multiples of 1024B; etc. Making your matrix a weird size doesn't play nice with hardware.
• Several instructions expect data to be aligned to a multiple of 16 or 32 or some other neat amount of bytes. For instance, PTX's matrix load instruction ldmatrix requires 16-byte alignment on every matrix row, failing which behavior is undefined. This means for example that a 33-column float32 matrix would require 12 extra bytes of padding per row, which wastes a lot of memory!

How about another method that plays nice with hardware and consumes no extra SMEM?

2. Permute matrix elements so every row and column gets spread out over all 32 banks.

Yeah, that's what this whole article is about.

Stated more mathematically, we have a 2D grid of size $M$ by $N$. We need to find some bijection $f$ from $\{0, ..., M-1\} \times \{0, ..., N-1\}$ to itself, such that:

• For all pairs of row indices $(r, r')$, $f(r, c)$ and $f(r', c)$ lie in different columns, and
• For all pairs of column indices $(c, c')$, $f(r, c)$ and $f(r, c')$ lie in different columns.

Several functions satisfy this property. A particularly elegant choice is:

$$f(r, c) = (r, c \oplus r)$$

where $\oplus$ means bitwise XOR. This is swizzling.

Why does that work?

Firstly, assume that $M$ and $N$ are powers of two, and $M \le N$. Let's prove all the previously claimed properties of $f$:

1. Elements originally in same row are always sent to different columns

Let's pick two distinct elements from row $r$, and call them $(r,c)$ and $(r,c')$ where $c \neq c'$.
$(r, c)$ gets sent to $(r, c \oplus r)$, and $(r, c')$ gets sent to $(r, c' \oplus r)$.
We want to show that $c \oplus r \neq c' \oplus r$.

Recall: $a \oplus b = 0$ if and only if $a = b$.

So if $(c \oplus r) \oplus (c' \oplus r) \neq 0$, we are done.

Recall: XOR is self-inverse, associative and commutative.

So $(c \oplus r) \oplus (c' \oplus r) = c \oplus c'$. Since $c \neq c'$, we know $c \oplus c' \neq 0$. Done!

2. Elements originally in same column are always sent to different columns

Similarly, pick two distinct elements from column $c$, and call them $(r, c)$ and $(r', c)$ where $r \neq r'$. These get sent to $(r, c \oplus r)$ and $(r', c \oplus r')$ respectively. Again, we want to show that $c \oplus r \neq c \oplus r'$.

Let's XOR the new column indices together like before:

$(c \oplus r) \oplus (c \oplus r') = r \oplus r'$. Since $r \neq r'$, we are done.

3. $f$ is a bijection

Since row indices don't change, and transformed column indices are pairwise distinct by (2), it suffices to show that $c \oplus r \le N-1$.

Recall that $M$ and $N$ are both powers of 2, and $M \le N$. Let $M = 2^a$ and $N = 2^b$. So $r$ and $c$ may be represented by at most $a$ and $b$ bits respectively, where $a \le b$.

We know $0 \le r \le M-1$ and $0 \le c \le N-1$. Suppose $c \oplus r \gt N-1$. This means that $c \oplus r$ occupies more than $b$ bits.

But $c$ occupies at most $b$ bits. This means that XORing with $r$ must have flipped a bit in $c$ past the $b$-th least significant position. But $r$ occupies at most $a$ bits, and $a \le b$. So $r$ doesn't contain enough bits to flip anything past the $b$-th least significant bit of $c$. Contradiction.

CUTLASS swizzling notation

Swizzling is so useful that it's been built into CUTLASS, Nvidia's open-source library for high-performance matrix multiplication. [3] At first glance, their implementation might look a bit inscrutable: [4]

class Swizzle:
    def __init__(self, bits, base, shift):
        assert bits >= 0
        assert base >= 0
        assert abs(shift) >= bits
        self.bits = bits
        self.base = base
        self.shift = shift
        bit_msk = (1 << bits) - 1
        self.yyy_msk = bit_msk << (base + max(0, shift))
        self.zzz_msk = bit_msk << (base - min(0, shift))

    def __call__(self, offset):
        return offset ^ shiftr(offset & self.yyy_msk, self.shift)
    
    # ...

Hmm, lots of bitwise arithmetic. I see the XOR, but what's with all the shifting? Maybe the accompanying comment will shed some light on the situation:

0bxxxxxxxxxxxxxxxYYYxxxxxxxZZZxxxx
                              ^--^  Base is the number of least-sig bits to keep constant
                 ^-^       ^-^      Bits is the number of bits in the mask
                   ^---------^      Shift is the distance to shift the YYY mask
                                      (pos shifts YYY to the right, neg shifts YYY to the left)

e.g. Given
0bxxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxx
the result is
0bxxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxx where AA = ZZ xor YY

This is starting to make a bit of sense. Notice the following correspondence:

• Our swizzling definition replaces column index $c$ by $c \oplus r$.
• Nvidia's swizzling definition replaces bits ZZ by ZZ $\oplus$ YY.

So YY is simply the bit representation of our row index $r$, and ZZ is the bit representation of our column index $c$.

This means that the bit sequence YYZZ represents the offset of $(r,c)$ in our grid as traced in raster order:

In the above diagram, two bits are used for the row index YY, and three bits for the column index ZZZ. (Remember, our number of rows and columns must both be powers of 2.)

If we want to XOR bits YY with bits ZZZ, then we need to right-shift YY by 3 bits such that it overlaps with the end of ZZZ. Per the Python implementation, we set shift to 3. And since bits is the length of the XOR mask, we set it to the length of YY, which is 2.

In general:

• bits is the number of bits used for the row index Y, and
• shift is the number of bits used for the column index Z.

Of course, you could generate other swizzling patterns by deviating from this, but then the unique-column property would not be guaranteed. For instance, a swizzle pattern made for a 2x8 grid would repeat columns four times if used on an 8x8 grid. (We see examples of this later in the article.)

What about base?

Base is the number of least-sig bits to keep constant.

Interesting. Why do we need this?

Some clues are given by the matrix load/store instructions that operate on SMEM. For instance, the ldmatrix instruction receives the starting address of a matrix row, where each row contains either:

• 8 elements of 16 bits each, or
• 16 elements of 8 bits each

and is therefore exactly 16 bytes wide. Further, documentation states that:

Consecutive instances of row need not be stored contiguously in memory.

So each row of 16 bytes needs to be contiguous in SMEM, but we can place the rows wherever we want, as long as the within-row layout doesn't change.

In other words: the 4 least significant bits of every element's address must be kept constant. Setting base = 4 does exactly this.

Another way of thinking about base is that it sets the atomicity of the swizzle operation:

• Swizzle(2, 0, 3) works on a grid of 2² = 4 rows and 2³ = 8 columns. Each grid element represents 2⁰ = 1 byte.
• Swizzle(2, 4, 3) is the same, except each grid element now represents 2⁴ = 16 bytes.

Now we have a complete definition of the swizzle functor Swizzle(bits, base, shift):

• bits = no. of bits for row index
• shift = no. of bits for column index
• base = log₂ of swizzle atomicity.

32B, 64B and 128B swizzles

Now we can understand how to construct the canonical 32B, 64B and 128B swizzle modes, and why they're named that way.

Swizzle 32B

Swizzle 32B is constructed by Swizzle(1, 4, 3), and it looks like this:

Recall that bits=1 means we replace the LSB of the column index by XORing it with the LSB of the row index. For even rows (0, 2, 4, 6), the row index ends in a 0 bit, so the column index doesn't change ($x \oplus 0 = x$). But for odd rows (1, 3, 5, 7), the row index ends in a 1 bit, and so the last bit of every column index is flipped. This means that in odd rows, every column $2n$ gets swapped with column $2n+1$. So row 1, which used to read $[8, 9, 10, 11, 12, 13, 14, 15]$, now reads $[9, 8, 11, 10, 13, 12, 15, 14]$.

In the above diagram, each cell is a 16-byte row consumed by ldmatrix. Each column in the non-swizzled layout is a 8x8 sub-matrix of two-byte elements; I've color-coded the columns so we can tell where each sub-matrix ends up after swizzling. If we follow the colors, we find that Swizzle 32B repeats itself after two 16B rows. Is that why it's called Swizzle 32B? Or is it because any two 16B cells of the same color can be read/written without bank conflicts? Who knows.

Anyway, a single call to ldmatrix.m8n8.b16 will consume all elements of one single color in the original grid. Without swizzling, we'd have to load eight 16B rows from the same column, which means an 8-way bank conflict, and therefore an 8-cycle wait. Swizzle 32B slightly improves this to a 4-cycle wait by spreading every color across two columns. But we can do better.

Swizzle 64B

Similar XOR math to the last section should yield that the above pattern is none other than Swizzle(2, 4, 3). (Replace the 2 LSBs of the column index by XORing them with the 2 LSBs of the row index.) Swizzle 64B repeats itself every 4 rows of 16B.

If we follow the colors again, we find that one ldmatrix.m8n8.b16 call will load from four distinct columns (= 16 SMEM banks), reducing our bank-conflict degree to 2. But as you probably guessed, we can do yet better.

Swizzle 128B

This is Swizzle 128B, constructed by Swizzle(3, 4, 3), and of course it repeats itself after every 8 rows of 16B. (You'll just have to trust me on this one; I didn't have the patience to draw and label and color-code another 8x8 grid.) One call to ldmatrix.m8n8.b16 now reads from eight distinct columns, hitting all 32 SMEM banks and eliminating all bank conflicts. Perfect.

Sample swizzled matmul

That's all well and good, but as the CUTLASS programming guidelines very wisely say:

"Performance requires it" implies measurement.

So let's do some measurement. Below, I benchmark various matrix multiplication CUDA kernels, with some basic optimizations like vectorized GMEM accesses, block-tiling and tensor cores so they aren't stupidly slow. All kernels are totally identical apart from their swizzle mode.

We compute $C = AB$ where $A \in \mathbb{R}^{m \times k}$, $B \in \mathbb{R}^{k \times n}$, $C \in \mathbb{R}^{m \times n}$, and $m=n=k=8192$. All benchmarks were run on an A100 80GB PCIe.

#ifndef TURING_TC_NO_PIPELINE_GENERIC_CUH
#define TURING_TC_NO_PIPELINE_GENERIC_CUH

#include "../../global.cuh"
#include "../../swizzle.cuh"
#include <cuda.h>
#include <cuda_fp16.h>

namespace turing_tc_no_pipeline_generic {

#define elems_per_copy(T) ((16) / (sizeof(T)))

enum class CTATileLayout {
  Default,
  Padded,
  Swizzle_32B,
  Swizzle_64B,
  Swizzle_128B,
};

template <CTATileLayout Layout> constexpr bool is_swizzled = false;
template <> constexpr bool is_swizzled<CTATileLayout::Swizzle_32B> = true;
template <> constexpr bool is_swizzled<CTATileLayout::Swizzle_64B> = true;
template <> constexpr bool is_swizzled<CTATileLayout::Swizzle_128B> = true;

template <CTATileLayout> struct get_swizzle;
template <> struct get_swizzle<CTATileLayout::Swizzle_32B> {
  static constexpr Swizzle func = Swizzle<1, 4, 3>{};
};
template <> struct get_swizzle<CTATileLayout::Swizzle_64B> {
  static constexpr Swizzle func = Swizzle<2, 4, 3>{};
};
template <> struct get_swizzle<CTATileLayout::Swizzle_128B> {
  static constexpr Swizzle func = Swizzle<3, 4, 3>{};
};

// WARNING: Padded output will not work because output staging tile size exceeds
// Turing max SMEM.

const int BLOCK_M = 128, BLOCK_N = 128, BLOCK_K = 64;
const int WARP_M = 32, WARP_N = 64;
const int MMA_M = 16, MMA_N = 8, MMA_K = 8;
const int SKEW_HALFS = 8; // only used for padded layout

const int nwarps_m = CDIV(BLOCK_M, WARP_M);
const int nwarps_n = CDIV(BLOCK_N, WARP_N);
const int mma_k_steps = CDIV(BLOCK_K, MMA_K);
const int mma_per_wt_m = CDIV(WARP_M, MMA_M);
const int mma_per_wt_n = CDIV(WARP_N, MMA_N);
const int n_threads = nwarps_m * nwarps_n * WARPSIZE;

const int epc_half = elems_per_copy(__half);
const int epc_float = elems_per_copy(float);
static_assert((BLOCK_M * BLOCK_K) % epc_half == 0,
              "A CTA block size not divisible by number of elements per "
              "vectorized copy");
static_assert((BLOCK_N * BLOCK_K) % epc_half == 0,
              "B CTA block size not divisible by number of elements per "
              "vectorized copy");

struct alignas(8) Half4 {
  __half2 x01, x23;
};
static_assert(sizeof(Half4) == 8, "half4 should have size of 8 bytes");
static_assert(alignof(Half4) == 8, "half4 should be aligned to 8 bytes");

using FragmentA = uint32_t[2];
using FragmentB = uint32_t[1];
using FragmentC = float[4];

__device__ __forceinline__ uint32_t generic_to_shared(const void *const p) {
  return static_cast<uint32_t>(__cvta_generic_to_shared(p));
}

template <typename T>
__device__ __forceinline__ std::ptrdiff_t elems_to_bytes(T *const arr,
                                                         const size_t offset) {
  return offset * sizeof(T);
}

// NOTE: overloaded for const T* and non-const T*
// FIXME: there has to be a better way...
__device__ __forceinline__ void *
offset_ptr_by_bytes(void *const base, const std::ptrdiff_t offset) {
  return static_cast<char *>(base) + offset;
}
__device__ __forceinline__ const void *
offset_ptr_by_bytes(const void *const base, const std::ptrdiff_t offset) {
  return static_cast<const char *>(base) + offset;
}

template <CTATileLayout InputLayout> struct get_input_ld {
  static constexpr int value = []() {
    if constexpr (InputLayout == CTATileLayout::Padded) {
      return BLOCK_K + SKEW_HALFS;
    } else {
      return BLOCK_K;
    }
  }();
};

template <CTATileLayout OutputLayout> struct get_output_ld {
  static constexpr int value = []() {
    if constexpr (OutputLayout == CTATileLayout::Padded) {
      return BLOCK_N + SKEW_HALFS;
    } else {
      return BLOCK_N;
    }
  }();
};

template <CTATileLayout InputLayout>
__device__ __forceinline__ void
load_input_tile_iter(const __half *const in_global, __half *const in_smem,
                     const int i, const int K) {
  const int cta_row = i / BLOCK_K;
  const int cta_col = i % BLOCK_K;

  if constexpr (is_swizzled<InputLayout>) {
    static __device__ auto swizzle = get_swizzle<InputLayout>::func;
    const std::ptrdiff_t swizzled_offset_bytes =
        swizzle(elems_to_bytes(in_smem, i));
    *reinterpret_cast<uint4 *>(
        offset_ptr_by_bytes(in_smem, swizzled_offset_bytes)) =
        *reinterpret_cast<const uint4 *>(&in_global[cta_row * K + cta_col]);
  } else {
    *reinterpret_cast<uint4 *>(
        &in_smem[cta_row * get_input_ld<InputLayout>::value + cta_col]) =
        *reinterpret_cast<const uint4 *>(&in_global[cta_row * K + cta_col]);
  }
}

template <CTATileLayout Layout, typename T>
__device__ __forceinline__ T *ptr_to_element(T *const arr, int offset_elems) {
  if constexpr (is_swizzled<Layout>) {
    static __device__ auto swizzle = get_swizzle<Layout>::func;
    const std::ptrdiff_t swizzled_offset_bytes =
        swizzle(elems_to_bytes(arr, offset_elems));
    return static_cast<T *>(offset_ptr_by_bytes(arr, swizzled_offset_bytes));
  } else {
    return &arr[offset_elems];
  }
}

template <CTATileLayout OutputLayout>
__device__ __forceinline__ void
store_output_tile_iter(const float *const out_smem, __half *const out_global,
                       const int i, const int N) {
  const int cta_row = i / BLOCK_N;
  const int cta_col = i % BLOCK_N;

  float4 floats = *reinterpret_cast<const float4 *>(
      ptr_to_element<OutputLayout>(out_smem, i));
  Half4 halfs{.x01 = __floats2half2_rn(floats.x, floats.y),
              .x23 = __floats2half2_rn(floats.z, floats.w)};

  *reinterpret_cast<uint2 *>(&out_global[cta_row * N + cta_col]) =
      *reinterpret_cast<uint2 *>(&halfs);
}

template <CTATileLayout ALayout, CTATileLayout BLayout,
          CTATileLayout OutputLayout>
__global__ void __launch_bounds__(n_threads, 1)
    matmul_kernel(const int M, const int N, const int K,
                  const __half *__restrict A, const __half *__restrict B,
                  __half *__restrict C) {
  extern __shared__ char smem[];
  const int A_smem_bytes =
      BLOCK_M * get_input_ld<ALayout>::value * sizeof(__half);

  const int block_size = blockDim.x * blockDim.y;
  const int warp_idx = threadIdx.x / WARPSIZE;
  const int lane_idx = threadIdx.x % WARPSIZE;
  const int warp_row = warp_idx / nwarps_n;
  const int warp_col = warp_idx % nwarps_n;

  const int a_ldmatrix_row = lane_idx & 0b1111;
  const int b_ldmatrix_row = lane_idx & 0b0111;

  const int groupID = lane_idx >> 2;
  const int threadID_in_group = lane_idx % 4;
  const int row_c01 = groupID;
  const int row_c23 = groupID + 8;
  const int col_c02 = threadID_in_group * 2;

  A += blockIdx.x * BLOCK_M * K;
  B += blockIdx.y * BLOCK_N * K;
  C += blockIdx.x * BLOCK_M * N + blockIdx.y * BLOCK_N;

  // Zero out accumulator
  FragmentC frag_acc[mma_per_wt_m][mma_per_wt_n];
  for (int i = 0; i < mma_per_wt_m; i++) {
    for (int j = 0; j < mma_per_wt_n; j++) {
      for (int k = 0; k < 4; k++) {
        frag_acc[i][j][k] = 0.0f;
      }
    }
  }

  __half *const A_cta = reinterpret_cast<__half *>(smem);
  __half *const B_cta = reinterpret_cast<__half *>(smem + A_smem_bytes);

  for (int mainloop_k = 0; mainloop_k < K; mainloop_k += BLOCK_K) {
    // Load CTA tiles
    for (int i = epc_half * threadIdx.x; i < BLOCK_M * BLOCK_K;
         i += epc_half * block_size) {
      load_input_tile_iter<ALayout>(A, A_cta, i, K);
    }
    for (int i = epc_half * threadIdx.x; i < BLOCK_N * BLOCK_K;
         i += epc_half * block_size) {
      load_input_tile_iter<BLayout>(B, B_cta, i, K);
    }
    __syncthreads(); // make sure CTA tiles loaded

    A += BLOCK_K;
    B += BLOCK_K;

    FragmentA frag_a[mma_per_wt_m];
    FragmentB frag_b[mma_per_wt_n];

    for (int wt_k = 0; wt_k < mma_k_steps; wt_k++) {
      for (int wt_m = 0; wt_m < mma_per_wt_m; wt_m++) {
        uint32_t *a_rmem = reinterpret_cast<uint32_t *>(&frag_a[wt_m]);

        const std::ptrdiff_t A_row_offset =
            (wt_m * MMA_M + warp_row * WARP_M + a_ldmatrix_row) *
                get_input_ld<ALayout>::value +
            (wt_k * MMA_K);

        asm volatile(
            "ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n"
            : "=r"(a_rmem[0]), "=r"(a_rmem[1])
            : "r"(
                generic_to_shared(ptr_to_element<ALayout>(A_cta, A_row_offset)))
            : "memory");
      }
      for (int wt_n = 0; wt_n < mma_per_wt_n; wt_n++) {
        uint32_t *b_rmem = reinterpret_cast<uint32_t *>(&frag_b[wt_n]);

        const std::ptrdiff_t B_row_offset =
            (wt_n * MMA_N + warp_col * WARP_N + b_ldmatrix_row) *
                get_input_ld<BLayout>::value +
            (wt_k * MMA_K);

        asm volatile("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n"
                     : "=r"(b_rmem[0])
                     : "r"(generic_to_shared(
                         ptr_to_element<BLayout>(B_cta, B_row_offset)))
                     : "memory");
      }

      for (int wt_m = 0; wt_m < mma_per_wt_m; wt_m++) {
        for (int wt_n = 0; wt_n < mma_per_wt_n; wt_n++) {
          const uint32_t *a = reinterpret_cast<const uint32_t *>(&frag_a[wt_m]);
          const uint32_t *b = reinterpret_cast<const uint32_t *>(&frag_b[wt_n]);
          float *c = reinterpret_cast<float *>(&frag_acc[wt_m][wt_n]);
          asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32"
                       "{%0, %1, %2, %3},"
                       "{%4, %5},"
                       "{%6},"
                       "{%0, %1, %2, %3};\n"
                       : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
                       : "r"(a[0]), "r"(a[1]), "r"(b[0]));
        }
      }
    }
    __syncthreads();
  } // mainloop end

  float *const out_cta = reinterpret_cast<float *>(smem);
  for (int wt_m = 0; wt_m < mma_per_wt_m; wt_m++) {
    for (int wt_n = 0; wt_n < mma_per_wt_n; wt_n++) {
      const uint32_t *c =
          reinterpret_cast<const uint32_t *>(frag_acc[wt_m][wt_n]);

      const std::ptrdiff_t offset01_elems =
          (warp_row * WARP_M + wt_m * MMA_M + row_c01) *
              get_output_ld<OutputLayout>::value +
          (warp_col * WARP_N + wt_n * MMA_N + col_c02);
      const std::ptrdiff_t offset23_elems =
          (warp_row * WARP_M + wt_m * MMA_M + row_c23) *
              get_output_ld<OutputLayout>::value +
          (warp_col * WARP_N + wt_n * MMA_N + col_c02);

      asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" ::"r"(generic_to_shared(
                       ptr_to_element<OutputLayout>(out_cta, offset01_elems))),
                   "r"(c[0]), "r"(c[1])
                   : "memory");
      asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" ::"r"(generic_to_shared(
                       ptr_to_element<OutputLayout>(out_cta, offset23_elems))),
                   "r"(c[2]), "r"(c[3])
                   : "memory");
    }
  }
  __syncthreads();

  for (int i = epc_float * threadIdx.x; i < BLOCK_M * BLOCK_N;
       i += epc_float * block_size) {
    store_output_tile_iter<OutputLayout>(out_cta, C, i, N);
  }
}

template <CTATileLayout ALayout, CTATileLayout BLayout,
          CTATileLayout OutputLayout>
void matmul(const Problem &problem, MatmulData &matmul_data) {
  auto kernel = matmul_kernel<ALayout, BLayout, OutputLayout>;

  const int A_smem_bytes =
      BLOCK_M * get_input_ld<ALayout>::value * sizeof(__half);
  const int B_smem_bytes =
      BLOCK_N * get_input_ld<BLayout>::value * sizeof(__half);
  const int out_smem_bytes =
      BLOCK_M * get_output_ld<OutputLayout>::value * sizeof(float);

  const size_t smem_bytes =
      std::max(A_smem_bytes + B_smem_bytes, out_smem_bytes);
  static_assert(smem_bytes <= (64 << 10),
                "Required SMEM bytes exceeds Turing max");

  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
                       smem_bytes);

  const dim3 n_blocks(CDIV(problem.M, BLOCK_M), CDIV(problem.N, BLOCK_N));
  kernel<<<n_blocks, n_threads, smem_bytes>>>(problem.M, problem.N, problem.K,
                                              matmul_data.A, matmul_data.B,
                                              matmul_data.C);
}

} // namespace turing_tc_no_pipeline_generic

#endif // TURING_TC_NO_PIPELINE_GENERIC_CUH

Results

Input swizzle

If we swizzle only the input tiles for $A$ and $B$, we obtain the following result:

As expected, wider swizzling modes give us more FLOP/s.

More interestingly, the gains diminish as we move from 32B to 64B to 128B. If bank-conflict degree halves with every doubling of the swizzle width, then our max possible SMEM bandwidth correspondingly doubles. But our FLOP/s don't show this same exponential gain, because as we progressively eliminate bank conflicts, other operations begin to dominate computation time. Nsight Compute profiling supports this:

With no-swizzle and 32B, memory I/O instruction queue waiting (MIO Throttle) accounts for the vast majority of stall cycles. But at 64B and above, it becomes nearly negligible, with GMEM latency (Long Scoreboard) and math instructions (Math Pipe Throttle) dominating instead. This means that when we eliminate SMEM bank conflicts, we spend a larger percentage of our time doing math and accessing GMEM.

Output swizzle

What happens if we swizzle the output $C$ instead?

Wow, that's kinda sad. What happened here (or not)?

Clearly, eliminating bank conflicts for our $C$ tile does practically nothing. In contrast, eliminating them for $A$ and $B$ almost triples our FLOP/s. The difference boils down to two reasons:

1. SMEM loads of $A$ and $B$ account for far more SMEM traffic than stores of $C$.

If we load more than we store, then slow loads will hurt more than slow stores. Nsight Compute can tell us exactly how much we read and write to SMEM:

Loads total 51.81GB; stores total 17.45GB. That's a 3x difference.

Why do we load almost exactly 3x as many bytes as we store?

Great question. Our SMEM accesses are as follows:

1. Store $A$ and $B$ tiles from GMEM -> SMEM
2. Load $A$ and $B$ tiles from SMEM -> RMEM for matmul
3. Store $C_{FP32}$ tile from RMEM -> SMEM after matmul
4. Load $C_{FP32}$ tile from SMEM -> RMEM for downcast from FP32 to FP16. (After this, FP16 tile is written from RMEM to GMEM.)

Naively, it might seem that all of $A$, $B$ and $C$ are stored to and loaded from SMEM exactly once each, and so our load and store byte counts should be equal. As you may have suspected, this is not true. The below diagram should roughly explain it:

In short: $A$ and $B$ tiles are loaded multiple times from SMEM to RMEM. The exact number of times each tile is loaded depends on the size of our warp tile relative to our SMEM tile.

Since we have 2 warp tiles per SMEM tile row and 4 warp tiles per SMEM tile column, $A$ tiles are loaded twice while $B$ tiles are loaded 4x. But our $C$ output tile is only stored once. So we load input elements 6 times for each output element.

But remember: $A$ and $B$ are in FP16, while the $C$ SMEM tile is in FP32. So we load 3 input bytes per output byte.

But this 3x difference doesn't explain the total absence of any effect of output swizzling.

The key is to realize where our SMEM stores are coming from. Almost all of them are incurred by stores of $A$ and $B$ from GMEM->SMEM, rather than stores of $C$ from RMEM->SMEM. (We only compute each output tile once, while tiles of $A$ and $B$ must be reloaded across different thread blocks.) Shared stores total 17.45GB, but writes of $C_{FP32}$ only account for (8192×8192) elements × 4 bytes per element = 268MB, which is... 1.5% of that. They might as well be a rounding error.

Why do we need to reload the same A and B tiles multiple times from GMEM to SMEM?

Another great question. This design is precipitated by the following constraints:

• There's no way to fit the whole of $A$ and $B$ into SMEM at once. Each of them is 134 MB, but modern GPUs only have a few hundred KB of SMEM at most. [5]
• Each thread block only computes one output tile.
• SMEM is only shared within a single thread block.
• We need to use the same tiles of $A$ and $B$ for multiple output tiles of $C$.

Given these constraints, we have no choice but to reload the same $A$ and $B$ tiles across multiple thread blocks. As for the number of reloads, a similar diagram applies as earlier:

So for our example where our $C_{FP32}$ SMEM tile is 128×128, then $A$ is copied 8192÷128=64 times from GMEM to SMEM; same for $B$. Therefore:

Total bytes loaded from GMEM->SMEM
= 64 times of combined A and B size in GMEM
= 64 × (8192 × 8192) elements × 2 bytes per element × 2 (A and B)
= 17.18 GB

which very closely matches our Nsight Compute counter.

The effect of duplicate SMEM->RMEM loads stacks multiplicatively on top of the duplicate GMEM->SMEM stores: each SMEM tile is stored from GMEM->SMEM multiple times, then those same SMEM tiles are read from SMEM->RMEM multiple times. Hence the number of SMEM bytes loaded is 3x of the SMEM bytes stored, which in turn is 64x of the size of A and B in GMEM.

There is another reason why $A$ and $B$ conflicts hurt more:

2. Our $C$ store conflicts are of lower degree than our $A$ and $B$ loads.

Nsight Compute tells us:

The memory access pattern for shared stores might not be optimal and causes on average a 4.5-way bank conflict across all 34603008 shared store requests.

So for a kernel with non-swizzled input and output, our SMEM loads are 8-way conflicted, while our SMEM stores are only 4(ish)-way conflicted.

This initially surprised me. The PTX reference documents the $C$ fragment layout for mma.sync.m16n8k8 as follows:

And our RMEM->SMEM store of $C_{FP32}$ looks like:

asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" ::"r"(generic_to_shared(
                 ptr_to_element<OutputLayout>(out_cta, offset01_elems))),
             "r"(c[0]), "r"(c[1])
             : "memory");

where the second 64-bit write has been omitted for brevity (it behaves the same as the first anyway).

Each thread writes two 4-byte floats in a single instruction. Each row of the $C_{FP32}$ tile is 128 elements × 4 bytes per element = 512 bytes apart. Since this distance is a multiple of 128, all elements in the same column of the $C_{FP32}$ tile map to the same SMEM bank. And since every 4th thread writes to the same columns (i.e. banks), our bank-conflict degree should be 32÷4=8 just like for our SMEM loads. Right?

Wrong. This explanation fails to account for a crucial hardware nuance: SMEM accesses are split into 128-byte transactions. Each warp writes 8 bytes per thread × 32 threads = 256 bytes. These 256 bytes are split into two transactions of 128 bytes each, and so only half the warp's shared writes are being concurrently executed. This correspondingly halves the average bank-conflict degree for shared stores.

Conclusion / appendix / post-article word vomit

Writing this article sent me down an insane rabbit hole, and I hope you enjoyed diving into it with me. I estimate that only about half of the noteworthy things I encountered while writing actually made it into the article. Honorable mentions:

__launch_bounds__ for avoiding unintentional register reuse

Initially, I compared all swizzled kernels to a padded kernel, the latter of which solves bank conflicts by padding input and output rows in SMEM. Initially, the padded kernel was unexpectedly slower than the swizzled kernel, despite doing less math and also avoiding all bank conflicts. By wading through SASS with the help of GPT 5.6 Sol, it was eventually determined that ptxas decided to reuse registers during the GMEM->SMEM load for the padded kernel but not the swizzled kernel. This was ostensibly some compiler heuristic trying to reduce register pressure, but what it actually ended up doing was to introduce a crap ton of extra read-after-write dependencies that totally prevented ILP between GMEM loads. Additionally:

• The padded kernel was way under the register budget anyway.
• The kernel's occupancy was limited by SMEM usage, not register pressure. Reducing the registers per thread would not have given me any extra CTAs per SM.
• Long Scoreboard stalls (GMEM latency) remain the dominant bottleneck of the kernel (refer to the stall cycle reasons graph). To additionally slow down GMEM loading would be terrible for performance. Indeed, the bug gave me a 25% slowdown on a Titan RTX.

ptxas, what are you doing?! Anyways, explicitly telling the compiler that I only intended to launch one CTA per SM (by setting attribute __launch_bounds__(n_threads, 1)) fixed the issue, allowing the kernel to use as many damn registers as it wanted. This was both a skill issue on my part (not knowing about __launch_bounds__) and a badly tuned compiler heuristic that tried to save registers where completely unnecessary.

sm75 uarch bug(?)

This was a pain. On all Turing cards I tested, 128B swizzle was unexpectedly slower than 64B. Somehow, Long Scoreboard stall cycles increased more than MIO and Short Scoreboard stalls fell! This can't be explained by "GMEM latency is no longer hidden by SMEM latency" because it doesn't account for why total stall cycles increased—that explanation is only plausible if stall cycles remained the same.

This bug disappeared on sm80 and all later generations. (I tested using a Tesla T4, Titan RTX, RTX 3090, RTX 4090, RTX 4060 mobile, RTX 3060 LHR, A100 80GB PCIe, and H100 96GB PCIe.) I suspect a Turing-specific microarchitectural bug: SMEM and GMEM accesses share a lot of the same hardware machinery, so maybe those shared hardware units become more contended on Turing than on later generations after bank conflicts disappear? I don't know. If any Nvidia engineers wanna explain an eight-year-old uarch that has long reached EOL, feel free.

Acknowledgments

Aleksa Gordic's article on GPU matrix multiplication, which got me into the whole GPU shtick.

Yang Yifan's article on swizzling for tensor core MMA ops has excellent diagrams of the 32B, 64B and 128B swizzling modes.

Simon Veitner's article on CuTe's swizzling implementation explained some of CuTe swizzling notation, which, combined with Yang Yifan's post, gave me the eureka moment of why CuTe defines its swizzle function like that.

Footnotes

[1] ^ A streaming multiprocessor (SM) is akin to a CPU core, or an 'execution unit' on other GPUs.

[2] ^ Technically, it's a set of 32 threads whose indices only differ in their 5 least significant bits. So warp 0 would comprise threads 0–31; warp 1 would comprise threads 32–63, etc.

[3] ^ Yes, I know this is technically from CuTe, which is the tensor-manipulation part of CUTLASS.

[4] ^ The real C++ implementation is located here. I decided against including it in the main post because, like most C++ code, it looks like ass:

template <int BBits, int MBase, int SShift = BBits>
struct Swizzle
{
  static constexpr int num_bits = BBits;
  static constexpr int num_base = MBase;
  static constexpr int num_shft = SShift;

  static_assert(num_base >= 0,             "MBase must be positive.");
  static_assert(num_bits >= 0,             "BBits must be positive.");
  static_assert(abs(num_shft) >= num_bits, "abs(SShift) must be more than BBits.");

  using bit_msk = cute::constant<int, (1 << num_bits) - 1>;
  using yyy_msk = cute::constant<int, bit_msk{} << (num_base + max(0,num_shft))>;
  using zzz_msk = cute::constant<int, bit_msk{} << (num_base - min(0,num_shft))>;
  using msk_sft = cute::constant<int, num_shft>;

  static constexpr uint32_t swizzle_code = uint32_t(yyy_msk::value | zzz_msk::value);

  template <class Offset>
  CUTE_HOST_DEVICE constexpr static
  auto
  apply(Offset const& offset)
  {
    return offset ^ shiftr(offset & yyy_msk{}, msk_sft{});   // ZZZ ^= YYY
  }

  // ...
};

[5] ^ "Actually, Cerebras—" Shut up.