[Mlir-commits] [mlir] [mlir][linalg] Hoist masked vector transfer pairs (PR #221481)
Hugo Pompougnac
llvmlistbot at llvm.org
Thu Sep 17 05:40:24 PDT 2026
https://github.com/qaco updated https://github.com/llvm/llvm-project/pull/221481
>From f7f59dcb33c78845fc025ee395f1eacc490a8112 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/2] [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 c6a05bb88b42169a7107eefc6231d8d3204505e6 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/2] 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
More information about the Mlir-commits
mailing list