FForge
← Labsmlir-linalg-02Linalg and transforms●●○○

Bufferization — where hidden copies come from

mlir-linalg/02_bufferization.mlirneeds: mlir-opt2 steps
0/2 steps
what this lab teaches

Tensors are values, memory is not. At some point every pipeline runs one-shot-bufferize, which assigns a buffer to every tensor value. For each op that WRITES a tensor (the `outs` operand of a linalg op, tensor.insert, ...) it asks one question: "Can I write into the SAME buffer as the destination operand, or would that clobber a value somebody still reads?" Yes -> in place, zero cost. No -> memref.alloc + memref.copy, then write. This is DESTINATION-PASSING STYLE (DPS): linalg ops take their output buffer as an operand (`outs(...)`) precisely so the bufferizer has a candidate to reuse. The frontend's job (triton-shared, torch-mlir) is to pick destinations that CAN be reused. When one can't, you pay a copy of the whole tensor, often inside a hot loop. A surprise memref.copy is the first thing to grep for when a kernel is slower than its FLOP count says. a vendor MLIR backend even has passes just to remove leftovers (copy canonicalization, buffer forwarding).

evidence to leave with
The final IR grepped for alloc/copy, with each remaining buffer explained or eliminated
walk it, one step at a time
run this
$ mlir-opt 02_bufferization.mlir --one-shot-bufferize="bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map" | grep -E 'func.func|memref.copy|memref.alloc'
what it printed on the machine that wrote this lab
  func.func @inplace(%arg0: memref<8xf32>) -> memref<8xf32> {
  func.func @needs_copy(%arg0: memref<8xf32>) -> (memref<8xf32>, memref<8xf32>) {
    %alloc = memref.alloc() alignment = 64 : memref<8xf32>
    memref.copy %arg0, %alloc : memref<8xf32> to memref<8xf32>
  func.func @fresh_dest(%arg0: memref<8xf32>) -> (memref<8xf32>, memref<8xf32>) {
    %alloc = memref.alloc() alignment = 64 : memref<8xf32>

(A) @inplace: no alloc, no copy. The input's buffer is updated directly. (B) @needs_copy: alloc + COPY, because %t is returned after being used as a destination; writing in place would change the second result. (C) @fresh_dest: alloc, no copy. `tensor.empty` says "I don't care about the contents", so the fresh buffer needs no initialisation. Same semantics as (B), one full memory pass cheaper.

run this
$ mlir-opt 02_bufferization.mlir --one-shot-bufferize="bufferize-function-boundaries test-analysis-only print-conflicts" 2>&1 | grep C_0
what it printed on the machine that wrote this lab
  func.func @needs_copy(%arg0: tensor<8xf32> {bufferization.access = "read"}) -> (tensor<8xf32>, tensor<8xf32>) attributes {"C_0[DEF: bbArg 0]"} {
    %0 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%arg0 : tensor<8xf32>) attrs =  {"C_0[CONFL-WRITE: 0]", __inplace_operands_attr__ = ["false"]} {
    return {"C_0[READ: 1]", __equivalent_func_args__ = [-1, 0], __inplace_operands_attr__ = ["true", "true"]} %0, %arg0 : tensor<8xf32>, tensor<8xf32>
`test-analysis-only` runs the decision logic without rewriting and
annotates every op. __inplace_operands_attr__ = ["false"] marks the operand
that got a copy. The C_0 tags name the conflict: DEF = the buffer comes from
bbArg 0 (%t), CONFL-WRITE = the generic writes it, READ: 1 = `return`
operand #1 still reads the OLD value.
That's your whole explanation, from the compiler, in one line.
take it further

1. In (B), return only %r. Does the copy disappear? 2. Put (A)'s generic inside an scf.for with iter_args. Does it stay in place across iterations? (Read about "equivalent" buffers in mlir/docs/Bufferization.md in llvm-project.) 3. Find where a vendor MLIR backend runs bufferization. grep -rn bufferize a production backend’s lib/ ============================================================

read the source · mlir-linalg/02_bufferization.mlir
#id = affine_map<(i) -> (i)> // (A) in place: %t is dead after the generic, so its buffer can be reused. func.func @inplace(%t: tensor<8xf32>) -> tensor<8xf32> { %r = linalg.generic {indexing_maps = [#id], iterator_types = ["parallel"]} outs(%t : tensor<8xf32>) { ^bb0(%o: f32): %d = arith.addf %o, %o : f32 linalg.yield %d : f32 } -> tensor<8xf32> return %r : tensor<8xf32> } // (B) NOT in place: %t is read AGAIN after being "overwritten" -> copy needed. func.func @needs_copy(%t: tensor<8xf32>) -> (tensor<8xf32>, tensor<8xf32>) { %r = linalg.generic {indexing_maps = [#id], iterator_types = ["parallel"]} outs(%t : tensor<8xf32>) { ^bb0(%o: f32): %d = arith.addf %o, %o : f32 linalg.yield %d : f32 } -> tensor<8xf32> return %r, %t : tensor<8xf32>, tensor<8xf32> } // (C) destination-passing style done right: a fresh tensor.empty as destination. func.func @fresh_dest(%t: tensor<8xf32>) -> (tensor<8xf32>, tensor<8xf32>) { %e = tensor.empty() : tensor<8xf32> %r = linalg.generic {indexing_maps = [#id, #id], iterator_types = ["parallel"]} ins(%t : tensor<8xf32>) outs(%e : tensor<8xf32>) { ^bb0(%i: f32, %o: f32): %d = arith.addf %i, %i : f32 linalg.yield %d : f32 } -> tensor<8xf32> return %r, %t : tensor<8xf32>, tensor<8xf32> }
◇
Finish the lab: write the evidence note

A lab is done when you can say what it showed on your machine — not when you have read it.