[Mlir-commits] [mlir] [mlir][Linalg] Fix output block arg remapping in DecomposeLinalgOp (PR #238168)
Cheng-Chieh Liu
llvmlistbot at llvm.org
Mon Nov 2 16:52:26 PST 2026
https://github.com/jay87406 created https://github.com/llvm/llvm-project/pull/238168
`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
>From be99510ce1124d3c2c6595e5837dda25ce6eb834 Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Mon, 2 Nov 2026 22:56:37 +0000
Subject: [PATCH] [mlir][Linalg] Fix output block arg remapping in
DecomposeLinalgOp
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
---
.../Linalg/Transforms/DecomposeLinalgOps.cpp | 9 ++-
mlir/test/Dialect/Linalg/decompose-ops.mlir | 63 ++++++++++++++++---
2 files changed, 61 insertions(+), 11 deletions(-)
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]]
More information about the Mlir-commits
mailing list