[Mlir-commits] [mlir] [mlir][tensor] Fold masked transfer_read through extract_slice (PR #228140)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tensor

Author: Hsin-Tang Cheng (dodochicha)

<details>
<summary>Changes</summary>

Allow fold-tensor-subset-ops to fold vector.transfer_read operations with an explicit mask while preserving the mask, padding, and in_bounds attribute.

Fixes #<!-- -->227993

---
Full diff: https://github.com/llvm/llvm-project/pull/228140.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp (+4-4) 
- (modified) mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp b/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp
index 5f838eae95de9e8..eba38d6515c1f18 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp
@@ -78,8 +78,6 @@ static LogicalResult preconditionsFoldExtractOrInsertWithTransferOp(
     ExtractOrInsertOp extractOrInsertSliceOp) {
   if (xferOp.hasOutOfBoundsDim())
     return rewriter.notifyMatchFailure(xferOp, "out of bounds transfer dim");
-  if (xferOp.getMask())
-    return rewriter.notifyMatchFailure(xferOp, "masked transfer");
   if (!extractOrInsertSliceOp.hasUnitStride()) {
     return rewriter.notifyMatchFailure(
         xferOp, "non-1 stride insert/extract, requires keeping track of "
@@ -119,7 +117,8 @@ TransferReadOfExtractSliceOpFolder::matchAndRewriteMaskableOp(
           readOp.getPermutationMap(), extractSliceOp.getSourceType().getRank(),
           extractSliceOp.getDroppedDims())),
       readOp.getPadding(),
-      /*mask=*/Value(), readOp.getInBoundsAttr());
+      /*mask=*/readOp.getMask(),
+      readOp.getInBoundsAttr());
   if (maskOp)
     newOp = mlir::vector::maskOperation(rewriter, newOp, maskOp.getMask());
   return newOp->getResults()[0];
@@ -131,7 +130,8 @@ LogicalResult InsertSliceOfTransferWriteOpFolder::matchAndRewrite(
                      .template getDefiningOp<vector::TransferWriteOp>();
   if (!writeOp)
     return rewriter.notifyMatchFailure(insertSliceOp, "not a transfer_write");
-
+  if (writeOp.getMask())
+    return rewriter.notifyMatchFailure(writeOp, "masked transfer");
   LogicalResult preconditionResult =
       preconditionsFoldExtractOrInsertWithTransferOp(rewriter, writeOp,
                                                      insertSliceOp);
diff --git a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir
index b7112eed47b3381..8d13f66e9588005 100644
--- a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir
+++ b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir
@@ -476,3 +476,26 @@ func.func @extract_slice_non_one_stride(
   %1 = tensor.extract_slice %0[%offset1] [%size1] [%stride1] : tensor<?xf32> to tensor<?xf32>
   return %1: tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @transfer_read_of_extract_slice_with_mask
+// CHECK-SAME: %[[SRC:.*]]: tensor<16xf32>, %[[MASK:.*]]: vector<4xi1>, %[[PAD:.*]]: f32
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK-NOT: tensor.extract_slice
+// CHECK: vector.transfer_read %[[SRC]][%[[C4]]], %[[PAD]], %[[MASK]] {in_bounds = [true]} : tensor<16xf32>, vector<4xf32>
+func.func @transfer_read_of_extract_slice_with_mask(
+    %src: tensor<16xf32>,
+    %mask: vector<4xi1>,
+    %padding: f32) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+
+  %slice = tensor.extract_slice %src[4] [4] [1]
+      : tensor<16xf32> to tensor<4xf32>
+
+  %0 = vector.transfer_read %slice[%c0], %padding, %mask
+      {in_bounds = [true]}
+      : tensor<4xf32>, vector<4xf32>
+
+  return %0 : vector<4xf32>
+}
\ No newline at end of file

``````````

</details>


https://github.com/llvm/llvm-project/pull/228140