CUDA Backend Development

This guide shows how to find the source files involved in extending or debugging the CUDA backend. We will follow T.cuda.thread_idx("x") from the GPU add-one example to its IR operator and CUDA output, then locate the compilation and runtime entry points.

The source tree

CUDA support has three main parts:

Directory	Responsibility
include/tvm/backend/cuda/	Public C++ operator declarations, attributes, and backend contracts.
src/backend/cuda/	Native operators, lowering passes, code generation, and runtime support.
python/tvm_tirx/backend/cuda/	Script namespaces, PTX instruction tables, tile patterns, compiler configuration, and profiling orchestration.

Trace a CUDA operation

The add-one kernel uses T.cuda.thread_idx("x") to get each thread’s position within its block. This call connects three parts of the backend:

TVMScript: T.cuda.thread_idx("x")
    -> IR call: tirx.cuda.thread_idx, argument "x", result type int32
    -> CUDA C++: ((int)threadIdx.x)

The script namespace is installed by register_backend() in python/tvm_tirx/backend/cuda/__init__.py. python/tvm_tirx/backend/cuda/script.py exposes the registered CUDA operators under T.cuda.

The operator itself is registered in src/backend/cuda/op/target_builtin.cc. Its definition requires an axis string and infers an int32 result. TVMScript constructs an ir.Call referring to that operator, with "x" as its argument.

CodeGenCUDA::Dispatch_ in src/backend/cuda/codegen/codegen_cuda.cc recognizes the operator and emits ((int)threadIdx.x). These locations let us trace the script interface, argument rules, and generated expression separately when investigating an incorrect result.

PTX instruction calls use Python code-generation callbacks. The CUDA source generator looks them up through tirx.intrinsics.cuda.get_codegen, registered in python/tvm_tirx/backend/cuda/codegen/registry.py. The PTX TableGen guide follows that path for add.rn.f32.

Compilation and launch

The shared pipeline in python/tvm_tirx/tirx/compilation_pipeline.py lowers tensor accesses and separates device computation from host code. SplitHostDevice extracts a device region into a kernel and leaves a launch in the host function. MakePackedAPI prepares the host entry point that accepts the input and output tensors.

Once the CUDA source has been generated, python/tvm_tirx/backend/cuda/compiler.py compiles it with NVCC or NVRTC. The runtime implementation in src/backend/cuda/runtime/cuda_module.cc loads the device module and launches kernels with their grid, block, and shared-memory settings. The Compiling and Inspecting guide shows how to inspect the generated source.

Where to make a change

The following locations provide starting points for common changes:

Task	Starting point
Add a PTX instruction form	python/tvm_tirx/backend/cuda/ptx/table.py and tests/python/tirx/codegen/test_ptx_dialect.py.
Add a tensor instruction pattern	python/tvm_tirx/backend/cuda/tile/patterns/ and tests/python/tirx/tile/cuda/.
Change source compilation	python/tvm_tirx/backend/cuda/compiler.py, python/tvm_tirx/backend/cuda/backend_config.py, and python/tvm_tirx/support/nvcc.py.
Change CUDA axis semantics	src/backend/cuda/layout/axis.cc and python/tvm_tirx/backend/cuda/axis.py.
Change device loading or launching	src/backend/cuda/runtime/ and src/backend/cuda/launch/.
Change profiling annotation lowering	LowerIket in src/backend/cuda/transform/lower_iket.cc.

The CUDA backend API documents the public interfaces in these modules.