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

[Jianhui Li]

Jianhui-Li wrote:

> > I would like to point out that padding and masking are not completely alternatives. They address the dynamic shape issue from different angles and can be applied together. Say for matmul, 1) Padding gives a static storage tile, which the fixed-shape op / its vector.contract lowering requires. 2) Masking restricts the compute to the valid extent N, so the computation over padded columns can be skipped.
> 
> I've been re-reading this statement and I'm still not sure I fully understand it. I may be misunderstanding what you mean here. In particular, what do you mean by "the fixed-shape op / its vector.contract lowering requires"?

I meant the padding is needed to pad every iteration to same storage size to fit fixed-shaped op like vector.contract. My point is the padding is to ensure static storage size so it can be lowered to hardware resource uniformly, either slm or registers. 

> > 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.

Exactly. There is a value for vector IR to keep the semantics even if the tensor is padded, since the hardware (or software lowering) can utilize the mask semantics to reduce the computation. 

> > They address the dynamic shape issue from different angles and can be applied together.
> 
> Note, we use masking for both statically and dynamically shaped tensors, so I'm not yet sure whether the static-vs-dynamic distinction is important here. What specific cases are you trying to address? It might be easier if we focus on those.

The issue we run into is that even vectorization can handle dynamic shape, but it won't pass bufferization properly. That's why we take the approach to pad it (functionally correct) first, and then add a mask to carry the mask semantics at vector IR level as an optimization.  

>Indeed and IMO we should not be designing vectorization based on invalid input.
>+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.

I agree that we should start with correct IR.  It is likely due to an issue in the current linalg padding mechanism: it blindly pad 0 even for these reduction ops and later being canceled out with the preceding unpad.  I just submitted an PR to address that. With that PR, the incorrectly padded IR issue should be gone.  https://github.com/llvm/llvm-project/pull/216517

> I've tried to highlight the parts that are unclear to me and suggest alternatives, such as tensor.extract_slice, that from a distance seem like a fairly straightforward solution
>Would it be possible to extract the non-padded slice of the tensor and feed that to the vectorizer? That should trigger the masking vectorization you are looking for without any flag?
>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.

Thanks for the ideas! I actually tried it and directly feed the non-padded slice extracted to vectorizer as suggested. It indeed triggers the masking vectorization!  @akroviakov  I think potentially we can use this approach to produce masked vector op on top of a valid padded linag IR (fixed by PR216517).  



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