tilelang.contrib.nvcc module#

Utility to invoke nvcc compiler in the system

tilelang.contrib.nvcc.callback_libdevice_path(arch)#
tilelang.contrib.nvcc.compile_cuda(code, target_format='ptx', arch=None, options=None, path_target=None, verbose=False)#

Compile cuda code with NVCC from env.

Parameters:

  • code (str) – The cuda code.

  • target_format (str) – The target format of nvcc compiler.

  • arch (str) – The cuda architecture.

  • options (str or list of str) – The additional options.

  • path_target (str, optional) – Output file.

Returns:

cubin – The bytearray of the cubin

Return type:

bytearray

tilelang.contrib.nvcc.find_cuda_path()#

Utility function to find cuda path

Returns:

path – Path to cuda root.

Return type:

str

tilelang.contrib.nvcc.get_cuda_version(cuda_path=None)#

Utility function to get cuda version

Parameters:

cuda_path (Optional[str]) – Path to cuda root. If None is passed, will use find_cuda_path() as default.

Returns:

version – The cuda version

Return type:

float

tilelang.contrib.nvcc.have_cudagraph()#

Either CUDA Graph support is provided

tilelang.contrib.nvcc.have_fp16(compute_version)#

Either fp16 support is provided in the compute capability or not

Parameters:

compute_version (str) – compute capability of a GPU (e.g. “6.0”)

tilelang.contrib.nvcc.have_int8(compute_version)#

Either int8 support is provided in the compute capability or not

Parameters:

compute_version (str) – compute capability of a GPU (e.g. “6.1”)

tilelang.contrib.nvcc.have_tensorcore(compute_version=None, target=None)#

Either TensorCore support is provided in the compute capability or not

Parameters:

  • compute_version (str, optional) – compute capability of a GPU (e.g. “7.0”).

  • target (tvm.target.Target, optional) – The compilation target, will be used to determine arch if compute_version isn’t specified.

tilelang.contrib.nvcc.parse_compute_version(compute_version)#

Parse compute capability string to divide major and minor version

Parameters:

compute_version (str) – compute capability of a GPU (e.g. “6.0”)

Returns:

  • major (int) – major version number

  • minor (int) – minor version number