[Mlir-commits] [mlir] [MLIR][Vector] Implement XferOp To {Load|Store}Lowering as MaskableOpRewritePattern (PR #92892)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue May 21 03:39:54 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hugo Trachino (nujaa)
<details>
<summary>Changes</summary>
Implements `TransferReadToVectorLoadLowering` and `TransferWriteToVectorStoreLowering` as a `MaskableOpRewritePattern`. Allowing to exit gracefully when such use of a xferOp is inside a vector::MaskOp
Split of https://github.com/llvm/llvm-project/pull/90835
---
Full diff: https://github.com/llvm/llvm-project/pull/92892.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp (+34-25)
- (modified) mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir (+35)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
index c59012266ceb3..9418a087c4367 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
@@ -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.
@@ -479,7 +483,7 @@ struct TransferReadToVectorLoadLowering
return rewriter.notifyMatchFailure(read, "out-of-bounds needs mask");
// Create vector load op.
- Operation *loadOp;
+ Operation *res;
if (read.getMask()) {
if (read.getVectorType().getRank() != 1)
// vector.maskedload operates on 1-D vectors.
@@ -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->getResults()[0];
}
std::optional<unsigned> maxTransferRank;
@@ -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.
@@ -639,14 +643,19 @@ 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())
+ .getBase();
+ return Value();
} else {
- rewriter.replaceOpWithNewOp<vector::StoreOp>(
- write, write.getVector(), write.getSource(), write.getIndices());
+ rewriter
+ .create<vector::StoreOp>(write.getLoc(), write.getVector(),
+ write.getSource(), write.getIndices())
+ .getBase();
+ return Value();
}
- return success();
}
std::optional<unsigned> maxTransferRank;
diff --git a/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir b/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
index 2f2bdcaab5b3e..a789aac717dab 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
@@ -392,6 +392,41 @@ func.func @transfer_2D_masked(%mem : memref<?x?xf32>, %mask : vector<2x4xi1>) ->
return %res : vector<2x4xf32>
}
+// transfer_read/write are 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 @masked_transfer_2D(
+// CHECK-SAME: %[[MEM:.*]]: memref<8x8xf32>,
+// CHECK-SAME: %[[IDX:.*]]: index,
+// CHECK-SAME: %[[MASK:.*]]: vector<2x4xi1>) -> memref<8x8xf32>
+// CHECK-NOT: vector.load
+// CHECK-NOT: vector.store
+// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[MEM]][%[[IDX]], %[[IDX]]]{{.*}} : memref<8x8xf32>, vector<2x4xf32> } : vector<2x4xi1> -> vector<2x4xf32>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[MEM]][%[[IDX]], %[[IDX]]]{{.*}}: vector<2x4xf32>, memref<8x8xf32> } : vector<2x4xi1>
+
+func.func @masked_transfer_2D(%mem : memref<8x8xf32>, %i : index, %mask : vector<2x4xi1>) -> memref<8x8xf32> {
+ %cf0 = arith.constant 0.0 : f32
+ %read = vector.mask %mask { vector.transfer_read %mem[%i, %i], %cf0 {in_bounds = [true, true]} : memref<8x8xf32>, vector<2x4xf32> } : vector<2x4xi1> -> vector<2x4xf32>
+ vector.mask %mask {vector.transfer_write %read, %mem[%i, %i] {in_bounds = [true, true]} : vector<2x4xf32>, memref<8x8xf32> } : vector<2x4xi1>
+ return %mem : memref<8x8xf32>
+}
+
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%root : !transform.any_op {transform.readonly}) {
%func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">
``````````
</details>
https://github.com/llvm/llvm-project/pull/92892
More information about the Mlir-commits
mailing list