[Mlir-commits] [mlir] [mlir][vector] Migrate drop-lead-unit-dim to shape_cast (PR #196206)

[Andrzej Warzyński]

================
@@ -270,53 +259,104 @@ func.func @cast_away_contraction_leading_one_dims_to_dot_product(%arg0: vector<6
 }
 
 // -----
+
+// CHECK-DAG: #[[$DOT_MAP:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: #[[$SCALAR_MAP:.*]] = affine_map<(d0) -> ()>
+
+// CHECK-LABEL: cast_away_masked_contraction_with_rank1_acc
+//  CHECK-NEXT:   %[[RHS:.+]] = vector.shape_cast %{{.*}} : vector<1x64xf32> to vector<64xf32>
+//  CHECK-NEXT:   %[[ACC:.+]] = vector.extract %{{.*}}[0] : f32 from vector<1xf32>
+//  CHECK-NEXT:   %[[MASK:.+]] = vector.shape_cast %{{.*}} : vector<64x1xi1> to vector<64xi1>
+//  CHECK-NEXT:   %[[DOT:.+]] = vector.mask %[[MASK]] {
+//  CHECK-SAME:     vector.contract {indexing_maps = [#[[$DOT_MAP]], #[[$DOT_MAP]], #[[$SCALAR_MAP]]], iterator_types = ["reduction"], kind = #vector.kind<add>}
+//  CHECK-SAME:     %{{.*}}, %[[RHS]], %[[ACC]] : vector<64xf32>, vector<64xf32> into f32
+//  CHECK-SAME:   } : vector<64xi1> -> f32
+//  CHECK-NEXT:   %[[RES:.+]] = vector.broadcast %[[DOT]] : f32 to vector<1xf32>
+//  CHECK-NEXT:   return %[[RES]] : vector<1xf32>
+
+func.func @cast_away_masked_contraction_with_rank1_acc(%arg0: vector<64xf32>, %arg1: vector<1x64xf32>, %arg2: vector<1xf32>, %mask: vector<64x1xi1>) -> vector<1xf32> {
+  %0 = vector.mask %mask {
+    vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<64xf32>, vector<1x64xf32> into vector<1xf32>
+  } : vector<64x1xi1> -> vector<1xf32>
+  return %0 : vector<1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: do_not_cast_away_contraction_with_scalable_rank1_acc
+//  CHECK-NOT: vector.shape_cast
+//  CHECK-NOT: vector.extract
+//  CHECK-NOT: vector.broadcast
+//  CHECK-NEXT: vector.contract
+//  CHECK-NEXT: return
+
+func.func @do_not_cast_away_contraction_with_scalable_rank1_acc(%arg0: vector<64xf32>, %arg1: vector<[1]x64xf32>, %arg2: vector<[1]xf32>) -> vector<[1]xf32> {
----------------
banach-space wrote:

[nit] Could you use more standard prefix for negative tests? For example `negative` (instead of `do_not`)? See https://mlir.llvm.org/getting_started/TestingGuide/#test-naming-convention

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