[Mlir-commits] [mlir] [MLIR][Linalg] Add static shape masking option in vectorization (PR #214812)

[Diego Caballero]

dcaballe wrote:

> I would like to point out that padding and masking are not completely alternatives.

That is fair. I do see cases where we want to pad for alignment and still mask because we really do not want to process the padded elements and the hardware is capable of that. However, we have used padding and masking as "either... or..."
so far, so enabling them together properly may take some work.

> > If padding is done before masking, then the input has incorrect semantics until
> > masking is added, because the input is used for two different neutral element
> > reductions.
>
> Indeed and IMO we should not be designing vectorization based on invalid input.
> That would be a very slippery slope.

+1, yes, in the motivating example the result is wrong if masking is not applied. The payload Linalg has to be correct on its own, independently of vectorization.

> this suggests that either this cannot be represented directly at this level in
> Linalg

At this level, you are effectively feeding different B's to the matmul so either:
1. Duplicate the pad and matmul op (and hope for them to be optimized away later on)
2. Introduce corrective ops that reset the padded values after the matmul (and hope for them to be optimized later).

> Do you have any rough idea how this should be done - we can do the padding on
> linalg, but we also want to trigger the vectorization process to add a mask on top
> of selective operations (after padding)?

With the current proposal, we wouldn't be able to support cases where different tensors have a different number of padded elements. *IF* we want to support padding + masking properly, we would need an input representation that allows modeling that and can work from a valid input. One option might be using `scf.if` (or something else) in Linalg to model predication, then turn those predicates into masked ops in the vectorizer.

For example (just an idea):

```mlir
  ^bb0(%in: f32, %acc: f32):
    %j = linalg.index 1 : index
    %valid = arith.cmpi ult, %j, %extent : index
    %r = scf.if %valid -> (f32) {
      %t = arith.maximumf %in, %acc : f32
      scf.yield %t : f32
    } else {
      scf.yield %acc : f32
    }
    linalg.yield %r : f32
```

What is clear to me is that the current padding representation (or Linalg, in general) falls a bit short to model what you want correctly and that's probably the first think we should fix.

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