gemm on tensor cores

Tensor cores are units on NVIDIA GPUs since Volta that multiply small matrix fragments in one instruction, at a much higher rate than ordinary multiply-adds. This kernel is the f32 gemm with one attribute added, @tensorcore, and a different range for TILE_K.

On an RTX 2080 SUPER at M = N = K = 4096, phobos-kbench measured this kernel at 21.4 TFLOP/s, 47% of the card's tensor-core peak with f32 accumulation. The cuBLAS run beside it is cublasSgemm_v2 in f32 on CUDA cores (11.0 TFLOP/s), another precision, so the two do not compare like for like. From results/results.csv, committed 2026-10-08; all kernels are on the front page.

@autotune(TILE_M in [32, 256], TILE_N in [32, 256], TILE_K in [16, 64])
@launch(256)
@pipeline
@tensorcore
@aligned(M = TILE_M, N = TILE_N, K = TILE_K)
kernel gemm(A: tensor<f32>[M, K],
            B: tensor<f32>[K, N],
            C: tensor<f32>[M, N],
            alpha: f32,
            beta: f32) {
  let pm = ctx.block.x
  let pn = ctx.block.y

  var acc: tile<f32>[TILE_M, TILE_N] = 0.0

  for kt in range(0, K, TILE_K) {
    var a = A[pm * TILE_M :+ TILE_M, kt :+ TILE_K]
    var b = B[kt :+ TILE_K, pn * TILE_N :+ TILE_N]
    acc += dot(a, b)
  }

  let c_old = C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N]
  C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N] = alpha * acc + beta * c_old
}

examples/gemm_fp16tc_fp32acc.phOpen full view

Stepstep 1 of 9

CTA (0, 0)All CTAs

TilesLanes

sm_75sm_80

• global memory
• shared memory
• threads and lanes
• extra 128-byte lines, masked elements

What changes

The tensors are still f32, and the body is the same as the f32 gemm, line for line. Two lines differ:

• @tensorcore runs dot on tensor cores. The operands are rounded to f16, and the products are accumulated in f32. That is what the file name says: fp16 tensor cores, fp32 accumulation. The rounding loses precision in the inputs, not in the sums.
• TILE_K in [16, 64] searches 16, 32 and 64. Tensor cores work on fragments whose sides are multiples of 16, so a k-slice of 4 or 8 would not qualify.

When the tensor cores are used

@tensorcore is a request, not a guarantee. A dot runs on tensor cores only when every tile dimension and the k-slice are multiples of 16 and the CTA’s warps tile the grid of 16x16 fragments. Here every TILE_M and TILE_N in the search space (32 to 256) and every TILE_K (16 to 64) is a multiple of 16, and @launch(256) gives 8 warps. A dot that does not qualify falls back to the f32 FMA path; run with PHOBOS_LOG=debug to see which dot fell back and why.

Two instruction paths

The compiler has two ways to drive the tensor cores, and which one it takes depends on the target:

Target	Path
64-bit indices, sm_75 or newer	mma.sync with operands loaded by ldmatrix from swizzled shared memory. The fragment is m16n8k8 on sm_75 and m16n8k16 from sm_80.
otherwise	WMMA, with m16n16k16 fragments.

phobos-compile and the runtime compile this kernel with 64-bit indices on sm_80 and newer, where it takes mma.sync. On sm_75 they keep 32-bit indices, so it takes WMMA. @tensorcore(wmma) forces WMMA everywhere.

From sm_80, with 64-bit indices, the copies into shared memory can also use cp.async, which writes shared memory directly from global memory without passing through registers. That suits the double-buffered k loop that @pipeline asks for.

To see the difference, print the MLIR for both targets:

cargo run -p phobos-lang --example emit -- examples/gemm_fp16tc_fp32acc.ph
PHOBOS_CHIP=sm_80 PHOBOS_INDEX_BITS=64 cargo run -p phobos-lang --example emit -- examples/gemm_fp16tc_fp32acc.ph

The emit example does not widen indices on its own, which is why the second command sets PHOBOS_INDEX_BITS.

Run it

cargo run -r -p phobos-kbench -- --bench gemm_fp16tc_fp32acc

phobos-kbench autotunes it at M = N = K = 4096, like the f32 gemm. Its --peak-fp16tcf32acc flag sets the tensor-core peak the CSV compares against. See dot and tensor cores and @tensorcore for the rules in full.