tvm.script.ir_builder

tvm.script.ir_builder#

Shared TVMScript construction APIs and lazy language variant builders.

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

Hold the exact emitted value; location handling preserves this receipt.

Parameters:: value (Any) – Object already emitted by a language variant’s builder. The receipt retains this object without copying it. Binding or emitting the receipt must not emit the object again.

value#

The same emitted object, available for identity checks and source-span attachment.

Type:: Any

class tvm.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.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.

with_source_span(span)#

Attach span to IR nodes constructed in the nested scope.

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

Parameters:: span (tvm.ir.Span) – 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.script.ir_builder.IRModuleFrame(source_span, global_vars, functions, attrs, global_infos)#

tvm.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.script.ir_builder.at_(span: SpanEntry | Span | None, value: _T) → _T#

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

Parameters:

span (SpanEntry, Span or None) – Source location to compose with the active construction context. None leaves the value unchanged.
value (Any) – Native object, AlreadyEmitted receipt, or list/tuple of native objects to annotate. A receipt’s contained object receives the span.

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 statement held by the builder itself. Keep the original Python facade as well, including callable objects and frames. Unsupported objects and ordinary Python values pass through unchanged.

tvm.script.ir_builder.check_well_formed_(module: IRModule) → None#

Validate a completed module after every body has been finalized.

Parameters:: module (IRModule) – Completed native module, including all resolved forward references and mixed- language variant members.
Returns:: No value or mutation; success means the existing native validators accepted the program.
Return type:: None

Notes

Requires completed IR, with no active construction frame. Root coordination invokes opaque whole-module hooks supplied by language variant initialization through tvm.script.register_module_validator. Each hook owns concrete types, eligibility and cross-function validation, including captured/preexisting members. No source decorator inventory selects the hooks. Validator exceptions propagate unchanged; an empty hook list raises RuntimeError instead of establishing validity. check_well_formed=False omits the generated call entirely.

# Source
@I.ir_module
class Module:
    pass
# Generated builder, after module frame exit
I.check_well_formed_(module)

tvm.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.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.script.ir_builder.dtype(value: str | dtype) → DataTypeImm#

Construct a DataType-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.script.ir_builder.dynamic(name: str, dtype: str = 'int64', *, span: SpanEntry | Span | None = None) → 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 int64.
span (SpanEntry, Span or None) – 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.script.ir_builder.ir_module() → IRModuleFrame#

Start a ir_module frame.

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

tvm.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.

Returns:

node – The corresponding IR Node

Return type:

Node

Note

If the created node is instance of AttrsNode, 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.ir.make_node("ir.IntImm", dtype="int32", value=10, span=None)
assert isinstance(x, tvm.tirx.IntImm)
assert x.value == 10

tvm.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.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.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.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.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.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.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 span.
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.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.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.script.ir_builder.with_at_group_(location: SpanEntry | Span | None, thunk: Callable[[], _T], *, attach_result: bool = True) → _T#

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

Parameters:

location (SpanEntry, Span or None) – Source context for the call. None, or the absence of an 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 spans regardless of attach_result.

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, span=None)#

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.
span (SpanEntry, Span or None, optional) – Source location attached to the constructed IR; None 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: str | Expr, args: Expr, out_ty: DTensorType | list[DTensorType]) → Call#

Distributed version of call_tir

Parameters:

func (Union[str, Expr]) – The destination-passing-style function, can be ExternFunc or Function.
args (Expr) – The ordered distributed-tensor and primitive input arguments. These correspond positionally to the leading parameters of the Function.
out_ty (Union[DTensorType, List[DTensorType]]) – The type information of the call_tir output. It should be a single or a list of DTensorType. Each one denotes the type information of a returned distributed tensor.

Returns:

ret – A call node for the call_tir operator.

Return type:

relax.Call

tvm.relax.script.ir_builder.distributed.call_tir_local_view(gvar: GlobalVar, args: Expr, out_ty: DTensorType | list[DTensorType] | None = None, *, ty_args=None, ty=None, span=None) → Call#

relax.Call a tirx.function and return the output. The function should be a worker-local function that is actually executed on each worker, instead of the unpartitioned function. The output of this operator is DTensor or a tuple of DTensors.

Parameters:

gvar (GlobalVar) – The GlobalVar referring to a tirx Function.
args (Expr) – The ordered distributed-tensor and primitive input arguments. These correspond positionally to the leading parameters of the Function.
out_ty (Union[DTensorType, List[DTensorType]]) – The type information of the call_tir output. It should be a single or a list of DTensorType. Each one denotes the type information of a returned tensor.

Returns:

ret – A call node for the call_tir_local_view operator.

Return type:

relax.Call

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.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, span=None) → 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.script.ir_builder

Contents

tvm.script.ir_builder#

tvm.script.ir_builder#

tvm.relax.script.ir_builder.distributed#