CUDA Thread Block Swizzle

Introduction

The execution order of thread blocks in a CUDA kernel affects what data is resident in the L2 cache at any given moment, which in turn influences memory access efficiency and overall kernel performance, especially when memory access becomes the bottleneck of the kernel. The execution order can be controlled using thread block swizzle algorithms, which determine the mapping from program IDs to tile locations. Different thread block swizzle algorithms can have different data access and cache utilization patterns, resulting in CUDA kernel performance variations.

In this blog post, I would like to quickly introduce several common thread block swizzle algorithms, simulate their L2 cache efficiency using a simple model, benchmark the performances of GEMM kernels using different thread block swizzle algorithms, and analyze the results.

CUDA Thread Block Swizzle

Unlike the brain-twisting shared memory swizzle, thread block swizzle is relatively straightforward to understand. We could like to use a $8 \times 8$ GEMM tiled grid for the output matrix for illustration purposes. The implementations of different thread block swizzle algorithms will be presented in the benchmark script later.

Linear Swizzle

The swizzled index of a $8 \times 8$ grid in a row-major order would be:

$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 01 & 02 & 03 & 04 & 05 & 06 & 07 \\ \hline
08 & 09 & 10 & 11 & 12 & 13 & 14 & 15 \\ \hline
16 & 17 & 18 & 19 & 20 & 21 & 22 & 23 \\ \hline
24 & 25 & 26 & 27 & 28 & 29 & 30 & 31 \\ \hline
32 & 33 & 34 & 35 & 36 & 37 & 38 & 39 \\ \hline
40 & 41 & 42 & 43 & 44 & 45 & 46 & 47 \\ \hline
48 & 49 & 50 & 51 & 52 & 53 & 54 & 55 \\ \hline
56 & 57 & 58 & 59 & 60 & 61 & 62 & 63 \\ \hline
\end{array}
$$

Grouped 2D Panel Swizzle

The swizzled index of a $8 \times 8$ grid using a grouped 2D panel swizzle with a group size of 4 would be:

$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 04 & 08 & 12 & 16 & 20 & 24 & 28 \\ \hline
01 & 05 & 09 & 13 & 17 & 21 & 25 & 29 \\ \hline
02 & 06 & 10 & 14 & 18 & 22 & 26 & 30 \\ \hline
03 & 07 & 11 & 15 & 19 & 23 & 27 & 31 \\ \hline
32 & 36 & 40 & 44 & 48 & 52 & 56 & 60 \\ \hline
33 & 37 & 41 & 45 & 49 & 53 & 57 & 61 \\ \hline
34 & 38 & 42 & 46 & 50 & 54 & 58 & 62 \\ \hline
35 & 39 & 43 & 47 & 51 & 55 & 59 & 63 \\ \hline
\end{array}
$$

This is the most common thread block swizzle algorithm used for GEMM in practice, apart from the linear swizzle, which is usually default.

Morton / Z-order Swizzle

The swizzled index of a $8 \times 8$ grid using a Morton / Z-order swizzle would be:

$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 02 & 08 & 10 & 32 & 34 & 40 & 42 \\ \hline
01 & 03 & 09 & 11 & 33 & 35 & 41 & 43 \\ \hline
04 & 06 & 12 & 14 & 36 & 38 & 44 & 46 \\ \hline
05 & 07 & 13 & 15 & 37 & 39 & 45 & 47 \\ \hline
16 & 18 & 24 & 26 & 48 & 50 & 56 & 58 \\ \hline
17 & 19 & 25 & 27 & 49 & 51 & 57 & 59 \\ \hline
20 & 22 & 28 & 30 & 52 & 54 & 60 & 62 \\ \hline
21 & 23 & 29 & 31 & 53 & 55 & 61 & 63 \\ \hline
\end{array}
$$

Note that Morton / Z-order swizzle only has a bijective mapping for grid sizes that are powers of two.

Bitwise XOR Swizzle

The swizzled index of a $8 \times 8$ grid using a bitwise XOR swizzle would be:

$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 09 & 18 & 27 & 36 & 45 & 54 & 63 \\ \hline
01 & 08 & 19 & 26 & 37 & 44 & 55 & 62 \\ \hline
02 & 11 & 16 & 25 & 38 & 47 & 52 & 61 \\ \hline
03 & 10 & 17 & 24 & 39 & 46 & 53 & 60 \\ \hline
04 & 13 & 22 & 31 & 32 & 41 & 50 & 59 \\ \hline
05 & 12 & 23 & 30 & 33 & 40 & 51 & 58 \\ \hline
06 & 15 & 20 & 29 & 34 & 43 & 48 & 57 \\ \hline
07 & 14 & 21 & 28 & 35 & 42 & 49 & 56 \\ \hline
\end{array}
$$

Note that bitwise XOR swizzle only has a bijective mapping for grid sizes that are powers of two.

CUDA Thread Block Swizzle Benchmark

The L2 cache hit rate of different thread block swizzle algorithms can be simulated using a simple offline model. The performances of each swizzle algorithm applied on GEMM kernels can also be benchmarked.

triton_gemm_swizzle_bench.py

