[Mlir-commits] [mlir] [mlir][scf] Preserve forall shared output order in promote (PR #222144)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Sep 8 14:02:34 PDT 2026
https://github.com/ParkHanbum created https://github.com/llvm/llvm-project/pull/222144
Build promoted results in shared output order instead of the textual
order of `scf.forall.in_parallel` operations.
This preserves the correct result mapping when the orders differ.
Fixes #222143
Realted https://github.com/llvm/llvm-project/pull/189162/changes/a89111e04967c203488fa3472e06b74dc46570ca
Assisted-by: Codex
>From aabd8c7b0c1b547094cc19e01b076d9316ac6e9d Mon Sep 17 00:00:00 2001
From: hanbeom <kese111 at gmail.com>
Date: Wed, 9 Sep 2026 05:51:15 +0900
Subject: [PATCH 1/2] [mlir][scf] Add test for reordered forall results
---
mlir/test/Dialect/SCF/canonicalize.mlir | 42 +++++++++++++++++++++++++
1 file changed, 42 insertions(+)
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index e89ad45867a57..42ca10e91b1b2 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2361,3 +2361,45 @@ func.func @fold_tensor_cast_into_forall_non_sequential_writes(
// %0#0 contains %arg1 data; %0#1 contains %arg0 data.
return %0#0, %0#1 : tensor<?x32xf32>, tensor<?x32xf32>
}
+
+// -----
+
+// The order of insertions is the reverse of the shared output order.
+// CHECK-LABEL: func.func @promote_reversed_outputs(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<2xf32>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<2xf32>,
+// CHECK-SAME: %[[ARG2:.*]]: tensor<2xf32>,
+// CHECK-SAME: %[[ARG3:.*]]: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
+// CHECK: return %[[ARG1]], %[[ARG0]] : tensor<2xf32>, tensor<2xf32>
+// CHECK: }
+func.func @promote_reversed_outputs(%a: tensor<2xf32>, %b: tensor<2xf32>, %init_a: tensor<2xf32>, %init_b: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
+ %r:2 = scf.forall (%i) in (1) shared_outs(%x = %init_a, %y = %init_b) -> (tensor<2xf32>, tensor<2xf32>) {
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %b into %y[0] [2] [1] : tensor<2xf32> into tensor<2xf32>
+ tensor.parallel_insert_slice %a into %x[0] [2] [1] : tensor<2xf32> into tensor<2xf32>
+ }
+ }
+ return %r#0, %r#1 : tensor<2xf32>, tensor<2xf32>
+}
+
+// -----
+
+// Partial insertions must preserve both the destination and result association.
+// CHECK-LABEL: func.func @promote_partial_outputs(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<2xf32>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<2xf32>,
+// CHECK-SAME: %[[ARG2:.*]]: tensor<4xf32>,
+// CHECK-SAME: %[[ARG3:.*]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+// CHECK: %[[INSERT_SLICE_0:.*]] = tensor.insert_slice %[[ARG1]] into %[[ARG3]][2] [2] [1] : tensor<2xf32> into tensor<4xf32>
+// CHECK: %[[INSERT_SLICE_1:.*]] = tensor.insert_slice %[[ARG0]] into %[[ARG2]][0] [2] [1] : tensor<2xf32> into tensor<4xf32>
+// CHECK: return %[[INSERT_SLICE_0]], %[[INSERT_SLICE_1]] : tensor<4xf32>, tensor<4xf32>
+// CHECK: }
+func.func @promote_partial_outputs(%a: tensor<2xf32>, %b: tensor<2xf32>, %init_a: tensor<4xf32>, %init_b: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+ %r:2 = scf.forall (%i) in (1) shared_outs(%x = %init_a, %y = %init_b) -> (tensor<4xf32>, tensor<4xf32>) {
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %b into %y[2] [2] [1] : tensor<2xf32> into tensor<4xf32>
+ tensor.parallel_insert_slice %a into %x[0] [2] [1] : tensor<2xf32> into tensor<4xf32>
+ }
+ }
+ return %r#0, %r#1 : tensor<4xf32>, tensor<4xf32>
+}
>From ac4da4c105e0784eaeb2e9e21012535e284b2994 Mon Sep 17 00:00:00 2001
From: hanbeom <kese111 at gmail.com>
Date: Wed, 9 Sep 2026 05:55:00 +0900
Subject: [PATCH 2/2] [mlir][scf] Preserve forall shared output order in
promote
Build promoted results in shared output order instead of the textual
order of `scf.forall.in_parallel` operations.
This preserves the correct result mapping when the orders differ.
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 50 ++++++++++++++++---------
mlir/test/Dialect/SCF/canonicalize.mlir | 6 +--
2 files changed, 35 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 3388fa490f996..1596898bb2dda 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -760,6 +760,24 @@ void mlir::scf::promote(RewriterBase &rewriter, scf::ForallOp forallOp) {
bbArgReplacements.append(forallOp.getOutputs().begin(),
forallOp.getOutputs().end());
+ // Collect insert ops in shared output order before inlining replaces the
+ // output block arguments. Keep a separate list for each output.
+ SmallVector<SmallVector<tensor::ParallelInsertSliceOp, 1>> yieldingOps;
+ yieldingOps.reserve(forallOp.getNumResults());
+
+ for (BlockArgument outputArg : forallOp.getRegionIterArgs()) {
+ auto &outputOps = yieldingOps.emplace_back();
+
+ for (Operation *op : forallOp.getCombiningOps(outputArg)) {
+ auto insertOp = dyn_cast<tensor::ParallelInsertSliceOp>(op);
+ if (!insertOp || insertOp->getParentOp() != terminator.getOperation() ||
+ insertOp.getDest() != outputArg)
+ continue;
+
+ outputOps.push_back(insertOp);
+ }
+ }
+
// Move the loop body operations to the loop's containing block.
rewriter.inlineBlockBefore(forallOp.getBody(), forallOp->getBlock(),
forallOp->getIterator(), bbArgReplacements);
@@ -768,26 +786,22 @@ void mlir::scf::promote(RewriterBase &rewriter, scf::ForallOp forallOp) {
rewriter.setInsertionPointAfter(forallOp);
SmallVector<Value> results;
results.reserve(forallOp.getResults().size());
- for (auto &yieldingOp : terminator.getYieldingOps()) {
- auto parallelInsertSliceOp =
- dyn_cast<tensor::ParallelInsertSliceOp>(yieldingOp);
- if (!parallelInsertSliceOp)
- continue;
-
- Value dst = parallelInsertSliceOp.getDest();
- Value src = parallelInsertSliceOp.getSource();
- if (llvm::isa<TensorType>(src.getType())) {
- results.push_back(tensor::InsertSliceOp::create(
- rewriter, forallOp.getLoc(), dst.getType(), src, dst,
- parallelInsertSliceOp.getOffsets(), parallelInsertSliceOp.getSizes(),
- parallelInsertSliceOp.getStrides(),
- parallelInsertSliceOp.getStaticOffsets(),
- parallelInsertSliceOp.getStaticSizes(),
- parallelInsertSliceOp.getStaticStrides()));
- } else {
- llvm_unreachable("unsupported terminator");
+ for (auto [output, insertOps] :
+ llvm::zip_equal(forallOp.getOutputs(), yieldingOps)) {
+ Value result = output;
+
+ for (tensor::ParallelInsertSliceOp insertOp : insertOps) {
+ result = tensor::InsertSliceOp::create(
+ rewriter, forallOp.getLoc(), result.getType(),
+ insertOp.getSource(), result,
+ insertOp.getOffsets(), insertOp.getSizes(), insertOp.getStrides(),
+ insertOp.getStaticOffsets(), insertOp.getStaticSizes(),
+ insertOp.getStaticStrides());
}
+
+ results.push_back(result);
}
+
rewriter.replaceAllUsesWith(forallOp.getResults(), results);
// Erase the old terminator and the loop.
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index 42ca10e91b1b2..04a5939cdf694 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2370,7 +2370,7 @@ func.func @fold_tensor_cast_into_forall_non_sequential_writes(
// CHECK-SAME: %[[ARG1:.*]]: tensor<2xf32>,
// CHECK-SAME: %[[ARG2:.*]]: tensor<2xf32>,
// CHECK-SAME: %[[ARG3:.*]]: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
-// CHECK: return %[[ARG1]], %[[ARG0]] : tensor<2xf32>, tensor<2xf32>
+// CHECK: return %[[ARG0]], %[[ARG1]] : tensor<2xf32>, tensor<2xf32>
// CHECK: }
func.func @promote_reversed_outputs(%a: tensor<2xf32>, %b: tensor<2xf32>, %init_a: tensor<2xf32>, %init_b: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
%r:2 = scf.forall (%i) in (1) shared_outs(%x = %init_a, %y = %init_b) -> (tensor<2xf32>, tensor<2xf32>) {
@@ -2390,8 +2390,8 @@ func.func @promote_reversed_outputs(%a: tensor<2xf32>, %b: tensor<2xf32>, %init_
// CHECK-SAME: %[[ARG1:.*]]: tensor<2xf32>,
// CHECK-SAME: %[[ARG2:.*]]: tensor<4xf32>,
// CHECK-SAME: %[[ARG3:.*]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
-// CHECK: %[[INSERT_SLICE_0:.*]] = tensor.insert_slice %[[ARG1]] into %[[ARG3]][2] [2] [1] : tensor<2xf32> into tensor<4xf32>
-// CHECK: %[[INSERT_SLICE_1:.*]] = tensor.insert_slice %[[ARG0]] into %[[ARG2]][0] [2] [1] : tensor<2xf32> into tensor<4xf32>
+// CHECK: %[[INSERT_SLICE_0:.*]] = tensor.insert_slice %[[ARG0]] into %[[ARG2]][0] [2] [1] : tensor<2xf32> into tensor<4xf32>
+// CHECK: %[[INSERT_SLICE_1:.*]] = tensor.insert_slice %[[ARG1]] into %[[ARG3]][2] [2] [1] : tensor<2xf32> into tensor<4xf32>
// CHECK: return %[[INSERT_SLICE_0]], %[[INSERT_SLICE_1]] : tensor<4xf32>, tensor<4xf32>
// CHECK: }
func.func @promote_partial_outputs(%a: tensor<2xf32>, %b: tensor<2xf32>, %init_a: tensor<4xf32>, %init_b: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
More information about the Mlir-commits
mailing list