[Mlir-commits] [mlir] [mlir][scf] Allow parallel loop fusion across side-effect-free regions (PR #221586)

[purnima shrivastava]

https://github.com/purnima-nlp created https://github.com/llvm/llvm-project/pull/221586


Allow `scf-parallel-loop-fusion` to fuse across side-effect-free region operations such as `scf.if`.
`isMemoryEffectFree` already checks nested regions for recursive memory effects, so the explicit `op.getNumRegions() == 0` restriction was unnecessarily blocking valid fusion opportunities.

Added tests for:

* empty `scf.if`
* pure `scf.if`
* side-effecting `scf.if` that must still block fusion

Fixes #221526


>From 5b449882720f0257d9155d83611e32199999e029 Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Sun, 6 Sep 2026 21:16:50 +0530
Subject: [PATCH] [mlir][scf] Allow parallel loop fusion across
 side-effect-free regions

---
 .../SCF/Transforms/ParallelLoopFusion.cpp     |  3 +-
 .../Dialect/SCF/parallel-loop-fusion.mlir     | 99 +++++++++++++++++++
 2 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 397920565c5ba..b6a7dfe1188af 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -1072,8 +1072,7 @@ void mlir::scf::naivelyFuseParallelOps(
         }
         continue;
       }
-      // TODO: Handle region side effects properly.
-      noSideEffects &= isMemoryEffectFree(&op) && op.getNumRegions() == 0;
+      noSideEffects &= isMemoryEffectFree(&op);
     }
     for (MutableArrayRef<ParallelOp> ploops : ploopChains) {
       for (int i = 0, e = ploops.size(); i + 1 < e; ++i)
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index 9472e9ebb9d22..631c6764015a3 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -50,6 +50,105 @@ func.func @fuse_ops_between(%A: f32, %B: f32) -> f32 {
 
 // -----
 
+func.func @fuse_across_empty_region(%A: memref<16xf32>, %B: memref<16xf32>, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %value = arith.constant 6.0 : f32
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    memref.store %value, %A[%i] : memref<16xf32>
+    scf.reduce
+  }
+  scf.if %cond {
+  } else {
+  }
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    %v = memref.load %A[%i] : memref<16xf32>
+    memref.store %v, %B[%i] : memref<16xf32>
+    scf.reduce
+  }
+  return
+}
+// CHECK-LABEL: func @fuse_across_empty_region
+// CHECK-NOT:   scf.parallel
+// CHECK:       scf.if
+// CHECK:       scf.parallel
+// CHECK-NEXT:    memref.store
+// CHECK-NEXT:    %[[V:.*]] = memref.load
+// CHECK-NEXT:    memref.store %[[V]]
+// CHECK-NEXT:    scf.reduce
+// CHECK-NEXT:  }
+// CHECK-NOT:   scf.parallel
+// CHECK:       return
+
+// -----
+
+func.func @fuse_across_pure_region(%A: memref<16xf32>, %B: memref<16xf32>, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %value = arith.constant 6.0 : f32
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    memref.store %value, %A[%i] : memref<16xf32>
+    scf.reduce
+  }
+  scf.if %cond {
+    %x = arith.addi %c0, %c1 : index
+  }
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    %v = memref.load %A[%i] : memref<16xf32>
+    memref.store %v, %B[%i] : memref<16xf32>
+    scf.reduce
+  }
+  return
+}
+// CHECK-LABEL: func @fuse_across_pure_region
+// CHECK-NOT:   scf.parallel
+// CHECK:       scf.if
+// CHECK-NEXT:    %{{.*}} = arith.addi
+// CHECK-NEXT:  }
+// CHECK-NEXT:  scf.parallel
+// CHECK-NEXT:    memref.store
+// CHECK-NEXT:    %[[V:.*]] = memref.load
+// CHECK-NEXT:    memref.store %[[V]]
+// CHECK-NEXT:    scf.reduce
+// CHECK-NEXT:  }
+// CHECK-NOT:   scf.parallel
+// CHECK:       return
+
+// -----
+
+func.func @do_not_fuse_across_region_store(%A: memref<16xf32>, %B: memref<16xf32>, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %value = arith.constant 6.0 : f32
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    memref.store %value, %A[%i] : memref<16xf32>
+    scf.reduce
+  }
+  scf.if %cond {
+    memref.store %value, %A[%c0] : memref<16xf32>
+  }
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    %v = memref.load %A[%i] : memref<16xf32>
+    memref.store %v, %B[%i] : memref<16xf32>
+    scf.reduce
+  }
+  return
+}
+// CHECK-LABEL: func @do_not_fuse_across_region_store
+// CHECK:       scf.parallel
+// CHECK:         memref.store
+// CHECK:       scf.if
+// CHECK-NEXT:     memref.store
+// CHECK:       scf.parallel
+// CHECK-NEXT:    %[[V:.*]] = memref.load
+// CHECK-NEXT:    memref.store %[[V]]
+// CHECK:       return
+
+// -----
+
 func.func @fuse_two(%A: memref<2x2xf32>, %B: memref<2x2xf32>) {
   %c2 = arith.constant 2 : index
   %c0 = arith.constant 0 : index