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

[llvmlistbot at llvm.org]

================
@@ -526,3 +526,31 @@ func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8x
   %resT = vector.transpose %contract, [0, 2, 1] : vector<2x4x8xf32> to vector<2x8x4xf32>
   return %resT : vector<2x8x4xf32>
 }
+
+// -----
+
+#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: %[[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]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
----------------
StarOne01 wrote:

```suggestion
//  CHECK-SAME: %[[LHS]], %[[RHS]], %[[CST]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
```

Missing the constant.

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