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

[Hsin-Tang Cheng]

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

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

>From 3fe5fb2187c5a7e2cc51173caf4be75d88ba74cb Mon Sep 17 00:00:00 2001
From: dodochicha <kevincheng60303 at gmail.com>
Date: Fri, 2 Oct 2026 00:50:55 +0800
Subject: [PATCH] [mlir][tensor] Fold masked transfer_read through
 extract_slice

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
---
 .../Tensor/Transforms/FoldTensorSubsetOps.cpp |  8 +++----
 .../Tensor/fold-tensor-subset-ops.mlir        | 23 +++++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp b/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp
index 5f838eae95de9..eba38d6515c1f 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 b7112eed47b33..8d13f66e95880 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