Compile CUDA code with NVCC or NVRTC. Parameters ---------- code : str The CUDA code. target_format : str The target format of the compiler ("ptx", "cubin", or "fatbin"). arch : str The CUDA architecture. options : str or list of str The ad
(
code, target_format=None, arch=None, options=None, path_target=None, compiler="nvrtc"
)
| 32 | |
| 33 | |
| 34 | def compile_cuda( |
| 35 | code, target_format=None, arch=None, options=None, path_target=None, compiler="nvrtc" |
| 36 | ): |
| 37 | """Compile CUDA code with NVCC or NVRTC. |
| 38 | |
| 39 | Parameters |
| 40 | ---------- |
| 41 | code : str |
| 42 | The CUDA code. |
| 43 | |
| 44 | target_format : str |
| 45 | The target format of the compiler ("ptx", "cubin", or "fatbin"). |
| 46 | |
| 47 | arch : str |
| 48 | The CUDA architecture. |
| 49 | |
| 50 | options : str or list of str |
| 51 | The additional options. |
| 52 | |
| 53 | path_target : str, optional |
| 54 | Output file. |
| 55 | |
| 56 | compiler : str, optional |
| 57 | Compiler backend: "nvrtc" (default) or "nvcc". |
| 58 | This can be set by the TVM_CUDA_COMPILE_MODE environment variable. |
| 59 | |
| 60 | Returns |
| 61 | ------- |
| 62 | res_binary : bytearray |
| 63 | The bytearray of the compiled binary (ptx/cubin/fatbin). |
| 64 | |
| 65 | Notes |
| 66 | ----- |
| 67 | - NVRTC is a "runtime" compilation library and can be faster for JIT compilation. |
| 68 | - NVRTC requires cuda-bindings: pip install cuda-bindings |
| 69 | """ |
| 70 | use_nvshmem = "#include <nvshmem.h>" in code or "#include <nvshmemx.h>" in code |
| 71 | |
| 72 | if compiler == "nvcc": |
| 73 | result = _compile_cuda_nvcc(code, target_format, arch, options, path_target, use_nvshmem) |
| 74 | elif compiler == "nvrtc": |
| 75 | result = _compile_cuda_nvrtc(code, target_format, arch, options, path_target, use_nvshmem) |
| 76 | else: |
| 77 | raise ValueError(f"CUDA compiler must be 'nvcc' or 'nvrtc', got: {compiler}") |
| 78 | |
| 79 | return result |
| 80 | |
| 81 | |
| 82 | def _compile_cuda_nvcc( |
no test coverage detected
searching dependent graphs…