@swizzle

An ldmatrix load of an int8 tensor-core fragment reads 16 bytes from each of eight consecutive rows of a shared tile. When the tile’s rows are a multiple of 128 bytes long, those eight 16-byte pieces start in the same bank, and the load serializes.

The compiler avoids that by default with a swizzled layout: each row’s 16-byte granules are permuted by the row’s index mod 8, so the eight rows of a load land in eight different banks. Unlike @padstage, this costs no shared memory. A tile that is swizzled is never padded.

@swizzle turns the default into an assertion, and @swizzle(off) turns it off for the kernel.

@swizzle
@swizzle(on)
@swizzle(off)

@swizzle and @swizzle(on) are the same. See the Attributes table in the spec.

Which tiles it reaches

A shared tile is swizzled when all of these hold:

• It is rank 2, with a static shape and an element of at most 4 bytes.
• Its rows are a multiple of 128 bytes.
• The scaled int8 dot, acc += f32(dot_t(a, w)) * s1 * s2, reads it, directly or through slices.
• Every other access to it can follow the permutation: a copy into it, a staged loop’s refills, and store or word runs that stay inside one 16-byte granule.

Anything else touching the tile, such as copying it out or reading one of its elements, leaves it unswizzled. The layout is invisible to the source: the tile keeps its shape, and every access the compiler allows is permuted for you.

A minimal kernel

@launch(128)
@aligned(M = 64, N = 64, K = 128)
kernel proj(A: tensor<i8>[M, K], AS: tensor<f32>[M, KB],
            W: tensor<i8>[N, K], WS: tensor<f32>[KB, N],
            C: tensor<f32>[M, N]) {
  let pm = ctx.block.x
  let pn = ctx.block.y
  var w: tile<i8>[64, 128]
  var acc: tile<f32>[64, 64] = 0.0
  for kt in range(0, K, 128) {
    var a = A[pm * 64 :+ 64, kt :+ 128]
    w = W[pn * 64 :+ 64, kt :+ 128]
    barrier()
    let a0 = a[0 :+ 64, 0 :+ 32]
    let w0 = w[0 :+ 64, 0 :+ 32]
    let da = AS[pm * 64 :+ 64, kt / 32 :+ 1]
    let ws = WS[kt / 32 :+ 1, pn * 64 :+ 64]
    acc += f32(dot_t(a0, w0)) * ws * da
  }
  C[pm * 64 :+ 64, pn * 64 :+ 64] = acc
}

Both a and w are [64, 128] i8, 128 bytes a row, and the scaled dot reads both through slices, so both are swizzled without being asked.

When it cannot

Under @swizzle, compilation fails if a tile the swizzle would help cannot take it, naming the tile and the access in the way, or if there is no such tile at all. Here a copy reads w out, which would see the permuted bytes:

@launch(128)
@swizzle
@aligned(M = 64, N = 64, K = 128)
kernel proj(A: tensor<i8>[M, K], AS: tensor<f32>[M, KB],
            W: tensor<i8>[N, K], WS: tensor<f32>[KB, N],
            C: tensor<f32>[M, N], E: tensor<i8>[N, K]) {
  let pm = ctx.block.x
  let pn = ctx.block.y
  var w: tile<i8>[64, 128]
  var acc: tile<f32>[64, 64] = 0.0
  for kt in range(0, K, 128) {
    w = W[pn * 64 :+ 64, kt :+ 128]
    barrier()
    let a0 = A[pm * 64 :+ 64, kt :+ 32]
    let w0 = w[0 :+ 64, 0 :+ 32]
    let da = AS[pm * 64 :+ 64, kt / 32 :+ 1]
    let ws = WS[kt / 32 :+ 1, pn * 64 :+ 64]
    acc += f32(dot_t(a0, w0)) * ws * da
    E[pn * 64 :+ 64, kt :+ 128] = w
  }
  C[pm * 64 :+ 64, pn * 64 :+ 64] = acc
}

`@swizzle` cannot swizzle tile `w`: something reads it that cannot apply the permutation

Without the attribute the same kernel compiles, with w laid out plainly.

Interactions

The swizzle replaces the padding @padstage would add to a staged tile, which matters most under @pipeline: double-buffered stages that fit 48 KB unpadded may not fit padded. The GGUF backend’s IQ2_XXS prompt projection, written in Phobos, relies on it to fit on sm_80 and later.