"""Triton GEMM thread-block swizzle benchmark.

Benchmarks a single Triton FP16/FP32 GEMM kernel across several
program-id -> output-tile ("swizzle") orderings and compares each one's
real, measured throughput against an offline L2 cache hit-rate prediction.

Background
----------
A GEMM kernel launches one thread block ("program") per output tile. The
order in which program ids are mapped to (M, N) tile coordinates does not
affect correctness, but it determines which tiles are resident in the L2
cache (and which DRAM rows are open) at any given moment during execution.
Reordering that mapping to favor cache reuse is a standard GEMM
optimization (see Triton's own matmul tutorial). This script benchmarks
five such swizzle families using self-contained CPU pid-mapping functions
and a matching offline L2 hit-rate model, so the predicted ranking can be
checked against real GPU throughput for the same problem:
  - Linear (Row-Major / Col-Major)
  - Grouped 2D Panel (configurable group size)
  - Morton / Z-order
  - Bitwise XOR

Persistent vs. non-persistent launch (--persistent)
----------------------------------------------------
The kernel always loops internally over a fixed, statically-strided subset
of tiles (`tile_id = pid, pid + NUM_PROGRAMS, pid + 2*NUM_PROGRAMS, ...`),
so the two launch styles are really just two choices for NUM_PROGRAMS:
  - Default (non-persistent): NUM_PROGRAMS == total_tiles, i.e. one program
    per output tile, one loop iteration per program. Real concurrency is
    whatever the GPU's block scheduler decides at runtime; the offline L2
    predictor has to *assume* a concurrency figure (measured occupancy).
  - `--persistent`: NUM_PROGRAMS == min(total_tiles, measured occupancy),
    i.e. exactly as many programs as can be concurrently resident, each
    looping over several tiles. The concurrently-active tile set at loop
    iteration i is then *exactly* `{c + i*NUM_PROGRAMS : c in
    [0, NUM_PROGRAMS)}` -- no more scheduler-order guesswork, no
    launch/ramp-up overhead from dispatching new blocks, no partial "tail
    wave" at the end.
Empirically, this determinism does not meaningfully change the measured
row-major-vs-col-major gap or the predicted/measured correlation: the
predictor's blind spot is that it never models DRAM row-buffer locality or
bank-level contention, and that gap is present regardless of how precisely
the concurrent tile set is known.

Methodology
-----------
1. `check_correctness` validates every swizzle mode against torch.matmul on
   a small, fixed grid before any benchmarking runs.
2. `measure_concurrent_capacity` compiles the kernel once (without
   executing it) to derive real occupancy (`num_sms * blocks_per_sm`) from
   its register and shared-memory footprint, rather than assuming a value.
3. `predict_l2_hit_rate` runs an offline, byte-granular LRU cache
   simulation to estimate each swizzle's L2 hit rate for the given problem
   shape and concurrency (`min(total_tiles, measured_capacity)`).
4. Each swizzle mode is timed with triton.testing.do_bench, which clears
   the L2 cache before every timed repetition, so measurements always
   start cold and are not contaminated by a previous algorithm's residual
   state.
5. A Spearman rank correlation between predicted hit rate and measured
   TFLOPS is reported as a sanity check for the offline model.

Known limitations
-----------------
- Morton and XOR are only bijective (visit every output tile exactly once)
  when the grid (ceil(M/BLOCK_M) x ceil(N/BLOCK_N)) is square and a power
  of two; otherwise they are automatically skipped with a warning (see
  `build_algorithm_list`).
- The L2 hit-rate model is a coarse, whole-tile LRU trace. It does not
  model DRAM row-buffer locality, bank-level queuing/contention, or memory
  coalescing, so a higher predicted hit rate does not always translate
  into higher measured throughput -- in particular, for a compute-bound
  configuration, throughput can be flat across a wide range of predicted
  hit rates.
- With `--persistent`, tile-to-program assignment is static (round-robin
  by tile_id), not dynamic/atomic-counter based, so there is no runtime
  load balancing across programs for edge-tile masking imbalance.

Example
-------
  python triton_gemm_swizzle_bench.py --M 8192 --N 8192 --K 8192 \\
      --block-m 128 --block-n 128 --block-k 32 \\
      --group-sizes 1 2 4 8 16 32 64 --dtype fp16 [--persistent]
"""

import argparse
from collections import OrderedDict
import math
import subprocess
import sys
from typing import Any, Callable, Dict, List, Tuple

import torch
import triton
import triton.language as tl
from triton.runtime import driver

# --- Swizzle mode ids shared between the Triton kernel and the CPU predictor ---
MODE_ROW_MAJOR = 0
MODE_COL_MAJOR = 1
MODE_PANEL = 2
MODE_MORTON = 3
MODE_XOR = 4

TORCH_DTYPES = {"fp16": torch.float16, "fp32": torch.float32}
TL_DTYPES = {"fp16": tl.float16, "fp32": tl.float32}
DTYPE_BYTES = {"fp16": 2, "fp32": 4}


# --- Hardware auto-detection ---


