Module gemm_tile 

The CUDA GEMM tiling lattice — the search space our optimizing compiler sweeps to climb from the naive one-thread-per-output DP4A kernel toward the cuBLAS ceiling.

A CudaTile is one point: a thread block computes a bm×bn output tile by staging bk-deep slices of A and B through shared memory (so each byte is reused across the tile instead of re-read from global memory — the ~20× “tiling/reuse” lever), and each thread holds a tm×tn register accumulator. A point is legal only if it tiles the block evenly and fits the GPU’s shared-memory + thread budget. The autotuner enumerates the legal lattice, measures each point on the device, and locks the fastest — the same divisor-lattice discipline as the Tenstorrent tune_tiling.

Structs

CudaTile

A point in the GEMM tiling lattice.

CutlassI8Gemm

A CUTLASS int8 tensorop-GEMM config. cuBLAS’s own int8 GEMM kernel IS a CUTLASS kernel — disassembly of cublasGemmEx (CUDA_R_8I, TN, 2048^3, sm_86) shows it dispatch cutlass_80_tensorop_i16832gemm_s8_256x128_64x3_tn_align16. Because CUTLASS is open source, our compiler can emit the same kernel and reach cuBLAS parity (measured 99.6-100.0% head-to-head), instead of approximating it with a hand-written kernel that tops out at ~90% (ptxas will not schedule our from-scratch source into the perfect pointer-increment addressing CUTLASS’s templates coax out of it).

CutlassNvfp4Gemm

A CUTLASS Blackwell (sm_100) NVFP4 block-scaled GEMM config — the GPU MODE nvfp4_gemm leaderboard target. Operands are FP4 (e2m1) with per-16-element FP8 (ue4m3) block scale factors; accumulate FP32; output FP16. cuBLAS/torch._scaled_mm’s own NVFP4 path is a CUTLASS kernel (as with int8), so our compiler emits the same CollectiveBuilder GEMM and then pulls ahead by consuming the pre-permuted GPU scale factors directly (the reference re-derives them from a CPU tensor inside the timed region — free speedup) + epilogue fusion. NOTE: sm_100 / NVFP4 requires CUTLASS >= 3.7 and a Blackwell GPU; validate on the leaderboard runner (popcorn submit --mode test --gpu B200), not on sm_86 hardware.

GpuLimits

The per-block hardware budget the lattice is bounded by.

WmmaTile

A pipelined-WMMA configuration: block tile BM×BN, K-stage BK, warp grid WARPS_M×WARPS_N.

Enums

ComputePrimitive

Which compute hardware a GEMM kernel targets.

CutlassEpilogue

A fused pointwise epilogue our compiler folds INTO the GEMM — the thing cuBLAS cannot do in-kernel. cuBLAS must write the GEMM result to global memory and run a SECOND kernel for the activation (an extra full-output read+write); fusing it here removes that round-trip and BEATS a cuBLAS+elementwise baseline.

GemmLayout

GEMM operand layout (BLAS transpose convention). The int8 tensor cores are fastest in Tn.

HmmaKernel

Which half-precision WMMA GEMM kernel is the FASTEST legal choice for a shape.

MmaKind

The tensor-core (or CUDA-core) MMA family a GEMM lowers to, keyed on operand PRECISION. Int8 (Dp4aI8 CUDA-core, WmmaI8 IMMA) already exist; HmmaBf16/HmmaF16 add the half-precision tensor-core path. (FP8/NVFP4 join later via the mma.sync/tcgen05 PTX path on Hopper/Blackwell.)

Constants

AMPERE_SM_REGS

SM-level register file on Ampere (sm_86): 65,536 32-bit registers per SM. A tile needs regs_per_thread · threads ≤ this to place at least one block per SM.

Functions

cpasync_stages_fit

A stages-deep pipeline fits iff stages shared buffers (each = the single-stage footprint) fit smem.

cuda_tile_lattice

