tvm_tirx.tirx.script.jit.JITFunction

tvm_tirx.tirx.script.jit.JITFunction#

class tvm_tirx.tirx.script.jit.JITFunction(func, check_well_formed=True, persistent=False, private=False, *, definition_scope=None, builder=None)#

Bases: object

Top-level kernel decorator with compile-time .specialize() params.

Parses the function body lazily: parsing is deferred until .specialize() supplies concrete values for the params annotated as T.constexpr. The return type of .specialize() is a tvm_tirx.tirx.Function, identical in type to what @T.function produces today.

Constexpr params are removed from the resulting Function’s parameter list; their values are baked into the IR (e.g. into T.Tensor((M, K), ...) shape annotations and into the body).

Optional parameters accept annotation | None, typing.Optional[annotation] and T.Optional(annotation). Selecting None preserves their parameter slots with ir.NoneType; the body sees Python None for compile-time guards. Use from __future__ import annotations for deferred tensor annotations.

Parameters:

  • func (FunctionType)

  • check_well_formed (bool)

  • persistent (bool)

  • private (bool)

  • definition_scope (Mapping[str, Any] | None)

  • builder (object | None)

specialize(**const_args)#

Build a Function by binding constexprs and absent optional params.

Parameters:

**const_args (Any) – One hashable value per T.constexpr-annotated parameter. An optional parameter may additionally be supplied as None to retain its ABI slot with ir.NoneType. Omitting an optional parameter keeps it as a normal runtime parameter.

Returns:

Function – A concrete TIRx Function, identical in type to the output of @T.function. Repeated selections reuse the cached function.

Raises:

TypeError – If a name is not specializable, a required constexpr value is absent, an optional parameter is assigned anything other than None, or a constexpr value is not hashable.

Return type:

Function

Notes

No supplied kwargs still selects JIT construction with an empty map when no constexpr defaults exist. Selected None optional values retain parameters and skip their original annotations. Source and builder exceptions propagate unchanged.