def detect_gpu_hardware() -> Dict[str, Any]:
    """Queries the local system for SM count and L2 cache size.

    Inspects PyTorch `_C.cuda.DeviceProperties` attributes with priority on
    `L2_cache_size`, falling back to `nvidia-smi` for the SM count/GPU name
    and to fixed defaults if neither source is available.
    """
    fallback_sms = 108
    fallback_l2_mb = 50.0  # Standard usable CUDA global memory L2 limit

    warnings = []
    gpu_name = "Unknown / CPU Fallback"
    detected_sms = None
    detected_l2_mb = None

    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        gpu_name = props.name
        detected_sms = getattr(props, "multi_processor_count", None)

        for attr in [
                "L2_cache_size",
                "l2_cache_size",
                "l2CacheSize",
                "l2_cache_size_bytes",
        ]:
            if hasattr(props, attr):
                l2_bytes = getattr(props, attr)
                if l2_bytes > 0:
                    detected_l2_mb = l2_bytes / (1024 * 1024)
                    break

    if detected_sms is None or gpu_name == "Unknown / CPU Fallback":
        try:
            cmd = [
                "nvidia-smi",
                "--query-gpu=name,multiprocessor_count",
                "--format=csv,noheader,nounits",
            ]
            output = subprocess.check_output(cmd).decode(
                "utf-8").strip().split(",")
            if len(output) >= 2:
                gpu_name = output[0].strip()
                if detected_sms is None:
                    detected_sms = int(output[1].strip())
        except Exception:
            pass

    if detected_sms is None:
        detected_sms = fallback_sms
        warnings.append(
            f"Unable to query SM count from platform. Using default fallback:"
            f" {fallback_sms} SMs.")

    if detected_l2_mb is None:
        detected_l2_mb = fallback_l2_mb
        if gpu_name == "Unknown / CPU Fallback":
            warnings.append(
                "Unable to query GPU platform or L2 cache size. Using default global"
                f" limit: {fallback_l2_mb:.2f} MB.")

    return {
        "num_sms": detected_sms,
        "l2_cache_size_mb": detected_l2_mb,
        "gpu_name": gpu_name,
        "warnings": warnings,
    }


# --- CPU reference pid-mapping functions (mirrored by get_swizzled_pid below) ---


def linear_row_major(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]:
    return pid // grid_n, pid % grid_n


def linear_col_major(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]:
    return pid % grid_m, pid // grid_m


def panel_swizzle(pid: int,
                   grid_m: int,
                   grid_n: int,
                   group_size_m: int = 8) -> Tuple[int, int]:
    num_pid_in_group = group_size_m * grid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * group_size_m
    group_size_m_actual = min(grid_m - first_pid_m, group_size_m)
    by = (pid % num_pid_in_group) // group_size_m_actual
    bx = first_pid_m + (pid % group_size_m_actual)
    return bx, by


def morton_swizzle(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]:
    bx, by = 0, 0
    for i in range(16):
        bx |= ((pid >> (2 * i)) & 1) << i
        by |= ((pid >> (2 * i + 1)) & 1) << i
    return bx % grid_m, by % grid_n


def xor_swizzle(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]:
    linear_x = pid % grid_n
    linear_y = pid // grid_n
    bx = (linear_x ^ linear_y) % grid_m
    by = linear_y % grid_n
    return bx, by


@triton.jit
def get_swizzled_pid(pid, num_pid_m, num_pid_n, SWIZZLE_MODE: tl.constexpr,
                      GROUP_SIZE_M: tl.constexpr):
    """Maps a linear tile id to (pid_m, pid_n) output-tile coordinates.

    SWIZZLE_MODE selects the ordering (0=row-major, 1=col-major, 2=grouped
    2D panel, 3=Morton/Z-order, 4=bitwise XOR); GROUP_SIZE_M is only used by
    the grouped-panel mode. Mirrors the CPU reference functions above
    (linear_row_major, linear_col_major, panel_swizzle, morton_swizzle,
    xor_swizzle) so the offline L2 predictor and the actual kernel agree on
    tile visitation order.
    """
    if SWIZZLE_MODE == 0:  # Linear row-major
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n
    elif SWIZZLE_MODE == 1:  # Linear col-major
        pid_m = pid % num_pid_m
        pid_n = pid // num_pid_m
    elif SWIZZLE_MODE == 2:  # Grouped 2D panel
        num_pid_in_group = GROUP_SIZE_M * num_pid_n
        group_id = pid // num_pid_in_group
        first_pid_m = group_id * GROUP_SIZE_M
        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
        pid_n = (pid % num_pid_in_group) // group_size_m
        pid_m = first_pid_m + (pid % group_size_m)
    elif SWIZZLE_MODE == 3:  # Morton / Z-order
        pid_m = 0
        pid_n = 0
        for i in range(16):
            pid_m |= ((pid >> (2 * i)) & 1) << i
            pid_n |= ((pid >> (2 * i + 1)) & 1) << i
        pid_m = pid_m % num_pid_m
        pid_n = pid_n % num_pid_n
    else:  # Bitwise XOR
        linear_x = pid % num_pid_n
        linear_y = pid // num_pid_n
        pid_m = (linear_x ^ linear_y) % num_pid_m
        pid_n = linear_y % num_pid_n
    return pid_m, pid_n


