[WIP] SLM pipelined GEMM design

[jiyang1011]

Description

This RFC is mainly to supply a proposal of copy API from G->S. And consider whether there is possibility to gain benefit when SLM is introduced into GEMM pipeline, which needs experiments to confirm this argument.

Type

• Bug - [ ] Feature - [ ] Performance - [ ] Refactor

Testing

• Tests pass - [ ] Xe12 - [ ] Xe20

Performance

Metric	Before	After

References

Fixes #

Checklist

• Copyright - [ ] Co-pilot Review - [ ] Deprecated APIs not used

[Antonyvance]

@jiyang1011 Could you please add the following
What the new APIs
Which examples should be modified to show the new APIs
Pipeline: Create a pseudo code to show the flow.
Thank you.

[rolandschulz]

Do you mean specialization? If not what's specification? If so why would there be duplicated data movement without specialization?

[jiyang1011]

I mean the warp specialization in Nvidia terminology. First we use a part of subgroups to loads data from G->S, and another part of subgroups to matrix math. Secondly supposed that all of subgroups with the work flow (G->S, S->R, R->dpas), obviously there is no gain compared with the existing pipeline. At last we certainly could use the prefetch like method to load data to SLM to avoid duplicated data movement, but for most cases, prefetch all has overlapped data to satisfy the cache line limitation.

[jiyang1011]

I fix the typo

[rolandschulz]

This isn't true in general, is it? For cases without dequantizing the cooperative prefetch makes sure that the global->L1 transfer also only happens once per WG. Of course we can't reuse dequant work with the prefetch pipeline.

[jiyang1011]

Yes, you are correct. Most of case prefetch has the same effect. So there is chance to gain benefit with SLM pipeline when dequantizing / reordering / gathering etc. exists.

[rolandschulz]

This isn't true, is it? Subgroups only copy the tiles they need which is only the same if subgroups do need the same data. And if they need the same data a SLM pipeline will also need to copy from SLM to reg multiple times.

[rolandschulz]

It isn't at all clear that producer-consumer would be beneficial. Probably needs experiments or a good analysis based on prior data.

[jiyang1011]

Yes, we need more experiments to confirm this argument

[sanchitintel]

@jiyang1011, can you please use scaledMMs for experimentation on BMG, as in #633?

I had not used subgroup specialization for my experiments, BTW.

On BMG, in each mainloop iteration, I had tried all subgroups loading equal parts of the B workgroup tile (32 subgroups loaded 16x16 tiles for a 256x32 workgroup-level B tile), converting to BF16/FP16 VNNI-16 format, applying scales, and caching in SLM. However, I used inefficient SLM <-> registers transfers.

I'll also try :

1. loading equal portions of the subgroup-level MMA tile a specific number of subgroups need. For example, if 4 subgroups need the same 32x32 tile for MMA, then they will each load 8x32 or 16x16 tiles, convert to BF16/FP16 VNNI-16 format, apply scales, and then cache them in SLM.
2. Using more efficient SLM <-> registers transfers, as Ji Yang described in this PR.

Thanks!

[jiyang1011]

@jiyang1011 Could you please add the following What the new APIs Which examples should be modified to show the new APIs Pipeline: Create a pseudo code to show the flow. Thank you.

Hi Antony, we have the necessary components. Please refer to these tests

sycl-tla/test/unit/cute/intel_xe/copy_scatter.cpp

Line 204 in 4cdea5a

void copy_kernel_local(TensorS S, TensorD D, TiledCopy Op) {

and https://github.com/intel/sycl-tla/blob/main/test/unit/cute/intel_xe/copy_1d.cpp#L50
block 2d copy and univseralCopy/1-D copy can be integrated into G->S copy operator to meet the demand.