Backend Internals

Backend Internals#

TIRx keeps the core IR and compiler machinery under tvm.tirx. A target backend owns its script namespaces, intrinsic builders, dispatch variants, pipeline additions, and code-generation support under tvm.backend.

Loading and registration#

tvm.backend.load imports a backend and calls its registration hook. A backend can register the following target-owned behavior:

  1. TVMScript namespaces such as Tx.cuda, Tx.ptx, and Tx.nki;

  2. tile-primitive dispatch implementations for that target kind;

  3. target tags and, when the target owns one, compilation-pipeline entry points; and

  4. code-generation helpers that translate backend calls to target source.

Registration imports some modules only for their side effects. Those modules are implementation surfaces, not additional kernel-authoring APIs.

CUDA ownership#

tvm.backend.cuda is divided by compiler responsibility:

Module

Responsibility

script and ptx

Construct the Tx.cuda, Tx.ptx, compatibility, and NVSHMEM namespaces.

op

Define IR builders used by the CUDA script namespaces.

tile_primitive

Register CUDA implementations of common Tx.tile operations.

codegen and cpp

Register source-generation callbacks and CUDA C++ helpers.

transforms

Provide CUDA-specific compiler passes.

target_tags

Register named NVIDIA targets.

lang and iket

Provide reusable kernel utilities and profiling orchestration.

The PTX namespace is table-driven. A table entry defines the legal modifier and operand forms; construction creates a tirx.ptx.* call, and CUDA codegen emits the corresponding instruction. See Direct PTX Instructions for the supported public forms.

Trainium ownership#

tvm.backend.trn follows the same boundary. script and op construct the Tx.nki surface, tile_primitive registers target dispatches, layout and transform lower Trainium-specific memory mappings, and pipeline assembles the Trainium pass sequence. target_tags registers the named AWS Trainium targets.

Public integration hooks are listed in Backend Extension API; target-facing CUDA and Trainium utilities are listed in CUDA Authoring and Support APIs and Trainium Authoring and Support APIs.

Contents