tvm_tirx.script.ir_builder

tvm_tirx.script.ir_builder#

Shared IR classes are documented in tvm_tirx.ir.

Shared TVMScript construction APIs and lazy language variant builders.

class tvm_tirx.script.ir_builder.AlreadyEmitted(value: _T)#

Hold the exact native result of an already emitted builder operation.

Parameters:: value (T) – Native result retained without copying or emitting it. Void operator APIs wrap their original Call; statement APIs wrap the emitted statement. Allocation APIs wrap the bound Var or its scalar TensorLoad view.

value#

The same object, available for identity checks and source-location attachment. Builder consumers must check for this receipt before FFI conversion, so they do not emit the object again.

Type:: T

Notes

FFI conversion exposes only the wrapped object, not its emission history.

class tvm_tirx.script.ir_builder.IRBuilder#

A shared construction context for any language variant’s IR.

Examples

Enter the builder before using a language variant’s construction frames. Each completed frame contributes to the result returned by get.

from tvm_tirx.script.ir_builder import IRBuilder

def build(frame, emit_body):
    with IRBuilder() as builder:
        with frame:
            emit_body()
    return builder.get()

static current() → IRBuilder#

Get the current IRBuilder put in the with-scope.

Returns:: builder – The current IRBuilder.
Return type:: IRBuilder

static is_in_scope() → bool#

See if the current thread-local scope has an IRBuilder.

Returns:: Whether the current thread-local scope has an IRBuilder
Return type:: bool

get() → Object#: Get the constructed IR.

line_trace_scope(*, exclude_paths=())#

Track Python source lines while executing builder code.

TVM and Python implementation frames are excluded automatically. exclude_paths supplies additional implementation directories. User helpers retain their own line and the caller’s location. Nested scopes and exceptions restore the preceding trace and location state. Locations cover whole lines; use the parser for expression ranges. Tracing only supplies locations: script-facing void calls emit under any active builder, independently of this scope.

with_loc(loc)#

Attach loc to IR nodes constructed in the nested scope.

Nested scopes are retained as a CallSiteLoc when they describe distinct source ranges, such as a TVMScript inline expansion.

Parameters:: loc (tvm_tirx.ir.Location) – The frontend source range active in the nested scope.

static name(s: str, v: Any) → Any#

Set the name of an object.

Parameters:

s (str) – The name of the object.
v (Any) – The object to name.

Returns:

v – The same object with the name set.

Return type:

Any

static name_many(s: list[str], vs: list[Any]) → list[Any]#

Set the name of a list of objects.

Parameters:

s (List[str]) – The names of the objects.
vs (List[Any]) – The objects to name.

Returns:

vs – The same objects with the names set.

Return type:

List[Any]

class tvm_tirx.script.ir_builder.IRModuleFrame(loc, global_vars, functions, attrs, global_infos)#

tvm_tirx.script.ir_builder.Lambda(parameter_types, function, *, ret_type=None)#

Build a shared staging lambda from explicit types and a Python callable.

Scalar constructors such as T.float32 may be used as parameter types. An optional return annotation checks the body type without inserting casts.

tvm_tirx.script.ir_builder.at_(loc: LocationEntry | Location | None, value: _T) → _T#

Attach source context to the same IR node, emission receipt, or frame.

Parameters:

loc (LocationEntry, Location or None, optional) – Source location to compose with the active construction context. None leaves the value unchanged; stored locations remain non-nullable. UnknownLoc adds no source context.
value (Any) – Native object, AlreadyEmitted receipt, or list/tuple of native objects to annotate. A receipt’s contained object receives the loc.

Returns:

The exact value object, including its original receipt or container. With no active builder, the value is returned without modification.

Return type:

Any

Notes

Native mutation annotates the wrapped object in the builder IR. For an operator receipt this is the Call inside its Evaluate statement. Keep the original Python facade as well, including callable objects and frames. Unsupported objects and ordinary Python values pass through unchanged.

tvm_tirx.script.ir_builder.constexpr(value: object) → NoReturn#

Mark host control syntax or a JIT specialization annotation.

