[Mlir-commits] [mlir] [MLIR][Vector] Broadcast folding creates invalid scalar-result vector.contract (PR #230256)

[Atharva Ajmera]

================
@@ -493,6 +493,38 @@ func.func @contract_broadcast_would_have_no_reduction_dim_pair(%arg0 : vector<1x
   return %result : vector<1xf32>
 }
 
+// -----
+
+// Test that CombineContractBroadcast is not combining this case, as that would
+// result in the reduction iterators of the other operand becoming unpaired and
+// produce an invalid scalar-result vector.contract.
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> ()>
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> ()>
+
+// CHECK-LABEL: contract_broadcast_one_sided_unit_dim_reduction
+//  CHECK-SAME: (%[[ARG0:.+]]: vector<2xbf16>, %[[ARG1:.+]]: bf16)
+//  CHECK: %[[CST:.+]] = arith.constant 0.0{{.*}}: bf16
+//  CHECK: %[[LHS:.+]] = vector.shape_cast %[[ARG0]] : vector<2xbf16> to vector<1x1x2xbf16>
+//  CHECK: %[[RHS:.+]] = vector.broadcast %[[ARG1]] : bf16 to vector<1x1x2xbf16>
+//  CHECK: vector.contract
+//  CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
+//  CHECK-SAME: iterator_types = ["reduction", "reduction", "reduction"]
+//  CHECK-SAME: %[[LHS]], %[[RHS]], %[[CST]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
+func.func @contract_broadcast_one_sided_unit_dim_reduction(%arg0 : vector<2xbf16>, %arg1 : bf16) -> bf16 {
----------------
atharvaajmera wrote:

Thanks! Have changed the test name and will take care of the testing guide from now on.

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