[mlir][bug]`linalg.map` one-shot bufferization drops a read of `outs`
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Minimal reproducer
Save the following as `manual-min.mlir`. The `linalg.map` region reads `%out`,the element corresponding to its `outs` tensor, and writes it to an observable buffer. Therefore, the required result is `11`.
```mlir
module {
func.func @func1() {
%c0 = arith.constant 0 : index
%initial = arith.constant -214 : i16
%yielded = arith.constant 9363 : i16
%input = arith.constant dense<0> : tensor<1xi16>
%init = arith.constant dense<11> : tensor<1xi16>
%buffer = memref.alloc() : memref<1xi16>
memref.store %initial, %buffer[%c0] : memref<1xi16>
%mapped = linalg.map ins(%input : tensor<1xi16>) outs(%init : tensor<1xi16>)
(%in: i16, %out: i16) {
memref.store %out, %buffer[%c0] : memref<1xi16>
linalg.yield %yielded : i16
}
%result = memref.load %buffer[%c0] : memref<1xi16>
%extended = arith.extsi %result : i16 to i64
vector.print %extended : i64
return
}
}
```
## Complete verification
Run the following in the directory containing `manual-min.mlir`. Both paths start from that exact file and have the same lowering sequence. The only difference is the one extra pass at the start of the trigger path.
```bash
BUILD=/12T/qin/Fuzzing/MLIR_b3/llvm-project/build
# Control path: no generalization.
$BUILD/bin/mlir-opt manual-min.mlir \
--one-shot-bufferize="bufferize-function-boundaries copy-before-write unknown-type-conversion=identity-layout-map" \
--convert-linalg-to-loops --convert-scf-to-cf --convert-index-to-llvm \
--finalize-memref-to-llvm --convert-func-to-llvm --convert-vector-to-llvm \
--convert-arith-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts \
-o control.mlir
# Trigger path: the only additional pass.
$BUILD/bin/mlir-opt manual-min.mlir \
--linalg-morph-ops=named-to-generic \
--one-shot-bufferize="bufferize-function-boundaries copy-before-write unknown-type-conversion=identity-layout-map" \
--convert-linalg-to-loops --convert-scf-to-cf --convert-index-to-llvm \
--finalize-memref-to-llvm --convert-func-to-llvm --convert-vector-to-llvm \
--convert-arith-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts \
-o trigger.mlir
for kind in control trigger; do
$BUILD/bin/mlir-translate --mlir-to-llvmir "$kind.mlir" -o "$kind.ll"
sed 's/@func1/@main/' "$kind.ll" > "$kind.main.ll"
$BUILD/bin/clang -lm "$kind.main.ll" \
"$BUILD/lib/libmlir_c_runner_utils.so.23.0git" \
"$BUILD/lib/libmlir_float16_utils.so.23.0git" \
"$BUILD/lib/libmlir_runner_utils.so.23.0git" -lstdc++ -o "$kind.out"
printf '%s: ' "$kind"
LD_LIBRARY_PATH="$BUILD/lib" "./$kind.out"
done
```
Observed output on the environment below:
```text
control: 0
trigger: 11
```
The control path reads uninitialized memory, so its value can vary with the allocator and host. It must not replace the required initial value, `11`.
`--linalg-generalize-named-ops` (the deprecated spelling) produces the same trigger IR as `--linalg-morph-ops=named-to-generic`.
## Root cause
After one-shot bufferization, the control path allocates `%alloc_0` without initializing it, then the map region loads it as `%init`. The trigger path adds:
```mlir
memref.copy %constant_11, %alloc_0 : memref<1xi16> to memref<1xi16>
```
In `mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td`, `MapOp::payloadUsesValueFromOperand` returns `false` for every DPS init, even when the matching region block argument (`%out` here) has uses. This makes one-shot bufferization incorrectly conclude that the init tensor does
not need materialization.
The method should use `!getMatchingBlockArgument(opOperand).use_empty()` for init operands as well, or otherwise return `true` when the init block argument is read.
Contributor guide
Research direction
Start with mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td and inspect MapOp::payloadUsesValueFromOperand, then reproduce the issue with manual-min.mlir and the two mlir-opt paths provided. Add or update a regression test covering a read of the linalg.map outs block argument, and verify one-shot bufferization materializes the initial value so the trigger path observes 11.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100