changed 128 items per CTA to 256 items

[weihanmines]

No description provided.

[sryap]

@amathews-amd @dllehr-amd Thanks for submitting the PR. Can you provide the reason as to why this would help improve the performance of TBE? Which cases that this change will improve the performance? Does it improve performance of TBE on both AMD GPUs and NVIDIA GPUs? Thanks

[weihanmines]

Hi @sryap, This change is specific to AMD GPUs because of the thread organization (an execution unit consists of 64 threads on AMD GPUs, while there are 32 threads, named warp, on NVDIA GPUs). Therefore, 256 elements are processed by each execution unit on AMD GPUs. This reduces the size of accumulator shown in this link, https://github.com/ROCmSoftwarePlatform/FBGEMM/blob/main/fbgemm_gpu/codegen/embedding_forward_split_template.cu#L238. Most importantly by setting items_per_warp to 256, the unnecessary branches and computation are avoided. Please see the following link, https://github.com/ROCmSoftwarePlatform/FBGEMM/blob/main/fbgemm_gpu/codegen/embedding_forward_split_template.cu#L283-L284. Please let me know if you have any questions. Thank you.

[weihanmines]

@sryap BTW, CTA should be replaced with warp (or basic unit of execution) to be more precise.

[sryap]

Thanks for your explanation @weihanmines. This will reduce the number of registers allocated in the thread block but I don't think it will reduce the number of branches. Do you agree that the performance improvement is from the reduced number of registers which will allow for more thread blocks to be pipelined on a CU? If so, can you please help me do the analysis to verify that this is true (e.g., collecting some performance counter)? If not, can you please show some evidence to support your theory (other than just the improved TBE time)?

Since this is the specific change for AMD GPUs, can you please wrap the change with #ifdef __HIP_PLATFORM_HCC__?

Another ask: could you please revise the description of the PR to be more informative?

Thank you very much!

[weihanmines]

Yes, it will reduce the number of registers being used. I am not sure how many registers are allocated on AMD GPUs. I could not find a compilation option to report that. We have planed to use MIPerf to profile the new code to take a look at new metrics. However, the server was not available for some reason in the past week. It will reduce the branches and computations. Please take a closer look at the link I sent in my previous reply.
We already have commits which take care of your suggestion and they are waiting to be merged. I will write up a more detailed description once we have profiling data ready.
Thanks.

[sryap]

No worries. Please share your findings when you have them.

I think it will change the number of branches because there are two conditions: i < kMaxVecsPerThread and 4 * kWarpSize * i + threadIdx.x * 4 < D. The number of branches is controlled by 4 * kWarpSize * i + threadIdx.x * 4 < D. I think i < kMaxVecsPerThread is there for correctness check. For example, for D=256, the number of branches is constant irrespective to the kMaxVecsPerThread value. Please correct me if I'm wrong.

[weihanmines]

Hi @sryap, we have confirmed that the occupancy stays the same with/without the changes. So I don't think that the performance gain is due to less registers after the change. The reasons should be less branches and computation, I believe.

[amathews-amd]

merged in pytorch#1240