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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: purnima shrivastava (purnima-nlp)

<details>
<summary>Changes</summary>


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


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


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp (+1-2) 
- (modified) mlir/test/Dialect/SCF/parallel-loop-fusion.mlir (+99) 


``````````diff
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

``````````

</details>


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