[Mlir-commits] [mlir] [MLIR][SCF] Support permutation-based parallel loop fusion (PR #203207)
Dmitriy Smirnov
llvmlistbot at llvm.org
Fri Jun 12 08:09:27 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/4] [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/4] 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/4] 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/4] 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
More information about the Mlir-commits
mailing list