Repository navigation
Add ahead-of-time compilation support to cuda builder #186
Description
From my knowledge/understanding, cuda_builder only supports JIT compilation. It would be beneficial for user adoption and performance for larger kernel sizes if we provided support for AOT compilation. Not sure if anything other than changes to cuda_builder would be necessary. Thanks @jorge-ortega for helping me flesh out this idea a little more.
Activity
jorge-ortega commentedon Apr 3, 2025
This post from Nvidia provides more context. Today, CudaBuilder uses the nvvm backend to compile crates to PTX. The host then loads and JITs through the driver API. Either the backend or CudaBuilder could pass the generated PTX to
ptxasto creates AOT compiled cubins that could be loaded by the driver.LegNeato commentedon Apr 4, 2025
I think it would be more idiomatic to treat this this as a different target or a a feature of the target like
target-cpu=native. Thoughts?adamcavendish commentedon Apr 4, 2025
I personally think that it's just a part of pipeline where Rust - PTX - fatbin. Maybe we should support a more complete pipeline and let user devide to what part should the builder build until?
jorge-ortega commentedon Apr 4, 2025
I agree that this is more of a build pipeline option. Our current pipeline is disjointed, and users have to glue the ptx into their host binaries themselves. We should target getting fatbins embedded into the final host binary to match what nvcc does. That's different from what's being asked here but does get us a step towards that.
LegNeato commentedon Apr 4, 2025
Sure, but I'm thinking for future integration in rustc...I actually think this maps pretty close to
crate_type!LegNeato commentedon Apr 8, 2025
Possibly useful techniques: https://github.com/calebzulawski/multiversion
adamcavendish commentedon Apr 9, 2025
For large language model optimizations, there are a lot of kernels that are written specialized for a specific NVidia card and using CPU to select based on the user input and the card used.
multiversionis probably useful but not flexible enough. Anyway, people can come back to the default match case to take back control and flexibility.LegNeato commentedon Apr 9, 2025
Right, that is why I said "techniques" rather than saying it is useful on its own 😁.
I think it is most idiomatic to use
crate_typefor jit vs AOT and target features for device-specific features (includingtarget-gpu=native), similar to what rustc uses for CPU (target-cpu) andrust-gpuuses for capabilities and extensions (https://github.com/Rust-GPU/rust-gpu/blob/698f10ac14b7c952394ac5620004e4e973308902/crates/spirv-std/src/arch.rs#L151).