Parameters:: value (object) – Source expression to evaluate with ordinary Python semantics in a supported control-flow position. The marker itself can also appear as an annotation identifying a JIT specialization parameter.
Raises:: TypeError – If invoked directly instead of being recognized in parsed source.

Notes

The parser recognizes this marker through its fixed namespace path and removes it before execution. Host operators retain ordinary Python behavior. Direct invocation raises TypeError; no builder frame or IR is constructed.

# Source
if I.constexpr(enabled):
    T.evaluate(1)
# Builder
if enabled:
    X.emit_(X.evaluate(1))

tvm_tirx.script.ir_builder.decl_function(func_name: str, func_signature: BaseFunc) → GlobalVar#

Parameters:

func_name (str) – The function unique name.
func_signature (BaseFunc) – A function signature or completed function. A signature reserves the global variable for later native frame completion; a completed function is installed directly.

Note

Native function frames complete previously declared bodies while retaining the declared global variable identity.

Returns:: gv – The corresponding GlobalVar.
Return type:: GlobalVar

tvm_tirx.script.ir_builder.dtype(value: str | dtype) → DataTypeImm#

Construct a tvm_ffi.dtype-valued constant for an ordinary operation argument.

Unlike scalar constructors such as T.float32, this constructs a dtype value, rather than a scalar value or variable of that dtype.

tvm_tirx.script.ir_builder.dynamic(name: str, dtype: str = 'int32', *, loc: LocationEntry | Location = UnknownLoc()) → Var#

Create a fresh primitive symbolic variable, independently of builder scope.

Parameters:

name (str) – The symbol’s display name. Repeated names do not share identity.
dtype (str) – Primitive dtype, defaulting to int32.
loc (LocationEntry or Location, optional) – Source location of the symbol.

Returns:

A fresh symbol. Reuse this object to share dimensions across annotations and function bodies, including outside an I.ir_module definition.

Return type:

Var

tvm_tirx.script.ir_builder.ir_module() → IRModuleFrame#

Start a ir_module frame.

Returns:: frame – The constructed frame.
Return type:: IRModuleFrame

tvm_tirx.script.ir_builder.make_node(type_key, /, **kwargs)#

Make a new IR node by its type key and fields

Parameters:

type_key (str) – The type key of the node.
**kwargs (dict) – The fields of the node. For ir.DictAttrs, these are the dictionary entries.

Returns:

node – The corresponding IR Node

Return type:

Node

Note

If the created node is instance of AttrsObj, then the creator function will also run bound checks and default value setup as supported by Attrs.

Example

The following code constructs a IntImm object

x = tvm_tirx.ir.make_node(
    "ir.IntImm", ty=tvm_tirx.ir.PrimType("int32"), value=10, loc=tvm_tirx.ir.UnknownLoc()
)
assert isinstance(x, tvm_tirx.ir.IntImm)
assert x.value == 10

tvm_tirx.script.ir_builder.meta_var(value: T) → T#

Return a Python metadata value without binding, naming or relocating it.

Parameters:

value (T) – Any host or IR object, including an unpackable sequence.

Returns:

T – The exact input object. No frame is required, no IR is emitted, and existing names and source locations are retained.
.. code:: python – # Source and generated Python (the value retains its identity) value = I.meta_var(existing_value) a, b = I.meta_var((left, right))

tvm_tirx.script.ir_builder.module_attrs(attrs: dict[str, Object], allow_overwrite=False) → None#: Specify the attrs of the ir_module frame. :param attrs: The module attrs. :type attrs: Dict[str, Object] :param allow_overwrite: Whether allow overwrite the existing attrs. :type allow_overwrite: bool

tvm_tirx.script.ir_builder.module_get_attr(attr_key: str) → Object | None#

Get the specified attr of the ir_module frame. :param attr_key: The key of the attr to be retrieved. :type attr_key: str

Returns:: attr – The specified module attr or None if not found.
Return type:: Optional[Object]

tvm_tirx.script.ir_builder.module_global_infos(global_infos: dict[str, list[GlobalInfo]]) → None#

Specify the global infos of the ir_module frame.

Parameters:: global_infos (Dict[str, List[GlobalInfo]]) – The module global infos.

