Function tvm_callback_cuda_compile

python/tvm/support/nvcc.py:852–892

Compile CUDA code using the configured backend (nvcc or nvrtc). This callback is invoked by TVM's C++ backend during CUDA module compilation. By default, uses nvrtc to generate cubin. The current target is fetched inside the callback (via ``tvm.target.Target.current(allow_none=Tru

(code)

@tvm_ffi.register_global_func
def tvm_callback_cuda_compile(code):
    """
    Compile CUDA code using the configured backend (nvcc or nvrtc).

    This callback is invoked by TVM's C++ backend during CUDA module compilation.
    By default, uses nvrtc to generate cubin.  The current target is fetched
    inside the callback (via ``tvm.target.Target.current(allow_none=True)``)
    so the caller does not need to push/pop a target scope around the
    invocation.

    Environment Variables
    ---------------------
    TVM_CUDA_COMPILE_MODE : str
        Compiler backend: "nvrtc" (default) or "nvcc"
        - "nvrtc": Use NVRTC via cuda-bindings for faster JIT, generates cubin
        - "nvcc": Use nvcc subprocess, generates fatbin
    TVM_KERNEL_DUMP : str
        If set, dump generated CUDA/intermediate files and append "-lineinfo" so profilers can
        correlate SASS back to the dumped source.

    Parameters
    ----------
    code : str
        CUDA source code to compile

    Returns
    -------
    bytes
        Compiled binary (fatbin for nvcc, cubin for nvrtc)
    """
    # The current Target is fetched inside compile_cuda via
    # tvm.target.Target.current(allow_none=True) when arch is unset; the
    # caller no longer needs to push/pop a target scope.
    compiler = os.environ.get("TVM_CUDA_COMPILE_MODE", "nvrtc").lower()

    if compiler == "nvrtc":
        return compile_cuda(code, target_format="cubin", compiler="nvrtc")
    if compiler == "nvcc":
        return compile_cuda(code, target_format="fatbin", compiler="nvcc")

    raise ValueError(f"Invalid TVM_CUDA_COMPILE_MODE: {compiler}. Expected 'nvcc' or 'nvrtc'.")


@tvm_ffi.register_global_func("tvm_callback_libdevice_path")