[Mlir-commits] [mlir] [mlir] [vector] Allow element/broadcast re-order for FMA if the source is a vector (PR #211208)

[Andrzej Warzyński]

banach-space wrote:

I've updated the title - IMHO it makes the intent clearer.

> Thanks for the contribution!
> 
> I'm not sure what you are trying to achieve. The intent of this pass is to minimize redundant computation on the same data. For your example:
> 
> ```
>   %0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>
>   %1 = vector.broadcast %arg1 : f32 to vector<3x4xf32>
>   %2 = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
>   %3 = vector.fma %0, %1, %2 : vector<3x4xf32>
> ```
> 
> I would expect that broadcasts on `%arg0` and `%arg2` are removed but `%arg1` is broadcasted to `<4xf32>`. Then FMA would happen on `<4xf32>`, followed by the broadcast to `<3x4xf32>`.

@dcaballe , your example is a bit complex and from what I can tell, there's no logic to support that. However, the pattern will currently also fail for even basic examples (note, no scalars at all):

```mlir
func.func @fma_vector_broadcast_source(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<4xf32>) -> vector<3x4xf32> {
  %0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>
  %1 = vector.broadcast %arg1 : vector<4xf32> to vector<3x4xf32>
  %2 = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
  %3 = vector.fma %0, %1, %2 : vector<3x4xf32>
  return %3 : vector<3x4xf32>
}
```

IIUC, the intent of this PR is to unblock trivial examples like the one above. I agree with your comment that the following example should also be supported:
```mlir
  %0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>
  %1 = vector.broadcast %arg1 : f32 to vector<3x4xf32>
  %2 = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
  %3 = vector.fma %0, %1, %2 : vector<3x4xf32>
```

That said, this PR seems fine to me as is, hence approving.

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