[Mlir-commits] [mlir] [mlir][vector] use `hoistRedundantVectorTransfers` in both `linalg` and `affine` (PR #205280)
Federico Bruzzone
llvmlistbot at llvm.org
Tue Jun 23 00:46:54 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/205280
>From 0d615ef089d40b040e62faee74ad9441286a53ee Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 23 Jun 2026 09:24:01 +0200
Subject: [PATCH] [mlir][vector] use `hoistRedundantVectorTransfers` in both
`linalg` and `affine`
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Vector/Transforms/VectorTransforms.h | 16 +
.../Dialect/Affine/Transforms/CMakeLists.txt | 1 +
.../Affine/Transforms/SuperVectorize.cpp | 9 +-
.../Dialect/Linalg/Transforms/Hoisting.cpp | 254 +---------------
mlir/lib/Dialect/Vector/CMakeLists.txt | 1 +
.../Dialect/Vector/Hoisting/CMakeLists.txt | 19 ++
.../Vector/Hoisting/VectorHoisting.cpp | 280 ++++++++++++++++++
.../Dialect/Vector/Transforms/CMakeLists.txt | 1 +
mlir/test/Dialect/Linalg/hoisting.mlir | 100 +++++++
9 files changed, 431 insertions(+), 250 deletions(-)
create mode 100644 mlir/lib/Dialect/Vector/Hoisting/CMakeLists.txt
create mode 100644 mlir/lib/Dialect/Vector/Hoisting/VectorHoisting.cpp
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
index e815e026305fa..362f18d76c611 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
@@ -132,6 +132,22 @@ struct VscaleRange {
void eliminateVectorMasks(IRRewriter &rewriter, FunctionOpInterface function,
std::optional<VscaleRange> vscaleRange = {});
+/// Hoist vector.transfer_read / vector.transfer_write pairs with loop-invariant
+/// indices out of loops by rewriting the loop with iter_args so the accumulated
+/// vector stays in a register across loop iterations rather than being
+/// round-tripped through memory.
+///
+/// The function runs loop-invariant code motion first so that subviews and
+/// padding constants become loop-invariant before the hoisting check. View-
+/// like bases (e.g., memref.subview) are accepted provided the view's source
+/// memref has no other uses inside the loop.
+///
+/// When `verifyNonZeroTrip` is true, hoisting is skipped for loops that cannot
+/// be proven to have a non-zero trip count, avoiding speculative memory
+/// accesses.
+void hoistRedundantVectorTransfers(Operation *root,
+ bool verifyNonZeroTrip = false);
+
} // namespace vector
} // namespace mlir
diff --git a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
index 9d912139810b2..09c903b77ee52 100644
--- a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
@@ -44,6 +44,7 @@ add_mlir_dialect_library(MLIRAffineTransforms
MLIRTransformUtils
MLIRValueBoundsOpInterface
MLIRVectorDialect
+ MLIRVectorHoisting
MLIRVectorUtils
)
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 3158a113a7600..b35a31d4cbe38 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -22,7 +22,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
@@ -1687,6 +1687,13 @@ vectorizeLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ vectorization result:\n"
<< *state.opVectorReplacement[rootLoop]);
+ // Hoist vector.transfer_read/write accumulator pairs out of non-vectorized
+ // inner loops (e.g., the k-reduction loop in matmul). After vectorization,
+ // C's read/write sit inside the k-loop with loop-invariant indices; hoisting
+ // them and adding an iter_arg keeps the C tile in a NEON register across
+ // k-iterations instead of reloading from memory every iteration.
+ vector::hoistRedundantVectorTransfers(state.opVectorReplacement[rootLoop]);
+
// Finish this vectorization pattern.
state.finishVectorizationPattern(rootLoop);
return success();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index a573b4a54dbba..264a5289fc16b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -12,16 +12,10 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
-#include "mlir/Analysis/SliceAnalysis.h"
-#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Affine/Utils.h"
-#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
-#include "mlir/IR/Dominance.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
#include "llvm/Support/Debug.h"
@@ -161,248 +155,10 @@ 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;
- while (changed) {
- changed = false;
- // First move loop invariant ops outside of their loop. This needs to be
- // done before as we cannot move ops without interrupting the function walk.
- root->walk(
- [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
-
- // Find all loops that are certain to have non zero trip count. Any loops
- // that are not part of this set cannot be hoisted from, since hoisting from
- // a potentially zero trip count loop may cause a vector transfer to be
- // executed when it shouldn't be.
- llvm::DenseSet<LoopLikeOpInterface> definiteNonZeroTripCountLoops;
- if (verifyNonZeroTrip) {
- root->walk([&](LoopLikeOpInterface loopLike) {
- std::optional<SmallVector<OpFoldResult>> lbs =
- loopLike.getLoopLowerBounds();
- std::optional<SmallVector<OpFoldResult>> ubs =
- loopLike.getLoopUpperBounds();
- // If loop bounds cannot be found, assume possibly zero trip count.
- if (!lbs || !ubs)
- return;
-
- // Otherwise, use ValueBounds to find the maximum lower bound and
- // minimum upper bound. If the bounds are found, and maxLb is less
- // than the minUb, then the loop will not have zero trip count.
- for (auto [lb, ub] : llvm::zip_equal(lbs.value(), ubs.value())) {
- FailureOr<int64_t> maxLb =
- ValueBoundsConstraintSet::computeConstantBound(
- presburger::BoundType::UB, lb,
- /*stopCondition=*/nullptr,
- ValueBoundsOptions{/*closedUB=*/true});
- if (failed(maxLb))
- return;
- FailureOr<int64_t> minUb =
- ValueBoundsConstraintSet::computeConstantBound(
- presburger::BoundType::LB, ub);
- if (failed(minUb))
- return;
- if (minUb.value() <= maxLb.value())
- return;
- definiteNonZeroTripCountLoops.insert(loopLike);
- }
- });
- }
-
- root->walk([&](vector::TransferReadOp transferRead) {
- if (!isa<MemRefType>(transferRead.getShapedType()))
- return WalkResult::advance();
-
- LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
- << *transferRead.getOperation() << "\n");
- auto loop = dyn_cast<LoopLikeOpInterface>(transferRead->getParentOp());
- LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()
- << "\n");
- if (!isa_and_nonnull<scf::ForOp, affine::AffineForOp>(loop))
- return WalkResult::advance();
-
- if (verifyNonZeroTrip && !definiteNonZeroTripCountLoops.contains(loop)) {
- LLVM_DEBUG(DBGS() << "Loop may have zero trip count: " << *loop
- << "\n");
- return WalkResult::advance();
- }
-
- LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()
- << "\n");
-
- SetVector<Operation *> forwardSlice;
- getForwardSlice(transferRead.getOperation(), &forwardSlice);
-
- // Look for the last TransferWriteOp in the forwardSlice of
- // `transferRead` that operates on the same memref.
- vector::TransferWriteOp transferWrite;
- for (auto *sliceOp : llvm::reverse(forwardSlice)) {
- auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);
- if (!candidateWrite ||
- candidateWrite.getBase() != transferRead.getBase())
- continue;
- transferWrite = candidateWrite;
- }
-
- // All operands of the TransferRead must be defined outside of the loop.
- for (auto operand : transferRead.getOperands())
- if (!loop.isDefinedOutsideOfLoop(operand))
- 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);
- return WalkResult::advance();
- }
-
- LLVM_DEBUG(DBGS() << "Candidate: " << *transferWrite.getOperation()
- << "\n");
-
- // Approximate aliasing by checking that:
- // 1. indices, vector type and permutation map are the same (i.e., the
- // transfer_read/transfer_write ops are matching),
- // 2. source operands for transfer.{read|write} do not originate from
- // nor have users that are Ops implementing ViewLikeOpInterface.
- // 3. no other operations in the loop access the same memref except
- // for transfer_read/transfer_write accessing statically disjoint
- // slices.
-
- // Check 1.
- if (transferRead.getIndices() != transferWrite.getIndices() ||
- transferRead.getVectorType() != transferWrite.getVectorType() ||
- transferRead.getPermutationMap() != transferWrite.getPermutationMap())
- return WalkResult::advance();
-
- // Check 2. Note, since both xfer Ops share the source, we only need to
- // look at one of them.
- auto base = transferRead.getBase();
- auto *source = base.getDefiningOp();
- if (source) {
- // NOTE: We treat `memref.assume_alignment` as a special case.
- //
- // The idea is that it is safe to look past AssumeAlignmemtOp (i.e.
- // MemRef _before_ alignment) iff:
- // 1. It has exactly two uses (these have to be the xfer Ops
- // being looked at).
- // 2. The original MemRef has only one use (i.e.
- // AssumeAlignmentOp).
- //
- // Relaxing these conditions will most likely require proper alias
- // analysis.
- if (auto assume = dyn_cast<memref::AssumeAlignmentOp>(source)) {
- Value memPreAlignment = assume.getMemref();
- auto numInLoopUses =
- llvm::count_if(base.getUses(), [&loop](OpOperand &use) {
- return loop->isAncestor(use.getOwner());
- });
-
- if (numInLoopUses && memPreAlignment.hasOneUse())
- source = memPreAlignment.getDefiningOp();
- }
- if (isa_and_nonnull<ViewLikeOpInterface>(source))
- return WalkResult::advance();
- }
-
- if (llvm::any_of(base.getUsers(), llvm::IsaPred<ViewLikeOpInterface>))
- return WalkResult::advance();
-
- // Check 3.
- // TODO: may want to memoize this information for performance but it
- // likely gets invalidated often.
- DominanceInfo dom(loop);
- if (!dom.properlyDominates(transferRead.getOperation(), transferWrite))
- return WalkResult::advance();
- for (auto &use : transferRead.getBase().getUses()) {
- if (!loop->isAncestor(use.getOwner()))
- continue;
- if (use.getOwner() == transferRead.getOperation() ||
- use.getOwner() == transferWrite.getOperation())
- continue;
- if (auto transferWriteUse =
- dyn_cast<vector::TransferWriteOp>(use.getOwner())) {
- if (!vector::isDisjointTransferSet(
- cast<VectorTransferOpInterface>(*transferWrite),
- cast<VectorTransferOpInterface>(*transferWriteUse),
- /*testDynamicValueUsingBounds=*/true))
- return WalkResult::advance();
- } else if (auto transferReadUse =
- dyn_cast<vector::TransferReadOp>(use.getOwner())) {
- if (!vector::isDisjointTransferSet(
- cast<VectorTransferOpInterface>(*transferWrite),
- cast<VectorTransferOpInterface>(*transferReadUse),
- /*testDynamicValueUsingBounds=*/true))
- return WalkResult::advance();
- } else {
- // Unknown use, we cannot prove that it doesn't alias with the
- // transferRead/transferWrite operations.
- return WalkResult::advance();
- }
- }
-
- // Hoist read before.
- loop.moveOutOfLoop(transferRead);
-
- // Hoist write after.
- transferWrite->moveAfter(loop);
-
- // Rewrite `loop` with new yields by cloning and erase the original
- // loop.
- IRRewriter rewriter(transferRead.getContext());
- NewYieldValuesFn yieldFn = [&](OpBuilder &b, Location loc,
- ArrayRef<BlockArgument> newBBArgs) {
- return SmallVector<Value>{transferWrite.getVector()};
- };
-
- auto maybeNewLoop = loop.replaceWithAdditionalYields(
- rewriter, transferRead.getVector(),
- /*replaceInitOperandUsesInLoop=*/true, yieldFn);
- if (failed(maybeNewLoop))
- return WalkResult::interrupt();
-
- transferWrite.getValueToStoreMutable().assign(
- maybeNewLoop->getOperation()->getResults().back());
- changed = true;
- // Need to interrupt and restart because erasing the loop messes up
- // the walk.
- return WalkResult::interrupt();
- });
- }
+ // Run LICM first to expose loop-invariant operands (e.g. subview ops, padding
+ // values) that would otherwise block transfer-pair hoisting.
+ root->walk([](LoopLikeOpInterface loop) { moveLoopInvariantCode(loop); });
+ mlir::vector::hoistRedundantVectorTransfers(root, verifyNonZeroTrip);
}
diff --git a/mlir/lib/Dialect/Vector/CMakeLists.txt b/mlir/lib/Dialect/Vector/CMakeLists.txt
index 34613b0a36e91..17a7bba553d7a 100644
--- a/mlir/lib/Dialect/Vector/CMakeLists.txt
+++ b/mlir/lib/Dialect/Vector/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(Hoisting)
add_subdirectory(IR)
add_subdirectory(Interfaces)
add_subdirectory(Transforms)
diff --git a/mlir/lib/Dialect/Vector/Hoisting/CMakeLists.txt b/mlir/lib/Dialect/Vector/Hoisting/CMakeLists.txt
new file mode 100644
index 0000000000000..84d64d1aa6c4b
--- /dev/null
+++ b/mlir/lib/Dialect/Vector/Hoisting/CMakeLists.txt
@@ -0,0 +1,19 @@
+add_mlir_dialect_library(MLIRVectorHoisting
+ VectorHoisting.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Vector/Transforms
+
+ LINK_LIBS PUBLIC
+ MLIRAffineDialect
+ MLIRAffineAnalysis
+ MLIRIR
+ MLIRMemRefDialect
+ MLIRSCFDialect
+ MLIRSideEffectInterfaces
+ MLIRTransforms
+ MLIRValueBoundsOpInterface
+ MLIRVectorDialect
+ MLIRVectorInterfaces
+ MLIRVectorUtils
+ )
diff --git a/mlir/lib/Dialect/Vector/Hoisting/VectorHoisting.cpp b/mlir/lib/Dialect/Vector/Hoisting/VectorHoisting.cpp
new file mode 100644
index 0000000000000..4a716e84b1ca1
--- /dev/null
+++ b/mlir/lib/Dialect/Vector/Hoisting/VectorHoisting.cpp
@@ -0,0 +1,280 @@
+//===- VectorHoisting.cpp - Hoist redundant vector transfer operations ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements hoisting of redundant vector.transfer_read /
+// vector.transfer_write pairs out of loops. The transformation detects pairs
+// with loop-invariant indices that act as accumulator loads/stores and rewrites
+// the loop with iter_args so the accumulator stays in a register across
+// iterations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
+#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "vector-hoisting"
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+
+using namespace mlir;
+using namespace mlir::vector;
+
+/// Return true if there are no aliasing uses of `transferRead`'s memref in
+/// `loop` other than the read itself and disjoint transfer reads. The function
+/// walks up through a chain of view-like ops to reach the underlying memref
+/// before collecting users.
+static bool noAliasingUseInLoop(vector::TransferReadOp transferRead,
+ LoopLikeOpInterface loop) {
+ Value source = transferRead.getBase();
+
+ // Walk up through view-like ops to find the real source memref.
+ 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 (!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::vector::hoistRedundantVectorTransfers(Operation *root,
+ bool verifyNonZeroTrip) {
+ bool changed = true;
+ while (changed) {
+ changed = false;
+
+ // Collect loops whose trip count is provably non-zero when requested.
+ // Hoisting from a loop with a potentially zero trip count would execute a
+ // transfer unconditionally when it should be skipped.
+ llvm::DenseSet<LoopLikeOpInterface> definiteNonZeroTripCountLoops;
+ if (verifyNonZeroTrip) {
+ root->walk([&](LoopLikeOpInterface loopLike) {
+ std::optional<SmallVector<OpFoldResult>> lbs =
+ loopLike.getLoopLowerBounds();
+ std::optional<SmallVector<OpFoldResult>> ubs =
+ loopLike.getLoopUpperBounds();
+ if (!lbs || !ubs)
+ return;
+ for (auto [lb, ub] : llvm::zip_equal(lbs.value(), ubs.value())) {
+ FailureOr<int64_t> maxLb =
+ ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::UB, lb,
+ /*stopCondition=*/nullptr,
+ ValueBoundsOptions{/*closedUB=*/true});
+ if (failed(maxLb))
+ return;
+ FailureOr<int64_t> minUb =
+ ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::LB, ub);
+ if (failed(minUb))
+ return;
+ if (minUb.value() <= maxLb.value())
+ return;
+ definiteNonZeroTripCountLoops.insert(loopLike);
+ }
+ });
+ }
+
+ root->walk([&](vector::TransferReadOp transferRead) {
+ if (!isa<MemRefType>(transferRead.getShapedType()))
+ return WalkResult::advance();
+
+ LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
+ << *transferRead.getOperation() << "\n");
+ auto loop = dyn_cast<LoopLikeOpInterface>(transferRead->getParentOp());
+ LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()
+ << "\n");
+ if (!isa_and_nonnull<scf::ForOp, affine::AffineForOp>(loop))
+ return WalkResult::advance();
+
+ if (verifyNonZeroTrip && !definiteNonZeroTripCountLoops.contains(loop)) {
+ LLVM_DEBUG(DBGS() << "Loop may have zero trip count: " << *loop
+ << "\n");
+ return WalkResult::advance();
+ }
+
+ LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()
+ << "\n");
+
+ SetVector<Operation *> forwardSlice;
+ getForwardSlice(transferRead.getOperation(), &forwardSlice);
+
+ // Look for the last TransferWriteOp in the forward slice of
+ // `transferRead` that operates on the same memref.
+ vector::TransferWriteOp transferWrite;
+ for (auto *sliceOp : llvm::reverse(forwardSlice)) {
+ auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);
+ if (!candidateWrite ||
+ candidateWrite.getBase() != transferRead.getBase())
+ continue;
+ transferWrite = candidateWrite;
+ }
+
+ // All operands of the TransferRead must be defined outside of the loop.
+ for (auto operand : transferRead.getOperands())
+ if (!loop.isDefinedOutsideOfLoop(operand))
+ return WalkResult::advance();
+
+ // Only hoist transfer_read / transfer_write pairs and singleton
+ // transfer_reads for now.
+ if (!transferWrite) {
+ if (noAliasingUseInLoop(transferRead, loop))
+ loop.moveOutOfLoop(transferRead);
+ return WalkResult::advance();
+ }
+
+ LLVM_DEBUG(DBGS() << "Candidate: " << *transferWrite.getOperation()
+ << "\n");
+
+ // Approximate aliasing by checking that:
+ // 1. indices, vector type and permutation map are the same (i.e., the
+ // transfer_read/transfer_write ops are matching),
+ // 2. source operands for transfer.{read|write} do not originate from
+ // nor have users that are Ops implementing ViewLikeOpInterface
+ // (with a relaxation for view-like ops whose parent memref has no
+ // other accesses inside the loop),
+ // 3. no other operations in the loop access the same memref except
+ // for transfer_read/transfer_write accessing statically disjoint
+ // slices.
+
+ // Check 1.
+ if (transferRead.getIndices() != transferWrite.getIndices() ||
+ transferRead.getVectorType() != transferWrite.getVectorType() ||
+ transferRead.getPermutationMap() != transferWrite.getPermutationMap())
+ return WalkResult::advance();
+
+ // Check 2. Note, since both xfer ops share the base, we only need to
+ // look at one of them.
+ auto base = transferRead.getBase();
+ auto *source = base.getDefiningOp();
+ if (source) {
+ // Special-case memref.assume_alignment: it is safe to look through it
+ // when (a) it has exactly two in-loop uses (the xfer pair), and (b)
+ // the underlying memref has a single use (the assume_alignment itself).
+ if (auto assume = dyn_cast<memref::AssumeAlignmentOp>(source)) {
+ Value memPreAlignment = assume.getMemref();
+ auto numInLoopUses =
+ llvm::count_if(base.getUses(), [&loop](OpOperand &use) {
+ return loop->isAncestor(use.getOwner());
+ });
+ if (numInLoopUses && memPreAlignment.hasOneUse())
+ source = memPreAlignment.getDefiningOp();
+ }
+
+ // For view-like ops (e.g., memref.subview), rather than bailing
+ // unconditionally, allow hoisting when:
+ // (a) no other view-like op derives from the same source (aliasing),
+ // (b) the source memref has no non-view direct uses inside the loop.
+ // Only one level of view indirection is handled; chained views bail
+ // conservatively.
+ if (auto viewLike = dyn_cast_if_present<ViewLikeOpInterface>(source)) {
+ Value parent = viewLike.getViewSource();
+ // Chained view: bail conservatively.
+ if (parent.getDefiningOp<ViewLikeOpInterface>())
+ return WalkResult::advance();
+ for (auto &use : parent.getUses()) {
+ Operation *user = use.getOwner();
+ if (user == source)
+ continue;
+ // Another view from the same parent can alias.
+ if (isa<ViewLikeOpInterface>(user))
+ return WalkResult::advance();
+ // A direct use of the parent inside the loop is a conflict.
+ if (loop->isAncestor(user))
+ return WalkResult::advance();
+ }
+ }
+ }
+
+ if (llvm::any_of(base.getUsers(), llvm::IsaPred<ViewLikeOpInterface>))
+ return WalkResult::advance();
+
+ // Check 3.
+ DominanceInfo dom(loop);
+ if (!dom.properlyDominates(transferRead.getOperation(), transferWrite))
+ return WalkResult::advance();
+ for (auto &use : transferRead.getBase().getUses()) {
+ if (!loop->isAncestor(use.getOwner()))
+ continue;
+ if (use.getOwner() == transferRead.getOperation() ||
+ use.getOwner() == transferWrite.getOperation())
+ continue;
+ if (auto transferWriteUse =
+ dyn_cast<vector::TransferWriteOp>(use.getOwner())) {
+ if (!vector::isDisjointTransferSet(
+ cast<VectorTransferOpInterface>(*transferWrite),
+ cast<VectorTransferOpInterface>(*transferWriteUse),
+ /*testDynamicValueUsingBounds=*/true))
+ return WalkResult::advance();
+ } else if (auto transferReadUse =
+ dyn_cast<vector::TransferReadOp>(use.getOwner())) {
+ if (!vector::isDisjointTransferSet(
+ cast<VectorTransferOpInterface>(*transferWrite),
+ cast<VectorTransferOpInterface>(*transferReadUse),
+ /*testDynamicValueUsingBounds=*/true))
+ return WalkResult::advance();
+ } else {
+ return WalkResult::advance();
+ }
+ }
+
+ // Hoist read before the loop.
+ loop.moveOutOfLoop(transferRead);
+
+ // Hoist write after the loop.
+ transferWrite->moveAfter(loop);
+
+ // Rewrite `loop` with new yields by cloning and erase the original loop.
+ IRRewriter rewriter(transferRead.getContext());
+ NewYieldValuesFn yieldFn = [&](OpBuilder &b, Location loc,
+ ArrayRef<BlockArgument> newBBArgs) {
+ return SmallVector<Value>{transferWrite.getVector()};
+ };
+
+ auto maybeNewLoop = loop.replaceWithAdditionalYields(
+ rewriter, transferRead.getVector(),
+ /*replaceInitOperandUsesInLoop=*/true, yieldFn);
+ if (failed(maybeNewLoop))
+ return WalkResult::interrupt();
+
+ transferWrite.getValueToStoreMutable().assign(
+ maybeNewLoop->getOperation()->getResults().back());
+ changed = true;
+ return WalkResult::interrupt();
+ });
+ }
+}
diff --git a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
index 112a1db6fe93b..0a0241192d05f 100644
--- a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
@@ -53,6 +53,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
MLIRTensorDialect
MLIRTransforms
MLIRVectorDialect
+ MLIRVectorHoisting
MLIRVectorInterfaces
MLIRVectorUtils
)
diff --git a/mlir/test/Dialect/Linalg/hoisting.mlir b/mlir/test/Dialect/Linalg/hoisting.mlir
index aa0b97a4787fa..45d1955716e8c 100644
--- a/mlir/test/Dialect/Linalg/hoisting.mlir
+++ b/mlir/test/Dialect/Linalg/hoisting.mlir
@@ -1171,3 +1171,103 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+///----------------------------------------------------------------------------------------
+/// Test that vector.transfer_read / vector.transfer_write pairs whose base is a
+/// memref.subview (a ViewLikeOpInterface op) are correctly hoisted when the
+/// subview's source memref has no other accesses inside the loop.
+///
+/// This exercises the relaxed check 2 in hoistRedundantVectorTransfers: rather
+/// than bailing unconditionally when the base is a view-like op, the function
+/// now looks through one level of view to verify the parent memref is not
+/// otherwise accessed inside the loop.
+///----------------------------------------------------------------------------------------
+
+// CHECK-LABEL: func.func @hoist_xfer_pair_subview_base(
+// CHECK-SAME: %[[MEM:[a-zA-Z0-9_]+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[OI:[a-zA-Z0-9_]+]]: index, %[[OJ:[a-zA-Z0-9_]+]]: index,
+// CHECK-SAME: %[[LB:[a-zA-Z0-9_]+]]: index, %[[UB:[a-zA-Z0-9_]+]]: index,
+// CHECK-SAME: %[[STEP:[a-zA-Z0-9_]+]]: index) {
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][%[[OI]], %[[OJ]]]
+// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]][%[[C0]], %[[C0]]], %[[PAD]]
+// CHECK: %[[LOOP:.*]] = scf.for %{{.*}} = %[[LB]] to %[[UB]] step %[[STEP]]
+// CHECK-SAME: iter_args(%[[ACC:.*]] = %[[READ]]) -> (vector<4x4xf32>) {
+// CHECK: %[[USE:.*]] = "val_use"(%[[ACC]]) : (vector<4x4xf32>) -> vector<4x4xf32>
+// CHECK: scf.yield %[[USE]] : vector<4x4xf32>
+// CHECK: }
+// CHECK: vector.transfer_write %[[LOOP]], %[[SV]][%[[C0]], %[[C0]]]
+func.func @hoist_xfer_pair_subview_base(
+ %mem: memref<?x?xf32>, %offset_i: index, %offset_j: index,
+ %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %sv = memref.subview %mem[%offset_i, %offset_j][4, 4][1, 1]
+ : memref<?x?xf32> to memref<4x4xf32, strided<[?, 1], offset: ?>>
+ scf.for %i = %lb to %ub step %step {
+ %r0 = vector.transfer_read %sv[%c0, %c0], %pad
+ : memref<4x4xf32, strided<[?, 1], offset: ?>>, vector<4x4xf32>
+ %u0 = "val_use"(%r0) : (vector<4x4xf32>) -> vector<4x4xf32>
+ vector.transfer_write %u0, %sv[%c0, %c0]
+ : vector<4x4xf32>, memref<4x4xf32, strided<[?, 1], offset: ?>>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.hoist_redundant_vector_transfers %0
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Same as @hoist_xfer_pair_subview_base but the surrounding loop is an
+// affine.for, verifying that the ViewLike fix applies to both scf.for and
+// affine.for (both implement LoopLikeOpInterface).
+
+// CHECK-LABEL: func.func @hoist_xfer_pair_subview_base_affine(
+// CHECK-SAME: %[[MEM:[a-zA-Z0-9_]+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[OI:[a-zA-Z0-9_]+]]: index, %[[OJ:[a-zA-Z0-9_]+]]: index) {
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][%[[OI]], %[[OJ]]]
+// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]][%[[C0]], %[[C0]]], %[[PAD]]
+// CHECK: %[[LOOP:.*]] = affine.for %{{.*}} = 0 to 16
+// CHECK-SAME: iter_args(%[[ACC:.*]] = %[[READ]]) -> (vector<4x4xf32>) {
+// CHECK: %[[USE:.*]] = "val_use"(%[[ACC]]) : (vector<4x4xf32>) -> vector<4x4xf32>
+// CHECK: affine.yield %[[USE]] : vector<4x4xf32>
+// CHECK: }
+// CHECK: vector.transfer_write %[[LOOP]], %[[SV]][%[[C0]], %[[C0]]]
+func.func @hoist_xfer_pair_subview_base_affine(
+ %mem: memref<?x?xf32>, %offset_i: index, %offset_j: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %sv = memref.subview %mem[%offset_i, %offset_j][4, 4][1, 1]
+ : memref<?x?xf32> to memref<4x4xf32, strided<[?, 1], offset: ?>>
+ affine.for %i = 0 to 16 {
+ %r0 = vector.transfer_read %sv[%c0, %c0], %pad
+ : memref<4x4xf32, strided<[?, 1], offset: ?>>, vector<4x4xf32>
+ %u0 = "val_use"(%r0) : (vector<4x4xf32>) -> vector<4x4xf32>
+ vector.transfer_write %u0, %sv[%c0, %c0]
+ : vector<4x4xf32>, memref<4x4xf32, strided<[?, 1], offset: ?>>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.hoist_redundant_vector_transfers %0
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list