@triton.jit
def gemm_swizzle_kernel(
    a_ptr,
    b_ptr,
    c_ptr,
    M,
    N,
    K,
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    SWIZZLE_MODE: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    DTYPE: tl.constexpr,
    NUM_PROGRAMS: tl.constexpr,
):
    """Masked, tiled GEMM (C = A @ B) with a swappable program-id swizzle.

    Exactly NUM_PROGRAMS programs are launched; each one statically owns
    tile ids `pid, pid + NUM_PROGRAMS, pid + 2*NUM_PROGRAMS, ...`. When
    NUM_PROGRAMS == total_tiles this degenerates to one tile per program
    (non-persistent, one loop iteration each); when NUM_PROGRAMS is capped
    at the GPU's measured concurrent capacity, each program loops over
    several tiles (persistent). Accumulates in fp32 regardless of input
    dtype and casts down to DTYPE (shared by A, B, and C) on store;
    correctness is identical across SWIZZLE_MODE values, only tile
    visitation order (and therefore cache/memory behavior) changes.
    """
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_tiles = num_pid_m * num_pid_n

    for tile_id in range(pid, num_tiles, NUM_PROGRAMS):
        pid_m, pid_n = get_swizzled_pid(tile_id, num_pid_m, num_pid_n,
                                         SWIZZLE_MODE, GROUP_SIZE_M)

        offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        offs_k = tl.arange(0, BLOCK_K)

        a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[
            None, :] * stride_ak
        b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[
            None, :] * stride_bn

        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
        for k in range(0, tl.cdiv(K, BLOCK_K)):
            k_remaining = K - k * BLOCK_K
            a_mask = (offs_am[:, None] < M) & (offs_k[None, :] < k_remaining)
            b_mask = (offs_k[:, None] < k_remaining) & (offs_bn[None, :] < N)
            a = tl.load(a_ptrs, mask=a_mask, other=0.0)
            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
            acc = tl.dot(a, b, acc)
            a_ptrs += BLOCK_K * stride_ak
            b_ptrs += BLOCK_K * stride_bk

        c = acc.to(DTYPE)
        offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        c_ptrs = c_ptr + offs_cm[:, None] * stride_cm + offs_cn[
            None, :] * stride_cn
        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
        tl.store(c_ptrs, c, mask=c_mask)


def measure_concurrent_capacity(device: torch.device,
                                 block_m: int,
                                 block_n: int,
                                 block_k: int,
                                 dtype: str = "fp16",
                                 num_warps: int = 4,
                                 num_stages: int = 3) -> int:
    """Compiles the kernel once (never executes it) and derives resident
    blocks/SM from its actual register and shared-memory footprint (same
    technique as the Triton softmax tutorial), instead of guessing/accepting
    an occupancy value on the CLI. Only compiles, so tiny dummy tensors are
    enough regardless of the real M/N/K of the benchmark. Returns the number
    of concurrently-resident CTAs across the whole GPU
    (`num_sms * blocks_per_sm`).
    """
    properties = driver.active.utils.get_device_properties(device.index)
    num_sms = properties["multiprocessor_count"]
    num_regs = properties["max_num_regs"]
    size_smem = properties["max_shared_mem"]
    warp_size = properties["warpSize"]
    # Not exposed by triton's get_device_properties on all versions; torch has it reliably.
    max_threads_per_sm = torch.cuda.get_device_properties(
        device.index).max_threads_per_multi_processor

    torch_dtype = TORCH_DTYPES[dtype]
    a = torch.empty((block_m, block_k), device=device, dtype=torch_dtype)
    b = torch.empty((block_k, block_n), device=device, dtype=torch_dtype)
    c = torch.empty((block_m, block_n), device=device, dtype=torch_dtype)

    kernel = gemm_swizzle_kernel.warmup(
        a,
        b,
        c,
        block_m,
        block_n,
        block_k,
        a.stride(0),
        a.stride(1),
        b.stride(0),
        b.stride(1),
        c.stride(0),
        c.stride(1),
        BLOCK_M=block_m,
        BLOCK_N=block_n,
        BLOCK_K=block_k,
        SWIZZLE_MODE=MODE_ROW_MAJOR,
        GROUP_SIZE_M=1,
        DTYPE=TL_DTYPES[dtype],
        NUM_PROGRAMS=1,
        num_warps=num_warps,
        num_stages=num_stages,
        grid=(1, ),
    )
    kernel._init_handles()
    n_regs = kernel.n_regs
    kernel_smem = kernel.metadata.shared

    reg_occupancy = num_regs // (n_regs * warp_size * num_warps)
    smem_occupancy = (size_smem //
                      kernel_smem) if kernel_smem > 0 else reg_occupancy
    thread_occupancy = max_threads_per_sm // (num_warps * warp_size)
    blocks_per_sm = max(1, min(reg_occupancy, smem_occupancy,
                               thread_occupancy))

    return num_sms * blocks_per_sm


def gemm_swizzle(a: torch.Tensor,
                  b: torch.Tensor,
                  swizzle_mode: int,
                  block_m: int,
                  block_n: int,
                  block_k: int,
                  num_programs: int,
                  group_size_m: int = 8,
                  dtype: str = "fp16") -> torch.Tensor:
    """Allocates the output tensor and launches exactly `num_programs`
    programs for one swizzle mode. Returns the (M, N) result tensor.
    """
    M, K = a.shape
    K2, N = b.shape
    assert K == K2
    c = torch.empty((M, N), device=a.device, dtype=TORCH_DTYPES[dtype])
    grid = (num_programs, )
    gemm_swizzle_kernel[grid](
        a,
        b,
        c,
        M,
        N,
        K,
        a.stride(0),
        a.stride(1),
        b.stride(0),
        b.stride(1),
        c.stride(0),
        c.stride(1),
        BLOCK_M=block_m,
        BLOCK_N=block_n,
        BLOCK_K=block_k,
        SWIZZLE_MODE=swizzle_mode,
        GROUP_SIZE_M=group_size_m,
        DTYPE=TL_DTYPES[dtype],
        NUM_PROGRAMS=num_programs,
    )
    return c


