[Mlir-commits] [mlir] [mlir][linalg] Hoist masked vector transfer pairs (PR #221481)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Sep 5 11:47:49 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Hugo Pompougnac (qaco)
<details>
<summary>Changes</summary>
This PR depends on #<!-- -->220244. The last commit (5959ef99a42ea949d0e8093190827d47a3fe7cba) is the only one relevant.
Extend `hoistRedundantVectorTransfers` to hoist `vector.mask`-wrapped transfer pairs (and lone reads), moving the whole mask op to the loop boundary and pairing a masked read only with a masked write carrying the same mask.
Assisted-by: Claude Code (Anthropic)
---
Patch is 145.05 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/221481.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h (+12)
- (modified) mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp (+73-52)
- (modified) mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp (+129)
- (added) mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir (+1571)
- (modified) mlir/test/Dialect/Linalg/hoisting.mlir (-1065)
``````````diff
diff --git a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
index c93da6ebada7c..4ee4f0ed68b5d 100644
--- a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
+++ b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
@@ -198,6 +198,18 @@ LogicalResult resolveSourceIndicesRankReducingSubview(
/// negative.
bool hasNegativeStaticStride(MemRefType memRefTy);
+/// Return "true" when no other access nested in `scope` can alias `base`.
+///
+/// The underlying buffer is found by looking through the view ops defining
+/// `base`, and its transitive in-`scope` users are visited. A user conflicts
+/// unless it is effect-free, in `excludedOps`, or each of its memref operands
+/// that resolves to the buffer touches a slice disjoint from `base`'s static
+/// footprint (or, when `readsAreSafe`, is only read). `excludedOps` are left
+/// out of the analysis, typically because the caller intends to relocate them.
+bool hasNoAliasingAccessInScope(Value base, Operation *scope,
+ ArrayRef<Operation *> excludedOps = {},
+ bool readsAreSafe = false);
+
} // namespace memref
} // namespace mlir
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index a573b4a54dbba..622b264797399 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -17,6 +17,8 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.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/Utils/Utils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -161,40 +163,6 @@ void mlir::linalg::hoistRedundantVectorBroadcasts(RewriterBase &rewriter,
}
}
-static bool noAliasingUseInLoop(vector::TransferReadOp transferRead,
- LoopLikeOpInterface loop) {
- Value source = transferRead.getBase();
-
- // Skip view-like Ops and retrive the actual soruce Operation
- while (auto viewLike = source.getDefiningOp<ViewLikeOpInterface>()) {
- if (viewLike.getViewDest() != source) {
- break;
- }
- source = viewLike.getViewSource();
- }
-
- llvm::SmallVector<Operation *, 32> users(source.getUsers().begin(),
- source.getUsers().end());
- llvm::SmallDenseSet<Operation *, 32> processed;
- while (!users.empty()) {
- Operation *user = users.pop_back_val();
- // If the user has already been processed skip.
- if (!processed.insert(user).second)
- continue;
- if (auto viewLike = dyn_cast<ViewLikeOpInterface>(user)) {
- Value viewDest = viewLike.getViewDest();
- users.append(viewDest.getUsers().begin(), viewDest.getUsers().end());
- continue;
- }
- if (isMemoryEffectFree(user) || isa<vector::TransferReadOp>(user))
- continue;
- if (!loop->isAncestor(user))
- continue;
- return false;
- }
- return true;
-}
-
void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
bool verifyNonZeroTrip) {
bool changed = true;
@@ -249,9 +217,14 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
<< *transferRead.getOperation() << "\n");
- auto loop = dyn_cast<LoopLikeOpInterface>(transferRead->getParentOp());
- LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()
- << "\n");
+ // A masked transfer is hoisted as its enclosing `vector.mask` op; that op
+ // is the unit whose parent must be the loop.
+ Operation *readUnit = transferRead;
+ if (auto mask = dyn_cast<vector::MaskOp>(transferRead->getParentOp()))
+ if (mask.getMaskableOp() == transferRead.getOperation())
+ readUnit = mask.getOperation();
+ auto loop = dyn_cast<LoopLikeOpInterface>(readUnit->getParentOp());
+ LLVM_DEBUG(DBGS() << "Parent op: " << *readUnit->getParentOp() << "\n");
if (!isa_and_nonnull<scf::ForOp, affine::AffineForOp>(loop))
return WalkResult::advance();
@@ -265,13 +238,25 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
<< "\n");
SetVector<Operation *> forwardSlice;
- getForwardSlice(transferRead.getOperation(), &forwardSlice);
+ getForwardSlice(readUnit, &forwardSlice);
+ // A forward slice does not follow a value out of a `vector.mask` region,
+ // so re-slice from each enclosing mask op to reach the write.
+ for (unsigned i = 0; i < forwardSlice.size(); ++i) {
+ auto mask = forwardSlice[i]->getParentOfType<vector::MaskOp>();
+ if (mask && forwardSlice.insert(mask.getOperation()))
+ getForwardSlice(mask.getOperation(), &forwardSlice);
+ }
// Look for the last TransferWriteOp in the forwardSlice of
- // `transferRead` that operates on the same memref.
+ // `transferRead` that operates on the same memref. A masked write appears
+ // as the `vector.mask` op wrapping it.
vector::TransferWriteOp transferWrite;
for (auto *sliceOp : llvm::reverse(forwardSlice)) {
auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);
+ if (!candidateWrite)
+ if (auto mask = dyn_cast<vector::MaskOp>(sliceOp))
+ candidateWrite =
+ dyn_cast_or_null<vector::TransferWriteOp>(mask.getMaskableOp());
if (!candidateWrite ||
candidateWrite.getBase() != transferRead.getBase())
continue;
@@ -282,17 +267,37 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
for (auto operand : transferRead.getOperands())
if (!loop.isDefinedOutsideOfLoop(operand))
return WalkResult::advance();
+ // A masked read is hoisted as its `vector.mask` op; its mask must be
+ // loop-invariant too.
+ auto readMask = dyn_cast<vector::MaskOp>(readUnit);
+ if (readMask && !loop.isDefinedOutsideOfLoop(readMask.getMask()))
+ return WalkResult::advance();
// Only hoist transfer_read / transfer_write pairs and singleton
// transfer_reads for now.
if (!transferWrite) {
- // Make sure there are no other accesses to the memref before
- // hoisting transfer_read.
- if (noAliasingUseInLoop(transferRead, loop))
- loop.moveOutOfLoop(transferRead);
+ // Hoisting a lone read is safe as long as no aliasing write remains in
+ // the loop; other reads never conflict.
+ if (memref::hasNoAliasingAccessInScope(transferRead.getBase(), loop,
+ /*excludedOps=*/{},
+ /*readsAreSafe=*/true))
+ loop.moveOutOfLoop(readUnit);
return WalkResult::advance();
}
+ // The write is hoisted with its own `vector.mask` op, if any. A masked
+ // read must pair with a masked write carrying the same mask (and an
+ // unmasked read with an unmasked write).
+ Operation *writeUnit = transferWrite;
+ if (auto mask = dyn_cast<vector::MaskOp>(transferWrite->getParentOp()))
+ if (mask.getMaskableOp() == transferWrite.getOperation())
+ writeUnit = mask.getOperation();
+ auto writeMask = dyn_cast<vector::MaskOp>(writeUnit);
+ if (static_cast<bool>(readMask) != static_cast<bool>(writeMask))
+ return WalkResult::advance();
+ if (readMask && readMask.getMask() != writeMask.getMask())
+ return WalkResult::advance();
+
LLVM_DEBUG(DBGS() << "Candidate: " << *transferWrite.getOperation()
<< "\n");
@@ -314,6 +319,17 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
// Check 2. Note, since both xfer Ops share the source, we only need to
// look at one of them.
auto base = transferRead.getBase();
+ // Whether hoisting is safe despite a view base. Computed lazily and
+ // cached since it is only consulted when a view is present.
+ std::optional<bool> viewAliasingIsSafeCache;
+ auto viewAliasingIsSafe = [&]() {
+ if (!viewAliasingIsSafeCache) {
+ Operation *hoistedPair[] = {transferRead, transferWrite};
+ viewAliasingIsSafeCache =
+ memref::hasNoAliasingAccessInScope(base, loop, hoistedPair);
+ }
+ return *viewAliasingIsSafeCache;
+ };
auto *source = base.getDefiningOp();
if (source) {
// NOTE: We treat `memref.assume_alignment` as a special case.
@@ -337,18 +353,22 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
if (numInLoopUses && memPreAlignment.hasOneUse())
source = memPreAlignment.getDefiningOp();
}
- if (isa_and_nonnull<ViewLikeOpInterface>(source))
+ if (isa_and_nonnull<ViewLikeOpInterface>(source) &&
+ !viewAliasingIsSafe())
return WalkResult::advance();
}
- if (llvm::any_of(base.getUsers(), llvm::IsaPred<ViewLikeOpInterface>))
+ if (llvm::any_of(base.getUsers(), llvm::IsaPred<ViewLikeOpInterface>) &&
+ !viewAliasingIsSafe())
return WalkResult::advance();
// Check 3.
// TODO: may want to memoize this information for performance but it
// likely gets invalidated often.
+ // Compare the hoisted units (the `vector.mask` ops when masked) since the
+ // inner transfers live in sibling mask regions and never dominate.
DominanceInfo dom(loop);
- if (!dom.properlyDominates(transferRead.getOperation(), transferWrite))
+ if (!dom.properlyDominates(readUnit, writeUnit))
return WalkResult::advance();
for (auto &use : transferRead.getBase().getUses()) {
if (!loop->isAncestor(use.getOwner()))
@@ -377,14 +397,15 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
}
}
- // Hoist read before.
- loop.moveOutOfLoop(transferRead);
+ // Hoist read before (with its mask op, if masked).
+ loop.moveOutOfLoop(readUnit);
- // Hoist write after.
- transferWrite->moveAfter(loop);
+ // Hoist write after (with its mask op, if masked).
+ writeUnit->moveAfter(loop);
// Rewrite `loop` with new yields by cloning and erase the original
- // loop.
+ // loop. The carried value is the read unit's result (the masked read
+ // result when masked) and the yielded value is what the write stores.
IRRewriter rewriter(transferRead.getContext());
NewYieldValuesFn yieldFn = [&](OpBuilder &b, Location loc,
ArrayRef<BlockArgument> newBBArgs) {
@@ -392,7 +413,7 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
};
auto maybeNewLoop = loop.replaceWithAdditionalYields(
- rewriter, transferRead.getVector(),
+ rewriter, readUnit->getResult(0),
/*replaceInitOperandUsesInLoop=*/true, yieldFn);
if (failed(maybeNewLoop))
return WalkResult::interrupt();
diff --git a/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt b/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt
index f8c9f8a7cf3d2..da7c365aa4a0b 100644
--- a/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt
+++ b/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt
@@ -8,4 +8,5 @@ add_mlir_dialect_library(MLIRMemRefUtils
MLIRIR
MLIRAffineDialect
MLIRArithUtils
+ MLIRSideEffectInterfaces
)
diff --git a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
index bc019d601dcd9..372d11488a81d 100644
--- a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
+++ b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
@@ -14,6 +14,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/STLExtras.h"
@@ -351,5 +352,133 @@ bool hasNegativeStaticStride(MemRefType memRefTy) {
});
}
+namespace {
+// TODO: This footprint modelling and its logic overlap with
+// `SubsetOpInterface`, which is unfortunately tensor-only today; if it gains
+// memref support, look into unifying with it to improve code reuse.
+//
+/// Footprint of a memref value relative to a `root` memref reached through a
+/// chain of rank-preserving, unit-stride `memref.subview` ops: per `root`
+/// dimension, the `[offset, offset + size)` interval the value covers (a `size`
+/// may be dynamic). The root is the underlying buffer if the whole chain
+/// composes, otherwise the result of the first subview (walking toward the
+/// root) that cannot be composed. Two footprints are therefore comparable only
+/// when they resolve to the *same* root value.
+struct SubviewFootprint {
+ Value root;
+ SmallVector<int64_t> offsets;
+ SmallVector<int64_t> sizes;
+};
+} // namespace
+
+/// Resolve `base` through a `memref.subview` chain to a footprint, or
+/// `std::nullopt` if it cannot be modelled. Composition stops at the first
+/// subview with a dynamic offset, a non-unit stride, or a rank reduction; that
+/// subview's own result becomes the footprint root. This keeps two slices that
+/// share a common (possibly dynamically-offset) base comparable through their
+/// static offsets relative to that base. Sizes may be dynamic:
+/// such dimensions simply cannot be used to prove disjointness.
+static std::optional<SubviewFootprint> resolveSubviewFootprint(Value base) {
+ auto baseTy = dyn_cast<MemRefType>(base.getType());
+ if (!baseTy)
+ return std::nullopt;
+ SubviewFootprint footprint;
+ // Size is fixed by `base` itself; offset accumulates the static offsets of
+ // the composable subviews walked toward the root.
+ footprint.offsets.assign(baseTy.getRank(), 0);
+ footprint.sizes.assign(baseTy.getShape().begin(), baseTy.getShape().end());
+
+ Value cur = base;
+ while (auto sv = cur.getDefiningOp<SubViewOp>()) {
+ ArrayRef<int64_t> staticOffsets = sv.getStaticOffsets();
+ ArrayRef<int64_t> staticStrides = sv.getStaticStrides();
+ // Rank-reducing or non-static-offset / non-unit-stride subviews cannot be
+ // composed: stop here and use `sv`'s result as the root.
+ if (sv.getSourceType().getRank() != sv.getType().getRank() ||
+ llvm::any_of(staticOffsets, ShapedType::isDynamic) ||
+ llvm::any_of(staticStrides, [](int64_t s) { return s != 1; }))
+ break;
+ for (unsigned d = 0, e = staticOffsets.size(); d < e; ++d)
+ footprint.offsets[d] += staticOffsets[d];
+ cur = sv.getSource();
+ }
+ // The root must be a genuine buffer or a `memref.subview` result, not another
+ // kind of view (collapse/expand/reshape) whose coordinate remapping would
+ // invalidate the per-dimension offset bookkeeping.
+ if (auto *def = cur.getDefiningOp())
+ if (isa<ViewLikeOpInterface>(def) && !isa<SubViewOp>(def))
+ return std::nullopt;
+ footprint.root = cur;
+ return footprint;
+}
+
+/// True if two footprints provably cover disjoint memory: they must share the
+/// same root and be separated along at least one statically-sized dimension.
+static bool areFootprintDisjoint(const SubviewFootprint &a,
+ const SubviewFootprint &b) {
+ if (a.root != b.root)
+ return false;
+ for (unsigned d = 0, e = a.offsets.size(); d < e; ++d) {
+ // A dimension can prove disjointness only when both extents are static.
+ if (ShapedType::isDynamic(a.sizes[d]) || ShapedType::isDynamic(b.sizes[d]))
+ continue;
+ int64_t aHi = a.offsets[d] + a.sizes[d];
+ int64_t bHi = b.offsets[d] + b.sizes[d];
+ if (aHi <= b.offsets[d] || bHi <= a.offsets[d])
+ return true;
+ }
+ return false;
+}
+
+bool hasNoAliasingAccessInScope(Value base, Operation *scope,
+ ArrayRef<Operation *> excludedOps,
+ bool readsAreSafe) {
+ auto baseMemref = dyn_cast<MemrefValue>(base);
+ if (!baseMemref)
+ return false;
+ Value buffer = skipViewLikeOps(baseMemref);
+ // Footprint of `base`, if modellable; enables the disjointness escape.
+ std::optional<SubviewFootprint> baseFp = resolveSubviewFootprint(base);
+
+ // Visit the transitive users of the buffer, following views.
+ SmallVector<Operation *> worklist(buffer.getUsers().begin(),
+ buffer.getUsers().end());
+ SmallPtrSet<Operation *, 8> processed;
+ while (!worklist.empty()) {
+ Operation *user = worklist.pop_back_val();
+ if (!processed.insert(user).second)
+ continue;
+ // Views do not access memory; follow them to their own users.
+ if (auto viewLike = dyn_cast<ViewLikeOpInterface>(user)) {
+ Value viewDest = viewLike.getViewDest();
+ worklist.append(viewDest.getUsers().begin(), viewDest.getUsers().end());
+ continue;
+ }
+ if (isMemoryEffectFree(user) || llvm::is_contained(excludedOps, user))
+ continue;
+ if (!scope->isAncestor(user))
+ continue;
+ // In-`scope`, memory-effecting op: each of its operands that resolves to
+ // the buffer conflicts unless it is a provably-disjoint slice or (when
+ // allowed) is only read.
+ for (Value operand : user->getOperands()) {
+ auto slice = dyn_cast<MemrefValue>(operand);
+ if (!slice || skipViewLikeOps(slice) != buffer)
+ continue;
+ if (baseFp) {
+ std::optional<SubviewFootprint> sliceFp =
+ resolveSubviewFootprint(slice);
+ if (sliceFp && areFootprintDisjoint(*baseFp, *sliceFp))
+ continue;
+ }
+ if (readsAreSafe && isa<MemoryEffectOpInterface>(user) &&
+ !hasEffect<MemoryEffects::Write>(user, slice))
+ continue;
+ return false;
+ }
+ }
+ return true;
+}
+
} // namespace memref
} // namespace mlir
diff --git a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
new file mode 100644
index 0000000000000..d9be635d98ae7
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
@@ -0,0 +1,1571 @@
+// RUN: mlir-opt -transform-interpreter -canonicalize --split-input-file %s | FileCheck %s
+
+///----------------------------------------------------------------------------------------
+/// Tests for vector.transfer_read + vector.transfer_write pairs
+///
+/// * Nested inside a single loop
+// * Indices are constant
+///----------------------------------------------------------------------------------------
+
+// The most basic example - hoisting is safe.
+
+// CHECK-LABEL: func.func @hoist_basic_vector_xfer_pair(
+// CHECK-SAME: %[[MEM:[a-zA-Z0-9]+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[LB:[a-zA-Z0-9]+]]: index,
+// CHECK-SAME: %[[UB:[a-zA-Z0-9]+]]: index,
+// CHECK-SAME: %[[STEP:[a-zA-Z0-9]+]]: index) {
+func.func @hoist_basic_vector_xfer_pair(
+ %mem: memref<?x?xf32>, %lb : index, %ub : index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[READ:.*]] = vector.transfer_read %[[MEM]][%[[C0]], %[[C0]]], %[[PAD]] : memref<?x?xf32>, vector<1xf32>
+// CHECK: %[[SCF:.*]] = scf.for %[[I:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] iter_args(%[[INIT:.*]] = %[[READ]]) -> (vector<1xf32>) {
+// CHECK: %[[VAL_6:.*]] = "test.val_use"(%[[INIT]]) : (vector<1xf32>) -> vector<1xf32>
+// CHECK: scf.yield %[[VAL_6]] : vector<1xf32>
+// CHECK: }
+// CHECK: vector.transfer_write %[[SCF]], %[[MEM]][%[[C0]], %[[C0]]] : vector<1xf32>, memref<?x?xf32>
+ scf.for %i = %lb to %ub step %step {
+ %r0 = vector.transfer_read %mem[%c0, %c0], %pad: memref<?x?xf32>, vector<1xf32>
+ %u0 = "test.val_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
+ vector.transfer_write %u0, %mem[%c0, %c0] : vector<1xf32>, memref<?x?xf32>
+ }
+ return
+}
+
+module attributes {transform.with_named_se...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/221481
More information about the Mlir-commits
mailing list