Skip to content

[MLIR][Vector] Implement XferOp To {Load|Store}Lowering as MaskableOpRewritePattern #92892

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 18, 2024
Merged

[MLIR][Vector] Implement XferOp To {Load|Store}Lowering as MaskableOpRewritePattern #92892

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
515c26b
[MLIR][Vector] Implement TransferOp To LoadStoreLowering as MaskableO…
nujaa May 13, 2024
2199bee
fixup MR comments
nujaa May 21, 2024
3705157
Update tests
nujaa May 22, 2024
eb9f46c
Add comment and move test.
nujaa May 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 31 additions & 25 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -423,20 +423,24 @@ namespace {
/// result type.
/// - The permutation map doesn't perform permutation (broadcasting is allowed).
struct TransferReadToVectorLoadLowering
: public OpRewritePattern<vector::TransferReadOp> {
: public MaskableOpRewritePattern<vector::TransferReadOp> {
TransferReadToVectorLoadLowering(MLIRContext *context,
std::optional<unsigned> maxRank,
PatternBenefit benefit = 1)
: OpRewritePattern<vector::TransferReadOp>(context, benefit),
: MaskableOpRewritePattern<vector::TransferReadOp>(context, benefit),
maxTransferRank(maxRank) {}

LogicalResult matchAndRewrite(vector::TransferReadOp read,
PatternRewriter &rewriter) const override {
FailureOr<mlir::Value>
matchAndRewriteMaskableOp(vector::TransferReadOp read,
MaskingOpInterface maskOp,
PatternRewriter &rewriter) const override {
if (maxTransferRank && read.getVectorType().getRank() > *maxTransferRank) {
return rewriter.notifyMatchFailure(
read, "vector type is greater than max transfer rank");
}

if (maskOp)
return rewriter.notifyMatchFailure(read, "Masked case not supported");
SmallVector<unsigned> broadcastedDims;
// Permutations are handled by VectorToSCF or
// populateVectorTransferPermutationMapLoweringPatterns.
Expand Down Expand Up @@ -479,7 +483,7 @@ struct TransferReadToVectorLoadLowering
return rewriter.notifyMatchFailure(read, "out-of-bounds needs mask");

// Create vector load op.
Operation *loadOp;
Operation *res;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why rename this from loadOp?

Suggested change
Operation *res;
Operation *loadOp;

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because it can be a BroadcastOp. from l.507

if (read.getMask()) {
if (read.getVectorType().getRank() != 1)
// vector.maskedload operates on 1-D vectors.
Expand All @@ -489,24 +493,20 @@ struct TransferReadToVectorLoadLowering

Value fill = rewriter.create<vector::SplatOp>(
read.getLoc(), unbroadcastedVectorType, read.getPadding());
loadOp = rewriter.create<vector::MaskedLoadOp>(
res = rewriter.create<vector::MaskedLoadOp>(
read.getLoc(), unbroadcastedVectorType, read.getSource(),
read.getIndices(), read.getMask(), fill);
} else {
loadOp = rewriter.create<vector::LoadOp>(
res = rewriter.create<vector::LoadOp>(
read.getLoc(), unbroadcastedVectorType, read.getSource(),
read.getIndices());
}

// Insert a broadcasting op if required.
if (!broadcastedDims.empty()) {
rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
read, read.getVectorType(), loadOp->getResult(0));
} else {
rewriter.replaceOp(read, loadOp->getResult(0));
}

return success();
if (!broadcastedDims.empty())
res = rewriter.create<vector::BroadcastOp>(
read.getLoc(), read.getVectorType(), res->getResult(0));
return res->getResult(0);
}

std::optional<unsigned> maxTransferRank;
Expand Down Expand Up @@ -575,19 +575,23 @@ struct VectorStoreToMemrefStoreLowering
/// - The permutation map is the minor identity map (neither permutation nor
/// broadcasting is allowed).
struct TransferWriteToVectorStoreLowering
: public OpRewritePattern<vector::TransferWriteOp> {
: public MaskableOpRewritePattern<vector::TransferWriteOp> {
TransferWriteToVectorStoreLowering(MLIRContext *context,
std::optional<unsigned> maxRank,
PatternBenefit benefit = 1)
: OpRewritePattern<vector::TransferWriteOp>(context, benefit),
: MaskableOpRewritePattern<vector::TransferWriteOp>(context, benefit),
maxTransferRank(maxRank) {}

LogicalResult matchAndRewrite(vector::TransferWriteOp write,
PatternRewriter &rewriter) const override {
FailureOr<mlir::Value>
matchAndRewriteMaskableOp(vector::TransferWriteOp write,
MaskingOpInterface maskOp,
PatternRewriter &rewriter) const override {
if (maxTransferRank && write.getVectorType().getRank() > *maxTransferRank) {
return rewriter.notifyMatchFailure(
write, "vector type is greater than max transfer rank");
}
if (maskOp)
return rewriter.notifyMatchFailure(write, "Masked case not supported");

// Permutations are handled by VectorToSCF or
// populateVectorTransferPermutationMapLoweringPatterns.
Expand Down Expand Up @@ -639,14 +643,16 @@ struct TransferWriteToVectorStoreLowering
<< write;
});

rewriter.replaceOpWithNewOp<vector::MaskedStoreOp>(
write, write.getSource(), write.getIndices(), write.getMask(),
write.getVector());
rewriter.create<vector::MaskedStoreOp>(
write.getLoc(), write.getSource(), write.getIndices(),
write.getMask(), write.getVector());
} else {
rewriter.replaceOpWithNewOp<vector::StoreOp>(
write, write.getVector(), write.getSource(), write.getIndices());
rewriter.create<vector::StoreOp>(write.getLoc(), write.getVector(),
write.getSource(), write.getIndices());
}
return success();
// There's no return value for StoreOps. Use Value() to signal success to
// matchAndRewrite.
return Value();
}

std::optional<unsigned> maxTransferRank;
Expand Down
17 changes: 17 additions & 0 deletions mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,23 @@ func.func @transfer_to_load(%mem : memref<8x8xf32>, %i : index) -> vector<4xf32>
return %res : vector<4xf32>
}

// Masked transfer_read/write inside are NOT lowered to vector.load/store
// CHECK-LABEL: func @masked_transfer_to_load(
// CHECK-SAME: %[[MEM:.*]]: memref<8x8xf32>,
// CHECK-SAME: %[[IDX:.*]]: index,
// CHECK-SAME: %[[MASK:.*]]: vector<4xi1>) -> memref<8x8xf32>
// CHECK-NOT: vector.load
// CHECK-NOT: vector.store
// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %arg0[%[[IDX]], %[[IDX]]]{{.*}} : memref<8x8xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
// CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[MEM]][%[[IDX]], %[[IDX]]]{{.*}} : vector<4xf32>, memref<8x8xf32> } : vector<4xi1>

func.func @masked_transfer_to_load(%mem : memref<8x8xf32>, %i : index, %mask : vector<4xi1>) -> memref<8x8xf32> {
%cf0 = arith.constant 0.0 : f32
%read = vector.mask %mask {vector.transfer_read %mem[%i, %i], %cf0 {in_bounds = [true]} : memref<8x8xf32>, vector<4xf32>} : vector<4xi1> -> vector<4xf32>
vector.mask %mask {vector.transfer_write %read, %mem[%i, %i] {in_bounds = [true]} : vector<4xf32>, memref<8x8xf32> } : vector<4xi1>
return %mem : memref<8x8xf32>
}

// n-D results are also supported.
// CHECK-LABEL: func @transfer_2D(
// CHECK-SAME: %[[MEM:.*]]: memref<8x8xf32>,
Expand Down
Loading