[Mlir-commits] [mlir] [mlir][vector] Preserve the inner mask when lowering masked transfer_… (PR #225643)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

When a `vector.mask` wraps a `vector.transfer_write` that already has a mask, `-lower-vector-mask` replaces the write's mask with the outer mask. This can write lanes that the original transfer mask disabled.

Intersect the two masks with `arith.andi` when the transfer already has a mask. Keep the existing lowering for transfers without an inner mask. Add a regression test that checks the resulting `transfer_write` uses the intersection.

Fixes: #<!-- -->225244

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp (+4-1) 
- (modified) mlir/test/Dialect/Vector/lower-vector-mask.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp
index 7730c4e7c950a..cd86484489ae6 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp
@@ -242,12 +242,15 @@ struct MaskedTransferWriteOpPattern
                             PatternRewriter &rewriter) const override {
     Type resultType =
         writeOp.getResult() ? writeOp.getResult().getType() : Type();
+    Value mask = maskingOp.getMask();
+    if (Value innerMask = writeOp.getMask())
+      mask = arith::AndIOp::create(rewriter, writeOp.getLoc(), mask, innerMask);
 
     // Replace the `vector.mask` operation.
     rewriter.replaceOpWithNewOp<TransferWriteOp>(
         maskingOp.getOperation(), resultType, writeOp.getVector(),
         writeOp.getBase(), writeOp.getIndices(), writeOp.getPermutationMap(),
-        maskingOp.getMask(), writeOp.getInBounds());
+        mask, writeOp.getInBounds());
     return success();
   }
 };
diff --git a/mlir/test/Dialect/Vector/lower-vector-mask.mlir b/mlir/test/Dialect/Vector/lower-vector-mask.mlir
index a8a1164e2f762..3d2127274eed5 100644
--- a/mlir/test/Dialect/Vector/lower-vector-mask.mlir
+++ b/mlir/test/Dialect/Vector/lower-vector-mask.mlir
@@ -34,6 +34,26 @@ func.func @vector_transfer_write_on_memref(%val: vector<16xf32>, %t0: memref<?xf
 
 // -----
 
+func.func @vector_transfer_write_with_inner_mask(%value: vector<4xi32>, %dest: memref<4xi32>, %inner: vector<4xi1>, %outer: vector<4xi1>) {
+  %c0 = arith.constant 0 : index
+  vector.mask %outer {
+    vector.transfer_write %value, %dest[%c0], %inner {in_bounds = [true]}
+      : vector<4xi32>, memref<4xi32>
+  } : vector<4xi1>
+  return
+}
+
+// CHECK-LABEL: func.func @vector_transfer_write_with_inner_mask(
+// CHECK-SAME: %[[VALUE:.*]]: vector<4xi32>,
+// CHECK-SAME: %[[DEST:.*]]: memref<4xi32>,
+// CHECK-SAME: %[[INNER:.*]]: vector<4xi1>,
+// CHECK-SAME: %[[OUTER:.*]]: vector<4xi1>) {
+// CHECK-NOT: vector.mask
+// CHECK: %[[COMBINED:.*]] = arith.andi %[[OUTER]], %[[INNER]] : vector<4xi1>
+// CHECK: vector.transfer_write %[[VALUE]], %[[DEST]]{{.*}}, %[[COMBINED]] {in_bounds = [true]} : vector<4xi32>, memref<4xi32>
+
+// -----
+
 func.func @vector_transfer_write_on_tensor(%val: vector<16xf32>, %t0: tensor<?xf32>, %idx: index, %m0: vector<16xi1>) -> tensor<?xf32> {
   %res = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
   return %res : tensor<?xf32>

``````````

</details>


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