# --- L2 hit-rate predictor (byte-granular LRU cache-residency simulation) ---


class ByteLRUCache:
    """Byte-capacity LRU cache used to trace-simulate L2 tile residency.

    Each cached entry is identified by an arbitrary hashable key (here, an
    ('A'|'B', tile_index, k_index) tuple) and tracks its own byte size, so
    tiles of different shapes (A vs. B) can share one capacity budget.
    """

    def __init__(self, capacity_bytes: int) -> None:
        self.capacity_bytes = capacity_bytes
        self.cache: "OrderedDict[Any, int]" = OrderedDict(
        )  # key -> size in bytes
        self.current_bytes = 0
        self.hits = 0
        self.misses = 0

    def access(self, key: Any, size_bytes: int) -> None:
        if key in self.cache:
            self.hits += 1
            self.cache.move_to_end(key)
        else:
            self.misses += 1
            self.cache[key] = size_bytes
            self.current_bytes += size_bytes
            while self.current_bytes > self.capacity_bytes and len(
                    self.cache) > 1:
                _, evicted_size = self.cache.popitem(last=False)
                self.current_bytes -= evicted_size

    @property
    def hit_rate(self) -> float:
        """Cumulative hit rate (percentage) across all access() calls so far."""
        total = self.hits + self.misses
        return (self.hits / total * 100) if total > 0 else 0.0


def predict_l2_hit_rate(fn: Callable[..., Tuple[int, int]],
                         kwargs: Dict[str, Any], M: int, N: int, K: int,
                         block_m: int, block_n: int, block_k: int,
                         l2_cache_size_mb: float, concurrency: int,
                         dtype_bytes: int) -> float:
    """Estimates the L2 cache hit rate for one swizzle mode via an offline,
    byte-granular LRU trace simulation, generalized to non-square A/B tile
    shapes.

    `fn` is a CPU pid-mapping function (e.g. linear_row_major, panel_swizzle)
    taking (pid, grid_m, grid_n, **kwargs) and returning (bx, by). Tile ids
    are grouped into waves of `concurrency` tiles (`min(total_tiles,
    measured_capacity)`) to approximate/reproduce real concurrent execution;
    within each wave, every tile's A/B sub-blocks are accessed for every
    K-iteration and replayed through an LRU cache sized to
    `l2_cache_size_mb`.

    This is a coarse model: it does not account for DRAM row-buffer
    locality, bank-level contention, or memory coalescing, so its output
    should be read as a directional signal, not a precise throughput
    predictor.
    """
    grid_m = math.ceil(M / block_m)
    grid_n = math.ceil(N / block_n)
    grid_k = math.ceil(K / block_k)
    total_tiles = grid_m * grid_n
    cache = ByteLRUCache(l2_cache_size_mb * 1024 * 1024)
    a_tile_bytes = block_m * block_k * dtype_bytes
    b_tile_bytes = block_k * block_n * dtype_bytes

    tiles = [fn(pid, grid_m, grid_n, **kwargs) for pid in range(total_tiles)]
    for wave_start in range(0, total_tiles, concurrency):
        wave_tiles = tiles[wave_start:wave_start + concurrency]
        for k in range(grid_k):
            for bx, by in wave_tiles:
                cache.access(("A", bx, k), a_tile_bytes)
                cache.access(("B", k, by), b_tile_bytes)
    return cache.hit_rate


def rank_correlation(values_a: List[float], values_b: List[float]) -> float:
    """Spearman rank correlation with no external dependencies."""

    def ranks(values: List[float]) -> List[float]:
        order = sorted(range(len(values)), key=lambda i: values[i])
        r = [0.0] * len(values)
        for rank, idx in enumerate(order):
            r[idx] = rank
        return r

    ra, rb = ranks(values_a), ranks(values_b)
    n = len(values_a)
    mean_a, mean_b = sum(ra) / n, sum(rb) / n
    cov = sum((ra[i] - mean_a) * (rb[i] - mean_b) for i in range(n))
    var_a = sum((x - mean_a)**2 for x in ra)
    var_b = sum((x - mean_b)**2 for x in rb)
    if var_a == 0 or var_b == 0:
        return 0.0
    return cov / math.sqrt(var_a * var_b)


def _is_power_of_two(n: int) -> bool:
    """True if n is a positive power of two."""
    return n > 0 and (n & (n - 1)) == 0


AlgorithmEntry = Tuple[str, Callable[..., Tuple[int, int]], int, Dict[str,
                                                                       Any],
                       int]