tvm_tirx.script.ir_builder.module_member_(name: str, value: Any) → Any#

Parameters:

name (str) – Source class member identifier.
value (Any) – Already-evaluated class member value; a BaseFunc is declared and defined as a module function.

Returns:

The reserved GlobalVar for a concrete BaseFunc, otherwise the exact original value.

Return type:

Any

Notes

Concrete function registration requires an active module frame and follows native duplicate/type checks. Other values need no frame and cause no IR emission or renaming. Shared source class setup runs before function signatures so module attributes/global info are available.

# Source
class Module:
    helper = existing_function
# Generated builder, inside I.ir_module
helper = I.module_member_("helper", existing_function)

tvm_tirx.script.ir_builder.module_set_attr(attr_key: str, attr_value: Object | None, allow_overwrite: bool = False) → None#: Set the specified attr of the ir_module frame. :param attr_key: The key of the attr to be set. :type attr_key: str :param attr_value: The value of the attr to be set. :type attr_value: Optional[Object] :param allow_overwrite: Whether allow overwrite the existing attr. :type allow_overwrite: bool

tvm_tirx.script.ir_builder.resolve_global_info_(content: str) → GlobalInfo#

Resolve a named global-info selector in the current module context.

Parameters:: content (str) – A named-list selector such as “mesh[0]”. The shared implementation has no device-specific syntax; language variants may extend this hook.
Returns:: The exact registered metadata object, without a new source location.
Return type:: GlobalInfo

Notes

Resolution uses the nearest active native module frame. Missing context or malformed selectors raise ValueError; missing names and indices propagate KeyError and IndexError. Non-string inputs raise TypeError. tvm_tirx.script.ir_builder.resolve_global_info_args() invokes this hook only for string arguments and preserves concrete constructor inputs.

@resolve_global_info_args("metadata", resolver=resolve_global_info_)
def CustomType(metadata):
    return make_type(metadata)

CustomType(metadata="mesh[0]")

tvm_tirx.script.ir_builder.resolve_global_info_args(*fields: str, resolver: Callable[[str], Any]) → Callable[[Callable[[...], Any]], Callable[[...], Any]]#

Resolve selected string arguments before calling a builder operation.

Parameters:

*fields (str) – Names of positional-only, positional-or-keyword, or keyword-only parameters. Repeated names are resolved once. Omitted arguments use their declared defaults.
resolver (Callable[[str], Any]) – Explicit callback receiving each selected string and returning its replacement. The callback owns selector syntax, lookup scope, and errors. Non-string values pass through with their identity preserved; containers are not decoded recursively.

Returns:

Decorator preserving the callable’s signature, name, and documentation. The signature is inspected once when decorating, then reused for argument binding.

Return type:

Callable

Raises:

ValueError – If a selected name is absent or names a variadic parameter.

Notes

All argument expressions are evaluated once in ordinary Python order before binding, resolution, and the callable body. Positional, keyword, unpacked, and aliased calls share this behavior. Selected string defaults are resolved on every call. Exceptions from binding, the resolver, and the callable propagate unchanged.

@resolve_global_info_args("device", resolver=lookup_device)
def tensor(shape, device="default"):
    return make_tensor(shape, device)

tvm_tirx.script.ir_builder.with_at_group_(location: LocationEntry | Location | None, thunk: Callable[[], _T], *, attach_result: bool = True) → _T#

Evaluate once under a location, optionally attaching it to the same result.

Parameters:

location (LocationEntry, Location) – Source context for the call. None, UnknownLoc, or no active builder leaves construction context unchanged.
thunk (Callable[[], Any]) – Zero-argument callable evaluated exactly once inside that context.
attach_result (bool, optional) – Attach the location to the returned object with at_(). Defaults to True. False supplies construction context only, leaving explicit result attachment to a later operation.

Returns:

The exact result of thunk, with its receipt or container preserved.

Return type:

Any

Notes

The prior source context is restored even if the callable raises; its exception propagates unchanged. Frames created during the call retain their construction locations regardless of attach_result.

tvm_tirx.script.ir_builder.Func#: alias of FuncType

tvm_tirx.script.ir_builder.Ptr#: alias of PtrType

