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

[llvmlistbot at llvm.org]

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

>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
+}