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

[purnima shrivastava]

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

`scf-parallel-loop-fusion` compares index-producing operations with
operands in their original order. Consequently, equivalent accesses such
as `%iv + 1` and `1 + %iv` are not recognized as accessing the same
location, preventing otherwise legal loop fusion.

Handle commuted `arith.addi` expressions in `valsAreEquivalent` while
requiring matching result types and operation attributes.

Add a regression test verifying that parallel loops using commuted index
expressions are fused.

Fixes #218614.

>From a2600d408cfbc328131bf2317426e2de1e50eed1 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     | 12 ++++++++
 .../Dialect/SCF/parallel-loop-fusion.mlir     | 30 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 397920565c5ba..265fa6c8a39d0 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -145,6 +145,18 @@ static bool valsAreEquivalent(Value val1, Value val2,
     return false;
   if (!isMemoryEffectFree(val1DefOp) || !isMemoryEffectFree(val2DefOp))
     return false;
+
+  // Handle commuted integer additions.
+  if (auto addOp1 = dyn_cast<arith::AddIOp>(val1DefOp)) {
+    if (auto addOp2 = dyn_cast<arith::AddIOp>(val2DefOp);
+        addOp2 &&
+        val1DefOp->getAttrDictionary() == val2DefOp->getAttrDictionary() &&
+        val1.getType() == val2.getType() &&
+        valsAreEquivalent(addOp1.getLhs(), addOp2.getRhs(), loopsIVsMap) &&
+        valsAreEquivalent(addOp1.getRhs(), addOp2.getLhs(), loopsIVsMap))
+      return true;
+  }
+
   return OperationEquivalence::isEquivalentTo(
       val1DefOp, val2DefOp,
       [&](Value v1, Value v2) {
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index 9472e9ebb9d22..48c7276318d0d 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
\ No newline at end of file