Enumerate the legal tiling lattice for gpu: the divisor space of block tiles × K-stages × thread tiles. Only points that pass CudaTile::is_legal are returned, so every element is directly emittable.

default_wmma_tile

A sensible default rung-3 config: 128×128 block, BK=32, 8 warps (2×4), so each warp owns a 64×32 warp-tile = 4×2 = 8 fragments. This is the classic WMMA GEMM shape.

emit_cutlass_i8_gemm

Plain GEMM (parity with cuBLAS) — the un-fused kernel.

emit_cutlass_i8_gemm_bench

Full head-to-head parity program: our compiler-emitted CUTLASS kernel vs cublasGemmEx, same process, interleaved best-of-100, varied-fill oracle-checked. Prints PARITY <pct> OURS <gops> CUBLAS <gops>.

emit_cutlass_i8_gemm_ep

Emit the CUTLASS int8 TN GEMM kernel our compiler chooses — the same one cuBLAS runs (with an optional fused epilogue). Produces the using Gemm = ... instantiation + an extern "C" run_gemm_k launcher. Requires the CUTLASS headers on the include path (-I <cutlass>/include).

emit_cutlass_i8_gemm_relu_bench

Full head-to-head program: our FUSED GEMM+ReLU vs the baseline cuBLAS forces (GEMM then a SEPARATE relu kernel). Both are oracle-checked against relu(A*B) and against each other; prints FUSED SPEEDUP <x>x (our_time vs cublas+relu_time). This is where we BEAT cuBLAS — via fusion.

emit_cutlass_nvfp4_gemm

Emit our compiler’s CUTLASS Blackwell NVFP4 block-scaled GEMM as a torch-extension launcher for load_inline on the Blackwell leaderboard runner. Consumes the pre-permuted FP8 block scales directly.

emit_cutlass_nvfp4_group_gemm

Emit our compiler’s CUTLASS Blackwell NVFP4 grouped GEMM (MoE) — all G experts in ONE kernel launch via CUTLASS GroupProblemShape + per-group pointer arrays. This is the fusion edge on the GPU MODE nvfp4_group_gemm leaderboard: the reference launches G separate GEMMs; we launch one grouped kernel over all experts (no per-expert launch overhead, one wave of tensor-core work). Validate on the B200 runner.

emit_gemm_program

Emit a full runnable GEMM program, choosing the compute primitive from the problem + hardware. When the GPU has tensor cores and the shape is 16-aligned, this emits the WMMA (tensor-core) kernel; otherwise the tiled DP4A (CUDA-core) kernel driven by tile. Both paths are byte-exact vs the scalar oracle.

emit_mma_tn_colb_bench

emit_mma_tn_colb_program

emit_mma_tn_cpasync_bench

emit_mma_tn_cpasync_program

emit_mma_tn_i8

emit_mma_tn_i8_bench

emit_mma_tn_i8_program

emit_mma_tn_k32_bench

emit_mma_tn_k32_program

emit_mma_tn_k32db_bench

emit_mma_tn_k32db_program

emit_mma_tn_k32swz_bench

emit_mma_tn_k32swz_kernel

Emit JUST the swizzled-k32 kernel (no host main) — so the compiler can compile it to a cubin and analyze its SASS (register pressure, occupancy, tensor-core utilisation) via codegen_sass::SassReport.

emit_mma_tn_k32swz_program

emit_mma_tn_k64_bench

emit_mma_tn_k64_kernel

Emit JUST the BK=64 champion kernel (no host main) — for the head-to-head cuBLAS parity harness and SASS analysis. Renames the entry to the given symbol so it can be linked beside cublas in one program.

emit_mma_tn_k64_program

emit_mma_tn_ldm_cfg_bench

Emit the ldmatrix+mma.sync.s8 TN kernel for a specific config so the frontier can sweep block/warp/stage shapes on-device. stages shared buffers (cp.async N-stage). WN must be a multiple of 8, WM of 16.

emit_mma_tn_ldm_swizzle_bench

