Double buffering improvements

[giuseros]

• Split the LDS reads and MFMA/WMMA into two independent loops
• Have them into two separate stages (so that they can be executed in parallel)

This is to make our pipeline similar to what CK is doing in:

• https://github.com/ROCm/composable_kernel/blob/6d073d31bbc7d39d8b170d549f2af61970378150/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v4.hpp

[krzysz00]

Doesn't this make BlockwiseGemm kinda obselete?

[krzysz00]

If you're comfortable leaving dead code running around, I don't see any deep reason not to land this.

Since it doesn't look like it makes sense to e2e test this, can we get a lit test for this new structure? Maybe some sort of internal forceDoubleBufferForTesting attribute

[giuseros]

Re: BlockwiseGemm: I don't think so, because in BlockwiseGemm the intuition is that we do everything at the same time
Re: DeadCode: it depends. We have introduced other things (like multi-buffers) that due to backend/hw limitations cannot be used (yet). I see this as progressively having a better scheduling, so we need to cope with some unused code until we completely sort this out. On the other hand, we can surely add an option where we enable all this and test it

[krzysz00]

Re BlockwiseGemm, yeah, I misread the comment

Re testing ... yeah, I'd like a lit test capturing the code structure you're expecting in the double-buffering case after gridwise-gemm-to-blockwise

[manupak]

This is fine with me for now as I had this long standing opinion, the interface for BlockwiseGemm is too limiting -- the fact that it requires LDS buffers to be passed in as inputs.

Long-term :

What would be better (as most other compilers IREE, Triton) is we do BlockwiseGemm take register input with the layout (i.e. tid, iter --> blockwise_tensor). Then let the lowering decides whether it needs to swizzle stuff with whatever way (DPP, LDS, etc) is suitable. We could even define a op : blockwise_copy to encapsulate layout changes.

In the case of user-defined buffering (i.e. what we currently do in gridwise lowering of gemm), user is free to also load directly into the appropiate layout a BlockwiseGemm would prefer. So this would strictly limit the BlockwiseGemm to be just "compute" part of the split you are doing here.

[manupak]

@giuseros should we close this ? I dont think this is happening. correct?