[Mlir-commits] [mlir] [mlir][Linalg] Fix output block arg remapping in DecomposeLinalgOp (PR #238168)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Nov 2 16:53:16 PST 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Cheng-Chieh Liu (jay87406)
<details>
<summary>Changes</summary>
`DecomposeLinalgOp` splits a `linalg.generic` into a peeled op for the
first statement and a residual op for the rest. The residual op takes
the results of the peeled op as extra inputs, placed between the
original inputs and outputs, but uses of the original block arguments
are remapped by position. An output block argument used in the
residual body therefore ends up on one of the extra inputs.
For an f32 body like `%0 = arith.mulf %in, %in` followed by
`%1 = arith.addf %0, %out`, the residual op computes `addf(%0, %0)`
instead of `addf(%0, %out)`. The IR verifies and silently computes
wrong values. When the types differ the IR becomes invalid instead,
e.g. the fold assertion in #<!-- -->237535. The existing
`destination_passing_style` test encoded the same wrong mapping: its
residual `mulf` read the first output instead of the second one.
Shift the output block argument indices by the number of extra inputs,
correct the `destination_passing_style` checks, and add a regression
test.
Fixes #<!-- -->237535
---
Full diff: https://github.com/llvm/llvm-project/pull/238168.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp (+8-1)
- (modified) mlir/test/Dialect/Linalg/decompose-ops.mlir (+53-10)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
index 2fb5d11f95a371..6b38e5fe58572b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
@@ -324,11 +324,18 @@ DecomposeLinalgOp::matchAndRewrite(GenericOp genericOp,
/// In the split operations, replace block arguments uses that refer to
/// original operation to the block arguments of the newly created operation.
+ /// The residual operation has the results of the peeled operation as extra
+ /// inputs, placed before the original outputs, so the block arguments of the
+ /// original outputs are shifted by the number of extra inputs.
unsigned origNumInputs = genericOp.getNumDpsInputs();
+ unsigned numExtraInputs = residualGenericOp.getNumDpsInputs() - origNumInputs;
for (const auto &inputBlockArg :
llvm::enumerate(genericOp.getBody()->getArguments())) {
+ unsigned residualArgNum = inputBlockArg.index();
+ if (residualArgNum >= origNumInputs)
+ residualArgNum += numExtraInputs;
Value residualOpReplacementArg =
- residualGenericOpBody->getArgument(inputBlockArg.index());
+ residualGenericOpBody->getArgument(residualArgNum);
rewriter.replaceUsesWithIf(
inputBlockArg.value(), residualOpReplacementArg, [&](OpOperand &use) {
return use.getOwner()->getBlock() == residualGenericOpBody;
diff --git a/mlir/test/Dialect/Linalg/decompose-ops.mlir b/mlir/test/Dialect/Linalg/decompose-ops.mlir
index 43bd42528f34f8..ae3f491f79083a 100644
--- a/mlir/test/Dialect/Linalg/decompose-ops.mlir
+++ b/mlir/test/Dialect/Linalg/decompose-ops.mlir
@@ -380,7 +380,7 @@ func.func @destination_passing_style(
// CHECK-SAME: %[[ARG11:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG12:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG13:[a-zA-Z0-9_]+]]: f32
-// CHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG10]], %[[ARG12]]
+// CHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG10]], %[[ARG13]]
// CHECK-NEXT: linalg.yield %[[ARG11]], %[[S2]]
// CHECK: return %[[GENERIC1]]#0, %[[GENERIC2]]#1
@@ -403,19 +403,17 @@ func.func @destination_passing_style(
// CANONICALIZECHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.addf %[[ARG4]], %[[ARG5]]
// CANONICALIZECHECK-NEXT: linalg.yield %[[S1]]
-// CANONICALIZECHECK: %[[GENERIC2:.+]]:2 = linalg.generic
-// CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP1]], #[[MAP1]], #[[MAP3]]]
+// CANONICALIZECHECK: %[[GENERIC2:.+]] = linalg.generic
+// CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]]
// CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"]
-// CANONICALIZECHECK-SAME: ins(%[[ARG1]], %[[GENERIC1]] :
-// CANONICALIZECHECK-SAME: outs(%[[ARG2]], %[[ARG3]] :
+// CANONICALIZECHECK-SAME: ins(%[[ARG1]] :
+// CANONICALIZECHECK-SAME: outs(%[[ARG3]] :
// CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32
-// CANONICALIZECHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: f32
-// CANONICALIZECHECK-SAME: %[[ARG7:[a-zA-Z0-9_]+]]: f32
-// CANONICALIZECHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG4]], %[[ARG6]]
-// CANONICALIZECHECK-NEXT: linalg.yield %[[ARG5]], %[[S2]]
-// CANONICALIZECHECK: return %[[GENERIC1]], %[[GENERIC2]]#1
+// CANONICALIZECHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG4]], %[[ARG5]]
+// CANONICALIZECHECK-NEXT: linalg.yield %[[S2]]
+// CANONICALIZECHECK: return %[[GENERIC1]], %[[GENERIC2]]
// -----
@@ -452,3 +450,48 @@ func.func @first_statement_no_result(%arg0 : tensor<27xi32>) -> tensor<27xi32> {
// CANONICALIZECHECK: arith.addi
// CANONICALIZECHECK: linalg.yield
// CANONICALIZECHECK: return %[[GENERIC]]
+
+// -----
+
+// The residual op gets the results of the peeled op as extra inputs, placed
+// before the original outputs, so uses of an original output block argument
+// must be remapped past those extra inputs.
+#map0 = affine_map<(d0) -> (d0)>
+func.func @use_output_block_arg(%arg0 : tensor<4xf32>, %arg1 : tensor<4xf32>)
+ -> tensor<4xf32> {
+ %r = linalg.generic {
+ indexing_maps = [#map0, #map0],
+ iterator_types = ["parallel"]}
+ ins(%arg0 : tensor<4xf32>)
+ outs(%arg1 : tensor<4xf32>) {
+ ^bb0(%b0 : f32, %b1 : f32):
+ %0 = arith.mulf %b0, %b0 : f32
+ %1 = arith.addf %0, %b1 : f32
+ linalg.yield %1 : f32
+ } -> tensor<4xf32>
+ return %r : tensor<4xf32>
+}
+// CHECK: func.func @use_output_block_arg(
+// CHECK: %[[GENERIC0:.+]]:2 = linalg.generic
+// CHECK: arith.mulf
+// CHECK: %[[GENERIC1:.+]] = linalg.generic
+// CHECK-SAME: ins(%{{.+}}, %[[GENERIC0]]#1 :
+// CHECK-NEXT: ^bb0(
+// CHECK-SAME: %[[B0:[a-zA-Z0-9_]+]]: f32
+// CHECK-SAME: %[[B1:[a-zA-Z0-9_]+]]: f32
+// CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32
+// CHECK-NEXT: %[[S:.+]] = arith.addf %[[B1]], %[[B2]] : f32
+// CHECK-NEXT: linalg.yield %[[S]]
+// CHECK: return %[[GENERIC1]]
+
+// CANONICALIZECHECK: func.func @use_output_block_arg(
+// CANONICALIZECHECK: %[[GENERIC0:.+]] = linalg.generic
+// CANONICALIZECHECK: arith.mulf
+// CANONICALIZECHECK: %[[GENERIC1:.+]] = linalg.generic
+// CANONICALIZECHECK-SAME: ins(%[[GENERIC0]] :
+// CANONICALIZECHECK-NEXT: ^bb0(
+// CANONICALIZECHECK-SAME: %[[B0:[a-zA-Z0-9_]+]]: f32
+// CANONICALIZECHECK-SAME: %[[B1:[a-zA-Z0-9_]+]]: f32
+// CANONICALIZECHECK-NEXT: %[[S:.+]] = arith.addf %[[B0]], %[[B1]] : f32
+// CANONICALIZECHECK-NEXT: linalg.yield %[[S]]
+// CANONICALIZECHECK: return %[[GENERIC1]]
``````````
</details>
https://github.com/llvm/llvm-project/pull/238168
More information about the Mlir-commits
mailing list