[Mlir-commits] [mlir] [mlir][scf] Preserve forall shared output order in promote (PR #222144)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-scf

Author: hanbeom (ParkHanbum)

<details>
<summary>Changes</summary>

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


---
Full diff: https://github.com/llvm/llvm-project/pull/222144.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+32-18) 
- (modified) mlir/test/Dialect/SCF/canonicalize.mlir (+42) 


``````````diff
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 e89ad45867a57..04a5939cdf694 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 %[[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>) {
+    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 %[[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>) {
+  %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>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/222144