[Mlir-commits] [mlir] [MLIR] Refactor mask compression logic when emulating `vector.maskedload` ops (PR #116520)

[Andrzej Warzyński]

================
@@ -123,6 +121,31 @@ func.func @check_unaligned_create_mask_static_i2(%passthru: vector<7xi2>) -> vec
 
 // -----
 
+// This test is similar to @vector_constant_mask_maskedload_i2, but the mask is multi-dimensional.
+func.func @vector_constant_mask_maskedload_i2_multidim(%passthru: vector<5xi2>) -> vector<5xi2> {
+  %0 = memref.alloc() : memref<4x3x5xi2>
+  %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>
+}
+
+// CHECK-LABEL: func @vector_constant_mask_maskedload_i2_multidim(
+// CHECK-SAME:   %[[PASSTHRU:[a-zA-Z0-9]+]]
+// CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<15xi8>
+// CHECK: %[[ORIG_MASK:.+]] = vector.constant_mask [2, 2] : vector<3x5xi1>
+// CHECK: %[[EXT_ORIG_MASK:.+]] = vector.extract %[[ORIG_MASK]][1]
----------------
banach-space wrote:

[nit 1] `EXT_ORIG_MASK` -> `ORIG_MASK_VEC_1`. This is extracting row/vector at index 1, "extraction" is already encoded in `vector.extract` (i.e. this info is already present). By using `VEC_1` instead of `EXT`, you are encoding some _new_ info ;-)
 
[nit 2] Rather than "{PREFIX}_OLD_VAR_NAME", I'd stick to "OLD_VAR_SUFFIX" (see e.g. `PTH_PADDED_UPCAST` or  `MASKLOAD_DOWNCAST)`.

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