[Mlir-commits] [mlir] [mlir][Vector] Reject scalable masks in maskedload/maskedstore emulation (PR #223316)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Cheng-Chieh Liu (jay87406)

<details>
<summary>Changes</summary>

`VectorMaskedLoadOpConverter` and `VectorMaskedStoreOpConverter` unroll into a fixed number of `scf.if` blocks based on `maskVType.getShape()[0]`. For a scalable mask like `vector<[4]xi1>` that's only the static minimum, not the actual (`vscale`-scaled) element count, so lanes past the fixed 4 iterations were silently dropped whenever `vscale > 1`.

Reject scalable masks the same way the existing rank check does, and add `lit` coverage for both ops.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp (+6) 
- (modified) mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir (+19) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
index 3ca08f3099bd3..53df3d3c48c6a 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
@@ -56,6 +56,9 @@ struct VectorMaskedLoadOpConverter final
     if (maskVType.getShape().size() != 1)
       return rewriter.notifyMatchFailure(
           maskedLoadOp, "expected vector.maskedstore with 1-D mask");
+    if (maskVType.isScalable())
+      return rewriter.notifyMatchFailure(
+          maskedLoadOp, "cannot unroll a scalable mask");
 
     Location loc = maskedLoadOp.getLoc();
     int64_t maskLength = maskVType.getShape()[0];
@@ -125,6 +128,9 @@ struct VectorMaskedStoreOpConverter final
     if (maskVType.getShape().size() != 1)
       return rewriter.notifyMatchFailure(
           maskedStoreOp, "expected vector.maskedstore with 1-D mask");
+    if (maskVType.isScalable())
+      return rewriter.notifyMatchFailure(
+          maskedStoreOp, "cannot unroll a scalable mask");
 
     Location loc = maskedStoreOp.getLoc();
     int64_t maskLength = maskVType.getShape()[0];
diff --git a/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir b/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
index 7c492781c8c55..76646a21c899b 100644
--- a/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
+++ b/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
@@ -145,3 +145,22 @@ func.func @vector_maskedstore_rank0(%arg0: memref<f32>, %arg3: vector<1xi1>, %ar
   vector.maskedstore %arg0[], %arg3, %arg4 : memref<f32>, vector<1xi1>, vector<1xf32>
   return
 }
+
+// A scalable mask cannot be unrolled at compile time, so these must not be
+// converted.
+
+// CHECK-LABEL:  @vector_maskedload_scalable
+//       CHECK:  vector.maskedload
+func.func @vector_maskedload_scalable(%arg0: memref<?xf32>, %mask: vector<[4]xi1>, %pass_thru: vector<[4]xf32>) -> vector<[4]xf32> {
+  %idx_0 = arith.constant 0 : index
+  %0 = vector.maskedload %arg0[%idx_0], %mask, %pass_thru : memref<?xf32>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+  return %0 : vector<[4]xf32>
+}
+
+// CHECK-LABEL:  @vector_maskedstore_scalable
+//       CHECK:  vector.maskedstore
+func.func @vector_maskedstore_scalable(%arg0: memref<?xf32>, %mask: vector<[4]xi1>, %value: vector<[4]xf32>) {
+  %idx_0 = arith.constant 0 : index
+  vector.maskedstore %arg0[%idx_0], %mask, %value : memref<?xf32>, vector<[4]xi1>, vector<[4]xf32>
+  return
+}

``````````

</details>


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