[Mlir-commits] [mlir] [MLIR][Vector] Add fast-math flags to vector.multi_reduction (PR #233546)

[Md Abdullah Shahneous Bari]

mshahneo wrote:

> > > Thank you! Before proceeding, I have a high-level question - do you require `vector.multi_reduction`? Why not use `vector.contract` instead?
> > 
> > 
> > For now yes, the requirement is driven by frontend frameworks. This specific one came from torch lowering to linalg generic to `vector.multi_reduction`. As far as I understand current linalg vectorization still lowers some ops to `vector.multi_reduction`.
> 
> Would you mind adding your example to #234238? Many thanks!

Do you want me add this as test? Or as a PR description?

Here is our case. It comes from the fused-attention payload in llvm/lighthouse (`lighthouse/ingress/mlir_gen/gpu_attention_payload.py`). Softmax needs the row max of the scores, and later the row sum of `exp(s - max)`. Each is a plain reduction with one input (simplified to 2-D here):

```mlir
%row_max = linalg.generic {
    indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],
    iterator_types = ["parallel", "reduction"]}
    ins(%scores : tensor<128x64xf32>) outs(%init : tensor<128xf32>) {
  ^bb0(%s: f32, %acc: f32):
    %m = arith.maximumf %s, %acc : f32
    linalg.yield %m : f32
} -> tensor<128xf32>
```

Linalg vectorization turns it into:

```mlir
%r = vector.multi_reduction <maximumf>, %s, %acc [1] : vector<128x64xf32> to vector<128xf32>
```

There is no multiplication, so this is not a contraction. Only the named contraction ops implement `ContractionOpInterface`, and the vectorizer takes the `vectorizeAsLinalgContraction` path only for them (and only with `createNamedContraction`). To write it as a `vector.contract` we would need a second operand of all ones and a multiply by it, which a later pass would then have to remove. 

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