[Target] Add __launch_bounds__ directive as part of the CUDA code generation

[ArmageddonKnight]

Short Summary

Sometimes, when executing CUDA kernels, we might encounter the error CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES (e.g., here). This happens because the nvcc compiler allocates too many registers per thread. In the case when we launch the CUDA kernel using too many threads, the GPU will notice that the CUDA kernel requests more registers than what are available on the chip and therefore refuse to launch the kernel.

This hence implies that we need a way of telling nvcc what to expect in terms of the number of threads per block. Luckily, the __launch_bounds__ directive can help us achieve what we want. In this patch, we add __launch_bounds__ as part of the CUDA code generation procedure. __launch_bounds__ will be automatically printed if it is detected that the number of threads per block is a constant integer value. Passing this information to nvcc allows it to spill registers if needed, which might hurt performance, but is still better than having a CUDA kernel that is not functional.

Q & A

Q: Would this affect the AutoTVM and the auto-scheduler submodule?

A: No. Although in those cases the number of threads keeps changing at each trial, the number will be set to a constant when it comes to the code generation phase. Furthermore, in the case when the number of threads per block is not a constant, __launch_bounds__ will simply not be printed.

Any feedback on this patch is appreciated. @comaniac @icemelon @yzhliu @yidawang

[comaniac]

LGTM. @Hzfengsy @Laurawly @masahi @tqchen it would be better if you folks could also take a look, as this will change the generated CUDA codes for all existing workloads.

[junrushao]

I am very much in favor of this approach. Thanks Bojian!

[junrushao]

The syntax should be __launch_bounds__(maxThreadsPerBlock) or __launch_bounds__(maxThreadsPerBlock, minBlocksPerMultiprocessor).

However, in the failing test, I am seeing:

grid=(16,8,1),  block=(32,2,2)
...
__launch_bounds__(1)

Would you like to double check? @ArmageddonKnight

[ArmageddonKnight]

@junrushao1994 Thanks for letting me know. I had a look into this issue. The problem is caused by assigning threadIdx.x to iv->thread_tag rather than iv->var->name_hint (and therefore, the extractor is unable to correctly extract the number of threads per block). To address this issue, I extend the extractor to cover thread_tag's as well. At the same time, if the number of threads per block is extracted as 1, then __launch_bounds__ will NOT be printed.

I ran the test case again and it works locally.

[junrushao]

Thanks Bojian! It makes perfect sense to me, which reminds me of a similar bug we encountered in TensorIR lowering :-)

[ArmageddonKnight]

@vinx13 @comaniac @junrushao1994 FYI, the patch has passed all the checks. Could you please merge it if possible?

[junrushao]

Thanks Bojian! It’s really super useful improvement in cuda codegen

[ArmageddonKnight]

Thanks @junrushao1994