Phobos

Back to the example's page

gemm_fp32

The gemm kernel, as the compiler built it for sm_75 at commit 5a313cd: TILE_K = 4, TILE_M = 32, TILE_N = 32, K = 64, M = 64, N = 64; a grid of 2 x 2 CTAs, 256 threads each.

@autotune(TILE_M in [32, 256], TILE_N in [32, 256], TILE_K in [4, 32])
@launch(256)
@pipeline
@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_fp32.ph

The gemm kernel, stepped through as the compiler built it for sm_75A 2 x 2 grid of CTAs with 256 threads each. The drawing shows the tensors A, B, C in global memory with the picked CTA's current boxes, shared memory as a byte ruler of 2048 bytes with each buffer at its planned offset, the barriers kept and elided, and how the CTA's threads and warp 0's lanes take part. The readout below describes each step.

CTA (0, 0)
  • global memory
  • shared memory
  • threads and lanes
  • extra 128-byte lines, masked elements

Keys: left and right arrows step, space plays and pauses; play runs on from the last step to the first. The address of this page keeps the step, the CTA and the zoom.