tvm_tirx.script.ir_builder.String#: alias of StringType

tvm_tirx.script.ir_builder.Prim#: alias of PrimType

tvm_tirx.script.ir_builder.Opaque#: alias of OpaqueType

tvm_tirx.script.ir_builder.Any#: alias of AnyType

tvm_tirx.script.ir_builder.Missing#: alias of MissingType

tvm.relax.script.ir_builder.distributed#

Public distributed Relax construction and operator exports.

tvm.relax.script.ir_builder.distributed.DTensor(shape=None, dtype=None, device_mesh=None, placement='', *, ndim=-1, loc: LocationEntry | Location = UnknownLoc())#

Construct a Relax distributed tensor type.

Parameters:

shape (Expr or sequence of Expr, optional) – Global tensor shape, or None when unknown. Symbolic dimensions are expressions over explicit variables, such as those created with I.dynamic.
dtype (str or PrimType, optional) – Element type; None leaves the element type unknown.
device_mesh (DeviceMesh or str, optional) – Concrete mesh or module metadata selector. None creates an empty mesh placeholder. Strings require an active module builder. Use postponed annotations to defer resolution until function construction.
placement (Placement or str, optional) – Distribution placement. Text, including the default empty string, is parsed with Placement.from_text.
ndim (int, optional) – Global rank when shape is unknown; -1 means unknown rank.
loc (LocationEntry or Location, optional) – Source location attached to the constructed IR; UnknownLoc leaves it unspecified.

Returns:

result – The constructed distributed type. String selectors outside an active module raise ValueError.

Return type:

DTensorType

tvm.relax.script.ir_builder.distributed.call_tir(func, args, *, ty_args, attrs=None, ty=None, loc: Location = UnknownLoc(), **kwargs) → Call#

relax.Call a destination-passing TIR function and allocate its output tensors.

func is the function’s GlobalVar and args contains its ordered inputs. ty_args contains one output type: a TensorType or a TupleType of tensor results. Omitted ty uses the registered result inference; explicit types and locs are forwarded unchanged.

tvm.relax.script.ir_builder.distributed.call_tir_local_view(*args, attrs=None, ty_args=None, loc: Location = UnknownLoc(), ty=None, **kwargs)#: Construct a call to relax.dist.call_tir_local_view.

tvm.relax.script.ir_builder.distributed.const(value: bool | int | float | ndarray | Tensor, ty: DTensorType) → GenericConst#

Create a distributed constant value with the specified tensor type.

Parameters:

value (bool, int, float, numpy.ndarray or tvm_tirx.runtime.Tensor) – The constant value.
ty (DTensorType) – The distributed tensor type, including its tensor dtype, device mesh and placement.

Returns:

res – The constant carrying the supplied distributed tensor type.

Return type:

GenericConst

Notes

Python scalars and NumPy values are converted to the dtype specified by ty.tensor_ty. The dtype is not inferred from the Python value.

tvm.relax.script.ir_builder.distributed.device_mesh(shape: Shape, device_ids: list[int] | Range) → DeviceMesh#

Create a device mesh expression. :param shape: The shape of the device mesh. :type shape: Shape :param device_ids: Represents the device id in the mesh :type device_ids: Union[List[int], Range]

Returns:: res – The device mesh.
Return type:: DeviceMesh

tvm.relax.script.ir_builder.distributed.redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int, *, ty=None, loc: Location = UnknownLoc()) → Expr#

Slice tensor into several parts along one axis,: and each worker takes one part. input.ty.shape[axis] % num_workers == 0 is required. Each worker must have an identical copy of the input. This is a specialized version of redistribute op.

Parameters:

input (relax.Expr) – The buffer to be sliced into equal parts.
num_worker (int) – The number of workers, i.e. the number of parts the given buffer should be sliced into.
axis (int) – The axis of the tensor to be sliced.

Returns:

result – Sliced Tensor kept by each device.

Return type:

relax.Expr

tvm_tirx.script.ir_builder

Contents

tvm_tirx.script.ir_builder#

tvm_tirx.script.ir_builder#

tvm.relax.script.ir_builder.distributed#