diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt index c6ef6ed86e0d9..0aabdaf667b9d 100644 --- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt @@ -39,5 +39,6 @@ add_mlir_dialect_library(MLIRTensorTransforms MLIRTilingInterface MLIRTransforms MLIRVectorDialect + MLIRVectorUtils MLIRValueBoundsOpInterface ) diff --git a/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp b/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp index 3b8d3708bb731..5396531922aab 100644 --- a/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp +++ b/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp @@ -18,6 +18,7 @@ #include "mlir/Dialect/Tensor/Transforms/Transforms.h" #include "mlir/Dialect/Utils/IndexingUtils.h" #include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/Dialect/Vector/Utils/VectorUtils.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" @@ -48,12 +49,14 @@ static Value getTensorOperand(tensor::InsertSliceOp op) { namespace { /// Merge extract_slice operation with load/transferRead operation. class TransferReadOfExtractSliceOpFolder final - : public OpRewritePattern { + : public vector::MaskableOpRewritePattern { public: - using OpRewritePattern::OpRewritePattern; + using MaskableOpRewritePattern::MaskableOpRewritePattern; - LogicalResult matchAndRewrite(vector::TransferReadOp readOp, - PatternRewriter &rewriter) const override; + FailureOr + matchAndRewriteMaskableOp(vector::TransferReadOp readOp, + vector::MaskingOpInterface maskOp, + PatternRewriter &rewriter) const override; }; /// Merge insert_slice operation with store/transferWriteOp operation. @@ -84,8 +87,10 @@ static LogicalResult preconditionsFoldExtractOrInsertWithTransferOp( return success(); } -LogicalResult TransferReadOfExtractSliceOpFolder::matchAndRewrite( - vector::TransferReadOp readOp, PatternRewriter &rewriter) const { +FailureOr +TransferReadOfExtractSliceOpFolder::matchAndRewriteMaskableOp( + vector::TransferReadOp readOp, vector::MaskingOpInterface maskOp, + PatternRewriter &rewriter) const { auto extractSliceOp = getTensorOperand(readOp).getDefiningOp(); if (!extractSliceOp) @@ -95,7 +100,7 @@ LogicalResult TransferReadOfExtractSliceOpFolder::matchAndRewrite( preconditionsFoldExtractOrInsertWithTransferOp(rewriter, readOp, extractSliceOp); if (failed(preconditionResult)) - return preconditionResult; + return rewriter.notifyMatchFailure(readOp, "Failed preconditions"); SmallVector indices(readOp.getIndices().begin(), readOp.getIndices().end()); @@ -105,15 +110,17 @@ LogicalResult TransferReadOfExtractSliceOpFolder::matchAndRewrite( extractSliceOp.getMixedStrides(), extractSliceOp.getDroppedDims(), indices, sourceIndices); - rewriter.replaceOpWithNewOp( - readOp, readOp.getVectorType(), extractSliceOp.getSource(), sourceIndices, + Operation *newOp = rewriter.create( + readOp.getLoc(), readOp.getVectorType(), extractSliceOp.getSource(), + sourceIndices, AffineMapAttr::get(expandDimsToRank( readOp.getPermutationMap(), extractSliceOp.getSourceType().getRank(), extractSliceOp.getDroppedDims())), readOp.getPadding(), /*mask=*/Value(), readOp.getInBoundsAttr()); - - return success(); + if (maskOp) + newOp = mlir::vector::maskOperation(rewriter, newOp, maskOp.getMask()); + return newOp->getResults()[0]; } LogicalResult InsertSliceOfTransferWriteOpFolder::matchAndRewrite( diff --git a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops-into-vector-transfers.mlir b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops-into-vector-transfers.mlir index 6213db3956f9a..e562cf8efc356 100644 --- a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops-into-vector-transfers.mlir +++ b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops-into-vector-transfers.mlir @@ -75,6 +75,21 @@ func.func @transfer_read_of_extract_slice_non_leading_rank_reduction(%t : tensor return %1 : vector<5x6xf32> } +// CHECK-LABEL: func @masked_transfer_read_of_extract_slice +// CHECK-SAME: %[[t:.*]]: tensor, %[[s1:.*]]: index, %[[s2:.*]]: index +// CHECK-DAG: %[[m:.*]] = vector.create_mask{{.*}} : vector<5x6xi1> +// CHECK-DAG: %[[a:.*]] = affine.apply {{.*}}[[s1]] +// CHECK: vector.mask %[[m]] { vector.transfer_read %[[t]]{{.*}}: tensor, vector<5x6xf32> } : vector<5x6xi1> -> vector<5x6xf32> +func.func @masked_transfer_read_of_extract_slice(%t : tensor, %s1 : index, %s2 : index) -> vector<5x6xf32> { + %c3 = arith.constant 3 : index + %c4 = arith.constant 4 : index + %cst = arith.constant 0.0 : f32 + %0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor to tensor<10x?xf32> + %mask = vector.create_mask %c3, %c4 : vector<5x6xi1> + %1 = vector.mask %mask {vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<10x?xf32>, vector<5x6xf32>} : vector<5x6xi1> -> vector<5x6xf32> + return %1 : vector<5x6xf32> +} + // CHECK-LABEL: func @insert_slice_of_transfer_write( // CHECK-SAME: %[[t1:.*]]: tensor, %[[v:.*]]: vector<5x6xf32>, %[[s:.*]]: index // CHECK: %[[c3:.*]] = arith.constant 3 : index