Compilation

tvm_tirx.tirx.compile

tvm_tirx.tirx.compile(mod: Function | IRModule, target: str | Target | None = None, pipeline: str | Pass | None = 'default') → Executable

Compile a function or module to an executable for the target.

Parameters:

• mod (Union[Function, IRModule]) – The input to be compiled.

• target (Optional[Union[str, Target]]) – The target for compilation, including backend configuration defaults.

• pipeline (Union[None, str, tvm_tirx.ir.transform.Pass]) – The pipeline to use for compilation.

Returns:

An executable combining both host and device code.

Return type:

tvm_tirx.codegen.Executable

tvm_tirx.tirx.compilation_pipeline

The TIR backend compilation pipeline.

tvm_tirx.tirx.compilation_pipeline.default_tir_pipeline(*, target=None)

Return tirx_pipeline() with its standard lowering stages.

Parameters:

target (tvm_tirx.target.Target or None, optional) – Target to bind before lowering; None preserves existing function targets.

Returns:

stages – Lowering pass, host-finalization factory, and device-finalization factory.

Return type:

tuple

tvm_tirx.tirx.compilation_pipeline.tirx_pipeline(*, prepare_only=False, target=None)

Prepare TIRx functions and optionally lower their module-wide ABI.

Dispatch tile instructions, lower layouts, flatten tensors, simplify indices and loops, and verify memory accesses.

Parameters:

• prepare_only (bool, optional) – Stop before finalize_tir_pipeline() when true. Defaults to false.

• target (tvm_tirx.target.Target or None, optional) – Target to bind before lowering; None preserves existing function targets.

Returns:

stages – Lowering pass, host-finalization factory, and device-finalization factory.

Return type:

tuple

tvm_tirx.tirx.compilation_pipeline.finalize_tir_pipeline()

Split host/device functions, legalize low-precision types and lower the ABI.

Run after dialect-specific preparation and target binding. Annotates the entry function, lowers CUDA helpers and creates packed entry wrappers where required.

Returns:

passes – Module-wide finalization passes.

Return type:

tvm_tirx.ir.transform.Sequential

tvm_tirx.tirx.compilation_pipeline.finalize_host_passes()

Lower runtime builtins and target intrinsics in host functions.

Run after ABI lowering and host/device splitting, with a host target on each function.

Returns:

passes – Host finalization passes.

Return type:

tvm_tirx.ir.transform.Sequential

tvm_tirx.tirx.compilation_pipeline.finalize_device_passes()

Simplify device statements and lower target intrinsics.

Run after host/device splitting, with a target on each device function. Intrinsic expansion uses the target’s registered lowering rules.

Returns:

passes – Device finalization passes.

Return type:

tvm_tirx.ir.transform.Sequential

tvm_tirx.tirx.compilation_pipeline.register_tir_pipeline(name: str, pipeline_factory) → None

Register a factory returning lowering and host/device finalization stages.

Factories accepting an explicit target keyword own target binding. They must lower S-TIR thread bindings before binding the complete module and specializing helpers. Other factories receive inputs with targets already bound; they cannot classify calls inside unlowered S-TIR thread bindings.

Parameters:

• name (str) – Name used to select the pipeline; replaces any existing registration.

• pipeline_factory (Callable) – Factory returning the lowering pass and host/device finalization factories.

tvm_tirx.tirx.compilation_pipeline.get_tir_pipeline(name: str | None = None, *, target=None, **kwargs)

Construct a registered pipeline by name.

Parameters:

• name (str or None, optional) – Registered pipeline name. None selects "default".

• target (tvm_tirx.target.Target or None, optional) – Target to bind during lowering, after thread launches become explicit.

• **kwargs – Additional keyword arguments forwarded to the registered factory.

Returns:

stages – Lowering pass, host-finalization factory, and device-finalization factory.

Return type:

tuple

tvm_tirx.tirx.compilation_pipeline.register_default_tir_pipeline_selector(selector) → None

Register a callback selecting a pipeline for a module and target.

Parameters:

selector (Callable) – Callback accepting (module, target) and returning a pipeline tuple or None to try the next selector. Callbacks accepting target_to_bind pass it to the selected factory; other callbacks inspect the bound module, which is reused without binding it a second time.

tvm_tirx.tirx.compilation_pipeline.get_default_tir_pipeline(target: Target | None, mod: IRModule | None = None, *, target_to_bind: Target | None = None)

Select the default pipeline for the input module and target.

Parameters:

• target (tvm_tirx.target.Target or None) – Target supplied to dialect selectors to choose the pipeline.

• mod (tvm_tirx.ir.IRModule or None, optional) – Module inspected by dialect selectors. None uses the default registration.

• target_to_bind (tvm_tirx.target.Target or None, optional) – Target to bind during lowering; None preserves existing function targets.

Returns:

stages – Selected lowering pass, host-finalization factory, and device-finalization factory. Falls back to the "default" registration if no selector matches.

Return type:

tuple