CUDA backend: compile

[zcbenz]

This PR is split from #1983.

Adds a JitModule class that wraps NVRTC APIs, and implements the Compiled primitive.

The NVRTC API allows passing a list of headers for search in #include, so unlike the Metal backend that uses a script to preprocess the kernel files, the CUDA backend takes a platform-independent approach by embedding the files under backend/cuda/kernels/ into char arrays in a generated header, and then pass the list to NVRTC.

There are also a few implementation details caused limitations of NVRTC:

• The cuLaunchKernel API which is used for launching compiled kernel, requires the arguments to be passed as an array of pointers, to avoid errors like passing pointers to temporary values and to simplify the usage, the passed args are copied to a list of variants managed by JitModule.
• Getting a kernel function from compiled module requires querying with mangled C++ function name, while NVRTC provides APIs to get mangled name from C++ name, the API can only be used after compiling the program. So when storing the compiled module in disk cache, we have to also store the map of kernel names.

[awni]

This is a bit of a nit, but the directory layout has a confusing naming scheme. It seems like kernels/ is not actually where the "kernels" go.. but where some of the utilities for the kernels go.

Can we reconsider the name of this subdirectory? Or alternatively just flatten the directory structure all into cuda and explicitly list the sources we want to include here?

[zcbenz]

The files under kernels/ must be compatible to JIT compilation, having more limitations than normal CUDA code, the biggest difference is that they can not contain any host-only code. I used "kernels" as name to keep consistency with Metal backend, a more precise name could be "device"?

[awni]

Yea I like device more. That's a great suggestion!

[zcbenz]

I added a commit to rename kernels/ to device/, it is going to cause lots of conflicts with other PRs so you might want to merged that one later.

[awni]

Thanks! I'll manage the conflicts.. should be fairly straightforward.