Phobos

Back to Collectives and shuffles

lanes_matvec

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

@launch(128)
@aligned(K = 128)
kernel lanes_matvec(W: tensor<f32>[N, K], x: tensor<f32>[1, K], y: tensor<f32>[1, N]) {
  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] * x[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 lanes_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 0 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.