emit_mma_tn_ldm_swizzle_program

The swizzled program (varied-fill oracle correctness).

emit_mma_tn_ldmatrix_bench

emit_mma_tn_ldmatrix_program

emit_tiled_dp4a

Emit a shared-memory + register-blocked tiled DP4A GEMM kernel for tile. Each block stages a bm×bk slice of A and a bk×bn slice of B into shared memory (every byte reused across the block — the reuse the naive one-thread-per-output kernel throws away), and each thread accumulates a tm×tn register tile with __dp4a (4 signed-int8 MACs/instruction) over the K-stage. Computes C(M×N) = A(M×K)·B(K×N) in int8 with int32 accumulation. The blocking constexprs are baked from tile, so a different lattice point yields a different kernel — genuine compiler output, not a fixed template.

emit_tiled_dp4a_bench

Emit a TIMING bench for the tiled kernel at M×K×N: fills A/B on device, warms up, times 30 launches with CUDA events, and prints GOPS <value> (2·M·N·K / seconds). Used by the autotuner to measure each lattice point on the real device and lock the fastest.

emit_tiled_dp4a_program

Emit a COMPLETE self-contained tiled-DP4A program for shape M×K×N: the tiled kernel plus a host main that fills A/B deterministically, launches gemm_tiled_dp4a<<<(N/BN, M/BM), NTHREAD>>>, and prints the int32 result row-major (one value per line) — so its output is directly comparable to a scalar oracle. Requires the shape to be tile-aligned (M%bm==0, N%bn==0, K%bk==0).

emit_wmma

Emit a WMMA tensor-core GEMM kernel for a tensor-core MmaKind — the SAME 16×16×16 fragment loop as the int8 IMMA path, over the operand/accumulator types the precision fixes (bf16/f16 → f32 accumulate). Row-major A(M×K)·B(K×N). (Not for Dp4aI8, which is a CUDA-core kernel; use emit_tiled_dp4a.)

emit_wmma_basic_bench

Timing bench for the BASIC (one-warp-per-16×16-tile) WMMA kernel at kind — prints GOPS <value>.

emit_wmma_cpasync_bench

Timing bench for the cp.async double-buffered WMMA kernel at kind + tile + stages — prints GOPS.

emit_wmma_cpasync_i8_bench_for

emit_wmma_cpasync_i8_bench_staged

emit_wmma_cpasync_i8_for

emit_wmma_cpasync_i8_program_for

emit_wmma_cpasync_i8_program_staged

Program/bench for a specific pipeline DEPTH (stages shared buffers). 2 = double-buffer; 3-4 = deeper pipeline that keeps the tensor cores fed further ahead (needs stages × smem to fit).

emit_wmma_cpasync_program

Checksum-matching cp.async GEMM program for a tensor-core kind (bf16/f16) — Σ C ≡ the scalar oracle on int8-range operands (exact in bf16/f16, wide f32 accumulate). Launches the double-buffered gemm_wmma_cpasync.

emit_wmma_i8

Emit a WMMA int8 tensor-core GEMM kernel: one warp per 16×16 output tile, accumulating over K in 16-wide fragments with wmma::mma_sync (the IMMA path). Row-major A(M×K)·B(K×N) → int32 C(M×N).

emit_wmma_i8_bench

Timing bench for the WMMA int8 kernel: warms up, times 30 launches, prints GOPS <value>.

emit_wmma_i8_program

Full self-contained WMMA int8 program for M×K×N: kernel + host main that fills A/B deterministically (same fill as the tiled program, so the SAME scalar oracle applies), launches gemm_wmma_i8<<<(N/16, M/16), 32>>> (one warp per 16×16 tile), and prints int32 C row-major.

emit_wmma_program

Full self-contained WMMA program for M×K×N at MmaKind kind — fills A/B with small integers (exact in bf16/f16/int8), runs one warp per 16×16 tile, and prints Σ C (a double). Because the operands are exact small integers and the accumulate is wide (f32/int32), Σ C is an exact integer, so a scalar integer GEMM is the ≡ oracle regardless of precision.

