llvm / llvm/llvm-project

[mlir][bug]`linalg.map` one-shot bufferization drops a read of `outs`

Open
#213,241 2 comments 0 reactions 0 assignees View on GitHub
mlir
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.