[Mlir-commits] [mlir] [MLIR][Linalg] Recompute linalg.broadcast dimensions when flattening (PR #213641)

[Andrzej Warzyński]

================
@@ -71,6 +71,28 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK-LABEL: func.func @broadcast_rank0_named_tensor(
+// CHECK-SAME:                         %[[ARG0:.*]]: tensor<i32>,
+// CHECK-SAME:                         %[[ARG1:.*]]: tensor<32x2xi32>
+// CHECK-NEXT:    %[[FLATTENED:.*]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1]]
+// CHECK-NEXT:    %[[FLATTENED_RESULT:.*]] = linalg.broadcast ins(%[[ARG0]] : tensor<i32>) outs(%[[FLATTENED]] : tensor<64xi32>) dimensions = [0]
+// CHECK:         %[[RESULT:.*]] = tensor.expand_shape %[[FLATTENED_RESULT]] {{\[}}[0, 1]] output_shape [32, 2] : tensor<64xi32> into tensor<32x2xi32>
+func.func @broadcast_rank0_named_tensor(%arg0: tensor<i32>, %arg1: tensor<32x2xi32>) -> tensor<32x2xi32> {
----------------
banach-space wrote:

```suggestion
func.func @broadcast_as_named_rank0_tensor(%arg0: tensor<i32>, %arg1: tensor<32x2xi32>) -> tensor<32x2xi32> {
```

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