Unnecessary amdgpu.lds_barrier in GEMM kernel
- Dominant language
- Python
- Stars
- 59
- Forks
- 32
- PR merge metrics
- No merged PRs in 30d
Description
While running moe_test.py the MoE kernel calls into the GEMM defined by get_wave_gemm_kernel(). The emitted MLIR for this GEMM contains redundant amdgpu.lds_barrier instructions which likely introduce unnecessary latency.
Specifically, the first barrier in the inner loop is unnecessary:
```
%15:4 = scf.for %arg3 = %c0 to %c32 step %c1 iter_args(%arg4 = %cst_0, %arg5 = %cst_0, %arg6 = %cst_0, %arg7 = %cst_0) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) {
%29 = affine.apply #map9()[%arg3, %thread_id_x]
%30 = vector.maskedload %1[%2, %29], %4, %cst : memref<1x1024xbf16, strided<[1024, 1], offset: ?>>, vector<8xi1>, vector<8xbf16> into vector<8xbf16>
amdgpu.lds_barrier
vector.store %30, %view_1[%2, %5] : memref<64x36xbf16, #gpu.address_space>, vector<8xbf16>
%31 = vector.maskedload %0[%6, %29], %8, %cst : memref<511x1024xbf16, strided<[1024, 1], offset: ?>>, vector<8xi1>, vector<8xbf16> into vector<8xbf16>
vector.store %31, %view[%2, %5] : memref<64x36xbf16, #gpu.address_space>, vector<8xbf16>
amdgpu.lds_barrier
%32 = vector.load %view[%9, %10] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%33 = vector.load %view[%9, %11] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%34 = vector.load %view[%12, %10] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%35 = vector.load %view[%12, %11] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%36 = vector.load %view_1[%13, %10] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%37 = vector.load %view_1[%13, %11] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%38 = vector.load %view_1[%14, %10] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%39 = vector.load %view_1[%14, %11] : memref<64x36xbf16, #gpu.address_space>, vector<4xbf16>
%40 = amdgpu.mfma %36 * %32 + %arg4 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
%41 = amdgpu.mfma %37 * %33 + %40 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
%42 = amdgpu.mfma %36 * %34 + %arg5 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
%43 = amdgpu.mfma %37 * %35 + %42 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
%44 = amdgpu.mfma %38 * %32 + %arg6 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
%45 = amdgpu.mfma %39 * %33 + %44 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
%46 = amdgpu.mfma %38 * %34 + %arg7 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
%47 = amdgpu.mfma %39 * %35 + %46 {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xbf16>, vector<4xbf16>, vector<4xf32>
scf.yield %41, %43, %45, %47 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>
}
```
Contributor guide
Assessment
This issue has not been assessed yet.