[Mlir-commits] [mlir] [MLIR][SCF] Support permutation-based parallel loop fusion (PR #203207)
Dmitriy Smirnov
llvmlistbot at llvm.org
Tue Jun 16 06:31:03 PDT 2026
================
@@ -726,81 +740,216 @@ 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;
-// Interchange loops of the parallel loop, if there are just two loops
-static std::optional<ParallelOp> interchangeLoops(OpBuilder &builder,
- ParallelOp &loop) {
+ // 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;
+}
- 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()) {
- 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;
}
-/// 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) {
- Block *block1 = firstPloop.getBody();
- Block *block2 = secondPloop.getBody();
- IRMapping firstToSecondPloopIndices;
- firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
+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;
+ }
+};
- 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]) {
- firstToSecondPloopIndices.clear();
- firstToSecondPloopIndices.map(block1->getArguments(),
- llvm::reverse(block2->getArguments()));
- if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
- mayAlias, builder))
- return;
- auto newLoop = interchangeLoops(builder, secondPloop);
- secondPloop->erase();
- secondPloop = *newLoop;
- block2 = secondPloop.getBody();
- } else {
- return;
+template <>
+struct llvm::DenseMapInfo<LoopIV> {
+ 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. 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 permBudget = 120) {
+ // Check preconditions
+ if (firstPloop.getNumLoops() < 2 ||
+ firstPloop.getNumLoops() != secondPloop.getNumLoops())
+ return {};
+
+ SmallVector<LoopIV> firstIVs(firstPloop.getNumLoops());
+ SmallVector<LoopIV> secondIVs(secondPloop.getNumLoops());
+ llvm::SmallSetVector<LoopIV, 6> unique;
+ 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];
+ 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());
+ 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))
+ if (val)
+ indices.push_back(idx);
+
+ // Not a permutation shortcut
+ if (indices.size() == 1)
+ return {};
+
+ // Initialize with identity permutations
+ SmallVector<int64_t> basic(firstIVs.size());
+ std::iota(basic.begin(), basic.end(), 0);
+
+ if (indices.empty() && unique.size() == firstIVs.size())
+ return {};
+
+ 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);
+ break;
+ }
+ }
+ }
+
+ // Not a permutation
+ if (indices.size() != remaps.size())
+ return {};
+
+ // compose permutation indices
+ for (auto [from, to] : zip(indices, remaps)) {
+ basic[from] = to;
+ }
+
+ LLVM_DEBUG(llvm::dbgs()
----------------
d-smirnov wrote:
Amended
https://github.com/llvm/llvm-project/pull/203207
More information about the Mlir-commits
mailing list