[Mlir-commits] [mlir] 05f7078 - [mlir][scf] Fuse parallel loops with equal constant bounds (#218450)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Aug 25 00:56:36 PDT 2026
Author: Hamza Qureshi
Date: 2026-08-25T09:56:30+02:00
New Revision: 05f70782cc14770f0250de0468ab1f9196b4391a
URL: https://github.com/llvm/llvm-project/commit/05f70782cc14770f0250de0468ab1f9196b4391a
DIFF: https://github.com/llvm/llvm-project/commit/05f70782cc14770f0250de0468ab1f9196b4391a.diff
LOG: [mlir][scf] Fuse parallel loops with equal constant bounds (#218450)
Two `scf.parallel` loops that run over the same range were not being
fused if
their bounds came from different `arith.constant` ops.
Both loops here go from 0 to 16 by 1, so they should fuse, but they do
not:
```mlir
%c0 = arith.constant 0 : index
%c16 = arith.constant 16 : index
%c1 = arith.constant 1 : index
scf.parallel (%i) = (%c0) to (%c16) step (%c1) { ... }
%c0_0 = arith.constant 0 : index
%c16_1 = arith.constant 16 : index
%c1_2 = arith.constant 1 : index
scf.parallel (%i) = (%c0_0) to (%c16_1) step (%c1_2) { ... }
```
The reason is that `equalIterationSpaces()` compared the bounds with
`std::equal` over `Value`s, which only asks "is this the same SSA
value?".
`%c16` and `%c16_1` are two different values, so the check failed even
though
both are 16.
Now the check also accepts two operands that are constants with the same
value. This is what the TODO on that code already asked for:
```
- // TODO: Extend this to support aliases and equal constants.
+ // TODO: Extend this to support aliases.
```
Operands that are not constants are still compared by identity, so loops
with
different dynamic bounds are still not fused.
Note that running `-cse` or `-canonicalize` first also fixes this,
because it
merges the duplicate constants. I did not want the pass to rely on that
having
run, since the two iteration spaces are equal either way.
Tests: two cases that should fuse (all bounds equal constants, and a mix
of
shared values and equal constants), and two that should not (constants
with
different values, and different dynamic bounds).
Fixes #218318
Added:
Modified:
mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 1e879d81f3559..397920565c5ba 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -21,6 +21,7 @@
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
@@ -65,10 +66,23 @@ static bool equalIterationSpaces(ParallelOp firstPloop,
if (firstPloop.getNumLoops() != secondPloop.getNumLoops())
return false;
+ // Two bounds match if they are the same value, or if both are constants
+ // holding the same value. The latter matters because equivalent bounds are
+ // often materialized by distinct `arith.constant` ops, which leaves the
+ // iteration spaces equal even though the SSA values
diff er.
auto matchOperands = [&](const OperandRange &lhs,
const OperandRange &rhs) -> bool {
- // TODO: Extend this to support aliases and equal constants.
- return std::equal(lhs.begin(), lhs.end(), rhs.begin());
+ // TODO: Extend this to support aliases.
+ return std::equal(lhs.begin(), lhs.end(), rhs.begin(),
+ [](Value lhsValue, Value rhsValue) {
+ if (lhsValue == rhsValue)
+ return true;
+ std::optional<int64_t> lhsConst =
+ getConstantIntValue(lhsValue);
+ std::optional<int64_t> rhsConst =
+ getConstantIntValue(rhsValue);
+ return lhsConst && rhsConst && *lhsConst == *rhsConst;
+ });
};
return matchOperands(firstPloop.getLowerBound(),
secondPloop.getLowerBound()) &&
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index 8a630713c5be6..9472e9ebb9d22 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -1643,3 +1643,121 @@ func.func @fuse_chain_after_interchanged_reduction(
// CHECK: }
// CHECK-NOT: scf.parallel
// CHECK: return
+
+// -----
+
+// Bounds that are equal constants describe the same iteration space even when
+// they are materialized by distinct `arith.constant` ops, so the loops fuse.
+
+func.func @fuse_equal_constant_bounds(%A: memref<16xf32>, %B: memref<16xf32>) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 1.000000e+00 : f32
+ scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+ memref.store %cst, %A[%i] : memref<16xf32>
+ scf.reduce
+ }
+ %c0_0 = arith.constant 0 : index
+ %c16_1 = arith.constant 16 : index
+ %c1_2 = arith.constant 1 : index
+ scf.parallel (%i) = (%c0_0) to (%c16_1) step (%c1_2) {
+ %0 = memref.load %A[%i] : memref<16xf32>
+ memref.store %0, %B[%i] : memref<16xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @fuse_equal_constant_bounds
+// CHECK: scf.parallel ([[I:%.*]]) =
+// CHECK-NEXT: memref.store {{.*}}{{\[}}[[I]]{{\]}}
+// CHECK-NEXT: [[V:%.*]] = memref.load {{.*}}{{\[}}[[I]]{{\]}}
+// CHECK-NEXT: memref.store [[V]], {{.*}}{{\[}}[[I]]{{\]}}
+// CHECK-NEXT: scf.reduce
+// CHECK-NEXT: }
+// CHECK-NOT: scf.parallel
+
+// -----
+
+// Only some of the bounds need to be shared SSA values; the rest may be equal
+// constants defined separately.
+
+func.func @fuse_partially_shared_constant_bounds(%A: memref<4x4xf32>,
+ %B: memref<4x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 1.000000e+00 : f32
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c4, %c4) step (%c1, %c1) {
+ memref.store %cst, %A[%i, %j] : memref<4x4xf32>
+ scf.reduce
+ }
+ %c4_0 = arith.constant 4 : index
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c4, %c4_0) step (%c1, %c1) {
+ %0 = memref.load %A[%i, %j] : memref<4x4xf32>
+ memref.store %0, %B[%i, %j] : memref<4x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @fuse_partially_shared_constant_bounds
+// CHECK: scf.parallel ([[I:%.*]], [[J:%.*]]) =
+// CHECK-NEXT: memref.store {{.*}}{{\[}}[[I]], [[J]]{{\]}}
+// CHECK-NEXT: [[V:%.*]] = memref.load {{.*}}{{\[}}[[I]], [[J]]{{\]}}
+// CHECK-NEXT: memref.store [[V]], {{.*}}{{\[}}[[I]], [[J]]{{\]}}
+// CHECK-NEXT: scf.reduce
+// CHECK-NEXT: }
+// CHECK-NOT: scf.parallel
+
+// -----
+
+// Distinct constants holding
diff erent values are still
diff erent iteration
+// spaces and must not fuse.
+
+func.func @do_not_fuse_unequal_constant_bounds(%A: memref<16xf32>,
+ %B: memref<16xf32>) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 1.000000e+00 : f32
+ scf.parallel (%i) = (%c0) to (%c16) step (%c1) {
+ memref.store %cst, %A[%i] : memref<16xf32>
+ scf.reduce
+ }
+ %c8 = arith.constant 8 : index
+ scf.parallel (%i) = (%c0) to (%c8) step (%c1) {
+ %0 = memref.load %A[%i] : memref<16xf32>
+ memref.store %0, %B[%i] : memref<16xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_unequal_constant_bounds
+// CHECK: scf.parallel
+// CHECK: scf.parallel
+
+// -----
+
+// A non-constant bound cannot be compared by value, so distinct dynamic bounds
+// must not fuse even when they may happen to be equal at runtime.
+
+func.func @do_not_fuse_distinct_dynamic_bounds(%A: memref<16xf32>,
+ %B: memref<16xf32>,
+ %ub0: index, %ub1: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 1.000000e+00 : f32
+ scf.parallel (%i) = (%c0) to (%ub0) step (%c1) {
+ memref.store %cst, %A[%i] : memref<16xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%ub1) step (%c1) {
+ %0 = memref.load %A[%i] : memref<16xf32>
+ memref.store %0, %B[%i] : memref<16xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_distinct_dynamic_bounds
+// CHECK: scf.parallel
+// CHECK: scf.parallel
More information about the Mlir-commits
mailing list