[Mlir-commits] [mlir] [mlir][vector] Fix masked load/store emulation for rank-0 memrefs (PR #173325)

[Andrzej Warzyński]

================
@@ -123,3 +123,21 @@ func.func @vector_maskedstore_with_alignment(%arg0 : memref<4x5xf32>, %arg1 : ve
   vector.maskedstore %arg0[%idx_0, %idx_4], %mask, %arg1 { alignment = 8 } : memref<4x5xf32>, vector<4xi1>, vector<4xf32>
   return
 }
+
+// CHECK-LABEL:  @vector_masked_load_rank0
+// CHECK-SAME: (%[[ARG0:.*]]: memref<f32, strided<[], offset: ?>>
+// CHECK: memref.load %[[ARG0]][]
+func.func @vector_masked_load_rank0(%arg0: memref<f32, strided<[], offset: ?>>, %arg3: vector<1xi1>,
+                                   %arg4: vector<1xf32>) -> vector<1xf32> {
+  %0 = vector.maskedload %arg0[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<1xi1>, vector<1xf32> into vector<1xf32>
+  return %0: vector<1xf32>
+}
+
+// CHECK-LABEL:  @vector_masked_store_rank0
+// CHECK-SAME: (%[[ARG0:.*]]: memref<f32, strided<[], offset: ?>>
+// CHECK: memref.store %{{.*}}, %[[ARG0]][]
+func.func @vector_masked_store_rank0(%arg0: memref<f32, strided<[], offset: ?>>, %arg3: vector<1xi1>,
+                                   %arg4: vector<1xf32>) {
+  vector.maskedstore %arg0[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<1xi1>, vector<1xf32>
+  return
+}
----------------
banach-space wrote:

1. Change test function name to match the other tests.
2. Remove explicit `strided` and `offset` to reduce noise.

```suggestion
// CHECK-LABEL:  @vector_maskedload_rank0
// CHECK-SAME: (%[[ARG0:.*]]: memref<f32>
// CHECK: memref.load %[[ARG0]][]
func.func @vector_maskedload_rank0(%arg0: memref<f32>, %arg3: vector<1xi1>,
                                   %arg4: vector<1xf32>) -> vector<1xf32> {
  %0 = vector.maskedload %arg0[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<1xi1>, vector<1xf32> into vector<1xf32>
  return %0: vector<1xf32>
}

// CHECK-LABEL:  @vector_maskedstore_rank0
// CHECK-SAME: (%[[ARG0:.*]]: memref<f32>
// CHECK: memref.store %{{.*}}, %[[ARG0]][]
func.func @vector_maskedstore_rank0(%arg0: memref<f32>, %arg3: vector<1xi1>,
                                   %arg4: vector<1xf32>) {
  vector.maskedstore %arg0[], %arg3, %arg4 : memref<f32>, vector<1xi1>, vector<1xf32>
  return
}
```

Could you also check that there's `vector.insert` and `vector.extract` Ops generated? Thanks!

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