def build_algorithm_list(group_sizes: List[int], grid_m: int,
                          grid_n: int) -> List[AlgorithmEntry]:
    """Builds the list of (name, cpu_fn, swizzle_mode, kwargs, group_size_m)
    tuples to benchmark for the given grid shape.

    Always includes Linear Row-Major/Col-Major and one Grouped 2D Panel
    entry per value in `group_sizes`. Morton/Z-order and Bitwise XOR are
    included only when (grid_m, grid_n) is square and a power of two --
    their bit tricks are not bijective otherwise (see module docstring), so
    they are skipped with a printed warning rather than silently producing
    incorrect results.
    """
    algorithms = [
        ("Linear (Row-Major)", linear_row_major, MODE_ROW_MAJOR, {}, 1),
        ("Linear (Col-Major)", linear_col_major, MODE_COL_MAJOR, {}, 1),
    ]
    for g in group_sizes:
        algorithms.append((
            f"Grouped 2D Panel (g={g})",
            panel_swizzle,
            MODE_PANEL,
            {
                "group_size_m": g
            },
            g,
        ))

    # Morton/XOR's bit tricks (interleaving, XOR-then-mod) are only bijective
    # (visit every tile exactly once) when the grid is square and a power of
    # two. Otherwise they silently skip some tiles and recompute others.
    if grid_m == grid_n and _is_power_of_two(grid_m):
        algorithms.append(
            ("Morton / Z-Order", morton_swizzle, MODE_MORTON, {}, 1))
        algorithms.append(("Bitwise XOR", xor_swizzle, MODE_XOR, {}, 1))
    else:
        print(
            f"[WARNING] Skipping Morton/XOR: grid {grid_m}x{grid_n} is not a "
            "square power-of-two, so those swizzles would drop/duplicate "
            "tiles and produce incorrect results.\n",
            file=sys.stderr,
        )
    return algorithms


def check_correctness(block_m: int,
                       block_n: int,
                       block_k: int,
                       dtype: str = "fp16") -> None:
    """Verifies every swizzle mode against torch.matmul, on a small, fixed
    square power-of-two grid (so Morton/XOR are well-defined) that doesn't
    compete with the (possibly huge) benchmark tensors for VRAM. One
    program is launched per tile here (num_programs == total_tiles), so
    each program's loop runs exactly once regardless of --persistent.
    """
    torch.manual_seed(0)
    grid_size = 4  # power of two, square: keeps every swizzle mode valid
    m, n, k = block_m * grid_size, block_n * grid_size, max(block_k * 2, 64)
    torch_dtype = TORCH_DTYPES[dtype]
    a = torch.randn((m, k), device="cuda", dtype=torch_dtype)
    b = torch.randn((k, n), device="cuda", dtype=torch_dtype)
    ref = torch.matmul(a, b)
    # fp32 tl.dot uses TF32 tensor cores by default (~10-bit mantissa, same
    # precision class as fp16), so it needs a loose tolerance too, with extra
    # margin for its worst-case (tail) accumulation error.
    rtol, atol = (1e-2, 1e-2) if dtype == "fp16" else (3e-2, 3e-2)
    num_programs = grid_size * grid_size

    algorithms = build_algorithm_list([1, 2, 4], grid_size, grid_size)
    for name, _, mode, _kwargs, group_size_m in algorithms:
        out = gemm_swizzle(a, b, mode, block_m, block_n, block_k,
                            num_programs, group_size_m, dtype)
        torch.testing.assert_close(out,
                                    ref,
                                    rtol=rtol,
                                    atol=atol,
                                    msg=f"{name} produced incorrect results")
        del out
    del a, b, ref
    torch.cuda.empty_cache()
    print(f"Correctness check passed for all swizzle modes ({dtype}).\n")


def run_benchmark(args: argparse.Namespace, hw_config: Dict[str,
                                                             Any]) -> None:
    """Runs the full benchmark: correctness check, concurrency measurement,
    and one predicted-hit-rate + timed-throughput measurement per swizzle
    mode, printed as a summary table followed by a rank-correlation sanity
    check.
    """
    check_correctness(args.block_m, args.block_n, args.block_k, args.dtype)

    torch.manual_seed(0)
    torch_dtype = TORCH_DTYPES[args.dtype]
    a = torch.randn((args.M, args.K), device="cuda", dtype=torch_dtype)
    b = torch.randn((args.K, args.N), device="cuda", dtype=torch_dtype)

    concurrent_capacity = measure_concurrent_capacity(a.device, args.block_m,
                                                       args.block_n,
                                                       args.block_k,
                                                       args.dtype)

    grid_m = math.ceil(args.M / args.block_m)
    grid_n = math.ceil(args.N / args.block_n)
    total_tiles = grid_m * grid_n
    # Real achievable concurrency either way; the predictor uses this
    # regardless of launch style (see module docstring).
    concurrency = min(total_tiles, concurrent_capacity)
    num_programs = concurrency if args.persistent else total_tiles
    algorithms = build_algorithm_list(args.group_sizes, grid_m, grid_n)

    print("=" * 92)
    print("TRITON GEMM SWIZZLE BENCHMARK")
    print("=" * 92)
    print(f"Hardware Platform  : {hw_config['gpu_name']}")
    print(f"Streaming Multiprocs: {hw_config['num_sms']} SMs")
    print(f"L2 Cache Capacity  : {hw_config['l2_cache_size_mb']:.2f} MB")
    print(f"Matrix Dimensions  : {args.M} x {args.N} x {args.K}")
    print(f"Tile Size          : {args.block_m}x{args.block_n}x{args.block_k}")
    print(f"Data Type          : {args.dtype}")
    print(f"Total Output Tiles : {total_tiles}")
    print(f"Measured Concurrent Capacity: {concurrent_capacity} CTAs")
    print("Kernel Mode        : "
          f"{'Persistent' if args.persistent else 'Non-Persistent'}"
          f" ({num_programs} programs launched)")
    print("=" * 92)
    print(f"{'Algorithm':<28} | {'Pred. L2 Hit %':>14} | {'Time (ms)':>10} | "
          f"{'TFLOPS':>8}")
    print("-" * 92)

    hit_rates, tflops_list = [], []
    for name, fn, mode, kwargs, group_size_m in algorithms:
        hit_rate = predict_l2_hit_rate(fn, kwargs, args.M, args.N, args.K,
                                        args.block_m, args.block_n,
                                        args.block_k,
                                        hw_config["l2_cache_size_mb"],
                                        concurrency, DTYPE_BYTES[args.dtype])

        def bench_fn(mode=mode, group_size_m=group_size_m):
            return gemm_swizzle(a, b, mode, args.block_m, args.block_n,
                                 args.block_k, num_programs, group_size_m,
                                 args.dtype)

        ms = triton.testing.do_bench(bench_fn,
                                      warmup=args.warmup,
                                      rep=args.rep)
        tflops = (2 * args.M * args.N * args.K) / (ms * 1e-3) / 1e12

        print(
            f"{name:<28} | {hit_rate:>13.2f}% | {ms:>10.3f} | {tflops:>8.2f}")
        hit_rates.append(hit_rate)
        tflops_list.append(tflops)

    print("-" * 92)
    corr = rank_correlation(hit_rates, tflops_list)
    print("Spearman rank correlation (predicted L2 hit rate vs measured "
          f"TFLOPS): {corr:.3f}")
    print(
        "(+1 = higher predicted hit rate always faster, -1 = always slower, 0 = uncorrelated)"
    )


