[AMD] Hoist Q out of the loop for FA optimization

[oplavsic]

Move writing to LDS and reading from LDS right after the loading of a tensor from global memory. This PR does reordering by considering 2 possible patterns depending on whether writing to LDS is done using an optional local_alloc argument or a local_store instruction: 1) load -> local_alloc -> local_store -> local_load, 2) load -> local_alloc -> local_load.

[sjw36]

This change seems to undo some of the ordering that the new pipeliner sets. For instance, the pipeliner will issue the global loads for all buffers first, then write those to shared memory. Is there a perf benefit to interleaving these (ie. load(A)->local_store(A) - load(B)->local_store(B)?

[sjw36]

What about all the pointer logic that these ops depend on?
You must capture/move the cone of logic for all operands of the localStore.
The local_store may also be in the loop and it may not be valid to hoist it.

[sjw36]

Check localAlloc->getBlock() == localStore->getBlock() == loadOp->getBlock() to verify in the same block.

And check below where the logic cone is captured (line 130).

[antiagainst]

+1

[zhanglx13]

I think this is for the Q tensor in FA kernels. And for Q the load is outside of the loop so stream-pipeline will not touch it.
Then if you see local_alloc/load/store inside the loop, you can find its definingOp. And if the definingOp is outside the loop, instead of inside the loop or being a blockArg, we know this is for Q.
More generally, such local_xxx can be hoisted out of the loop since we don't want to 1) use more LDS, and 2) load from LDS at every iteration of the loop.

[sjw36]

I think this is for the Q tensor in FA kernels. And for Q the load is outside of the loop so stream-pipeline will not touch it. Then if you see local_alloc/load/store inside the loop, you can find its definingOp. And if the definingOp is outside the loop, instead of inside the loop or being a blockArg, we know this is for Q. More generally, such local_xxx can be hoisted out of the loop since we don't want to 1) use more LDS, and 2) load from LDS at every iteration of the loop.

Ok, all the logic must be checked to verify loop invariance. This can be checked when gathering the cone.

[antiagainst]

Thanks! A few nits. Also +1 to what Simon said. I think we are starting to see conflicting needs so we should restructure this pass. I'm fine landing this fix first and then doing that next.

[antiagainst]

+1

[oplavsic]

Thanks for detailed review! I addressed @antiagainst code cleanup comments and added some more checks to better verify that the code is affecting only these very specific patterns. @sjw36 I'm not sure I quite follow all of the comments. The way I see it, this code affects only two very particular cases. It should not undo anything that pipeliner has done. If the tensor is pipelined and argument to local_load (or even local_store) is the block argument, this code will bail (on getDefiningOp(); if(SomeOp) ... checks that are in the code). So the code should only affect tensors that are not pipelined. Moreover, it affects only cases where localAlloc has 1 or 2 uses (local_load, or local_load and local_store), and again, the code bails if it's not exactly this pattern with this number of uses. So I'm not sure I understand what pointer arithmetic can be broken here. Pipeliner lit tests are passing. Is there some particular test case or config in general you are worried about that I can run and see in more detail (by examining TTGIR) cases that you are worried about?

[sjw36]

Code reviewed with @oplavsic, the pattern makes sense and captures all necessary ops.
Minor nits and please add check for specifically moving local_ ops out of loop.