[Mlir-commits] [mlir] [MLIR][SCF] Support permutation-based parallel loop fusion (PR #203207)
Dmitriy Smirnov
llvmlistbot at llvm.org
Tue Jun 16 06:30:33 PDT 2026
https://github.com/d-smirnov updated https://github.com/llvm/llvm-project/pull/203207
>From 70a7f819abaa68a086a52dc5deefa428f7508ffa Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Wed, 10 Jun 2026 18:04:34 +0100
Subject: [PATCH 1/7] [MLIR][SCF] Support permutation-based parallel loop
fusion
Improve SCF parallel-loop fusion for loops with permuted iteration spaces.
Allow fusion after rewriting the second loop using an arbitrary
permutation of its iteration space. When multiple axes have identical
bounds and steps, also enumerate additional candidate remaps for those
equal axes.
Change-Id: Ie84f14aff5e935e51a70125610e07136f153535f
---
.../SCF/Transforms/ParallelLoopFusion.cpp | 214 ++++++++++++++++--
.../Dialect/SCF/parallel-loop-fusion.mlir | 105 +++++++++
2 files changed, 294 insertions(+), 25 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index bdaade59e99df..ddefb08a21b14 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -20,6 +20,7 @@
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
@@ -28,6 +29,7 @@
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "llvm/ADT/STLExtras.h"
@@ -35,6 +37,7 @@
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/TypeSwitch.h"
+#include <numeric>
#include <optional>
#include <tuple>
@@ -733,26 +736,33 @@ static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
firstToSecondPloopIndices, mayAlias, b);
}
-// Interchange loops of the parallel loop, if there are just two loops
-static std::optional<ParallelOp> interchangeLoops(OpBuilder &builder,
- ParallelOp &loop) {
-
- if (loop.getNumLoops() != 2)
+// Returns new parallel loop where two loops matching indices param are
+// interchanged
+static std::optional<ParallelOp>
+interchangeLoops(OpBuilder &builder, ParallelOp &loop,
+ const ArrayRef<int64_t> &indices) {
+ assert(loop.getNumLoops() == indices.size());
+ if (loop.getNumLoops() < 2)
return std::nullopt;
- OpBuilder::InsertionGuard guard(builder);
-
// Replace the parallel loop with the same parallel loop.
builder.setInsertionPoint(loop);
- auto newOp = ParallelOp::create(builder, loop.getLoc(), loop.getLowerBound(),
- loop.getUpperBound(), loop.getStep(),
+ SmallVector<Value> newLB =
+ applyPermutation(SmallVector<Value>(loop.getLowerBound()), indices);
+ SmallVector<Value> newUB =
+ applyPermutation(SmallVector<Value>(loop.getUpperBound()), indices);
+ SmallVector<Value> newStep =
+ applyPermutation(SmallVector<Value>(loop.getStep()), indices);
+ auto newOp = ParallelOp::create(builder, loop.getLoc(), newLB, newUB, newStep,
loop.getInitVals(), nullptr);
- IRMapping mapping;
auto ivs = loop.getInductionVars();
- auto newIvs = newOp.getInductionVars();
- for (auto [iv, riv] : llvm::zip(ivs, llvm::reverse(newIvs))) {
+ SmallVector<Value> newIvs = applyPermutation(
+ newOp.getInductionVars(), invertPermutationVector(indices));
+ IRMapping mapping;
+ for (auto [iv, riv] : llvm::zip(ivs, newIvs)) {
mapping.map(iv, riv);
}
+
// Copy parallel loop body
builder.setInsertionPoint(&(newOp.getBody()->front()));
for (auto &o : loop.getRegion().front().without_terminator()) {
@@ -761,6 +771,156 @@ static std::optional<ParallelOp> interchangeLoops(OpBuilder &builder,
return newOp;
}
+struct LoopIV {
+ Value lBound, uBound, step;
+ bool operator!=(LoopIV const &other) const { return !(*this == other); }
+ bool operator==(LoopIV const &other) const {
+ return lBound == other.lBound && uBound == other.uBound &&
+ step == other.step;
+ }
+};
+
+template <>
+struct llvm::DenseMapInfo<LoopIV> {
+ static inline LoopIV getEmptyKey() {
+ auto e = DenseMapInfo<mlir::Value>::getEmptyKey();
+ return {e, e, e}; // Must be impossible to occur naturally
+ }
+ static inline LoopIV getTombstoneKey() {
+ auto t = DenseMapInfo<mlir::Value>::getTombstoneKey();
+ return {t, t, t}; // Must be impossible to occur naturally
+ }
+
+ static inline bool isEqual(const LoopIV &lhs, const LoopIV &rhs) {
+ return (lhs == rhs);
+ }
+
+ static inline unsigned getHashValue(const LoopIV &val) {
+ return llvm::hash_combine(
+ DenseMapInfo<mlir::Value>::getHashValue(val.lBound),
+ DenseMapInfo<mlir::Value>::getHashValue(val.uBound),
+ DenseMapInfo<mlir::Value>::getHashValue(val.step));
+ }
+};
+
+// Returns vector of candidate permutation indices vectors,
+// can be empty
+static SmallVector<SmallVector<int64_t>>
+computeCandidateInterchangePermutations(ParallelOp &firstPloop,
+ ParallelOp &secondPloop) {
+ SmallVector<SmallVector<int64_t>> extraResults;
+
+ // Check preconditions
+ if (firstPloop.getNumLoops() < 2 ||
+ firstPloop.getNumLoops() != secondPloop.getNumLoops())
+ return extraResults;
+
+ SmallVector<LoopIV> firstIVs(firstPloop.getNumLoops());
+ SmallVector<LoopIV> secondIVs(secondPloop.getNumLoops());
+ llvm::SmallSetVector<LoopIV, 6> unique;
+ for (unsigned index = 0; index < firstPloop.getNumLoops(); ++index) {
+ firstIVs[index].lBound = firstPloop.getLowerBound()[index];
+ firstIVs[index].uBound = firstPloop.getUpperBound()[index];
+ firstIVs[index].step = firstPloop.getStep()[index];
+ secondIVs[index].lBound = secondPloop.getLowerBound()[index];
+ secondIVs[index].uBound = secondPloop.getUpperBound()[index];
+ secondIVs[index].step = secondPloop.getStep()[index];
+ unique.insert(firstIVs[index]);
+ }
+
+ SmallVector<bool> diffIVs(firstPloop.getNumLoops());
+ std::transform(firstIVs.begin(), firstIVs.end(), secondIVs.begin(),
+ diffIVs.begin(), std::not_equal_to());
+
+ SmallVector<int64_t> indices;
+ for (auto [idx, val] : enumerate(diffIVs))
+ if (val)
+ indices.push_back(idx);
+
+ // Not a permutation shortcut
+ if (indices.size() == 1)
+ return extraResults;
+
+ // Initialize with identity permutations
+ SmallVector<int64_t> result(firstIVs.size());
+ std::iota(result.begin(), result.end(), 0);
+
+ if (indices.empty() && unique.size() == firstIVs.size())
+ return extraResults;
+
+ if (indices.size() > 1) {
+ // Determine whether the iteration space of the first loop is a permutation
+ // of the second and collect remaps.
+ SmallVector<int64_t> remaps;
+ for (auto fIdx : indices) {
+ for (auto sIdx : indices) {
+ // can be remapped
+ if (fIdx != sIdx && firstIVs[fIdx] == secondIVs[sIdx] &&
+ remaps.end() == std::find(remaps.begin(), remaps.end(), sIdx)) {
+ remaps.push_back(sIdx);
+ }
+ }
+ }
+
+ // Not a permutation
+ if (indices.size() != remaps.size())
+ return extraResults;
+
+ // compose permutation indices
+ for (auto [from, to] : zip(indices, remaps)) {
+ result[from] = to;
+ }
+
+ extraResults.push_back(result);
+
+ // All axes are unique, no further permutatons needed
+ if (unique.size() == firstIVs.size()) {
+ return extraResults;
+ }
+ }
+
+ //
+ // Permute equal axes
+ assert(unique.size() != firstIVs.size() &&
+ "Expected at least two equal axes");
+
+ // Collect equal axes to groups
+ SmallVector<SmallVector<int64_t>> groups;
+ for (auto iv : unique) {
+ SmallVector<int64_t> group;
+ for (unsigned index = 0; index < firstIVs.size(); ++index) {
+ if (firstIVs[index] == iv)
+ group.push_back(index);
+ }
+ if (group.size() > 1)
+ groups.push_back(std::move(group));
+ }
+
+ // Permute axes groups
+ SmallVector<SmallVector<int64_t>> rmpdGroups(groups);
+ bool repeat = false;
+ do {
+ repeat = false;
+ for (auto const &[group, groupRemaps] : zip(groups, rmpdGroups)) {
+ repeat |= std::next_permutation(groupRemaps.begin(), groupRemaps.end());
+ if (repeat)
+ break;
+ }
+
+ if (repeat) {
+ SmallVector<int64_t> extra(result);
+ for (auto const &[group, groupRemaps] : zip(groups, rmpdGroups)) {
+ for (auto [from, to] : zip(group, groupRemaps))
+ extra[from] = result[to];
+ }
+ if (result != extra)
+ extraResults.push_back(std::move(extra));
+ }
+ } while (repeat);
+
+ return extraResults;
+}
+
/// Prepend operations of firstPloop's body into secondPloop's body.
/// Update secondPloop with new loop.
static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
@@ -773,27 +933,31 @@ static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
mayAlias, builder)) {
- // If second parallel loop consists of two loops of same iteration space
- // then exchange these loops and re-asses the possibility of fusion.
- if (secondPloop.getNumLoops() == 2 &&
- secondPloop.getUpperBound()[0] == secondPloop.getUpperBound()[1] &&
- secondPloop.getLowerBound()[0] == secondPloop.getLowerBound()[1] &&
- secondPloop.getStep()[0] == secondPloop.getStep()[1]) {
+ // If iteration space of the second parallel loop is a permutation of the
+ // first one then interchange iteration space of the second parallel loop
+ // and re-asses possibility of fusion.
+ for (auto &perms :
+ computeCandidateInterchangePermutations(firstPloop, secondPloop)) {
+ OpBuilder::InsertionGuard guard(builder);
+ auto newLoop = interchangeLoops(builder, secondPloop, perms);
firstToSecondPloopIndices.clear();
firstToSecondPloopIndices.map(block1->getArguments(),
- llvm::reverse(block2->getArguments()));
- if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
- mayAlias, builder))
- return;
- auto newLoop = interchangeLoops(builder, secondPloop);
+ newLoop->getBody()->getArguments());
+ if (!isFusionLegal(firstPloop, *newLoop, firstToSecondPloopIndices,
+ mayAlias, builder)) {
+ newLoop->erase();
+ continue;
+ }
+
secondPloop->erase();
secondPloop = *newLoop;
block2 = secondPloop.getBody();
- } else {
- return;
+ goto fuseLabel;
}
+ return;
}
+fuseLabel:
DominanceInfo dom;
// We are fusing first loop into second, make sure there are no users of the
// first loop results between loops.
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index ac6ab11963bd5..54c418143546e 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -1249,3 +1249,108 @@ func.func @test_fuse_interchanged_loops(%arg0: memref<1x64xf32>) {
// CHECK-LABEL: func @test_fuse_interchanged_loops
// CHECK: scf.parallel
// CHECK-NOT: scf.parallel
+
+// -----
+
+func.func @fuse_three_cycle_permutation(
+ %out: memref<2x3x5xf32>) {
+ %A = memref.alloc() : memref<2x3x5xf32>
+ %tmp = memref.alloc() : memref<2x3x5xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c5 = arith.constant 5 : index
+ %cst = arith.constant 1.0 : f32
+
+ scf.parallel (%i, %j, %k) = (%c0, %c0, %c0) to (%c2, %c3, %c5) step (%c1, %c1, %c1) {
+ %a = memref.load %A[%i, %j, %k] : memref<2x3x5xf32>
+ %b = arith.addf %a, %cst : f32
+ memref.store %b, %tmp[%i, %j, %k] : memref<2x3x5xf32>
+ scf.reduce
+ }
+
+ scf.parallel (%k2, %i2, %j2) = (%c0, %c0, %c0) to (%c5, %c2, %c3) step (%c1, %c1, %c1) {
+ %t = memref.load %tmp[%i2, %j2, %k2] : memref<2x3x5xf32>
+ memref.store %t, %out[%i2, %j2, %k2] : memref<2x3x5xf32>
+ scf.reduce
+ }
+ return
+}
+
+// CHECK-LABEL: func @fuse_three_cycle_permutation
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[C5:.*]] = arith.constant 5 : index
+// CHECK: %[[CST:.*]] = arith.constant 1.
+
+// CHECK: scf.parallel (%[[I:.*]], %[[J:.*]], %[[K:.*]]) = (%[[C0]], %[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C2]], %[[C3]], %[[C5]]) step (%[[C1]], %[[C1]], %[[C1]]) {
+// CHECK: %[[A_ELT:.*]] = memref.load %{{.*}}%[[I]], %[[J]], %[[K]]] : memref<2x3x5xf32>
+// CHECK: %[[B_ELT:.*]] = arith.addf %[[A_ELT]], %[[CST]] : f32
+// CHECK: memref.store %[[B_ELT]], %{{.*}}%[[I]], %[[J]], %[[K]]] : memref<2x3x5xf32>
+// CHECK-NOT: scf.parallel
+// CHECK: %[[T:.*]] = memref.load %{{.*}}%[[I]], %[[J]], %[[K]]] : memref<2x3x5xf32>
+// CHECK: memref.store %[[T]], %{{.*}}%[[I]], %[[J]], %[[K]]] : memref<2x3x5xf32>
+// CHECK: scf.reduce
+// CHECK: }
+// CHECK-NOT: scf.parallel
+
+// -----
+
+func.func @fuse_duplicate_axes_permutation(
+%out : memref<2x2x3x3xf32>) {
+ %tmp = memref.alloc() : memref<2x2x3x3xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %v = arith.constant 1.0 : f32
+
+ // First loop: canonical order (i, j, k, l)
+ scf.parallel (%i, %j, %k, %l) = (%c0, %c0, %c0, %c0)
+ to (%c2, %c2, %c3, %c3) step (%c1, %c1, %c1, %c1) {
+ memref.store %v, %tmp[%i, %j, %k, %l] : memref<2x2x3x3xf32>
+ scf.reduce
+ }
+
+ // Second loop iteration space is a permutation of the first:
+ // positions are (k2, l2, j2, i2) with extents (3, 3, 2, 2).
+ //
+ // The body is written so that the "right" correspondence is:
+ // i -> i2 (pos 3)
+ // j -> j2 (pos 2)
+ // k -> k2 (pos 0)
+ // l -> l2 (pos 1)
+ //
+ // i.e. permutation [3, 2, 0, 1] if interpreted as newPos -> oldPos.
+ scf.parallel (%k2, %l2, %j2, %i2) = (%c0, %c0, %c0, %c0)
+ to (%c3, %c3, %c2, %c2) step (%c1, %c1, %c1, %c1) {
+ %t = memref.load %tmp[%i2, %j2, %k2, %l2] : memref<2x2x3x3xf32>
+ memref.store %t, %out[%i2, %j2, %k2, %l2] : memref<2x2x3x3xf32>
+ scf.reduce
+ }
+ return
+}
+
+// CHECK-LABEL: func @fuse_duplicate_axes_permutation
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[CST:.*]] = arith.constant 1.
+
+// CHECK: scf.parallel (%[[I:.*]], %[[J:.*]], %[[K:.*]], %[[L:.*]]) = (%[[C0]], %[[C0]], %[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C2]], %[[C2]], %[[C3]], %[[C3]]) step (%[[C1]], %[[C1]], %[[C1]], %[[C1]]) {
+
+// CHECK: memref.store %[[CST]], %{{.*}}{{\[}}%[[I]], %[[J]], %[[K]], %[[L]]{{\]}} : memref<2x2x3x3xf32>
+
+// CHECK-NOT: scf.parallel
+// CHECK: %[[T:.*]] = memref.load %{{.*}}{{\[}}%[[I]], %[[J]], %[[K]], %[[L]]{{\]}} : memref<2x2x3x3xf32>
+// CHECK: memref.store %[[T]], %{{.*}}{{\[}}%[[I]], %[[J]], %[[K]], %[[L]]{{\]}} : memref<2x2x3x3xf32>
+
+// CHECK: scf.reduce
+// CHECK: }
+// CHECK-NOT: scf.parallel
>From 89e590c45c01990e1d4f227368217755f23608a1 Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Thu, 11 Jun 2026 12:36:18 +0100
Subject: [PATCH 2/7] Fixed build issues
Change-Id: I2c6b21fb8f196df0435459a9aec4991e41f15830
---
mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp | 9 ---------
1 file changed, 9 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index ddefb08a21b14..ee860891c6945 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -782,15 +782,6 @@ struct LoopIV {
template <>
struct llvm::DenseMapInfo<LoopIV> {
- static inline LoopIV getEmptyKey() {
- auto e = DenseMapInfo<mlir::Value>::getEmptyKey();
- return {e, e, e}; // Must be impossible to occur naturally
- }
- static inline LoopIV getTombstoneKey() {
- auto t = DenseMapInfo<mlir::Value>::getTombstoneKey();
- return {t, t, t}; // Must be impossible to occur naturally
- }
-
static inline bool isEqual(const LoopIV &lhs, const LoopIV &rhs) {
return (lhs == rhs);
}
>From 4c5c0c803474f741d5dc645d73797ac907119009 Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Thu, 11 Jun 2026 13:15:59 +0100
Subject: [PATCH 3/7] Fixed build issues
Change-Id: I08d893cc50cddd83434dda820d6f8bb3d3295a0e
---
mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index ee860891c6945..0415bcf3a16d8 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -821,7 +821,7 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
SmallVector<bool> diffIVs(firstPloop.getNumLoops());
std::transform(firstIVs.begin(), firstIVs.end(), secondIVs.begin(),
- diffIVs.begin(), std::not_equal_to());
+ diffIVs.begin(), std::not_equal_to<LoopIV>());
SmallVector<int64_t> indices;
for (auto [idx, val] : enumerate(diffIVs))
>From 5ea94c790af142597723db7e9c32ddae0ded48ab Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Fri, 12 Jun 2026 16:00:18 +0100
Subject: [PATCH 4/7] Addressed comments
Change-Id: Idfc9b57d8e8649d9f5660a9725190b169a11426f
---
.../SCF/Transforms/ParallelLoopFusion.cpp | 25 +++--
.../Dialect/SCF/parallel-loop-fusion.mlir | 101 ++++++++++++++++++
2 files changed, 117 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 0415bcf3a16d8..63795234c00a9 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -764,9 +764,11 @@ interchangeLoops(OpBuilder &builder, ParallelOp &loop,
}
// Copy parallel loop body
- builder.setInsertionPoint(&(newOp.getBody()->front()));
- for (auto &o : loop.getRegion().front().without_terminator()) {
- builder.clone(o, mapping);
+ auto b = OpBuilder::atBlockBegin(newOp.getBody());
+ for (auto &o : loop.getNumReductions()
+ ? loop.getBodyRegion().front()
+ : loop.getBodyRegion().front().without_terminator()) {
+ b.clone(o, mapping);
}
return newOp;
}
@@ -798,7 +800,8 @@ struct llvm::DenseMapInfo<LoopIV> {
// can be empty
static SmallVector<SmallVector<int64_t>>
computeCandidateInterchangePermutations(ParallelOp &firstPloop,
- ParallelOp &secondPloop) {
+ ParallelOp &secondPloop,
+ int axesLimit = 5) {
SmallVector<SmallVector<int64_t>> extraResults;
// Check preconditions
@@ -809,7 +812,7 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
SmallVector<LoopIV> firstIVs(firstPloop.getNumLoops());
SmallVector<LoopIV> secondIVs(secondPloop.getNumLoops());
llvm::SmallSetVector<LoopIV, 6> unique;
- for (unsigned index = 0; index < firstPloop.getNumLoops(); ++index) {
+ for (unsigned index : llvm::seq(firstPloop.getNumLoops())) {
firstIVs[index].lBound = firstPloop.getLowerBound()[index];
firstIVs[index].uBound = firstPloop.getUpperBound()[index];
firstIVs[index].step = firstPloop.getStep()[index];
@@ -820,8 +823,9 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
}
SmallVector<bool> diffIVs(firstPloop.getNumLoops());
- std::transform(firstIVs.begin(), firstIVs.end(), secondIVs.begin(),
- diffIVs.begin(), std::not_equal_to<LoopIV>());
+ llvm::transform(
+ llvm::zip(firstIVs, secondIVs), diffIVs.begin(),
+ [](auto const &pair) { return std::get<0>(pair) != std::get<1>(pair); });
SmallVector<int64_t> indices;
for (auto [idx, val] : enumerate(diffIVs))
@@ -879,9 +883,11 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
SmallVector<SmallVector<int64_t>> groups;
for (auto iv : unique) {
SmallVector<int64_t> group;
- for (unsigned index = 0; index < firstIVs.size(); ++index) {
- if (firstIVs[index] == iv)
+ for (unsigned index : llvm::seq(firstIVs.size())) {
+ if (axesLimit && firstIVs[index] == iv) {
+ axesLimit--;
group.push_back(index);
+ }
}
if (group.size() > 1)
groups.push_back(std::move(group));
@@ -940,6 +946,7 @@ static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
continue;
}
+ secondPloop.replaceAllUsesWith(newLoop->getResults());
secondPloop->erase();
secondPloop = *newLoop;
block2 = secondPloop.getBody();
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index 54c418143546e..0d66e15a2b05a 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -1354,3 +1354,104 @@ func.func @fuse_duplicate_axes_permutation(
// CHECK: scf.reduce
// CHECK: }
// CHECK-NOT: scf.parallel
+
+// -----
+
+func.func @fuse_interchanged_reductions(%A: memref<2x3xf32>,
+ %B: memref<2x3xf32>) -> (f32, f32) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %init1 = arith.constant 1.0 : f32
+ %init2 = arith.constant 2.0 : f32
+ %res1 = scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c3)
+ step (%c1, %c1) init(%init1) -> f32 {
+ %A_elem = memref.load %A[%i, %j] : memref<2x3xf32>
+ scf.reduce(%A_elem : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %1 = arith.addf %lhs, %rhs : f32
+ scf.reduce.return %1 : f32
+ }
+ }
+ %res2 = scf.parallel (%j2, %i2) = (%c0, %c0) to (%c3, %c2)
+ step (%c1, %c1) init(%init2) -> f32 {
+ %B_elem = memref.load %B[%i2, %j2] : memref<2x3xf32>
+ scf.reduce(%B_elem : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %1 = arith.mulf %lhs, %rhs : f32
+ scf.reduce.return %1 : f32
+ }
+ }
+ return %res1, %res2 : f32, f32
+}
+
+// CHECK-LABEL: func @fuse_interchanged_reductions
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[INIT1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK: %[[INIT2:.*]] = arith.constant 2.000000e+00 : f32
+// CHECK: %[[RES:.*]]:2 = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C2]], %[[C3]]) step (%[[C1]], %[[C1]])
+// CHECK-SAME: init (%[[INIT1]], %[[INIT2]]) -> (f32, f32) {
+// CHECK: %[[AELT:.*]] = memref.load %{{.*}}{{\[}}%[[I]], %[[J]]{{\]}} : memref<2x3xf32>
+// CHECK: %[[BELT:.*]] = memref.load %{{.*}}{{\[}}%[[I]], %[[J]]{{\]}} : memref<2x3xf32>
+// CHECK: scf.reduce(%[[AELT]], %[[BELT]] : f32, f32) {
+// CHECK: ^bb0
+// CHECK: ^bb0
+// CHECK: return %[[RES]]#0, %[[RES]]#1 : f32, f32
+
+// -----
+
+func.func @fuse_three_cycle_reductions(%A: memref<2x3x5xf32>,
+ %B: memref<2x3x5xf32>) -> (f32, f32) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c5 = arith.constant 5 : index
+ %init1 = arith.constant 1.0 : f32
+ %init2 = arith.constant 2.0 : f32
+
+ %res1 = scf.parallel (%i, %j, %k) = (%c0, %c0, %c0)
+ to (%c2, %c3, %c5) step (%c1, %c1, %c1) init(%init1) -> f32 {
+ %a = memref.load %A[%i, %j, %k] : memref<2x3x5xf32>
+ scf.reduce(%a : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ scf.reduce.return %sum : f32
+ }
+ }
+
+ %res2 = scf.parallel (%k2, %i2, %j2) = (%c0, %c0, %c0)
+ to (%c5, %c2, %c3) step (%c1, %c1, %c1) init(%init2) -> f32 {
+ %b = memref.load %B[%i2, %j2, %k2] : memref<2x3x5xf32>
+ scf.reduce(%b : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %prod = arith.mulf %lhs, %rhs : f32
+ scf.reduce.return %prod : f32
+ }
+ }
+
+ return %res1, %res2 : f32, f32
+}
+
+// CHECK-LABEL: func @fuse_three_cycle_reductions
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[C5:.*]] = arith.constant 5 : index
+// CHECK: %[[INIT1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK: %[[INIT2:.*]] = arith.constant 2.000000e+00 : f32
+// CHECK: %[[RES:.*]]:2 = scf.parallel (%[[I:.*]], %[[J:.*]], %[[K:.*]]) = (%[[C0]], %[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C2]], %[[C3]], %[[C5]]) step (%[[C1]], %[[C1]], %[[C1]])
+// CHECK-SAME: init (%[[INIT1]], %[[INIT2]]) -> (f32, f32)
+// CHECK: %[[AELT:.*]] = memref.load %{{.*}}{{\[}}%[[I]], %[[J]], %[[K]]{{\]}} : memref<2x3x5xf32>
+// CHECK: %[[BELT:.*]] = memref.load %{{.*}}{{\[}}%[[I]], %[[J]], %[[K]]{{\]}} : memref<2x3x5xf32>
+// CHECK: scf.reduce(%[[AELT]], %[[BELT]] : f32, f32) {
+// CHECK: ^bb0
+// CHECK: ^bb0
+// CHECK: return %[[RES]]#0, %[[RES]]#1 : f32, f32
>From b161e59802d69f8faab6011364ddba01ba4bb55c Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Mon, 15 Jun 2026 13:26:23 +0100
Subject: [PATCH 5/7] Addressed comments-2
Change-Id: I89a4a886f142684b309ef5b569ecc9793136539a
---
.../SCF/Transforms/ParallelLoopFusion.cpp | 141 +++++++++++-------
.../Dialect/SCF/parallel-loop-fusion.mlir | 46 ++++++
2 files changed, 130 insertions(+), 57 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 63795234c00a9..4345532593613 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -37,9 +37,11 @@
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Debug.h"
#include <numeric>
#include <optional>
#include <tuple>
+#define DEBUG_TYPE "parallel-loop-fusion"
namespace mlir {
#define GEN_PASS_DEF_SCFPARALLELLOOPFUSION
@@ -56,6 +58,15 @@ static bool hasNestedParallelOp(ParallelOp ploop) {
return walkResult.wasInterrupted();
}
+#ifndef NDEBUG
+template <class T>
+static inline std::string toString(T &cnt) {
+ std::stringstream ss;
+ std::copy(cnt.begin(), cnt.end(), std::ostream_iterator<int64_t>(ss, ""));
+ return ss.str();
+}
+#endif
+
/// Verify equal iteration spaces.
static bool equalIterationSpaces(ParallelOp firstPloop,
ParallelOp secondPloop) {
@@ -797,17 +808,17 @@ struct llvm::DenseMapInfo<LoopIV> {
};
// Returns vector of candidate permutation indices vectors,
-// can be empty
+// can be empty. Caps the number of extra candidate permutations
+// explored to avoid combinatorial explosion. This makes the search
+// intentionally incomplete.
static SmallVector<SmallVector<int64_t>>
computeCandidateInterchangePermutations(ParallelOp &firstPloop,
ParallelOp &secondPloop,
- int axesLimit = 5) {
- SmallVector<SmallVector<int64_t>> extraResults;
-
+ int permBudget = 120) {
// Check preconditions
if (firstPloop.getNumLoops() < 2 ||
firstPloop.getNumLoops() != secondPloop.getNumLoops())
- return extraResults;
+ return {};
SmallVector<LoopIV> firstIVs(firstPloop.getNumLoops());
SmallVector<LoopIV> secondIVs(secondPloop.getNumLoops());
@@ -834,14 +845,14 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
// Not a permutation shortcut
if (indices.size() == 1)
- return extraResults;
+ return {};
// Initialize with identity permutations
- SmallVector<int64_t> result(firstIVs.size());
- std::iota(result.begin(), result.end(), 0);
+ SmallVector<int64_t> basic(firstIVs.size());
+ std::iota(basic.begin(), basic.end(), 0);
if (indices.empty() && unique.size() == firstIVs.size())
- return extraResults;
+ return {};
if (indices.size() > 1) {
// Determine whether the iteration space of the first loop is a permutation
@@ -853,24 +864,26 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
if (fIdx != sIdx && firstIVs[fIdx] == secondIVs[sIdx] &&
remaps.end() == std::find(remaps.begin(), remaps.end(), sIdx)) {
remaps.push_back(sIdx);
+ break;
}
}
}
// Not a permutation
if (indices.size() != remaps.size())
- return extraResults;
+ return {};
// compose permutation indices
for (auto [from, to] : zip(indices, remaps)) {
- result[from] = to;
+ basic[from] = to;
}
- extraResults.push_back(result);
+ LLVM_DEBUG(llvm::dbgs()
+ << "Collected basic permutations: " << toString(basic) << "\n");
// All axes are unique, no further permutatons needed
if (unique.size() == firstIVs.size()) {
- return extraResults;
+ return {basic};
}
}
@@ -880,14 +893,13 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
"Expected at least two equal axes");
// Collect equal axes to groups
+ SmallVector<SmallVector<int64_t>> extraResults{basic};
SmallVector<SmallVector<int64_t>> groups;
for (auto iv : unique) {
SmallVector<int64_t> group;
for (unsigned index : llvm::seq(firstIVs.size())) {
- if (axesLimit && firstIVs[index] == iv) {
- axesLimit--;
+ if (firstIVs[index] == iv)
group.push_back(index);
- }
}
if (group.size() > 1)
groups.push_back(std::move(group));
@@ -895,8 +907,8 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
// Permute axes groups
SmallVector<SmallVector<int64_t>> rmpdGroups(groups);
- bool repeat = false;
- do {
+ bool repeat = true;
+ while (repeat && permBudget) {
repeat = false;
for (auto const &[group, groupRemaps] : zip(groups, rmpdGroups)) {
repeat |= std::next_permutation(groupRemaps.begin(), groupRemaps.end());
@@ -905,58 +917,31 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
}
if (repeat) {
- SmallVector<int64_t> extra(result);
+ SmallVector<int64_t> extra(basic);
for (auto const &[group, groupRemaps] : zip(groups, rmpdGroups)) {
for (auto [from, to] : zip(group, groupRemaps))
- extra[from] = result[to];
+ extra[from] = basic[to];
}
- if (result != extra)
+ if (basic != extra) {
+ LLVM_DEBUG(llvm::dbgs() << "Collected extra permutations: "
+ << toString(extra) << "\n");
+
extraResults.push_back(std::move(extra));
+ permBudget--;
+ }
}
- } while (repeat);
+ }
return extraResults;
}
/// Prepend operations of firstPloop's body into secondPloop's body.
/// Update secondPloop with new loop.
-static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
- OpBuilder builder,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
+static void applyLoopFusion(ParallelOp &firstPloop, ParallelOp &secondPloop,
+ OpBuilder &builder) {
+ DominanceInfo dom;
Block *block1 = firstPloop.getBody();
Block *block2 = secondPloop.getBody();
- IRMapping firstToSecondPloopIndices;
- firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
-
- if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
- mayAlias, builder)) {
- // If iteration space of the second parallel loop is a permutation of the
- // first one then interchange iteration space of the second parallel loop
- // and re-asses possibility of fusion.
- for (auto &perms :
- computeCandidateInterchangePermutations(firstPloop, secondPloop)) {
- OpBuilder::InsertionGuard guard(builder);
- auto newLoop = interchangeLoops(builder, secondPloop, perms);
- firstToSecondPloopIndices.clear();
- firstToSecondPloopIndices.map(block1->getArguments(),
- newLoop->getBody()->getArguments());
- if (!isFusionLegal(firstPloop, *newLoop, firstToSecondPloopIndices,
- mayAlias, builder)) {
- newLoop->erase();
- continue;
- }
-
- secondPloop.replaceAllUsesWith(newLoop->getResults());
- secondPloop->erase();
- secondPloop = *newLoop;
- block2 = secondPloop.getBody();
- goto fuseLabel;
- }
- return;
- }
-
-fuseLabel:
- DominanceInfo dom;
// We are fusing first loop into second, make sure there are no users of the
// first loop results between loops.
for (Operation *user : firstPloop->getUsers())
@@ -1013,6 +998,48 @@ static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
secondPloop = newSecondPloop;
}
+/// Check fusion pre-conditions and call fusion if it is possible
+static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
+ OpBuilder builder,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ Block *block1 = firstPloop.getBody();
+ Block *block2 = secondPloop.getBody();
+ IRMapping firstToSecondPloopIndices;
+ firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
+
+ if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
+ mayAlias, builder)) {
+ // If iteration space of the second parallel loop is a permutation of the
+ // first one then interchange iteration space of the second parallel loop
+ // and re-asses possibility of fusion.
+ for (auto &perms :
+ computeCandidateInterchangePermutations(firstPloop, secondPloop)) {
+ OpBuilder::InsertionGuard guard(builder);
+ LLVM_DEBUG(llvm::dbgs()
+ << "Applied permutation: " << toString(perms) << "\n");
+
+ auto newLoop = interchangeLoops(builder, secondPloop, perms);
+ firstToSecondPloopIndices.clear();
+ firstToSecondPloopIndices.map(block1->getArguments(),
+ newLoop->getBody()->getArguments());
+ if (!isFusionLegal(firstPloop, *newLoop, firstToSecondPloopIndices,
+ mayAlias, builder)) {
+ LLVM_DEBUG(llvm::dbgs() << "Rejected: " << newLoop << "\n");
+
+ newLoop->erase();
+ continue;
+ }
+
+ secondPloop.replaceAllUsesWith(newLoop->getResults());
+ secondPloop->erase();
+ applyLoopFusion(firstPloop, *newLoop, builder);
+ break;
+ }
+ return;
+ }
+ applyLoopFusion(firstPloop, secondPloop, builder);
+}
+
void mlir::scf::naivelyFuseParallelOps(
Region ®ion, llvm::function_ref<bool(Value, Value)> mayAlias) {
OpBuilder b(region);
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index 0d66e15a2b05a..e737d99d59bf5 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -1455,3 +1455,49 @@ func.func @fuse_three_cycle_reductions(%A: memref<2x3x5xf32>,
// CHECK: ^bb0
// CHECK: ^bb0
// CHECK: return %[[RES]]#0, %[[RES]]#1 : f32, f32
+
+// -----
+
+// Two duplicate axis groups are interleaved: the first loop has iteration
+// extents (2, 3, 2, 3), while the second loop visits the same space as
+// (3, 2, 3, 2). Fusion should find the permutation that maps the second loop
+// back to the first loop order and then fold both bodies into one loop.
+func.func @fuse_interleaved_duplicate_axes_permutation(
+ %out: memref<2x3x2x3xf32>) {
+ %tmp = memref.alloc() : memref<2x3x2x3xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %v = arith.constant 1.0 : f32
+
+ scf.parallel (%a, %b, %c, %d) = (%c0, %c0, %c0, %c0)
+ to (%c2, %c3, %c2, %c3) step (%c1, %c1, %c1, %c1) {
+ memref.store %v, %tmp[%a, %b, %c, %d] : memref<2x3x2x3xf32>
+ scf.reduce
+ }
+
+ scf.parallel (%b2, %a2, %d2, %c2v) = (%c0, %c0, %c0, %c0)
+ to (%c3, %c2, %c3, %c2) step (%c1, %c1, %c1, %c1) {
+ %x = memref.load %tmp[%a2, %b2, %c2v, %d2] : memref<2x3x2x3xf32>
+ memref.store %x, %out[%a2, %b2, %c2v, %d2] : memref<2x3x2x3xf32>
+ scf.reduce
+ }
+ return
+}
+
+// CHECK-NAME: func @fuse_interleaved_duplicate_axes_permutation
+// CHECK: %[[TMP:.*]] = memref.alloc() : memref<2x3x2x3xf32>
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK: scf.parallel (%[[A:.*]], %[[B:.*]], %[[C:.*]], %[[D:.*]]) = (%[[C0]], %[[C0]], %[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C2]], %[[C3]], %[[C2]], %[[C3]])
+// CHECK-SAME: step (%[[C1]], %[[C1]], %[[C1]], %[[C1]])
+// CHECK: memref.store %[[CST]], %[[TMP]]{{\[}}%[[A]], %[[B]], %[[C]], %[[D]]{{\]}} : memref<2x3x2x3xf32>
+// CHECK: %[[V0:.*]] = memref.load %[[TMP]]{{\[}}%[[A]], %[[B]], %[[C]], %[[D]]{{\]}} : memref<2x3x2x3xf32>
+// CHECK: memref.store %[[V0]], %{{.*}}{{\[}}%[[A]], %[[B]], %[[C]], %[[D]]{{\]}} : memref<2x3x2x3xf32>
+// CHECK: scf.reduce
+// CHECK-NOT: scf.parallel
>From e4f69c82b9bbdea05e3c72ff801f14131e9079ac Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Tue, 16 Jun 2026 10:52:21 +0100
Subject: [PATCH 6/7] Addressed comments-3
Change-Id: I006d48a82e9e2df8b2cf81923f595f3b8df88e0d
---
.../SCF/Transforms/ParallelLoopFusion.cpp | 26 +++----
.../Dialect/SCF/parallel-loop-fusion.mlir | 67 ++++++++++++++++++-
2 files changed, 80 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 4345532593613..7b1de79e030a9 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -740,11 +740,20 @@ static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
const IRMapping &firstToSecondPloopIndices,
llvm::function_ref<bool(Value, Value)> mayAlias,
OpBuilder &b) {
- return !hasNestedParallelOp(firstPloop) &&
- !hasNestedParallelOp(secondPloop) &&
- equalIterationSpaces(firstPloop, secondPloop) &&
- noIncompatibleDataDependencies(firstPloop, secondPloop,
- firstToSecondPloopIndices, mayAlias, b);
+ if (hasNestedParallelOp(firstPloop) || hasNestedParallelOp(secondPloop) ||
+ !equalIterationSpaces(firstPloop, secondPloop) ||
+ !noIncompatibleDataDependencies(firstPloop, secondPloop,
+ firstToSecondPloopIndices, mayAlias, b))
+ return false;
+
+ // We are fusing first loop into second, make sure there are no users of the
+ // first loop results between loops.
+ DominanceInfo dom;
+ for (Operation *user : firstPloop->getUsers()) {
+ if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
+ return false;
+ }
+ return true;
}
// Returns new parallel loop where two loops matching indices param are
@@ -939,15 +948,8 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
/// Update secondPloop with new loop.
static void applyLoopFusion(ParallelOp &firstPloop, ParallelOp &secondPloop,
OpBuilder &builder) {
- DominanceInfo dom;
Block *block1 = firstPloop.getBody();
Block *block2 = secondPloop.getBody();
- // We are fusing first loop into second, make sure there are no users of the
- // first loop results between loops.
- for (Operation *user : firstPloop->getUsers())
- if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
- return;
-
ValueRange inits1 = firstPloop.getInitVals();
ValueRange inits2 = secondPloop.getInitVals();
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index e737d99d59bf5..bc165bcaa014b 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -1486,7 +1486,7 @@ func.func @fuse_interleaved_duplicate_axes_permutation(
return
}
-// CHECK-NAME: func @fuse_interleaved_duplicate_axes_permutation
+// CHECK-LABEL: func @fuse_interleaved_duplicate_axes_permutation
// CHECK: %[[TMP:.*]] = memref.alloc() : memref<2x3x2x3xf32>
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[C1:.*]] = arith.constant 1 : index
@@ -1501,3 +1501,68 @@ func.func @fuse_interleaved_duplicate_axes_permutation(
// CHECK: memref.store %[[V0]], %{{.*}}{{\[}}%[[A]], %[[B]], %[[C]], %[[D]]{{\]}} : memref<2x3x2x3xf32>
// CHECK: scf.reduce
// CHECK-NOT: scf.parallel
+
+// -----
+
+// The first fusion candidate needs loop interchange to pass dependency checks,
+// but fusion must still be abandoned because the first loop result is used by
+// an operation between the loops.
+func.func @permuted_dominance_bail_chain(%a: memref<2x3xf32>,
+ %out: memref<2x3xf32>) -> f32 {
+ %tmp = memref.alloc() : memref<2x3xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %init = arith.constant 1.0 : f32
+ %res = scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c3)
+ step (%c1, %c1) init(%init) -> f32 {
+ %elt = memref.load %a[%i, %j] : memref<2x3xf32>
+ scf.reduce(%elt : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ scf.reduce.return %sum : f32
+ }
+ }
+ %between = arith.addf %res, %init : f32
+ scf.parallel (%j2, %i2) = (%c0, %c0) to (%c3, %c2)
+ step (%c1, %c1) {
+ memref.store %between, %tmp[%i2, %j2] : memref<2x3xf32>
+ scf.reduce
+ }
+ scf.parallel (%i3, %j3) = (%c0, %c0) to (%c2, %c3)
+ step (%c1, %c1) {
+ %x = memref.load %tmp[%i3, %j3] : memref<2x3xf32>
+ memref.store %x, %out[%i3, %j3] : memref<2x3xf32>
+ scf.reduce
+ }
+ return %between : f32
+}
+
+// CHECK-LABEL: func @permuted_dominance_bail_chain
+// CHECK: %[[TMP:.*]] = memref.alloc() : memref<2x3xf32>
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK: %[[RES:.*]] = scf.parallel (%[[A:.*]], %[[B:.*]]) = (%[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C2]], %[[C3]])
+// CHECK-SAME: step (%[[C1]], %[[C1]]) init (%[[CST]]) -> f32 {
+// CHECK: %[[V2:.*]] = memref.load %arg0{{\[}}%[[A]], %[[B]]{{\]}} : memref<2x3xf32>
+// CHECK: scf.reduce(%[[V2]] : f32) {
+// CHECK: ^bb0
+// CHECK: arith.addf
+// CHECK: scf.reduce.return
+// CHECK: }
+// CHECK: }
+// CHECK: %[[BETWEEN:.*]] = arith.addf %[[RES]], %[[CST]] : f32
+// CHECK: scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C3]], %[[C2]]) step (%[[C1]], %[[C1]]) {
+// CHECK: memref.store %[[BETWEEN]], %[[TMP]]{{\[}}%[[J]], %[[I]]{{\]}} : memref<2x3xf32>
+// CHECK: %[[V0:.*]] = memref.load %[[TMP]]{{\[}}%[[J]], %[[I]]{{\]}} : memref<2x3xf32>
+// CHECK: memref.store %[[V0]], %{{.*}}{{\[}}%[[J]], %[[I]]{{\]}} : memref<2x3xf32>
+// CHECK: scf.reduce
+// CHECK: }
+// CHECK-NOT: scf.parallel
+// CHECK: return %[[BETWEEN]]
>From 9b3bf7368ae3f21d00b7becfe9ac6dbf3361a55d Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Tue, 16 Jun 2026 14:29:35 +0100
Subject: [PATCH 7/7] Addressed comments-4
Change-Id: Ia4cc4b7a809049a8eb4770b0f089e6510734d7ce
---
.../SCF/Transforms/ParallelLoopFusion.cpp | 77 +++++++++----------
.../Dialect/SCF/parallel-loop-fusion.mlir | 77 +++++++++++++++++++
2 files changed, 112 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 7b1de79e030a9..1e879d81f3559 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -36,8 +36,9 @@
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/InterleavedRange.h"
-#include "llvm/Support/Debug.h"
+#include "llvm/Support/DebugLog.h"
#include <numeric>
#include <optional>
#include <tuple>
@@ -58,15 +59,6 @@ static bool hasNestedParallelOp(ParallelOp ploop) {
return walkResult.wasInterrupted();
}
-#ifndef NDEBUG
-template <class T>
-static inline std::string toString(T &cnt) {
- std::stringstream ss;
- std::copy(cnt.begin(), cnt.end(), std::ostream_iterator<int64_t>(ss, ""));
- return ss.str();
-}
-#endif
-
/// Verify equal iteration spaces.
static bool equalIterationSpaces(ParallelOp firstPloop,
ParallelOp secondPloop) {
@@ -887,8 +879,8 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
basic[from] = to;
}
- LLVM_DEBUG(llvm::dbgs()
- << "Collected basic permutations: " << toString(basic) << "\n");
+ LDBG() << "Collected basic permutations: "
+ << llvm::interleaved_array(basic);
// All axes are unique, no further permutatons needed
if (unique.size() == firstIVs.size()) {
@@ -932,8 +924,8 @@ computeCandidateInterchangePermutations(ParallelOp &firstPloop,
extra[from] = basic[to];
}
if (basic != extra) {
- LLVM_DEBUG(llvm::dbgs() << "Collected extra permutations: "
- << toString(extra) << "\n");
+ LDBG() << "Collected extra permutations: "
+ << llvm::interleaved_array(extra);
extraResults.push_back(std::move(extra));
permBudget--;
@@ -1009,37 +1001,38 @@ static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
IRMapping firstToSecondPloopIndices;
firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
- if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
- mayAlias, builder)) {
- // If iteration space of the second parallel loop is a permutation of the
- // first one then interchange iteration space of the second parallel loop
- // and re-asses possibility of fusion.
- for (auto &perms :
- computeCandidateInterchangePermutations(firstPloop, secondPloop)) {
- OpBuilder::InsertionGuard guard(builder);
- LLVM_DEBUG(llvm::dbgs()
- << "Applied permutation: " << toString(perms) << "\n");
-
- auto newLoop = interchangeLoops(builder, secondPloop, perms);
- firstToSecondPloopIndices.clear();
- firstToSecondPloopIndices.map(block1->getArguments(),
- newLoop->getBody()->getArguments());
- if (!isFusionLegal(firstPloop, *newLoop, firstToSecondPloopIndices,
- mayAlias, builder)) {
- LLVM_DEBUG(llvm::dbgs() << "Rejected: " << newLoop << "\n");
-
- newLoop->erase();
- continue;
- }
+ if (isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
+ mayAlias, builder)) {
+ applyLoopFusion(firstPloop, secondPloop, builder);
+ return;
+ }
- secondPloop.replaceAllUsesWith(newLoop->getResults());
- secondPloop->erase();
- applyLoopFusion(firstPloop, *newLoop, builder);
- break;
+ // If iteration space of the second parallel loop is a permutation of the
+ // first one then interchange iteration space of the second parallel loop
+ // and re-asses possibility of fusion.
+ for (auto &perms :
+ computeCandidateInterchangePermutations(firstPloop, secondPloop)) {
+ OpBuilder::InsertionGuard guard(builder);
+ LDBG() << "Applied permutation: " << llvm::interleaved_array(perms);
+
+ auto newLoop = interchangeLoops(builder, secondPloop, perms);
+ firstToSecondPloopIndices.clear();
+ firstToSecondPloopIndices.map(block1->getArguments(),
+ newLoop->getBody()->getArguments());
+ if (!isFusionLegal(firstPloop, *newLoop, firstToSecondPloopIndices,
+ mayAlias, builder)) {
+ LDBG() << "Rejected: " << newLoop;
+
+ newLoop->erase();
+ continue;
}
- return;
+
+ secondPloop.replaceAllUsesWith(newLoop->getResults());
+ secondPloop->erase();
+ secondPloop = *newLoop;
+ applyLoopFusion(firstPloop, secondPloop, builder);
+ break;
}
- applyLoopFusion(firstPloop, secondPloop, builder);
}
void mlir::scf::naivelyFuseParallelOps(
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index bc165bcaa014b..fa18817d076a8 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -1566,3 +1566,80 @@ func.func @permuted_dominance_bail_chain(%a: memref<2x3xf32>,
// CHECK: }
// CHECK-NOT: scf.parallel
// CHECK: return %[[BETWEEN]]
+
+// -----
+
+// The first pair only fuses after the second loop is interchanged.
+func.func @fuse_chain_after_interchanged_reduction(
+ %a: memref<2x3xf32>, %out: memref<2x3xf32>) -> (f32, f32) {
+ %tmp = memref.alloc() : memref<2x3xf32>
+ %mid = memref.alloc() : memref<2x3xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %init1 = arith.constant 1.0 : f32
+ %init2 = arith.constant 2.0 : f32
+
+ %r1 = scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c3)
+ step (%c1, %c1) init(%init1) -> f32 {
+ %x = memref.load %a[%i, %j] : memref<2x3xf32>
+ memref.store %x, %tmp[%i, %j] : memref<2x3xf32>
+ scf.reduce(%x : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ scf.reduce.return %sum : f32
+ }
+ }
+ %r2 = scf.parallel (%j2, %i2) = (%c0, %c0) to (%c3, %c2)
+ step (%c1, %c1) init(%init2) -> f32 {
+ %x = memref.load %tmp[%i2, %j2] : memref<2x3xf32>
+ memref.store %x, %mid[%i2, %j2] : memref<2x3xf32>
+ scf.reduce(%x : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %prod = arith.mulf %lhs, %rhs : f32
+ scf.reduce.return %prod : f32
+ }
+ }
+ scf.parallel (%i3, %j3) = (%c0, %c0) to (%c2, %c3)
+ step (%c1, %c1) {
+ %y = memref.load %mid[%i3, %j3] : memref<2x3xf32>
+ memref.store %y, %out[%i3, %j3] : memref<2x3xf32>
+ scf.reduce
+ }
+ return %r1, %r2 : f32, f32
+}
+
+// CHECK-LABEL: func @fuse_chain_after_interchanged_reduction
+// CHECK: %[[TMP:.*]] = memref.alloc() : memref<2x3xf32>
+// CHECK: %[[MID:.*]] = memref.alloc() : memref<2x3xf32>
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[INIT1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK: %[[INIT2:.*]] = arith.constant 2.000000e+00 : f32
+
+// CHECK: %[[RES:.*]]:2 = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[C0]], %[[C0]]) to (%[[C2]], %[[C3]])
+// CHECK-SAME: step (%[[C1]], %[[C1]]) init (%[[INIT1]], %[[INIT2]]) -> (f32, f32) {
+// CHECK: %[[X1:.*]] = memref.load %{{.*}}{{\[}}%[[I]], %[[J]]{{\]}} : memref<2x3xf32>
+// CHECK: memref.store %[[X1]], %[[TMP]]{{\[}}%[[I]], %[[J]]{{\]}} : memref<2x3xf32>
+// CHECK: %[[X2:.*]] = memref.load %[[TMP]]{{\[}}%[[I]], %[[J]]{{\]}} : memref<2x3xf32>
+// CHECK: memref.store %[[X2]], %[[MID]]{{\[}}%[[I]], %[[J]]{{\]}} : memref<2x3xf32>
+// CHECK: scf.reduce
+// CHECK: ^bb0
+// CHECK: arith.addf
+// CHECK: scf.reduce.return
+// CHECK: }, {
+// CHECK: ^bb0
+// CHECK: arith.mulf
+// CHECK: scf.reduce.return
+// CHECK: }
+// CHECK: }
+// CHECK: scf.parallel (%[[I3:.*]], %[[J3:.*]]) = (%[[C0]], %[[C0]]) to (%[[C2]], %[[C3]]) step (%[[C1]], %[[C1]]) {
+// CHECK: %[[Y1:.*]] = memref.load %[[MID]]{{\[}}%[[I3]], %[[J3]]{{\]}} : memref<2x3xf32>
+// CHECK: memref.store %[[Y1]], %{{.*}}{{\[}}%[[I3]], %[[J3]]{{\]}} : memref<2x3xf32>
+// CHECK: scf.reduce
+// CHECK: }
+// CHECK-NOT: scf-parallel
+// CHECK: return
More information about the Mlir-commits
mailing list