emit_wmma_tiled_bench

Timing bench for the PIPELINED (shared-memory, warp-tiled) WMMA kernel at kind + tile — prints GOPS.

emit_wmma_tiled_i8

Emit the pipelined-WMMA kernel with the default config.

emit_wmma_tiled_i8_bench

Emit a timing harness (default config): 30 launches, prints GOPS <value>.

emit_wmma_tiled_i8_bench_for

Emit a bench for a specific config (for the config sweep).

emit_wmma_tiled_i8_for

Emit just the pipelined-WMMA kernel for tile (structural / nvcc-compile checks).

emit_wmma_tiled_i8_program

Emit a full runnable program (default config). Same deterministic fill as the other GEMM programs.

emit_wmma_tiled_i8_program_for

Emit a program for a specific config (for the config sweep).

emit_wmma_tiled_program

Checksum-matching pipelined-WMMA GEMM program for a tensor-core kind (bf16/f16) — the shared-memory, warp-tiled kernel with the matmul emitters’ fill (a[p]=(p%7)-3, b[p]=(p%5)-2) and the Σ C integer checksum, so on int8-range operands (exact in bf16/f16) it equals the scalar oracle byte-for-byte.

feasible_tiles

The full feasible tile grid — every legal tile that fits Ampere’s smem + register file. PROVABLY COMPLETE over this parameter space: nothing legal is silently dropped.

lattice_best_bf16_bench

The knob lattice’s own best bf16 rung for a shape, as a runnable benchmark: the kernel select_hmma_kernel picks (cp.async when its tile divides the shape and fits, else the basic WMMA kernel). This is the backend-native baseline a lowering measures itself against — the fair opponent is the best thing this lattice can already emit, not its worst.

pareto_frontier

CONSTRUCT the measurement frontier: the Pareto-optimal feasible tiles — those NOT dominated by any other feasible tile on the three physical axes (DRAM traffic ↓, SMEM traffic ↓, occupancy ↑). This is the honest “let it construct the answer” step: we do NOT pick a winner and we assume NO weighting — we only discard tiles that are provably worse on every axis (they cannot be optimal), then hand the survivors to the GPU, which crowns the winner by measurement. The frontier is a small superset of the true optimum; the machine, not us, resolves the un-modelable physics (bank conflicts, scheduling) that separates the survivors.

select_hmma_kernel

Select the fastest LEGAL half-precision kernel for (kind, shape, arch): cp.async when the default tile divides the shape and its bf16/f16 footprint fits static shared memory on an Ampere+ part, else the basic kernel. Callers pass a shape already known WMMA-legal (16-aligned) via select_mma.

select_layout

Select the layout to emit FROM THE PROBLEM: if the caller’s data is already TN (or we are free to choose the staging), prefer the tensor-core-native Tn for max throughput; if the data is fixed as Nn and a transpose would cost more than it saves, stay in Nn (where our kernel already beats cuBLAS). caller_is_tn = the operands are already available contraction-contiguous. This is the tradeoff knob, not a fixed choice.

select_mma

Select the MMA kind for a (precision, shape, sm_arch). Tensor-core WMMA when the shape is 16-aligned and the arch provides the precision’s HMMA/IMMA; otherwise the always-legal DP4A CUDA-core fallback — we never emit a tensor-core kernel on a shape or GPU that cannot run it.

select_primitive

SELECT the compute primitive FROM THE PROBLEM: use the tensor cores when the GPU has them AND the shape is WMMA-legal (16-aligned); otherwise fall back to DP4A on the CUDA cores. Correct by construction — we never emit a WMMA kernel on a shape or GPU that cannot run it. (Which of the two is actually faster when both are legal is settled by measurement on the frontier, but WMMA is the strictly-more-powerful default.)

wmma_legal

A shape is WMMA-legal iff every dimension is a multiple of the int8 fragment size (16).