Design of Regularization Pass

[manupak]

The expectation of the pass

Now that we support input and output fusions, the distinction between readers and writers becoms irrelevant. What matters is we need to maintain the invariant of "no views between the alloc and all users where one is a RockFusingOp". Correct me here if I am wrong ? No point reading beyond this point if this assumption is wrong.

How this could be a pattern rewrite ?

So basically it feels to me like memref.alloc & RockFusingOp op rewrite to me which finally achieves "no views between the alloc and all users where one is a RockFusingOp"

Ok, but how do you keep fusing towards functions inputs/outputs ?

So this is where I think readers and writers become important.
The way to achieve "no views between the alloc and all users where one is a RockFusingOp":

if RockFusingOp is a reading/writing user of the alloc, we do two steps:
R1) We invert view between alloc and RockFusingOp append them onto writing/reading users of the alloc
This is memref.alloc rewrite (see diagram below). Then we annotate writing users of the allocs as RockFusingOps

Input

  graph TD;
      NAlloc[memref.alloc]-->NRdView1[RdView1];
      NRdView1[RdView1]-->NFusor[RockFusingOp];
      NWr1[WriterOp1] --> NWrView1[WrView1];
      NWrView1[WrView1] -->  NAlloc[memref.alloc];

Loading

Output

  graph TD;
      NAlloc[memref.alloc_new]-->NFusor[RockFusingOp];
      NWr1[WriterOp1_FusingOp] --> NWrView1[WrView1];
      NWrView1[WrView1] --> NInvRdView1_1[inv_RdView1];
      NInvRdView1_1[inv_RdView1] -->  NAlloc[memref.alloc_new];

Loading

R2) Then we do RockFusingOp rewrite to conform to "no views between the alloc and all users where one is a RockFusingOp". So that d be inverting all the views between RockFusingOps and the alloc but as a RockFusingOp rewrite (not a memref.alloc rewrite).

So we end up with :

  graph TD;
      NAlloc[memref.alloc_new]-->NFusor[RockFusingOp];
      NWr1[WriterOp1_FusingOp] -->  NAlloc[memref.alloc_new];
      NWr1NewViews --> NWr1[WriterOp1_FusingOp];

Loading

Then we back to doing R1, because now there are new RockFusingOps that violate "no views between the alloc and all users where one is a RockFusingOp" and follow by R2. We do this until now new RockFusingOp are not idenfied.

Summary

I think #1590 is doing the a similiar thing in effect but it feels more stateful and Im doubtful about stability of that design. Hence Im proposing this which is more inline with PatternRewrite style of writing passes.

I d like to hear @krzysz00 thoughts here -- at least why we went away from PatternRewrite style of writing passes. I agree we might need AlignTilingRewriter style one where we signal things to be rewritten.

[krzysz00]

Another way to phrase the goal here is "you never need to call invertTransforms when AlignTiling is processing a threadwise_read_into/input fusion" + "you never need to call invertTransforms when AlignTiling is processing a threadwise_write_all/output fusion"

This allows us to, in cases where things are decently straightforward, to look through generics, copies, and so on before align-tiling in order to analyze vectorization and thus not lose performance because there's a transpose hiding behind a type conversion generic, for example.

But in order to do that, we need to decide if each linalg.generic is going to be input-fused or output-fused, which is a whole-function analysis that you need to keep state for. And pattern rewrites are supposed to be stateless.

[manupak]

pattern rewrites are supposed to be stateless.

Not really. https://mlir.llvm.org/docs/PatternRewriter/#initialization
See above for a description.
You can recursively (maybe with bounded recursion limits) apply patterns as long as they are needed

[manupak]

Another way to phrase the goal here is "you never need to call invertTransforms when AlignTiling is processing a threadwise_read_into/input fusion" + "you never need to call invertTransforms when AlignTiling is processing a threadwise_write_all/output fusion"

I d argue thats the wrong usecase of this pass. My renewed understanding of "why" this pass exist is not as a sidekick to AlignTiling.
It is to be able to query vectorization info on loads and stores in gridwise lowering. All these rotations are done for that reason.