You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/09/21 20:54:11 UTC
(tvm) branch tvmscript-ast-only-transpiler updated (78291ea37b -> c86de02152)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch tvmscript-ast-only-transpiler
in repository https://gitbox.apache.org/repos/asf/tvm.git
discard 78291ea37b Complete missing Relax value annotations before normalization
discard c4ede00b36 Route global calls through ordinary callable overloads
discard c6ba25159f Resolve named outputs through explicit completed builder frames
discard f7de503247 Remove redundant import lint suppression
discard 063ba1330b Build keyword expressions eagerly with conditional runtime evaluation
discard e073ec33b6 Share source metadata and make span instrumentation optional
discard 801caf7498 Collect registered declarations in the shared source prescan
discard a72f369da1 Register parameter declarations in the parser protocol
discard 5b7969b327 Own syntax registration in the parser protocol
discard 1a23bc7059 Lower TVMScript syntax through explicit builder programs
new c86de02152 [TVMScript] Stage an opt-in AST-only parser with builder-owned construction
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (78291ea37b)
\
N -- N -- N refs/heads/tvmscript-ast-only-transpiler (c86de02152)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
include/tvm/relax/script/builder/ir.h | 18 +-
python/tvm/ir/__init__.py | 11 +-
python/tvm/ir/expr.py | 33 +-
python/tvm/ir/type.py | 11 -
python/tvm/relax/script/__init__.py | 2 +-
python/tvm/relax/script/builder/__init__.py | 249 +--
python/tvm/relax/script/builder/v2.py | 1615 ++++++++++++++++++
python/tvm/s_tir/tensor_intrin/rocm.py | 11 +-
python/tvm/script/__init__.py | 6 +-
python/tvm/script/ir_builder/__init__.py | 7 +-
python/tvm/script/ir_builder/base.py | 142 +-
python/tvm/script/ir_builder/construction.py | 447 ++++-
python/tvm/script/ir_builder/ir/__init__.py | 3 +-
python/tvm/script/ir_builder/protocol.py | 374 ++--
python/tvm/script/ir_builder/type_var_frame.py | 307 +++-
python/tvm/script/parser/__init__.py | 38 +-
python/tvm/script/parser/annotations.py | 622 +++++++
python/tvm/script/parser/diagnostics.py | 31 +-
python/tvm/script/parser/frontend.py | 676 +++++---
python/tvm/script/parser/functions.py | 237 +++
python/tvm/script/{ => parser}/ir.py | 6 +-
python/tvm/script/parser/protocol.py | 225 ---
python/tvm/script/parser/transform.py | 782 +++++++++
.../tvm/script/{parser => parser_v2}/__init__.py | 27 +
.../script/{parser => parser_v2}/diagnostics.py | 0
.../tvm/script/{parser => parser_v2}/expression.py | 36 +-
.../tvm/script/{parser => parser_v2}/frontend.py | 445 +++--
python/tvm/script/parser_v2/protocol.py | 394 +++++
.../tvm/script/{parser => parser_v2}/transpile.py | 300 +++-
python/tvm/tirx/script/builder/__init__.py | 390 +----
python/tvm/tirx/script/builder/comparison.py | 82 +
python/tvm/tirx/script/builder/ir.py | 80 +-
python/tvm/tirx/script/builder/v2.py | 1805 ++++++++++++++++++++
src/relax/analysis/type_analysis.cc | 12 +-
src/relax/backend/vm/vm_shape_lower.cc | 17 +-
src/relax/ir/block_builder.cc | 6 +-
src/relax/ir/dependent_type.cc | 9 +-
src/relax/ir/expr.cc | 9 +-
src/relax/script/builder/ir.cc | 49 +-
src/relax/script/printer/tir.cc | 6 +-
src/relax/transform/static_plan_block_memory.cc | 20 +-
src/script/printer/ir/ir.cc | 3 +-
src/script/printer/utils.h | 42 +-
src/tirx/script/printer/expr.cc | 7 +-
src/tirx/transform/lower_tvm_builtin.cc | 59 +-
tests/python/relax/test_analysis_well_formed.py | 2 +-
.../relax/test_backend_dispatch_sort_scan.py | 2 -
.../relax/test_backend_transform_shape_lower.py | 1 -
tests/python/relax/test_blockbuilder_core.py | 3 +-
tests/python/relax/test_builder_annotations.py | 199 ---
tests/python/relax/test_expr.py | 5 +-
tests/python/relax/test_shape_i32.py | 147 --
tests/python/relax/test_tvmscript_parser.py | 73 +-
tests/python/relax/test_tvmscript_type_vars.py | 40 -
.../python/tvmscript/test_builder_short_circuit.py | 246 ---
tests/python/tvmscript/test_parser.py | 352 +---
.../python/tvmscript/test_tvmscript_parser_tir.py | 23 +-
.../python/tvmscript/test_tvmscript_printer_tir.py | 4 +-
58 files changed, 7723 insertions(+), 3025 deletions(-)
create mode 100644 python/tvm/relax/script/builder/v2.py
create mode 100644 python/tvm/script/parser/annotations.py
create mode 100644 python/tvm/script/parser/functions.py
rename python/tvm/script/{ => parser}/ir.py (87%)
delete mode 100644 python/tvm/script/parser/protocol.py
create mode 100644 python/tvm/script/parser/transform.py
copy python/tvm/script/{parser => parser_v2}/__init__.py (69%)
copy python/tvm/script/{parser => parser_v2}/diagnostics.py (100%)
rename python/tvm/script/{parser => parser_v2}/expression.py (87%)
copy python/tvm/script/{parser => parser_v2}/frontend.py (53%)
create mode 100644 python/tvm/script/parser_v2/protocol.py
rename python/tvm/script/{parser => parser_v2}/transpile.py (86%)
create mode 100644 python/tvm/tirx/script/builder/comparison.py
create mode 100644 python/tvm/tirx/script/builder/v2.py
delete mode 100644 tests/python/relax/test_builder_annotations.py
delete mode 100644 tests/python/relax/test_shape_i32.py
delete mode 100644 tests/python/tvmscript/test_builder_short_circuit.py
(tvm) 01/01: [TVMScript] Stage an opt-in AST-only parser with builder-owned construction
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tvmscript-ast-only-transpiler
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit c86de02152ab936a0b822e6b0cf26004e94143b4
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Sep 21 20:52:56 2026 +0000
[TVMScript] Stage an opt-in AST-only parser with builder-owned construction
---
include/tvm/script/ir_builder/base.h | 38 +-
python/tvm/relax/script/builder/v2.py | 1615 +++++++++++++++++++++
python/tvm/script/ir_builder/base.py | 142 +-
python/tvm/script/ir_builder/construction.py | 647 +++++++++
python/tvm/script/ir_builder/type_var_frame.py | 343 +++++
python/tvm/script/parser_v2/__init__.py | 79 ++
python/tvm/script/parser_v2/diagnostics.py | 84 ++
python/tvm/script/parser_v2/expression.py | 246 ++++
python/tvm/script/parser_v2/frontend.py | 738 ++++++++++
python/tvm/script/parser_v2/protocol.py | 394 ++++++
python/tvm/script/parser_v2/transpile.py | 1506 ++++++++++++++++++++
python/tvm/tirx/script/builder/comparison.py | 82 ++
python/tvm/tirx/script/builder/v2.py | 1805 ++++++++++++++++++++++++
src/relax/script/builder/frame.cc | 2 +-
src/script/ir_builder/base.cc | 6 +
src/tirx/script/builder/frame.cc | 2 +-
16 files changed, 7723 insertions(+), 6 deletions(-)
diff --git a/include/tvm/script/ir_builder/base.h b/include/tvm/script/ir_builder/base.h
index 417e9b0c38..bbe460dee5 100644
--- a/include/tvm/script/ir_builder/base.h
+++ b/include/tvm/script/ir_builder/base.h
@@ -129,6 +129,29 @@ class IRBuilderFrame : public ffi::ObjectRef {
}
};
+/*! \brief Per-function symbols, retained across declaration and definition.
+ * The native builder stack owns active scope; the FunctionRecord owns this
+ * frame between entries. symbols maps source names to canonical primitive Vars.
+ * It starts empty, survives re-entry, and is never shared between functions.
+ */
+class TypeVarFrameNode : public IRBuilderFrameNode {
+ public:
+ ffi::Map<ffi::String, tvm::Var> symbols;
+
+ static void RegisterReflection() {
+ ffi::reflection::ObjectDef<TypeVarFrameNode>().def_ro("symbols", &TypeVarFrameNode::symbols);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.TypeVarFrame", TypeVarFrameNode,
+ IRBuilderFrameNode);
+};
+
+/*! \brief Managed reference to the shared function symbol frame. */
+class TypeVarFrame : public IRBuilderFrame {
+ public:
+ TypeVarFrame() : IRBuilderFrame(ffi::make_object<TypeVarFrameNode>()) {}
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypeVarFrame, IRBuilderFrame, TypeVarFrameNode);
+};
+
////////////////////////////// IRBuilder //////////////////////////////
/*!
@@ -183,7 +206,7 @@ class IRBuilderNode : public ffi::Object {
template <typename TFrame>
inline ffi::Optional<TFrame> FindFrame() const;
/*!
- * \brief Get the frame on top of the stack `this->frames` if its type is `TFrame`.
+ * \brief Get the top semantic frame if its type is `TFrame`; skip symbol-only frames.
* \tparam TFrame The assumed type of the last frame on stack.
* \return The frame if the stack is non-empty and the top of the stack is of type `TFrame`.
* Otherwise std::nullopt.
@@ -197,6 +220,13 @@ class IRBuilderNode : public ffi::Object {
*/
template <typename TObjectRef>
inline TObjectRef Get() const;
+ /*! \brief Whether any semantic construction frame remains (symbol frames are transparent). */
+ bool HasConstructionFrames() const {
+ for (const auto& frame : frames) {
+ if (!frame->IsInstance<TypeVarFrameNode>()) return true;
+ }
+ return false;
+ }
/*! \brief Push a frontend source span for IR constructed in the nested scope. */
void PushSourceSpan(Span span);
/*! \brief Pop the innermost frontend source span. */
@@ -296,8 +326,10 @@ inline ffi::Optional<TFrame> IRBuilderNode::FindFrame() const {
template <typename TFrame>
inline ffi::Optional<TFrame> IRBuilderNode::GetLastFrame() const {
using TFrameNode = typename TFrame::ContainerType;
- if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
- return frames.back().as_or_throw<TFrame>();
+ for (auto it = frames.rbegin(); it != frames.rend(); ++it) {
+ if ((*it)->IsInstance<TFrameNode>()) return (*it).template as_or_throw<TFrame>();
+ // Symbol state is administrative, not a lexical IR construction scope.
+ if (!(*it)->IsInstance<TypeVarFrameNode>()) break;
}
return std::nullopt;
}
diff --git a/python/tvm/relax/script/builder/v2.py b/python/tvm/relax/script/builder/v2.py
new file mode 100644
index 0000000000..233bd64645
--- /dev/null
+++ b/python/tvm/relax/script/builder/v2.py
@@ -0,0 +1,1615 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Concrete Relax construction operations over the shared native builder stack."""
+
+# pylint: disable=wildcard-import,redefined-builtin,invalid-name
+import builtins as _python
+import numbers as _numbers
+
+import tvm_ffi as _ffi
+
+from tvm import ir as _ir
+from tvm import relax as _relax
+from tvm import tirx as _tir
+from tvm.relax.distributed import DeviceMesh as _DeviceMesh
+from tvm.relax.distributed import DTensorType as _DTensorType
+from tvm.relax.distributed import Placement as _Placement
+from tvm.relax.distributed import device_mesh as device_mesh
+from tvm.script.ir_builder import IRBuilder as _IRBuilder
+from tvm.script.ir_builder import ir as _I
+from tvm.script.ir_builder.base import MISSING as _MISSING
+from tvm.script.ir_builder.base import _construction_span
+from tvm.script.ir_builder.base import _frame_result as _named_frame_result
+from tvm.script.ir_builder.base import at as _at
+from tvm.script.ir_builder.base import source_span as _source_span
+from tvm.script.ir_builder.type_var_frame import TypeVarDecl as _TypeVarDecl
+from tvm.script.ir_builder.type_var_frame import TypeVarFrame as _TypeVarFrame
+from tvm.script.ir_builder.type_var_frame import resolve_type_var
+from tvm.script.parser_v2.protocol import expr_str_args as _expr_str_args
+
+from . import _ffi_api
+from . import distributed as dist
+from . import frame as _frame
+from . import ir as _native
+from .distributed.ir import _lookup_device_mesh
+from .ir import * # noqa: F403
+
+
+@_expr_str_args("shape", introduce=True, dtype="int64", scalar_strings=False)
+def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
+ """Construct a tensor type from concrete shape dimensions.
+
+ Parameters
+ ----------
+ shape : sequence, str, or None
+ Concrete dimensions, None for unknown shape, or a bare dtype string when dtype is
+ omitted.
+ dtype : str or None
+ Element dtype; None means unknown.
+ vdevice : ir.VDevice, str, or None
+ Concrete device or target[:index] lookup in the active module.
+ ndim : int
+ Rank when shape is unknown; -1 means unknown.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ relax.TensorType or ir.Type
+ Concrete tensor type, or MissingType for unresolved eager annotation strings.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate. Shape dimensions follow the native Relax int64 requirement.
+ Device-name lookup needs an active module. Marked unresolved strings/TypeVars outside a
+ builder produce MissingType; active builders require concrete symbols.
+
+ Examples
+ --------
+ >>> Tensor((2, 3), "float32")
+ R.Tensor((2, 3), dtype="float32")
+ """
+ if isinstance(shape, _python.str) and dtype is None:
+ dtype, shape = shape, None
+ if isinstance(vdevice, _python.str):
+ target, *qualifiers = vdevice.split(":", 2)
+ vdevice = _I.lookup_vdevice(target, int(qualifiers[0]) if qualifiers else 0)
+ return _relax.TensorType(shape, dtype, vdevice, ndim, _source_span(span))
+
+
+@_expr_str_args("shape", introduce=True, dtype="int64", scalar_strings=False)
+def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1, span=None):
+ """Construct a distributed tensor type from concrete dimensions.
+
+ Parameters
+ ----------
+ shape : sequence or None
+ Concrete dimensions or unknown shape.
+ dtype : str or None
+ Element dtype or unknown dtype.
+ device_mesh : DeviceMesh, str, or None
+ Concrete mesh, module mesh name, or an empty default mesh.
+ placement : Placement or str
+ Placement object or text; empty text is the default.
+ ndim : int
+ Rank for unknown shape; -1 means unknown.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ DTensorType or ir.Type
+ Distributed tensor type, or MissingType for unresolved eager strings.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate. Named device meshes require module context; concrete shape
+ dimensions follow Tensor constraints.
+
+ Examples
+ --------
+ >>> ty = DTensor((2, 3), "float32", mesh, "S[0]")
+ """
+ if device_mesh is None:
+ device_mesh = _DeviceMesh([], _ir.Range(0, 1))
+ elif isinstance(device_mesh, _python.str):
+ device_mesh = _lookup_device_mesh(device_mesh)
+ if isinstance(placement, _python.str):
+ placement = _Placement.from_text(placement)
+ return _DTensorType(Tensor(shape, dtype, ndim=ndim), device_mesh, placement, _source_span(span))
+
+
+Range = _ir.Range
+
+
+@_expr_str_args("values", introduce=True, dtype="int64")
+def Shape(values=None, ndim=-1, *, span=None):
+ """Construct a shape type from concrete dimensions.
+
+ Parameters
+ ----------
+ values : sequence or None
+ Concrete int64 dimension expressions, or None for unknown values.
+ ndim : int
+ Rank when values are unknown; -1 means unknown.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ relax.ShapeType or ir.Type
+ Concrete shape type, or MissingType for unresolved eager strings.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate. Active handwritten builders require concrete dimension
+ expressions.
+
+ Examples
+ --------
+ >>> ty = Shape((2, 3))
+ """
+ return _relax.ShapeType(values, ndim, _source_span(span))
+
+
+def _type(value):
+ if value is None:
+ return _ir.TupleType([])
+ if callable(value):
+ value = value()
+ if _ir.is_prim_expr(value) or isinstance(value, _TypeVarDecl):
+ value = value.ty
+ if not isinstance(value, _ir.Type):
+ raise TypeError(f"Expected a concrete type, got {type(value).__name__}")
+ return value
+
+
+def Callable(params=None, ret=None, purity=None, derive_func=None, *, span=None):
+ """Construct a concrete or opaque Relax function type.
+
+ Parameters
+ ----------
+ params : type, sequence, or None
+ Parameter annotations. None selects an opaque callable.
+ ret : object or None
+ Return annotation; None denotes an empty tuple for explicit signatures.
+ purity : bool or None
+ Purity flag; None defaults true for explicit parameters, false for opaque signatures.
+ derive_func : object or None
+ Opaque-call return-type derivation function.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ relax.FuncType
+ Function type with normalized concrete parameter and return types.
+
+ Raises
+ ------
+ ValueError
+ derive_func is supplied for a nonopaque callable.
+ TypeError
+ An annotation cannot be converted to a concrete type.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate.
+
+ Examples
+ --------
+ >>> ty = Callable([Tensor((2,), "float32")], Tensor((2,), "float32"))
+ """
+ if purity is None:
+ purity = params is not None
+ if params is None:
+ return _relax.FuncType.opaque_func(
+ ret=None if ret is None else _type(ret),
+ derive_func=derive_func,
+ purity=purity,
+ span=_source_span(span),
+ )
+ if derive_func is not None:
+ raise ValueError("A derivation function requires an opaque callable")
+ if not isinstance(params, list | _python.tuple):
+ params = [params]
+ return _relax.FuncType(
+ [_type(param) for param in params], _type(ret), purity, _source_span(span)
+ )
+
+
+def Tuple(*fields, span=None):
+ """Construct a tuple type while preserving missing component types.
+
+ Parameters
+ ----------
+ fields : object
+ Positional annotations or one list/tuple of annotations. Empty input creates an empty
+ tuple type.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ ir.TupleType
+ Tuple of normalized concrete or missing field types.
+
+ Raises
+ ------
+ TypeError
+ A field cannot be converted to an IR type.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate. Primitive expressions, declaration descriptors, and
+ zero-argument annotation constructors are accepted. MissingType fields do not collapse the
+ tuple.
+
+ Examples
+ --------
+ >>> ty = Tuple(Prim("int32"), Object())
+ """
+ if len(fields) == 1 and isinstance(fields[0], list | _python.tuple):
+ fields = fields[0]
+ return _ir.TupleType([_type(field) for field in fields], _source_span(span))
+
+
+def Prim(dtype, *, span=None):
+ """Construct a primitive scalar type.
+
+ Parameters
+ ----------
+ dtype : str or DataType
+ Concrete scalar element dtype.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ ir.PrimType
+ Primitive type.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate. span is accepted for the common builder interface; PrimType
+ stores no source span.
+
+ Examples
+ --------
+ >>> Prim("int32")
+ int32
+ """
+ return _ir.PrimType(dtype)
+
+
+Prim.__tvm_parameter_dtype__ = "dtype"
+
+
+def Object(*, span=None):
+ """Construct the unconstrained Relax value type.
+
+ Parameters
+ ----------
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ relax.AnyType
+ Unconstrained value type with the requested source span.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate.
+
+ Examples
+ --------
+ >>> ty = Object()
+ """
+ return _relax.AnyType(_source_span(span))
+
+
+Any = Object
+
+
+is_type_var = _ir.is_prim_var
+
+
+def type_var(name, *, dtype=None, span=None):
+ """Construct a fresh standalone primitive symbol.
+
+ Parameters
+ ----------
+ name : str
+ Source name of the new variable.
+ dtype : str or None
+ Primitive dtype; None retains the explicit-constructor default int64.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ ir.Var
+ Fresh variable with the requested name and type.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Creates no construction frame and retains no per-function state. Native constructor
+ validation errors propagate. This does not register a canonical symbol. resolve_type_var
+ instead uses the active function frame and defaults newly introduced symbols to
+ int64.
+
+ Examples
+ --------
+ >>> n = type_var("n", dtype="int64")
+ """
+ return _ir.Var(name, "int64" if dtype is None else dtype, _source_span(span))
+
+
+class _Frame:
+ """Retain source metadata and exports around an existing native frame."""
+
+ def __init__(self, native, span=None):
+ # Each wrapper owns one native construction frame and source location.
+ # result starts empty, records finalized lexical exports on exit, and
+ # never outlives its construction region or stores another function's
+ # symbols (those belong to the separate TypeVarFrame).
+ self.native = native
+ self.span = span
+ self.result = {}
+
+ def __getattr__(self, name):
+ return getattr(self.native, name)
+
+ @property
+ def reference(self):
+ """Return the stable module or local function reference after declaration."""
+ if isinstance(self.native, _frame.FunctionFrame):
+ local_var = self.native.local_var
+ return local_var if local_var is not None else self.native.global_var
+ raise AttributeError("This frame does not declare a function")
+
+ def __enter__(self):
+ with _construction_span(self.span):
+ self.native.__enter__()
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ with _construction_span(self.span):
+ self.native.__exit__(exc_type, exc_value, traceback)
+ if exc_type is None:
+ if isinstance(self.native, _frame.BindingBlockFrame):
+ self.result = {var.name: var for var in self.native.output_vars}
+ elif (
+ isinstance(self.native, _frame.FunctionFrame) and self.native.local_var is not None
+ ):
+ self.result = {self.native.name: self.native.local_var}
+ elif isinstance(self.native, _frame.IfFrame):
+ self.result = {self.native.var_name: self.native.var}
+ return False
+
+
+def frame_result(completed_frame, name):
+ """Read one named export from an explicitly supplied completed region.
+
+ Parameters
+ ----------
+ completed_frame : builder frame or dict
+ Completed Relax context or explicit host branch export dictionary.
+ name : str
+ Original source binding name.
+
+ Returns
+ -------
+ object
+ Exported value, including None, or ir_builder.base.MISSING when absent.
+
+ Raises
+ ------
+ TypeError
+ name is not a string.
+
+ Notes
+ -----
+ A conditional exports its native designated result; dataflow exports declared outputs; local
+ functions export their reference. Other region-local names stay hidden. This is read-only
+ and consults no ambient frame state; nested results remain independent.
+
+ Examples
+ --------
+ >>> frame_result({"y": 3}, "y")
+ 3
+ """
+ return _named_frame_result(completed_frame, name)
+
+
+def function(is_pure=True, is_private=False, *, local=False, reference=None, span=None):
+ """Create a native Relax function definition context.
+
+ Parameters
+ ----------
+ is_pure : bool
+ Whether the function is pure; defaults True.
+ is_private : bool
+ Whether a module function is private; defaults False.
+ local : bool
+ Whether to define a nested function; defaults False.
+ reference : ir.Var or None
+ Previously declared local function reference; required when local is True.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ context manager
+ Definition context retaining its native function and exports.
+
+ Raises
+ ------
+ ValueError
+ A local definition has no declared reference.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ The returned context wraps an existing native frame. Enter it under an active IRBuilder;
+ exit finalizes its result and releases the native scope. No ambient completed-frame cache is
+ used.
+
+ Examples
+ --------
+ >>> with function():
+ ... return_(None)
+ """
+ if local:
+ if reference is None:
+ raise ValueError("A local function requires its declared reference")
+ return _Frame(_ffi_api.LocalFunction(is_pure, reference), span)
+ return _Frame(_native.function(is_pure, is_private), span)
+
+
+def decl_function(is_pure=True, is_private=False, *, local=False, span=None):
+ """Create a bodyless Relax function declaration context.
+
+ Parameters
+ ----------
+ is_pure : bool
+ Function purity; defaults True.
+ is_private : bool
+ Module-function privacy; defaults False.
+ local : bool
+ Whether to declare a nested function; defaults False.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ context manager
+ Declaration context exposing its reference after successful exit.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ The returned context wraps an existing native frame. Enter it under an active IRBuilder;
+ exit finalizes its result and releases the native scope. No ambient completed-frame cache is
+ used. Signature operations match definition operations; the function body is supplied later.
+
+ Examples
+ --------
+ >>> with decl_function() as declared:
+ ... x = arg("x", Object())
+ ... func_ret_type(Object())
+ """
+ return _Frame(_ffi_api.DeclFunction(is_pure, is_private, local), span)
+
+
+def arg(name, ty, *, span=None):
+ """Add a parameter while retaining a supplied variable identity.
+
+ Parameters
+ ----------
+ name : str
+ Original parameter name.
+ ty : ir.Var or object
+ Cached parameter variable or concrete annotation accepted by the type adapter.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ ir.Var
+ Parameter registered on the active function frame.
+
+ Raises
+ ------
+ TypeError
+ The annotation cannot be converted to a type.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Mutates the active declaration/definition signature. A supplied variable is reused, not
+ cloned; source instrumentation is temporary.
+
+ Examples
+ --------
+ >>> x = arg("x", Tensor((2,), "float32"))
+ """
+ with _construction_span(span):
+ if isinstance(ty, _ir.Var):
+ return _ffi_api.ArgVar(name, ty)
+ return _at(span, _native.arg(name, _type(ty)))
+
+
+def func_ret_type(ret_ty):
+ """Set the active function signature return type.
+
+ Parameters
+ ----------
+ ret_ty : object
+ Concrete annotation, primitive expression, or zero-argument constructor.
+
+ Returns
+ -------
+ None
+ The active frame retains the normalized return type.
+
+ Raises
+ ------
+ TypeError
+ The annotation cannot be converted to a type.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Requires an active declaration or definition. No additional frame is entered.
+
+ Examples
+ --------
+ >>> func_ret_type(Object())
+ """
+ return _native.func_ret_type(_type(ret_ty))
+
+
+func_ret_ty = func_ret_type
+
+
+def dataflow(*, span=None):
+ """Create a dataflow context with explicit finalized exports.
+
+ Parameters
+ ----------
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ context manager
+ Region whose outputs are accessible through frame_result after exit.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ The returned context wraps an existing native frame. Enter it under an active IRBuilder;
+ exit finalizes its result and releases the native scope. No ambient completed-frame cache is
+ used. Native output operations determine which variables escape.
+
+ Examples
+ --------
+ >>> with dataflow() as region:
+ ... y = bind_(value, name="y")
+ ... output(y)
+ """
+ return _Frame(_native.dataflow(), span)
+
+
+def If(condition, *, span=None):
+ """Create a conditional region with a designated named result.
+
+ Parameters
+ ----------
+ condition : relax.Expr
+ Scalar boolean condition accepted by the native Relax If builder.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ context manager
+ Conditional context exposing its designated result after exit.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ The returned context wraps an existing native frame. Enter it under an active IRBuilder;
+ exit finalizes its result and releases the native scope. No ambient completed-frame cache is
+ used. Then and Else construct branch bodies; both must provide compatible designated
+ outputs.
+
+ Examples
+ --------
+ >>> with If(condition) as region:
+ ... with Then():
+ ... bind_(left, name="y")
+ ... with Else():
+ ... bind_(right, name="y")
+ """
+ return _Frame(_native.If(condition), span)
+
+
+def Then(*, span=None):
+ """Create the true branch of the active conditional.
+
+ Parameters
+ ----------
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ context manager
+ Native true-branch construction context.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ The returned context wraps an existing native frame. Enter it under an active IRBuilder;
+ exit finalizes its result and releases the native scope. No ambient completed-frame cache is
+ used. Requires an enclosing If frame.
+
+ Examples
+ --------
+ >>> with Then():
+ ... bind_(value, name="y")
+ """
+ return _Frame(_native.Then(), span)
+
+
+def Else(*, span=None):
+ """Create the false branch of the active conditional.
+
+ Parameters
+ ----------
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ context manager
+ Native false-branch construction context.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ The returned context wraps an existing native frame. Enter it under an active IRBuilder;
+ exit finalizes its result and releases the native scope. No ambient completed-frame cache is
+ used. Requires an enclosing If frame.
+
+ Examples
+ --------
+ >>> with Else():
+ ... bind_(value, name="y")
+ """
+ return _Frame(_native.Else(), span)
+
+
+def _check_unterminated():
+ for active_frame in reversed(_IRBuilder.current().frames):
+ if isinstance(active_frame, _frame.FunctionFrame):
+ if active_frame.output is not None:
+ raise ValueError("A Relax operation cannot follow an unconditional return")
+ break
+
+
+def _value(value, ty=None):
+ if isinstance(value, _python.tuple):
+ return _relax.utils.convert_to_expr(value)
+ if isinstance(value, _numbers.Number):
+ if isinstance(ty, _ir.PrimType):
+ return _relax.prim_value(value, dtype=ty.dtype)
+ return _relax.const(value)
+ return value
+
+
+def bind_(
+ value=_MISSING,
+ *,
+ ty=None,
+ name=None,
+ span=None,
+ name_span=None,
+ previous=_MISSING,
+ declaration=False,
+ frame_value=False,
+):
+ """Emit a Relax binding or retain a named frame-owned value.
+
+ Parameters
+ ----------
+ value : object
+ Expression, Python value, symbolic declaration descriptor, or MISSING.
+ ty : object or None
+ Optional annotation or zero-argument annotation constructor.
+ name : str or None
+ Optional source binding name.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+ name_span : ir.Span, tuple, or None
+ Binding-name range; None uses span.
+ previous : object
+ Previously resolved symbol for compatible declarations; defaults MISSING.
+ declaration : bool
+ Validate/reuse a primitive symbol declaration; defaults False.
+ frame_value : bool
+ Name a frame-owned value without emitting a binding; defaults False.
+
+ Returns
+ -------
+ object
+ Emitted variable, canonical symbol, or unchanged Python value.
+
+ Raises
+ ------
+ ValueError
+ Initializer is missing or an operation follows an unconditional return.
+ TypeError
+ A declaration or match-cast annotation is incompatible.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Expression inputs emit in the active native function/block. TypeVarDecl updates the
+ nearest symbol frame. Existing native Emit rules fill MissingType on the original RHS
+ before normalization and check concrete annotations. This operation adds no recursive
+ annotation validation or propagation to tuple fields or call arguments.
+
+ Examples
+ --------
+ >>> y = bind_(value, ty=Object(), name="y")
+ """
+ _check_unterminated()
+ name_span = _source_span(span if name_span is None else name_span)
+ if isinstance(value, _TypeVarDecl):
+ return _TypeVarFrame.current().resolve(name, value.ty, span=name_span)
+ if frame_value:
+ if isinstance(value, _python.list | _python.tuple | _ir.Array):
+ for index, item in enumerate(value):
+ bind_(
+ item,
+ name=None if name is None else f"{name}_{index}",
+ span=_source_span(span),
+ name_span=name_span,
+ frame_value=True,
+ )
+ elif isinstance(value, _ir.Var):
+ if name is not None:
+ _IRBuilder.name(name, value)
+ _at(name_span if name_span is not None else span, value)
+ return value
+ if declaration:
+ if not _ir.is_prim_var(value):
+ raise TypeError("A symbol declaration requires a concrete primitive variable")
+ if ty is not None and not _ffi.structural_equal(_type(ty), value.ty):
+ raise TypeError("The symbol declaration has an incompatible type")
+ if previous is not _MISSING:
+ if not _ir.is_prim_var(previous) or not _ffi.structural_equal(previous.ty, value.ty):
+ raise TypeError("The symbol declaration has an incompatible signature dtype")
+ return previous
+ if name is not None:
+ _IRBuilder.name(name, value)
+ return _at(name_span if name_span is not None else span, value)
+ if value is _MISSING:
+ raise ValueError("Relax bindings require an initializer")
+ if isinstance(value, _I.meta_var):
+ return value.value
+ ty = None if ty is None else _type(ty)
+ value = _value(value, ty)
+ with _construction_span(span):
+ if isinstance(value, _relax.MatchCast):
+ if ty is not None and not _ffi.structural_equal(ty, value.ty):
+ raise TypeError("The binding annotation differs from the match-cast type")
+ result = _ffi_api.EmitMatchCastV2(value.value, value.ty, name_span)
+ elif isinstance(value, _relax.Expr):
+ result = _ffi_api.EmitV2(value, ty, name_span)
+ else:
+ return value
+ if name is not None:
+ _IRBuilder.name(name, result)
+ return _at(name_span if name_span is not None else span, result)
+
+
+def emit_(value, *, span=None):
+ """Emit a void expression statement.
+
+ Parameters
+ ----------
+ value : relax.Expr or None
+ Void expression, or None from an already emitted effect-only operation.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ None
+ Emits the expression when needed.
+
+ Raises
+ ------
+ TypeError
+ The value is neither a Relax expression nor None.
+ ValueError
+ The emitted expression has a nonvoid type.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ None is a no-op. Other inputs create an underscore binding in the active function/block and
+ must have empty tuple type.
+
+ Examples
+ --------
+ >>> emit_(None)
+ """
+ if value is None:
+ return
+ if not isinstance(value, _relax.Expr):
+ raise TypeError(f"Unsupported expression statement value: {type(value).__name__}")
+ result = bind_(value, name="_", span=span)
+ if not isinstance(result.ty, _ir.TupleType) or len(result.ty.fields) != 0:
+ raise ValueError(
+ "Non-void expressions must be bound to a variable; "
+ f"expression of type {result.ty} was used as a statement"
+ )
+
+
+def return_(value=None, *, span=None):
+ """Record a function result without exiting Python construction.
+
+ Parameters
+ ----------
+ value : object or None
+ Concrete result or convertible Python value; None denotes an empty tuple.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ None
+ Updates the active native function result.
+
+ Raises
+ ------
+ ValueError
+ An unconditional return has already been recorded.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Subsequent ordinary Relax binding operations are rejected. This operation enters no new
+ construction frame.
+
+ Examples
+ --------
+ >>> return_(value)
+ """
+ _check_unterminated()
+ with _construction_span(span):
+ if value is None:
+ value = _relax.Tuple([])
+ _native.func_ret_value(_value(value))
+
+
+def match_cast(value, ty, *, span=None):
+ """Construct a match-cast descriptor for bind_ to consume.
+
+ Parameters
+ ----------
+ value : object
+ Concrete or convertible Relax value; None is rejected.
+ ty : object
+ Concrete target annotation.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ relax.MatchCast
+ Unemitted match-cast with an anonymous result variable.
+
+ Raises
+ ------
+ ValueError
+ value is None.
+ TypeError
+ The annotation cannot be converted to a type.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Constructs the descriptor without emitting a binding or entering a frame. bind_ registers it
+ in the active block.
+
+ Examples
+ --------
+ >>> y = bind_(match_cast(value, Tensor((2,), "float32")), name="y")
+ """
+ if value is None:
+ raise ValueError("The match-cast value cannot be None")
+ ty = _type(ty)
+ return _relax.MatchCast(_ir.Var("", ty), _value(value), ty, _source_span(span))
+
+
+def unpack(value):
+ """Project an IR tuple with known arity or preserve host iteration.
+
+ Parameters
+ ----------
+ value : object
+ Tuple literal, tuple-typed expression, or arbitrary host iterable.
+
+ Returns
+ -------
+ object
+ Python tuple of fields/projections, or the original value.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Creates projections only; no binding or construction frame is entered.
+
+ Examples
+ --------
+ >>> unpack((1, 2))
+ (1, 2)
+ """
+ if isinstance(value, _relax.Tuple):
+ return _python.tuple(value.fields)
+ if isinstance(value, _relax.Expr) and isinstance(value.ty, _ir.TupleType):
+ return _python.tuple(_relax.TupleGetItem(value, i) for i in range(len(value.ty.fields)))
+ return value
+
+
+def assert_(condition, message="", *, span=None):
+ """Emit a runtime assertion with construction-time diagnostic text.
+
+ Parameters
+ ----------
+ condition : relax.Expr
+ Runtime boolean condition.
+ message : str
+ Construction-time message; defaults to empty text.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ None
+ Emits a void assertion operation in the active block.
+
+ Raises
+ ------
+ TypeError
+ message is not a string.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Uses existing native assertion and void-binding rules; enters no new construction frame.
+
+ Examples
+ --------
+ >>> assert_(condition, "shape mismatch")
+ """
+ if not isinstance(message, _python.str):
+ raise TypeError("An assertion message must be construction-time text")
+ with _construction_span(span):
+ emit_(_at(span, _native.assert_op(condition, format=message)), span=span)
+
+
+def For(*args, span=None, **kwargs):
+ """Reject imperative for loops in the Relax expression dialect.
+
+ Parameters
+ ----------
+ args : object
+ Rejected iteration arguments.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+ kwargs : object
+ Rejected iteration options.
+
+ Returns
+ -------
+ None
+ Never returns normally.
+
+ Raises
+ ------
+ TypeError
+ Relax does not support imperative for loops.
+
+ Notes
+ -----
+ This rejection changes no builder state and enters no frame.
+
+ Examples
+ --------
+ >>> # For(...) raises TypeError in Relax.
+ """
+ raise TypeError("Relax does not support imperative for loops")
+
+
+def break_(*, span=None):
+ """Reject break in the Relax expression dialect.
+
+ Parameters
+ ----------
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ None
+ Never returns normally.
+
+ Raises
+ ------
+ TypeError
+ Relax does not support break.
+
+ Notes
+ -----
+ This rejection changes no builder state and enters no frame.
+
+ Examples
+ --------
+ >>> # break_(...) raises TypeError in Relax.
+ """
+ raise TypeError("Relax does not support break")
+
+
+def continue_(*, span=None):
+ """Reject continue in the Relax expression dialect.
+
+ Parameters
+ ----------
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ None
+ Never returns normally.
+
+ Raises
+ ------
+ TypeError
+ Relax does not support continue.
+
+ Notes
+ -----
+ This rejection changes no builder state and enters no frame.
+
+ Examples
+ --------
+ >>> # continue_(...) raises TypeError in Relax.
+ """
+ raise TypeError("Relax does not support continue")
+
+
+def setitem(target, index, value, *, span=None):
+ """Reject indexed assignment in the Relax expression dialect.
+
+ Parameters
+ ----------
+ target : object
+ Rejected assignment target.
+ index : object
+ Rejected subscript.
+ value : object
+ Rejected assigned value.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ None
+ Never returns normally.
+
+ Raises
+ ------
+ TypeError
+ Relax does not support indexed assignment.
+
+ Notes
+ -----
+ This rejection changes no builder state and enters no frame.
+
+ Examples
+ --------
+ >>> # setitem(...) raises TypeError in Relax.
+ """
+ raise TypeError("Relax does not support indexed assignment")
+
+
+__all__ = [
+ *_native.__all__,
+ "Any",
+ "Callable",
+ "DTensor",
+ "For",
+ "for_",
+ "frame_result",
+ "resolve_type_var",
+ "Object",
+ "Prim",
+ "Range",
+ "Shape",
+ "Tensor",
+ "Tuple",
+ "assert_",
+ "break_",
+ "continue_",
+ "bind_",
+ "decl_function",
+ "device_mesh",
+ "dist",
+ "emit_",
+ "is_type_var",
+ "match_cast",
+ "return_",
+ "setitem",
+ "type_var",
+ "unpack",
+]
+
+
+def _logical_pair(lhs, rhs, operation, primitive, python_operation):
+ if not isinstance(lhs, _ir.Expr) and not isinstance(rhs, _ir.Expr):
+ return python_operation(lhs, rhs)
+ if _ir.is_prim_expr(lhs) or _ir.is_prim_expr(rhs):
+ return primitive(lhs, rhs)
+ return operation(_value(lhs), _value(rhs))
+
+
+def logical_and(*values):
+ """Construct conjunction of concrete host, primitive, or tensor values.
+
+ Parameters
+ ----------
+ values : object
+ One or more host values, primitive boolean expressions, or boolean tensors.
+
+ Returns
+ -------
+ object
+ Host value, primitive And expression, or Relax logical_and call.
+
+ Raises
+ ------
+ TypeError
+ No operands are supplied.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ This combines eagerly constructed operands elementwise for tensors. It enters no frame and
+ emits no binding. Invalid IR conditions may fail later during normalization.
+
+ Examples
+ --------
+ >>> logical_and(True, False)
+ """
+ if not values:
+ raise TypeError("logical_and requires at least one operand")
+ result = values[0]
+ for value in values[1:]:
+ result = _logical_pair(result, value, _relax.op.logical_and, _tir.And, lambda a, b: a and b)
+ return result
+
+
+def logical_or(*values):
+ """Construct disjunction of concrete host, primitive, or tensor values.
+
+ Parameters
+ ----------
+ values : object
+ One or more host values, primitive boolean expressions, or boolean tensors.
+
+ Returns
+ -------
+ object
+ Host value, primitive Or expression, or Relax logical_or call.
+
+ Raises
+ ------
+ TypeError
+ No operands are supplied.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ This combines eagerly constructed operands elementwise for tensors. It enters no frame and
+ emits no binding. Invalid IR conditions may fail later during normalization.
+
+ Examples
+ --------
+ >>> logical_or(False, True)
+ """
+ if not values:
+ raise TypeError("logical_or requires at least one operand")
+ result = values[0]
+ for value in values[1:]:
+ result = _logical_pair(result, value, _relax.op.logical_or, _tir.Or, lambda a, b: a or b)
+ return result
+
+
+def logical_not(value):
+ """Negate a host or IR boolean without coercing IR to Python bool.
+
+ Parameters
+ ----------
+ value : object
+ Host value, primitive boolean expression, or boolean tensor.
+
+ Returns
+ -------
+ bool or ir.Expr
+ Host bool, primitive Not, or Relax logical_not call.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Enters no frame, emits no binding, and retains no state. Host truth-conversion errors
+ propagate; invalid IR types may fail during normalization.
+
+ Examples
+ --------
+ >>> logical_not(True)
+ False
+ """
+ if _ir.is_prim_expr(value):
+ return _tir.Not(value)
+ if isinstance(value, _ir.Expr):
+ return _relax.op.logical_not(value)
+ return not value
+
+
+def select(condition, true_value, false_value):
+ """Construct elementwise selection or select already-built host values.
+
+ Parameters
+ ----------
+ condition : object
+ Host truth value, primitive boolean, or boolean tensor.
+ true_value : object
+ Already constructed value for true entries.
+ false_value : object
+ Already constructed value for false entries.
+
+ Returns
+ -------
+ object
+ Selected host value, primitive Select, or Relax where call.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ This elementwise operation makes no runtime short-circuit promise. It enters no frame and
+ emits no binding. Native type errors may surface during normalization.
+
+ Examples
+ --------
+ >>> select(True, 1, 2)
+ 1
+ """
+ if _ir.is_prim_expr(condition):
+ return _tir.Select(condition, true_value, false_value)
+ if isinstance(condition, _ir.Expr):
+ return _relax.op.where(condition, _value(true_value), _value(false_value))
+ return true_value if condition else false_value
+
+
+__all__ += ["logical_and", "logical_not", "logical_or", "select"]
+
+
+for_ = For
+
+
+def if_then_else_(condition, true_value, false_value):
+ """Construct scalar conditional evaluation from eagerly constructed operands.
+
+ Parameters
+ ----------
+ condition : object
+ Host truth value, primitive boolean, or scalar boolean Relax expression.
+ true_value : object
+ Already constructed true arm.
+ false_value : object
+ Already constructed false arm.
+
+ Returns
+ -------
+ object
+ Selected host value, primitive conditional intrinsic, or Relax If.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Python constructs both arms before calling this operation. The resulting IR selects its arm
+ at runtime. Relax normalization keeps branch calls inside their branch. Arm types must be
+ compatible; invalid conditions/types can fail during normalization. No outer bindings, new
+ frames, or persistent state are created.
+
+ Examples
+ --------
+ >>> if_then_else_(True, 1, 2)
+ 1
+ """
+ if isinstance(condition, _ffi.ObjectConvertible):
+ condition = condition.asobject()
+ if not isinstance(condition, _ir.Expr):
+ return true_value if condition else false_value
+ true_value = (
+ true_value.asobject() if isinstance(true_value, _ffi.ObjectConvertible) else true_value
+ )
+ false_value = (
+ false_value.asobject() if isinstance(false_value, _ffi.ObjectConvertible) else false_value
+ )
+ if _ir.is_prim_expr(condition) and all(
+ _ir.is_prim_expr(value)
+ if isinstance(value, _ir.Expr)
+ else isinstance(value, _numbers.Number)
+ for value in (true_value, false_value)
+ ):
+ return _tir.if_then_else(condition, true_value, false_value)
+ return _relax.If(condition, _value(true_value), _value(false_value))
+
+
+def _chain_binding(variable, value, body):
+ if _ir.is_prim_expr(value) and _ir.is_prim_expr(body):
+ return _tir.Let(variable, value, body)
+ return _relax.SeqExpr([_relax.BindingBlock([_relax.VarBinding(variable, value)])], body)
+
+
+def and_(*values, chain=None):
+ """Construct scalar conjunction with runtime short-circuit evaluation.
+
+ Parameters
+ ----------
+ values : object
+ One or more host values, primitive booleans, or scalar boolean tensors.
+ chain : tuple or None
+ Optional original operands for a comparison chain. Must contain len(values)+1 entries;
+ values must be their ordered comparisons.
+
+ Returns
+ -------
+ object
+ Host value, primitive expression, or Relax If/SeqExpr.
+
+ Raises
+ ------
+ TypeError
+ No operands are supplied.
+ ValueError
+ The chain operand count is invalid.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Python eagerly constructs every argument. chain retains shared operand identity and
+ introduces progressively scoped IR bindings so each middle operand executes once, before its
+ first comparison. Neither later operands nor conditional branches are emitted in the outer
+ scope. No callbacks, persistent state, or new construction frames are used. Invalid scalar
+ tensor conditions may fail during normalization.
+
+ Examples
+ --------
+ >>> and_(True, False)
+ False
+ """
+ if not values:
+ raise TypeError("and_ requires at least one operand")
+ if chain is not None:
+ from tvm.tirx.script.builder.comparison import _comparison_chain
+
+ return _comparison_chain(values, chain, and_, _chain_binding)
+ result = values[-1]
+ for value in reversed(values[:-1]):
+ result = if_then_else_(
+ value, result, False if isinstance(value, _ir.Expr | _ffi.ObjectConvertible) else value
+ )
+ return result
+
+
+def or_(*values):
+ """Construct scalar disjunction with runtime short-circuit evaluation.
+
+ Parameters
+ ----------
+ values : object
+ One or more host values, primitive boolean expressions, or boolean tensors.
+
+ Returns
+ -------
+ object
+ Host value, primitive conditional expression, or Relax If.
+
+ Raises
+ ------
+ TypeError
+ No operands are supplied.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Python eagerly constructs all arguments. Compiled evaluation uses the right operand only
+ when the left is false. Tensor conditions must be scalar booleans. No outer bindings, new
+ frames, or persistent state are created. Invalid IR conditions may fail later during
+ normalization.
+
+ Examples
+ --------
+ >>> or_(False, True)
+ """
+ if not values:
+ raise TypeError("or_ requires at least one operand")
+ result = values[-1]
+ for value in reversed(values[:-1]):
+ result = if_then_else_(
+ value, True if isinstance(value, _ir.Expr | _ffi.ObjectConvertible) else value, result
+ )
+ return result
+
+
+def not_(value):
+ """Negate a host or IR boolean without coercing IR to Python bool.
+
+ Parameters
+ ----------
+ value : object
+ Host value, primitive boolean expression, or boolean tensor.
+
+ Returns
+ -------
+ bool or ir.Expr
+ Host bool, primitive Not, or Relax logical_not call.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Enters no frame, emits no binding, and retains no state. Host truth-conversion errors
+ propagate; invalid IR types may fail during normalization.
+
+ Examples
+ --------
+ >>> not_(True)
+ False
+ """
+ return logical_not(value)
+
+
+__all__ += ["and_", "if_then_else_", "not_", "or_"]
diff --git a/python/tvm/script/ir_builder/base.py b/python/tvm/script/ir_builder/base.py
index 56e78b9072..bc045e7d22 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -17,11 +17,12 @@
"""A generic IRBuilder across the TVM stack"""
from collections.abc import Callable
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
from typing import Any
from tvm_ffi import register_object as _register_object
+from tvm import ir
from tvm.runtime import Object as _Object
from . import _ffi_api
@@ -229,3 +230,142 @@ class IRBuilder(_Object):
"""
assert len(s) == len(vs)
return [IRBuilder.name(i, v) for i, v in zip(s, vs)]
+
+
+# V2 absence is distinct from legacy protocol state; no function data is retained.
+class _Missing:
+ def __repr__(self):
+ return "MISSING"
+
+
+MISSING = _Missing()
+
+
+def source_span(location):
+ """Materialize a source range without retaining source-unit state.
+
+ Parameters
+ ----------
+ location : ir.Span, tuple, or None
+ Optional source range accepted by source_span; None creates no source metadata.
+
+ Returns
+ -------
+ ir.Span or None
+ The existing span, a span with exact supplied coordinates, or None.
+
+ Raises
+ ------
+ TypeError or ValueError
+ The location cannot be unpacked or its fields are invalid.
+
+ Notes
+ -----
+ Generated tuples reuse one parser-created SourceName per source unit. Lines are one-based
+ and columns are zero-based UTF-8 byte offsets. Filename strings remain available to
+ handwritten callers. None is a no-op and accesses no builder; this operation enters no
+ frame.
+
+ Examples
+ --------
+ >>> source_span(None) is None
+ True
+ """
+ if location is None or isinstance(location, ir.Span):
+ return location
+ source_name, line, end_line, column, end_column = location
+ if isinstance(source_name, str):
+ source_name = ir.SourceName(source_name)
+ return ir.Span(source_name, line, end_line, column, end_column)
+
+
+@contextmanager
+def _construction_span(span):
+ """Apply a builder operation's span through the existing native span stack.
+
+ This private construction helper is used only by builders, never emitted by
+ the transpiler. ``span`` accepts the source_span input forms. None yields
+ directly without IR instrumentation. Otherwise it pushes/pops the existing
+ builder span stack when active, and retains no context after exit. Without
+ a builder it still materializes diagnostic locations.
+ Exceptions retain their original type and receive __tvm_script_location__
+ as a plain location tuple only when an inner operation has not already
+ supplied a more precise range. Diagnostics need never inspect an IR object.
+ Span-construction and nested-operation errors propagate unchanged.
+ """
+ if span is None:
+ yield
+ return
+ span = source_span(span)
+ context = (
+ IRBuilder.current().with_source_span(span)
+ if span is not None and IRBuilder.is_in_scope()
+ else nullcontext()
+ )
+ try:
+ with context:
+ yield
+ except Exception as error:
+ if span is not None and not hasattr(error, "__tvm_script_location__"):
+ diagnostic_span = span.spans[-1] if isinstance(span, ir.SequentialSpan) else span
+ error.__tvm_script_location__ = (
+ str(diagnostic_span.source_name.name),
+ diagnostic_span.line,
+ diagnostic_span.end_line,
+ diagnostic_span.column,
+ diagnostic_span.end_column,
+ )
+ raise
+
+
+def at(span, value):
+ """Attach a source range to a concrete expression when a builder is active.
+
+ Parameters
+ ----------
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+ value : object
+ Already constructed IR expression or arbitrary Python value.
+
+ Returns
+ -------
+ object
+ The same value, with missing expression source metadata filled when supported.
+
+ Raises
+ ------
+ TypeError or ValueError
+ A source range used for an IR expression is invalid.
+
+ Notes
+ -----
+ None returns the value without source instrumentation. Ordinary Python values and values
+ outside a builder pass through. An active builder temporarily uses its native source-span
+ stack; no construction frame or persistent cache is created. This is a value operation, not
+ a context manager or callback.
+
+ Examples
+ --------
+ >>> at(None, 3)
+ 3
+ """
+ if span is not None and isinstance(value, ir.Expr) and IRBuilder.is_in_scope():
+ with _construction_span(span):
+ return IRBuilder.current()._set_current_source_span(value)
+ return value
+
+
+def _frame_result(frame, name):
+ """Read one explicit export without consulting ambient construction state.
+
+ Dialects pass a completed frame whose result is its finalized export map,
+ or a Python-branch dictionary containing host lexical results. The map stays
+ private to its owner; a missing name returns MISSING instead of leaking a
+ branch-local binding. No frame is entered and no result is cached globally.
+ """
+ if not isinstance(name, str):
+ raise TypeError("A frame result name must be a string")
+ exports = frame if isinstance(frame, dict) else getattr(frame, "result", {})
+ return exports.get(name, MISSING)
diff --git a/python/tvm/script/ir_builder/construction.py b/python/tvm/script/ir_builder/construction.py
new file mode 100644
index 0000000000..a830f815bc
--- /dev/null
+++ b/python/tvm/script/ir_builder/construction.py
@@ -0,0 +1,647 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Runtime owners for generated builder programs.
+
+This module is deliberately on the builder side: concrete parameters, return
+annotations, module references, symbol identities and source spans never enter
+transpiler state. Generated Python calls these operations in source order.
+"""
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+from tvm import ir
+
+from . import ir as I
+from .base import MISSING, IRBuilder, source_span
+from .type_var_frame import TypeVarDecl, TypeVarFrame
+
+# Process-wide owner-supplied callbacks, initialized unset and replaced only by
+# explicit entry-module registration. Neither table contains per-parse state.
+_OPAQUE_FACTORY = None
+_MODULE_ADAPTER = None
+
+
+def register_opaque_factory(factory, *, module_adapter=None):
+ """Register construction and adaptation callbacks for Python module functions.
+
+ Parameters
+ ----------
+ factory : callable
+ Receives name, Python function, source text, and concrete span; returns a native module
+ function.
+ module_adapter : callable or None
+ Optional callback receiving completed module, original class, and base tuple; returns
+ the public module.
+
+ Returns
+ -------
+ None
+ Replaces the registered callbacks without executing function bodies.
+
+ Raises
+ ------
+ TypeError
+ A supplied callback is not callable.
+
+ Notes
+ -----
+ Registration is process-wide and retains callbacks until replaced. No per-parse state is
+ retained. None disables module adaptation.
+
+ Examples
+ --------
+ >>> # Frontend initialization registers its concrete opaque-function factory.
+ >>> register_opaque_factory(factory, module_adapter=adapter)
+ """
+ global _OPAQUE_FACTORY, _MODULE_ADAPTER
+ if not callable(factory) or (module_adapter is not None and not callable(module_adapter)):
+ raise TypeError("Module construction callbacks must be callable")
+ _OPAQUE_FACTORY, _MODULE_ADAPTER = factory, module_adapter
+
+
+class FunctionRecord:
+ """Own one function's declaration and definition construction state.
+
+ Parameters
+ ----------
+ builder : module
+ Registered construction namespace exposing decl_function, function,
+ func_name, arg and func_ret_type.
+ name : str
+ Original source function name.
+ options : dict
+ Evaluated keyword options accepted by this builder's function frames.
+ location : tuple or None
+ (SourceName or filename, line, end_line, column, end_column); columns are UTF-8 offsets.
+ captures : dict or None
+ Opaque lexical namespace captured by the generated factory. Existing
+ concrete primitive symbols retain identity; other host values are ignored.
+ The dictionary is consumed during initialization and is not retained.
+ capture_names : tuple[str]
+ Names mentioned by signature syntax; only these enclosing bindings may
+ seed the symbol frame. Unrelated outer symbols never constrain locals.
+ parameters : tuple[str]
+ Source parameter names that shadow outer symbolic captures; default empty.
+ local : bool
+ True for a nested function; its builder must support local references.
+
+ Notes
+ -----
+ builder/name/options/span/local are immutable construction configuration.
+ symbols is one TypeVarFrame, re-entered for declaration and definition;
+ its name-to-symbol dictionary is never shared with another function.
+ params is an insertion-ordered name-to-concrete-parameter dict populated by
+ parameter() only during declaration and reused without recreating identities.
+ signature_symbols retains the names introduced by parameter annotations,
+ excluding return-only free symbols; it is updated after each parameter.
+ return_type starts MISSING and is set by returns(); reference/function start
+ None and are populated after declaration/definition respectively. All state
+ belongs to this generated-program invocation and dies with its result graph.
+
+ Examples
+ --------
+ >>> record = FunctionRecord(builder, "main", {})
+ >>> # Enter a ModuleProgram before declaring or defining the record.
+ """
+
+ def __init__(
+ self,
+ builder,
+ name,
+ options,
+ location=None,
+ *,
+ local=False,
+ captures=None,
+ parameters=(),
+ capture_names=(),
+ ):
+ self.builder, self.name, self.options = builder, name, dict(options)
+ self.span, self.local = source_span(location), local
+ self.symbols = TypeVarFrame()
+ # Explicit lexical primitive symbols are captures, not new declarations.
+ # Parameter names shadow them; all concrete inspection stays builder-side.
+ for captured_name in capture_names:
+ value = (captures or {}).get(captured_name)
+ if (
+ captured_name not in parameters
+ and ir.is_prim_var(value)
+ and value.name in ("", captured_name)
+ ):
+ self.symbols.bind(captured_name, value)
+ self.params = {}
+ self.signature_symbols = set(self.symbols.symbols)
+ self.return_type = MISSING
+ self.reference = self.function = None
+
+ @contextmanager
+ def declaration(self):
+ """Enter the retained symbol frame and a native declaration frame.
+
+ Yields
+ ------
+ FunctionRecord
+ This record, ready for parameter and return-type operations.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Requires an active module builder. Frames leave in reverse order on normal and
+ exceptional exit. The finalized reference is retained only on success; symbols remain
+ available for definition.
+
+ Examples
+ --------
+ >>> with record.declaration():
+ ... record.parameter("x", builder.Prim("int32"))
+ """
+ mode = {"local": True} if self.local else {}
+ with self.symbols:
+ with self.builder.decl_function(**self.options, **mode, span=self.span) as frame:
+ self.builder.func_name(self.name)
+ yield self
+ self.reference = frame.reference
+
+ def parameter(self, name, annotation, location=None):
+ """Construct and retain a parameter in the active declaration.
+
+ Parameters
+ ----------
+ name : str
+ Original source parameter name.
+ annotation : object
+ Evaluated type, primitive variable, declaration descriptor, or zero-argument
+ annotation callable.
+ location : ir.Span, tuple, or None
+ Optional source range accepted by source_span; None creates no source metadata.
+
+ Returns
+ -------
+ ir.Var
+ Concrete parameter retained in source order.
+
+ Raises
+ ------
+ TypeError
+ A primitive declaration conflicts with a canonical symbol.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Primitive parameters reuse canonical symbol identity. Each successful parameter
+ refreshes signature_symbols; params is reused during definition rather than
+ reconstructing variables.
+
+ Examples
+ --------
+ >>> with record.declaration():
+ ... x = record.parameter("x", builder.Prim("int32"))
+ """
+ if callable(annotation) and not isinstance(annotation, ir.Expr | ir.Type):
+ annotation = annotation()
+ if isinstance(annotation, TypeVarDecl):
+ annotation = self.symbols.resolve(name, annotation.ty, span=source_span(location))
+ if isinstance(annotation, ir.PrimType):
+ annotation = self.symbols.resolve(name, annotation, span=source_span(location))
+ elif ir.is_prim_var(annotation):
+ annotation = self.symbols.bind(name, annotation)
+ value = self.builder.arg(name, annotation, span=source_span(location))
+ self.params[name] = value
+ self.signature_symbols = set(self.symbols.symbols)
+ return value
+
+ def predeclare(self, name, annotation, location=None):
+ """Reserve an explicit symbol type before dependent signature shapes.
+
+ Parameters
+ ----------
+ name : str
+ Original symbol name.
+ annotation : object
+ Primitive type, dtype spelling, declaration descriptor, variable, or zero-argument
+ constructor.
+ location : ir.Span, tuple, or None
+ Optional source range accepted by source_span; None creates no source metadata.
+
+ Returns
+ -------
+ ir.Var
+ Canonical primitive symbol.
+
+ Raises
+ ------
+ TypeError
+ The annotation is not primitive or conflicts with a prior declaration.
+ ValueError
+ The symbol name is invalid.
+
+ Notes
+ -----
+ Called during declaration. Registered symbols count as signature bindings, including
+ explicit body declarations in parameterless functions; return-only symbols remain
+ distinguishable.
+
+ Examples
+ --------
+ >>> with record.declaration():
+ ... n = record.predeclare("n", "int64")
+ """
+ if callable(annotation) and not isinstance(annotation, ir.Expr | ir.Type):
+ annotation = annotation()
+ if isinstance(annotation, TypeVarDecl):
+ annotation = annotation.ty
+ if ir.is_prim_var(annotation):
+ value = self.symbols.bind(name, annotation)
+ else:
+ value = self.symbols.resolve(name, annotation, span=source_span(location))
+ # A body declaration is already an explicit binding, even in a function
+ # with no parameters. Only symbols newly introduced by the return
+ # annotation itself are unbound return-only symbols.
+ self.signature_symbols = set(self.symbols.symbols)
+ return value
+
+ def symbol(self, name):
+ """Read a declared symbol after signature construction.
+
+ Parameters
+ ----------
+ name : str
+ Source name already introduced by the signature.
+
+ Returns
+ -------
+ ir.Var
+ Retained canonical symbol.
+
+ Raises
+ ------
+ KeyError
+ The symbol was not declared.
+
+ Notes
+ -----
+ This read-only operation does not require the retained frame to be active and does not
+ create missing symbols.
+
+ Examples
+ --------
+ >>> n = record.symbol("n")
+ """
+ return self.symbols.symbols[name]
+
+ def capture(self, name, fallback):
+ """Resolve a free annotation name against retained signature symbols.
+
+ Parameters
+ ----------
+ name : str
+ Source identifier.
+ fallback : object
+ Opaque enclosing lexical value.
+
+ Returns
+ -------
+ object
+ Existing canonical symbol, or fallback unchanged.
+
+ Notes
+ -----
+ No symbols or frames are introduced and no fallback is cached.
+
+ Examples
+ --------
+ >>> value = record.capture("n", outer_n)
+ """
+ return self.symbols.symbols.get(name, fallback)
+
+ def returns(self, annotation):
+ """Set the return annotation on the record and active native frame.
+
+ Parameters
+ ----------
+ annotation : object
+ Evaluated concrete return annotation accepted by the dialect.
+
+ Returns
+ -------
+ None
+ Retains the annotation for definition.
+
+ Raises
+ ------
+ ValueError
+ The return annotation introduced a symbol not bound by the signature.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ The same concrete annotation is reused in definition. Annotation evaluation itself has
+ no exactly-once guarantee.
+
+ Examples
+ --------
+ >>> with record.declaration():
+ ... record.returns(builder.Object())
+ """
+ introduced = set(self.symbols.symbols) - self.signature_symbols
+ if introduced:
+ raise ValueError(
+ f"Return annotation introduces unbound symbol {sorted(introduced)[0]!r}"
+ )
+ self.return_type = annotation
+ self.builder.func_ret_type(annotation)
+
+ def define(self, body):
+ """Run a generated body with retained parameters under a definition frame.
+
+ Parameters
+ ----------
+ body : callable
+ Receives parameters in source order; Python closures and defaults retain lexical
+ captures.
+
+ Returns
+ -------
+ ir.BaseFunc
+ Constructed function, also retained as function on this record.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Re-enters the same symbol frame and reuses parameter identities and return type.
+ Requires an active module builder. Host callback exceptions propagate and scopes clean
+ up normally.
+
+ Examples
+ --------
+ >>> function = record.define(lambda x: builder.return_(x))
+ """
+ mode = {"local": True, "reference": self.reference} if self.local else {}
+ with self.symbols:
+ with self.builder.function(**self.options, **mode, span=self.span) as frame:
+ self.builder.func_name(self.name)
+ for name, value in self.params.items():
+ self.builder.arg(name, value)
+ if self.return_type is not MISSING:
+ self.builder.func_ret_type(self.return_type)
+ body(*self.params.values())
+ self.function = frame.function
+ self.function.__name__ = self.name
+ return self.function
+
+
+class ModuleProgram:
+ """Own module construction while a generated builder program executes.
+
+ Parameters
+ ----------
+ name : str or None
+ Source module class name; None for a standalone function.
+ original : type or None
+ Optional original class for the registered module adapter.
+ bases : tuple
+ Evaluated source base classes; default empty.
+
+ Notes
+ -----
+ builder and frame own the active native builder/module scopes from entry to
+ exit. namespace is a fresh SimpleNamespace of declared references for source
+ Class.member expressions. python_functions stores original host callables by
+ source name until exit; result is None until successful native finalization.
+ No state is global or shared between invocations.
+
+ Examples
+ --------
+ >>> with ModuleProgram("Example") as program:
+ ... program.member("factor", 2)
+ 2
+ """
+
+ def __init__(self, name=None, original=None, bases=()):
+ self.name, self.original, self.bases = name, original, bases
+ self.builder = IRBuilder()
+ self.frame = None
+ self.namespace = SimpleNamespace()
+ self.python_functions = {}
+ self.result = None
+
+ def __enter__(self):
+ self.builder.__enter__()
+ self.frame = I.ir_module()
+ self.frame.__enter__()
+ return self
+
+ def __exit__(self, *error):
+ try:
+ self.frame.__exit__(*error)
+ if error[0] is None:
+ self.result = self.builder.get()
+ if self.python_functions:
+ self.result.pyfuncs = self.python_functions
+ if self.name is not None:
+ self.result.__name__ = self.name
+ if _MODULE_ADAPTER is not None:
+ self.result = _MODULE_ADAPTER(self.result, self.original, self.bases)
+ finally:
+ self.builder.__exit__(*error)
+
+ def reserve(self, name):
+ """Reserve and expose a stable named module function reference.
+
+ Parameters
+ ----------
+ name : str
+ Source function identifier.
+
+ Returns
+ -------
+ ir.GlobalVar
+ Native reference, also exposed on namespace.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Requires this program to be entered. The native module validates duplicate or
+ incompatible declarations.
+
+ Examples
+ --------
+ >>> with ModuleProgram("Example") as program:
+ ... reference = program.reserve("main")
+ """
+ reference = I.reserve_function(name)
+ setattr(self.namespace, name, reference)
+ return reference
+
+ def member(self, name, value):
+ """Register a class assignment in the generated module namespace.
+
+ Parameters
+ ----------
+ name : str
+ Source member identifier.
+ value : object
+ Evaluated Python value or concrete IR function.
+
+ Returns
+ -------
+ object
+ Module function reference or the unchanged host value.
+
+ Raises
+ ------
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ IR functions are declared and defined in the active module; other values remain opaque
+ namespace members.
+
+ Examples
+ --------
+ >>> with ModuleProgram("Example") as program:
+ ... program.member("factor", 2)
+ 2
+ """
+ if isinstance(value, ir.BaseFunc):
+ reference = I.decl_function(name, value)
+ I.def_function(name, value)
+ value = reference
+ setattr(self.namespace, name, value)
+ return value
+
+ def python(self, name, function, source, location=None):
+ """Register a host function without executing its body.
+
+ Parameters
+ ----------
+ name : str
+ Original source function name.
+ function : callable
+ Original Python function.
+ source : str
+ Original source text supplied to the registered factory.
+ location : ir.Span, tuple, or None
+ Optional source range accepted by source_span; None creates no source metadata.
+
+ Returns
+ -------
+ ir.GlobalVar
+ Opaque function reference, also exposed on namespace.
+
+ Raises
+ ------
+ ValueError
+ No opaque function factory has been registered.
+ tvm.error.TVMError
+ The active native frame rejects the operation or its concrete types.
+
+ Notes
+ -----
+ Uses the process-registered factory inside the active module. Original host callables
+ are retained by this program and transferred to the resulting module on successful exit.
+
+ Examples
+ --------
+ >>> reference = program.python("helper", helper, source)
+ """
+ if _OPAQUE_FACTORY is None:
+ raise ValueError("No opaque Python function constructor has been registered")
+ opaque = _OPAQUE_FACTORY(name, function, source, source_span(location))
+ reference = I.decl_function(name, opaque)
+ I.def_function(name, opaque)
+ self.python_functions[name] = function
+ setattr(self.namespace, name, reference)
+ return reference
+
+
+def require_defined(value, name):
+ """Return an exported value or report an undefined source name.
+
+ Parameters
+ ----------
+ value : object
+ Opaque lexical export, possibly MISSING.
+ name : str
+ Original source identifier for diagnostics.
+
+ Returns
+ -------
+ object
+ The value unchanged, including None.
+
+ Raises
+ ------
+ NameError
+ The value is MISSING.
+
+ Notes
+ -----
+ This operation enters no frame and retains no lexical state.
+
+ Examples
+ --------
+ >>> require_defined(None, "x") is None
+ True
+ """
+ if value is MISSING:
+ raise NameError(f"name {name!r} is not defined")
+ return value
+
+
+def is_python_bool(value):
+ """Check for a host boolean without invoking truth conversion.
+
+ Parameters
+ ----------
+ value : object
+ Host value or builder expression.
+
+ Returns
+ -------
+ bool
+ True exactly for Python bool values.
+
+ Notes
+ -----
+ Pure predicate used to select construction-time control flow without inspecting IR in
+ generated code.
+
+ Examples
+ --------
+ >>> is_python_bool(True)
+ True
+ >>> is_python_bool(1)
+ False
+ """
+ return isinstance(value, bool)
diff --git a/python/tvm/script/ir_builder/type_var_frame.py b/python/tvm/script/ir_builder/type_var_frame.py
new file mode 100644
index 0000000000..f70cede1f3
--- /dev/null
+++ b/python/tvm/script/ir_builder/type_var_frame.py
@@ -0,0 +1,343 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Function-owned symbolic type variables on the native builder frame stack."""
+
+from dataclasses import dataclass
+from functools import wraps
+from typing import TypeVar
+
+import tvm_ffi
+
+from tvm import ir
+
+from . import _ffi_api
+from .base import IRBuilder, IRBuilderFrame, source_span
+
+
+@dataclass(frozen=True)
+class TypeVarDecl:
+ """Describe an unbound primitive symbolic declaration.
+
+ Parameters
+ ----------
+ ty : ir.PrimType or str
+ Concrete primitive type or dtype spelling, converted to PrimType.
+
+ Raises
+ ------
+ TypeError
+ The supplied type is not primitive.
+
+ Notes
+ -----
+ The immutable ty field stores only the concrete type. Construction creates no symbol and
+ enters no frame; a dialect bind_ determines the canonical symbol identity.
+
+ Examples
+ --------
+ >>> TypeVarDecl("int32").ty
+ int32
+ """
+
+ ty: object
+
+ def __post_init__(self):
+ ty = ir.PrimType(self.ty) if isinstance(self.ty, str) else self.ty
+ if not isinstance(ty, ir.PrimType):
+ raise TypeError("A symbolic declaration requires a primitive type")
+ object.__setattr__(self, "ty", ty)
+
+
+@tvm_ffi.register_object("script.ir_builder.TypeVarFrame")
+class TypeVarFrame(IRBuilderFrame):
+ """Retain canonical symbols across one function declaration and definition.
+
+ Notes
+ -----
+ The zero-argument constructor creates a native symbols map from names to ir.Var. Enter this
+ frame outside the native function frame with an active IRBuilder. Exit pops the native frame
+ stack but retains symbols, so the same instance can be re-entered for definition. Blocks
+ reuse it; nested functions own separate instances. Nearest active lookup provides shadowing
+ without a Python global stack. Exceptional cleanup follows IRBuilderFrame.
+
+ Examples
+ --------
+ >>> from tvm.script.ir_builder import IRBuilder
+ >>> with IRBuilder():
+ ... with TypeVarFrame() as symbols:
+ ... n = symbols.resolve("n")
+ """
+
+ def __init__(self):
+ self.__init_handle_by_constructor__(_ffi_api.TypeVarFrame)
+
+ @staticmethod
+ def current():
+ """Find the nearest active function symbol frame.
+
+ Returns
+ -------
+ TypeVarFrame
+ Nearest matching frame on the current native builder stack.
+
+ Raises
+ ------
+ ValueError
+ No builder or symbol frame is active.
+
+ Notes
+ -----
+ This lookup changes no state and never falls back to a different builder.
+
+ Examples
+ --------
+ >>> with IRBuilder():
+ ... with TypeVarFrame() as symbols:
+ ... assert TypeVarFrame.current().same_as(symbols)
+ """
+ for frame in reversed(IRBuilder.current().frames):
+ if isinstance(frame, TypeVarFrame):
+ return frame
+ raise ValueError("Symbol resolution requires an active TypeVarFrame")
+
+ def resolve(self, name, ty=None, *, span=None):
+ """Resolve or introduce one canonical primitive symbol.
+
+ Parameters
+ ----------
+ name : str
+ Nonempty source binding name.
+ ty : ir.PrimType, str, or None
+ Explicit primitive type. None reuses an existing type or creates int64.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines,
+ and start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ ir.Var
+ The canonical variable retained by this frame.
+
+ Raises
+ ------
+ ValueError
+ The name is empty or not a string.
+ TypeError
+ The supplied type is nonprimitive.
+
+ Notes
+ -----
+ First use adds an entry to symbols and attaches the supplied span. Reuse preserves
+ identity and the original span. The map survives frame exit and declaration/definition
+ re-entry. Semantic consistency is left to the existing well-formedness pass.
+
+ Examples
+ --------
+ >>> symbols = TypeVarFrame()
+ >>> n = symbols.resolve("n", "int64")
+ >>> assert symbols.resolve("n").same_as(n)
+ """
+ if not isinstance(name, str) or not name:
+ raise ValueError("A symbolic variable requires a nonempty string name")
+ if isinstance(ty, str):
+ ty = ir.PrimType(ty)
+ if ty is not None and not isinstance(ty, ir.PrimType):
+ raise TypeError("A symbolic variable requires a primitive type")
+ if name in self.symbols:
+ return self.symbols[name]
+ value = ir.Var(name, ir.PrimType("int64") if ty is None else ty, source_span(span))
+ _ffi_api.TypeVarFrameSetSymbol(self, name, value)
+ return value
+
+ def bind(self, name, symbol):
+ """Register an existing primitive variable or reuse its canonical binding.
+
+ Parameters
+ ----------
+ name : str
+ Nonempty source binding name.
+ symbol : ir.Var
+ Primitive variable to retain under the source binding name.
+
+ Returns
+ -------
+ ir.Var
+ The previously registered symbol, or the supplied symbol.
+
+ Raises
+ ------
+ ValueError
+ The binding name is empty or is not a string.
+ TypeError
+ The supplied value is not a primitive variable.
+
+ Notes
+ -----
+ An anonymous newly registered variable receives its source name. This method enters no
+ frame and retains the canonical variable across re-entry; it never clones the variable.
+
+ Examples
+ --------
+ >>> symbols = TypeVarFrame()
+ >>> n = symbols.resolve("n")
+ >>> assert symbols.bind("n", n).same_as(n)
+ """
+ if not isinstance(name, str) or not name:
+ raise ValueError("A symbolic variable requires a nonempty string name")
+ if not ir.is_prim_var(symbol):
+ raise TypeError("A symbolic binding requires a primitive Var")
+ if name in self.symbols:
+ return self.resolve(name, symbol.ty)
+ if not symbol.name:
+ IRBuilder.name(name, symbol)
+ _ffi_api.TypeVarFrameSetSymbol(self, name, symbol)
+ return symbol
+
+
+def resolve_type_var(name, *, dtype=None, span=None):
+ """Resolve a symbol in the nearest active function frame.
+
+ Parameters
+ ----------
+ name : str
+ Nonempty source binding name.
+ dtype : ir.PrimType, str, or None
+ Explicit type; None reuses an existing type or defaults new symbols to int64.
+ span : ir.Span, tuple, or None
+ Optional source range. A tuple contains SourceName (or filename), start/end lines, and
+ start/end UTF-8 byte columns; None adds no metadata.
+
+ Returns
+ -------
+ ir.Var
+ Canonical symbol owned by the active TypeVarFrame.
+
+ Raises
+ ------
+ ValueError
+ No symbol frame is active or the name is invalid.
+ TypeError
+ The explicit type is nonprimitive or incompatible.
+
+ Notes
+ -----
+ This operation updates only the nearest frame's symbol map, retaining identity through
+ declaration and definition. It enters no additional frame.
+
+ Examples
+ --------
+ >>> with IRBuilder():
+ ... with TypeVarFrame():
+ ... n = resolve_type_var("n")
+ """
+ return TypeVarFrame.current().resolve(name, dtype, span=span)
+
+
+def wrap_expression_constructor(constructor, call_signature, policy, *, as_type=False):
+ """Adapt an eager constructor using parser-owned expression-string metadata.
+
+ Parameters
+ ----------
+ constructor : callable
+ Concrete annotation constructor to invoke.
+ call_signature : inspect.Signature
+ Signature used to bind arguments before inspecting marked fields.
+ policy : ExprStrPolicy
+ Immutable parser metadata specifying fields and scalar-string handling.
+ as_type : bool
+ False returns a callable; True creates an annotation-class surface supporting unions.
+
+ Returns
+ -------
+ callable or type
+ Adapter preserving the constructor signature and public documentation.
+
+ Raises
+ ------
+ TypeError
+ Argument binding fails, a TypeVar has constraints/bounds, or unresolved strings occur in
+ an active builder.
+ ValueError
+ A TypeVar requires resolution but no symbol frame is active.
+
+ Notes
+ -----
+ Outside a builder, unresolved marked strings and TypeVars yield composable MissingType.
+ Inside a builder, TypeVars resolve through the nearest native TypeVarFrame. The closure
+ retains configuration only, never invocation results or symbols; the factory owns no
+ metadata registry. Errors from the concrete constructor propagate.
+
+ Examples
+ --------
+ >>> import inspect
+ >>> from tvm.script.parser_v2.protocol import ExprStrPolicy
+ >>> # expr_str_args registers metadata and calls this factory internally.
+ """
+ fields = policy.fields
+
+ def unresolved(value, nested=False):
+ if isinstance(value, str):
+ return nested or policy.scalar_strings
+ if isinstance(value, TypeVar):
+ return True
+ if isinstance(value, tuple | list):
+ return any(unresolved(item, True) for item in value)
+ return False
+
+ @wraps(constructor)
+ def invoke(*args, **kwargs):
+ bound = call_signature.bind(*args, **kwargs)
+ if IRBuilder.is_in_scope():
+ # typing.TypeVar is ordinary eager Python metadata. Resolve it
+ # here, never in the syntax-only transpiler.
+ def resolve(value):
+ if isinstance(value, TypeVar):
+ if value.__bound__ is not None or value.__constraints__:
+ raise TypeError("A symbolic TypeVar cannot have constraints or a bound")
+ return TypeVarFrame.current().resolve(value.__name__)
+ if isinstance(value, tuple):
+ return tuple(resolve(item) for item in value)
+ if isinstance(value, list):
+ return [resolve(item) for item in value]
+ return value
+
+ for field in fields:
+ if field in bound.arguments:
+ bound.arguments[field] = resolve(bound.arguments[field])
+ if any(unresolved(bound.arguments[field]) for field in fields if field in bound.arguments):
+ if IRBuilder.is_in_scope():
+ raise TypeError(
+ "Builder expression arguments require concrete symbols, not strings"
+ )
+ return ir.Type.missing()
+ return constructor(*bound.args, **bound.kwargs)
+
+ result = invoke
+ if as_type:
+ # The class is an annotation surface, not an IR or proxy type.
+ # __new__ returns the concrete construction result (or MissingType).
+ result = type(
+ constructor.__name__,
+ (),
+ {
+ "__new__": lambda cls, *args, **kwargs: invoke(*args, **kwargs),
+ "__signature__": call_signature,
+ "__doc__": constructor.__doc__,
+ "__module__": constructor.__module__,
+ },
+ )
+ return result
diff --git a/python/tvm/script/parser_v2/__init__.py b/python/tvm/script/parser_v2/__init__.py
new file mode 100644
index 0000000000..173fd06500
--- /dev/null
+++ b/python/tvm/script/parser_v2/__init__.py
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Translate original Python ASTs into calls on a registered construction namespace.
+
+Shared I infrastructure owns builder lifetime, spans and module identities; each
+context selects X from reverse-registered function metadata. Calls return concrete
+values, constructor metadata describes syntax, and compiled AST locations remain
+those of the original source. Entry modules register policies; this parser never
+imports their namespaces.
+"""
+
+import importlib
+import sys
+
+# Process-wide fixed export names, never per-function construction state. Lazy
+# entry loading permits builder imports to register syntax policy without a
+# parser/frontend/builder-construction import cycle.
+_FRONTEND_EXPORTS = (
+ "_NAMESPACES",
+ "from_source",
+ "ir_module",
+ "make_decorator",
+ "make_helper",
+ "parse",
+ "pyfunc",
+ "register_namespace",
+)
+__all__ = [name for name in _FRONTEND_EXPORTS if not name.startswith("_")]
+
+
+_initialized = False
+
+
+def _initialize():
+ """Register existing dialect entry identities for explicit v2 parsing only."""
+ global _initialized
+ if _initialized:
+ return
+ from tvm.script.parser import frontend as legacy
+ from tvm.script.ir_builder import construction
+ from tvm.tirx.script import builder as tir_builder
+ from tvm.tirx.script.builder import v2 as tir_v2
+ from tvm.relax import script as relax_entry
+ from tvm.relax.script import builder as relax_builder
+ from tvm.relax.script.builder import v2 as relax_v2
+ from . import frontend, protocol
+
+ protocol.register_builder(tir_builder, tir_v2)
+ protocol.register_builder(relax_builder, relax_v2)
+ frontend._NAMESPACES.update(legacy._NAMESPACES)
+ construction.register_opaque_factory(
+ relax_entry._opaque_function, module_adapter=relax_entry._python_module
+ )
+ _initialized = True
+
+
+def __getattr__(name):
+ if name in _FRONTEND_EXPORTS:
+ _initialize()
+ frontend = importlib.import_module(f"{__name__}.frontend")
+ return getattr(frontend, name)
+ frontend = sys.modules.get(f"{__name__}.frontend")
+ if frontend is not None and name in getattr(frontend, "_NAMESPACES", {}):
+ return frontend._NAMESPACES[name]
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/python/tvm/script/parser_v2/diagnostics.py b/python/tvm/script/parser_v2/diagnostics.py
new file mode 100644
index 0000000000..f8563b9142
--- /dev/null
+++ b/python/tvm/script/parser_v2/diagnostics.py
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Source-located construction errors without replacing their original traceback."""
+
+import ast
+import linecache
+import traceback
+
+from tvm.error import DiagnosticError
+
+
+def diagnostic_error(error, compiler):
+ """Return a DiagnosticError highlighting the original source operation.
+
+ error is the caught Python/builder exception. compiler supplies the source
+ filename and original AST for fallback coordinates. Builder errors carry
+ plain location tuples; otherwise Python traceback/SyntaxError coordinates
+ identify the innermost original operation. Python 3.11+ supplies precise
+ traceback columns, while earlier versions fall back to the source line.
+ No frame or IR object is inspected; the caller attaches error as __cause__.
+ """
+ filename = compiler.filename
+ location = getattr(error, "__tvm_script_location__", None)
+ if location is not None:
+ filename, start, end, column, end_column = location
+ elif isinstance(error, SyntaxError) and error.lineno:
+ start = error.lineno
+ end = error.end_lineno or start
+ column = max((error.offset or 1) - 1, 0)
+ end_column = max((error.end_offset or column + 2) - 1, column + 1)
+ else:
+ frames = [
+ frame
+ for frame in traceback.extract_tb(error.__traceback__)
+ if frame.filename == filename
+ ]
+ if frames:
+ frame = frames[-1]
+ start, end = frame.lineno, getattr(frame, "end_lineno", None) or frame.lineno
+ column = getattr(frame, "colno", None) or 0
+ end_column = getattr(frame, "end_colno", None)
+ if getattr(frame, "end_lineno", None) is None:
+ # Python before 3.11 supplies only a line. Recover the original
+ # statement range from source AST, without an IR span context or
+ # executing/reparsing a generated program. Prefer the narrowest
+ # statement beginning on that line (a nested body over its def).
+ candidates = [
+ node
+ for node in ast.walk(compiler.tree)
+ if isinstance(node, ast.stmt) and node.lineno == start
+ ]
+ if candidates:
+ node = min(candidates, key=lambda item: item.end_lineno - item.lineno)
+ end, column, end_column = node.end_lineno, node.col_offset, node.end_col_offset
+ else:
+ node = compiler.tree.body[-1]
+ start, end = node.lineno, node.lineno
+ column, end_column = node.col_offset, None
+ lines = [f"{filename}:{start}: {type(error).__name__}: {error}"]
+ for number in range(start, end + 1):
+ source = linecache.getline(filename, number).rstrip("\n")
+ first = column if number == start else len(source) - len(source.lstrip())
+ last = end_column if number == end and end_column is not None else len(source)
+ lines.extend(
+ (
+ f" {number} | {source}",
+ " " * (len(str(number)) + 4 + first) + "^" * max(last - first, 1),
+ )
+ )
+ return DiagnosticError("\n".join(lines))
diff --git a/python/tvm/script/parser_v2/expression.py b/python/tvm/script/parser_v2/expression.py
new file mode 100644
index 0000000000..21a10c695d
--- /dev/null
+++ b/python/tvm/script/parser_v2/expression.py
@@ -0,0 +1,246 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Syntax-only rewriting of registered expression-string arguments.
+
+Existing AST nodes retain all four Python location fields. Parsed strings use a
+UTF-8 byte-offset map through the original literal spelling, including escapes
+and physical newlines. If a literal cannot be recovered (for example a synthetic
+AST or adjacent concatenated literals), ranges conservatively fall back to the
+literal's complete original range. Rewriting never evaluates source expressions.
+"""
+
+import ast
+import copy
+import inspect
+import linecache
+import re
+
+from . import protocol
+
+
+class _LiteralParser:
+ # filename is immutable for one expression rewrite and owns no source/IR cache.
+ def __init__(self, filename):
+ self.filename = filename
+
+ def _string_expression(self, node):
+ try:
+ expression = ast.parse(node.value, mode="eval").body
+ except SyntaxError as error:
+ raise SyntaxError(f"Invalid annotation expression: {error.msg}") from error
+ # Example: in R.Tensor(("n + 1",), "float32"), the generated resolve
+ # call for n retains the byte range of n inside the quoted literal, not
+ # the whole constructor. A triple-quoted physical newline advances lineno;
+ # an escaped \n advances decoded input but maps back to the escape's real
+ # source bytes. Generated call wrappers copy these four mapped fields.
+ source = "".join(linecache.getlines(self.filename))
+ literal = ast.get_source_segment(source, node) if source else None
+ positions = self._literal_positions(literal, node) if literal else None
+ lines = node.value.splitlines(keepends=True)
+ for inner in ast.walk(expression):
+ if not hasattr(inner, "lineno"):
+ continue
+ for line_field, column_field in (
+ ("lineno", "col_offset"),
+ ("end_lineno", "end_col_offset"),
+ ):
+ line, column = getattr(inner, line_field), getattr(inner, column_field)
+ offset = len("".join(lines[: line - 1]).encode("utf-8")) + column
+ if positions is not None and offset in positions:
+ line, column = positions[offset]
+ else:
+ # Without a reliable spelling map (synthetic or concatenated
+ # literals), preserve the literal's entire real source range.
+ # Never invent physical lines from decoded escape sequences.
+ line = getattr(node, line_field)
+ column = getattr(node, column_field)
+ setattr(inner, line_field, line)
+ setattr(inner, column_field, column)
+ return expression
+
+ @staticmethod
+ def _literal_positions(literal, node):
+ """Map decoded expression byte offsets back through the literal's escapes."""
+ match = re.match("(?i:([rub]*))([\"'])", literal)
+ if match is None:
+ return None
+ prefix, quote = match.groups()
+ width = 3 if literal[len(prefix) :].startswith(quote * 3) else 1
+ start, stop = len(prefix) + width, len(literal) - width
+ delimiter = quote * width
+ positions, decoded, offset = {}, "", 0
+ index = start
+
+ def location(raw_index):
+ before = literal[:raw_index]
+ line = node.lineno + before.count("\n")
+ column = len(before.rsplit("\n", 1)[-1].encode("utf-8"))
+ return line, column + (node.col_offset if line == node.lineno else 0)
+
+ while index < stop:
+ end = index + 1
+ if literal[index] == "\\" and "r" not in prefix.lower():
+ escape = re.match(
+ r"\\(?:N\{[^}]*\}|u[0-9a-fA-F]{4}|U[0-9a-fA-F]{8}|x[0-9a-fA-F]{2}|[0-7]{1,3}|\r?\n|.)",
+ literal[index:stop],
+ )
+ if escape:
+ end = index + len(escape.group())
+ piece = literal[index:end]
+ try:
+ value = ast.literal_eval(prefix + delimiter + piece + delimiter)
+ except (SyntaxError, ValueError):
+ return None
+ positions[offset] = location(index)
+ for char in value:
+ offset += len(char.encode("utf-8"))
+ positions[offset] = location(end)
+ decoded += value
+ index = end
+ return positions if decoded == node.value else None
+
+
+def rewrite_expression(node, resolve, namespace, filename, *, annotation=False):
+ """Copy an expression and rewrite registered expression-string fields.
+
+ Parameters
+ ----------
+ node : ast.expr
+ Original source expression. It is never mutated.
+ resolve : callable
+ Static namespace lookup for Name and Attribute nodes. Returns host
+ namespaces or callables, or None, without executing descriptors.
+ namespace : str
+ Generated builder binding used for ``resolve_type_var`` calls.
+ filename : str
+ Original compilation filename and linecache source key.
+ annotation : bool, optional
+ Parse a quoted whole annotation before constructor-field rewriting.
+ Default is False. True also emits lexical bindings for names
+ introduced by signature expression strings.
+
+ Returns
+ -------
+ ast.expr
+ New expression AST preserving original source ranges.
+
+ Raises
+ ------
+ SyntaxError
+ If a marked string is not a Python expression.
+ TypeError or ValueError
+ If a registered constructor's signature cannot be inspected.
+
+ Notes
+ -----
+ Rewriting never evaluates annotations or enters builder frames. Temporary
+ literal and symbol transformers belong to this call only. String-derived
+ nodes map decoded UTF-8 offsets back through the original literal spelling;
+ unrecoverable spellings use the literal's full source range.
+
+ Examples
+ --------
+ A registered shape string becomes a builder resolution call::
+
+ R.Tensor(("n + 1",), "float32")
+ # Rewritten shape: (X.resolve_type_var("n") + 1,)
+
+ The unmarked dtype string remains unchanged.
+ """
+ literals = _LiteralParser(filename)
+
+ class Symbols(ast.NodeTransformer):
+ # Pattern: "n + m * 2" -> X.resolve_type_var("n") + X.resolve_type_var("m") * 2.
+ def visit_Attribute(self, current):
+ # A recognized namespace in T.max/T.Cast is a Python namespace,
+ # not an expression-string symbol. Keep its attribute path intact.
+ if resolve(current.value) is not None:
+ return current
+ return self.generic_visit(current)
+
+ def visit_Call(self, current):
+ # Constructor/operator names keep normal Python callable resolution;
+ # symbol-bearing operands still resolve through the active builder.
+ if resolve(current.func) is None:
+ current.func = self.visit(current.func)
+ current.args = [self.visit(value) for value in current.args]
+ for keyword in current.keywords:
+ keyword.value = self.visit(keyword.value)
+ return current
+
+ def visit_Name(self, current):
+ call = ast.copy_location(
+ ast.Call(
+ ast.Attribute(ast.Name(namespace, ast.Load()), "resolve_type_var", ast.Load()),
+ [ast.Constant(current.id)],
+ [],
+ ),
+ current,
+ )
+ if annotation:
+ # Signature strings introduce lexical names at their actual
+ # evaluation point: ("n + 1", n) works; (n, "n") stays unbound.
+ # This generated assignment expression contains only a builder
+ # call and preserves the original symbol's source range.
+ result = ast.copy_location(
+ ast.NamedExpr(ast.Name(current.id, ast.Store()), call), current
+ )
+ result._tvm_signature_binding = True
+ return result
+ return call
+
+ class Rewrite(ast.NodeTransformer):
+ def field(self, current, policy, nested=False):
+ if isinstance(current, ast.Tuple | ast.List):
+ current.elts = [self.field(value, policy, True) for value in current.elts]
+ return current
+ if isinstance(current, ast.Constant) and isinstance(current.value, str):
+ if nested or policy.scalar_strings:
+ return Symbols().visit(literals._string_expression(current))
+ return self.visit(current)
+
+ def visit_Call(self, current):
+ # Pattern: R.Tensor(("n + 1",), dtype) ->
+ # R.Tensor((X.resolve_type_var("n") + 1,), dtype).
+ constructor = resolve(current.func)
+ policy = protocol.expr_str_policy(constructor)
+ parameters = []
+ if policy is not None:
+ parameters = [
+ p.name
+ for p in inspect.signature(constructor).parameters.values()
+ if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
+ ]
+ current.func = self.visit(current.func)
+ current.args = [
+ self.field(value, policy)
+ if index < len(parameters) and parameters[index] in policy.fields
+ else self.visit(value)
+ for index, value in enumerate(current.args)
+ ]
+ for keyword in current.keywords:
+ keyword.value = (
+ self.field(keyword.value, policy)
+ if policy and keyword.arg in policy.fields
+ else self.visit(keyword.value)
+ )
+ return current
+
+ result = copy.deepcopy(node)
+ if annotation and isinstance(result, ast.Constant) and isinstance(result.value, str):
+ result = literals._string_expression(result)
+ return ast.fix_missing_locations(Rewrite().visit(result))
diff --git a/python/tvm/script/parser_v2/frontend.py b/python/tvm/script/parser_v2/frontend.py
new file mode 100644
index 0000000000..4b787d7391
--- /dev/null
+++ b/python/tvm/script/parser_v2/frontend.py
@@ -0,0 +1,738 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Source acquisition and declaration/body execution for registered builders."""
+
+import __future__
+
+import ast
+import copy
+import inspect
+import linecache
+import sys
+import textwrap
+from functools import wraps
+from types import SimpleNamespace
+from typing import TypeVar
+
+from tvm.error import DiagnosticError
+from tvm.ir import SourceName
+from tvm.script.ir_builder import base, construction
+
+from . import protocol as syntax_protocol
+from .diagnostics import diagnostic_error
+from .transpile import IRBuilderTranspiler, NameCollector
+
+# Runtime callables are injected by identity; this namespace owns no IR state.
+_EXECUTION = SimpleNamespace(
+ _at=base.at,
+ MISSING=base.MISSING,
+ require_defined=construction.require_defined,
+ is_python_bool=construction.is_python_bool,
+ slice=slice,
+ locals=locals,
+)
+_NAMESPACES = {}
+
+
+def register_namespace(alias, namespace):
+ """Register a host namespace for source-text entry points.
+
+ Parameters
+ ----------
+ alias : str
+ Python identifier used to refer to the namespace in source text.
+ namespace : object
+ Module or object bound to ``alias``.
+
+ Returns
+ -------
+ None
+
+ Notes
+ -----
+ Registration replaces the process-wide alias entry. Each new `Compiler`
+ copies this table; existing compilations are unaffected. This operation
+ enters no builder frame.
+ """
+ _NAMESPACES[alias] = namespace
+
+
+def _closure_values(function):
+ values = {}
+ for name, cell in zip(function.__code__.co_freevars, function.__closure__ or ()):
+ try:
+ values[name] = cell.cell_contents
+ except ValueError:
+ # Recursive and later-bound locals are empty until the helper is used.
+ pass
+ return values
+
+
+def _capture(obj):
+ target = obj if inspect.isfunction(obj) else None
+ module = inspect.getmodule(obj)
+ env = dict(vars(module)) if module is not None else {}
+ env.update(getattr(target, "__globals__", {}))
+ if target is not None:
+ env.update(_closure_values(target))
+ filename = inspect.getsourcefile(obj)
+ # Deferred annotations may be the only use of an enclosing local, so Python
+ # need not put that value in the function's closure cells.
+ frames = inspect.stack()
+ try:
+ for info in reversed(frames):
+ if info.filename == filename:
+ env.update(info.frame.f_locals)
+ finally:
+ del frames
+ if inspect.isclass(obj):
+ env.update(vars(obj))
+ return env
+
+
+def _inside_class(function):
+ frame = inspect.currentframe().f_back
+ try:
+ while frame is not None:
+ local = frame.f_locals
+ if local.get("__module__") == function.__module__ and "__qualname__" in local:
+ return True
+ if frame.f_code.co_filename == function.__code__.co_filename:
+ return False
+ frame = frame.f_back
+ finally:
+ del frame
+ return False
+
+
+def make_decorator(builder, *, option_map=None, defaults=None):
+ """Create and register a function decorator for a construction namespace.
+
+ Parameters
+ ----------
+ builder : object
+ Namespace implementing the function construction protocol.
+ option_map : mapping of str to str, optional
+ Public option names mapped to builder keyword names. Default is None,
+ interpreted as an empty mapping; supplied entries are copied.
+ defaults : mapping of str to object, optional
+ Default builder keyword values. Default is None, interpreted as an
+ empty mapping; supplied entries are copied.
+
+ Returns
+ -------
+ decorator : callable
+ Callable supporting ``@decorator``, ``@decorator(**options)``, and
+ ``decorator(function)``.
+
+ Raises
+ ------
+ ValueError
+ When the returned decorator receives a non-function positional value.
+ DiagnosticError
+ When standalone construction fails during `parse` execution.
+
+ Notes
+ -----
+ Class members retain their Python functions until module construction.
+ Standalone functions immediately transpile and execute a builder program.
+ Annotations must be safe to re-evaluate: eager MissingType placeholders
+ are not cached, and source annotations execute in declaration frames.
+ Registration persists for the lifetime of the returned decorator.
+ """
+ mapping, default_options = dict(option_map or {}), dict(defaults or {})
+
+ def decorator(function=None, **options):
+ if function is not None and not inspect.isfunction(function):
+ raise ValueError("Construction decorators require a function or keyword options")
+
+ def apply(function):
+ function.__tvm_function_info__ = decorator.__tvm_function_info__
+ function.__tvm_function_options__ = options
+ if _inside_class(function):
+ return function
+ return parse(function, _capture(function))
+
+ return apply(function) if function is not None else apply
+
+ return syntax_protocol.register_function(
+ decorator, builder, option_map=mapping, defaults=default_options
+ )
+
+
+def make_helper(builder, *, preserve_return=True, late_binding=False):
+ """Create a decorator for helpers executed in a caller's builder frames.
+
+ Parameters
+ ----------
+ builder : object
+ Namespace implementing construction operations for the helper body.
+ preserve_return : bool, optional
+ Keep helper returns as ordinary Python control flow. Default is True.
+ late_binding : bool, optional
+ Refresh captured closure cells on each call. Default is False.
+
+ Returns
+ -------
+ decorator : callable
+ Accepts a function directly or keyword options. The ``hygienic``
+ option defaults to True and snapshots the definition environment;
+ False captures the calling environment on each invocation. Other
+ options remain metadata for the namespace consumer.
+
+ Raises
+ ------
+ ValueError
+ When the returned decorator receives a non-function positional value.
+ TypeError
+ When a helper invocation cannot bind its Python signature.
+
+ Notes
+ -----
+ Each invocation binds arguments and defaults, transpiles the original
+ body, and returns its result. It shares the caller's active construction
+ frames instead of declaring an IR function. Source acquisition,
+ compilation, and builder exceptions propagate to the caller.
+ """
+
+ def decorator(function=None, **options):
+ if function is not None and not inspect.isfunction(function):
+ raise ValueError("Construction decorators require a function or keyword options")
+
+ def apply(function):
+ definition_env = _capture(function)
+
+ @wraps(function)
+ def invoke(*args, **kwargs):
+ bound = inspect.signature(function).bind(*args, **kwargs)
+ bound.apply_defaults()
+ environment = (
+ {**definition_env, **(_closure_values(function) if late_binding else {})}
+ if options.get("hygienic", True)
+ else _capture(function)
+ )
+ compiler = Compiler(function, environment)
+ node = compiler.tree.body[0]
+ return compiler.run_statements(
+ node.body,
+ builder,
+ {**compiler.env, **bound.arguments},
+ set(bound.arguments),
+ preserve_return=preserve_return,
+ )
+
+ invoke.__tvm_construction_helper__ = (builder, options)
+ return invoke
+
+ return apply(function) if function is not None else apply
+
+ return decorator
+
+
+def pyfunc(function):
+ """Mark a Python function for opaque registration in a module.
+
+ Parameters
+ ----------
+ function : callable
+ Python function supporting attribute assignment.
+
+ Returns
+ -------
+ callable
+ The same function, with its registration marker attached.
+
+ Raises
+ ------
+ AttributeError
+ If the supplied object does not support the marker attribute.
+
+ Notes
+ -----
+ The function body remains ordinary Python. Module construction delegates
+ registration to the builder runtime. This decorator enters no frame and
+ preserves callable identity for the function's lifetime.
+ """
+ function.__tvm_python_function__ = True
+ return function
+
+
+syntax_protocol.register_function(pyfunc, None, python=True)
+
+
+class Compiler:
+ """Acquire source and execute a location-preserving builder program.
+
+ Parameters
+ ----------
+ source : str or function or type
+ Source text, Python function, or Python class to compile.
+ env : mapping of str to object, optional
+ Opaque lexical bindings overriding registered source aliases.
+ Default is None, interpreted as an empty mapping.
+ filename : str, optional
+ Override the original filename. Default is None, which uses the
+ inspected filename for objects and ``"<str>"`` for source text.
+ track_span : bool, optional
+ Emit shared source metadata and IR location instrumentation.
+ Default is True. False preserves Python AST locations and tracebacks.
+
+ Raises
+ ------
+ OSError
+ If inspection cannot recover the object's source.
+ TypeError
+ If the source object cannot be inspected.
+ SyntaxError
+ If the acquired text is not valid Python syntax.
+
+ Notes
+ -----
+ ``env``, ``original``, and the fresh source ``tree`` belong to this
+ compilation. ``filename`` and ``compile_flags`` retain the source's file
+ and annotation mode. ``name_map`` holds reserved identifiers and prefix
+ counters shared by all generated functions. ``builder_name`` and
+ ``infrastructure_name`` are reserved aliases.
+
+ With tracking enabled, one SourceName metadata object is held in ``env``
+ under ``source_name_binding`` for the complete source unit. Disabling
+ tracking skips that object and generated span instrumentation. No
+ concrete IR construction state is retained; `build` passes its result
+ opaquely to the caller.
+ """
+
+ def __init__(self, source, env=None, filename=None, *, track_span: bool = True):
+ self.env = {"TypeVar": TypeVar, "tvm": sys.modules.get("tvm"), **_NAMESPACES, **(env or {})}
+ self.original = source
+ self.track_span = track_span
+ members = vars(source).values() if inspect.isclass(source) else (source,)
+ self.compile_flags = 0
+ for member in members:
+ code = getattr(member, "__code__", None)
+ if code is not None:
+ self.compile_flags |= code.co_flags & __future__.annotations.compiler_flag
+ if isinstance(source, str):
+ text = source
+ self.filename = filename or "<str>"
+ start, indent = 1, 0
+ linecache.cache[self.filename] = (
+ len(text),
+ None,
+ text.splitlines(keepends=True),
+ self.filename,
+ )
+ else:
+ lines, start = inspect.getsourcelines(source)
+ text = "".join(lines)
+ self.filename = filename or inspect.getsourcefile(source)
+ indent = len(lines[0]) - len(lines[0].lstrip())
+ self.tree = ast.parse(textwrap.dedent(text), self.filename)
+ if start != 1:
+ ast.increment_lineno(self.tree, start - 1)
+ if indent:
+ for node in ast.walk(self.tree):
+ if hasattr(node, "col_offset"):
+ node.col_offset += indent
+ node.end_col_offset += indent
+ # env is the opaque execution namespace; the transformer receives only
+ # host namespace/callable bindings needed to identify registered metadata.
+ # name_map is shared by every nested factory/body in this compilation unit.
+ self.name_map = dict.fromkeys(self.env, 0)
+ NameCollector(self.name_map).visit(self.tree)
+ self.builder_name = self.fresh()
+ self.infrastructure_name = self.fresh()
+ # One source metadata object is shared by this unit and nested factories.
+ # This narrow metadata exception never constructs Expr, Type or Span.
+ # The hygienic binding cannot collide with source or captured names.
+ self.source_name_binding = self.fresh() if track_span else None
+ if track_span:
+ self.env[self.source_name_binding] = SourceName(self.filename)
+
+ def fresh(self, prefix="_t"):
+ """Allocate a generated identifier without changing any source name.
+
+ Parameters
+ ----------
+ prefix : str, optional
+ Identifier prefix. Default is ``"_t"``.
+
+ Returns
+ -------
+ str
+ An unused generated identifier.
+
+ Notes
+ -----
+ Updates this compilation's shared ``name_map``. Reserved names have
+ value zero; prefix entries hold the next counter. The allocator remains
+ shared across every function in the source unit.
+ """
+ counter = self.name_map.get(prefix, 0)
+ while f"{prefix}{counter}" in self.name_map:
+ counter += 1
+ name = f"{prefix}{counter}"
+ self.name_map[prefix] = counter + 1
+ self.name_map[name] = 0
+ return name
+
+ def span_ast(self, node):
+ """Emit location data for builders to materialize during execution.
+
+ Parameters
+ ----------
+ node : ast.AST
+ Original node carrying line, end-line, and UTF-8 column ranges.
+
+ Returns
+ -------
+ ast.expr
+ Tuple expression referencing the unit's shared SourceName binding,
+ or a None constant when span tracking is disabled.
+
+ Notes
+ -----
+ The generated expression inherits the full source range through
+ ``ast.copy_location``. This method constructs no IR Span object.
+ """
+ if not self.track_span:
+ return ast.copy_location(ast.Constant(None), node)
+ return ast.copy_location(
+ ast.Tuple(
+ [
+ ast.Name(self.source_name_binding, ast.Load()),
+ *[
+ ast.Constant(value)
+ for value in (
+ node.lineno,
+ node.end_lineno,
+ node.col_offset,
+ node.end_col_offset,
+ )
+ ],
+ ],
+ ast.Load(),
+ ),
+ node,
+ )
+
+ def transformer(self, builder_name=None, **options):
+ """Create a syntax transformer sharing this unit's name allocator.
+
+ Parameters
+ ----------
+ builder_name : str, optional
+ Injected builder alias. Default is None, which selects this unit's
+ ``builder_name``.
+ **options
+ Additional `IRBuilderTranspiler` syntax-handler configuration.
+
+ Returns
+ -------
+ IRBuilderTranspiler
+ New transformer with local binding and export analysis.
+
+ Raises
+ ------
+ TypeError
+ If options are unsupported or duplicate internally supplied options.
+
+ Notes
+ -----
+ Only host namespace and callable metadata are exposed to the transformer.
+ The name allocator is shared with this compiler; no IR state is shared.
+ """
+ metadata = {
+ name: value
+ for name, value in self.env.items()
+ if inspect.ismodule(value)
+ or inspect.isfunction(value)
+ or inspect.isclass(value)
+ or type(value).__module__ == "types"
+ }
+ return IRBuilderTranspiler(
+ self.filename,
+ metadata,
+ builder_name or self.builder_name,
+ self.infrastructure_name,
+ self.span_ast,
+ None,
+ name_map=self.name_map,
+ track_span=self.track_span,
+ **options,
+ )
+
+ def function_kind(self, node, env=None, *, allow_python=False):
+ """Read registered function metadata and unevaluated option syntax.
+
+ Parameters
+ ----------
+ node : ast.FunctionDef
+ Original function definition with construction decorator syntax.
+ env : mapping, optional
+ Unused compatibility argument. Default is None.
+ allow_python : bool, optional
+ Permit unregistered ordinary Python helpers. Default is False.
+
+ Returns
+ -------
+ info : FunctionDecoratorInfo
+ Registered metadata or an ordinary-Python fallback when permitted.
+ options : ast.Dict
+ Unevaluated construction options.
+
+ Raises
+ ------
+ SyntaxError
+ If construction metadata is absent and Python helpers are disallowed,
+ or a construction decorator supplies positional options.
+
+ Notes
+ -----
+ This lookup evaluates neither annotations nor option expressions.
+ """
+ return self.transformer().function_metadata(node, allow_python=allow_python)
+
+ def run_statements(self, body, builder, env, bound_names, *, preserve_return=False):
+ """Compile and execute a helper body through AST-only lowering.
+
+ Parameters
+ ----------
+ body : list of ast.stmt
+ Nonempty original helper body.
+ builder : object
+ Registered construction namespace for the helper.
+ env : mapping of str to object
+ Opaque execution bindings, copied for this invocation.
+ bound_names : iterable of str
+ Names of supplied Python parameters.
+ preserve_return : bool, optional
+ Retain ordinary Python return semantics. Default is False.
+
+ Returns
+ -------
+ object
+ Opaque result of the generated helper.
+
+ Raises
+ ------
+ SyntaxError
+ If the helper body cannot be translated or compiled.
+
+ Notes
+ -----
+ Execution shares the caller's active builder frames and this unit's name
+ allocator. Source nodes are copied before translation. Builder and host
+ execution exceptions propagate unchanged.
+ """
+ namespace = dict(env)
+ namespace.update({self.builder_name: builder, self.infrastructure_name: _EXECUTION})
+ transformer = self.transformer(signature_names=bound_names, preserve_return=preserve_return)
+ statements = transformer.transform_statements(body)
+ names = sorted(name for name in bound_names if name in namespace)
+ helper_name = self.fresh()
+ helper = ast.copy_location(
+ ast.FunctionDef(
+ helper_name,
+ ast.arguments(
+ posonlyargs=[],
+ args=[ast.arg(name) for name in names],
+ kwonlyargs=[],
+ kw_defaults=[],
+ defaults=[],
+ ),
+ statements or [ast.Pass()],
+ [],
+ None,
+ ),
+ body[0],
+ )
+ if "type_params" in ast.FunctionDef._fields:
+ helper.type_params = []
+ module = ast.fix_missing_locations(ast.Module([helper], []))
+ exec(
+ compile(module, self.filename, "exec", flags=self.compile_flags, dont_inherit=True),
+ namespace,
+ )
+ return namespace[helper_name](*(namespace[name] for name in names))
+
+ def build(self):
+ """Compile original-location builder AST and execute it.
+
+ Returns
+ -------
+ object
+ Opaque result produced by the builder program.
+
+ Raises
+ ------
+ SyntaxError
+ If the source cannot be lowered to supported construction syntax
+ or the generated program cannot be compiled.
+
+ Notes
+ -----
+ Imports establish host namespace bindings before metadata lookup. The
+ remaining source expressions execute in the generated builder program;
+ execution exceptions propagate to `parse` for diagnostic conversion.
+
+ Compilation uses the original filename and AST coordinates, never
+ unparse/reparse. Replaced nodes inherit all four location fields, while
+ ``fix_missing_locations`` fills only absent fields. Python exceptions
+ therefore retain the source file and original expression range; column
+ detail requires Python 3.11 or later. Generated names and injected
+ bindings remain local to this compiler and its execution namespace.
+ """
+ # Source-text imports establish host namespace bindings before static
+ # metadata lookup. Only import statements execute here; annotations and
+ # construction expressions remain exclusively in the generated program.
+ imports = [
+ copy.deepcopy(node)
+ for node in self.tree.body[:-1]
+ if isinstance(node, ast.Import | ast.ImportFrom)
+ ]
+ if imports:
+ exec(
+ compile(ast.Module(imports, []), self.filename, "exec", dont_inherit=True), self.env
+ )
+ if (
+ inspect.isfunction(self.original)
+ and syntax_protocol.function_info(self.original) is not None
+ ):
+ # Direct decorator application (T.prim_func(host_function)) has no
+ # decorator in source AST. Inject only its registered host metadata;
+ # already evaluated option values stay opaque execution bindings.
+ decorator_name = self.fresh()
+ self.env[decorator_name] = self.original
+ keywords = []
+ for key, value in getattr(self.original, "__tvm_function_options__", {}).items():
+ option_name = self.fresh()
+ self.env[option_name] = value
+ keywords.append(ast.keyword(key, ast.Name(option_name, ast.Load())))
+ root = self.tree.body[-1]
+ root.decorator_list = [
+ ast.copy_location(
+ ast.Call(ast.Name(decorator_name, ast.Load()), [], keywords), root
+ )
+ ]
+ runtime, original = self.fresh(), self.fresh()
+ bindings = {
+ runtime: construction,
+ original: self.original if inspect.isclass(self.original) else None,
+ self.infrastructure_name: _EXECUTION,
+ }
+ transformed, result = self.transformer().program(self.tree, runtime, original, bindings)
+ namespace = {**self.env, **bindings}
+ exec(
+ compile(
+ transformed, self.filename, "exec", flags=self.compile_flags, dont_inherit=True
+ ),
+ namespace,
+ )
+ return namespace[result]
+
+
+def parse(source, extra_vars=None, *, filename=None, track_span: bool = True, **options):
+ """Transpile and execute a source string, Python function, or Python class.
+
+ Parameters
+ ----------
+ source : str or function or type
+ Original source text or inspectable Python object.
+ extra_vars : mapping of str to object, optional
+ Lexical bindings overriding captured values. Default is None,
+ interpreted as an empty mapping.
+ filename : str, optional
+ Source filename override. Default is None, which uses the inspected
+ filename for objects and ``"<str>"`` for text.
+ track_span : bool, optional
+ Enable shared source metadata and IR location instrumentation.
+ Default is True. False retains Python source locations only.
+ **options
+ Accepted for entry-point compatibility; construction policy comes
+ from registered source decorators, not these options.
+
+ Returns
+ -------
+ object
+ Opaque result of the generated builder program.
+
+ Raises
+ ------
+ DiagnosticError
+ If transpilation or host/builder execution fails. An existing
+ DiagnosticError is preserved; other execution errors gain original
+ source ranges.
+ OSError
+ If source inspection cannot recover the supplied object's text.
+ TypeError
+ If the source object cannot be inspected.
+ SyntaxError
+ If initial source parsing fails before program execution.
+
+ Notes
+ -----
+ Each call owns a fresh compiler and lexical environment. Declaration and
+ definition frames are entered only during generated execution. Source
+ acquisition errors propagate directly, before diagnostic conversion.
+ """
+ env = {} if isinstance(source, str) else _capture(source)
+ env.update(extra_vars or {})
+ compiler = Compiler(source, env, filename, track_span=track_span)
+ try:
+ return compiler.build()
+ except DiagnosticError:
+ raise
+ except Exception as error:
+ raise diagnostic_error(error, compiler) from error
+
+
+def ir_module(module=None, **options):
+ """Decorate a Python class with two-phase module construction.
+
+ Parameters
+ ----------
+ module : type, optional
+ Class to compile immediately. Default is None, which returns a
+ decorator awaiting a class.
+ **options
+ Keyword arguments forwarded to `parse`.
+
+ Returns
+ -------
+ object or callable
+ Generated module result when a class is supplied, otherwise a class
+ decorator.
+
+ Raises
+ ------
+ DiagnosticError
+ If module transpilation or builder execution fails.
+
+ Notes
+ -----
+ Class host bindings are captured before transpilation. Generated execution
+ declares all registered signatures before defining their bodies. Source
+ acquisition errors and frame lifetime follow `parse`.
+ """
+
+ def apply(module):
+ return parse(module, _capture(module), **options)
+
+ return apply(module) if module is not None else apply
+
+
+from_source = parse
diff --git a/python/tvm/script/parser_v2/protocol.py b/python/tvm/script/parser_v2/protocol.py
new file mode 100644
index 0000000000..a67e6caf5e
--- /dev/null
+++ b/python/tvm/script/parser_v2/protocol.py
@@ -0,0 +1,394 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Parser-owned syntax policies and function-decorator registration.
+
+This module stores only host callable identities and flat syntax metadata.
+Concrete eager annotation behavior is delegated to a builder-owned adapter;
+no IR definition, concrete annotation result, symbol or frame is stored here.
+"""
+
+from inspect import signature
+from typing import Any, NamedTuple
+
+
+class ExprStrPolicy(NamedTuple):
+ """Immutable syntax policy for registered constructor arguments.
+
+ Parameters
+ ----------
+ fields : tuple of str
+ Parameter names whose string values represent expressions.
+ introduce : bool, optional
+ Legacy symbol-introduction metadata. Default is False.
+ dtype : object, optional
+ Legacy explicit symbol-type metadata. Default is None.
+ scalar_strings : bool, optional
+ Interpret bare strings as expressions. Default is True. Nested
+ strings in marked fields are always treated as expressions.
+ implicit_dtype : object, optional
+ Legacy implicit symbol-type metadata. Default is None.
+ compound_declarations : bool, optional
+ Legacy compound-expression metadata. Default is False.
+
+ Notes
+ -----
+ This tuple record lives with its registered callable across transpilation
+ passes. It stores no symbols, eager results, or function-local state.
+ Legacy metadata never drives transpiler symbol construction or type
+ validation; builders own those operations. Construction enters no frame.
+ """
+
+ fields: tuple[str, ...]
+ introduce: bool = False
+ dtype: Any = None
+ scalar_strings: bool = True
+ implicit_dtype: Any = None
+ compound_declarations: bool = False
+
+
+# Process-wide registry: callable identity -> immutable syntax policy. Dialect
+# imports register once; aliases share identities. No per-function entries or
+# evaluation results are cached, and transpilers only read this table.
+_EXPR_STR_POLICIES = {}
+
+
+def expr_str_policy(constructor):
+ """Look up the expression-string policy for a host value.
+
+ Parameters
+ ----------
+ constructor : object
+ Python value whose registration or attached policy is queried.
+
+ Returns
+ -------
+ ExprStrPolicy or None
+ Registered or attached policy, or None if unregistered or unhashable.
+ Legacy attached policies may use the equivalent legacy record type.
+
+ Notes
+ -----
+ This read-only lookup neither imports dialects nor calls the constructor.
+ Returned policies are shared registration state; no builder frame is
+ entered. Exceptions from custom attribute access propagate.
+ """
+ try:
+ return _EXPR_STR_POLICIES.get(constructor) or getattr(
+ constructor, "__tvm_expression_args__", None
+ )
+ except TypeError:
+ return None
+
+
+def expr_str_args(
+ *fields,
+ introduce=False,
+ dtype=None,
+ scalar_strings=True,
+ implicit_dtype=None,
+ compound_declarations=False,
+ as_type=False,
+):
+ """Register expression-string fields and wrap eager constructor calls.
+
+ Parameters
+ ----------
+ *fields : str
+ Parameter names whose string values denote source expressions.
+ introduce : bool, optional
+ Legacy syntax metadata. Default is False; builders own introduction.
+ dtype : object, optional
+ Legacy explicit symbol-type metadata. Default is None.
+ scalar_strings : bool, optional
+ Interpret bare strings as expressions. Default is True. False
+ preserves literal shorthand such as ``Tensor("float32")``; nested
+ strings in tuples and lists remain expressions.
+ implicit_dtype : object, optional
+ Legacy implicit-type metadata. Default is None; newly resolved
+ symbols default to int64 in the builder's TypeVarFrame.
+ compound_declarations : bool, optional
+ Legacy compound-expression metadata. Default is False.
+ as_type : bool, optional
+ Preserve use as an annotation class, including Python type unions.
+ Default is False. True returns a class invoking the wrapped callable;
+ construction still returns that callable's result.
+
+ Returns
+ -------
+ decorator : callable
+ Registers the constructor's policy and returns its eager-call adapter.
+
+ Raises
+ ------
+ ValueError
+ When applying the decorator to a constructor whose signature cannot
+ be inspected or does not contain a requested field.
+ TypeError
+ When the constructor cannot be inspected, or a wrapped call cannot
+ bind its signature or has unresolved fields inside an active builder.
+
+ Notes
+ -----
+ Both original and wrapped callable identities retain the policy in the
+ process-wide registry. Builders own eager annotation behavior: unresolved
+ strings or typing.TypeVar values in marked fields produce MissingType
+ outside active builders. Concrete arguments call the original constructor.
+ Handwritten builders require concrete expressions. Calls are never cached;
+ annotations must be safe to re-evaluate.
+
+ Examples
+ --------
+ Register a shape parameter while leaving the dtype string literal intact::
+
+ @expr_str_args("shape", scalar_strings=False)
+ def tensor_type(shape, dtype="float32"):
+ return make_tensor_type(shape, dtype)
+ """
+
+ def decorate(constructor):
+ call_signature = signature(constructor)
+ unknown = set(fields).difference(call_signature.parameters)
+ if unknown:
+ raise ValueError(f"Unknown expression argument fields: {sorted(unknown)}")
+ policy = ExprStrPolicy(
+ tuple(fields),
+ bool(introduce),
+ dtype,
+ bool(scalar_strings),
+ implicit_dtype,
+ bool(compound_declarations),
+ )
+
+ # Registration is syntax-only. Builders own eager construction and the
+ # active-frame/MissingType decisions behind this generic wrapper factory.
+ from tvm.script.ir_builder.type_var_frame import wrap_expression_constructor
+
+ result = wrap_expression_constructor(constructor, call_signature, policy, as_type=as_type)
+ result.__tvm_expression_args__ = policy
+ _EXPR_STR_POLICIES[result] = policy
+ _EXPR_STR_POLICIES[constructor] = policy
+ return result
+
+ return decorate
+
+
+# Compatibility spellings share the registry and wrapper implementation.
+ExpressionArguments = ExprStrPolicy
+expression_args = expr_str_args
+
+
+class DeclarationArguments(NamedTuple):
+ """Immutable metadata for scalar-constructor predeclarations.
+
+ Parameters
+ ----------
+ value_parameter : str
+ Optional value parameter whose absence denotes a declaration.
+ dtype : object, optional
+ Explicit primitive dtype metadata. Default is None.
+
+ Notes
+ -----
+ `register_type_var_decl` attaches this record to its callable. Aliases
+ and transpilation passes share it for the callable's lifetime; it stores
+ no symbols or construction results and enters no builder frame.
+ """
+
+ value_parameter: str
+ dtype: Any = None
+
+
+def register_type_var_decl(constructor, *, value_parameter="expr", dtype=None):
+ """Register a constructor that can declare a type variable.
+
+ Parameters
+ ----------
+ constructor : callable
+ Scalar constructor supporting attribute assignment.
+ value_parameter : str, optional
+ Optional value parameter. Default is ``"expr"``; omission of this
+ parameter identifies declaration syntax.
+ dtype : object, optional
+ Explicit primitive dtype metadata. Default is None. A string dtype
+ also permits predeclaration of direct zero-argument body calls.
+
+ Returns
+ -------
+ callable
+ The original constructor with attached `DeclarationArguments`.
+
+ Raises
+ ------
+ AttributeError
+ If the constructor cannot store the registration attribute.
+
+ Notes
+ -----
+ The signature prepass reserves explicitly declared parameter types before
+ resolving earlier shape strings. The shared name prescan also recognizes
+ direct unconditional zero-argument body calls when dtype is a string,
+ emitting the spelling as builder predeclaration data. It excludes nested
+ control scopes and argument-bearing or effectful expressions. Original
+ calls and bindings retain their body order.
+
+ Registration persists with the callable. The constructor is neither
+ wrapped nor evaluated, and no builder frame is entered.
+ """
+ constructor.__tvm_type_var_decl__ = DeclarationArguments(value_parameter, dtype)
+ return constructor
+
+
+class FunctionDecoratorInfo(NamedTuple):
+ """Flat syntax registration for a source function decorator.
+
+ Parameters
+ ----------
+ builder : object or None
+ Opaque construction namespace, or None for ordinary Python functions.
+ option_map : dict of str to str, optional
+ Public option names mapped to builder option names. Default is None,
+ interpreted as empty. `register_function` copies supplied mappings.
+ defaults : dict of str to object, optional
+ Default builder keyword values. Default is None, interpreted as
+ empty. `register_function` copies supplied mappings.
+ python : bool, optional
+ Preserve the source body as ordinary Python. Default is False.
+
+ Notes
+ -----
+ These are the complete supported fields; no generic metadata dictionary
+ is retained. Records live with decorators across compilations. Consumers
+ must treat mappings as read-only. Records retain no frame, function
+ result, annotation result, or symbol state and enter no builder frames.
+ """
+
+ builder: Any
+ option_map: dict | None = None
+ defaults: dict | None = None
+ python: bool = False
+
+
+def register_function(decorator, builder, *, option_map=None, defaults=None, python=False):
+ """Register a decorator with explicit supported syntax options.
+
+ Parameters
+ ----------
+ decorator : callable
+ Function decorator supporting attribute assignment.
+ builder : object or None
+ Opaque construction namespace, or None when ``python=True``.
+ option_map : mapping of str to str, optional
+ Public option names mapped to builder keyword names. Default is None,
+ interpreted as an empty mapping. Entries are copied.
+ defaults : mapping of str to object, optional
+ Default builder keyword values. Default is None, interpreted as an
+ empty mapping. Entries are copied.
+ python : bool, optional
+ Preserve the source body as ordinary Python. Default is False.
+
+ Returns
+ -------
+ callable
+ The original decorator with attached `FunctionDecoratorInfo`.
+
+ Raises
+ ------
+ AttributeError
+ If the decorator cannot store the registration attribute.
+ TypeError or ValueError
+ If supplied option mappings cannot be converted to dictionaries.
+
+ Notes
+ -----
+ Registration replaces ``__tvm_function_info__`` for the callable's
+ lifetime. It executes no constructor or annotation and enters no frame.
+ Unknown keyword arguments are rejected by the explicit signature.
+ """
+ decorator.__tvm_function_info__ = FunctionDecoratorInfo(
+ builder, dict(option_map or {}), dict(defaults or {}), bool(python)
+ )
+ return decorator
+
+
+# V2 mappings are separate from legacy metadata; registering opt-in construction
+# never mutates the old parser's decorator or namespace objects.
+_LEGACY_BUILDERS = {}
+
+
+def register_builder(legacy_builder, builder):
+ """Register a v2 construction implementation for a legacy namespace.
+
+ Parameters
+ ----------
+ legacy_builder : object
+ Hashable namespace used by existing decorator metadata.
+ builder : object
+ Namespace supplying v2 construction operations.
+
+ Returns
+ -------
+ None
+
+ Raises
+ ------
+ TypeError
+ If the legacy namespace cannot be used as a dictionary key.
+
+ Notes
+ -----
+ The process-wide v2 registry retains this mapping across compilations.
+ Registration replaces only its v2 entry; legacy decorator and namespace
+ objects are unchanged. It enters no builder frame.
+ """
+ _LEGACY_BUILDERS[legacy_builder] = builder
+
+
+def function_info(decorator):
+ """Look up v2 metadata for a resolved function decorator.
+
+ Parameters
+ ----------
+ decorator : object
+ Already-resolved host value to inspect for registration metadata.
+
+ Returns
+ -------
+ FunctionDecoratorInfo or None
+ Existing v2 metadata, translated legacy metadata, or None if neither
+ registration exists. A translated record uses the v2 builder mapping
+ when one is registered for the legacy builder.
+
+ Notes
+ -----
+ This lookup neither calls the decorator nor enters a builder frame.
+ Custom attribute-access errors propagate. Returned option mappings are
+ shared registration state and must not be edited. Legacy metadata is
+ adapted without mutating its decorator or namespace.
+ """
+ info = getattr(decorator, "__tvm_function_info__", None)
+ if info is not None:
+ return info
+ legacy = getattr(decorator, "__tvm_function_kind__", None)
+ if legacy is None:
+ return None
+ metadata = legacy.metadata
+ return FunctionDecoratorInfo(
+ _LEGACY_BUILDERS.get(legacy.builder, legacy.builder),
+ metadata.get("option_map"),
+ metadata.get("defaults"),
+ metadata.get("python", False),
+ )
diff --git a/python/tvm/script/parser_v2/transpile.py b/python/tvm/script/parser_v2/transpile.py
new file mode 100644
index 0000000000..320195a4ad
--- /dev/null
+++ b/python/tvm/script/parser_v2/transpile.py
@@ -0,0 +1,1506 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Translate Python syntax into calls on a context-selected construction namespace.
+
+The translator owns ordering, lexical scopes, and original source locations.
+Construction operations and callable syntax metadata come from its environment;
+ordinary expressions remain nested Python expressions producing concrete values.
+"""
+
+import ast
+import builtins
+import copy
+import inspect
+
+from . import protocol
+from .expression import rewrite_expression
+
+
+class NameCollector(ast.NodeVisitor):
+ """Reserve identifiers and collect declaration syntax for one source unit.
+
+ Parameters
+ ----------
+ names : dict of str to int
+ Shared name-to-counter mapping, populated in place. The enclosing
+ compilation owns its lifetime.
+
+ Notes
+ -----
+ Visits reserve function, class, import, parameter, and expression-string
+ names before generated bindings are allocated. Source names are unchanged.
+
+ Each visited FunctionDef receives ``_tvm_type_var_declarations``: ordered
+ ``(name, expression AST, location AST, body-assignment flag)`` tuples.
+ Registered constructor metadata filters these syntax candidates during
+ emission. Signature candidates precede direct unconditional body candidates;
+ nested functions receive independent lists. Collection executes no source
+ expression and retains no IR construction state.
+ """
+
+ def __init__(self, names):
+ self.names = names
+
+ def visit_Name(self, node):
+ self.names.setdefault(node.id, 0)
+
+ def visit_arg(self, node):
+ self.names.setdefault(node.arg, 0)
+ self.generic_visit(node)
+
+ def visit_FunctionDef(self, node):
+ # Reserve definition names even if no ast.Name refers to them.
+ self.names.setdefault(node.name, 0)
+ node._tvm_type_var_declarations = [
+ (parameter.arg, parameter.annotation, parameter, False)
+ for parameter in [*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs]
+ if parameter.annotation is not None
+ ]
+ # Only direct statements belong to this function's declaration prescan.
+ # Do not descend through conditions, loops, with scopes or nested defs.
+ for statement in node.body:
+ if isinstance(statement, ast.Assign):
+ for target in statement.targets:
+ node._tvm_type_var_declarations.extend(
+ self._body_declarations(target, statement.value)
+ )
+ elif isinstance(statement, ast.AnnAssign) and statement.value is not None:
+ node._tvm_type_var_declarations.extend(
+ self._body_declarations(statement.target, statement.value)
+ )
+ elif isinstance(statement, ast.Return | ast.Raise):
+ break
+ self.generic_visit(node)
+
+ @staticmethod
+ def _body_declarations(target, value):
+ # Pattern: n = X.int64() or m, n = X.int64(), X.int64(). No argument
+ # expressions, starred unpacking or ordinary value-producing calls qualify.
+ if isinstance(target, ast.Name) and isinstance(value, ast.Call):
+ if not value.args and not value.keywords:
+ yield target.id, value, target, True
+ elif isinstance(target, ast.Tuple | ast.List) and isinstance(value, ast.Tuple | ast.List):
+ if len(target.elts) == len(value.elts):
+ for lhs, rhs in zip(target.elts, value.elts):
+ yield from NameCollector._body_declarations(lhs, rhs)
+
+ visit_AsyncFunctionDef = visit_FunctionDef
+
+ def visit_ClassDef(self, node):
+ self.names.setdefault(node.name, 0)
+ self.generic_visit(node)
+
+ def visit_alias(self, node):
+ self.names.setdefault(node.asname or node.name.split(".")[0], 0)
+
+ def visit_Constant(self, node):
+ if isinstance(node.value, str):
+ try:
+ expression = ast.parse(node.value, mode="eval")
+ except SyntaxError:
+ return
+ for inner in ast.walk(expression):
+ if isinstance(inner, ast.Name):
+ self.names.setdefault(inner.id, 0)
+
+
+class IRBuilderTranspiler(ast.NodeTransformer):
+ """Match original Python patterns and emit a builder-program AST.
+
+ Parameters
+ ----------
+ filename : str
+ Original source filename, used in Python compilation and diagnostics.
+ environment : dict of str to object
+ Host namespace and callable bindings for static metadata lookup.
+ No expression evaluation or concrete symbol construction occurs here.
+ builder_name : str
+ Collision-free injected alias for the current construction namespace.
+ infrastructure_name : str
+ Collision-free injected alias for shared execution helpers.
+ span : callable
+ Map an original AST node to an expression containing location data.
+ fresh : callable or None
+ Whole-unit name allocator accepting an optional prefix. None selects
+ `fresh_unique_name`. Source identifiers are never renamed.
+ signature_names : iterable of str, optional
+ Initially bound parameter names. Default is None, interpreted as empty.
+ nested_function : callable, optional
+ Syntax-only nested-function handler. Default is None. Whole-program
+ lowering temporarily installs its declaration/definition emitter.
+ preserve_return : bool, optional
+ Retain Python return semantics in construction helpers. Default is False.
+ name_map : dict of str to int, optional
+ Shared name reservations and counters. Default is None, which initializes
+ the mapping from environment names. `NameCollector` reserves source names
+ before whole-program lowering.
+ track_span : bool, optional
+ Emit IR source instrumentation. Default is True. False retains AST ranges
+ without expression ``at`` calls or builder ``span`` keywords.
+
+ Notes
+ -----
+ ``filename``, the span and name callbacks, and ``infrastructure_name`` belong
+ to one source unit. ``namespace_bindings`` maps visible source spellings to
+ host namespaces and callables; assignments and parameters shadow entries.
+ ``dialect_prefix`` selects the builder alias and is restored after functions.
+ ``bound`` holds identifier strings only. ``optional`` maps conditionally
+ exported names to AST reads of cached named outputs, so a name defined by
+ only one branch is checked when read.
+
+ Each lexical helper copies and restores ``bound``, ``namespace_bindings``,
+ and ``optional``. No concrete value or IR identity is stored. Return policy
+ and syntax hooks configure one body pass; nested scopes share the whole-unit
+ allocator without resetting it. Unsupported syntax raises SyntaxError during
+ translation rather than construction.
+
+ Examples
+ --------
+ Representative source patterns lower to builder calls::
+
+ v = e # v = X.bind_(e, name="v", span=location)
+ v: ty = e # X.bind_ also receives ty
+ return e # X.return_(e)
+ for i in range(n): ... # with X.for_(X.range_(n)) as i: ...
+
+ An annotation-only binding uses the builder MISSING initializer; a bare
+ return supplies no argument. Grid iteration retains its grid expression and
+ tuple target. Function factories declare signatures before invoking body
+ definitions; builders retain parameter and symbol identities. Marked strings
+ become ``resolve_type_var`` calls, while unmarked dtype strings remain
+ literal. No rendered Python is reparsed to recover locations.
+ """
+
+ def __init__(
+ self,
+ filename,
+ environment,
+ builder_name,
+ infrastructure_name,
+ span,
+ fresh,
+ signature_names=None,
+ nested_function=None,
+ preserve_return=False,
+ name_map=None,
+ track_span=True,
+ ):
+ self.filename = filename
+ self.track_span = track_span
+ self.namespace_bindings = dict(environment)
+ self.dialect_prefix = builder_name
+ self.infrastructure_name = infrastructure_name
+ self.span = span
+ self.fresh = fresh or self.fresh_unique_name
+ self.name_map = name_map if name_map is not None else dict.fromkeys(environment, 0)
+ self.bound = set(signature_names or ())
+ self.optional = {}
+ self.nested_function = nested_function
+ self.preserve_return = preserve_return
+
+ def fresh_unique_name(self, prefix="_t"):
+ """Allocate an unused identifier from the shared whole-unit name map.
+
+ Parameters
+ ----------
+ prefix : str, optional
+ Generated identifier prefix. Default is ``"_t"``.
+
+ Returns
+ -------
+ str
+ An identifier absent from the current reservations.
+
+ Notes
+ -----
+ Updates ``name_map`` with the reservation and next prefix counter. All
+ nested functions share this dictionary; no source name is renamed and
+ no lexical scope resets the allocator. No builder frame is entered.
+ """
+ counter = self.name_map.get(prefix, 0)
+ while f"{prefix}{counter}" in self.name_map:
+ counter += 1
+ name = f"{prefix}{counter}"
+ self.name_map[prefix] = counter + 1
+ self.name_map[name] = 0
+ return name
+
+ def transform_statements(self, body):
+ """Translate a sequence of original statements into builder syntax.
+
+ Parameters
+ ----------
+ body : sequence of ast.stmt
+ Original statements, deep-copied before visiting.
+
+ Returns
+ -------
+ list of ast.stmt
+ Lowered statements in source order. A source statement can expand
+ into several generated statements.
+
+ Raises
+ ------
+ SyntaxError
+ If a statement or expression uses unsupported construction syntax.
+
+ Notes
+ -----
+ Source nodes are unchanged. Syntactic name state is updated for later
+ statements in this body. Builder operations receive statement locations
+ through ``span=``; expression locations use value-wrapping ``at`` calls.
+ Only syntax callbacks run here; no construction frame is entered.
+ """
+ result = []
+ for statement in copy.deepcopy(body):
+ translated = self.visit(statement)
+ if translated is not None:
+ block = translated if isinstance(translated, list) else [translated]
+ result.extend(block)
+ return result
+
+ def _error(self, node, message):
+ raise SyntaxError(message, (self.filename, node.lineno, node.col_offset + 1, None))
+
+ @staticmethod
+ def _located(value, original):
+ return ast.copy_location(value, original)
+
+ def _name(self, name, original, store=False):
+ return self._located(ast.Name(name, ast.Store() if store else ast.Load()), original)
+
+ def _attribute(self, namespace, member, original):
+ return self._located(
+ ast.Attribute(self._name(namespace, original), member, ast.Load()), original
+ )
+
+ def _call(self, namespace, member, args, original, **keywords):
+ # Disabled tracking removes instrumentation, while copy_location below
+ # preserves Python diagnostics for the same generated computation.
+ if not self.track_span:
+ if namespace == self.infrastructure_name and member == "_at":
+ return args[1]
+ keywords.pop("span", None)
+ keywords.pop("name_span", None)
+ return self._located(
+ ast.Call(
+ self._attribute(namespace, member, original),
+ args,
+ [ast.keyword(arg=key, value=value) for key, value in keywords.items()],
+ ),
+ original,
+ )
+
+ def _operation(self, member, args, original, **keywords):
+ return self._call(
+ self.dialect_prefix, member, args, original, span=self.span(original), **keywords
+ )
+
+ def _statement(self, expression, original):
+ return self._located(ast.Expr(expression), original)
+
+ def _assign(self, name, value, original):
+ return self._located(ast.Assign([self._name(name, original, True)], value), original)
+
+ def _cache(self, value, original, prefix="value"):
+ name = self.fresh(prefix)
+ return self._assign(name, value, original), self._name(name, original)
+
+ def _resolve(self, node):
+ if isinstance(node, ast.Name):
+ return self.namespace_bindings.get(node.id, getattr(builtins, node.id, None))
+ if isinstance(node, ast.Attribute):
+ owner = self._resolve(node.value)
+ if owner is None:
+ return None
+ value = inspect.getattr_static(owner, node.attr, None)
+ if isinstance(value, staticmethod):
+ return value.__func__
+ if inspect.isfunction(value):
+ return (
+ value
+ if inspect.ismodule(owner) or inspect.isclass(owner)
+ else value.__get__(owner)
+ )
+ # Inspect syntax metadata without executing a source-level descriptor.
+ if hasattr(type(value), "__get__"):
+ return None
+ return value
+ return None
+
+ def _expression(self, original, *, attach_span=True):
+ node = copy.deepcopy(original)
+ if isinstance(getattr(node, "ctx", None), ast.Store):
+ return node
+ if isinstance(node, ast.Name) and node.id in self.optional:
+ value = self._call(
+ self.infrastructure_name,
+ "require_defined",
+ [copy.deepcopy(self.optional[node.id]), ast.Constant(node.id)],
+ node,
+ )
+ return self._call(
+ self.infrastructure_name, "_at", [self.span(original), value], original
+ )
+ node = rewrite_expression(node, self._resolve, self.dialect_prefix, self.filename)
+ if isinstance(node, ast.Await | ast.Yield | ast.YieldFrom) or (
+ isinstance(node, ast.NamedExpr) and not getattr(node, "_tvm_signature_binding", False)
+ ):
+ self._error(original, f"Unsupported expression: {type(node).__name__}")
+ # Pattern: not e -> X.not_(e). These operations construct IR eagerly;
+ # the builders, not Python truth testing, own compiled short-circuiting.
+ if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
+ node = self._call(self.dialect_prefix, "not_", [self._expression(node.operand)], node)
+ # Pattern: a and/or b -> X.and_(a, b) / X.or_(a, b). Python evaluates
+ # every operand left-to-right while building; compiled IR is conditional.
+ elif isinstance(node, ast.BoolOp):
+ method = "and_" if isinstance(node.op, ast.And) else "or_"
+ node = self._call(
+ self.dialect_prefix,
+ method,
+ [self._expression(value) for value in node.values],
+ node,
+ )
+ # Pattern: a if c else b -> X.if_then_else_(c, a, b). Both arms are
+ # constructed now; only the selected arm executes in compiled code.
+ elif isinstance(node, ast.IfExp):
+ node = self._call(
+ self.dialect_prefix,
+ "if_then_else_",
+ [
+ self._expression(node.test),
+ self._expression(node.body),
+ self._expression(node.orelse),
+ ],
+ node,
+ )
+ # Pattern: a < b < c -> X.and_(a < b, b < c, chain=(a, b, c)).
+ elif isinstance(node, ast.Compare) and len(node.ops) > 1:
+ node = self._compare_chain(node)
+ # Pattern: f"text{e}" -> preserve fragments and translate only expressions.
+ elif isinstance(node, ast.JoinedStr):
+ # JoinedStr's children must remain literal fragments/FormattedValue nodes.
+ for child in node.values:
+ if isinstance(child, ast.FormattedValue):
+ child.value = self._expression(child.value)
+ if child.format_spec is not None:
+ child.format_spec = self._format_spec(child.format_spec)
+ else:
+ self._expression_children(node)
+ # Pattern: f(args, keyword=value) remains an ordinary Python call.
+ # Recursive argument translation above preserves order and locations;
+ # callable overloads own concrete IR argument/result semantics.
+ if isinstance(node, ast.Starred) or (
+ isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store)
+ ):
+ return node
+ if not attach_span:
+ return node
+ return self._call(self.infrastructure_name, "_at", [self.span(original), node], original)
+
+ def _lambda(self, names, value, original):
+ arguments = ast.arguments(
+ posonlyargs=[],
+ args=[ast.arg(arg=name) for name in names],
+ vararg=None,
+ kwonlyargs=[],
+ kw_defaults=[],
+ kwarg=None,
+ defaults=[],
+ )
+ return self._located(ast.Lambda(arguments, value), original)
+
+ def _compare_chain(self, node):
+ # An immediately invoked Python lambda binds eagerly evaluated operands
+ # once in source order. It passes ordinary overloaded comparisons to the
+ # builder, never lazy callbacks. chain=(a, b, c) records syntax provenance so
+ # builders also bind shared middle operands once at compiled runtime,
+ # where duplicating an effectful expression would change semantics.
+ operands = [node.left, *node.comparators]
+ names = [self.fresh("operand") for _ in operands]
+ comparisons = [
+ self._located(
+ ast.Compare(self._name(left, node), [operation], [self._name(right, node)]), node
+ )
+ for left, operation, right in zip(names, node.ops, names[1:])
+ ]
+ value = self._call(
+ self.dialect_prefix,
+ "and_",
+ comparisons,
+ node,
+ chain=ast.Tuple([self._name(name, node) for name in names], ast.Load()),
+ )
+ return self._located(
+ ast.Call(
+ self._lambda(names, value, node),
+ [self._expression(value) for value in operands],
+ [],
+ ),
+ node,
+ )
+
+ def _format_spec(self, node):
+ for child in node.values:
+ if isinstance(child, ast.FormattedValue):
+ child.value = self._expression(child.value)
+ if child.format_spec is not None:
+ child.format_spec = self._format_spec(child.format_spec)
+ return node
+
+ def _expression_children(self, node):
+ for field, value in ast.iter_fields(node):
+ if isinstance(value, ast.expr):
+ setattr(node, field, self._expression(value))
+ elif isinstance(value, list):
+ for index, item in enumerate(value):
+ if isinstance(item, ast.expr):
+ value[index] = self._expression(item)
+ elif isinstance(item, ast.AST):
+ self._expression_children(item)
+ elif isinstance(value, ast.AST):
+ self._expression_children(value)
+
+ def _index(self, node):
+ if isinstance(node, ast.Slice):
+ fields = [
+ self._expression(value) if value is not None else ast.Constant(None)
+ for value in (node.lower, node.upper, node.step)
+ ]
+ return self._call(self.infrastructure_name, "slice", fields, node)
+ if isinstance(node, ast.Tuple):
+ return self._located(
+ ast.Tuple([self._index(value) for value in node.elts], ast.Load()), node
+ )
+ return self._expression(node)
+
+ def _bind(self, target, value, statement, ty=None, frame_value=False):
+ # Pattern: name = value -> X.bind_; Python name analysis tracks previous values.
+ if isinstance(target, ast.Name):
+ keywords = {"name": ast.Constant(target.id), "name_span": self.span(target)}
+ if frame_value:
+ keywords["frame_value"] = ast.Constant(True)
+ if ty is not None:
+ keywords["ty"] = ty
+ if target.id in self.bound:
+ keywords["previous"] = self._name(target.id, target)
+ elif target.id in self.optional:
+ keywords["previous"] = copy.deepcopy(self.optional[target.id])
+ self.bound.add(target.id)
+ self.optional.pop(target.id, None)
+ return [
+ self._assign(
+ target.id, self._operation("bind_", [value], statement, **keywords), target
+ )
+ ]
+ # Pattern: object.field = value -> X.setattr(object, "field", value).
+ if isinstance(target, ast.Attribute):
+ return [
+ self._statement(
+ self._operation(
+ "setattr",
+ [self._expression(target.value), ast.Constant(target.attr), value],
+ statement,
+ ),
+ statement,
+ )
+ ]
+ # Pattern: object[index] = value -> X.setitem(object, index, value).
+ if isinstance(target, ast.Subscript):
+ return [
+ self._statement(
+ self._operation(
+ "setitem",
+ [self._expression(target.value), self._index(target.slice), value],
+ statement,
+ ),
+ statement,
+ )
+ ]
+ if isinstance(target, ast.Tuple | ast.List):
+ # Each Python unpack finishes before visiting that level's targets. A nested
+ # unpack occurs only when reached, preserving assignment and failure order.
+ names = [self.fresh("unpack") for _ in target.elts]
+ pattern = []
+ for item, name in zip(target.elts, names):
+ temporary = self._name(name, item, True)
+ pattern.append(
+ self._located(ast.Starred(temporary, ast.Store()), item)
+ if isinstance(item, ast.Starred)
+ else temporary
+ )
+ unpack = self._call(self.dialect_prefix, "unpack", [value], target)
+ assignment = self._located(
+ ast.Assign([ast.Tuple(pattern, ast.Store())], unpack), target
+ )
+ result = [assignment]
+ for item, name in zip(target.elts, names):
+ item = item.value if isinstance(item, ast.Starred) else item
+ result.extend(
+ self._bind(
+ item,
+ self._name(name, item),
+ statement,
+ frame_value=frame_value,
+ )
+ )
+ return result
+ self._error(target, f"Unsupported assignment target: {type(target).__name__}")
+
+ def visit_Assign(self, node):
+ # Pattern: var = value -> var = X.bind_(value, name="var", span=...).
+ # Cache the RHS once before chained stores, index loads and nested unpacking.
+ cache, value = self._cache(self._expression(node.value), node.value)
+ result = [cache]
+ resolved = self._resolve(node.value)
+ for target in node.targets:
+ result.extend(self._bind(target, copy.deepcopy(value), node))
+ if isinstance(target, ast.Name):
+ if resolved is not None:
+ self.namespace_bindings[target.id] = resolved
+ else:
+ self.namespace_bindings.pop(target.id, None)
+ return result
+
+ def visit_AnnAssign(self, node):
+ # Pattern: var: ty -> var = X.bind_(I.MISSING, ty=ty, name="var", ...).
+ # Pattern: var: ty = value -> var = X.bind_(value, ty=ty, name="var", ...).
+ if not isinstance(node.target, ast.Name):
+ self._error(node.target, "An annotated binding requires a name")
+ result = []
+ if node.value is None:
+ value = self._attribute(self.infrastructure_name, "MISSING", node)
+ else:
+ cache, value = self._cache(self._expression(node.value), node.value)
+ result.append(cache)
+ annotation = rewrite_expression(
+ node.annotation, self._resolve, self.dialect_prefix, self.filename, annotation=True
+ )
+ result.extend(self._bind(node.target, value, node, self._expression(annotation)))
+ return result
+
+ def visit_AugAssign(self, node):
+ # Pattern: target += value -> load target once, then X.bind_/setitem/setattr.
+ result = []
+ if isinstance(node.target, ast.Name):
+ old, previous = self._cache(
+ self._expression(self._located(ast.Name(node.target.id, ast.Load()), node.target)),
+ node.target,
+ "old",
+ )
+ result.append(old)
+ value = self._located(ast.BinOp(previous, node.op, self._expression(node.value)), node)
+ value = self._call(self.infrastructure_name, "_at", [self.span(node), value], node)
+ return result + self._bind(node.target, value, node)
+ if not isinstance(node.target, ast.Subscript | ast.Attribute):
+ self._error(node.target, "An augmented assignment requires a name, attribute, or index")
+ base_stmt, base = self._cache(
+ self._expression(node.target.value), node.target.value, "base"
+ )
+ result.append(base_stmt)
+ if isinstance(node.target, ast.Attribute):
+ key, operation = ast.Constant(node.target.attr), "setattr"
+ load = self._located(
+ ast.Attribute(copy.deepcopy(base), node.target.attr, ast.Load()), node.target
+ )
+ else:
+ key_stmt, key = self._cache(self._index(node.target.slice), node.target.slice, "key")
+ result.append(key_stmt)
+ operation = "setitem"
+ load = self._located(
+ ast.Subscript(copy.deepcopy(base), copy.deepcopy(key), ast.Load()), node.target
+ )
+ old_stmt, old = self._cache(
+ self._call(
+ self.infrastructure_name, "_at", [self.span(node.target), load], node.target
+ ),
+ node.target,
+ "old",
+ )
+ result.append(old_stmt)
+ value = self._located(ast.BinOp(old, node.op, self._expression(node.value)), node)
+ value = self._call(self.infrastructure_name, "_at", [self.span(node), value], node)
+ result.append(self._statement(self._operation(operation, [base, key, value], node), node))
+ return result
+
+ def visit_Expr(self, node):
+ # Pattern: expression -> X.emit_(translated_expression).
+ return self._statement(self._operation("emit_", [self._expression(node.value)], node), node)
+
+ def visit_Return(self, node):
+ # Pattern: return value -> X.return_(value); return -> X.return_().
+ value = self._expression(node.value) if node.value is not None else None
+ if self.preserve_return:
+ return self._located(ast.Return(value), node)
+ return self._statement(
+ self._operation("return_", [] if value is None else [value], node), node
+ )
+
+ def visit_Break(self, node):
+ # Pattern: break -> X.break_().
+ return self._statement(self._operation("break_", [], node), node)
+
+ def visit_Continue(self, node):
+ # Pattern: continue -> X.continue_().
+ return self._statement(self._operation("continue_", [], node), node)
+
+ def visit_Assert(self, node):
+ # Pattern: assert condition, message -> X.assert_(condition, message).
+ message = self._expression(node.msg) if node.msg is not None else ast.Constant("")
+ return self._statement(
+ self._operation("assert_", [self._expression(node.test), message], node), node
+ )
+
+ @staticmethod
+ def _assigned_names(body):
+ names = set()
+
+ class Names(ast.NodeVisitor):
+ def visit_Name(self, node):
+ if isinstance(node.ctx, ast.Store):
+ names.add(node.id)
+
+ def visit_FunctionDef(self, node):
+ # Nested definitions introduce a name, not their body bindings.
+ names.add(node.name)
+
+ visit_AsyncFunctionDef = visit_FunctionDef
+
+ def visit_Lambda(self, node):
+ pass
+
+ visitor = Names()
+ for statement in body:
+ visitor.visit(statement)
+ return names
+
+ def _scope(self, body, original, prefix, initial=None, return_bindings=False):
+ outer_bound, outer_environment, outer_optional = (
+ self.bound,
+ self.namespace_bindings,
+ self.optional,
+ )
+ referenced = {
+ node.id
+ for statement in body
+ for node in ast.walk(statement)
+ if isinstance(node, ast.Name)
+ }
+ captures = sorted((outer_bound | outer_optional.keys()).intersection(referenced))
+ self.bound, self.namespace_bindings = set(outer_bound), dict(outer_environment)
+ self.optional = {
+ name: self._name(name, original)
+ if name in captures and not self.preserve_return
+ else value
+ for name, value in outer_optional.items()
+ }
+ prefix_statements = [] if initial is None else initial()
+ translated = prefix_statements + self.transform_statements(body)
+ if return_bindings and not self.preserve_return:
+ names = sorted(self._assigned_names(body))
+ values = [
+ self._name(name, original)
+ if name in self.bound
+ else copy.deepcopy(
+ self.optional.get(
+ name, self._attribute(self.infrastructure_name, "MISSING", original)
+ )
+ )
+ for name in names
+ ]
+ translated.append(
+ self._located(
+ ast.Return(ast.Dict([ast.Constant(name) for name in names], values)), original
+ )
+ )
+ self.bound, self.namespace_bindings, self.optional = (
+ outer_bound,
+ outer_environment,
+ outer_optional,
+ )
+ if self.preserve_return:
+ return translated or [self._located(ast.Pass(), original)]
+ # A helper isolates construction locals. Optional exports are captured as
+ # values or MISSING, and checked only when the original body reads them.
+ defaults = [
+ self._name(name, original)
+ if name in outer_bound
+ else copy.deepcopy(outer_optional[name])
+ for name in captures
+ ]
+ helper = self.fresh(prefix)
+ arguments = ast.arguments(
+ posonlyargs=[],
+ args=[ast.arg(arg=name) for name in captures],
+ vararg=None,
+ kwonlyargs=[],
+ kw_defaults=[],
+ kwarg=None,
+ defaults=defaults,
+ )
+ definition = self._located(
+ ast.FunctionDef(
+ helper, arguments, translated or [self._located(ast.Pass(), original)], [], None
+ ),
+ original,
+ )
+ if "type_params" in ast.FunctionDef._fields:
+ definition.type_params = []
+ invocation = self._located(ast.Call(self._name(helper, original), [], []), original)
+ return [definition, self._statement(invocation, original)]
+
+ def _exports(self, frame, candidates, original):
+ # Pattern: names assigned in a completed region -> one named builder
+ # request per candidate, e.g. y = X.frame_result(frame, "y"). Builders
+ # decide which names are designated outputs; absent results remain
+ # MISSING so reads cannot expose branch-local bindings accidentally.
+ result = []
+ for name in sorted(candidates):
+ assignment, value = self._cache(
+ self._call(
+ self.dialect_prefix,
+ "frame_result",
+ [self._name(frame, original), ast.Constant(name)],
+ original,
+ ),
+ original,
+ "output",
+ )
+ result.append(assignment)
+ condition = self._located(
+ ast.Compare(
+ copy.deepcopy(value),
+ [ast.IsNot()],
+ [self._attribute(self.infrastructure_name, "MISSING", original)],
+ ),
+ original,
+ )
+ result.append(
+ self._located(
+ ast.If(condition, [self._assign(name, copy.deepcopy(value), original)], []),
+ original,
+ )
+ )
+ if name not in self.bound:
+ self.optional[name] = copy.deepcopy(value)
+ return result
+
+ def _with(self, context, body, original, target=None):
+ return self._located(
+ ast.With([ast.withitem(context, target)], body or [ast.Pass()]), original
+ )
+
+ def visit_If(self, node):
+ # Pattern: if condition: ... else: ... -> host branch or X.If/Then/Else plus named exports.
+ frame = self.fresh("conditional")
+ condition_stmt, condition = self._cache(self._expression(node.test), node.test, "condition")
+ if self.preserve_return:
+ # Construction helpers retain Python return semantics in ordinary branches.
+ then_body = self._scope(node.body, node, "then")
+ else_body = self._scope(node.orelse, node, "else")
+ host = self._located(ast.If(copy.deepcopy(condition), then_body, else_body), node)
+ definitions = []
+ then_call, else_call = then_body, else_body
+ else:
+ then_def, then_invoke = self._scope(node.body, node, "then", return_bindings=True)
+ else_def, else_invoke = self._scope(node.orelse, node, "else", return_bindings=True)
+ definitions = [then_def, else_def]
+ then_call, else_call = [then_invoke], [else_invoke]
+ mapping_name = self.fresh("exports")
+ host = self._located(
+ ast.If(
+ copy.deepcopy(condition),
+ [self._assign(mapping_name, copy.deepcopy(then_invoke.value), node)],
+ [self._assign(mapping_name, copy.deepcopy(else_invoke.value), node)],
+ ),
+ node,
+ )
+ branches = [self._with(self._operation("Then", [], node), then_call, node)]
+ if node.orelse:
+ branches.append(self._with(self._operation("Else", [], node), else_call, node))
+ region = self._with(
+ self._operation("If", [copy.deepcopy(condition)], node),
+ branches,
+ node,
+ self._name(frame, node, True),
+ )
+ native = [region]
+ candidates = self._assigned_names(node.body + node.orelse)
+ if self.preserve_return:
+ exports = self._exports(frame, candidates, node)
+ native.extend(exports)
+ tail = []
+ else:
+ native.append(
+ self._assign(
+ mapping_name,
+ self._name(frame, node),
+ node,
+ )
+ )
+ tail = self._exports(mapping_name, candidates, node)
+ dispatch = self._located(
+ ast.If(
+ self._call(
+ self.infrastructure_name, "is_python_bool", [copy.deepcopy(condition)], node
+ ),
+ [host],
+ native,
+ ),
+ node,
+ )
+ return [condition_stmt, *definitions, dispatch, *tail]
+
+ def visit_For(self, node):
+ # Pattern: for i, j in T.grid(m, n) -> with X.for_(T.grid(m, n)) as (i, j).
+ # Pattern: for i in range(...) -> with X.for_(X.range_(...)) as i.
+ # Only the actual builtin binding is normalized; a shadowing callable stays intact.
+ if node.orelse:
+ self._error(node, "A construction loop does not support an else clause")
+ iterable = copy.deepcopy(node.iter)
+ if isinstance(iterable, ast.Call) and self._resolve(iterable.func) is builtins.range:
+ iterable.func = self._attribute(self.dialect_prefix, "range_", iterable.func)
+ old_bound, old_optional, old_environment = (
+ self.bound,
+ self.optional,
+ self.namespace_bindings,
+ )
+ self.bound, self.optional, self.namespace_bindings = (
+ set(old_bound),
+ dict(old_optional),
+ dict(old_environment),
+ )
+ for target in ast.walk(node.target):
+ if isinstance(target, ast.Name):
+ self.bound.add(target.id)
+ self.optional.pop(target.id, None)
+ self.namespace_bindings.pop(target.id, None)
+ body = self.transform_statements(node.body)
+ self.bound, self.optional, self.namespace_bindings = (
+ old_bound,
+ old_optional,
+ old_environment,
+ )
+ return self._with(
+ self._operation(
+ "for_",
+ [self._expression(iterable)],
+ node,
+ names=ast.Constant(node.target.id)
+ if isinstance(node.target, ast.Name)
+ else ast.Tuple(
+ [
+ ast.Constant(
+ "*" + item.value.id if isinstance(item, ast.Starred) else item.id
+ )
+ for item in node.target.elts
+ ],
+ ast.Load(),
+ ),
+ ),
+ body,
+ node,
+ copy.deepcopy(node.target),
+ )
+
+ def visit_While(self, node):
+ # Pattern: while condition: ... -> with X.While(condition): translated body.
+ if node.orelse:
+ self._error(node, "A construction loop does not support an else clause")
+ frame = self.fresh("loop")
+ region = self._with(
+ self._operation("While", [self._expression(node.test)], node),
+ self._scope(node.body, node, "body"),
+ node,
+ self._name(frame, node, True),
+ )
+ return [region, *self._exports(frame, self._assigned_names(node.body), node)]
+
+ def _bind_entered(self, target, value, original):
+ # Context/iteration targets introduce lexical names; they never reassign
+ # an outer mutable variable merely because its source spelling matches.
+ for item in ast.walk(target):
+ if isinstance(item, ast.Name):
+ self.bound.discard(item.id)
+ self.optional.pop(item.id, None)
+ return self._bind(target, value, original, frame_value=True)
+
+ def visit_With(self, node):
+ # Pattern: with context as target: ... -> builder context plus lexical exports.
+ item = node.items[0]
+ body = node.body
+ if len(node.items) > 1:
+ nested = self._located(ast.With(node.items[1:], node.body), node)
+ body = [nested]
+ manager, value = self.fresh("context"), self.fresh("entered")
+ cache = self._assign(manager, self._expression(item.context_expr), item.context_expr)
+ initial = (
+ None
+ if item.optional_vars is None
+ else lambda: self._bind_entered(
+ item.optional_vars, self._name(value, item.context_expr), node
+ )
+ )
+ region = self._with(
+ self._name(manager, node),
+ self._scope(body, node, "scope", initial),
+ node,
+ self._name(value, node, True),
+ )
+ return [cache, region, *self._exports(manager, self._assigned_names(body), node)]
+
+ def visit_FunctionDef(self, node):
+ # Pattern: nested def -> ordinary Python definition or complete registered builder program.
+ self.bound.add(node.name)
+ if self.nested_function is not None:
+ return self.nested_function(node)
+ # Ordinary construction helpers retain normal Python execution. Registered
+ # function-kind entry points are resolved by the enclosing compiler callback.
+ if any(
+ getattr(self._resolve(decorator), "__tvm_function_info__", None)
+ for decorator in node.decorator_list
+ ):
+ self._error(node, "A registered nested function requires a function compiler")
+ return node
+
+ def visit_Pass(self, node):
+ # Pattern: pass -> pass; it has no builder effect.
+ return node
+
+ def generic_visit(self, node):
+ if isinstance(node, ast.stmt):
+ self._error(node, f"Unsupported statement: {type(node).__name__}")
+ return super().generic_visit(node)
+
+ def function_metadata(self, node, *, allow_python=False):
+ """Read decorator metadata without evaluating source annotations.
+
+ Parameters
+ ----------
+ node : ast.FunctionDef
+ Original function definition with decorator syntax.
+ allow_python : bool, optional
+ Permit unregistered ordinary Python helpers. Default is False.
+
+ Returns
+ -------
+ info : FunctionDecoratorInfo
+ Registered metadata or ordinary-Python metadata when permitted.
+ options : ast.Dict
+ Unevaluated builder options after keyword-name mapping and defaults.
+
+ Raises
+ ------
+ SyntaxError
+ If a construction decorator uses positional options, or registration
+ is missing while ordinary Python functions are disallowed.
+
+ Notes
+ -----
+ Metadata lookup does not execute annotations or option expressions and
+ enters no builder frame. Returned option syntax belongs to this lowering
+ pass; registered metadata is shared and treated as read-only.
+ """
+ for decorator in node.decorator_list:
+ target = decorator.func if isinstance(decorator, ast.Call) else decorator
+ kind = protocol.function_info(self._resolve(target))
+ if kind is None:
+ continue
+ mapping = kind.option_map or {}
+ values = {key: ast.Constant(value) for key, value in (kind.defaults or {}).items()}
+ expansions = []
+ if isinstance(decorator, ast.Call):
+ if decorator.args:
+ self._error(decorator, "Function decorators accept keyword options only")
+ for keyword in decorator.keywords:
+ if keyword.arg is None:
+ expansions.append(keyword.value)
+ elif keyword.arg != "check_well_formed":
+ values[mapping.get(keyword.arg, keyword.arg)] = copy.deepcopy(keyword.value)
+ result = ast.Dict(
+ [ast.Constant(key) for key in values] + [None] * len(expansions),
+ list(values.values()) + expansions,
+ )
+ return kind, ast.copy_location(result, node)
+ if allow_python:
+ return protocol.FunctionDecoratorInfo(None, python=True), ast.Dict([], [])
+ self._error(node, f"Function {node.name!r} has no registered construction kind")
+
+ def function_program(self, node, runtime, bindings, *, local=False):
+ """Emit a declaration factory and deferred function-definition callback.
+
+ Parameters
+ ----------
+ node : ast.FunctionDef
+ Original function with a registered decorator.
+ runtime : str
+ Generated alias for the builder runtime.
+ bindings : dict of str to object
+ Caller-owned mapping receiving injected builder namespace aliases.
+ local : bool, optional
+ Request a nested-function reference from the builder. Default is False.
+
+ Returns
+ -------
+ tuple of (list of ast.stmt, str, str) or None
+ Factory statements, function-record binding, and body binding. None
+ indicates a registered ordinary Python function.
+
+ Raises
+ ------
+ SyntaxError
+ If signatures, decorators, or body syntax cannot be lowered.
+
+ Notes
+ -----
+ The method allocates generated names and restores body-local syntax state
+ following recursive translation. It neither executes annotations nor
+ creates IR. The generated factory enters a builder declaration frame;
+ its deferred callback defines the body only when invoked.
+ """
+ # Pattern: @X.function def f(a: ty) -> ret: body -> declaration factory + definition body.
+ # A factory gives each function an ordinary Python closure: parameters and
+ # symbols from sibling signatures never overwrite one another's captures.
+ kind, options = self.function_metadata(node)
+ if kind.python:
+ return None
+ builder = self.fresh("_t")
+ bindings[builder] = kind.builder
+ record, factory, body_name, captures = (self.fresh("_t") for _ in range(4))
+ original_builder = self.dialect_prefix
+ self.dialect_prefix = builder
+ record_expr = self._call(
+ runtime,
+ "FunctionRecord",
+ [self._name(builder, node), ast.Constant(node.name), options, self.span(node)],
+ node,
+ local=ast.Constant(local),
+ captures=self._name(captures, node),
+ parameters=ast.Tuple(
+ [ast.Constant(argument.arg) for argument in node.args.args], ast.Load()
+ ),
+ )
+ factory_body = [self._assign(record, record_expr, node)]
+ if node.args.posonlyargs or node.args.kwonlyargs or node.args.vararg or node.args.kwarg:
+ self._error(node, "IR signatures require ordinary named parameters")
+ declaration = []
+ introduced_names = set()
+ for parameter in getattr(node, "type_params", []):
+ # Pattern: def f[n: int](...) -> n = X.resolve_type_var("n")
+ # in this function's fresh symbol frame. Only scalar parameters fit
+ # the registered builder protocol; Python owns the syntax itself.
+ if not isinstance(parameter, getattr(ast, "TypeVar", ())):
+ self._error(parameter, "Only scalar type parameters are supported")
+ bound = getattr(parameter, "bound", None)
+ if bound is not None and self._resolve(bound) is not int:
+ self._error(parameter, "A symbolic type parameter bound must be int")
+ if getattr(parameter, "default_value", None) is not None:
+ self._error(parameter, "A symbolic type parameter cannot have a default")
+ declaration.append(
+ self._assign(
+ parameter.name,
+ self._call(
+ builder, "resolve_type_var", [ast.Constant(parameter.name)], parameter
+ ),
+ parameter,
+ )
+ )
+ # Pattern: f(A: X.Buffer(("n",)), n: X.int64), or a direct body
+ # n = X.int64(): reserve the declared type before signature shapes.
+ # The shared name prescan collects candidates; registered metadata alone
+ # selects declarations here. Body constructors are NOT executed early:
+ # emit their registered dtype spelling and let TypeVarFrame construct it.
+ # The original body call and uniform bind_ remain at their source position.
+ local_names = self._assigned_names(node.body) | {arg.arg for arg in node.args.args}
+ for name, annotation, location, in_body in getattr(node, "_tvm_type_var_declarations", ()):
+ target = annotation.func if isinstance(annotation, ast.Call) else annotation
+ if in_body and any(
+ isinstance(item, ast.Name) and item.id in local_names for item in ast.walk(target)
+ ):
+ continue # A local namespace/callable binding shadows outer metadata.
+ constructor = self._resolve(target)
+ metadata = getattr(constructor, "__tvm_type_var_decl__", None) or getattr(
+ constructor, "__tvm_declaration_args__", None
+ )
+ parameter_dtype = getattr(constructor, "__tvm_parameter_dtype__", None)
+ if in_body:
+ if metadata is None or not isinstance(metadata.dtype, str):
+ continue
+ value = ast.copy_location(ast.Constant(metadata.dtype), annotation)
+ elif metadata is not None or parameter_dtype is not None:
+ value = copy.deepcopy(annotation)
+ else:
+ continue
+ declaration.append(
+ self._statement(
+ self._call(
+ record,
+ "predeclare",
+ [ast.Constant(name), value, self.span(location)],
+ location,
+ ),
+ location,
+ )
+ )
+ for parameter in node.args.args:
+ if parameter.annotation is None:
+ self._error(parameter, f"Parameter {parameter.arg!r} requires an annotation")
+ annotation = rewrite_expression(
+ parameter.annotation, self._resolve, builder, self.filename, annotation=True
+ )
+ introduced_names.update(
+ call.args[0].value
+ for call in ast.walk(annotation)
+ if isinstance(call, ast.Call)
+ and isinstance(call.func, ast.Attribute)
+ and call.func.attr == "resolve_type_var"
+ and call.args
+ and isinstance(call.args[0], ast.Constant)
+ and isinstance(call.args[0].value, str)
+ )
+ declaration.append(
+ self._assign(
+ parameter.arg,
+ self._call(
+ record,
+ "parameter",
+ [
+ ast.Constant(parameter.arg),
+ self._expression(annotation),
+ self.span(parameter),
+ ],
+ parameter,
+ ),
+ parameter,
+ )
+ )
+ if node.returns is not None:
+ annotation = rewrite_expression(
+ node.returns, self._resolve, builder, self.filename, annotation=True
+ )
+ # Return expression strings also assign lexical names. Seed any
+ # enclosing binding before earlier parameter expressions read it.
+ introduced_names.update(
+ call.args[0].value
+ for call in ast.walk(annotation)
+ if isinstance(call, ast.Call)
+ and isinstance(call.func, ast.Attribute)
+ and call.func.attr == "resolve_type_var"
+ and call.args
+ and isinstance(call.args[0], ast.Constant)
+ and isinstance(call.args[0].value, str)
+ )
+ declaration.append(
+ self._statement(
+ self._call(record, "returns", [self._expression(annotation)], node.returns),
+ node.returns,
+ )
+ )
+ for name in sorted(introduced_names - {argument.arg for argument in node.args.args}):
+ # A real enclosing binding remains visible before a signature string
+ # assigns the same spelling; absent names still raise before first use.
+ condition = ast.Compare(ast.Constant(name), [ast.In()], [self._name(captures, node)])
+ assignment = self._assign(
+ name,
+ self._call(
+ record,
+ "capture",
+ [ast.Constant(name), self._call(captures, "get", [ast.Constant(name)], node)],
+ node,
+ ),
+ node,
+ )
+ factory_body.append(self._located(ast.If(condition, [assignment], []), node))
+ factory_body.append(
+ self._with(self._call(record, "declaration", [], node), declaration, node)
+ )
+ factory_body.append(
+ self._assign(node.name, self._attribute(record, "reference", node), node)
+ )
+
+ old_bound, old_optional, old_environment, old_nested = (
+ self.bound,
+ self.optional,
+ self.namespace_bindings,
+ self.nested_function,
+ )
+ self.bound, self.optional = {arg.arg for arg in node.args.args}, {}
+ self.namespace_bindings = dict(old_environment)
+ # Python determines local shadowing for the entire function, including
+ # uses before an assignment. Do not apply namespace/range policies there.
+ for name in self._assigned_names(node.body) | self.bound:
+ self.namespace_bindings.pop(name, None)
+ self.namespace_bindings[name] = None
+
+ def nested(child):
+ nested_kind, _ = self.function_metadata(child, allow_python=True)
+ if nested_kind.python:
+ return copy.deepcopy(child)
+ statements, child_record, child_body = self.function_program(
+ child, runtime, bindings, local=True
+ )
+ return [
+ *statements,
+ self._assign(child.name, self._attribute(child_record, "reference", child), child),
+ self._statement(
+ self._call(child_record, "define", [self._name(child_body, child)], child),
+ child,
+ ),
+ ]
+
+ self.nested_function = nested
+ statements = self.transform_statements(node.body)
+ self.bound, self.optional, self.namespace_bindings, self.nested_function = (
+ old_bound,
+ old_optional,
+ old_environment,
+ old_nested,
+ )
+ arguments = copy.deepcopy(node.args)
+ for argument in arguments.args:
+ argument.annotation = None
+ arguments.defaults = []
+ # Bare host TypeVars in annotations are resolved by constructor builders.
+ # Capture their resulting values for free body names without inspecting
+ # them here. Source locals/parameters keep Python's original shadowing.
+ annotation_names = {
+ item.id
+ for argument in node.args.args
+ if argument.annotation is not None
+ for item in ast.walk(argument.annotation)
+ if isinstance(item, ast.Name) and isinstance(item.ctx, ast.Load)
+ }
+ if node.returns is not None:
+ annotation_names.update(
+ item.id
+ for item in ast.walk(node.returns)
+ if isinstance(item, ast.Name) and isinstance(item.ctx, ast.Load)
+ )
+ # Only syntactically referenced signature names can capture enclosing
+ # symbols. An unrelated outer n must not constrain a local n declaration.
+ capture_names = (annotation_names | introduced_names) - {
+ parameter.name for parameter in getattr(node, "type_params", [])
+ }
+ record_expr.keywords.append(
+ ast.keyword(
+ "capture_names",
+ ast.Tuple([ast.Constant(name) for name in sorted(capture_names)], ast.Load()),
+ )
+ )
+ body_names = {
+ item.id
+ for statement in node.body
+ for item in ast.walk(statement)
+ if isinstance(item, ast.Name) and isinstance(item.ctx, ast.Load)
+ }
+ free_names = annotation_names & body_names - self._assigned_names(node.body) - {
+ arg.arg for arg in node.args.args
+ }
+ for name in sorted(free_names):
+ arguments.kwonlyargs.append(ast.arg(name))
+ arguments.kw_defaults.append(
+ self._call(
+ record,
+ "capture",
+ [ast.Constant(name), self._call(captures, "get", [ast.Constant(name)], node)],
+ node,
+ )
+ )
+ definition = self._located(
+ ast.FunctionDef(body_name, arguments, statements or [ast.Pass()], [], None), node
+ )
+ if "type_params" in ast.FunctionDef._fields:
+ definition.type_params = []
+ factory_body.append(definition)
+ factory_body.append(
+ self._located(
+ ast.Return(
+ ast.Tuple([self._name(record, node), self._name(body_name, node)], ast.Load())
+ ),
+ node,
+ )
+ )
+ factory_def = self._located(
+ ast.FunctionDef(
+ factory,
+ ast.arguments(
+ posonlyargs=[],
+ args=[ast.arg(captures)],
+ kwonlyargs=[],
+ kw_defaults=[],
+ defaults=[self._call(self.infrastructure_name, "locals", [], node)],
+ ),
+ factory_body,
+ [],
+ None,
+ ),
+ node,
+ )
+ if "type_params" in ast.FunctionDef._fields:
+ factory_def.type_params = []
+ unpack = self._located(
+ ast.Assign(
+ [
+ ast.Tuple(
+ [self._name(record, node, True), self._name(body_name, node, True)],
+ ast.Store(),
+ )
+ ],
+ ast.Call(self._name(factory, node), [], []),
+ ),
+ node,
+ )
+ self.dialect_prefix = original_builder
+ return [factory_def, unpack], record, body_name
+
+ def program(self, tree, runtime, original_name, bindings):
+ """Emit a complete module or function builder program from source AST.
+
+ Parameters
+ ----------
+ tree : ast.Module
+ Source module whose final statement is a function or module class.
+ runtime : str
+ Injected builder-runtime alias.
+ original_name : str
+ Injected binding for the original class object, if available.
+ bindings : dict of str to object
+ Caller-owned mapping receiving generated builder namespace aliases.
+
+ Returns
+ -------
+ module : ast.Module
+ Complete generated Python program, ready for compilation.
+ result_name : str
+ Generated binding containing the program's construction result.
+
+ Raises
+ ------
+ SyntaxError
+ If source structure, declarations, or function syntax is unsupported.
+
+ Notes
+ -----
+ Generated names share the enclosing unit's allocator; execution remains
+ separate from translation. Every wrapper inherits its source construct's
+ range, while original descendants retain finer ranges. Missing location
+ fields are filled without rendering and reparsing source code.
+ """
+ nodes = list(tree.body)
+ if not nodes or not isinstance(nodes[-1], ast.FunctionDef | ast.ClassDef):
+ self._error(tree.body[0], "Source must contain one function or module class")
+ root, prefix = nodes[-1], nodes[:-1]
+ is_module = isinstance(root, ast.ClassDef)
+ members = root.body if is_module else [root]
+ functions = [item for item in members if isinstance(item, ast.FunctionDef)]
+ if len({item.name for item in functions}) != len(functions):
+ self._error(root, "Duplicate function declaration")
+ program, result = self.fresh("_t"), self.fresh("_t")
+ context = self._call(
+ runtime,
+ "ModuleProgram",
+ [
+ ast.Constant(root.name if is_module else None),
+ self._name(original_name, root) if is_module else ast.Constant(None),
+ ast.Tuple(copy.deepcopy(root.bases) if is_module else [], ast.Load()),
+ ],
+ root,
+ )
+ body = [
+ self._assign(
+ item.name, self._call(program, "reserve", [ast.Constant(item.name)], item), item
+ )
+ for item in functions
+ ]
+ if is_module:
+ body.append(self._assign(root.name, self._attribute(program, "namespace", root), root))
+ for member in members:
+ if isinstance(member, ast.FunctionDef):
+ continue
+ # Pattern: class-level statement -> same host statement, with module
+ # member publication after named assignments. Builders classify values.
+ body.append(copy.deepcopy(member))
+ targets = (
+ member.targets
+ if isinstance(member, ast.Assign)
+ else [member.target]
+ if isinstance(member, ast.AnnAssign)
+ else []
+ )
+ for target in targets:
+ if isinstance(target, ast.Name):
+ body.append(
+ self._assign(
+ target.id,
+ self._call(
+ program,
+ "member",
+ [ast.Constant(target.id), self._name(target.id, target)],
+ target,
+ ),
+ target,
+ )
+ )
+ definitions = []
+ records = []
+ for function in functions:
+ kind, _ = self.function_metadata(function)
+ if kind.python:
+ # Pattern: @I.pyfunc def f(...): body -> unchanged host function
+ # plus opaque builder registration.
+ host = copy.deepcopy(function)
+ host.decorator_list = []
+ body.append(host)
+ body.append(
+ self._assign(
+ function.name,
+ self._call(
+ program,
+ "python",
+ [
+ ast.Constant(function.name),
+ self._name(function.name, function),
+ ast.Constant(ast.unparse(function)),
+ self.span(function),
+ ],
+ function,
+ ),
+ function,
+ )
+ )
+ continue
+ declaration, record, callback = self.function_program(function, runtime, bindings)
+ body.extend(declaration)
+ definitions.append(
+ self._statement(
+ self._call(record, "define", [self._name(callback, function)], function),
+ function,
+ )
+ )
+ records.append(record)
+ body.extend(definitions)
+ output = (
+ self._attribute(program, "result", root)
+ if is_module
+ else self._attribute(records[0], "function", root)
+ )
+ translated = [
+ *copy.deepcopy(prefix),
+ self._with(context, body, root, self._name(program, root, True)),
+ self._assign(result, output, root),
+ ]
+ return ast.fix_missing_locations(ast.Module(translated, [])), result
diff --git a/python/tvm/tirx/script/builder/comparison.py b/python/tvm/tirx/script/builder/comparison.py
new file mode 100644
index 0000000000..0e6007c56a
--- /dev/null
+++ b/python/tvm/tirx/script/builder/comparison.py
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Build shared operand bindings for eager comparison-chain expressions."""
+
+from tvm import ir
+
+
+def _comparison_chain(comparisons, operands, conjunction, bind):
+ """Bind original chain operands progressively around short-circuit tests.
+
+ All inputs have already been constructed eagerly. The operand tuple retains
+ pre-conversion expression identities: comparison overloads may cast one
+ middle operand differently on either side. Python object reuse alone does
+ not make its compiled value execute once. Fresh IR variables replace those
+ identities inside comparisons, and each binding occurs only when its
+ comparison is reached. The first two operands evaluate left to right.
+
+ conjunction and bind are builder implementation functions, not lazy source
+ callbacks. They construct control flow and a dialect-specific expression
+ binding. All maps/variables are local to this construction; no registry,
+ native frame, or source expression is mutated.
+ """
+ import tvm_ffi
+
+ if len(operands) != len(comparisons) + 1:
+ raise ValueError("A comparison chain requires one more operand than comparison")
+ replacements = []
+ bindings = []
+ for operand in operands:
+ if isinstance(operand, tvm_ffi.ObjectConvertible):
+ operand = operand.asobject()
+ if not isinstance(operand, ir.Expr) or isinstance(operand, ir.Var):
+ bindings.append(None)
+ continue
+ previous = next((var for value, var in replacements if value.same_as(operand)), None)
+ if previous is not None:
+ bindings.append(None)
+ continue
+ variable = ir.Var("chain_operand", operand.ty)
+ replacements.append((operand, variable))
+ bindings.append((variable, operand))
+
+ def replace(value, mutator):
+ for original, variable in replacements:
+ if value.same_as(original):
+ return variable
+ return mutator.default_mutate(value)
+
+ conditions = []
+ for comparison in comparisons:
+ if isinstance(comparison, tvm_ffi.ObjectConvertible):
+ comparison = comparison.asobject()
+ conditions.append(
+ tvm_ffi.structural_mutate(comparison, [(ir.Expr, replace)])
+ if isinstance(comparison, ir.Expr)
+ else comparison
+ )
+ result = conditions[-1]
+ for index in range(len(conditions) - 1, -1, -1):
+ if index < len(conditions) - 1:
+ result = conjunction(conditions[index], result)
+ if bindings[index + 1] is not None:
+ variable, value = bindings[index + 1]
+ result = bind(variable, value, result)
+ if bindings[0] is not None:
+ variable, value = bindings[0]
+ result = bind(variable, value, result)
+ return result
diff --git a/python/tvm/tirx/script/builder/v2.py b/python/tvm/tirx/script/builder/v2.py
new file mode 100644
index 0000000000..c94359deb5
--- /dev/null
+++ b/python/tvm/tirx/script/builder/v2.py
@@ -0,0 +1,1805 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Concrete TIRx construction operations over the shared native IRBuilder stack."""
+
+import builtins as _python
+from dataclasses import dataclass as _dataclass
+from dataclasses import field as _field
+from functools import partial as _partial
+from functools import wraps as _wraps
+
+import tvm_ffi as _ffi
+
+from tvm import ir as _ir
+from tvm import tirx as _tir
+from tvm.script.ir_builder import IRBuilder as _IRBuilder
+from tvm.script.ir_builder import ir as _I
+from tvm.script.ir_builder.base import MISSING as _MISSING
+from tvm.script.ir_builder.base import IRBuilderFrame as _NativeFrame
+from tvm.script.ir_builder.base import _construction_span
+from tvm.script.ir_builder.base import _frame_result as _named_frame_result
+from tvm.script.ir_builder.base import at as _at
+from tvm.script.ir_builder.base import source_span as _source_span
+from tvm.script.ir_builder.type_var_frame import TypeVarDecl as _TypeVarDecl
+from tvm.script.ir_builder.type_var_frame import TypeVarFrame as _TypeVarFrame
+from tvm.script.ir_builder.type_var_frame import resolve_type_var as resolve_type_var
+from tvm.script.parser_v2.protocol import expr_str_args as _expression_args
+from tvm.script.parser_v2.protocol import register_type_var_decl as _register_type_var_decl
+from tvm.tirx.lang.alloc_pool import SMEMPool as SMEMPool
+from tvm.tirx.lang.alloc_pool import TMEMPool as TMEMPool
+
+from . import _ffi_api
+from . import frame as _frame
+from . import ir as _native
+from . import tirx as tile # noqa: F401
+from .ir import * # noqa: F403
+from .ir import boolean as bool # pylint: disable=redefined-builtin
+from .tirx import cluster as cluster
+from .tirx import cta as cta
+from .tirx import thread as thread
+from .tirx import warp as warp
+from .tirx import warpgroup as warpgroup
+from .tirx import wg as wg
+from .utils import buffer_proxy as buffer_proxy
+from .utils import frame_scope as frame_scope
+from .utils import seq_scope as seq_scope
+
+is_type_var = _ir.is_prim_var
+
+
+def type_var(name, *, dtype=None, span=None):
+ """Construct an explicit standalone primitive symbol.
+
+ Parameters
+ ----------
+ name : str
+ Name of the new symbol.
+ dtype : str, optional
+ Primitive dtype; None uses the explicit-construction default int64.
+ span : ir.Span or tuple, optional
+ IR span or original-source location; None omits the span.
+
+ Returns
+ -------
+ ir.Var
+ Fresh primitive variable.
+
+ Raises
+ ------
+ TypeError
+ If the primitive dtype or source location is invalid.
+
+ Notes
+ -----
+ This operation enters no frame and retains no symbol state. Expression
+ strings use resolve_type_var, whose function-owned first-use default is i64.
+
+ Examples
+ --------
+ >>> n = type_var("n", dtype="int32")
+ """
+ return _ir.Var(name, "int64" if dtype is None else dtype, _source_span(span))
+
+
+@_expression_args(
+ "shape",
+ "strides",
+ "elem_offset",
+ "byte_offset",
+ introduce=True,
+ implicit_dtype="int32",
+ compound_declarations=True,
+ as_type=True,
+)
+def Buffer(
+ shape,
+ dtype="float32",
+ data=None,
+ strides=None,
+ elem_offset=None,
+ byte_offset=None,
+ scope="global",
+ align=0,
+ offset_factor=0,
+ layout="default",
+ allocated_addr=None,
+ buffer_name="",
+ *,
+ span=None,
+):
+ """Construct a concrete buffer annotation from resolved expressions.
+
+ Parameters
+ ----------
+ shape : sequence of PrimExpr
+ Buffer dimensions; expression strings are resolved by the transpiler.
+ dtype : str, optional
+ Element dtype, default "float32".
+ data : ir.Var, optional
+ Underlying pointer; None creates a fresh pointer.
+ strides : sequence of PrimExpr, optional
+ Element strides; None uses the native compact-layout default.
+ elem_offset : PrimExpr, optional
+ Element offset, using the native default when omitted.
+ byte_offset : PrimExpr, optional
+ Byte offset, using the native default when omitted.
+ scope : str, optional
+ Storage scope, default "global".
+ align : int, optional
+ Data alignment; zero requests the native default.
+ offset_factor : int, optional
+ Required offset divisibility; zero requests the native default.
+ layout : object, optional
+ Native layout specification, default "default".
+ allocated_addr : sequence, optional
+ Allocation metadata, absent by default.
+ buffer_name : str, optional
+ Buffer name, empty by default.
+ span : ir.Span or tuple, optional
+ IR span or original-source location; None omits the span.
+
+ Returns
+ -------
+ ir.Var or ir.Type
+ Concrete buffer variable, or a missing annotation type when an eager
+ annotation still contains unresolved expression strings.
+
+ Raises
+ ------
+ TypeError
+ If unresolved strings are used in an active handwritten builder or the
+ native constructor rejects the shape, layout or element type.
+
+ Notes
+ -----
+ Only shape, strides and offset fields are expression fields. Names, dtype
+ and scope remain literal strings. The call temporarily applies its span,
+ but enters no persistent construction frame.
+
+ Examples
+ --------
+ >>> annotation = Buffer((16, 32), "float32")
+ """
+ with _construction_span(span):
+ return _at(
+ span,
+ _native.buffer(
+ shape,
+ dtype,
+ data,
+ strides,
+ elem_offset,
+ byte_offset,
+ scope,
+ align,
+ offset_factor,
+ layout,
+ allocated_addr,
+ buffer_name,
+ ),
+ )
+
+
+buffer = Buffer
+
+
+def Ptr(dtype, storage_scope="global", *, span=None):
+ """Construct a pointer variable usable as a function annotation.
+
+ Parameters
+ ----------
+ dtype : str, ir.PrimType or callable
+ Element dtype or a zero-argument primitive type constructor.
+ storage_scope : str, optional
+ Pointer storage scope, default "global".
+ span : ir.Span or tuple, optional
+ Source location; None omits the span.
+
+ Returns
+ -------
+ ir.Var
+ Fresh pointer variable.
+
+ Raises
+ ------
+ TypeError
+ If the element type cannot describe a pointer.
+
+ Notes
+ -----
+ The span applies only during construction; no frame or symbol is retained.
+
+ Examples
+ --------
+ >>> pointer = Ptr("float32", "shared")
+ """
+ if callable(dtype) and not isinstance(dtype, _ir.Expr):
+ dtype = dtype()
+ if isinstance(dtype, _ir.Expr | _TypeVarDecl):
+ dtype = dtype.ty
+ if isinstance(dtype, _ir.PrimType):
+ dtype = dtype.dtype
+ with _construction_span(span):
+ return _at(span, _native.ptr(dtype, storage_scope))
+
+
+class _Frame:
+ """Preserve a frame's source range through native finalization."""
+
+ def __init__(self, native, span=None):
+ # Each wrapper owns one native construction frame and source location.
+ # result starts empty, records finalized lexical exports on exit, and
+ # never outlives its construction region or stores another function's
+ # symbols (those belong to the separate TypeVarFrame).
+ self.native = native
+ self.span = span
+ self.result = {}
+
+ def __enter__(self):
+ with _construction_span(self.span):
+ value = self.native.__enter__()
+ return self if value is self.native else value
+
+ def __exit__(self, *exc):
+ with _construction_span(self.span):
+ return self.native.__exit__(*exc)
+
+ @property
+ def reference(self):
+ """Return the stable module reference after signature finalization."""
+ return self.native.global_var
+
+ def __getattr__(self, name):
+ return getattr(self.native, name)
+
+
+def frame_result(completed_frame, name):
+ """Return a named export from an explicitly supplied completed region.
+
+ Parameters
+ ----------
+ completed_frame : object or dict
+ Completed builder context or host-branch result dictionary.
+ name : str
+ Original source variable name.
+
+ Returns
+ -------
+ object
+ Exported value, including None, or MISSING when no export exists.
+
+ Raises
+ ------
+ TypeError
+ If name is not a string.
+
+ Notes
+ -----
+ TIR statement frames do not export their internal bindings. Host branch
+ dictionaries expose explicitly supplied values. This operation changes no
+ result maps, enters no frames and retains no completed-frame state.
+
+ Examples
+ --------
+ >>> frame_result({"answer": 42}, "answer")
+ 42
+ """
+ return _named_frame_result(completed_frame, name)
+
+
+def function(*, private=False, s_tir=False, persistent=False, span=None):
+ """Create a native primitive-function definition frame.
+
+ Parameters
+ ----------
+ private : bool, optional
+ Omit public symbol export, default False.
+ s_tir : bool, optional
+ Construct S-TIR semantics, default False.
+ persistent : bool, optional
+ Enable persistent-function metadata, default False.
+ span : ir.Span or tuple, optional
+ Function source location; None omits the span.
+
+ Returns
+ -------
+ _Frame
+ Context manager over the native function body frame.
+
+ Raises
+ ------
+ ValueError
+ If the native builder rejects the enclosing construction context.
+
+ Notes
+ -----
+ Entry pushes the native frame; exit finalizes it with the supplied span.
+ The context retains only that frame, its source location and result mapping.
+
+ Examples
+ --------
+ >>> # Inside an IRBuilder:
+ >>> with function(private=True):
+ ... arg("n", int32)
+ """
+ with _construction_span(span):
+ return _Frame(_native.prim_func(private=private, s_tir=s_tir, persistent=persistent), span)
+
+
+def decl_function(*, private=False, s_tir=False, persistent=False, span=None):
+ """Create a native primitive-function declaration frame.
+
+ Parameters
+ ----------
+ private : bool, optional
+ Omit public symbol export, default False.
+ s_tir : bool, optional
+ Construct S-TIR semantics, default False.
+ persistent : bool, optional
+ Enable persistent-function metadata, default False.
+ span : ir.Span or tuple, optional
+ Function source location; None omits the span.
+
+ Returns
+ -------
+ _Frame
+ Context manager over the native bodyless signature frame.
+
+ Raises
+ ------
+ ValueError
+ If the native builder rejects the enclosing construction context.
+
+ Notes
+ -----
+ Entry pushes the native frame; exit finalizes it with the supplied span.
+ The context retains only that frame, its source location and result mapping.
+
+ Examples
+ --------
+ >>> # Inside an IRBuilder:
+ >>> with decl_function(private=True):
+ ... arg("n", int32)
+ """
+ with _construction_span(span):
+ return _Frame(_ffi_api.DeclFunction(private, s_tir, persistent), span)
+
+
+def arg(name, annotation, *, span=None):
+ """Register a concrete primitive-function parameter.
+
+ Parameters
+ ----------
+ name : str
+ Parameter name.
+ annotation : ir.Type, ir.Expr, TypeVarDecl or callable
+ Concrete annotation or zero-argument constructor.
+ span : ir.Span or tuple, optional
+ Parameter source location; None omits the span.
+
+ Returns
+ -------
+ ir.Var
+ Registered parameter shared by declaration and definition.
+
+ Raises
+ ------
+ TypeError
+ If the annotation is not supported by the native parameter builder.
+
+ Notes
+ -----
+ Requires an active primitive-function frame and modifies that signature.
+ S-TIR buffer annotations discard the TIRx layout when registering parameters.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> parameter = arg("n", int32)
+ """
+ if callable(annotation) and not isinstance(annotation, _ir.Expr):
+ annotation = annotation()
+ if isinstance(annotation, _TypeVarDecl):
+ annotation = annotation.ty
+ if isinstance(annotation, _ir.Type):
+ annotation = _ir.Var(name, annotation)
+ with _construction_span(span):
+ if _tir.is_buffer_var(annotation) and annotation.ty.layout is not None:
+ frames = _IRBuilder.current().frames
+ if _python.any(
+ isinstance(frame, _frame.PrimFuncFrame) and frame.s_tir for frame in frames
+ ):
+ ty = annotation.ty
+ annotation = _native.buffer(
+ ty.shape,
+ ty.dtype,
+ strides=ty.strides,
+ elem_offset=ty.elem_offset,
+ scope=ty.storage_scope,
+ align=ty.data_alignment,
+ offset_factor=ty.offset_factor,
+ layout=None,
+ allocated_addr=list(ty.allocated_addr),
+ buffer_name=name,
+ )
+ return _native.arg(name, _at(span, annotation))
+
+
+def func_ret_type(annotation, *, span=None):
+ """Set the active primitive function's return type.
+
+ Parameters
+ ----------
+ annotation : ir.Type, ir.Expr, TypeVarDecl or callable
+ Concrete return type or zero-argument type constructor.
+ span : ir.Span or tuple, optional
+ Annotation source location; None omits the span.
+
+ Returns
+ -------
+ None
+ The active function signature is updated in place.
+
+ Raises
+ ------
+ TypeError
+ If the annotation cannot represent a native return type.
+
+ Notes
+ -----
+ Requires an active primitive-function frame; no additional frame is retained.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> func_ret_type(int32)
+ """
+ if callable(annotation) and not isinstance(annotation, _ir.Expr):
+ annotation = annotation()
+ if isinstance(annotation, _ir.Expr | _TypeVarDecl):
+ annotation = annotation.ty
+ with _construction_span(span):
+ return _native.func_ret(annotation)
+
+
+def _name(value, name, span):
+ if name is not None:
+ _IRBuilder.name(name, value)
+ return _at(span, value)
+
+
+def _enter_concise(frame):
+ native = frame.native if isinstance(frame, _Frame) else frame
+ native.add_callback(_partial(frame.__exit__, None, None, None))
+ return frame.__enter__()
+
+
+def _as_expr(value):
+ if isinstance(value, _ffi.ObjectConvertible):
+ value = value.asobject()
+ if isinstance(value, _ir.Expr):
+ return value
+ if isinstance(value, str):
+ return _ir.StringImm(value)
+ if isinstance(value, list | tuple):
+ return _ir.Tuple([_as_expr(item) for item in value])
+ return _tir.const(value)
+
+
+def bind_(
+ value=_MISSING,
+ *,
+ ty=None,
+ name=None,
+ span=None,
+ name_span=None,
+ previous=_MISSING,
+ declaration=False,
+ frame_value=False,
+):
+ """Construct a named binding under the active primitive function's policy.
+
+ Parameters
+ ----------
+ value : object, optional
+ Concrete expression, native frame or TypeVarDecl. MISSING denotes an
+ annotation-only binding.
+ ty : ir.Type or callable, optional
+ Concrete annotation or zero-argument type constructor; None omits it.
+ name : str, optional
+ Source binding name; None leaves the native name unchanged.
+ span : ir.Span or tuple, optional
+ Statement source location; None omits the span.
+ name_span : ir.Span or tuple, optional
+ Name source location; None uses span.
+ previous : object, optional
+ Prior binding for reassignment, or MISSING for a fresh binding.
+ declaration : bool, optional
+ Check and reuse an explicit primitive declaration, default False.
+ frame_value : bool, optional
+ Name a context-manager result without allocating storage, default False.
+
+ Returns
+ -------
+ object
+ Constructed binding, canonical symbol or unchanged Python value.
+
+ Raises
+ ------
+ ValueError
+ If an initializer is missing or a buffer, pointer or axis is rebound.
+ TypeError
+ If declaration types disagree or an annotation is unsupported.
+
+ Notes
+ -----
+ Requires the relevant native function frame. Symbol declarations use the
+ nearest TypeVarFrame and retain identity across declaration and definition.
+ Scalar bindings may allocate storage and emit stores. Native frames may be
+ entered with a callback that finalizes them with the enclosing region.
+
+ Examples
+ --------
+ >>> # Inside a function and TypeVarFrame:
+ >>> n = bind_(int32(), name="n")
+ """
+ name_span = span if name_span is None else name_span
+ # Dtype constructors shared with the legacy builder return anonymous Vars.
+ # Block axes and environment threads are already owned by native frames;
+ # replacing their Vars would disconnect references from those registrations.
+ # Other anonymous declarations share the function-owned type-variable frame.
+ if ty is None and not frame_value and _ir.is_prim_var(value):
+ if _python.any(
+ (
+ isinstance(frame, _frame.SBlockFrame)
+ and _python.any(axis.var.same_as(value) for axis in frame.iter_vars)
+ )
+ or (
+ isinstance(frame, _frame.PrimFuncFrame)
+ and _python.any(thread.same_as(value) for thread in frame.env_threads)
+ )
+ for frame in _IRBuilder.current().frames
+ ):
+ return _name(value, name, name_span)
+ if not value.name:
+ return _TypeVarFrame.current().resolve(name, value.ty, span=name_span)
+ if isinstance(value, _TypeVarDecl):
+ return _TypeVarFrame.current().resolve(name, value.ty, span=name_span)
+ with _construction_span(span):
+ if frame_value:
+ if isinstance(value, _frame.SBlockFrame):
+ raise TypeError("A block does not introduce an as-target value")
+ if isinstance(value, _python.list | _python.tuple | _ir.Array):
+ for index, item in enumerate(value):
+ bind_(
+ item,
+ name=None if name is None else f"{name}_{index}",
+ span=span,
+ name_span=name_span,
+ frame_value=True,
+ )
+ elif isinstance(value, _ir.Var | _tir.IterVar | _tir.Layout):
+ _name(value, name, name_span)
+ elif isinstance(value, _ir.TensorLoad) and _tir.is_buffer_var(value.source):
+ _name(value.source, name, name_span)
+ return value
+ if previous is not _MISSING and _tir.is_buffer_var(previous):
+ shape = previous.ty.shape
+ if len(shape) == 1 and bool(shape[0] == 1):
+ if value is _MISSING:
+ raise ValueError("A reassignment requires an initializer")
+ _native.buffer_store(previous, value, [0])
+ return previous
+ if previous is not _MISSING and isinstance(getattr(previous, "ty", None), _ir.PointerType):
+ raise ValueError(f"Pointer variable {name!r} cannot be reassigned")
+ if previous is not _MISSING and (
+ _tir.is_buffer_var(previous)
+ or isinstance(previous, _tir.IterVar)
+ or _python.any(
+ isinstance(frame, _frame.SBlockFrame)
+ and _python.any(axis.var.same_as(previous) for axis in frame.iter_vars)
+ for frame in _IRBuilder.current().frames
+ )
+ ):
+ raise ValueError(f"Cannot rebind buffer or block axis {name!r}")
+ if declaration:
+ if not _ir.is_prim_var(value):
+ raise TypeError("A symbol declaration requires a concrete primitive variable")
+ if ty is not None:
+ annotation = ty() if callable(ty) else ty
+ annotation = (
+ annotation.ty if isinstance(annotation, _ir.Expr | _TypeVarDecl) else annotation
+ )
+ if not _ffi.structural_equal(annotation, value.ty):
+ raise TypeError("The symbol declaration has an incompatible type")
+ if previous is not _MISSING:
+ if not _ir.is_prim_var(previous) or not _ffi.structural_equal(
+ previous.ty, value.ty
+ ):
+ raise TypeError("The symbol declaration has an incompatible signature dtype")
+ return previous
+ return _name(value, name, name_span)
+ if previous is not _MISSING and isinstance(previous, _ir.TensorLoad):
+ if value is _MISSING:
+ raise ValueError("A reassignment requires an initializer")
+ _native.buffer_store(previous.source, value, list(previous.indices))
+ return previous
+ if isinstance(value, _I.meta_var):
+ return value.value
+ if isinstance(ty, _native.LocalVectorAnnotation):
+ if value is not _MISSING:
+ raise ValueError("Vector annotation does not support an initializer")
+ return _name(_native.alloc_local(ty.shape, ty.dtype), name, name_span)
+ if isinstance(ty, _native.LetAnnotation):
+ if value is _MISSING:
+ raise ValueError("An immutable binding requires an initializer")
+ value = _as_expr(value)
+ variable = _name(ty.as_var(rhs_dtype=value.ty), name, name_span)
+ _native.Bind(value, var=variable)
+ return variable
+ if ty is not None:
+ annotation = ty() if callable(ty) and not isinstance(ty, _ir.Expr) else ty
+ annotation = (
+ annotation.ty if isinstance(annotation, _ir.Expr | _TypeVarDecl) else annotation
+ )
+ if not isinstance(annotation, _ir.PrimType) or str(annotation) == "handle":
+ raise TypeError("Mutable scalar annotations require a primitive scalar type")
+ result = _native.local_scalar(str(annotation)).scalar
+ _name(result.source, name, name_span)
+ if value is not _MISSING:
+ _native.buffer_store(result.source, value, [0])
+ return result
+ if value is _MISSING:
+ raise ValueError("An uninitialized binding requires a scalar type annotation")
+ if (
+ isinstance(value, _ir.TensorLoad)
+ and _tir.is_buffer_var(value.source)
+ and not value.source.name
+ and len(value.source.ty.shape) == 1
+ and isinstance(value.source.ty.shape[0], _tir.IntImm)
+ and value.source.ty.shape[0].value == 1
+ ):
+ _name(value.source, name, name_span)
+ return value
+ if isinstance(value, _native.scalar_wrapper):
+ _name(value.scalar.source, name, name_span)
+ return value.scalar
+ if isinstance(value, _NativeFrame | _Frame):
+ return _name(_enter_concise(value), name, name_span)
+ if isinstance(value, list | tuple):
+ for index, item in enumerate(value):
+ bind_(item, name=None if name is None else f"{name}_{index}", span=span)
+ return value
+ if getattr(type(value), "_is_meta_class", False):
+ if name is not None:
+ _native.name_meta_class_value(name, value)
+ return value
+ if _tir.is_buffer_var(value) or isinstance(value, _tir.IterVar | _tir.Layout):
+ return _name(value, name, name_span)
+ if isinstance(value, _ir.Var) and not value.name:
+ return _name(value, name, name_span)
+ if isinstance(value, _ir.TensorRegion):
+ return value
+ if not isinstance(value, _ir.Expr | _python.int | _python.float | _python.bool | str):
+ return value
+ value = _as_expr(value)
+ if _ir.is_prim_expr(value):
+ result = _native.local_scalar(str(value.ty.dtype)).scalar
+ _name(result.source, name, name_span)
+ _native.buffer_store(result.source, value, [0])
+ return result
+ return _name(_native.Bind(value), name, name_span)
+
+
+def emit_(value, *, span=None):
+ """Consume an expression statement, including effect-only calls.
+
+ Parameters
+ ----------
+ value : object
+ Expression, native frame, scalar value, or None from an effect-only call.
+ span : ir.Span or tuple, optional
+ Statement source location; None omits the span.
+
+ Returns
+ -------
+ None
+ Effects are emitted into the active native statement frame.
+
+ Raises
+ ------
+ TypeError
+ If the native evaluator cannot represent the supplied value.
+
+ Notes
+ -----
+ Strings, variables and None need no emitted evaluation. Concise native frames
+ remain entered until their enclosing region finalizes its callbacks.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> emit_(0)
+ """
+ if value is None or isinstance(value, str | _ir.Var):
+ return
+ with _construction_span(span):
+ if isinstance(value, _NativeFrame | _Frame):
+ _enter_concise(value)
+ elif hasattr(value, "frames"):
+ for frame in value.frames:
+ _enter_concise(frame)
+ elif isinstance(value, _tir.BufferStore):
+ _native.buffer_store(value.buffer, value.value, value.indices)
+ else:
+ _native.evaluate(value)
+
+
+def setitem(target, key, value, *, span=None):
+ """Emit an indexed store using already evaluated operands.
+
+ Parameters
+ ----------
+ target : ir.Var
+ Destination buffer.
+ key : PrimExpr or sequence of PrimExpr
+ Buffer indices.
+ value : PrimExpr
+ Value to store.
+ span : ir.Span or tuple, optional
+ Statement source location; None omits the span.
+
+ Returns
+ -------
+ None
+ A store is appended to the active statement frame.
+
+ Raises
+ ------
+ TypeError
+ If the target, indices or value are invalid for the native store.
+
+ Notes
+ -----
+ Requires an active statement frame and retains no additional Python state.
+
+ Examples
+ --------
+ >>> # Inside a function frame with buffer A:
+ >>> setitem(A, [0], 1)
+ """
+ with _construction_span(span):
+ _native.buffer_store(target, value, key)
+
+
+def setattr(target, name, value, *, span=None):
+ """Store through a scalar attribute or update Python metadata.
+
+ Parameters
+ ----------
+ target : object
+ Object containing the attribute.
+ name : str
+ Attribute name.
+ value : object
+ Replacement value, optionally wrapped by meta_var.
+ span : ir.Span or tuple, optional
+ Store source location; None omits the span.
+
+ Returns
+ -------
+ None
+ The attribute or its existing scalar storage is updated.
+
+ Raises
+ ------
+ AttributeError
+ If Python attribute assignment is unsupported.
+ TypeError
+ If a native scalar store receives an incompatible value.
+
+ Notes
+ -----
+ Native scalar stores require an active statement frame. Ordinary metadata
+ updates follow Python attribute semantics and retain no builder state.
+
+ Examples
+ --------
+ >>> # Update an ordinary metadata object:
+ >>> setattr(metadata, "count", 4)
+ """
+ if isinstance(value, _I.meta_var):
+ _python.setattr(target, name, value.value)
+ return
+ previous = getattr(target, name, _MISSING)
+ if isinstance(previous, _native.scalar_wrapper):
+ previous = previous.scalar
+ buffer = previous.source if isinstance(previous, _ir.TensorLoad) else previous
+ if _tir.is_buffer_var(buffer):
+ shape = buffer.ty.shape
+ if len(shape) == 1 and bool(shape[0] == 1):
+ bind_(value, previous=previous, span=span)
+ return
+ _python.setattr(target, name, value)
+
+
+def return_(value=None, *, span=None):
+ """Emit a primitive-function return expression.
+
+ Parameters
+ ----------
+ value : PrimExpr, optional
+ Return expression; the default None is rejected.
+ span : ir.Span or tuple, optional
+ Statement source location; None omits the span.
+
+ Returns
+ -------
+ None
+ A return statement is appended to the active function body.
+
+ Raises
+ ------
+ TypeError
+ If value is None or cannot be converted to a native expression.
+
+ Notes
+ -----
+ Requires an active primitive-function frame. Subsequent unreachable source
+ statements remain in the constructed IR; this operation does not exit Python.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> return_(0)
+ """
+ if value is None:
+ raise TypeError("A primitive function return requires an expression")
+ with _construction_span(span):
+ _native.Return(_as_expr(value))
+
+
+def _require_loop():
+ for frame in reversed(_IRBuilder.current().frames):
+ if isinstance(frame, _frame.ForFrame | _frame.WhileFrame):
+ return
+ if isinstance(frame, _frame.PrimFuncFrame):
+ break
+ raise ValueError("Loop control requires an enclosing primitive loop")
+
+
+def break_(*, span=None):
+ """Emit a break targeting the nearest primitive loop.
+
+ Parameters
+ ----------
+ span : ir.Span or tuple, optional
+ Statement source location; None omits the span.
+
+ Returns
+ -------
+ None
+ Loop control is appended to the active statement frame.
+
+ Raises
+ ------
+ ValueError
+ If no loop exists inside the current primitive function.
+
+ Notes
+ -----
+ Only the native frame stack determines the target. No additional state or
+ frame is retained, and Python construction continues after this call.
+
+ Examples
+ --------
+ >>> # Inside a primitive loop:
+ >>> break_()
+ """
+ _require_loop()
+ with _construction_span(span):
+ _native.evaluate(_native.break_loop())
+
+
+def continue_(*, span=None):
+ """Emit a continue targeting the nearest primitive loop.
+
+ Parameters
+ ----------
+ span : ir.Span or tuple, optional
+ Statement source location; None omits the span.
+
+ Returns
+ -------
+ None
+ Loop control is appended to the active statement frame.
+
+ Raises
+ ------
+ ValueError
+ If no loop exists inside the current primitive function.
+
+ Notes
+ -----
+ Only the native frame stack determines the target. No additional state or
+ frame is retained, and Python construction continues after this call.
+
+ Examples
+ --------
+ >>> # Inside a primitive loop:
+ >>> continue_()
+ """
+ _require_loop()
+ with _construction_span(span):
+ _native.evaluate(_native.continue_loop())
+
+
+def assert_(condition, message="", *, span=None):
+ """Emit a flat native assertion with its source location.
+
+ Parameters
+ ----------
+ condition : PrimExpr
+ Boolean assertion condition.
+ message : str, sequence or tuple, optional
+ Message, message parts, or (error_kind, message_parts). The default is
+ an empty message with RuntimeError as its error kind.
+ span : ir.Span or tuple, optional
+ Assertion source location; None omits the span.
+
+ Returns
+ -------
+ None
+ An assertion is appended to the active statement frame.
+
+ Raises
+ ------
+ TypeError
+ If error metadata is malformed or the condition is invalid.
+
+ Notes
+ -----
+ Requires an active statement frame. The assertion's temporary native frame
+ is finalized before this operation returns.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> assert_(n > 0, "n must be positive")
+ """
+ kind = "RuntimeError"
+ if isinstance(message, tuple):
+ if len(message) != 2 or not isinstance(message[0], str):
+ raise TypeError("Assertion metadata must be (error_kind, message_parts)")
+ kind, message = message
+ if isinstance(message, list | tuple):
+ message = [str(part) for part in message]
+ if not isinstance(message, list | tuple):
+ message = [message]
+ with _construction_span(span):
+ with _native.Assert(condition, message, error_kind=kind):
+ pass
+
+
+def If(condition, *, span=None):
+ """Create a native if statement region.
+
+ Parameters
+ ----------
+ condition : PrimExpr
+ Scalar boolean condition.
+ span : ir.Span or tuple, optional
+ Region source location; None omits the span.
+
+ Returns
+ -------
+ _Frame
+ Context manager for the native if region.
+
+ Raises
+ ------
+ ValueError
+ If the enclosing native region is invalid for this frame.
+
+ Notes
+ -----
+ Entry pushes the native region and exit finalizes it. Internal bindings are
+ lexical to the region and are not implicitly exported.
+
+ Examples
+ --------
+ >>> # Inside the appropriate native statement context:
+ >>> with If(condition):
+ ... emit_(0)
+ """
+ with _construction_span(span):
+ return _Frame(_native.If(condition), span)
+
+
+def Then(*, span=None):
+ """Create a native then statement region.
+
+ Parameters
+ ----------
+ span : ir.Span or tuple, optional
+ Region source location; None omits the span.
+
+ Returns
+ -------
+ _Frame
+ Context manager for the native then region.
+
+ Raises
+ ------
+ ValueError
+ If the enclosing native region is invalid for this frame.
+
+ Notes
+ -----
+ Entry pushes the native region and exit finalizes it. Internal bindings are
+ lexical to the region and are not implicitly exported.
+
+ Examples
+ --------
+ >>> # Inside the appropriate native statement context:
+ >>> with Then():
+ ... emit_(0)
+ """
+ with _construction_span(span):
+ return _Frame(_native.Then(), span)
+
+
+def Else(*, span=None):
+ """Create a native else statement region.
+
+ Parameters
+ ----------
+ span : ir.Span or tuple, optional
+ Region source location; None omits the span.
+
+ Returns
+ -------
+ _Frame
+ Context manager for the native else region.
+
+ Raises
+ ------
+ ValueError
+ If the enclosing native region is invalid for this frame.
+
+ Notes
+ -----
+ Entry pushes the native region and exit finalizes it. Internal bindings are
+ lexical to the region and are not implicitly exported.
+
+ Examples
+ --------
+ >>> # Inside the appropriate native statement context:
+ >>> with Else():
+ ... emit_(0)
+ """
+ with _construction_span(span):
+ return _Frame(_native.Else(), span)
+
+
+@_dataclass(frozen=True)
+class _IterationSpec:
+ """Loop description; native frame construction belongs to for_.
+
+ kind is 'range' or 'grid'; arguments holds concrete bounds/extents, and dtype
+ is the optional grid dtype. These fields are immutable and may be consumed
+ repeatedly into distinct frames. _context starts/reset to None and holds
+ only the delegated for_ context during a direct handwritten ``with spec``.
+ It is owned by this descriptor, excluded from equality/repr, and cleared on
+ every exit. It adds no frame stack: entry/exit use the existing native stack.
+ Direct nested reuse of one descriptor is rejected; sequential reuse is safe.
+ """
+
+ kind: str
+ arguments: tuple
+ dtype: str | None = None
+ _context: object = _field(default=None, init=False, repr=False, compare=False)
+
+ def __enter__(self):
+ if self._context is not None:
+ raise ValueError("An iteration descriptor is already entered")
+ context = for_(self)
+ _python.object.__setattr__(self, "_context", context)
+ try:
+ return context.__enter__()
+ except BaseException:
+ _python.object.__setattr__(self, "_context", None)
+ raise
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ if self._context is None:
+ raise ValueError("An iteration descriptor is not entered")
+ try:
+ return self._context.__exit__(exc_type, exc_value, traceback)
+ finally:
+ _python.object.__setattr__(self, "_context", None)
+
+
+def grid(*extents, dtype=None):
+ """Describe Cartesian iteration without constructing native loop frames.
+
+ Parameters
+ ----------
+ *extents : int, PrimExpr or tuple
+ Dimension extents or (start, extent) pairs.
+ dtype : str, optional
+ Loop-variable dtype; None infers it from each extent.
+
+ Returns
+ -------
+ _IterationSpec
+ Reusable descriptor consumed by for_ or directly as a context manager.
+
+ Raises
+ ------
+ ValueError
+ If the same descriptor is directly entered recursively.
+
+ Notes
+ -----
+ Native bound validation occurs when consumed. Sequential reuse constructs
+ fresh native frames. A direct context yields the native variable sequence
+ and finalizes the same frames on exit. Handwritten bounds must be concrete.
+
+ Examples
+ --------
+ >>> loops = grid(16, 32)
+ """
+ return _IterationSpec("grid", extents, dtype)
+
+
+class _LoopFrame(_Frame):
+ """One loop frame plus optional source target names, consumed at entry.
+
+ names is None, one source identifier, or an immutable tuple of identifiers
+ with an optional '*name' unpack marker, owned by this loop wrapper.
+ It never enters symbol state: loop variables are lexical iteration values,
+ not function-wide TypeVars. Native variables are named once on entry.
+ """
+
+ def __init__(self, native, names, span):
+ super().__init__(native, span)
+ self.names = names
+
+ def __enter__(self):
+ values = super().__enter__()
+ if self.names is not None:
+ variables = values if isinstance(values, list | tuple | _ir.Array) else [values]
+ if isinstance(self.names, str):
+ names = (
+ [self.names]
+ if len(variables) == 1
+ else [f"{self.names}_{index}" for index in range(len(variables))]
+ )
+ else:
+ names = list(self.names)
+ stars = [index for index, name in enumerate(names) if name.startswith("*")]
+ if stars:
+ index = stars[0]
+ count = len(variables) - len(names) + 1
+ if count < 0:
+ raise ValueError("Loop target count differs from iteration dimensions")
+ prefix = names[index][1:]
+ names[index : index + 1] = [f"{prefix}_{item}" for item in range(count)]
+ if len(variables) != len(names):
+ raise ValueError("Loop target count differs from iteration dimensions")
+ # Python's original tuple target still performs unpacking; even
+ # a one-dimensional grid must return a one-element sequence.
+ values = variables
+ for name, value in zip(names, variables):
+ _IRBuilder.name(name, value)
+ return values
+
+
+def for_(iterable, *, names=None, span=None):
+ """Create native loop scope from a concrete iteration specification.
+
+ Parameters
+ ----------
+ iterable : _IterationSpec, ForFrame or range
+ Grid/range descriptor, native loop frame or Python range.
+ names : str or tuple of str, optional
+ Source target names, optionally including one '*name' unpack marker.
+ None preserves native names. Grouped targets receive numeric suffixes.
+ span : ir.Span or tuple, optional
+ Loop source location; None omits the span.
+
+ Returns
+ -------
+ _LoopFrame
+ Context manager yielding one serial variable or a grid variable sequence.
+
+ Raises
+ ------
+ TypeError
+ If iterable or names has an unsupported type.
+ ValueError
+ If unpack arity differs or more than one starred target is supplied.
+
+ Notes
+ -----
+ Entry pushes native loop frames and names their variables. Exit finalizes
+ them. No function-wide type-variable symbols are introduced by loop targets.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> with for_(range_(16), names="i") as i:
+ ... emit_(i)
+ """
+ if names is not None and not isinstance(names, str):
+ if not isinstance(names, tuple) or not _python.all(isinstance(name, str) for name in names):
+ raise TypeError("Loop names must be a source identifier or tuple of identifiers")
+ if _python.sum(name.startswith("*") for name in names) > 1:
+ raise ValueError("Loop targets may contain only one starred group")
+ with _construction_span(span):
+ if isinstance(iterable, _IterationSpec):
+ if iterable.kind == "grid":
+ iterable = _native.grid(*iterable.arguments, dtype=iterable.dtype)
+ else:
+ start, stop, step = iterable.arguments
+ iterable = _native.serial(start, stop, step=step)
+ if isinstance(iterable, _python.range):
+ iterable = _native.serial(iterable.start, iterable.stop, step=iterable.step)
+ if not isinstance(iterable, _frame.ForFrame):
+ raise TypeError("A primitive for loop requires an iteration specification")
+ return _LoopFrame(iterable, names, span)
+
+
+For = for_
+
+
+def While(condition, *, span=None):
+ """Create a native while statement region.
+
+ Parameters
+ ----------
+ condition : PrimExpr
+ Scalar boolean condition.
+ span : ir.Span or tuple, optional
+ Region source location; None omits the span.
+
+ Returns
+ -------
+ _Frame
+ Context manager for the native while region.
+
+ Raises
+ ------
+ ValueError
+ If the enclosing native region is invalid for this frame.
+
+ Notes
+ -----
+ Entry pushes the native region and exit finalizes it. Internal bindings are
+ lexical to the region and are not implicitly exported.
+
+ Examples
+ --------
+ >>> # Inside the appropriate native statement context:
+ >>> with While(condition):
+ ... emit_(0)
+ """
+ with _construction_span(span):
+ return _Frame(_native.While(condition), span)
+
+
+def unpack(value):
+ """Project a concrete IR tuple while preserving Python iteration.
+
+ Parameters
+ ----------
+ value : object
+ IR tuple, tuple-typed expression or ordinary Python iterable.
+
+ Returns
+ -------
+ object
+ Tuple of concrete fields or projections, or the unchanged Python value.
+
+ Raises
+ ------
+ TypeError
+ If a native tuple projection rejects an invalid type.
+
+ Notes
+ -----
+ No frame is entered and no construction state is retained.
+
+ Examples
+ --------
+ >>> unpack((1, 2))
+ (1, 2)
+ """
+ if isinstance(value, _ir.Tuple):
+ return _python.tuple(value.fields)
+ if isinstance(value, _ir.Expr) and isinstance(value.ty, _ir.TupleType):
+ return _python.tuple(_ir.TupleGetItem(value, i) for i in range(len(value.ty.fields)))
+ return value
+
+
+def alloc_scalar(dtype="float32", scope="global"):
+ """Allocate scalar storage and return its load expression.
+
+ Parameters
+ ----------
+ dtype : str, optional
+ Scalar dtype, default "float32".
+ scope : str, optional
+ Storage scope, default "global".
+
+ Returns
+ -------
+ PrimExpr
+ Concrete scalar load, unwrapped from the native scalar wrapper.
+
+ Raises
+ ------
+ TypeError
+ If the native allocation rejects the dtype or storage scope.
+
+ Notes
+ -----
+ Requires an active allocation-capable native frame. The allocated storage
+ belongs to that enclosing construction region.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> counter = alloc_scalar("int32", "local")
+ """
+ value = _native.alloc_scalar(dtype, scope)
+ return value.scalar if isinstance(value, _native.scalar_wrapper) else value
+
+
+def local_scalar(dtype="float32"):
+ """Allocate scalar storage in local memory.
+
+ Parameters
+ ----------
+ dtype : str, optional
+ Scalar dtype, default "float32".
+
+ Returns
+ -------
+ PrimExpr
+ Concrete load from the allocated scalar.
+
+ Raises
+ ------
+ TypeError
+ If the dtype cannot describe scalar storage.
+
+ Notes
+ -----
+ Requires an active allocation-capable frame. Storage belongs to the
+ containing native region; no separate Python state is retained.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> counter = local_scalar("int32")
+ """
+ return alloc_scalar(dtype, "local")
+
+
+def shared_scalar(dtype="float32"):
+ """Allocate scalar storage in shared memory.
+
+ Parameters
+ ----------
+ dtype : str, optional
+ Scalar dtype, default "float32".
+
+ Returns
+ -------
+ PrimExpr
+ Concrete load from the allocated scalar.
+
+ Raises
+ ------
+ TypeError
+ If the dtype cannot describe scalar storage.
+
+ Notes
+ -----
+ Requires an active allocation-capable frame. Storage belongs to the
+ containing native region; no separate Python state is retained.
+
+ Examples
+ --------
+ >>> # Inside a function frame:
+ >>> counter = shared_scalar("int32")
+ """
+ return alloc_scalar(dtype, "shared")
+
+
+@_expression_args(
+ "shape",
+ "strides",
+ "elem_offset",
+ introduce=True,
+ implicit_dtype="int32",
+ compound_declarations=True,
+)
+@_wraps(_native.match_buffer)
+def match_buffer(*args, **kwargs):
+ """Construct a native buffer match with resolved symbolic shape fields."""
+ return _native.match_buffer(*args, **kwargs)
+
+
+# Constructor identities carry syntax policy; aliases share it without wrappers.
+for _constructor in vars(_native).values():
+ if isinstance(_constructor, _native.DtypeConstructor):
+ _register_type_var_decl(_constructor, dtype=_constructor._dtype_str)
+del _constructor
+
+
+def range_(*args):
+ """Describe serial iteration using Python's range argument forms.
+
+ Parameters
+ ----------
+ *args : int or PrimExpr
+ Stop; start and stop; or start, stop and step. Start defaults to zero,
+ and omitted step retains the native semantic default of one.
+
+ Returns
+ -------
+ _IterationSpec
+ Descriptor consumed by for_ or directly as a context manager.
+
+ Raises
+ ------
+ TypeError
+ Unless one to three positional arguments are supplied.
+ ValueError
+ If a literal step is zero or the descriptor is directly re-entered.
+
+ Notes
+ -----
+ Creates no native frame or variable until consumed. Native loops validate
+ other bounds and steps. Sequential context-manager reuse creates fresh frames.
+
+ Examples
+ --------
+ >>> iteration = range_(0, 16, 2)
+ """
+ if len(args) == 1:
+ args = (0, args[0], None)
+ elif len(args) == 2:
+ args = (*args, None)
+ elif len(args) != 3:
+ raise TypeError("range expects one to three arguments")
+ if isinstance(args[2], _python.int) and args[2] == 0:
+ raise ValueError("range step cannot be zero")
+ return _IterationSpec("range", args)
+
+
+def logical_and(*values):
+ """Construct scalar or vector conjunction from eager operands.
+
+ Parameters
+ ----------
+ *values : object
+ Host values or primitive IR expressions, optionally ObjectConvertible.
+
+ Returns
+ -------
+ object
+ Selected host value or native scalar/vector boolean expression.
+
+ Raises
+ ------
+ TypeError
+ If no operands are supplied or native types are incompatible.
+
+ Notes
+ -----
+ Python has already constructed all arguments. Native scalar operands use
+ boolean IR nodes; vector operands use elementwise operators. No frame or
+ construction state persists after the call.
+
+ Examples
+ --------
+ >>> logical_and(True, False)
+ False
+ """
+ if not values:
+ raise TypeError("logical_and requires at least one operand")
+ values = [
+ value.asobject() if isinstance(value, _ffi.ObjectConvertible) else value for value in values
+ ]
+ result = values[0]
+ for value in values[1:]:
+ if not isinstance(result, _ir.Expr) and not isinstance(value, _ir.Expr):
+ result = result and value
+ else:
+ lhs, rhs = _as_expr(result), _as_expr(value)
+ result = _tir.And(lhs, rhs) if lhs.ty.is_scalar() and rhs.ty.is_scalar() else lhs & rhs
+ return result
+
+
+def logical_or(*values):
+ """Construct scalar or vector disjunction from eager operands.
+
+ Parameters
+ ----------
+ *values : object
+ Host values or primitive IR expressions, optionally ObjectConvertible.
+
+ Returns
+ -------
+ object
+ Selected host value or native scalar/vector boolean expression.
+
+ Raises
+ ------
+ TypeError
+ If no operands are supplied or native types are incompatible.
+
+ Notes
+ -----
+ Python has already constructed all arguments. Native scalar operands use
+ boolean IR nodes; vector operands use elementwise operators. No frame or
+ construction state persists after the call.
+
+ Examples
+ --------
+ >>> logical_or(True, False)
+ True
+ """
+ if not values:
+ raise TypeError("logical_or requires at least one operand")
+ values = [
+ value.asobject() if isinstance(value, _ffi.ObjectConvertible) else value for value in values
+ ]
+ result = values[0]
+ for value in values[1:]:
+ if not isinstance(result, _ir.Expr) and not isinstance(value, _ir.Expr):
+ result = result or value
+ else:
+ lhs, rhs = _as_expr(result), _as_expr(value)
+ result = _tir.Or(lhs, rhs) if lhs.ty.is_scalar() and rhs.ty.is_scalar() else lhs | rhs
+ return result
+
+
+def logical_not(value):
+ """Negate a host or IR value without testing IR truth in Python.
+
+ Parameters
+ ----------
+ value : object
+ Host value or boolean IR expression, optionally ObjectConvertible.
+
+ Returns
+ -------
+ bool or PrimExpr
+ Host boolean or primitive boolean expression.
+
+ Raises
+ ------
+ TypeError
+ If native operand types or host truth conversion are invalid.
+
+ Notes
+ -----
+ This operation enters no frame and retains no state.
+
+ Examples
+ --------
+ >>> logical_not(False)
+ True
+ """
+ if isinstance(value, _ffi.ObjectConvertible):
+ value = value.asobject()
+ return _tir.Not(value) if isinstance(value, _ir.Expr) else not value
+
+
+def select(condition, true_value, false_value):
+ """Construct a scalar conditional whose runtime evaluates one arm.
+
+ Parameters
+ ----------
+ condition : bool or PrimExpr
+ Host condition or scalar boolean expression, optionally ObjectConvertible.
+ true_value : object
+ Already constructed value selected by a true condition.
+ false_value : object
+ Already constructed value selected by a false condition.
+
+ Returns
+ -------
+ object
+ Selected host value or scalar primitive conditional expression.
+
+ Raises
+ ------
+ TypeError
+ If the condition is not scalar or IR arm types are incompatible.
+
+ Notes
+ -----
+ Python eagerly constructs all three arguments. The native conditional emits
+ control flow without binding both arms outside it. No frame or state persists.
+
+ Examples
+ --------
+ >>> select(True, 1, 2)
+ 1
+ """
+ if isinstance(condition, _ffi.ObjectConvertible):
+ condition = condition.asobject()
+ if not isinstance(condition, _ir.Expr):
+ return true_value if condition else false_value
+ return _tir.if_then_else(condition, true_value, false_value)
+
+
+def if_then_else_(condition, true_value, false_value):
+ """Construct a scalar conditional whose runtime evaluates one arm.
+
+ Parameters
+ ----------
+ condition : bool or PrimExpr
+ Host condition or scalar boolean expression, optionally ObjectConvertible.
+ true_value : object
+ Already constructed value selected by a true condition.
+ false_value : object
+ Already constructed value selected by a false condition.
+
+ Returns
+ -------
+ object
+ Selected host value or scalar primitive conditional expression.
+
+ Raises
+ ------
+ TypeError
+ If the condition is not scalar or IR arm types are incompatible.
+
+ Notes
+ -----
+ Python eagerly constructs all three arguments. The native conditional emits
+ control flow without binding both arms outside it. No frame or state persists.
+
+ Examples
+ --------
+ >>> if_then_else_(True, 1, 2)
+ 1
+ """
+ return select(condition, true_value, false_value)
+
+
+def and_(*values, chain=None):
+ """Construct left-to-right scalar conjunction with runtime short circuit.
+
+ Parameters
+ ----------
+ *values : object
+ Eagerly constructed host values or scalar boolean IR expressions.
+ chain : tuple, optional
+ Original comparison operands; None denotes an ordinary conjunction.
+ A chain must contain one more operand than values.
+
+ Returns
+ -------
+ object
+ Selected host value or scalar boolean expression.
+
+ Raises
+ ------
+ TypeError
+ If no values are supplied or IR operands are invalid.
+ ValueError
+ If the chain operand count is invalid.
+
+ Notes
+ -----
+ Runtime evaluation stops at the first false operand. For comparison chains,
+ native bindings evaluate each shared operand once when its comparison is
+ reached. All operands are eagerly constructed; no callbacks or frames persist.
+
+ Examples
+ --------
+ >>> and_(True, False)
+ False
+ """
+ if not values:
+ raise TypeError("and_ requires at least one operand")
+ if chain is not None:
+ from .comparison import _comparison_chain
+
+ return _comparison_chain(values, chain, and_, _tir.Let)
+ result = values[-1]
+ for value in reversed(values[:-1]):
+ result = if_then_else_(
+ value, result, False if isinstance(value, _ir.Expr | _ffi.ObjectConvertible) else value
+ )
+ return result
+
+
+def or_(*values):
+ """Construct left-to-right scalar disjunction with runtime short circuit.
+
+ Parameters
+ ----------
+ *values : object
+ Eagerly constructed host values or scalar boolean IR expressions.
+
+ Returns
+ -------
+ object
+ Selected host value or scalar boolean expression.
+
+ Raises
+ ------
+ TypeError
+ If no operands are supplied or IR operands have invalid types.
+
+ Notes
+ -----
+ Runtime evaluation stops at the first true operand. Python has already
+ constructed every argument. No native frame or state persists after the call.
+
+ Examples
+ --------
+ >>> or_(False, True)
+ True
+ """
+ if not values:
+ raise TypeError("or_ requires at least one operand")
+ result = values[-1]
+ for value in reversed(values[:-1]):
+ result = if_then_else_(
+ value, True if isinstance(value, _ir.Expr | _ffi.ObjectConvertible) else value, result
+ )
+ return result
+
+
+def not_(value):
+ """Negate a host or IR value without testing IR truth in Python.
+
+ Parameters
+ ----------
+ value : object
+ Host value or boolean IR expression, optionally ObjectConvertible.
+
+ Returns
+ -------
+ bool or PrimExpr
+ Host boolean or primitive boolean expression.
+
+ Raises
+ ------
+ TypeError
+ If native operand types or host truth conversion are invalid.
+
+ Notes
+ -----
+ This operation enters no frame and retains no state.
+
+ Examples
+ --------
+ >>> not_(False)
+ True
+ """
+ return logical_not(value)
diff --git a/src/relax/script/builder/frame.cc b/src/relax/script/builder/frame.cc
index 6d05e05dce..9a62dfa7cf 100644
--- a/src/relax/script/builder/frame.cc
+++ b/src/relax/script/builder/frame.cc
@@ -116,7 +116,7 @@ void FunctionFrameNode::ExitWithScope() {
TVM_FFI_CHECK(local_var.has_value(), ValueError)
<< "A local function definition requires its declared reference";
EmitVarBinding(tvm::relax::VarBinding(local_var.value(), func, span));
- } else if (builder->frames.empty()) {
+ } else if (!builder->HasConstructionFrames()) {
// Case 0. No outer frame, return function directly
TVM_FFI_CHECK(!builder->result.has_value(), ValueError)
<< "Builder.result has already been set";
diff --git a/src/script/ir_builder/base.cc b/src/script/ir_builder/base.cc
index 0090dcb7d3..5032d1d2d2 100644
--- a/src/script/ir_builder/base.cc
+++ b/src/script/ir_builder/base.cc
@@ -64,6 +64,7 @@ void AppendNormalizedSpan(const Span& span, std::vector<Span>* normalized) {
TVM_FFI_STATIC_INIT_BLOCK() {
IRBuilderFrameNode::RegisterReflection();
+ TypeVarFrameNode::RegisterReflection();
IRBuilderNode::RegisterReflection();
}
@@ -193,6 +194,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_method("script.ir_builder.IRBuilderFrameEnter", &IRBuilderFrameNode::EnterWithScope)
.def_method("script.ir_builder.IRBuilderFrameExit", &IRBuilderFrameNode::ExitWithScope)
.def_method("script.ir_builder.IRBuilderFrameAddCallback", &IRBuilderFrameNode::AddCallback)
+ .def("script.ir_builder.TypeVarFrame", []() { return TypeVarFrame(); })
+ .def("script.ir_builder.TypeVarFrameSetSymbol",
+ [](TypeVarFrame frame, ffi::String name, tvm::Var symbol) {
+ frame->symbols.Set(name, symbol);
+ })
.def("script.ir_builder.IRBuilder", []() { return IRBuilder(); })
.def_method("script.ir_builder.IRBuilderEnter", &IRBuilder::EnterWithScope)
.def_method("script.ir_builder.IRBuilderExit", &IRBuilder::ExitWithScope)
diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc
index 3fe8188e23..f00ccd9138 100644
--- a/src/tirx/script/builder/frame.cc
+++ b/src/tirx/script/builder/frame.cc
@@ -185,7 +185,7 @@ void PrimFuncFrameNode::ExitWithScope() {
}
function = func;
IRBuilder builder = IRBuilder::Current();
- if (builder->frames.empty()) {
+ if (!builder->HasConstructionFrames()) {
TVM_FFI_CHECK(!builder->result.has_value(), ValueError)
<< "Builder.result has already been set";
builder->result = func;