The grid and the pipeline

The kernel path teaches the grid as a software pipeline, at /s/pallas, and that teaching is true of the TPU backend. It is not what pallas_call promises. The design document says so in a single sentence, and the sentence is worth carrying around.

pallas_call also provides no guarantees on the order of loop iterations over the iteration space, just that every member of the iteration space will be looped over.

That line is from the Pallas design document, and the paragraph after it says where the order actually comes from: compilers like Triton and MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → have more specific operational semantics associated with the grid. Portability lives in the front end. Order lives in the backend.

On TPU the backend is specific indeed. The reference calls TPUs highly sequential machines and says the grid is generally not processed in parallel but sequentially, in lexicographic order. So a kernel that depends on order is correct here and is not portable, and you were told which of those two you were buying.

dimension_semantics is often described as a hint about your intent. It is closer to permission. Some TPU chips carry two TensorCores behind one device, each with its own VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu →, VREGs, SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu →, SREGs and compute units, sharing HBM. Using both means breaking the sequential grid guarantee, and the annotation is how you say which axis may be broken.

The rule of thumb the reference gives is mechanical: an axis is parallel unless the output window does not vary along it. Which is why the annotation always reads as some parallel axes followed by some arbitrary ones, and why a reduction axis is never parallel. Mark it wrong and you get a wrong answer rather than an error, because nothing about the shapes became invalid.

Two caveats are worth carrying. Partitioning across two cores often gives close to 2x, and can give much less when the per-step cost varies, since one core can be handed all the expensive steps and the other idles waiting. And the guide notes megacore as currently a v4 and v5p feature, where supplying the annotation elsewhere is a no-op, while omitting it entirely leaves a second core unused. This site's own matmul retune moved the annotation and the block shape in the same step, so that measurement does not isolate either one.

pl.pallas_call(
    add_matrices_kernel,
    out_shape=jax.ShapeDtypeStruct.like(x),
    in_specs=[block_spec, block_spec],
    out_specs=block_spec,
    grid=(2,),
    compiler_params=pltpu.CompilerParams(
        dimension_semantics=("parallel",)),
)(x, y)

The derivation in the software pipelining tutorial is short enough to hold in your head, and holding it is what turns "the grid is a pipeline" from a slogan into something you can predict from. Start with a loop that copies in, computes, copies out. Split each copy into a start and a wait so asynchrony is expressible. Give the staging buffer two slots so iteration i can compute out of one while i+1 fills the other. Push each copy-out wait as late as it can go. Re-roll the loop.

What falls out has a prologue that starts the first copy, a steady state where every iteration issues the next input copy before waiting on the current one, and an epilogue that drains the last write. The alternation between slots is the i % 2 you see in the pseudocode, and the reason the default buffer count is two.

# Prologue
copy_in_start(A[0], X[0])

# Main loop
for i in range(N):
  cur_slot = i % 2
  next_slot = (i + 1) % 2

  if i+1 < N:
    copy_in_start(A[i+1], X[next_slot])

  copy_in_wait(X[cur_slot])
  Y[cur_slot] = X[cur_slot] + 1
  copy_out_start(Y[cur_slot], A[i])

  if i > 0:
    copy_out_wait(Y[next_slot])

# Epilogue
copy_out_wait(Y[1])

The design document says BlockSpecs can be converted into pipeline schedules, and on this repo you can read the conversion instead of taking it on faith. Lowering the tiled matmul with debug=True puts the whole schedule into one attribute dictionary on the MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → function: the grid arrives as iteration_bounds, each spec arrives as a window_params entry carrying its window_bounds and a transform_indices function, and dimension_semantics rides in beside them.

Everything you wrote outside the kernel body is in that line, and nothing else is. The body became instructions; the carve became metadata that the pipeline emitter reads.

