Compile CUDA code using nvcc. Parameters ---------- code : str The CUDA source code. target_format : str, optional Output format: "ptx", "cubin", or "fatbin". arch : str, optional Target architecture. Auto-detected if None. options : str or list of st
(
code,
target_format=None,
arch=None,
options=None,
path_target=None,
use_nvshmem=False,
)
| 80 | |
| 81 | |
| 82 | def _compile_cuda_nvcc( |
| 83 | code, |
| 84 | target_format=None, |
| 85 | arch=None, |
| 86 | options=None, |
| 87 | path_target=None, |
| 88 | use_nvshmem=False, |
| 89 | ): |
| 90 | """Compile CUDA code using nvcc. |
| 91 | |
| 92 | Parameters |
| 93 | ---------- |
| 94 | code : str |
| 95 | The CUDA source code. |
| 96 | target_format : str, optional |
| 97 | Output format: "ptx", "cubin", or "fatbin". |
| 98 | arch : str, optional |
| 99 | Target architecture. Auto-detected if None. |
| 100 | options : str or list of str, optional |
| 101 | Additional nvcc options. |
| 102 | path_target : str, optional |
| 103 | Output file path. |
| 104 | |
| 105 | Returns |
| 106 | ------- |
| 107 | bytearray |
| 108 | Compiled binary data. |
| 109 | """ |
| 110 | # Check for NVSHMEM dependency |
| 111 | nvshmem_include_path, nvshmem_lib_path = None, None |
| 112 | if use_nvshmem: |
| 113 | # NOTE: we cannot check whether nvshmem is used based on whether |
| 114 | # the global function "runtime.nvshmem.cumodule_init" is defined. |
| 115 | # The reason is because that if the input code does not use any NVSHMEM functions |
| 116 | # while the global function is defined, using cubin to compile the |
| 117 | # code may cause a compilation error. |
| 118 | target_format = "cubin" |
| 119 | nvshmem_include_path, nvshmem_lib_path = find_nvshmem_paths() |
| 120 | |
| 121 | if arch is None: |
| 122 | # If None, then it will use `tvm.target.Target.current().arch`. |
| 123 | # Target arch could be a str like "sm_xx", or a list, such as |
| 124 | # [ |
| 125 | # "-gencode", "arch=compute_52,code=sm_52", |
| 126 | # "-gencode", "arch=compute_70,code=sm_70" |
| 127 | # ] |
| 128 | compute_version = "".join( |
| 129 | get_target_compute_version(Target.current(allow_none=True)).split(".") |
| 130 | ) |
| 131 | arch = ["-gencode", f"arch=compute_{compute_version},code=sm_{compute_version}"] |
| 132 | |
| 133 | temp = utils.tempdir() |
| 134 | file_name = "tvm_kernels" |
| 135 | if target_format is None and not use_nvshmem: |
| 136 | target_format = "ptx" |
| 137 | |
| 138 | tvm_kernel_dump = os.environ.get("TVM_KERNEL_DUMP", None) |
| 139 | if tvm_kernel_dump is not None: |
no test coverage detected
searching dependent graphs…