[Mlir-commits] [mlir] [mlir][scf] Refactor and improve ParallelLoopFusion (PR #179284)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Feb 2 08:56:36 PST 2026
https://github.com/fabrizio-indirli created https://github.com/llvm/llvm-project/pull/179284
Refactor and extend the scf::ParalleLoopFusion pass:
- Refactor code, rename functions and add comments to improve readability
- Make the dependency analysis safer by checking for read-after-write dependencies also with vector.load/store & vector.transfer_read/write ops, in addition to memref.load/store, and bail out when other unsupported ops with memory effects are found.
- Extend the cases when the fusion is applied: allow fusing also when one of the two loops reads/writes to memory through a full view/alias of the buffer (read/written by the dual operation in the other loop) that can be trivially resolved, including rank-reducing full subviews.
>From 643538986a36adbf6c800fe4a3bea32d37882967 Mon Sep 17 00:00:00 2001
From: Fabrizio Indirli <fabrizio.indirli at arm.com>
Date: Thu, 29 Jan 2026 16:49:00 +0000
Subject: [PATCH] [mlir][scf] Refactor and improve ParallelLoopFusion
Refactor and extend the scf::ParalleLoopFusion pass:
- Refactor code, rename functions and add comments to improve
readability
- Make the dependency analysis safer by checking for read-after-write
dependencies also with vector.load/store & vector.transfer_read/write ops,
in addition to memref.load/store, and bail out when other unsupported ops
with memory effects are found.
- Extend the cases when the fusion is applied: allow fusing also
when one of the two loops reads/writes to memory through a full
view/alias of the buffer (read/written by the dual operation in the
other loop) that can be trivially resolved, including
rank-reducing full subviews.
---
.../SCF/Transforms/ParallelLoopFusion.cpp | 479 +++++++++++++++---
.../Dialect/SCF/parallel-loop-fusion.mlir | 335 +++++++++++-
2 files changed, 738 insertions(+), 76 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 4ea832177c4f9..f3e841dfb5dcf 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -14,14 +14,19 @@
#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/TypeSwitch.h"
+
namespace mlir {
#define GEN_PASS_DEF_SCFPARALLELLOOPFUSION
#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
@@ -55,110 +60,442 @@ static bool equalIterationSpaces(ParallelOp firstPloop,
matchOperands(firstPloop.getStep(), secondPloop.getStep());
}
-/// Checks if the parallel loops have mixed access to the same buffers. Returns
-/// `true` if the first parallel loop writes to the same indices that the second
-/// loop reads.
-static bool haveNoReadsAfterWriteExceptSameIndex(
+/// Check if both operations are the same type of memory write op and
+/// write to the same memory location (same buffer and same indices).
+static bool opsWriteSameMemLocation(Operation *op1, Operation *op2) {
+ if (!op1 || !op2 || op1->getName() != op2->getName())
+ return false;
+ // support only these memory-writing ops for now
+ if (!isa<memref::StoreOp, vector::TransferWriteOp, vector::StoreOp>(op1))
+ return false;
+ bool opsAreIdentical =
+ llvm::TypeSwitch<Operation *, bool>(op1)
+ .Case([&](memref::StoreOp storeOp1) {
+ auto storeOp2 = cast<memref::StoreOp>(op2);
+ return (storeOp1.getMemRef() == storeOp2.getMemRef()) &&
+ (storeOp1.getIndices() == storeOp2.getIndices());
+ })
+ .Case([&](vector::TransferWriteOp writeOp1) {
+ auto writeOp2 = cast<vector::TransferWriteOp>(op2);
+ return (writeOp1.getBase() == writeOp2.getBase()) &&
+ (writeOp1.getIndices() == writeOp2.getIndices()) &&
+ (writeOp1.getMask() == writeOp2.getMask()) &&
+ (writeOp1.getValueToStore().getType() ==
+ writeOp2.getValueToStore().getType()) &&
+ (writeOp1.getInBounds() == writeOp2.getInBounds());
+ })
+ .Case([&](vector::StoreOp vecStoreOp1) {
+ auto vecStoreOp2 = cast<vector::StoreOp>(op2);
+ return (vecStoreOp1.getBase() == vecStoreOp2.getBase()) &&
+ (vecStoreOp1.getIndices() == vecStoreOp2.getIndices()) &&
+ (vecStoreOp1.getValueToStore().getType() ==
+ vecStoreOp2.getValueToStore().getType()) &&
+ (vecStoreOp1.getAlignment() == vecStoreOp2.getAlignment()) &&
+ (vecStoreOp1.getNontemporal() ==
+ vecStoreOp2.getNontemporal());
+ })
+ .Default([](Operation *) { return false; });
+ return opsAreIdentical;
+}
+
+/// Check if val1 (from the first parallel loop) and val2 (from the
+/// second) are equivalent, considering the mapping of induction variables from
+/// the first to the second parallel loop.
+static bool valsAreEquivalent(Value val1, Value val2,
+ const IRMapping &loopsIVsMap) {
+ if (val1 == val2 || loopsIVsMap.lookupOrDefault(val1) == val2 ||
+ loopsIVsMap.lookupOrDefault(val2) == val1)
+ return true;
+ Operation *val1DefOp = val1.getDefiningOp();
+ Operation *val2DefOp = val2.getDefiningOp();
+ if (!val1DefOp || !val2DefOp)
+ return false;
+ if (!isMemoryEffectFree(val1DefOp) || !isMemoryEffectFree(val2DefOp))
+ return false;
+ return OperationEquivalence::isEquivalentTo(
+ val1DefOp, val2DefOp,
+ [&](Value v1, Value v2) {
+ return success(loopsIVsMap.lookupOrDefault(v1) == v2 ||
+ loopsIVsMap.lookupOrDefault(v2) == v1);
+ },
+ /*markEquivalent=*/nullptr, OperationEquivalence::Flags::IgnoreLocations);
+}
+
+/// Return the base memref value used by the given memory op.
+template <typename OpTy>
+static Value getBaseMemref(OpTy op) {
+ return llvm::TypeSwitch<Operation *, Value>(op.getOperation())
+ .Case([&](memref::LoadOp load) { return load.getMemRef(); })
+ .Case([&](memref::StoreOp store) { return store.getMemRef(); })
+ .Case([&](vector::TransferReadOp read) { return read.getBase(); })
+ .Case([&](vector::TransferWriteOp write) { return write.getBase(); })
+ .Case([&](vector::LoadOp load) { return load.getBase(); })
+ .Case([&](vector::StoreOp store) { return store.getBase(); })
+ .Default([](Operation *) { return Value(); });
+}
+
+/// Recognize scalar memref.load of an element produced by a
+/// vector.transfer_write (optionally through a rank-reducing, unit-stride
+/// subview) of the same buffer. This covers the pattern where a vector write
+/// stores a full lane pack and a subsequent loop iterates over the lane
+/// dimension with scalar loads. EXAMPLE:
+/// vector.transfer_write %V, %arg[%x, %y, ..., 0] {in_bounds = [true]} :
+/// vector<4xf32>, memref<4xf32, strided<[1], offset: ?>>
+/// scf.for %iter = %c0 to %c4 step %c1 iter_args(...) -> (f32) {
+/// %0 = memref.load %arg[%x, %y, ..., %iter] : memref<1x128x16x4xf32>
+/// ...
+/// }
+///
+static bool loadMatchesVectorWrite(memref::LoadOp loadOp,
+ vector::TransferWriteOp writeOp,
+ const IRMapping &ivsMap) {
+ auto vecTy = dyn_cast<VectorType>(writeOp.getVector().getType());
+ if (!vecTy || vecTy.getRank() != 1)
+ return false;
+
+ Value base = writeOp.getBase();
+ MemrefValue baseMemref = nullptr;
+ SmallVector<OpFoldResult> offsets;
+ SmallVector<OpFoldResult> sizes;
+ auto ctx = loadOp.getContext();
+ if (auto subView = base.getDefiningOp<memref::SubViewOp>()) {
+ if (!subView.hasUnitStride())
+ return false;
+ baseMemref = cast<MemrefValue>(subView.getSource());
+ offsets = llvm::to_vector(subView.getMixedOffsets());
+ sizes = llvm::to_vector(subView.getMixedSizes());
+ } else {
+ baseMemref = dyn_cast<MemrefValue>(base);
+ if (!baseMemref)
+ return false;
+ // Fabricate rank-1 view matching the vector length at the end.
+ sizes = SmallVector<OpFoldResult>{
+ IntegerAttr::get(IndexType::get(ctx), vecTy.getDimSize(0))};
+ offsets = SmallVector<OpFoldResult>{
+ writeOp.getIndices().empty()
+ ? OpFoldResult(IntegerAttr::get(IndexType::get(ctx), 0))
+ : writeOp.getIndices().front()};
+ }
+
+ if (sizes.empty() || !isConstantIntValue(sizes.back(), vecTy.getDimSize(0)))
+ return false;
+
+ if (loadOp.getMemref() != baseMemref)
+ return false;
+
+ auto loadIndices = loadOp.getIndices();
+ if (loadIndices.size() != sizes.size())
+ return false;
+
+ // All leading dims size-1; offsets must match load indices.
+ for (unsigned i = 0; i + 1 < sizes.size(); ++i) {
+ if (!isConstantIntValue(sizes[i], 1))
+ return false;
+ if (auto attr = offsets[i].dyn_cast<Attribute>()) {
+ auto cst = dyn_cast<IntegerAttr>(attr);
+ if (!cst || cst.getInt() != 0 || !matchPattern(loadIndices[i], m_Zero()))
+ return false;
+ } else if (auto val = offsets[i].dyn_cast<Value>()) {
+ if (!valsAreEquivalent(val, loadIndices[i], ivsMap))
+ return false;
+ } else {
+ return false;
+ }
+ }
+
+ // transfer_write must start at lane 0 of the subview.
+ if (writeOp.getIndices().size() != 1 ||
+ !matchPattern(writeOp.getIndices().front(), m_Zero()))
+ return false;
+
+ // Last load index must be an scf.for induction variable iterating [0,
+ // vecLen).
+ auto laneIdx = dyn_cast_or_null<BlockArgument>(loadIndices.back());
+ auto forOp = laneIdx ? dyn_cast<scf::ForOp>(laneIdx.getOwner()->getParentOp())
+ : nullptr;
+ if (!forOp || laneIdx != forOp.getInductionVar())
+ return false;
+ auto lb = forOp.getLowerBound().getDefiningOp<arith::ConstantIndexOp>();
+ auto ub = forOp.getUpperBound().getDefiningOp<arith::ConstantIndexOp>();
+ auto step = forOp.getStep().getDefiningOp<arith::ConstantIndexOp>();
+ if (!lb || lb.value() != 0 || !step || step.value() != 1 || !ub)
+ return false;
+ if (ub.value() != vecTy.getDimSize(0))
+ return false;
+
+ return true;
+}
+
+/// Check if both operations access the same positions of the same
+/// buffer, but one of the two does it through a rank-reducing full subview of
+/// the buffer (the other's base). EXAMPLE:
+/// memref.store %a, %buf[%c0, %i, %j] : memref<1x2x2xf32>
+/// %alias = memref.subview %buf[0, 0, 0][1, 2, 2][1, 1, 1]: memref<1x2x2xf32>
+/// to memref<2x2xf32>
+/// %val = memref.load %alias[%i, %j] : memref<2x2xf32>
+template <typename OpTy1, typename OpTy2>
+static bool opsAccessSameIndicesViaRankReducingSubview(
+ OpTy1 op1, OpTy2 op2, const IRMapping &firstToSecondPloopIVsMap) {
+ auto base1 = cast<MemrefValue>(getBaseMemref(op1));
+ auto base2 = cast<MemrefValue>(getBaseMemref(op2));
+ if (!base1 || !base2)
+ return false;
+
+ auto accessThroughTrivialSubviewIsSame =
+ [](memref::SubViewOp subView, ValueRange subViewAccess,
+ ValueRange sourceAccess, const IRMapping &ivsMap) -> bool {
+ if (!subView.hasZeroOffset() || !subView.hasUnitStride())
+ return false;
+
+ MemRefType srcType = subView.getSourceType();
+ MemRefType resType = subView.getType();
+ unsigned srcRank = srcType.getRank();
+ unsigned resRank = resType.getRank();
+ if (sourceAccess.size() != srcRank || subViewAccess.size() != resRank)
+ return false;
+
+ auto staticSizes = subView.getStaticSizes();
+ auto droppedDims =
+ mlir::computeRankReductionMask(srcType.getShape(), resType.getShape());
+ if (!droppedDims || (droppedDims->size() != srcRank - resRank))
+ return false;
+
+ unsigned resPos = 0;
+ for (unsigned srcPos = 0; srcPos < srcRank; ++srcPos) {
+ if (droppedDims->contains(srcPos)) {
+ if (staticSizes[srcPos] != 1 ||
+ !matchPattern(sourceAccess[srcPos], m_Zero()))
+ return false;
+ continue;
+ }
+ if (resPos >= resRank || !valsAreEquivalent(subViewAccess[resPos],
+ sourceAccess[srcPos], ivsMap))
+ return false;
+ ++resPos;
+ }
+ return resPos == resRank;
+ };
+
+ // Case 1: op1 uses a subview of op2's base.
+ if (auto subView = base1.template getDefiningOp<memref::SubViewOp>();
+ subView &&
+ memref::isSameViewOrTrivialAlias(
+ base2, cast<MemrefValue>(subView.getSource())) &&
+ accessThroughTrivialSubviewIsSame(subView, op1.getIndices(),
+ op2.getIndices(),
+ firstToSecondPloopIVsMap))
+ return true;
+
+ // Case 2: op2 uses a subview of op1's base.
+ if (auto subView = base2.template getDefiningOp<memref::SubViewOp>();
+ subView &&
+ memref::isSameViewOrTrivialAlias(
+ base1, cast<MemrefValue>(subView.getSource())) &&
+ accessThroughTrivialSubviewIsSame(subView, op2.getIndices(),
+ op1.getIndices(),
+ firstToSecondPloopIVsMap))
+ return true;
+
+ return false;
+}
+
+/// Check if both memory read/write operations access the same indices
+/// (considering also the mapping of induction variables from the first to the
+/// second parallel loop).
+template <typename OpTy1, typename OpTy2>
+static bool opsAccessSameIndices(OpTy1 op1, OpTy2 op2,
+ const IRMapping &loopsIVsMap) {
+ auto indices1 = op1.getIndices();
+ auto indices2 = op2.getIndices();
+ if (indices1.size() != indices2.size())
+ return opsAccessSameIndicesViaRankReducingSubview(op1, op2, loopsIVsMap);
+ for (auto [idx1, idx2] : llvm::zip(indices1, indices2)) {
+ if (!valsAreEquivalent(idx1, idx2, loopsIVsMap))
+ return false;
+ }
+ return true;
+}
+
+/// Check if the loadOp reads from the same memory location (same buffer,
+/// same indices and same properties) as written by the storeOp.
+static bool loadsFromSameMemoryLocationWrittenBy(
+ Operation *loadOp, Operation *storeOp,
+ const IRMapping &firstToSecondPloopIVsMap) {
+ if (!loadOp || !storeOp)
+ return false;
+ // support only these memory-reading ops for now
+ if (!isa<memref::LoadOp, vector::TransferReadOp, vector::LoadOp>(loadOp))
+ return false;
+ bool accessSameMemory =
+ llvm::TypeSwitch<Operation *, bool>(loadOp)
+ .Case([&](memref::LoadOp memLoadOp) {
+ if (auto memStoreOp = dyn_cast<memref::StoreOp>(storeOp))
+ return opsAccessSameIndices(memLoadOp, memStoreOp,
+ firstToSecondPloopIVsMap);
+ if (auto vecWriteOp = dyn_cast<vector::TransferWriteOp>(storeOp))
+ return loadMatchesVectorWrite(memLoadOp, vecWriteOp,
+ firstToSecondPloopIVsMap);
+ return false;
+ })
+ .Case([&](vector::TransferReadOp vecReadOp) {
+ auto vecWriteOp = dyn_cast<vector::TransferWriteOp>(storeOp);
+ if (!vecWriteOp)
+ return false;
+ return opsAccessSameIndices(vecReadOp, vecWriteOp,
+ firstToSecondPloopIVsMap) &&
+ (vecReadOp.getMask() == vecWriteOp.getMask()) &&
+ (vecReadOp.getInBounds() == vecWriteOp.getInBounds());
+ })
+ .Case([&](vector::LoadOp vecLoadOp) {
+ auto vecStoreOp = dyn_cast<vector::StoreOp>(storeOp);
+ if (!vecStoreOp)
+ return false;
+ return opsAccessSameIndices(vecLoadOp, vecStoreOp,
+ firstToSecondPloopIVsMap) &&
+ (vecLoadOp.getAlignment() == vecStoreOp.getAlignment());
+ })
+ .Default([](Operation *) { return false; });
+ return accessSameMemory;
+}
+
+static Value getStoreOpTargetBuffer(Operation *op) {
+ return llvm::TypeSwitch<Operation *, Value>(op)
+ .Case([&](memref::StoreOp storeOp) { return storeOp.getMemRef(); })
+ .Case([&](vector::TransferWriteOp writeOp) { return writeOp.getBase(); })
+ .Case([&](vector::StoreOp vecStoreOp) { return vecStoreOp.getBase(); })
+ .Default([](Operation *) { return Value(); });
+}
+
+/// Check that the parallel loops have no mixed access to the same buffers.
+/// Return `true` if the second parallel loop does not read or write the buffers
+/// written by the first loop using different indices.
+static bool haveNoDataDependenciesExceptSameIndex(
ParallelOp firstPloop, ParallelOp secondPloop,
const IRMapping &firstToSecondPloopIndices,
llvm::function_ref<bool(Value, Value)> mayAlias) {
- DenseMap<Value, SmallVector<ValueRange, 1>> bufferStores;
- SmallVector<Value> bufferStoresVec;
- firstPloop.getBody()->walk([&](memref::StoreOp store) {
- bufferStores[store.getMemRef()].push_back(store.getIndices());
- bufferStoresVec.emplace_back(store.getMemRef());
- });
- auto walkResult = secondPloop.getBody()->walk([&](memref::LoadOp load) {
- Value loadMem = load.getMemRef();
+ // Map buffers to their store/write ops in the firstPloop
+ DenseMap<Value, SmallVector<Operation *>> bufferStoresInFirstPloop;
+ // Record all the memory buffers used in store/write ops found in firstPloop
+ llvm::SmallSetVector<Value, 4> buffersWrittenInFirstPloop;
+
+ // Walk the first parallel loop to collect all store/write ops and their
+ // target buffers
+ if (firstPloop.getBody()
+ ->walk([&](Operation *op) {
+ auto memOpInterf = dyn_cast_if_present<MemoryEffectOpInterface>(op);
+ // ignore ops that don't write to memory
+ if (!memOpInterf ||
+ (!memOpInterf.hasEffect<MemoryEffects::Write>() &&
+ !memOpInterf.hasEffect<MemoryEffects::Free>()))
+ return WalkResult::advance();
+
+ // only these memory-writing ops are supported for now:
+ // memref.store, vector.transfer_write, vector.store
+ Value storeOpBase = getStoreOpTargetBuffer(op);
+ if (!storeOpBase)
+ return WalkResult::interrupt();
+
+ // Expect the base operand to be a Memref
+ MemrefValue storeOpBaseMemref = dyn_cast<MemrefValue>(storeOpBase);
+ if (!storeOpBaseMemref)
+ return WalkResult::interrupt();
+ // Get the original memref buffer, skipping full view-like ops
+ Value buffer =
+ memref::skipFullyAliasingOperations(storeOpBaseMemref);
+ bufferStoresInFirstPloop[buffer].push_back(op);
+ buffersWrittenInFirstPloop.insert(buffer);
+ return WalkResult::advance();
+ })
+ .wasInterrupted())
+ return false;
+
+ // Walk the second parallel loop to check load/read ops against the stores
+ // collected from the first parallel loop: the loops can be fused only if in
+ // the 2nd loop there are no loads/stores from/yo the buffers written in the
+ // 1st loop, except when on the same exact memory location (same indices) as
+ // written in the 1st loop.
+ auto walkResult = secondPloop.getBody()->walk([&](Operation *loadOp) {
+ auto memOpInterf = dyn_cast_if_present<MemoryEffectOpInterface>(loadOp);
+ // ignore ops that don't read from memory
+ if (!memOpInterf || (!memOpInterf.hasEffect<MemoryEffects::Read>()))
+ return WalkResult::advance();
+ // support only these memory-reading ops for now
+ if (!isa<memref::LoadOp, vector::TransferReadOp, vector::LoadOp>(loadOp) ||
+ !isa<MemrefValue>(loadOp->getOperand(0)))
+ return WalkResult::interrupt();
+
+ MemrefValue loadOpBase = cast<MemrefValue>(loadOp->getOperand(0));
+ MemrefValue loadedOrigBuf = memref::skipFullyAliasingOperations(loadOpBase);
// Stop if the memref is defined in secondPloop body. Careful alias analysis
// is needed.
- auto *memrefDef = loadMem.getDefiningOp();
- if (memrefDef && memrefDef->getBlock() == load->getBlock())
+ auto *memrefDef = loadedOrigBuf.getDefiningOp();
+ if (memrefDef && secondPloop->isAncestor(memrefDef))
return WalkResult::interrupt();
- for (Value store : bufferStoresVec)
- if (store != loadMem && mayAlias(store, loadMem))
+ for (Value storedMem : buffersWrittenInFirstPloop)
+ if (storedMem != loadedOrigBuf && mayAlias(storedMem, loadedOrigBuf))
return WalkResult::interrupt();
- auto write = bufferStores.find(loadMem);
- if (write == bufferStores.end())
+ auto writeOpsIt = bufferStoresInFirstPloop.find(loadedOrigBuf);
+ if (writeOpsIt == bufferStoresInFirstPloop.end())
return WalkResult::advance();
+ // Store/write ops to this buffer in the firstPloop
+ auto &writeOps = writeOpsIt->second;
- // Check that at last one store was retrieved
- if (write->second.empty())
- return WalkResult::interrupt();
+ // If the first loop has no writes to this buffer, continue
+ if (writeOps.empty())
+ return WalkResult::advance();
- auto storeIndices = write->second.front();
+ Operation *writeOp = writeOps.front();
- // Multiple writes to the same memref are allowed only on the same indices
- for (const auto &othStoreIndices : write->second) {
- if (othStoreIndices != storeIndices)
- return WalkResult::interrupt();
- }
+ // In the first parallel loop, multiple writes to the same memref are
+ // allowed only on the same memory location
+ if (!llvm::all_of(writeOps, [&](Operation *otherWriteOp) {
+ return opsWriteSameMemLocation(writeOp, otherWriteOp);
+ }))
+ return WalkResult::interrupt();
- // Check that the load indices of secondPloop coincide with store indices of
- // firstPloop for the same memrefs.
- auto loadIndices = load.getIndices();
- if (storeIndices.size() != loadIndices.size())
+ // Check that the load in secondPloop reads from the same memory location as
+ // written by the corresponding store in firstPloop
+ if (!loadsFromSameMemoryLocationWrittenBy(loadOp, writeOp,
+ firstToSecondPloopIndices))
return WalkResult::interrupt();
- for (int i = 0, e = storeIndices.size(); i < e; ++i) {
- if (firstToSecondPloopIndices.lookupOrDefault(storeIndices[i]) !=
- loadIndices[i]) {
- auto *storeIndexDefOp = storeIndices[i].getDefiningOp();
- auto *loadIndexDefOp = loadIndices[i].getDefiningOp();
- if (storeIndexDefOp && loadIndexDefOp) {
- if (!isMemoryEffectFree(storeIndexDefOp))
- return WalkResult::interrupt();
- if (!isMemoryEffectFree(loadIndexDefOp))
- return WalkResult::interrupt();
- if (!OperationEquivalence::isEquivalentTo(
- storeIndexDefOp, loadIndexDefOp,
- [&](Value storeIndex, Value loadIndex) {
- if (firstToSecondPloopIndices.lookupOrDefault(storeIndex) !=
- firstToSecondPloopIndices.lookupOrDefault(loadIndex))
- return failure();
- else
- return success();
- },
- /*markEquivalent=*/nullptr,
- OperationEquivalence::Flags::IgnoreLocations)) {
- return WalkResult::interrupt();
- }
- } else {
- return WalkResult::interrupt();
- }
- }
- }
+
return WalkResult::advance();
});
+
return !walkResult.wasInterrupted();
}
-/// Analyzes dependencies in the most primitive way by checking simple read and
-/// write patterns.
-static LogicalResult
-verifyDependencies(ParallelOp firstPloop, ParallelOp secondPloop,
- const IRMapping &firstToSecondPloopIndices,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- if (!haveNoReadsAfterWriteExceptSameIndex(
+/// Check that in each loop there are no read ops on the buffers written
+/// by the other loop, except when reading from the same exact memory location
+/// (same indices) as written in the other loop.
+static bool noIncompatibleDataDependencies(
+ ParallelOp firstPloop, ParallelOp secondPloop,
+ const IRMapping &firstToSecondPloopIndices,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ if (!haveNoDataDependenciesExceptSameIndex(
firstPloop, secondPloop, firstToSecondPloopIndices, mayAlias))
- return failure();
+ return false;
IRMapping secondToFirstPloopIndices;
secondToFirstPloopIndices.map(secondPloop.getBody()->getArguments(),
firstPloop.getBody()->getArguments());
- return success(haveNoReadsAfterWriteExceptSameIndex(
- secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias));
+ return haveNoDataDependenciesExceptSameIndex(
+ secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias);
}
+/// Check if fusion of the two parallel loops is legal:
+/// i.e. no nested parallel loops, equal iteration spaces,
+/// and no incompatible data dependencies between the loops.
static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
const IRMapping &firstToSecondPloopIndices,
llvm::function_ref<bool(Value, Value)> mayAlias) {
return !hasNestedParallelOp(firstPloop) &&
!hasNestedParallelOp(secondPloop) &&
equalIterationSpaces(firstPloop, secondPloop) &&
- succeeded(verifyDependencies(firstPloop, secondPloop,
- firstToSecondPloopIndices, mayAlias));
+ noIncompatibleDataDependencies(firstPloop, secondPloop,
+ firstToSecondPloopIndices, mayAlias);
}
/// Prepends operations of firstPloop's body into secondPloop's body.
diff --git a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
index 0d4ea6f20e8d9..d28b5ee6b2feb 100644
--- a/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
+++ b/mlir/test/Dialect/SCF/parallel-loop-fusion.mlir
@@ -314,23 +314,23 @@ func.func @do_not_fuse_unmatching_read_write_patterns(
// -----
-func.func @do_not_fuse_loops_with_memref_defined_in_loop_bodies() {
+func.func @do_not_fuse_loops_with_nonfull_alias_defined_in_loop_bodies() {
%c2 = arith.constant 2 : index
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%buffer = memref.alloc() : memref<2x2xf32>
- scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c1) step (%c1, %c1) {
scf.reduce
}
- scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
- %A = memref.subview %buffer[%c0, %c0][%c2, %c2][%c1, %c1]
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c1) step (%c1, %c1) {
+ %A = memref.subview %buffer[%c0, %j][%c2, %c2][%c1, %c1]
: memref<2x2xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%A_elem = memref.load %A[%i, %j] : memref<?x?xf32, strided<[?, ?], offset: ?>>
scf.reduce
}
return
}
-// CHECK-LABEL: func @do_not_fuse_loops_with_memref_defined_in_loop_bodies
+// CHECK-LABEL: func @do_not_fuse_loops_with_nonfull_alias_defined_in_loop_bodies
// CHECK: scf.parallel
// CHECK: scf.parallel
@@ -604,6 +604,331 @@ func.func @do_not_fuse_affine_apply_to_non_ind_var(
// -----
+func.func @fuse_trivial_rank_reducing_subview() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c1fp = arith.constant 1.0 : f32
+ %buf = memref.alloc() : memref<1x2x2xf32>
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ memref.store %c1fp, %buf[%c0, %i, %j] : memref<1x2x2xf32>
+ scf.reduce
+ }
+ %sub = memref.subview %buf[0, 0, 0][1, 2, 2][1, 1, 1]
+ : memref<1x2x2xf32> to memref<2x2xf32>
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ %v = memref.load %sub[%i, %j] : memref<2x2xf32>
+ memref.store %v, %buf[%c0, %i, %j] : memref<1x2x2xf32>
+ scf.reduce
+ }
+ memref.dealloc %buf : memref<1x2x2xf32>
+ return
+}
+// CHECK-LABEL: func @fuse_trivial_rank_reducing_subview
+// CHECK: %[[BUF:.*]] = memref.alloc() : memref<1x2x2xf32>
+// CHECK: %[[SUB:.*]] = memref.subview %[[BUF]]
+// CHECK: scf.parallel
+// CHECK: memref.store {{.*}}, %[[BUF]]
+// CHECK: %[[L:.*]] = memref.load %[[SUB]]
+// CHECK: memref.store %[[L]], %[[BUF]]
+// CHECK-NOT: scf.parallel
+// CHECK: memref.dealloc %[[BUF]] : memref<1x2x2xf32>
+
+// -----
+
+func.func @do_not_fuse_nontrivial_subview_offset() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c1fp = arith.constant 1.0 : f32
+ %buf = memref.alloc() : memref<2x2x2xf32>
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ memref.store %c1fp, %buf[%c0, %i, %j] : memref<2x2x2xf32>
+ scf.reduce
+ }
+ %sub = memref.subview %buf[1, 0, 0][1, 2, 2][1, 1, 1]
+ : memref<2x2x2xf32> to memref<2x2xf32, strided<[2, 1], offset: 4>>
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ %v = memref.load %sub[%i, %j]
+ : memref<2x2xf32, strided<[2, 1], offset: 4>>
+ memref.store %v, %buf[%c0, %i, %j] : memref<2x2x2xf32>
+ scf.reduce
+ }
+ memref.dealloc %buf : memref<2x2x2xf32>
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_nontrivial_subview_offset
+// CHECK: scf.parallel
+// CHECK: scf.parallel
+
+// -----
+
+func.func @fuse_vector_load_store(%A: memref<4x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %vec0 = arith.constant dense<0.0> : vector<4xf32>
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ vector.store %vec0, %A[%i, %c0] : memref<4x4xf32>, vector<4xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ %v = vector.load %A[%i, %c0] : memref<4x4xf32>, vector<4xf32>
+ vector.store %v, %A[%i, %c0] : memref<4x4xf32>, vector<4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @fuse_vector_load_store
+// CHECK: scf.parallel (%[[I:.*]]) = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
+// CHECK: vector.store
+// CHECK: %[[V:.*]] = vector.load
+// CHECK: vector.store %[[V]]
+// CHECK-NOT: scf.parallel
+
+// -----
+
+func.func @do_not_fuse_vector_different_indices(%A: memref<4x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %vec0 = arith.constant dense<0.0> : vector<4xf32>
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ vector.store %vec0, %A[%i, %c0] : memref<4x4xf32>, vector<4xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ %j = affine.apply affine_map<(d0) -> (d0 + 1)>(%i)
+ %v = vector.load %A[%j, %c0] : memref<4x4xf32>, vector<4xf32>
+ vector.store %v, %A[%i, %c0] : memref<4x4xf32>, vector<4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_vector_different_indices
+// CHECK: scf.parallel
+// CHECK: scf.parallel
+
+// -----
+
+func.func @fuse_vector_transfer_same_indices(%A: memref<4x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %zero = arith.constant 0.0 : f32
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ %v = vector.transfer_read %A[%i, %c0], %zero {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : memref<4x4xf32>, vector<4xf32>
+ vector.transfer_write %v, %A[%i, %c0] {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<4xf32>, memref<4x4xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ %v = vector.transfer_read %A[%i, %c0], %zero {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : memref<4x4xf32>, vector<4xf32>
+ vector.transfer_write %v, %A[%i, %c0] {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<4xf32>, memref<4x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @fuse_vector_transfer_same_indices
+// CHECK: scf.parallel
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}]
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}]
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}]
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}]
+// CHECK-NOT: scf.parallel
+
+// -----
+
+func.func @do_not_fuse_vector_transfer_different_indices(%A: memref<4x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %zero = arith.constant 0.0 : f32
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ %v = vector.transfer_read %A[%i, %c0], %zero {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : memref<4x4xf32>, vector<4xf32>
+ vector.transfer_write %v, %A[%i, %c0] {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<4xf32>, memref<4x4xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+ %j = affine.apply affine_map<(d0) -> (d0 + 1)>(%i)
+ %v = vector.transfer_read %A[%j, %c0], %zero {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : memref<4x4xf32>, vector<4xf32>
+ vector.transfer_write %v, %A[%i, %c0] {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<4xf32>, memref<4x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_vector_transfer_different_indices
+// CHECK: scf.parallel
+// CHECK: scf.parallel
+
+// -----
+
+func.func @fuse_vector_transfer_with_subview(%A: memref<1x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %zero = arith.constant 0.0 : f32
+ %vec = arith.constant dense<1.0> : vector<4xf32>
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %sub = memref.subview %A[0, 0][1, 4][1, 1] : memref<1x4xf32> to memref<4xf32>
+ vector.transfer_write %vec, %sub[%c0] {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %sum = scf.for %k = %c0 to %c4 step %c1 iter_args(%acc = %zero) -> f32 {
+ %v = memref.load %A[%c0, %k] : memref<1x4xf32>
+ %n = arith.addf %v, %acc : f32
+ scf.yield %n : f32
+ }
+ memref.store %sum, %A[%c0, %c0] : memref<1x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @fuse_vector_transfer_with_subview
+// CHECK: scf.parallel
+// CHECK: vector.transfer_write
+// CHECK: scf.for
+// CHECK-NOT: scf.parallel
+
+// -----
+
+func.func @do_not_fuse_vector_transfer_nontrivial_subview(%A: memref<2x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %zero = arith.constant 0.0 : f32
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %v = vector.transfer_read %A[%c0, %i], %zero {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : memref<2x4xf32>, vector<1xf32>
+ vector.transfer_write %v, %A[%c0, %i] {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<1xf32>, memref<2x4xf32>
+ scf.reduce
+ }
+ %sub = memref.subview %A[1, 0][1, 4][1, 1] : memref<2x4xf32> to memref<4xf32, strided<[1], offset: 4>>
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %v = vector.transfer_read %sub[%i], %zero {in_bounds = [true]} : memref<4xf32, strided<[1], offset: 4>>, vector<1xf32>
+ vector.transfer_write %v, %sub[%i] {in_bounds = [true]} : vector<1xf32>, memref<4xf32, strided<[1], offset: 4>>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_vector_transfer_nontrivial_subview
+// CHECK: scf.parallel
+// CHECK: scf.parallel
+
+// -----
+
+func.func @do_not_fuse_vector_transfer_different_masks(%A: memref<1x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %zero = arith.constant 0.0 : f32
+ %mask_true = vector.create_mask %c1 : vector<1xi1>
+ %mask_false = vector.create_mask %c0 : vector<1xi1>
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %v = vector.transfer_read %A[%c0, %i], %zero, %mask_true {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : memref<1x4xf32>, vector<1xf32>
+ vector.transfer_write %v, %A[%c0, %i], %mask_true {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<1xf32>, memref<1x4xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %v = vector.transfer_read %A[%c0, %i], %zero, %mask_false {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : memref<1x4xf32>, vector<1xf32>
+ vector.transfer_write %v, %A[%c0, %i], %mask_false {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<1xf32>, memref<1x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_vector_transfer_different_masks
+// CHECK: scf.parallel
+// CHECK: scf.parallel
+
+// -----
+
+func.func @fuse_vector_transfer_subview_rank_reducing(%A: memref<1x4xf32>, %B: memref<1x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %zero = arith.constant 0.0 : f32
+ %vec = arith.constant dense<1.0> : vector<4xf32>
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %sub = memref.subview %A[%i, %c0][1, 4][1, 1] : memref<1x4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ vector.transfer_write %vec, %sub[%c0] {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]} : vector<4xf32>, memref<4xf32, strided<[1], offset: ?>>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %sum = scf.for %k = %c0 to %c4 step %c1 iter_args(%acc = %zero) -> f32 {
+ %v = memref.load %A[%i, %k] : memref<1x4xf32>
+ %n = arith.addf %v, %acc : f32
+ scf.yield %n : f32
+ }
+ memref.store %sum, %B[%i, %c0] : memref<1x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @fuse_vector_transfer_subview_rank_reducing
+// CHECK: scf.parallel
+// CHECK: vector.transfer_write
+// CHECK: scf.for
+// CHECK-NOT: scf.parallel
+
+// -----
+
+func.func @do_not_fuse_vector_transfer_subview_offset(%A: memref<1x4xf32>, %B: memref<1x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %zero = arith.constant 0.0 : f32
+ %vec = arith.constant dense<1.0> : vector<4xf32>
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %sub = memref.subview %A[%i, %c0][1, 4][1, 1] : memref<1x4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ vector.transfer_write %vec, %sub[%c0] {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]} : vector<4xf32>, memref<4xf32, strided<[1], offset: ?>>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %sum = scf.for %k = %c0 to %c4 step %c1 iter_args(%acc = %zero) -> f32 {
+ %v = memref.load %A[%i, %k] : memref<1x4xf32>
+ %n = arith.addf %v, %acc : f32
+ scf.yield %n : f32
+ }
+ // Read from an offset alias to prevent fusion.
+ %off = memref.subview %A[%i, %c1][1, 3][1, 1] : memref<1x4xf32> to memref<3xf32, strided<[1], offset: ?>>
+ %v0 = memref.load %off[%c0] : memref<3xf32, strided<[1], offset: ?>>
+ %res = arith.addf %sum, %v0 : f32
+ memref.store %res, %B[%i, %c0] : memref<1x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @do_not_fuse_vector_transfer_subview_offset
+// CHECK: scf.parallel
+// CHECK: scf.parallel
+
+// -----
+
+func.func @fuse_vector_transfer_no_subview(%A: memref<1x4xf32>, %B: memref<1x4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %zero = arith.constant 0.0 : f32
+ %vec = arith.constant dense<2.0> : vector<4xf32>
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ vector.transfer_write %vec, %A[%c0, %i] {permutation_map = affine_map<(d0, d1) -> (d1)>, in_bounds = [true]} : vector<4xf32>, memref<1x4xf32>
+ scf.reduce
+ }
+ scf.parallel (%i) = (%c0) to (%c1) step (%c1) {
+ %sum = scf.for %k = %c0 to %c4 step %c1 iter_args(%acc = %zero) -> f32 {
+ %v = memref.load %A[%c0, %k] : memref<1x4xf32>
+ %n = arith.addf %v, %acc : f32
+ scf.yield %n : f32
+ }
+ memref.store %sum, %B[%c0, %c0] : memref<1x4xf32>
+ scf.reduce
+ }
+ return
+}
+// CHECK-LABEL: func @fuse_vector_transfer_no_subview
+// CHECK: vector.transfer_write
+// CHECK: scf.for
+// CHECK-NOT: scf.parallel
+
+// -----
+
func.func @fuse_reductions_two(%A: memref<2x2xf32>, %B: memref<2x2xf32>) -> (f32, f32) {
%c2 = arith.constant 2 : index
%c0 = arith.constant 0 : index
More information about the Mlir-commits
mailing list