func.func @main(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: memref<256x256xbf16, #tpu.memory_space<vmem>>, %arg4: memref<256x256xbf16, #tpu.memory_space<vmem>>, %arg5: memref<256x256xbf16, #tpu.memory_space<vmem>>) attributes {dimension_semantics = [#tpu.dimension_semantics<arbitrary>, #tpu.dimension_semantics<arbitrary>, #tpu.dimension_semantics<arbitrary>], iteration_bounds = array<i64: 2, 2, 2>, scalar_prefetch = 0 : i64, scratch_operands = 0 : i64, window_params = [{transform_indices = @transform_0, window_bounds = array<i64: 256, 256>}, {transform_indices = @transform_1, window_bounds = array<i64: 256, 256>}, {transform_indices = @transform_2, window_bounds = array<i64: 256, 256>}]}

module @matmul_kernel {
  func.func @main(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: memref<256x256xbf16, #tpu.memory_space<vmem>>, %arg4: memref<256x256xbf16, #tpu.memory_space<vmem>>, %arg5: memref<256x256xbf16, #tpu.memory_space<vmem>>) attributes {dimension_semantics = [#tpu.dimension_semantics<arbitrary>, #tpu.dimension_semantics<arbitrary>, #tpu.dimension_semantics<arbitrary>], iteration_bounds = array<i64: 2, 2, 2>, scalar_prefetch = 0 : i64, scratch_operands = 0 : i64, window_params = [{transform_indices = @transform_0, window_bounds = array<i64: 256, 256>}, {transform_indices = @transform_1, window_bounds = array<i64: 256, 256>}, {transform_indices = @transform_2, window_bounds = array<i64: 256, 256>}]} {
    %c0_i32 = arith.constant 0 : i32
    %0 = arith.cmpi eq, %arg2, %c0_i32 : i32
    %1 = arith.extui %0 : i1 to i32
    %c0_i32_0 = arith.constant 0 : i32
    %2 = arith.cmpi ne, %1, %c0_i32_0 : i32
    scf.if %2 {
      %cst_8 = arith.constant 0.000000e+00 : bf16
      %10 = vector.broadcast %cst_8 : bf16 to vector<256x256xbf16>
      %c0_9 = arith.constant 0 : index
      %c0_10 = arith.constant 0 : index
      %11 = vector.load %arg5[%c0_9, %c0_10] : memref<256x256xbf16, #tpu.memory_space<vmem>>, vector<256x256xbf16>
      tpu.vector_store %arg5[%c0_9, %c0_10], %10 {strides = array<i32>} : memref<256x256xbf16, #tpu.memory_space<vmem>>, vector<256x256xbf16>,
    } else {
    }
    %c0 = arith.constant 0 : index
    %c0_1 = arith.constant 0 : index
    %3 = vector.load %arg5[%c0, %c0_1] : memref<256x256xbf16, #tpu.memory_space<vmem>>, vector<256x256xbf16>
    %c0_2 = arith.constant 0 : index
    %c0_3 = arith.constant 0 : index
    %4 = vector.load %arg3[%c0_2, %c0_3] : memref<256x256xbf16, #tpu.memory_space<vmem>>, vector<256x256xbf16>
    %c0_4 = arith.constant 0 : index
    %c0_5 = arith.constant 0 : index
    %5 = vector.load %arg4[%c0_4, %c0_5] : memref<256x256xbf16, #tpu.memory_space<vmem>>, vector<256x256xbf16>
    %cst = arith.constant dense<0.000000e+00> : vector<256x256xf32>
    %6 = tpu.matmul %4, %5, %cst {dimension_numbers = #tpu.dot_dimension_numbers<[1], [0], [0], [1], [0, 0, 1, 1], [], []>} : vector<256x256xbf16>, vector<256x256xbf16>, vector<256x256xf32> -> vector<256x256xf32>
    %7 = arith.truncf %6 : vector<256x256xf32> to vector<256x256xbf16>
    %8 = arith.addf %3, %7 : vector<256x256xbf16>
    %c0_6 = arith.constant 0 : index
    %c0_7 = arith.constant 0 : index
    %9 = vector.load %arg5[%c0_6, %c0_7] : memref<256x256xbf16, #tpu.memory_space<vmem>>, vector<256x256xbf16>
    tpu.vector_store %arg5[%c0_6, %c0_7], %8 {strides = array<i32>} : memref<256x256xbf16, #tpu.memory_space<vmem>>, vector<256x256xbf16>,
    return
  }
  func.func @transform_0(%arg0: i32, %arg1: i32, %arg2: i32) -> (i32, i32) {
    %c0_i32 = arith.constant 0 : i32
    return %arg0, %arg2 : i32, i32
  }
  func.func @transform_1(%arg0: i32, %arg1: i32, %arg2: i32) -> (i32, i32) {
    %c0_i32 = arith.constant 0 : i32
    return %arg2, %arg1 : i32, i32
  }
  func.func @transform_2(%arg0: i32, %arg1: i32, %arg2: i32) -> (i32, i32) {
    %c0_i32 = arith.constant 0 : i32
    return %arg0, %arg1 : i32, i32
  }
}

Buffer count is per argument. Pass pl.Buffered(buffer_count=n) as a BlockSpec's pipeline_mode and that input or output gets n slots instead of two, which is what you want when one operand's transfer is much longer than a single step of compute. The same object turns on lookahead prefetch with use_lookahead=True.

pltpu.emit_pipeline moves the whole mechanism inside the kernel body. Instead of one pipeline created at kernel entry, you construct pipelines where you need them, which is how a nested schedule gets written: an outer pipeline moving data between chips, an inner one moving it between HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu → and VMEM. It also carries the features that only exist at that level, dynamic block shapes among them.

A pipeline hides transfer under compute, so it stalls whenever a step has more transfer than compute to hide it behind. The site plays this as an instrument: EX·07 runs the same matmul as EX·02 with the overlap removed, and the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → sitting idle between loads is what the stall looks like when you can see it. The number that decides which animation you are living in is bytes per step over flops per step, against the chip's ridge.

Three other stalls have nothing to do with block size. A grid with very few steps pays its prologue and epilogue in full, since neither has a partner to hide behind, and a short grid is mostly prologue and epilogue. A grid reordered so that consecutive steps no longer touch the same input slice loses the skipped transfers that the previous order was quietly getting. And a partitioned axis whose steps cost wildly different amounts leaves one core idle while the other finishes.

The diagnosis is the same in every case and it lives in the profile: name the envelope, sum the compute ops inside it, and the difference is unhidden transfer time. The /l/tpu chapter walks that reading against a real capture, and the number it produces is the only honest answer to whether your pipeline is working.