[Mlir-commits] [mlir] [mlir][scf] Preserve reductions when collapsing parallel loops (PR #216538)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: crZhao (cr-zhao)

<details>
<summary>Changes</summary>

Fixes #<!-- -->216218.

`collapseParallelLoops` previously created the collapsed `scf.parallel`
without the original init values, erased the original `scf.reduce`, and then
erased the old loop without replacing its results. A parallel loop with a
reduction result therefore left live uses of the erased operation and crashed.

Preserve the original init values when creating the collapsed loop, move the
body together with its reduction terminator, and replace the old loop with the
new loop results through the rewriter. The empty reduction terminator generated
for loops without reductions is removed before moving the original body.

Add a regression test that collapses a two-dimensional parallel loop with a
reduction and uses the result outside the loop.

Testing:

- `ninja -C build -j4 mlir-opt`
- `mlir-opt ... single-parallel-loop-collapsing.mlir | FileCheck ...`
- Original reproducer from #<!-- -->216218
- Existing multi-dimensional parallel-loop-collapsing test、

Assisted by: codex


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


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+10-5) 
- (modified) mlir/test/Transforms/single-parallel-loop-collapsing.mlir (+24) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c158e624002bd..8c8bb608edcd4 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1158,8 +1158,8 @@ void mlir::collapseParallelLoops(
   // of the original induction value this represents. This is a normalized value
   // that is un-normalized already by the previous logic.
   auto newPloop = scf::ParallelOp::create(
-      rewriter, loc, lowerBounds, upperBounds, steps,
-      [&](OpBuilder &insideBuilder, Location, ValueRange ploopIVs) {
+      rewriter, loc, lowerBounds, upperBounds, steps, loops.getInitVals(),
+      [&](OpBuilder &insideBuilder, Location, ValueRange ploopIVs, ValueRange) {
         for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
           Value previous = ploopIVs[i];
           unsigned numberCombinedDimensions = combinedDimensions[i].size();
@@ -1187,11 +1187,16 @@ void mlir::collapseParallelLoops(
       });
 
   // Replace the old loop with the new loop.
-  loops.getBody()->back().erase();
+  // The builder creates an empty terminator when there are no reductions;
+  // replace it with the old loop's terminator together with the rest of the
+  // body. With reductions, no placeholder terminator is created.
+  if (!newPloop.getBody()->empty() &&
+      isa<scf::ReduceOp>(newPloop.getBody()->back()))
+    newPloop.getBody()->back().erase();
   newPloop.getBody()->getOperations().splice(
-      Block::iterator(newPloop.getBody()->back()),
+      newPloop.getBody()->getOperations().end(),
       loops.getBody()->getOperations());
-  loops.erase();
+  rewriter.replaceOp(loops, newPloop.getResults());
 }
 
 // Hoist the ops within `outer` that appear before `inner`.
diff --git a/mlir/test/Transforms/single-parallel-loop-collapsing.mlir b/mlir/test/Transforms/single-parallel-loop-collapsing.mlir
index 1ef787bec1bb3..f8a71f6603ade 100644
--- a/mlir/test/Transforms/single-parallel-loop-collapsing.mlir
+++ b/mlir/test/Transforms/single-parallel-loop-collapsing.mlir
@@ -27,3 +27,27 @@ func.func @collapse_to_single() {
 // CHECK:           scf.reduce
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return
+
+// CHECK-LABEL: func @collapse_with_reduction
+// CHECK-SAME:  (%[[INIT:.*]]: index)
+// CHECK:         %[[RESULT:.*]] = scf.parallel (%[[IV:.*]]) = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) init (%[[INIT]]) -> index {
+// CHECK:           %[[REM:.*]] = arith.remsi %[[IV]], %{{.*}} : index
+// CHECK:           %[[DIV:.*]] = arith.divsi %[[IV]], %{{.*}} : index
+// CHECK:           %[[SUM:.*]] = arith.addi %[[DIV]], %[[REM]] : index
+// CHECK:           scf.reduce(%[[SUM]] : index)
+// CHECK:         return %[[RESULT]] : index
+func.func @collapse_with_reduction(%init: index) -> index {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  %result = scf.parallel (%i, %j) = (%c0, %c0) to (%c10, %c10)
+      step (%c1, %c1) init (%init) -> index {
+    %sum = arith.addi %i, %j : index
+    scf.reduce(%sum : index) {
+    ^bb0(%lhs: index, %rhs: index):
+      %reduced = arith.addi %lhs, %rhs : index
+      scf.reduce.return %reduced : index
+    }
+  }
+  return %result : index
+}

``````````

</details>


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