Phobos

Back to @pipeline

gemm_fp16_pipeline

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

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


  var acc: tile<f16>[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_fp16.ph@pipeline added

The gemm kernel with @pipeline added, 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 14848 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.