Phobos

Back to Collectives and shuffles

staged_matvec

The staged_matvec kernel, as the compiler built it for sm_75 at commit b48331d: K = 256, N = 32; a grid of 8 CTAs, 128 threads each.

@autotune(K in [256])
@launch(128)
kernel staged_matvec(W: tensor<f32>[N, K], x: tensor<f32>[1, K], y: tensor<f32>[1, N]) {
  var xs = x[0 :+ 1, :]
  let row = ctx.block.x * ctx.warps + ctx.warp
  var acc: f32 = 0.0
  if row < N {
    for k in range(0, K, ctx.lanes * 4) {
      let at = k + ctx.lane * 4
      acc += sum(W[row, at :+ 4] * xs[0, at :+ 4])
    }
  }
  let total = sum<over ctx.lane>(acc)
  if ctx.lane == 0 && row < N {
    y[0, row] = total
  }
}

A kernel written in the docs

The staged_matvec kernel, stepped through as the compiler built it for sm_75A 8 grid of CTAs with 128 threads each. The drawing shows the tensors W, x, y in global memory with the picked CTA's current boxes, shared memory as a byte ruler of 1024 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
  • 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.