Skip to content
Phobos

ctx.block and ctx.blocks

An output tensor tiled into CTAs, one tile per ctx.blockSchematic. An output tensor C of M rows and N columns cut into 64 by 64 tiles, one per CTA. ctx.block.x picks a row of tiles and ctx.block.y a column, as gemm does. The grid is M and N divided by 64, rounded up, so the last row and column of tiles can reach past C; that part is masked. Picking a CTA outlines the tile it owns. A batch axis stacks one grid per batch along z.FIG / CTX / BLOCKSschematicctx.block.y: columns of C (N)ctx.block.x: rows of C (M)012340123456C[288, 416], 64 x 64 tilesctx.blocks.x = ceil(288 / 64) = 5ctx.blocks.y = ceil(416 / 64) = 7ctx.blocks.z = 1grid 5 x 7 = 35 CTAspicked CTActx.block.x = 1, ctx.block.y = 2ctx.block.z = 0owns C[64:128, 128:192]a whole tile inside C

C[288, 416] in 64 x 64 tiles needs a grid of 5 x 7 CTAs. CTA (1, 2) owns C[64:128, 128:192], a whole tile.

  • global memory
  • past C, masked
Schematic. C[pm * 64 :+ 64, pn * 64 :+ 64] with pm = ctx.block.x and pn = ctx.block.y. Click a tile, or use the arrow keys and Enter, to pick its CTA.

A launch is a grid of CTAs with up to three axes. ctx.block is this CTA’s position in the grid, and ctx.blocks is the grid’s extent. In the figure an output tensor is cut into tiles, one per CTA, and picking a CTA outlines the tile it owns.

Field Result
ctx.block.x, ctx.block.y, ctx.block.z The CTA’s index along a grid axis, from 0.
ctx.blocks.x, ctx.blocks.y, ctx.blocks.z The grid’s extent along that axis.

Both are uniform: every thread of a CTA reads the same ctx.block, and every thread of the grid reads the same ctx.blocks. Both are runtime values. The grid is chosen by whoever launches the kernel, so the kernel only learns it when it runs. An axis is required: ctx.block alone is an error.

See Indexing and synchronization in the spec.

Tiling an output

Section titled “Tiling an output”

The usual pattern gives each CTA one tile of the output, so the grid is the output’s shape divided by the tile shape, rounded up. A one-dimensional example, where CTA i owns elements i * BLOCK up to i * BLOCK + BLOCK:

@autotune(BLOCK in [256, 4096])
kernel add(x: tensor<f32>[N], y: tensor<f32>[N], out: tensor<f32>[N]) {
  let base = ctx.block.x * BLOCK
  out[base :+ BLOCK] = x[base :+ BLOCK] + y[base :+ BLOCK]
}

For N = 10000 and BLOCK = 1024 the launch needs 10 CTAs along x. The last one owns elements 9216 to 10239, of which only 9216 to 9999 exist. Phobos masks that tile against the tensor’s runtime extent: the reads past the end are 0 and the stores past the end are skipped. See Slicing, bounds and masking.

In two dimensions, .x and .y pick a row and a column of tiles. This is how the gemm example addresses its output:

let pm = ctx.block.x
let pn = ctx.block.y
// ...
C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N] = alpha * acc + beta * c_old

A grid of ceil(M / TILE_M) by ceil(N / TILE_N) CTAs then covers C once. Which grid axis maps to which tensor axis is the kernel’s choice; Phobos attaches no meaning to x, y or z.

A third axis

Section titled “A third axis”

ctx.block.z is often used for an axis that is not part of the tile at all, such as a batch or a head index, with the tile coordinates on x and y:

let head = ctx.block.z
let row = head * S + ctx.block.x * BR
let q = Q[row :+ BR, :]

Using the grid extent

Section titled “Using the grid extent”

ctx.blocks lets a kernel run correctly on a grid smaller than its data. In a grid-stride loop each CTA starts at its own tile and steps by the whole grid:

@autotune(BLOCK in [1024])
kernel scale(x: tensor<f32>[N], out: tensor<f32>[N], alpha: f32) {
  for base in range(ctx.block.x * BLOCK, N, ctx.blocks.x * BLOCK) {
    out[base :+ BLOCK] = alpha * x[base :+ BLOCK]
  }
}

The loop bounds are uniform, so every thread of a CTA runs the same iterations and the tile operations inside stay legal.

A @persistent kernel uses the same pattern with a grid sized to the card rather than to the data. grid_barrier makes every CTA of the grid wait for every other, and their global writes before it are visible to all of them after it. It requires a one-dimensional grid.