Implement asynchronous LDS loads for MI350

[avbokovoy]

This PR implements direct HBM->LDS stores in tbe inference kernel. There are 2 major changes:

1. Rows data isn't loaded in-place, instead we store pointers to global memory and store the actual data w.r.t. the predicate into LDS. In case predicate is false, we pre-allocate small chunk of static device memory of 16B once, fill it with zeros, and fallback to this chunk
2. HBM->LDS 16B loads are implemented for ROCm >= 7.0 and MI350. We can expand the support range to MI30* through 4B loads, however it doesn't bring any performance benefits because we'll have to introduce an overhead of addresses transposition and 4x more load operations. You can find out the reference implementation here: pytorch@fe52557.

Due to pre-7.2 ROCm features, we are forced to used assembly inline to get 16B loads to work, so manual synchronization was added. In case of ROCm >= 7.2, we use proper intrinsics to handle memory synchronization.

This change brings ~10% performance boost on average for weighted and unweighted cases. We may try to push it further by doing async loads for indices weights.

cc: @amirakb89 you might be interested

[aryaman-gupta]

I have taken a look at the PR and identified a couple of areas that looked tricky to me. It may make sense to double-check these to confirm that the logic is correct

[aryaman-gupta]

This enforces that the entire warp will load a contiguous chunk of memory from global to LDS. What happens when the row is not large enough, i.e., kWarpSize > NumUint4LoadsPerRow? As I understand it, this would assign different row_load_idx to different lanes in the wavefront

FBGEMM/fbgemm_gpu/codegen/inference/embedding_forward_quantized_split_nbit_kernel_template.cu

Line 153 in 856f1af

uint32_t row_load_idx = load_idx % NumUint4LoadsPerRow;

You might want to confirm that this case is correctly handled

[avbokovoy]

It will load 16 bytes (16 x 64 for the whole wave) from corresponding vector register (address is different from lane to lane) into LDS pointer with corresponding strides. Global memory doesn't have to be contiguous. The sanity of the loads are checked outside of this function and is handled with pred_guard. Tailing or OOB loads are redirected to zero_tile global memory chunk, which contains zeroes. It's then handled properly by kernel logic

[aryaman-gupta]

Does this mean that there is no wait instruction defined for ROCm version >= 7.2? Just wanted to confirm that this is intentional

[avbokovoy]

It should be handled by intrinsic, which assumes that this functions has a side effect

[aryaman-gupta]

It seems to me that in PackedMode, the smem_ptr passed to the different lanes in the cp_async_zfill_cg function is strided. However, the cp_async_zfill_cg function uses lane 0's smem_ptr and performs a contiguous memory read into that location. This seems suspicious to me, so I wanted to point it out. I suppose you have verified that the logic is correct @avbokovoy ?

[avbokovoy]

I guess this comment applies here as well:
#138 (comment)