Skip to content
Phobos

ctx.thread, ctx.warp, ctx.lane

One CTA as rows of warps, and the indices of one threadSchematic. A CTA of 256 threads drawn as 8 rows of 32 cells, one row per warp. Thread 113 is picked: it is lane 17 of warp 3, and 113 = 3 * 32 + 17. In the loop mode, a row of 128 elements of W is drawn as 4 rows of 32 under the lane columns; the loop for k in range(ctx.lane, K, ctx.lanes) gives lane 17 the elements 17, 49, 81 and 113, and each iteration the warp reads 32 neighboring elements.FIG / CTX / INDICESschematicwarp 0warp 1warp 2warp 3warp 4warp 5warp 6warp 70481220242817lanek 0k 32k 64k 96W[col, 0:128] for k in range(ctx.lane, K, ctx.lanes)thread = warp * lanes + lane113 = 3 * 32 + 17ctx.thread = 113ctx.warp = 113 / 32 = 3ctx.lane = 113 - 3 * 32 = 17

Thread 113 is lane 17 of warp 3: 113 = 3 * 32 + 17.

  • global memory
  • registers
Schematic. @launch(256) on NVIDIA, 32 lanes a warp. Hover or focus a thread, or use the arrow keys on the grid; the loop reads one row of W with K = 128.

Inside a CTA every thread has an index, and so does its warp and its place in the warp. The figure lays one CTA out as rows of warps; hovering a thread shows how its three indices relate.

Field Result Same in
ctx.thread The thread’s index in the CTA, 0 to ctx.threads - 1. nothing: it differs between lanes
ctx.warp ctx.thread / ctx.lanes, 0 to ctx.warps - 1. every lane of a warp
ctx.lane The thread’s index in its warp, 0 to ctx.lanes - 1. nothing: it differs between lanes

They are related by

ctx.thread = ctx.warp * ctx.lanes + ctx.lane

so with @launch(256) on NVIDIA, thread 113 is lane 17 of warp 3, and lane 0 of warp 5 is thread 160.

These are the only sources of per-thread values in Phobos. Everything a kernel computes is the same in every thread until one of them flows into it. The compiler tracks how far each value may differ, and the next page, Uniform, warp-uniform and varying, lists the rules that follow. The spec covers the fields under Indexing and synchronization and their consequences under Per-thread values.

The lane-strided loop

Section titled “The lane-strided loop”

The most common use of ctx.lane is a loop that starts at the lane and steps by the warp width. Lane 0 takes elements 0, 32, 64 and so on, lane 1 takes 1, 33, 65, and together the 32 lanes cover the row once:

var acc: f32 = 0.0
for k in range(ctx.lane, K, ctx.lanes) {
  acc += W[col, k] * A[0, k]
}

On each iteration the 32 lanes read 32 neighboring elements, so a warp’s loads fall on as few memory lines as possible. The loop variable k differs between lanes because its start does, and acc is a scalar var, which lives in each thread’s registers: after the loop every lane holds its own partial sum.

The same shape works over the whole CTA with ctx.thread and ctx.threads. This kernel sums each row of X, one CTA per row; each thread adds every 256th element, and a collective over the CTA folds the 256 partial sums:

@launch(256)
kernel row_sum(X: tensor<f32>[R, K], O: tensor<f32>[R, 1]) {
  let r = ctx.block.x
  var acc: f32 = 0.0
  for k in range(ctx.thread, K, ctx.threads) {
    acc += X[r, k]
  }
  let total = sum<over ctx.thread>(acc)
  if ctx.thread == 0 {
    O[r, 0] = total
  }
}

total is the same in every thread, but only thread 0 stores it, so the tensor sees one store instead of 256 identical ones. The collective is explained on Collectives and shuffles.

One warp per output

Section titled “One warp per output”

ctx.warp is the index to use when each warp takes one piece of work. A matvec with one warp per output column computes its column as

let col = ctx.block.x * ctx.warps + ctx.warp

so a CTA of 8 warps handles 8 consecutive columns and the grid needs ceil(N / 8) CTAs. col is warp-uniform: every lane of a warp agrees on it, which is exactly what a lane collective over that warp needs. An if col < N around the work skips the warps past the last column without splitting any warp.

A thread’s own run

Section titled “A thread’s own run”

A per-thread loop can also read several neighboring elements at once. A subscript with a constant :+ length inside sum(..), elem(..), word(..) or store(..) is the thread’s own run of elements, loaded with one vector access when the alignment can be proven:

for k in range(ctx.lane * 16, K, ctx.lanes * 16) {
  acc += f32(sum(i32(A[0, k :+ 16]) * i32(W[col, k :+ 16])))
}

Each lane now owns 16 consecutive i8 values per step instead of one. The details, including when this becomes a dp4a per four bytes, are on Collectives and shuffles.

What the indices do not change

Section titled “What the indices do not change”

A tile operation in a kernel that also uses the indices is unchanged: it still runs once per CTA, spread over all threads by the compiler. What changes is where it may appear. A tile operation inside if ctx.lane == 0 { .. } would be reached by only some threads, and the compiler rejects it. The rules are on the next page.