[Mlir-commits] [mlir] [MLIR][SCF] Handle commuted indices in parallel loop fusion (PR #219665)

[purnima shrivastava]

https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/219665

>From d2fe2072d55ab690953f53646df5aa4b4a1010c0 Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Sat, 29 Aug 2026 15:21:28 +0530
Subject: [PATCH] [MLIR][SCF] Handle commuted indices in parallel loop fusion

---
 .../SCF/Transforms/ParallelLoopFusion.cpp     | 15 +++++++++-
 .../Dialect/SCF/parallel-loop-fusion.mlir     | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 397920565c5ba..bf9cc7bb203f2 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -145,13 +145,26 @@ static bool valsAreEquivalent(Value val1, Value val2,
     return false;
   if (!isMemoryEffectFree(val1DefOp) || !isMemoryEffectFree(val2DefOp))
     return false;
+
+  auto checkCommutativeEquivalent = [&](ValueRange lhs,
+                                        ValueRange rhs) -> LogicalResult {
+    if (lhs.size() != 2 || rhs.size() != 2)
+      return failure();
+    if (valsAreEquivalent(lhs[0], rhs[0], loopsIVsMap) &&
+        valsAreEquivalent(lhs[1], rhs[1], loopsIVsMap))
+      return success();
+    return success(valsAreEquivalent(lhs[0], rhs[1], loopsIVsMap) &&
+                   valsAreEquivalent(lhs[1], rhs[0], loopsIVsMap));
+  };
+
   return OperationEquivalence::isEquivalentTo(
       val1DefOp, val2DefOp,
       [&](Value v1, Value v2) {
         return success(loopsIVsMap.lookupOrDefault(v1) == v2 ||
                        loopsIVsMap.lookupOrDefault(v2) == v1);
       },
-      /*markEquivalent=*/nullptr, OperationEquivalence::Flags::IgnoreLocations);
+      /*markEquivalent=*/nullptr, OperationEquivalence::Flags::IgnoreLocations,
+      checkCommutativeEquivalent);
 }
 
 /// If the `expr` value is the result of an integer addition of `base` and a
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index 9472e9ebb9d22..45cfb7c20e0ca 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -1761,3 +1761,33 @@ func.func @do_not_fuse_distinct_dynamic_bounds(%A: memref<16xf32>,
 // CHECK-LABEL: func @do_not_fuse_distinct_dynamic_bounds
 // CHECK:        scf.parallel
 // CHECK:        scf.parallel
+
+// -----
+
+func.func @fuse_commuted_indices(%arg0: memref<32xf32>,
+                                 %arg1: memref<32xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %cst = arith.constant 3.000000e+00 : f32
+
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    %index = arith.addi %i, %c1 : index
+    memref.store %cst, %arg0[%index] : memref<32xf32>
+    scf.reduce
+  }
+
+  scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+    %index = arith.addi %c1, %i : index
+    %value = memref.load %arg0[%index] : memref<32xf32>
+    memref.store %value, %arg1[%index] : memref<32xf32>
+    scf.reduce
+  }
+
+  return
+}
+
+// CHECK-LABEL: func.func @fuse_commuted_indices
+// CHECK:       scf.parallel
+// CHECK-NOT:   scf.parallel
+// CHECK:       return