[Mlir-commits] [mlir] [MLIR][Vector] Generalize broadcast lowering for single-element vector to nD (PR #206501)

[Andrzej Warzyński]

================
@@ -95,6 +95,42 @@ func.func @broadcast_stretch(%arg0: vector<1xf32>) -> vector<4xf32> {
   return %0 : vector<4xf32>
 }
 
+// CHECK-LABEL: func @broadcast_vec2d_from_vec1d_single_element
+// CHECK-SAME: %[[A:.*0]]: vector<1xf32>
+// CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : f32 from vector<1xf32>
+// CHECK:      %[[T1:.*]] = vector.broadcast %[[T0]] : f32 to vector<16x1xf32>
+// CHECK:      return %[[T1]] : vector<16x1xf32>
+
+func.func @broadcast_vec2d_from_vec1d_single_element(%arg0: vector<1xf32>) -> vector<16x1xf32> {
+  %0 = vector.broadcast %arg0 : vector<1xf32> to vector<16x1xf32>
+  return %0 : vector<16x1xf32>
+}
+
+// CHECK-LABEL: func @broadcast_vec2d_from_vec2d_single_element
+// CHECK-SAME: %[[A:.*0]]: vector<1x1xf32>
+// CHECK:      %[[T0:.*]] = vector.extract %[[A]][0, 0] : f32 from vector<1x1xf32>
+// CHECK:      %[[T1:.*]] = vector.broadcast %[[T0]] : f32 to vector<4x3xf32>
+// CHECK:      return %[[T1]] : vector<4x3xf32>
+
+func.func @broadcast_vec2d_from_vec2d_single_element(%arg0: vector<1x1xf32>) -> vector<4x3xf32> {
+  %0 = vector.broadcast %arg0 : vector<1x1xf32> to vector<4x3xf32>
+  return %0 : vector<4x3xf32>
+}
+
+// CHECK-LABEL: func @broadcast_single_elem_scalable_src
+// CHECK-SAME: %[[A:.*0]]: vector<[1]xf32>
+// CHECK:      %[[U0:.*]] = ub.poison : vector<4x[1]xf32>
+// CHECK:      %[[T0:.*]] = vector.insert %[[A]], %[[U0]] [0] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK:      %[[T1:.*]] = vector.insert %[[A]], %[[T0]] [1] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK:      %[[T2:.*]] = vector.insert %[[A]], %[[T1]] [2] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK:      %[[T3:.*]] = vector.insert %[[A]], %[[T2]] [3] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK:      return %[[T3]] : vector<4x[1]xf32>
+
+func.func @broadcast_single_elem_scalable_src(%arg0: vector<[1]xf32>) -> vector<4x[1]xf32> {
----------------
banach-space wrote:

[ultra-nit] We tend to use "unit scalable" for `[1]`
```suggestion
func.func @broadcast_vec2d_from_vec1d_unit_scalable_dim(%arg0: vector<[1]xf32>) -> vector<4x[1]xf32> {
```

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