def _validate_args(args: argparse.Namespace) -> None:
    """Validates CLI arguments up front, raising a clear error instead of a
    cryptic Triton compile failure or silently-wrong behavior."""
    for name in ("M", "N", "K", "warmup", "rep"):
        if getattr(args, name) <= 0:
            raise ValueError(
                f"--{name} must be positive, got {getattr(args, name)}")

    for name in ("block_m", "block_n", "block_k"):
        value = getattr(args, name)
        flag = name.replace("_", "-")
        if not _is_power_of_two(value):
            # tl.arange (used to build per-tile offsets) requires a power-of-two size.
            raise ValueError(f"--{flag} must be a power of two, got {value}")

    for g in args.group_sizes:
        if g <= 0:
            raise ValueError(f"--group-sizes values must be positive, got {g}")


def main() -> None:
    """CLI entry point: parses arguments, resolves hardware config
    (auto-detected via detect_gpu_hardware, with optional CLI overrides),
    and runs the benchmark.
    """
    hw_detected = detect_gpu_hardware()

    parser = argparse.ArgumentParser(
        description=(
            "Benchmark a Triton FP16/FP32 GEMM kernel across several "
            "thread-block swizzle orderings."),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--M", type=int, default=8192)
    parser.add_argument("--N", type=int, default=8192)
    parser.add_argument("--K", type=int, default=8192)
    parser.add_argument("--block-m", type=int, default=128)
    parser.add_argument("--block-n", type=int, default=128)
    parser.add_argument("--block-k", type=int, default=32)
    parser.add_argument(
        "--dtype",
        choices=["fp16", "fp32"],
        default="fp16",
        help="Data type shared by A, B, and C (fp32 uses the same tl.dot kernel path)",
    )
    parser.add_argument(
        "--group-sizes",
        type=int,
        nargs="+",
        default=[1, 2, 4, 8, 16, 32, 64],
        help="Group sizes to try for the Grouped 2D Panel swizzle",
    )
    parser.add_argument(
        "--persistent",
        action="store_true",
        help=("Launch exactly min(total_tiles, measured_capacity) "
              "persistent programs, each looping over several tiles, "
              "instead of one program per tile (see module docstring)"),
    )
    parser.add_argument(
        "--num-sms",
        type=int,
        default=None,
        help="Override SM count reported/used by the predictor",
    )
    parser.add_argument(
        "--l2-cache-mb",
        type=float,
        default=None,
        help="Override L2 cache size in MB used by the predictor",
    )
    parser.add_argument("--warmup", type=int, default=25)
    parser.add_argument("--rep", type=int, default=100)
    args = parser.parse_args()
    _validate_args(args)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA device required to run this benchmark.")

    final_sms = args.num_sms if args.num_sms is not None else hw_detected[
        "num_sms"]
    final_l2 = (args.l2_cache_mb if args.l2_cache_mb is not None else
                hw_detected["l2_cache_size_mb"])
    hw_config = {
        "num_sms": final_sms,
        "l2_cache_size_mb": final_l2,
        "gpu_name": hw_detected["gpu_name"],
    }

    run_benchmark(args, hw_config)


if __name__ == "__main__":
    main()

It turns out that it is very difficult to correlate the actual kernel performance with the predicted L2 cache hit rate for both persistent and non-persistent GEMM kernel implementations.

For example, for $8192 \times 8192 \times 8192$ GEMM, with tile sizes of $128 \times 128 \times 32$, in spite of the lower L2 cache hit rate predicted for the row-major linear swizzle, there is no significant performance degradation compared to the grouped 2D panel swizzle.

What’s also interesting is that the column-major linear swizzle performs significantly worse than the row-major linear swizzle, despite having the same predicted L2 cache hit rate. In row-major order, row blocks of matrix B, which are contiguous in memory, are accessed in each wave. In column-major order, however, column blocks of matrix A, which are strided in memory, are accessed. Note that because both matrices are at least 16-byte aligned, the memory accesses with both row-major order and column-major order are still fully coalesced and vectorized. The reason to this performance difference is probably at much lower level, which we could not model, such as the DRAM access patterns.

$ python triton_gemm_swizzle_bench.py --M 8192 --N 8192 --K 8192 --block-m 128 --block-n 128 --block-k 32
Correctness check passed for all swizzle modes (fp16).

============================================================================================
TRITON GEMM SWIZZLE BENCHMARK
============================================================================================
Hardware Platform  : NVIDIA GeForce RTX 5080
Streaming Multiprocs: 84 SMs
L2 Cache Capacity  : 64.00 MB
Matrix Dimensions  : 8192 x 8192 x 8192
Tile Size          : 128x128x32
Data Type          : fp16
Total Output Tiles : 4096
Measured Concurrent Capacity: 168 CTAs
Kernel Mode        : Non-Persistent (4096 programs launched)
============================================================================================
Algorithm                    | Pred. L2 Hit % |  Time (ms) |   TFLOPS
--------------------------------------------------------------------------------------------
Linear (Row-Major)           |         79.43% |      9.342 |   117.70
Linear (Col-Major)           |         79.43% |     11.548 |    95.21
Grouped 2D Panel (g=1)       |         79.43% |      9.383 |   117.18
Grouped 2D Panel (g=2)       |         79.52% |      9.486 |   115.90
Grouped 2D Panel (g=4)       |         85.57% |      9.476 |   116.03
Grouped 2D Panel (g=8)       |         92.13% |      9.464 |   116.18
Grouped 2D Panel (g=16)      |         95.27% |      9.482 |   115.95
Grouped 2D Panel (g=32)      |         88.06% |      9.485 |   115.92
Grouped 2D Panel (g=64)      |         79.43% |     11.636 |    94.49
Morton / Z-Order             |         89.39% |      9.590 |   114.66
Bitwise XOR                  |         79.43% |     11.530 |    95.36
--------------------------------------------------------------------------------------------
Spearman rank correlation (predicted L2 hit rate vs measured TFLOPS): 0.045
(+1 = higher predicted hit rate always faster, -1 = always slower, 0 = uncorrelated)

$ python triton_gemm_swizzle_bench.py --M 8192 --N 8192 --K 8192 --block-m 128 --block-n 128 --block-k 32 --persistent
Correctness check passed for all swizzle modes (fp16).

============================================================================================
TRITON GEMM SWIZZLE BENCHMARK
============================================================================================
Hardware Platform  : NVIDIA GeForce RTX 5080
Streaming Multiprocs: 84 SMs
L2 Cache Capacity  : 64.00 MB
Matrix Dimensions  : 8192 x 8192 x 8192
Tile Size          : 128x128x32
Data Type          : fp16
Total Output Tiles : 4096
Measured Concurrent Capacity: 168 CTAs
Kernel Mode        : Persistent (168 programs launched)
============================================================================================
Algorithm                    | Pred. L2 Hit % |  Time (ms) |   TFLOPS
--------------------------------------------------------------------------------------------
Linear (Row-Major)           |         79.43% |      9.574 |   114.85
Linear (Col-Major)           |         79.43% |     11.672 |    94.20
Grouped 2D Panel (g=1)       |         79.43% |      9.516 |   115.54
Grouped 2D Panel (g=2)       |         79.52% |      9.638 |   114.08
Grouped 2D Panel (g=4)       |         85.57% |      9.663 |   113.79
Grouped 2D Panel (g=8)       |         92.13% |      9.636 |   114.11
Grouped 2D Panel (g=16)      |         95.27% |      9.631 |   114.17
Grouped 2D Panel (g=32)      |         88.06% |      9.869 |   111.41
Grouped 2D Panel (g=64)      |         79.43% |     11.875 |    92.59
Morton / Z-Order             |         89.39% |      9.703 |   113.31
Bitwise XOR                  |         79.43% |     11.681 |    94.13
--------------------------------------------------------------------------------------------
Spearman rank correlation (predicted L2 hit rate vs measured TFLOPS): 0.064
(+1 = higher predicted hit rate always faster, -1 = always slower, 0 = uncorrelated)

Such performance differences might diminish if other benchmark configurations are used or a different compute platform is considered.

Conclusions

Predicting how thread block swizzle affects the actual performance of GEMM kernels is challenging. It might just be more pragmatic to empirically benchmark different swizzle strategies for specific kernels and configurations rather than relying solely on predicted L2 cache hit rates.

The order of thread block execution can significantly impact cache utilization and overall performance. In some other non-GEMM applications, sometimes simply switching between the row-major and column-major execution order can lead to noticeable performance differences.

References

• CUTLASS Tutorial: Persistent Kernels and Stream-K
• Matrix Multiplication - Triton Tutorial
• Optimizing Compute Shaders for L2 Locality using Thread-Group ID Swizzling