[Mlir-commits] [mlir] [mlir][linalg] Hoist masked vector transfer pairs (PR #221481)
Hugo Pompougnac
llvmlistbot at llvm.org
Mon Sep 28 01:10:20 PDT 2026
https://github.com/qaco updated https://github.com/llvm/llvm-project/pull/221481
>From 1cfb1cd1c5570db19897ada4b0070e7e032d8e77 Mon Sep 17 00:00:00 2001
From: Hugo Pompougnac <tech at hugo-pompougnac.fr>
Date: Sat, 5 Sep 2026 20:37:24 +0200
Subject: [PATCH 1/4] [mlir][linalg] Hoist masked vector transfer pairs
---
.../Dialect/Linalg/Transforms/Hoisting.cpp | 69 +++++-
.../hoist-redundant-vector-transfers.mlir | 226 ++++++++++++++++++
2 files changed, 282 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 36c3e8007ab64..81fe09eca1c16 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -217,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();
@@ -233,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;
@@ -250,6 +267,16 @@ 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 and passthru
+ // must be loop-invariant too.
+ auto readMask = dyn_cast<vector::MaskOp>(readUnit);
+ if (readMask) {
+ if (!loop.isDefinedOutsideOfLoop(readMask.getMask()))
+ return WalkResult::advance();
+ if (readMask.getPassthru() &&
+ !loop.isDefinedOutsideOfLoop(readMask.getPassthru()))
+ return WalkResult::advance();
+ }
// Only hoist transfer_read / transfer_write pairs and singleton
// transfer_reads for now.
@@ -259,10 +286,23 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
if (memref::hasNoAliasingAccessInScope(transferRead.getBase(), loop,
/*excludedOps=*/{},
/*readsAreSafe=*/true))
- loop.moveOutOfLoop(transferRead);
+ 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 (!!readMask != !!writeMask)
+ return WalkResult::advance();
+ if (readMask && readMask.getMask() != writeMask.getMask())
+ return WalkResult::advance();
+
LLVM_DEBUG(DBGS() << "Candidate: " << *transferWrite.getOperation()
<< "\n");
@@ -330,8 +370,10 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
// 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()))
@@ -360,14 +402,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) {
@@ -375,7 +418,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/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
index ced8840d5c892..75b6c4629c0c5 100644
--- a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
+++ b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
@@ -1443,3 +1443,229 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+///----------------------------------------------------------------------------------------
+/// Tests for hoisting masked (`vector.mask`-wrapped) transfer pairs
+///----------------------------------------------------------------------------------------
+
+// A masked read/write pair sharing the same mask is hoisted as a whole: the
+// masked read/write move to the loop boundary while the value is carried
+// unmasked through an iter_arg.
+
+// CHECK-LABEL: func.func @hoist_masked_transfer_pair(
+// CHECK-SAME: %[[MEM:[a-zA-Z0-9]+]]: memref<?xf32>,
+// CHECK-SAME: %[[MASK:[a-zA-Z0-9]+]]: vector<4xi1>,
+// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] {{.*}}transfer_read %[[MEM]]
+// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[IARG:.*]] = %[[READ]]) -> (vector<4xf32>) {
+// CHECK-NEXT: %[[USE:.*]] = "test.val_use"(%[[IARG]])
+// CHECK-NEXT: scf.yield %[[USE]]
+// CHECK-NEXT: }
+// CHECK: vector.mask %[[MASK]] {{.*}}transfer_write %[[FOR]], %[[MEM]]
+func.func @hoist_masked_transfer_pair(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
+ vector.mask %mask { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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
+ }
+}
+
+// -----
+
+// A lone masked read is hoisted as its `vector.mask` op.
+
+// CHECK-LABEL: func.func @hoist_masked_singleton_read(
+// CHECK: %[[READ:.*]] = vector.mask %{{.*}}transfer_read
+// CHECK: scf.for
+// CHECK: "test.some_use"(%[[READ]])
+func.func @hoist_masked_singleton_read(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ "test.some_use"(%r) : (vector<4xf32>) -> ()
+ }
+ 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
+ }
+}
+
+// -----
+
+// A masked read paired with an unmasked write must not be hoisted: the boundary
+// masking would not be preserved.
+
+// CHECK-LABEL: func.func @negative_hoist_masked_read_unmasked_write(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.mask {{.*}}transfer_read
+// CHECK: vector.transfer_write
+// CHECK: }
+func.func @negative_hoist_masked_read_unmasked_write(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
+ vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32>
+ }
+ 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
+ }
+}
+
+// -----
+
+// A masked pair carrying two distinct masks must not be hoisted.
+
+// CHECK-LABEL: func.func @negative_hoist_masked_pair_distinct_masks(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.mask {{.*}}transfer_read
+// CHECK: vector.mask {{.*}}transfer_write
+// CHECK: }
+func.func @negative_hoist_masked_pair_distinct_masks(%mem: memref<?xf32>, %mask0: vector<4xi1>, %mask1: vector<4xi1>, %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %r = vector.mask %mask0 { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
+ vector.mask %mask1 { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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
+ }
+}
+
+// -----
+
+// An unmasked read paired with a masked write must not be hoisted: the boundary
+// masking would not be preserved.
+
+// CHECK-LABEL: func.func @negative_hoist_unmasked_read_masked_write(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.transfer_read
+// CHECK: vector.mask {{.*}}transfer_write
+// CHECK: }
+func.func @negative_hoist_unmasked_read_masked_write(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %r = vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32>
+ %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
+ vector.mask %mask { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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
+ }
+}
+
+// -----
+
+// A masked pair whose mask is recomputed in the loop must not be hoisted: the
+// mask is not loop-invariant.
+
+// CHECK-LABEL: func.func @negative_hoist_masked_pair_loop_variant_mask(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.mask {{.*}}transfer_read
+// CHECK: vector.mask {{.*}}transfer_write
+// CHECK: }
+func.func @negative_hoist_masked_pair_loop_variant_mask(%mem: memref<?xf32>, %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %mask = vector.create_mask %i : vector<4xi1>
+ %r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
+ vector.mask %mask { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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
+ }
+}
+
+// -----
+
+// A masked read whose passthru is recomputed in the loop must not be hoisted:
+// the passthru is not loop-invariant, so the read value changes every iteration.
+
+// CHECK-LABEL: func.func @negative_hoist_masked_read_loop_variant_passthru(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.mask {{.*}}transfer_read
+// CHECK: vector.mask {{.*}}transfer_write
+// CHECK: }
+func.func @negative_hoist_masked_read_loop_variant_passthru(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %idx = arith.index_cast %i : index to i32
+ %f = arith.sitofp %idx : i32 to f32
+ %pt = vector.broadcast %f : f32 to vector<4xf32>
+ %r = vector.mask %mask, %pt { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
+ vector.mask %mask { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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
+ }
+}
>From 3d133c65fc9c3e466e4a738b0ce28c193e4311cd Mon Sep 17 00:00:00 2001
From: Hugo Pompougnac <tech at hugo-pompougnac.fr>
Date: Thu, 17 Sep 2026 14:39:33 +0200
Subject: [PATCH 2/4] After banach-space's review
---
.../Dialect/Linalg/Transforms/Hoisting.cpp | 61 ++++++++-----------
.../hoist-redundant-vector-transfers.mlir | 32 +++-------
2 files changed, 36 insertions(+), 57 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 81fe09eca1c16..5b703590da59b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -217,14 +217,16 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
<< *transferRead.getOperation() << "\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;
+ // When masked, the transfer_read is hoisted as its enclosing
+ // `vector.mask` op, so that op is what must sit directly under the loop.
+ Operation *readOrMaskedRead = 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");
+ readOrMaskedRead = mask.getOperation();
+ auto loop =
+ dyn_cast<LoopLikeOpInterface>(readOrMaskedRead->getParentOp());
+ LLVM_DEBUG(DBGS() << "Parent op: " << *readOrMaskedRead->getParentOp()
+ << "\n");
if (!isa_and_nonnull<scf::ForOp, affine::AffineForOp>(loop))
return WalkResult::advance();
@@ -238,25 +240,15 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
<< "\n");
SetVector<Operation *> 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);
- }
+ getForwardSlice(readOrMaskedRead, &forwardSlice);
// Look for the last TransferWriteOp in the forwardSlice of
- // `transferRead` that operates on the same memref. A masked write appears
- // as the `vector.mask` op wrapping it.
+ // `transferRead` that operates on the same memref. When masked, the write
+ // is the transfer_write nested in its `vector.mask` region, reached as a
+ // user of the read value.
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;
@@ -267,9 +259,9 @@ 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 and passthru
- // must be loop-invariant too.
- auto readMask = dyn_cast<vector::MaskOp>(readUnit);
+ // A masked read is hoisted together with the `vector.mask` op wrapping
+ // it, so that op's mask and passthru operands must be loop-invariant too.
+ auto readMask = dyn_cast<vector::MaskOp>(readOrMaskedRead);
if (readMask) {
if (!loop.isDefinedOutsideOfLoop(readMask.getMask()))
return WalkResult::advance();
@@ -286,18 +278,18 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
if (memref::hasNoAliasingAccessInScope(transferRead.getBase(), loop,
/*excludedOps=*/{},
/*readsAreSafe=*/true))
- loop.moveOutOfLoop(readUnit);
+ loop.moveOutOfLoop(readOrMaskedRead);
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;
+ Operation *writeOrMaskedWrite = 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);
+ writeOrMaskedWrite = mask.getOperation();
+ auto writeMask = dyn_cast<vector::MaskOp>(writeOrMaskedWrite);
if (!!readMask != !!writeMask)
return WalkResult::advance();
if (readMask && readMask.getMask() != writeMask.getMask())
@@ -370,10 +362,11 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
// 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.
+ // The read must dominate the write. Compare `readOrMaskedRead` and
+ // `writeOrMaskedWrite`, not the inner transfers: when masked, those live
+ // in separate `vector.mask` regions, so neither dominates the other.
DominanceInfo dom(loop);
- if (!dom.properlyDominates(readUnit, writeUnit))
+ if (!dom.properlyDominates(readOrMaskedRead, writeOrMaskedWrite))
return WalkResult::advance();
for (auto &use : transferRead.getBase().getUses()) {
if (!loop->isAncestor(use.getOwner()))
@@ -403,13 +396,13 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
}
// Hoist read before (with its mask op, if masked).
- loop.moveOutOfLoop(readUnit);
+ loop.moveOutOfLoop(readOrMaskedRead);
// Hoist write after (with its mask op, if masked).
- writeUnit->moveAfter(loop);
+ writeOrMaskedWrite->moveAfter(loop);
// Rewrite `loop` with new yields by cloning and erase the original
- // loop. The carried value is the read unit's result (the masked read
+ // loop. The carried value is `readOrMaskedRead`'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,
@@ -418,7 +411,7 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
};
auto maybeNewLoop = loop.replaceWithAdditionalYields(
- rewriter, readUnit->getResult(0),
+ rewriter, readOrMaskedRead->getResult(0),
/*replaceInitOperandUsesInLoop=*/true, yieldFn);
if (failed(maybeNewLoop))
return WalkResult::interrupt();
diff --git a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
index 75b6c4629c0c5..6975bbc659cc5 100644
--- a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
+++ b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
@@ -1463,9 +1463,7 @@ module attributes {transform.with_named_sequence} {
// CHECK-NEXT: scf.yield %[[USE]]
// CHECK-NEXT: }
// CHECK: vector.mask %[[MASK]] {{.*}}transfer_write %[[FOR]], %[[MEM]]
-func.func @hoist_masked_transfer_pair(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
- %c0 = arith.constant 0 : index
- %pad = arith.constant 0.0 : f32
+func.func @hoist_masked_transfer_pair(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
%u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
@@ -1492,9 +1490,7 @@ module attributes {transform.with_named_sequence} {
// CHECK: %[[READ:.*]] = vector.mask %{{.*}}transfer_read
// CHECK: scf.for
// CHECK: "test.some_use"(%[[READ]])
-func.func @hoist_masked_singleton_read(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
- %c0 = arith.constant 0 : index
- %pad = arith.constant 0.0 : f32
+func.func @hoist_masked_singleton_read(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
"test.some_use"(%r) : (vector<4xf32>) -> ()
@@ -1522,9 +1518,7 @@ module attributes {transform.with_named_sequence} {
// CHECK: vector.mask {{.*}}transfer_read
// CHECK: vector.transfer_write
// CHECK: }
-func.func @negative_hoist_masked_read_unmasked_write(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
- %c0 = arith.constant 0 : index
- %pad = arith.constant 0.0 : f32
+func.func @negative_hoist_masked_read_unmasked_write(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
%u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
@@ -1552,13 +1546,11 @@ module attributes {transform.with_named_sequence} {
// CHECK: vector.mask {{.*}}transfer_read
// CHECK: vector.mask {{.*}}transfer_write
// CHECK: }
-func.func @negative_hoist_masked_pair_distinct_masks(%mem: memref<?xf32>, %mask0: vector<4xi1>, %mask1: vector<4xi1>, %lb: index, %ub: index, %step: index) {
- %c0 = arith.constant 0 : index
- %pad = arith.constant 0.0 : f32
+func.func @negative_hoist_masked_pair_distinct_masks(%mem: memref<?xf32>, %mask_read: vector<4xi1>, %mask_write: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
- %r = vector.mask %mask0 { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ %r = vector.mask %mask_read { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
%u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
- vector.mask %mask1 { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ vector.mask %mask_write { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
}
return
}
@@ -1583,9 +1575,7 @@ module attributes {transform.with_named_sequence} {
// CHECK: vector.transfer_read
// CHECK: vector.mask {{.*}}transfer_write
// CHECK: }
-func.func @negative_hoist_unmasked_read_masked_write(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
- %c0 = arith.constant 0 : index
- %pad = arith.constant 0.0 : f32
+func.func @negative_hoist_unmasked_read_masked_write(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%r = vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32>
%u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
@@ -1614,9 +1604,7 @@ module attributes {transform.with_named_sequence} {
// CHECK: vector.mask {{.*}}transfer_read
// CHECK: vector.mask {{.*}}transfer_write
// CHECK: }
-func.func @negative_hoist_masked_pair_loop_variant_mask(%mem: memref<?xf32>, %lb: index, %ub: index, %step: index) {
- %c0 = arith.constant 0 : index
- %pad = arith.constant 0.0 : f32
+func.func @negative_hoist_masked_pair_loop_variant_mask(%mem: memref<?xf32>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%mask = vector.create_mask %i : vector<4xi1>
%r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
@@ -1646,9 +1634,7 @@ module attributes {transform.with_named_sequence} {
// CHECK: vector.mask {{.*}}transfer_read
// CHECK: vector.mask {{.*}}transfer_write
// CHECK: }
-func.func @negative_hoist_masked_read_loop_variant_passthru(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index) {
- %c0 = arith.constant 0 : index
- %pad = arith.constant 0.0 : f32
+func.func @negative_hoist_masked_read_loop_variant_passthru(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%idx = arith.index_cast %i : index to i32
%f = arith.sitofp %idx : i32 to f32
>From 5ad1f95e024e21078d2a00dc3f623d4624600594 Mon Sep 17 00:00:00 2001
From: Hugo Pompougnac <tech at hugo-pompougnac.fr>
Date: Mon, 28 Sep 2026 11:56:40 +0400
Subject: [PATCH 3/4] Handle single transfer_write's
---
.../mlir/Dialect/Linalg/Transforms/Hoisting.h | 6 ++
.../Dialect/Linalg/Transforms/Hoisting.cpp | 48 ++++++++++++++++
.../hoist-redundant-vector-transfers.mlir | 57 +++++++++++++++++++
3 files changed, 111 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
index 4edf432d9d97d..a415edb6a02b7 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
@@ -37,6 +37,12 @@ namespace linalg {
/// results in scf::ForOp yielding the value that originally transited through
/// memory.
///
+/// When `verifyNonZeroTrip` is true, a singleton transfer_write with no
+/// matching read is also sunk after the loop when its operands are
+/// loop-invariant and nothing else in the loop accesses its memref, since the
+/// loop is then proven to execute at least once and the store is not
+/// introduced spuriously.
+///
/// TODO: To further improve hoisting opportunities, fold aliasing memref
/// operations into respective vector.transfer{read|write} operations and
/// avoid using ops implementing ViewLikeOpInterface as the source for transfer
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 5b703590da59b..9671e6f3488e7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -423,5 +423,53 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
// the walk.
return WalkResult::interrupt();
});
+
+ // Sink a singleton transfer_write whose operands are loop-invariant and
+ // that is the only op touching its memref: it stores the same value to the
+ // same place every iteration. A lone write introduces a store the loop may
+ // never perform, so it is only sunk past a loop proven to run at least
+ // once.
+ if (changed || !verifyNonZeroTrip)
+ continue;
+ root->walk([&](vector::TransferWriteOp transferWrite) {
+ if (!isa<MemRefType>(transferWrite.getShapedType()))
+ return WalkResult::advance();
+
+ // When masked, the transfer_write is sunk as its enclosing `vector.mask`
+ // op, so that op is what must sit directly under the loop.
+ Operation *writeOrMaskedWrite = transferWrite;
+ if (auto mask = dyn_cast<vector::MaskOp>(transferWrite->getParentOp()))
+ if (mask.getMaskableOp() == transferWrite.getOperation())
+ writeOrMaskedWrite = mask.getOperation();
+ auto loop =
+ dyn_cast<LoopLikeOpInterface>(writeOrMaskedWrite->getParentOp());
+ if (!isa_and_nonnull<scf::ForOp, affine::AffineForOp>(loop) ||
+ !definiteNonZeroTripCountLoops.contains(loop))
+ return WalkResult::advance();
+
+ // All operands of the transfer_write must be defined outside of the loop.
+ for (auto operand : transferWrite.getOperands())
+ if (!loop.isDefinedOutsideOfLoop(operand))
+ return WalkResult::advance();
+ // A masked write is sunk with its `vector.mask` op, so that op's mask
+ // operand must be loop-invariant too.
+ if (auto writeMask = dyn_cast<vector::MaskOp>(writeOrMaskedWrite))
+ if (!loop.isDefinedOutsideOfLoop(writeMask.getMask()))
+ return WalkResult::advance();
+
+ // Sinking is only safe if nothing else in the loop accesses the memref:
+ // any other read or write could observe an intermediate state.
+ Operation *sunkWrite[] = {transferWrite};
+ if (!memref::hasNoAliasingAccessInScope(transferWrite.getBase(), loop,
+ sunkWrite,
+ /*readsAreSafe=*/false))
+ return WalkResult::advance();
+
+ writeOrMaskedWrite->moveAfter(loop);
+ changed = true;
+ // Need to interrupt and restart because moving the write messes up the
+ // walk.
+ return WalkResult::interrupt();
+ });
}
}
diff --git a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
index 6975bbc659cc5..7fd4f5588b890 100644
--- a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
+++ b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
@@ -1510,6 +1510,63 @@ module attributes {transform.with_named_sequence} {
// -----
+// A lone masked write whose operands are loop-invariant, sole accessor of its
+// memref, in a loop proven to run at least once, is sunk past the loop as its
+// `vector.mask` op (the emptied loop is then folded away).
+
+// CHECK-LABEL: func.func @hoist_masked_singleton_write(
+// CHECK-SAME: %[[MEM:[a-zA-Z0-9]+]]: memref<?xf32>,
+// CHECK-SAME: %[[MASK:[a-zA-Z0-9]+]]: vector<4xi1>,
+// CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]]: vector<4xf32>,
+// CHECK-NOT: scf.for
+// CHECK: vector.mask %[[MASK]] {{.*}}transfer_write %[[VEC]], %[[MEM]]
+func.func @hoist_masked_singleton_write(%mem: memref<?xf32>, %mask: vector<4xi1>, %vec: vector<4xf32>, %c0: index) {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 8 : index
+ %step = arith.constant 1 : index
+ scf.for %i = %lb to %ub step %step {
+ vector.mask %mask { vector.transfer_write %vec, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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 verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// A lone masked write must not be sunk past a loop that may not execute: the
+// store would run when the original loop body never did.
+
+// CHECK-LABEL: func.func @negative_hoist_masked_singleton_write_maybe_zero_trip(
+// CHECK: scf.for
+// CHECK: vector.mask %{{.*}}transfer_write
+func.func @negative_hoist_masked_singleton_write_maybe_zero_trip(%mem: memref<?xf32>, %mask: vector<4xi1>, %vec: vector<4xf32>, %lb: index, %ub: index, %step: index, %c0: index) {
+ scf.for %i = %lb to %ub step %step {
+ vector.mask %mask { vector.transfer_write %vec, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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 verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// A masked read paired with an unmasked write must not be hoisted: the boundary
// masking would not be preserved.
>From 02fa484533a5b78d03e3bb21f5abc4f5c86df69f Mon Sep 17 00:00:00 2001
From: Hugo Pompougnac <tech at hugo-pompougnac.fr>
Date: Mon, 28 Sep 2026 12:07:33 +0400
Subject: [PATCH 4/4] Test masked transfer hoisting invariance per op
---
.../hoist-redundant-vector-transfers.mlir | 55 ++++++++++++++-----
1 file changed, 41 insertions(+), 14 deletions(-)
diff --git a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
index 7fd4f5588b890..ba3fea01f21db 100644
--- a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
+++ b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
@@ -1653,20 +1653,18 @@ module attributes {transform.with_named_sequence} {
// -----
-// A masked pair whose mask is recomputed in the loop must not be hoisted: the
-// mask is not loop-invariant.
+// A singleton masked read whose mask is recomputed in the loop is not hoisted:
+// the mask is not loop-invariant.
-// CHECK-LABEL: func.func @negative_hoist_masked_pair_loop_variant_mask(
+// CHECK-LABEL: func.func @negative_hoist_masked_singleton_read_loop_variant_mask(
// CHECK: scf.for {{.*}} step %{{.*}} {
// CHECK: vector.mask {{.*}}transfer_read
-// CHECK: vector.mask {{.*}}transfer_write
// CHECK: }
-func.func @negative_hoist_masked_pair_loop_variant_mask(%mem: memref<?xf32>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
+func.func @negative_hoist_masked_singleton_read_loop_variant_mask(%mem: memref<?xf32>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%mask = vector.create_mask %i : vector<4xi1>
%r = vector.mask %mask { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
- %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
- vector.mask %mask { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ "test.some_use"(%r) : (vector<4xf32>) -> ()
}
return
}
@@ -1683,22 +1681,51 @@ module attributes {transform.with_named_sequence} {
// -----
-// A masked read whose passthru is recomputed in the loop must not be hoisted:
-// the passthru is not loop-invariant, so the read value changes every iteration.
+// A singleton masked write whose mask is recomputed in the loop is not sunk:
+// the mask is not loop-invariant.
-// CHECK-LABEL: func.func @negative_hoist_masked_read_loop_variant_passthru(
+// CHECK-LABEL: func.func @negative_hoist_masked_singleton_write_loop_variant_mask(
// CHECK: scf.for {{.*}} step %{{.*}} {
-// CHECK: vector.mask {{.*}}transfer_read
// CHECK: vector.mask {{.*}}transfer_write
// CHECK: }
-func.func @negative_hoist_masked_read_loop_variant_passthru(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
+func.func @negative_hoist_masked_singleton_write_loop_variant_mask(%mem: memref<?xf32>, %vec: vector<4xf32>, %c0: index) {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 8 : index
+ %step = arith.constant 1 : index
+ scf.for %i = %lb to %ub step %step {
+ %mask = vector.create_mask %i : vector<4xi1>
+ vector.mask %mask { vector.transfer_write %vec, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ }
+ 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 verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// A singleton masked read whose passthru is recomputed in the loop is not
+// hoisted: the passthru is not loop-invariant, so the read value changes every
+// iteration.
+
+// CHECK-LABEL: func.func @negative_hoist_masked_singleton_read_loop_variant_passthru(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.mask {{.*}}transfer_read
+// CHECK: }
+func.func @negative_hoist_masked_singleton_read_loop_variant_passthru(%mem: memref<?xf32>, %mask: vector<4xi1>, %lb: index, %ub: index, %step: index, %c0: index, %pad: f32) {
scf.for %i = %lb to %ub step %step {
%idx = arith.index_cast %i : index to i32
%f = arith.sitofp %idx : i32 to f32
%pt = vector.broadcast %f : f32 to vector<4xf32>
%r = vector.mask %mask, %pt { vector.transfer_read %mem[%c0], %pad : memref<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
- %u = "test.val_use"(%r) : (vector<4xf32>) -> vector<4xf32>
- vector.mask %mask { vector.transfer_write %u, %mem[%c0] : vector<4xf32>, memref<?xf32> } : vector<4xi1>
+ "test.some_use"(%r) : (vector<4xf32>) -> ()
}
return
}
More information about the Mlir-commits
mailing list