Port "sub-group transpose reduction" to default path

[victor-eds]

#2109 explores layout conversion in the advanced path to improve reduction performance (see #1637 for investigation). Porting this to the default path would involve a transformation similar to (after heuristics to check profitability):

1. Reshape input tensor so no data movement is needed and we can perform reduction of elements within the work-item tt.reshape
2. Perform reduction within the work-item tt.reduce
3. Convert layout so a transposition within the sub-group as explained in the investigation is performed triton_gpu.convert_layout
4. Finalize reduction (within work-item and possibly within the work-group) tt.reduce
5. Convert back to initial layout triton_gpu.convert_layout

Note 5 can be dropped in case the new layout is beneficial for performance.

[victor-eds]

Working on generating always NOP reshape ops.

[victor-eds]

Adding lit tests locally. Pass 1.0 in a good shape.

[victor-eds]

Pass working for repCluster[0] = 2 case. Working on repCluster[0] > 2 (real cases). Pass itself working fine with it (finalizing tests), will work on layout conversion optimizations too.

[victor-eds]

After integrating all of the changes we have up to this point. This is the reported speedup calculated as xetla_speedup_slm / xetla_speedup_baseline, with xetla_speedup_X = triton_mean_gbps / xetla_mean_gpbs from (slm vs. baseline). As this pass matches FA for now, we focus on that workload only. Note the pass can be extended to support more reductions and there is still room for improvement.

      Z   H  N_CTX  D_HEAD  CAUSAL  Speedup
0    1  16  16384     128   False     0.994840066037
1    1  16  16384     128    True     0.998617503047
2    1  32  16384      64   False     0.998388486551
3    1  32  16384      64    True    1.21945316547
4    2  16   8192     128   False     0.996481552534
5    2  16   8192     128    True    0.986066253947
6    2  32   8192      64   False    0.983668089797
7    2  32   8192      64    True    1.19552397853
8    4  16   4096     128   False    0.992287982604
9    4  16   4096     128    True    0.975264209835
10   4  32   4096      64   False    0.960293901543
11   4  32   4096      64    True    1.17439716888
12   8  16   2048     128   False    0.983495166614
13   8  16   2048     128    True    1.00029389437
14   8  32   2048      64   False    1.0007848101
15   8  32   2048      64    True   1.20433393254
16  16  16   1024     128   False    0.96880814982
17  16  16   1024     128    True    1.01326103123
18  16  32   1024      64   False   0.950964800075
19  16  32   1024      64    True   1.08050768492
20  32  16    512     128   False   0.979909264277
21  32  16    512     128    True   1.01801789953
22  32  32    512      64   False   0.954538727308
23  32  32    512      64    True   1.0598195142
24   4  48   1024      64   False   0.993554608016
25   4  48   1024      64    True   1.16116589681

Mean: 1.03248991302
Max: 1.21945316547
Min: 0.950964800075

[etiotto]

After integrating all of the changes we have up to this point. This is the reported speedup calculated as xetla_speedup_slm / xetla_speedup_baseline, with xetla_speedup_X = triton_mean_gbps / xetla_mean_gpbs from (slm vs. baseline). As this pass matches FA for now, we focus on that workload only. Note the pass can be extended to support more reductions and there is still room for improvement.

      Z   H  N_CTX  D_HEAD  CAUSAL  Speedup
0    1  16  16384     128   False     0.994840066037
1    1  16  16384     128    True     0.998617503047
2    1  32  16384      64   False     0.998388486551
3    1  32  16384      64    True    1.21945316547
4    2  16   8192     128   False     0.996481552534
5    2  16   8192     128    True    0.986066253947
6    2  32   8192      64   False    0.983668089797
7    2  32   8192      64    True    1.19552397853
8    4  16   4096     128   False    0.992287982604
9    4  16   4096     128    True    0.975264209835
10   4  32   4096      64   False    0.960293901543
11   4  32   4096      64    True    1.17439716888
12   8  16   2048     128   False    0.983495166614
13   8  16   2048     128    True    1.00029389437
14   8  32   2048      64   False    1.0007848101
15   8  32   2048      64    True   1.20433393254
16  16  16   1024     128   False    0.96880814982
17  16  16   1024     128    True    1.01326103123
18  16  32   1024      64   False   0.950964800075
19  16  32   1024      64    True   1.08050768492
20  32  16    512     128   False   0.979909264277
21  32  16    512     128    True   1.01801789953
22  32  32    512      64   False   0.954538727308
23  32  32    512      64    True   1.0598195142
24   4  48   1024      64   False   0.993554608016
25   4  48   1024      64    True   1.16116589681

Mean: 1.03248991302 Max: 1.21945316547 Min: 0.950964800075

Looks promising! Thanks @victor-eds .
I am looking forward to enabling this by default after you have a chance to test it on other workloads.

[victor-eds]

I am looking forward to enabling this by default after you have a chance to test it on other workloads.

I have filled a series of followup issues I'd like to address before a deeper performance investigation. I'll add these as a comment in this issue.

I would however call this issue done when the outstanding PRs are merged, as the new pass is better (and generates better code) than the one in the advanced path.

[victor-eds]

Followup issues to further improve performance:

• Detect sub-group shuffles with sliced dimensions #2555
• Optimize SLM allocations for sub-group transpose and shuffle layout conversions #2556
• Do not insert barriers in new sub-group shuffle layout conversions #2557
• Use linear layout to extend -tritonintelgpu-optimize-reduction-locality to other layouts #2563