[Mlir-commits] [mlir] [MLIR] `vector.constant_mask` to support unaligned cases (PR #116520)

[Andrzej Warzyński]

================
@@ -74,6 +74,48 @@ func.func @vector_cst_maskedload_i2(%passthru: vector<5xi2>) -> vector<3x5xi2> {
 
 // -----
 
+func.func @vector_constant_mask_maskedload_i2_multidim(%passthru: vector<5xi2>) -> vector<5xi2> {
+  %0 = memref.alloc() : memref<4x3x5xi2>
+  %cst = arith.constant dense<0> : vector<3x5xi2>
+  %mask = vector.constant_mask [2, 2] : vector<3x5xi1>
+  %ext_mask = vector.extract %mask[1] : vector<5xi1> from vector<3x5xi1>
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %1 = vector.maskedload %0[%c2, %c0, %c0], %ext_mask, %passthru :
+    memref<4x3x5xi2>, vector<5xi1>, vector<5xi2> into vector<5xi2>
+  return %1 : vector<5xi2>
+}
----------------
banach-space wrote:

I see an opportunity to reduce the number of CHECK lines :)

Since `@vector_cst_maskedload_i2` and `@vector_constant_mask_maskedload_i2_multidim` are nearly identical, I think it would be sufficient to focus on testing the differences between them.

Specifically:

* Remove `vector.insert` from `@vector_cst_maskedload_i2` (unless it's significant), so that it aligns closely with `@vector_constant_mask_maskedload_i2_multidim`. Update the `CHECK` lines accordingly.
* Rename `@vector_cst_maskedload_i2` to `@vector_constant_mask_maskedload_i2` for consistency.
* In `@vector_constant_mask_maskedload_i2_multidim`, focus primarily on testing the compressed mask generation logic, as the rest should already align with `@vector_cst_maskedload_i2`.

How does it sound?

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