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

[llvmlistbot at llvm.org]

https://github.com/cr-zhao created https://github.com/llvm/llvm-project/pull/216538

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


>From de49f548bc6dff0e36b1cb74eedd613fc0f9170c Mon Sep 17 00:00:00 2001
From: real-cpu <zhaochenrui757 at gmail.com>
Date: Sat, 15 Aug 2026 23:24:54 -0700
Subject: [PATCH] [mlir][scf] Preserve reductions when collapsing parallel
 loops

---
 mlir/lib/Dialect/SCF/Utils/Utils.cpp          | 16 +++++++++----
 .../single-parallel-loop-collapsing.mlir      | 24 +++++++++++++++++++
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c789b4c8904d3..c3f8795e3da82 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1154,8 +1154,9 @@ 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();
@@ -1183,11 +1184,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
+}