feat(ck): Extend and optimize Quant Gemm Kernel for Aiter
a8w8 (#8423)
## Motivation
JIRA ID : ROCM-20362
JIRA ID : ROCM-26313
The main purpose of this PR is to allow using CK and CK Tile GEMM for
Aiter a8w8 (row-col quantization) and improve its performance.
## Technical Details
### Multiple D for Aiter a8w8 with bias
* Support multiple D (bias) in Quant GEMM kernel
* Extend CShuffleEpilogue
* support int8 -> int32 WarpGemms with row-col quantization
* allow shuffling in fp32 before applying multiple D to prevent
precision loss
### Large tensors support
* Support large tensors in the Quant GEMM kernel by offsetting pointers
of matrices A, D and C.
This feature can be used when M is large, N and K are relatively small
and layout is RCR, it's currently enabled only row-col quantization.
* Allow broadcasting of D column vectors in the old CK's
`DeviceGemmMultiD_Xdl_CShuffle_V3` with large tensors, this case is used
to implement row-col quant scales in Aiter.
### Optimization and workarounds
* Use literal 0 as scales for unscaled 16x16x128 and 32x32x64 mfma: llvm
uses v_mfma for `__builtin_amdgcn_mfma_scale_f32_..._f8f6f4` instead of
v_mfma_scale only if scales are literal 0 values. These instruction
don't need loading scales and save vector registers. See
https://github.com/ROCm/llvm-project/blob/therock-7.13/llvm/lib/Target/AMDGPU/SIInstrInfo.td#L317-L327
* Add workaround for inefficient buffer_load to lds on 7.2
The 3rd argument of buffer_load_dwordx4 is a scalar register. But the
compiler generates a waterfall loop as if lanes can have a different
value, even though the original values comes from as scalar register.
* Use buffer_store_dwordx4 to store 8 bf16 values in epilogue instead of
two buffer_store_dwordx2
* Optimize eight waves pipeline:
* Improve instruction scheduling
* Remove unneeded barriers
* Use nontemporal store/load for C and D matrices in the Quant GEMM
kernel (they are used once per block but may consume cache that can be
better used for matrices A and B)
* Use more efficient padding in epilogue with CTransposed
## Test Plan
A new test is added for multiple D Quant Gemm
(`TestCkTileGemmRowColQuantMultiD/*.RowColQuantMultiDTest`):
```
ninja test_tile_gemm_quant_rowcol && bin/test_tile_gemm_quant_rowcol
```
Testing the large tensor support is not feasible with the current
testing infrastructure because of reference calculations on CPU - it
takes several minutes to run a single test case. Such cases are tested
manually in Aiter.
## Test Result
<!-- Briefly summarize test outcomes. -->
## Submission Checklist
- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
7a710c3542