scf.forall tiling

A loop with a destination

scf.forall is the parallel tile loop. On tensors it carries the result through shared_outs, and the terminator says how each iteration writes into it:

func.func @matmul_forall(%A: tensor<4x4xf32>, %B: tensor<4x4xf32>, %C: tensor<4x4xf32>)
    -> tensor<4x4xf32> {
  %0 = scf.forall (%i, %j) = (0, 0) to (4, 4) step (2, 2)
       shared_outs(%c = %C) -> (tensor<4x4xf32>) {
    %sA = tensor.extract_slice %A[%i, 0] [2, 4] [1, 1]
        : tensor<4x4xf32> to tensor<2x4xf32>
    %sB = tensor.extract_slice %B[0, %j] [4, 2] [1, 1]
        : tensor<4x4xf32> to tensor<4x2xf32>
    %sC = tensor.extract_slice %c[%i, %j] [2, 2] [1, 1]
        : tensor<4x4xf32> to tensor<2x2xf32>
    %t = linalg.matmul ins(%sA, %sB : tensor<2x4xf32>, tensor<4x2xf32>)
                       outs(%sC : tensor<2x2xf32>) -> tensor<2x2xf32>
    scf.forall.in_parallel {
      tensor.parallel_insert_slice %t into %c[%i, %j] [2, 2] [1, 1]
          : tensor<2x2xf32> into tensor<4x4xf32>
    }
  }
  return %0 : tensor<4x4xf32>
}

With explicit bounds and steps, %i and %j are element offsets, which is why they can be used as slice offsets directly. The other spelling, in (2, 2), is normalised: lower bounds zero and steps one, so the induction variables are tile indices and every offset needs iv * tile. Mixing the two shifts every tile by one tile size.

shared_outs binds the destination tensor to a block argument, and the loop yields one result per shared out. It is a synchronisation point as well, so nothing after it observes a partial tile.

mapping records the processing units the op is meant for. It is an attribute, not a launch: a lowering interprets it or drops it.

Reductions are the hard case

Tiling M and N needs no communication: each iteration owns a distinct tile of C. Tiling K does, because every iteration wants the same accumulator. The shape that works is a wider shared out with one partial-sum slice per iteration, plus a combining op after the loop. A split-reduction transform builds exactly that.

Things that break

Writing the result into %C instead of the block argument %c is the common mistake, and it fails loudly because the parallel insert must target a shared-out argument. A missing scf.forall.in_parallel fails verification instead of silently dropping the result. And a forall whose iteration count does not match the induction-variable list is a verifier error, not a runtime truncation, so read the diagnostic rather than the IR.