[mlir] [linalg] Fold broadcast/transpose into linalg.generic (#212415)
Currently we are able to fold broadcast/transpose into
linalg.elementwise. This patch extends the ability to fold
broadcast/transpose into linalg.generic.
For example,
```
%empty = tensor.empty() : tensor<8x16xf32>
%broadcasted = linalg.broadcast ins(%A : tensor<8xf32>) outs(%empty : tensor<8x16xf32>) dimensions = [1]
%result = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%broadcasted : tensor<8x16xf32>) outs(%B : tensor<8x16xf32>) {
^bb0(%in: f32, %out: f32):
%v = arith.addf %in, %in : f32
linalg.yield %v : f32
} -> tensor<8x16xf32>
```
we can fold the broadcast into:
```
%result = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0)>,
affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%A: tensor<8xf32>) outs(%B : tensor<8x16xf32>) {
^bb0(%in: f32, %out: f32):
%v = arith.addf %in, %in : f32
linalg.yield %v : f32
} -> tensor<8x16xf32>
```
For simplicity, we only consider all parallel linalg.generic right now.
AI assisted.
1a5fc1c6b6