Skip to content
Phobos

@aligned

@aligned is a promise about the data: a tensor dimension is a whole number of tiles. With it the compiler drops the masks on that dimension and may access its rows with 16-byte vector loads and stores.

@aligned(DIM = tile, ...)

DIM is a symbolic dimension of the kernel’s tensors. tile is a literal or an @autotune constant. See the Attributes table and Bounds and masking in the spec.

A minimal kernel

Section titled “A minimal kernel”
@autotune(TM in [32], TN in [128])
@aligned(M = TM, N = TN)
kernel scale(X: tensor<f32>[M, N], Y: tensor<f32>[M, N], alpha: f32) {
  let pm = ctx.block.x
  let pn = ctx.block.y
  var x = X[pm * TM :+ TM, pn * TN :+ TN]
  Y[pm * TM :+ TM, pn * TN :+ TN] = alpha * x
}

Without the attribute, every CTA’s copy and store would be masked against the runtime M and N, in case the CTA sits at the edge. With it, no CTA is at a ragged edge, and the slices are read and written whole.

without @aligned

@autotune(TILE_M in [32, 256], TILE_N in [32, 256], TILE_K in [4, 32])
@launch(256)
@pipeline
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@aligned droppedOpen full view

The gemm kernel with @aligned dropped, 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 8192 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

with @aligned

@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.phOpen full view

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

What it changes

Section titled “What it changes”
Without @aligned With @aligned(DIM = tile)
A slice along DIM indexed by ctx.block is masked against the runtime extent. No mask.
Extents are only assumed to be multiples of 4 elements. Rows along DIM may be accessed with 16-byte vector loads and stores.
A for over DIM with mma.sync accumulators stays masked. Unmasked.
A thread’s run in sum, elem or store is proven aligned only from its subscripts. The divisor joins the proof, so more runs become single vector accesses.

The before and after pair above shows the same kernel compiled both ways.

When the promise is broken

Section titled “When the promise is broken”

On a dynamic DIM, the promise is checked once at kernel entry. A launch that breaks it traps, and the next synchronize on the host fails, instead of the kernel writing past the end of a tensor. In a @persistent kernel that uses grid_barrier, the other blocks then wait at the barrier for good, so a broken promise there shows up as a hang rather than an error.

The arguments themselves are checked at compile time: a tile that is neither a literal nor an @autotune constant is an unknown constant, and a divisor must be positive.

Interactions

Section titled “Interactions”
  • @autotune: the common pairing is @aligned(M = TILE_M, N = TILE_N, K = TILE_K), so the promise follows whatever tile sizes the tuner picks. The host must then pad or pick shapes that keep the promise for every candidate.
  • @pipeline: a loop over a dynamic extent covered by @aligned is a whole number of tiles, so all of it can be pipelined, with no remainder iteration.
  • @tensorcore: mma.sync loops stay masked unless @aligned covers the dimension they walk.