[Mlir-commits] [mlir] [MLIR][SCF] Support permutation-based parallel loop fusion (PR #203207)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jun 11 02:19:45 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Dmitriy Smirnov (d-smirnov)
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/203207.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp (+189-25)
- (modified) mlir/test/Dialect/SCF/parallel-loop-fusion.mlir (+105)
``````````diff
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/203207
More information about the Mlir-commits
mailing list