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 17:05:29 UTC
(tvm) branch tvmscript-ast-only-transpiler created (now 1a23bc7059)
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
at 1a23bc7059 Lower TVMScript syntax through explicit builder programs
This branch includes the following new commits:
new 1a23bc7059 Lower TVMScript syntax through explicit builder programs
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.
(tvm) 01/01: Lower TVMScript syntax through explicit builder programs
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 1a23bc70593219f1289b7a13953f5ffe8b46e832
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Sep 21 17:04:38 2026 +0000
Lower TVMScript syntax through explicit builder programs
Move signature construction, symbolic identity, annotation policy and source spans into shared builder infrastructure. Keep source transformation limited to Python syntax and metadata, with direct statement lowering and original locations.
Support canonical i32 shape symbols through native Relax construction and runtime lowering while preserving explicit types in printed signatures.
---
include/tvm/script/ir_builder/base.h | 38 +-
python/tvm/ir/__init__.py | 11 +-
python/tvm/ir/type.py | 11 +
python/tvm/relax/script/__init__.py | 2 +-
python/tvm/relax/script/builder/__init__.py | 106 +-
python/tvm/script/__init__.py | 6 +-
python/tvm/script/{parser => }/ir.py | 6 +-
python/tvm/script/ir_builder/__init__.py | 3 +-
python/tvm/script/ir_builder/construction.py | 333 +++++
python/tvm/script/ir_builder/protocol.py | 331 ++++-
python/tvm/script/ir_builder/type_var_frame.py | 142 +++
python/tvm/script/parser/__init__.py | 1 -
python/tvm/script/parser/annotations.py | 622 ---------
python/tvm/script/parser/diagnostics.py | 17 +-
python/tvm/script/parser/expression.py | 226 ++++
python/tvm/script/parser/frontend.py | 632 ++++------
python/tvm/script/parser/functions.py | 237 ----
python/tvm/script/parser/transform.py | 782 ------------
python/tvm/script/parser/transpile.py | 1334 ++++++++++++++++++++
python/tvm/tirx/script/builder/__init__.py | 258 +++-
python/tvm/tirx/script/builder/ir.py | 80 +-
src/relax/analysis/type_analysis.cc | 12 +-
src/relax/backend/vm/vm_shape_lower.cc | 17 +-
src/relax/ir/dependent_type.cc | 9 +-
src/relax/ir/expr.cc | 9 +-
src/relax/script/builder/frame.cc | 2 +-
src/relax/script/printer/tir.cc | 6 +-
src/relax/transform/static_plan_block_memory.cc | 20 +-
src/script/ir_builder/base.cc | 6 +
src/script/printer/ir/ir.cc | 3 +-
src/script/printer/utils.h | 42 +-
src/tirx/script/builder/frame.cc | 2 +-
tests/python/relax/test_analysis_well_formed.py | 2 +-
.../relax/test_backend_dispatch_sort_scan.py | 12 +
.../relax/test_backend_transform_shape_lower.py | 26 +-
tests/python/relax/test_blockbuilder_core.py | 3 +-
tests/python/relax/test_expr.py | 5 +-
tests/python/relax/test_shape_i32.py | 147 +++
tests/python/relax/test_transform_normalize.py | 2 +
tests/python/relax/test_tvmscript_parser.py | 122 +-
tests/python/relax/test_tvmscript_type_vars.py | 40 +
tests/python/relax/test_vm_build.py | 5 +
tests/python/tvmscript/test_parser.py | 72 +-
43 files changed, 3459 insertions(+), 2283 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/ir/__init__.py b/python/tvm/ir/__init__.py
index 17e2c7f147..bc334f9630 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -33,7 +33,16 @@ from .base import (
# Register Type before Expr. Expr's reflected ``ty`` field otherwise creates
# an auto-generated Type wrapper before the concrete Python class is available.
-from .type import FuncType, OpaqueType, PointerType, PrimType, StringType, TupleType, Type
+from .type import (
+ MissingType,
+ FuncType,
+ OpaqueType,
+ PointerType,
+ PrimType,
+ StringType,
+ TupleType,
+ Type,
+)
from .expr import (
Call,
Constant,
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 43060f95a6..782261baeb 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -54,6 +54,17 @@ class Type(Node, Scriptable):
return self.is_(other)
+def MissingType():
+ """Return the concrete missing-type singleton.
+
+ This zero-argument constructor returns an ``ir.Type`` accepted by composite
+ type constructors, including TupleType. It shares identity and semantics
+ with ``Type.missing()`` and stores no deferred expression or callable.
+ It has no builder effects and does not raise for missing annotation data.
+ """
+ return Type.missing()
+
+
@tvm_ffi.register_object("ir.OpaqueType")
class OpaqueType(Type):
"""Type marker for opaque values that must be removed from finished IR."""
diff --git a/python/tvm/relax/script/__init__.py b/python/tvm/relax/script/__init__.py
index 97cbd7277e..1b1484950c 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -21,10 +21,10 @@ import sys as _sys
from tvm import relax as _relax
from tvm.relax.base_py_module import PyModuleFactory as _PyModuleFactory
+from tvm.script.ir_builder.construction import register_opaque_factory as _register_opaque_factory
from tvm.script.parser.frontend import make_decorator as _make_decorator
from tvm.script.parser.frontend import make_helper as _make_helper
from tvm.script.parser.frontend import register_namespace as _register_namespace
-from tvm.script.parser.functions import register_opaque_factory as _register_opaque_factory
from . import builder as _builder
from .builder import *
diff --git a/python/tvm/relax/script/builder/__init__.py b/python/tvm/relax/script/builder/__init__.py
index 234f279b5e..16e6386287 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/relax/script/builder/__init__.py
@@ -34,6 +34,9 @@ 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 import protocol as _protocol
+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 . import _ffi_api, distributed, frame, ir
from . import distributed as dist
@@ -45,13 +48,25 @@ from .ir import *
@_protocol.expression_args("shape", introduce=True, dtype="int64", scalar_strings=False)
def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
- """Construct a concrete tensor type from already resolved dimensions."""
+ """Construct a tensor type from resolved shape dimensions.
+
+ ``shape`` is a dimension sequence or None (unknown); a bare string abbreviates
+ dtype when dtype is omitted. ``dtype`` is a dtype string or None (unknown).
+ ``vdevice`` is None, a concrete virtual device, or a target[:index] lookup
+ string. ``ndim`` defaults to -1 (unknown). ``span`` accepts None, ir.Span, or
+ a source-location tuple. Returns TensorType without entering a frame; device
+ lookup requires module context, and native shape/type errors propagate.
+ Signed int32/int64 dimensions retain their canonical symbols directly;
+ representation conversion belongs to runtime shape lowering.
+ Nested shape strings/TypeVars produce MissingType during eager annotation
+ evaluation; handwritten calls in a builder require resolved expressions.
+ """
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, span)
+ return _relax.TensorType(shape, dtype, vdevice, ndim, _protocol.source_span(span))
@_protocol.expression_args("shape", introduce=True, dtype="int64", scalar_strings=False)
@@ -63,7 +78,9 @@ def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1,
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, span)
+ return _DTensorType(
+ Tensor(shape, dtype, ndim=ndim), device_mesh, placement, _protocol.source_span(span)
+ )
Range = _ir.Range
@@ -71,8 +88,16 @@ Range = _ir.Range
@_protocol.expression_args("values", introduce=True, dtype="int64")
def Shape(values=None, ndim=-1, *, span=None):
- """Construct a concrete shape type."""
- return _relax.ShapeType(values, ndim, span)
+ """Construct a shape type from concrete dimension expressions.
+
+ ``values`` is a dimension sequence or None (unknown); ``ndim`` defaults to
+ -1 (unknown). ``span`` is None, ir.Span, or a source-location tuple. Returns
+ ShapeType without entering a frame. Signed int32/int64 dimensions retain
+ their symbols directly. Native rank/value errors propagate.
+ Unresolved eager expression strings produce MissingType through the protocol
+ wrapper; handwritten calls in a builder require concrete expressions.
+ """
+ return _relax.ShapeType(values, ndim, _protocol.source_span(span))
def _type(value):
@@ -80,7 +105,7 @@ def _type(value):
return _ir.TupleType([])
if callable(value):
value = value()
- if _ir.is_prim_expr(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__}")
@@ -96,20 +121,30 @@ def Callable(params=None, ret=None, purity=None, derive_func=None, *, span=None)
ret=None if ret is None else _type(ret),
derive_func=derive_func,
purity=purity,
- span=span,
+ span=_protocol.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, span)
+ return _relax.FuncType(
+ [_type(param) for param in params], _type(ret), purity, _protocol.source_span(span)
+ )
def Tuple(*fields, span=None):
- """Construct a concrete tuple type."""
+ """Construct a tuple type, preserving missing individual field types.
+
+ ``fields`` accepts positional concrete types or one list/tuple of types;
+ zero fields denotes an empty tuple. Primitive expression annotations and
+ zero-argument type constructors are also accepted. ``span`` is None,
+ ir.Span, or a source-location tuple. Returns ir.TupleType without entering
+ a frame. MissingType fields remain fields rather than collapsing the tuple.
+ Unsupported field values raise TypeError; native type errors propagate.
+ """
if len(fields) == 1 and isinstance(fields[0], list | _python.tuple):
fields = fields[0]
- return _ir.TupleType([_type(field) for field in fields], span)
+ return _ir.TupleType([_type(field) for field in fields], _protocol.source_span(span))
def Prim(dtype, *, span=None):
@@ -122,7 +157,7 @@ Prim.__tvm_parameter_dtype__ = "dtype"
def Object(*, span=None):
"""Construct the unconstrained Relax value type."""
- return _relax.AnyType(span)
+ return _relax.AnyType(_protocol.source_span(span))
Any = Object
@@ -132,14 +167,26 @@ is_type_var = _ir.is_prim_var
def type_var(name, *, dtype=None, span=None):
- """Construct a signature symbol under Relax's default shape dtype policy."""
- return _ir.Var(name, "int64" if dtype is None else dtype, span)
+ """Construct an explicit standalone primitive symbol.
+
+ name is the source-name string; dtype is a dtype string or None (the legacy
+ explicit-construction default int64). span accepts None, an IR span or an
+ original-source location tuple. Returns a new primitive Var, entering no
+ frame and retaining no symbol state. Invalid types propagate constructor
+ errors. Generated expression strings use resolve_type_var instead, which
+ shares the active function frame and defaults newly introduced names to i32.
+ """
+ return _ir.Var(name, "int64" if dtype is None else dtype, _protocol.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 = {}
@@ -254,15 +301,36 @@ def bind_(
declaration=False,
frame_value=False,
):
- """Emit a binding, or name and preserve an existing frame-owned value."""
+ """Emit a Relax binding or preserve a named frame-owned value.
+
+ ``value`` is an expression, Python value, TypeVarDecl, or MISSING (default).
+ ``ty`` is an optional concrete annotation or zero-argument type constructor.
+ ``name`` optionally names the result. ``span`` and ``name_span`` accept None,
+ ir.Span, or source-location tuples; name_span defaults to span. ``previous``
+ optionally supplies a prior symbol for compatible declarations, otherwise
+ MISSING. ``declaration`` checks/reuses a primitive symbol; ``frame_value``
+ names an existing frame result without emitting a binding (both default
+ False). TypeVarDecl resolves a canonical symbol in the nearest TypeVarFrame.
+
+ Returns a named emitted Var, canonical symbol, or unchanged Python value.
+ Expression inputs emit bindings in the active native function/block; a
+ TypeVarDecl only updates the symbol frame. Missing initializers, incompatible
+ declaration/match-cast types, or operations after an unconditional return
+ raise ValueError/TypeError; native construction errors propagate.
+ """
_check_unterminated()
+ name_span = _protocol.source_span(span if name_span is None else name_span)
+ if isinstance(value, _TypeVarDecl):
+ if ty is not None and not _ffi.structural_equal(_type(ty), value.ty):
+ raise TypeError("The symbol declaration has an incompatible type")
+ 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=span,
+ span=_protocol.source_span(span),
name_span=name_span,
frame_value=True,
)
@@ -331,7 +399,7 @@ def match_cast(value, ty, *, span=None):
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, span)
+ return _relax.MatchCast(_ir.Var("", ty), _value(value), ty, _protocol.source_span(span))
def unpack(value):
@@ -377,6 +445,8 @@ __all__ = [
"Callable",
"DTensor",
"For",
+ "for_",
+ "resolve_type_var",
"Object",
"Prim",
"Range",
@@ -460,6 +530,4 @@ def _global_callee(function):
_protocol.register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee)
-from tvm.script.parser.annotations import enable_eager_constructors as _enable_eager_constructors
-
-_enable_eager_constructors(_sys.modules[__name__])
+for_ = For
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index b198a0eb10..dbf2d711a3 100644
--- a/python/tvm/script/__init__.py
+++ b/python/tvm/script/__init__.py
@@ -125,6 +125,8 @@ def _redirect_target(fullname: str) -> str | None:
Returns ``None`` if ``fullname`` is not a redirected name.
"""
+ if fullname == "tvm.script.parser.ir":
+ return "tvm.script.ir"
if fullname.startswith("tvm.script."):
# tvm.script.<dialect>[.subpath]
rest = fullname[len("tvm.script.") :]
@@ -222,7 +224,7 @@ def __getattr__(name: str) -> Any:
return module
if name == "ir":
# IR is foundational and shares the generic parser entry point.
- ir_parser = importlib.import_module("tvm.script.parser.ir")
+ ir_parser = importlib.import_module("tvm.script.ir")
globals()["ir"] = ir_parser
return ir_parser
if name in ("from_source", "parse"):
@@ -233,7 +235,7 @@ def __getattr__(name: str) -> Any:
return parse
if name == "ir_module":
# The foundational module entry does not use the dialect registry.
- ir_parser = importlib.import_module("tvm.script.parser.ir")
+ ir_parser = importlib.import_module("tvm.script.ir")
ir_module_value = ir_parser.ir_module
globals()["ir_module"] = ir_module_value
return ir_module_value
diff --git a/python/tvm/script/parser/ir.py b/python/tvm/script/ir.py
similarity index 87%
rename from python/tvm/script/parser/ir.py
rename to python/tvm/script/ir.py
index 9471ada2a0..d3d6b6d5a9 100644
--- a/python/tvm/script/parser/ir.py
+++ b/python/tvm/script/ir.py
@@ -21,8 +21,8 @@ import sys
from tvm.ir import Range as Range
from tvm.script.ir_builder.ir import * # noqa: F403
-from .frontend import ir_module as ir_module
-from .frontend import pyfunc as pyfunc
-from .frontend import register_namespace
+from .parser.frontend import ir_module as ir_module
+from .parser.frontend import pyfunc as pyfunc
+from .parser.frontend import register_namespace
register_namespace("I", sys.modules[__name__])
diff --git a/python/tvm/script/ir_builder/__init__.py b/python/tvm/script/ir_builder/__init__.py
index 56f54040d0..30f34d0a51 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -35,7 +35,8 @@ import importlib
from typing import Any
from .base import IRBuilder
-from .protocol import MISSING, _at, expression_args
+from .protocol import MISSING, ExprStrPolicy, _at, expr_str_args, expression_args
+from .type_var_frame import TypeVarDecl, TypeVarFrame
def __getattr__(name: str) -> Any:
diff --git a/python/tvm/script/ir_builder/construction.py b/python/tvm/script/ir_builder/construction.py
new file mode 100644
index 0000000000..5fec914630
--- /dev/null
+++ b/python/tvm/script/ir_builder/construction.py
@@ -0,0 +1,333 @@
+# 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 . import protocol
+from .base import IRBuilder
+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 of opaque Python module functions.
+
+ Parameters
+ ----------
+ factory : callable
+ Called with (name, Python callable, source text, concrete source span);
+ must return a concrete function accepted by the module builder.
+ module_adapter : callable or None
+ Optional (completed module, original class or None, tuple of bases)
+ callback returning the public module object. None disables adaptation.
+
+ Returns
+ -------
+ None
+ Registration replaces the process-wide callbacks without running bodies.
+
+ Raises
+ ------
+ TypeError
+ Either supplied callback is not callable.
+ """
+ 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
+ (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.
+
+ State
+ -----
+ 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.
+ """
+
+ 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 = protocol.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 = protocol.MISSING
+ self.reference = self.function = None
+
+ @contextmanager
+ def declaration(self):
+ """Enter this function's symbol and declaration frames, yielding self.
+
+ Frames leave in reverse order, including on errors. The finalized
+ reference is retained only after normal exit. Builder validation errors
+ propagate to the original source operation.
+ """
+ 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 one parameter in the active declaration.
+
+ name is its original string name; annotation is an evaluated builder
+ annotation (including a zero-argument declaration callable/descriptor).
+ location is optional original-source location data. Returns the concrete
+ parameter. Builder type/compatibility errors propagate unchanged.
+ """
+ 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=protocol.source_span(location)
+ )
+ if isinstance(annotation, ir.PrimType):
+ annotation = self.symbols.resolve(name, annotation, span=protocol.source_span(location))
+ elif ir.is_prim_var(annotation):
+ annotation = self.symbols.bind(name, annotation)
+ value = self.builder.arg(name, annotation, span=protocol.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 explicitly typed scalar parameter before dependent shapes.
+
+ name is its source name; annotation is a re-evaluated declaration callable,
+ TypeVarDecl or primitive type. location is optional source data. Returns
+ its canonical primitive symbol; raises TypeError for other annotations or
+ incompatible prior declarations. Runs inside declaration() only.
+ """
+ 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):
+ return self.symbols.bind(name, annotation)
+ return self.symbols.resolve(name, annotation, span=protocol.source_span(location))
+
+ def symbol(self, name):
+ """Return a declared symbol by source name after signature construction.
+
+ name is a string introduced by the signature's expression syntax. Returns
+ its concrete symbol without constructing one; raises KeyError if absent.
+ This operation does not require the retained frame to be currently active.
+ """
+ return self.symbols.symbols[name]
+
+ def capture(self, name, fallback):
+ """Resolve a generated body's free annotation name without introducing it.
+
+ name is a source identifier and fallback its opaque lexical value.
+ Returns an existing canonical symbol when the signature registered that
+ name, otherwise fallback unchanged. There are no frame or cache effects.
+ """
+ return self.symbols.symbols.get(name, fallback)
+
+ def returns(self, annotation):
+ """Set the evaluated return annotation and apply it to the active frame.
+
+ Returns None. The same concrete annotation is reused at definition;
+ source annotation evaluation itself has no exactly-once guarantee.
+ """
+ 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):
+ """Execute body with retained parameters under the definition frame.
+
+ body is a generated Python callable accepting parameters in source order;
+ ordinary Python closures/defaults preserve lexical captures. Returns the
+ constructed function and retains it in function. The shared symbol frame
+ is re-entered, never copied. Builder/host exceptions propagate unchanged.
+ """
+ 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 protocol.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.
+
+ State
+ -----
+ 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.
+ """
+
+ 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 a stable named function reference and expose it on namespace.
+
+ name is a source identifier. Returns the concrete module reference; the
+ active module validates duplicate/incompatible declarations.
+ """
+ reference = I.reserve_function(name)
+ setattr(self.namespace, name, reference)
+ return reference
+
+ def member(self, name, value):
+ """Register a source class assignment, returning its exposed value.
+
+ name is a source identifier and value an evaluated Python value. Concrete
+ functions become module declarations/definitions; all other values remain
+ opaque class namespace members. Native validation errors propagate.
+ """
+ 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.
+
+ name/source are original name/text, function is its Python callable, and
+ location optional original source data. Returns the opaque module
+ reference. Raises ValueError when no owner factory was registered.
+ """
+ if _OPAQUE_FACTORY is None:
+ raise ValueError("No opaque Python function constructor has been registered")
+ opaque = _OPAQUE_FACTORY(name, function, source, protocol.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
diff --git a/python/tvm/script/ir_builder/protocol.py b/python/tvm/script/ir_builder/protocol.py
index 5a971e24a2..b6247e3760 100644
--- a/python/tvm/script/ir_builder/protocol.py
+++ b/python/tvm/script/ir_builder/protocol.py
@@ -16,23 +16,27 @@
# under the License.
"""Shared construction metadata for explicit TVMScript builders.
-Metadata describes source syntax; decorated constructors still execute normally
-and return concrete values. Dialects register their own function kinds here, so
-translation consumes construction policies without importing their owners.
+Metadata describes source syntax. Decorated constructors return concrete values
+or a composable missing type during unresolved eager annotation evaluation.
+Dialects register their own policies here, so translation reads metadata without
+importing dialects or evaluating annotations.
"""
from builtins import locals as locals
from builtins import slice as slice
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
+from functools import wraps
from inspect import signature
-from typing import Any, NamedTuple
+from typing import Any, NamedTuple, TypeVar
from tvm import ir
from .base import IRBuilder
+# Process-wide absence marker for bindings, distinct from the concrete missing
+# IR type. It is immutable and shared; no function or evaluation data is stored.
class _Missing:
def __repr__(self):
return "MISSING"
@@ -41,8 +45,20 @@ class _Missing:
MISSING = _Missing()
-class ExpressionArguments(NamedTuple):
- """Syntax policy shared by a constructor and every alias of it."""
+class ExprStrPolicy(NamedTuple):
+ """Immutable syntax policy for registered constructor arguments.
+
+ ``fields`` names parameters whose strings represent expressions.
+ ``scalar_strings`` controls bare strings; nested strings are always marked.
+ fields is a tuple of parameter-name strings; scalar_strings defaults True.
+ introduce (False), dtype (None), implicit_dtype (None), and
+ compound_declarations (False) retain legacy registration metadata only.
+ They never drive transpiler symbol construction or type validation; those
+ operations belong exclusively to builders. Construction returns this tuple
+ record and has no builder effects; NamedTuple enforces field arity only.
+ Instances live as long as their registered callable, across all transpilation
+ passes, and hold no symbols, eager results, or function-local state.
+ """
fields: tuple[str, ...]
introduce: bool = False
@@ -52,32 +68,76 @@ class ExpressionArguments(NamedTuple):
compound_declarations: bool = False
-def expression_args(
+# 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):
+ """Return a callable's ExprStrPolicy, or None when it is unregistered.
+
+ ``constructor`` may be any Python value. This read-only lookup neither
+ imports dialects nor evaluates the value and has no builder/frame effects.
+ """
+ try:
+ return _EXPR_STR_POLICIES.get(constructor)
+ 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,
):
- """Mark constructor fields whose nested strings are source expressions.
-
- ``introduce`` permits signature/match scopes to introduce otherwise unknown
- names. ``dtype`` optionally selects the dialect's symbol-construction policy;
- it is metadata, never an evaluator or a replacement constructor. Set
- ``scalar_strings=False`` when a bare string is literal shorthand while
- strings nested in tuples/lists remain expressions. ``implicit_dtype`` selects
- only the default for names introduced by strings; explicit TypeVars retain
- the dialect default unless ``dtype`` is supplied. ``compound_declarations``
- also permits first-use declarations inside quoted compound expressions.
+ """Register expression-string fields and wrap eager constructor calls.
+
+ Parameters
+ ----------
+ fields : str
+ Named parameters whose string values denote source expressions. Unknown
+ parameter names raise ValueError when applying the decorator.
+ introduce : bool, optional
+ Legacy syntax metadata, default False. Builders own symbol introduction.
+ dtype : object, optional
+ Legacy explicit symbol-type metadata, default None.
+ scalar_strings : bool, optional
+ Whether bare strings are expressions, default True. False preserves
+ literal shorthand such as ``Tensor("float32")``; nested strings in
+ tuples/lists remain expressions.
+ implicit_dtype : object, optional
+ Legacy implicit-type metadata, default None; newly resolved symbols
+ default to int32 in TypeVarFrame.
+ compound_declarations : bool, optional
+ Legacy expression metadata, default False.
+ as_type : bool, optional
+ Preserve use as an annotation class (including Python type unions),
+ default False. True returns a class whose construction invokes the
+ wrapped callable; its instances are still the callable's return values.
+
+ Returns
+ -------
+ decorator : callable
+ Registers and wraps a constructor. Unresolved strings or typing.TypeVar
+ values in marked fields return a concrete ``ir.MissingType()`` during
+ eager Python annotation evaluation. Concrete arguments call the original
+ constructor. Inside an active builder, unresolved inputs raise TypeError:
+ handwritten builders require concrete expressions. Calls are never
+ cached and annotations must be safe to re-evaluate. Ordinary call-binding
+ errors propagate from inspect.Signature.bind.
"""
def decorate(constructor):
- parameters = signature(constructor).parameters
- unknown = set(fields).difference(parameters)
+ call_signature = signature(constructor)
+ unknown = set(fields).difference(call_signature.parameters)
if unknown:
raise ValueError(f"Unknown expression argument fields: {sorted(unknown)}")
- constructor.__tvm_expression_args__ = ExpressionArguments(
+ policy = ExprStrPolicy(
tuple(fields),
bool(introduce),
dtype,
@@ -85,59 +145,200 @@ def expression_args(
implicit_dtype,
bool(compound_declarations),
)
- return constructor
+
+ 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.
+ from .type_var_frame import TypeVarFrame
+
+ 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.MissingType()
+ 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__,
+ },
+ )
+ 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):
- """Callable syntax that declares a symbol when its value argument is absent."""
+ """Immutable scalar-constructor metadata used by signature predeclaration.
+
+ value_parameter names the optional value argument whose absence denotes a
+ declaration; dtype is the explicit primitive dtype. Stored on the callable
+ by register_declaration and shared across aliases and transpilation passes.
+ It contains no symbols or construction results.
+ """
value_parameter: str
dtype: Any = None
def register_declaration(constructor, *, value_parameter="expr", dtype=None):
- """Mark a concrete constructor's declaration form without wrapping the call."""
+ """Register metadata for a constructor's zero-argument declaration form.
+
+ ``constructor`` is a callable supporting attribute assignment;
+ ``value_parameter`` defaults to "expr" and names its optional value argument.
+ ``dtype`` is the explicit dtype spelling, default None. Returns the same
+ callable after attaching immutable DeclarationArguments; calls are never
+ wrapped/evaluated and no frame is entered. Unsupported attribute assignment
+ raises the ordinary Python error. Registration persists with the callable.
+ """
constructor.__tvm_declaration_args__ = DeclarationArguments(value_parameter, dtype)
return constructor
@dataclass(frozen=True)
class FunctionKind:
- """Construction namespace and policies registered by a function decorator."""
+ """Builder namespace and syntax options attached to a function decorator.
+
+ builder is the dialect construction namespace; metadata maps string policy
+ names to values (for example option_map/defaults). Registration owns this
+ record for the decorator's lifetime; transpilers read it across functions.
+ No concrete function, frame, annotation result, or symbol is stored here.
+ """
builder: Any
metadata: dict
def register_function(decorator, builder, **metadata):
- """Associate a decorator identity with its explicit construction namespace."""
+ """Associate a decorator with its builder namespace and syntax policies.
+
+ ``decorator`` must support attribute assignment; ``builder`` is its dialect
+ construction namespace and ``metadata`` contains named syntax policies.
+ Returns the same decorator after replacing its FunctionKind record. No
+ construction or annotation evaluation occurs. Unsupported attribute
+ assignment raises the ordinary Python error; metadata is retained for the
+ decorator's lifetime and must not be mutated by a transpilation pass.
+ """
decorator.__tvm_function_kind__ = FunctionKind(builder, metadata)
return decorator
def function_kind(decorator):
- """Return registered function metadata, without resolving namespace names."""
+ """Return a decorator's FunctionKind record, or None when unregistered.
+
+ ``decorator`` is any already resolved Python value. This lookup neither
+ evaluates it nor enters frames; ordinary custom attribute-access errors
+ propagate. Returned metadata is shared and must be treated as read-only.
+ """
return getattr(decorator, "__tvm_function_kind__", None)
+def source_span(location):
+ """Materialize an IR span from a source location, or preserve an existing span.
+
+ ``location`` is None, an ir.Span, or ``(filename, lineno, end_lineno,
+ col_offset, end_col_offset)`` using Python AST's one-based lines and
+ zero-based UTF-8 byte columns. The returned ir.Span retains those exact
+ columns, matching TVMScript's source-range convention.
+ None returns None. Invalid tuple shape/types propagate TypeError/ValueError
+ from unpacking or span construction. This function does not enter a frame.
+ """
+ if location is None or isinstance(location, ir.Span):
+ return location
+ filename, line, end_line, column, end_column = location
+ return ir.Span(ir.SourceName(filename), line, end_line, column, end_column)
+
+
@contextmanager
def span_context(span):
- """Use the existing span stack, preserving a failing operation's source range."""
- context = IRBuilder.current().with_source_span(span) if span is not None else nullcontext()
+ """Temporarily attach source context to nested builder operations.
+
+ ``span`` accepts the same None/ir.Span/location-tuple forms as source_span.
+ Yields None and pushes/pops the existing builder span stack when a builder
+ is active. 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.
+ """
+ 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_span__"):
- error.__tvm_script_span__ = span
+ 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 an expression and preserve ordinary Python values."""
- if span is not None and isinstance(value, ir.Expr):
+ """Attach source metadata to a concrete expression where possible.
+
+ ``span`` is None, ir.Span, or a location tuple; ``value`` is any Python
+ value. Returns the original value, attaching the normalized current source
+ range to IR expressions without a span when a builder is active. Ordinary
+ values and calls without a builder pass through. Temporarily enters the
+ span stack, never a construction frame; invalid spans propagate errors.
+ """
+ if span is not None and isinstance(value, ir.Expr) and IRBuilder.is_in_scope():
with span_context(span):
return IRBuilder.current()._set_current_source_span(value)
return value
@@ -147,24 +348,45 @@ _at = at
def frame_result(frame):
- """Read finalized lexical exports without imposing a dialect policy."""
+ """Return frame.result, or an empty mapping when no result is exposed.
+
+ frame is a finalized builder context object. The result maps source names
+ to exported values and is owned by the frame. This read-only operation
+ enters no frames; custom attribute-access errors propagate.
+ """
return getattr(frame, "result", {})
def require_defined(value, name):
- """Reject reads of names absent from a finalized lexical export set."""
+ """Return value unless it is MISSING, in which case raise NameError.
+
+ name is the original source identifier used in the error. value is an
+ opaque lexical export. No state is retained and no frame is entered.
+ """
if value is MISSING:
raise NameError(f"name {name!r} is not defined")
return value
def is_python_bool(value):
- """Identify an ordinary condition that selects a construction-time branch."""
+ """Return whether value is a Python bool, without invoking truth conversion.
+
+ value may be any host or builder value. This pure predicate lets generated
+ code choose construction-time control flow without inspecting IR classes.
+ """
return isinstance(value, bool)
def compare_chain(logical_and, operands, comparisons):
- """Evaluate each operand once, stopping after an ordinary false comparison."""
+ """Evaluate a comparison chain in source order with host short-circuiting.
+
+ logical_and is the builder binary conjunction callable; operands is a
+ nonempty sequence of zero-argument callbacks; comparisons has one binary
+ callback per adjacent operand pair. Returns a host bool or builder result.
+ Each operand runs at most once. A Python False stops further evaluation;
+ symbolic conjunction uses logical_and. Callback/arity errors propagate.
+ No frames or persistent state are created by this helper.
+ """
left = operands[0]()
result = True
for index, comparison in enumerate(comparisons):
@@ -178,7 +400,14 @@ def compare_chain(logical_and, operands, comparisons):
def logical_chain(operation, operands, short_circuit):
- """Preserve ordinary boolean short-circuiting while constructing symbolic operands."""
+ """Fold lazy operands with host boolean short-circuiting.
+
+ operation is the builder binary and/or callable; operands is a nonempty
+ sequence of zero-argument callbacks. short_circuit is False for and, True
+ for or. Returns the first short-circuit bool or final builder result.
+ Callback errors propagate. Only evaluation-local state is retained and
+ no frames are entered; symbolic values never undergo Python truth tests.
+ """
result = operands[0]()
for operand in operands[1:]:
if isinstance(result, bool) and result is short_circuit:
@@ -188,21 +417,43 @@ def logical_chain(operation, operands, short_circuit):
def select_lazy(operation, condition, true_value, false_value):
- """Select one ordinary boolean arm, or construct both symbolic arms."""
+ """Evaluate a conditional expression through lazy source callbacks.
+
+ operation accepts (condition, true_result, false_result); condition is an
+ already evaluated host/builder value. true_value and false_value are
+ zero-argument callbacks. A Python bool evaluates and returns one arm;
+ otherwise both arms run in source order and operation returns the result.
+ Callback errors propagate; this helper owns no frames or persistent state.
+ """
if isinstance(condition, bool):
return true_value() if condition else false_value()
return operation(condition, true_value(), false_value())
def register_call_kind(builder, value_type, adapter):
- """Register a concrete callable's construction policy on its builder namespace."""
+ """Register runtime adaptation of callable values for one builder namespace.
+
+ ``builder`` is a namespace supporting attributes, ``value_type`` a Python
+ type suitable for isinstance, and ``adapter`` maps matching values to
+ callable construction operations. Returns None and replaces that type's
+ entry in the namespace-owned __tvm_call_kinds__ dict. The dict maps types
+ to adapters, starts empty, and persists across functions; no values or
+ call results are cached. No frames are entered. Invalid attribute writes
+ raise Python errors; invalid types/adapters fail when consumed by callee.
+ """
policies = dict(getattr(builder, "__tvm_call_kinds__", {}))
policies[value_type] = adapter
builder.__tvm_call_kinds__ = policies
def callee(builder, value):
- """Resolve an already-evaluated callable without altering global call behavior."""
+ """Adapt a runtime callable according to one builder's registered policy.
+
+ ``builder`` is the construction namespace and ``value`` the already
+ evaluated call target. Returns the first matching adapter's result, or
+ value unchanged when no type matches. The operation reads namespace-owned
+ registration state, enters no frames itself, and propagates adapter errors.
+ """
for value_type, adapter in getattr(builder, "__tvm_call_kinds__", {}).items():
if isinstance(value, value_type):
return adapter(value)
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..ef34bfbb6b
--- /dev/null
+++ b/python/tvm/script/ir_builder/type_var_frame.py
@@ -0,0 +1,142 @@
+# 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
+
+import tvm_ffi
+
+from tvm import ir
+
+from . import _ffi_api
+from .base import IRBuilder, IRBuilderFrame
+from .protocol import source_span
+
+
+@dataclass(frozen=True)
+class TypeVarDecl:
+ """An unbound symbolic declaration consumed by a dialect's ``bind_``.
+
+ Parameters
+ ----------
+ ty : ir.Type or str
+ Explicit primitive type or dtype spelling. Dtype strings become
+ PrimType. Non-primitive types raise TypeError. The immutable ``ty``
+ field holds only that concrete type, never a symbol, name, or frame.
+ Construction has no builder effects; binding determines symbol identity.
+ """
+
+ 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 a function's declaration and definition.
+
+ The zero-argument constructor initializes a native ``symbols`` map from
+ source names (str) to ir.Var. The function construction record owns this
+ frame; entering pushes it on the existing native IRBuilder frame stack,
+ and normal exit pops it without clearing symbols. Re-enter the same frame
+ outside the function's native declaration/definition frame. Blocks reuse
+ the enclosing frame; nested functions allocate a distinct frame, so nearest
+ active lookup establishes shadowing. There is no Python global frame stack.
+
+ Requires an active IRBuilder when entered. On exceptional exit the enclosing
+ builder is discarded, following IRBuilderFrame's normal cleanup contract.
+ """
+
+ def __init__(self):
+ self.__init_handle_by_constructor__(_ffi_api.TypeVarFrame)
+
+ @staticmethod
+ def current():
+ """Return the nearest active TypeVarFrame on the native stack.
+
+ No arguments or mutations. Raises ValueError if no IRBuilder or no
+ function symbol frame is active; never falls back to another builder.
+ """
+ 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 variable.
+
+ ``name`` is a nonempty str. ``ty`` is a PrimType or dtype string, or
+ None to reuse any existing type and default a new symbol to int32.
+ ``span`` is None, ir.Span, or a protocol source-location tuple; it is
+ attached only when creating a symbol. Returns ir.Var and adds a map
+ entry on first use. Raises ValueError for an empty name, TypeError for
+ a non-primitive type or an explicit type incompatible with prior use.
+ The map survives frame exit and declaration/definition re-entry.
+ """
+ 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:
+ value = self.symbols[name]
+ if ty is not None and not tvm_ffi.structural_equal(value.ty, ty):
+ raise TypeError(f"Symbol {name!r} has an incompatible declared type")
+ return value
+ value = ir.Var(name, ir.PrimType("int32") 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 Var, or reuse a compatible prior Var.
+
+ ``name`` is its nonempty source binding name and ``symbol`` an ir.Var
+ with primitive type. Returns the prior canonical symbol when present,
+ otherwise stores/returns ``symbol``, assigning name to an anonymous Var.
+ A named Var must already match name. Raises TypeError on incompatible
+ type or nonprimitive input, ValueError on an empty name. No new frame
+ is entered; this frame's map retains the result across re-entry.
+ """
+ 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)
+ elif symbol.name != name:
+ raise ValueError("A symbolic binding must retain its declared name")
+ _ffi_api.TypeVarFrameSetSymbol(self, name, symbol)
+ return symbol
+
+
+def resolve_type_var(name, *, dtype=None, span=None):
+ """Resolve a symbolic name in the nearest active function frame.
+
+ ``name`` is a nonempty source name; ``dtype`` is an optional PrimType or
+ dtype string (None reuses existing type, defaulting new symbols to int32).
+ ``span`` accepts None, ir.Span, or a source-location tuple. Returns ir.Var,
+ retaining its identity in the function frame. Raises ValueError without an
+ active frame or valid name and TypeError for incompatible explicit types.
+ """
+ return TypeVarFrame.current().resolve(name, dtype, span=span)
diff --git a/python/tvm/script/parser/__init__.py b/python/tvm/script/parser/__init__.py
index 787d6d82a2..6d9941886d 100644
--- a/python/tvm/script/parser/__init__.py
+++ b/python/tvm/script/parser/__init__.py
@@ -23,7 +23,6 @@ those of the original source. Entry modules register policies; this parser never
imports their namespaces.
"""
-from . import ir
from .frontend import (
_NAMESPACES,
from_source,
diff --git a/python/tvm/script/parser/annotations.py b/python/tvm/script/parser/annotations.py
deleted file mode 100644
index ce63a781b3..0000000000
--- a/python/tvm/script/parser/annotations.py
+++ /dev/null
@@ -1,622 +0,0 @@
-# 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.
-"""Annotation evaluation and callable-owned expression-argument syntax.
-
-Constructors receive concrete values. This scope prepares designated symbols,
-rewrites only registered expression fields, and evaluates each annotation once.
-"""
-
-import ast
-import builtins
-import copy
-import inspect
-import linecache
-import re
-from functools import lru_cache, wraps
-from types import SimpleNamespace
-from typing import TypeVar
-
-from tvm.script.ir_builder import IRBuilder
-
-
-class AnnotationScope:
- """A function's canonical symbols and once-evaluated annotation values."""
-
- def __init__(self, env, builder, filename, span):
- self.env = dict(env)
- self.builder = builder
- self.filename = filename
- self.span = span
- self.symbols = {}
- self._evaluated = {}
- self._prepared = {}
- self._pending_parameters = {}
- self._parameter_annotations = {}
- self._unbound_parameters = set()
- self._shape_declarations = {}
- self._type_vars = {}
- self._counter = 0
- self._used_names = set(env)
-
- def _error(self, node, message):
- raise SyntaxError(
- message,
- (
- self.filename,
- node.lineno,
- node.col_offset + 1,
- linecache.getline(self.filename, node.lineno),
- ),
- )
-
- def _symbol(self, name, node, dtype=None, *, shadow=False):
- if not shadow and name in self.symbols:
- return self.symbols[name]
- value = self.env.get(name)
- if not shadow and name in self.env and not isinstance(value, TypeVar):
- if getattr(self.builder, "is_type_var", lambda value: False)(value):
- self.symbols[name] = value
- return value
- value = self.builder.type_var(name, dtype=dtype, span=self.span(node))
- self.symbols[name] = self.env[name] = value
- return value
-
- def _canonical_type_var(self, name, node, dtype=None):
- value = self.env.get(name)
- if not isinstance(value, TypeVar):
- return
- if value.__constraints__ or value.__bound__ is not None:
- self._error(node, "A symbolic TypeVar cannot have constraints or a bound")
- if value.__name__ != name:
- self._error(node, "A symbolic TypeVar binding must match its declared name")
- if value not in self._type_vars:
- self._type_vars[value] = self._symbol(name, node, dtype)
- else:
- self.symbols[name] = self.env[name] = self._type_vars[value]
-
- def prepare_type_params(self, type_params):
- """Introduce explicit host-supported PEP 695 parameters, shadowing outer names."""
- type_var_node = getattr(ast, "TypeVar", ())
- for parameter in type_params:
- if not isinstance(parameter, type_var_node):
- self._error(parameter, "Only scalar TypeVar parameters are supported")
- bound = getattr(parameter, "bound", None)
- if bound is not None and not (
- isinstance(bound, ast.Name)
- and self.env.get(bound.id, getattr(builtins, bound.id, None)) is int
- ):
- self._error(parameter, "A symbolic TypeVar bound must be int")
- if getattr(parameter, "default_value", None) is not None:
- self._error(parameter, "A symbolic TypeVar cannot have a default")
- self._symbol(parameter.name, parameter, shadow=True)
-
- def _resolve(self, node):
- # Inspect callable identity without executing the annotation or its arguments.
- if isinstance(node, ast.Name):
- return self.env.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__
- # Descriptor execution belongs to evaluation, not metadata discovery.
- return None if isinstance(value, property) else value
- return None
-
- @staticmethod
- def _arguments(call, constructor):
- try:
- parameters = list(inspect.signature(constructor).parameters.values())
- except (TypeError, ValueError):
- return {}
- positional = [
- p
- for p in parameters
- if p.kind
- in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
- ]
- fields = {
- parameter.name: value
- for parameter, value in zip(positional, call.args)
- if not isinstance(value, ast.Starred)
- }
- fields.update(
- {keyword.arg: keyword.value for keyword in call.keywords if keyword.arg is not None}
- )
- return fields
-
- def _cache_expression(self, node):
- name = f"__tvm_annotation_value_{self._counter}"
- self._counter += 1
- while name in self.env or name in self._used_names:
- name = f"__tvm_annotation_value_{self._counter}"
- self._counter += 1
- self.env[name] = self._eval(node)
- return ast.copy_location(ast.Name(name, ast.Load()), node)
-
- def prepare_parameters(self, arguments):
- """Prepare declared symbols and sequential scalar parameter annotations.
-
- Declaration constructors reserve identities before dependent annotations.
- Type constructors introduce their parameters in signature order. Cached
- dtype expressions are evaluated once. Bare symbolic strings declare shape
- names across the signature; compound expressions only reference them.
- """
- parameters = [*arguments.posonlyargs, *arguments.args, *arguments.kwonlyargs]
- self._used_names.update(
- node.id for node in ast.walk(arguments) if isinstance(node, ast.Name)
- )
- self._used_names.update(parameter.arg for parameter in parameters)
- for parameter in parameters:
- annotation = parameter.annotation
- if annotation is None:
- continue
- prepared = copy.deepcopy(annotation)
- if isinstance(prepared, ast.Constant) and isinstance(prepared.value, str):
- prepared = self._string_expression(prepared)
- self._prepared[id(annotation)] = prepared
- constructor = self._resolve(
- prepared.func if isinstance(prepared, ast.Call) else prepared
- )
- declaration = getattr(constructor, "__tvm_declaration_args__", None)
- dtype_field = getattr(constructor, "__tvm_parameter_dtype__", None)
- if declaration is not None:
- dtype = declaration.dtype
- elif dtype_field is not None and isinstance(prepared, ast.Call):
- dtype_node = self._arguments(prepared, constructor).get(dtype_field)
- if dtype_node is None:
- dtype = inspect.signature(constructor).parameters[dtype_field].default
- if dtype is inspect.Parameter.empty:
- continue
- else:
- cached = self._cache_expression(dtype_node)
- dtype = self.env[cached.id]
- for index, argument in enumerate(prepared.args):
- if argument is dtype_node:
- prepared.args[index] = cached
- for keyword in prepared.keywords:
- if keyword.value is dtype_node:
- keyword.value = cached
- else:
- continue
- if declaration is not None:
- self._symbol(parameter.arg, parameter, dtype, shadow=True)
- self._parameter_annotations[id(annotation)] = parameter.arg
- self._unbound_parameters.add(parameter.arg)
- else:
- self._pending_parameters[id(annotation)] = (parameter, dtype)
- for prepared in self._prepared.values():
- self.rewrite(prepared, introduce=True, collect_declarations=True)
- return self.symbols
-
- def _eval(self, node):
- expression = ast.Expression(body=node)
- ast.fix_missing_locations(expression)
- return eval(compile(expression, self.filename, "eval"), self.env) # pylint: disable=eval-used
-
- def evaluate(self, node, introduce=True):
- """Evaluate an annotation once in the prepared signature scope."""
- key = id(node)
- if key not in self._evaluated:
- self._unbound_parameters.discard(self._parameter_annotations.get(key))
- if key in self._pending_parameters:
- parameter, dtype = self._pending_parameters[key]
- if parameter.arg in self.symbols:
- self._error(
- parameter, "A later parameter cannot adopt an existing shape symbol"
- )
- self._symbol(parameter.arg, parameter, dtype, shadow=True)
- prepared = self._prepared.get(key, node)
- # A quoted whole annotation is ordinary Python annotation syntax.
- if isinstance(prepared, ast.Constant) and isinstance(prepared.value, str):
- prepared = self._string_expression(prepared)
- self._evaluated[key] = self._eval(self.rewrite(prepared, introduce=introduce))
- return self._evaluated[key]
-
- def _string_expression(self, node):
- try:
- expression = ast.parse(node.value, mode="eval").body
- except SyntaxError as error:
- self._error(node, f"Invalid annotation expression: {error.msg}")
- 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:
- column += node.col_offset + 1 if line == 1 else 0
- line += node.lineno - 1
- 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(self, node, *, introduce=False, collect_declarations=False):
- """Return a copied expression AST, registering new symbols in ``env``.
-
- Construction code must execute with this scope's updated environment.
- No assignment to a generated local is needed for newly introduced names.
- """
- scope = self
-
- class Rewrite(ast.NodeTransformer):
- def __init__(self):
- self.allow_names = False
- self.in_string = False
- self.dtype = None
- self.declaration_dtype = None
- self.compound_declarations = False
-
- def visit_Name(self, current):
- if (
- self.allow_names
- and self.in_string
- and self.compound_declarations
- and isinstance(current.ctx, ast.Load)
- and current.id not in scope.env
- and not hasattr(builtins, current.id)
- ):
- scope._shape_declarations.setdefault(
- current.id, (current, self.declaration_dtype)
- )
- if collect_declarations:
- return current
- if isinstance(current.ctx, ast.Load):
- if not self.in_string and current.id in scope._unbound_parameters:
- scope._error(current, f"Parameter {current.id!r} is not yet bound")
- if (
- self.allow_names
- and not self.in_string
- and current.id not in scope.env
- and not hasattr(builtins, current.id)
- ):
- scope._error(current, f"Name {current.id!r} is not defined")
- if isinstance(scope.env.get(current.id), TypeVar) and not introduce:
- scope._error(
- current, "A TypeVar must be introduced in a signature or match scope"
- )
- scope._canonical_type_var(current.id, current, self.dtype)
- if self.allow_names and current.id in scope.env:
- scope._symbol(current.id, current, self.dtype)
- if self.in_string:
- scope._unbound_parameters.discard(current.id)
- elif (
- self.allow_names
- and self.in_string
- and current.id in scope._shape_declarations
- ):
- declaration, dtype = scope._shape_declarations[current.id]
- scope._symbol(current.id, declaration, dtype)
- return current
-
- def visit_Attribute(self, current):
- old_allow = self.allow_names
- self.allow_names = False
- current.value = self.visit(current.value)
- self.allow_names = old_allow
- return current
-
- def expression_field(self, current, metadata, *, nested=False):
- old_allow, old_dtype, old_string = self.allow_names, self.dtype, self.in_string
- old_compound, old_declaration_dtype = (
- self.compound_declarations,
- self.declaration_dtype,
- )
- self.allow_names = introduce and metadata.introduce
- self.dtype = metadata.dtype
- self.compound_declarations = metadata.compound_declarations
- self.declaration_dtype = (
- metadata.dtype if metadata.dtype is not None else metadata.implicit_dtype
- )
- try:
- if isinstance(current, ast.List | ast.Tuple):
- current.elts = [
- self.expression_field(value, metadata, nested=True)
- for value in current.elts
- ]
- return current
- if isinstance(current, ast.Constant) and isinstance(current.value, str):
- if nested or metadata.scalar_strings:
- current = scope._string_expression(current)
- self.in_string = True
- if self.allow_names and isinstance(current, ast.Name):
- scope._shape_declarations.setdefault(
- current.id,
- (
- current,
- self.dtype
- if self.dtype is not None
- else metadata.implicit_dtype,
- ),
- )
- return self.visit(current)
- finally:
- self.allow_names, self.dtype, self.in_string = old_allow, old_dtype, old_string
- self.compound_declarations, self.declaration_dtype = (
- old_compound,
- old_declaration_dtype,
- )
-
- def visit_Call(self, current):
- constructor = scope._resolve(current.func)
- metadata = getattr(constructor, "__tvm_expression_args__", None)
- fields = scope._arguments(current, constructor) if metadata else {}
- marked = (
- {id(value) for name, value in fields.items() if name in metadata.fields}
- if metadata
- else set()
- )
- old_allow = self.allow_names
- self.allow_names = False
- current.func = self.visit(current.func)
- self.allow_names = old_allow
- current.args = [
- self.expression_field(value, metadata)
- if id(value) in marked
- else self.visit(value)
- for value in current.args
- ]
- for keyword in current.keywords:
- keyword.value = (
- self.expression_field(keyword.value, metadata)
- if id(keyword.value) in marked
- else self.visit(keyword.value)
- )
- return current
-
- return ast.fix_missing_locations(Rewrite().visit(copy.deepcopy(node)))
-
-
-@lru_cache(maxsize=128)
-def _source_tree(filename, text):
- return ast.parse(text, filename)
-
-
-def _signature_at(frame, builder):
- if IRBuilder.is_in_scope():
- return None
- text = "".join(linecache.getlines(frame.f_code.co_filename))
- if not text:
- return None
- tree = _source_tree(frame.f_code.co_filename, text)
- env = {**frame.f_globals, **frame.f_locals}
- scope = AnnotationScope(env, builder, frame.f_code.co_filename, lambda node: None)
- for node in ast.walk(tree):
- if not isinstance(node, ast.FunctionDef):
- continue
- if not (node.lineno <= frame.f_lineno <= node.body[0].lineno):
- continue
- if frame.f_code.co_name == node.name:
- continue
- decorators = [
- item.func if isinstance(item, ast.Call) else item for item in node.decorator_list
- ]
- if any(getattr(scope._resolve(item), "__tvm_function_kind__", None) for item in decorators):
- return node, scope
- return None
-
-
-def take_eager_annotations(function):
- """Transfer signature construction state from the defining Python frame."""
- frame = inspect.currentframe().f_back
- try:
- while frame is not None:
- states = frame.f_locals.get("__tvm_eager_annotations__", {})
- scope = states.pop((function.__name__, function.__code__.co_filename), None)
- if scope is not None:
- if not states:
- frame.f_locals.pop("__tvm_eager_annotations__", None)
- function.__tvm_eager_annotations__ = scope
- return
- frame = frame.f_back
- finally:
- del frame
-
-
-def enable_eager_constructors(builder, *, classes=()):
- """Resolve registered symbolic arguments into concrete eager annotations.
-
- Constructor metadata and signature declarations are identical in both entry
- modes. Python evaluates eager annotations once; the parsing decorator retains
- their values and signature symbols for declaration and definition.
- """
- wrappers = {}
-
- def wrap(constructor):
- metadata = constructor.__tvm_expression_args__
- signature = inspect.signature(constructor)
- as_type = constructor.__name__ in classes
-
- @wraps(constructor)
- def invoke(*args, **kwargs):
- if as_type:
- args = args[1:]
- bound = signature.bind(*args, **kwargs)
- caller = inspect.currentframe().f_back
- state = None
- try:
- found = _signature_at(caller, builder)
- if found is not None:
- node, fresh_scope = found
- states = caller.f_locals.setdefault("__tvm_eager_annotations__", {})
- key = (node.name, caller.f_code.co_filename)
- state = states.get(key)
- site = (node.lineno, caller.f_lasti)
- # A host argument may fail before entering any constructor.
- # Re-entering a definition therefore starts a fresh signature,
- # even when its decorator never got a chance to consume it.
- if state is not None and (
- site[0] != state._eager_site[0] or site[1] <= state._eager_site[1]
- ):
- state = None
- if state is None:
- state = states[key] = fresh_scope
- for parameter in node.args.args:
- if parameter.annotation is not None:
- state.rewrite(
- parameter.annotation, introduce=True, collect_declarations=True
- )
- state._eager_site = site
- scope = state
- else:
- scope = AnnotationScope(
- {**caller.f_globals, **caller.f_locals},
- builder,
- caller.f_code.co_filename,
- lambda node: None,
- )
- finally:
- del caller
-
- def symbolic(value, nested=False):
- if isinstance(value, TypeVar):
- return True
- if isinstance(value, str):
- return nested or metadata.scalar_strings
- if isinstance(value, tuple | list):
- return any(symbolic(item, True) for item in value)
- return False
-
- if not any(symbolic(bound.arguments.get(name)) for name in metadata.fields):
- try:
- return constructor(*args, **kwargs)
- except Exception:
- if state is not None:
- states.pop(key, None)
- raise
- counter = 0
-
- def fresh(value):
- nonlocal counter
- while True:
- name = f"__tvm_eager_value_{counter}"
- counter += 1
- if name not in scope.env:
- scope.env[name] = value
- return ast.Name(name, ast.Load())
-
- def expression(value):
- if isinstance(value, str):
- return ast.Constant(value)
- if isinstance(value, tuple | list):
- kind = ast.Tuple if isinstance(value, tuple) else ast.List
- return kind([expression(item) for item in value], ast.Load())
- if isinstance(value, TypeVar):
- scope.env[value.__name__] = value
- return ast.Name(value.__name__, ast.Load())
- return fresh(value)
-
- call = ast.Call(
- fresh(constructor),
- [],
- [
- ast.keyword(
- name, expression(value) if name in metadata.fields else fresh(value)
- )
- for name, value in bound.arguments.items()
- ],
- )
- ast.fix_missing_locations(call)
- scope.rewrite(call, introduce=True, collect_declarations=True)
- try:
- return scope._eval(scope.rewrite(call, introduce=True))
- except Exception as error:
- if state is None:
- raise
- states.pop(key, None)
- from .diagnostics import diagnostic_error
-
- raise diagnostic_error(
- error,
- SimpleNamespace(
- filename=scope.filename, tree=ast.Module(body=[node], type_ignores=[])
- ),
- ) from error
-
- if as_type:
- # A construction class also supports ordinary Python type unions.
- # __new__ returns the concrete IR object, never an annotation proxy.
- return type(
- constructor.__name__,
- (),
- {
- "__new__": staticmethod(invoke),
- "__signature__": signature,
- "__module__": constructor.__module__,
- "__doc__": constructor.__doc__,
- "__tvm_expression_args__": metadata,
- },
- )
- return invoke
-
- for name, constructor in list(vars(builder).items()):
- if getattr(constructor, "__tvm_expression_args__", None) is not None:
- if constructor not in wrappers:
- wrappers[constructor] = wrap(constructor)
- setattr(builder, name, wrappers[constructor])
diff --git a/python/tvm/script/parser/diagnostics.py b/python/tvm/script/parser/diagnostics.py
index 430cac796a..25e0ab660a 100644
--- a/python/tvm/script/parser/diagnostics.py
+++ b/python/tvm/script/parser/diagnostics.py
@@ -23,12 +23,19 @@ from tvm.error import DiagnosticError
def diagnostic_error(error, compiler):
- """Render the innermost original operation and retain the original cause."""
+ """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
- span = getattr(error, "__tvm_script_span__", None)
- if span is not None:
- start, end = span.line, span.end_line
- column, end_column = span.column, span.end_column
+ 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
diff --git a/python/tvm/script/parser/expression.py b/python/tvm/script/parser/expression.py
new file mode 100644
index 0000000000..c8b31c0669
--- /dev/null
+++ b/python/tvm/script/parser/expression.py
@@ -0,0 +1,226 @@
+# 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 tvm.script.ir_builder 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, rewriting only registered expression-string fields.
+
+ Parameters
+ ----------
+ node : ast.expr
+ Original source expression. It is never mutated.
+ resolve : callable
+ Static namespace lookup for Name/Attribute nodes; returns registered
+ Python namespaces/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
+ If true, parse a quoted whole annotation before constructor-field rewriting.
+
+ Returns
+ -------
+ ast.expr
+ New AST preserving original ranges; no annotation is executed here.
+
+ Raises
+ ------
+ SyntaxError
+ A marked string is not a Python expression.
+ """
+ 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/frontend.py b/python/tvm/script/parser/frontend.py
index ffb9208d67..c613835855 100644
--- a/python/tvm/script/parser/frontend.py
+++ b/python/tvm/script/parser/frontend.py
@@ -22,39 +22,29 @@ import ast
import copy
import inspect
import linecache
+import sys
import textwrap
-from dataclasses import dataclass, field
from functools import wraps
-from types import SimpleNamespace
from typing import TypeVar
-import tvm
-from tvm import ir
from tvm.error import DiagnosticError
-from tvm.script.ir_builder import IRBuilder, protocol
-from tvm.script.ir_builder import ir as I
+from tvm.script.ir_builder import construction, protocol
-from .annotations import AnnotationScope, take_eager_annotations
from .diagnostics import diagnostic_error
-from .functions import (
- FunctionGroup,
- adapt_python_module,
- attach_python,
- declare_python,
- is_python_function,
-)
+from .transpile import IRBuilderTranspiler, NameCollector
_NAMESPACES = {}
def register_namespace(alias, namespace):
- """Let an entry module supply a source-level namespace without reverse imports."""
- _NAMESPACES[alias] = namespace
-
+ """Register a host namespace for source-text entry points.
-def _resolve(node, env, filename):
- expression = ast.Expression(copy.deepcopy(node))
- return eval(compile(ast.fix_missing_locations(expression), filename, "eval"), env)
+ alias is a Python identifier string; namespace is its module/object binding.
+ Returns None. The process-wide alias table is updated (last registration wins)
+ and copied into each new Compiler; existing compilations are unaffected.
+ This registers syntax metadata access only and enters no builder frame.
+ """
+ _NAMESPACES[alias] = namespace
def _closure_values(function):
@@ -106,7 +96,20 @@ def _inside_class(function):
def make_decorator(builder, *, option_map=None, defaults=None):
- """Create a parsing decorator whose construction policy belongs to its caller."""
+ """Create and register a function decorator for a construction namespace.
+
+ builder exposes the function construction protocol. option_map optionally
+ maps public keyword names to builder keyword names; defaults optionally maps
+ builder keyword names to default values. Both mappings default to empty and
+ are copied. Returns a callable usable as @decorator, @decorator(**options),
+ or decorator(function). Non-function positional arguments raise ValueError.
+
+ 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 original source annotations execute inside declaration frames.
+ Builder validation and source errors propagate through parse diagnostics.
+ """
mapping, default_options = dict(option_map or {}), dict(defaults or {})
def decorator(function=None, **options):
@@ -114,7 +117,6 @@ def make_decorator(builder, *, option_map=None, defaults=None):
raise ValueError("Construction decorators require a function or keyword options")
def apply(function):
- take_eager_annotations(function)
function.__tvm_function_kind__ = decorator.__tvm_function_kind__
function.__tvm_function_options__ = options
if _inside_class(function):
@@ -129,7 +131,20 @@ def make_decorator(builder, *, option_map=None, defaults=None):
def make_helper(builder, *, preserve_return=True, late_binding=False):
- """Create a construction helper, optionally retaining live Python closure cells."""
+ """Create a decorator for helpers executed inside a caller's builder frames.
+
+ builder is a construction namespace. preserve_return (default True) keeps
+ helper returns as ordinary Python control flow. late_binding (default False)
+ refreshes captured closure cells at each invocation. Returns a decorator
+ accepting a function or keyword options; hygienic=True snapshots its lexical
+ environment, while hygienic=False captures the calling environment each time.
+ Other options are retained as helper metadata for the namespace consumer.
+
+ Invocations bind ordinary Python arguments/defaults, transpile the original
+ body and return its result. The helper shares the caller's active construction
+ frames; it does not declare a new IR function. Invalid decorator arguments
+ raise ValueError; Python signature binding and builder errors propagate.
+ """
def decorator(function=None, **options):
if function is not None and not inspect.isfunction(function):
@@ -166,7 +181,12 @@ def make_helper(builder, *, preserve_return=True, late_binding=False):
def pyfunc(function):
- """Mark a function whose body and execution remain ordinary Python."""
+ """Mark and return a Python function for opaque registration in a module.
+
+ function must permit function attributes. Its body is never transpiled;
+ module construction delegates registration to the builder runtime. This
+ decorator enters no frame and preserves the original callable identity.
+ """
function.__tvm_python_function__ = True
return function
@@ -174,22 +194,23 @@ def pyfunc(function):
protocol.register_function(pyfunc, None, python=True)
-@dataclass
-class Signature:
- node: ast.FunctionDef
- kind: protocol.FunctionKind
- options: dict
- scope: AnnotationScope
- params: dict = field(default_factory=dict)
- result_type: object = protocol.MISSING
- reference: object = None
+class Compiler:
+ """Acquire source and execute a location-preserving builder program.
+ source is a source string, Python function, or Python class. env is an optional
+ mapping of opaque lexical bindings, overriding registered source aliases.
+ filename optionally overrides the original source filename ("<str>" for text).
+ Source inspection, parsing and compilation errors propagate to parse().
-class Compiler:
- """Build from a copied original AST, retaining file and range information."""
+ env and original belong to this compilation only. tree is a fresh source AST;
+ filename and compile_flags retain the source's file and annotation mode.
+ name_map holds reserved identifiers and next-prefix counters, shared across
+ all generated functions. builder_name/infrastructure_name are reserved aliases.
+ This entry owns no IR state: build() passes its execution result opaquely.
+ """
def __init__(self, source, env=None, filename=None):
- self.env = {"TypeVar": TypeVar, "tvm": tvm, **_NAMESPACES, **(env or {})}
+ self.env = {"TypeVar": TypeVar, "tvm": sys.modules.get("tvm"), **_NAMESPACES, **(env or {})}
self.original = source
members = vars(source).values() if inspect.isclass(source) else (source,)
self.compile_flags = 0
@@ -220,296 +241,120 @@ class Compiler:
if hasattr(node, "col_offset"):
node.col_offset += indent
node.end_col_offset += indent
- self.used_names = {n.id for n in ast.walk(self.tree) if isinstance(n, ast.Name)}
- self.used_names.update(n.arg for n in ast.walk(self.tree) if isinstance(n, ast.arg))
- self.used_names.update(
- n.name for n in ast.walk(self.tree) if isinstance(n, ast.FunctionDef | ast.ClassDef)
- )
- self.used_names.update(self.env)
- self.counter = 0
- self.function_kinds = {}
- self.spans = []
- self.span_indices = {}
- self.source_name = ir.SourceName(self.filename)
- self.span_name = self.fresh("spans")
- self.builder_name = self.fresh("builder")
- self.infrastructure_name = self.fresh("infrastructure")
-
- def fresh(self, prefix):
- while True:
- self.counter += 1
- name = f"__script_{prefix}_{self.counter}"
- if name not in self.used_names:
- self.used_names.add(name)
- return name
-
- def span(self, node):
- key = (node.lineno, node.end_lineno, node.col_offset, node.end_col_offset)
- if key not in self.span_indices:
- self.span_indices[key] = len(self.spans)
- self.spans.append(ir.Span(self.source_name, *key))
- return self.spans[self.span_indices[key]]
+ # 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()
+
+ def fresh(self, prefix="_t"):
+ """Allocate a generated identifier without changing any source name.
+
+ prefix is an identifier prefix, default '_t'. Returns an unused identifier;
+ name_map stores reserved names (value 0) and next counters for prefixes.
+ Its lifetime is this complete module/function compilation 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):
- self.span(node)
- index = self.span_indices[
- (node.lineno, node.end_lineno, node.col_offset, node.end_col_offset)
- ]
+ """Emit location data only; builders materialize source spans at execution.
+
+ node supplies original line/end-line and UTF-8 column/end-column fields.
+ The returned AST tuple keeps those values plus the original filename;
+ generated wrappers inherit the full node range via copy_location.
+ """
return ast.copy_location(
- ast.Subscript(ast.Name(self.span_name, ast.Load()), ast.Constant(index), ast.Load()),
+ ast.Tuple(
+ [
+ ast.Constant(value)
+ for value in (
+ self.filename,
+ node.lineno,
+ node.end_lineno,
+ node.col_offset,
+ node.end_col_offset,
+ )
+ ],
+ ast.Load(),
+ ),
node,
)
- def function_kind(self, node, env, *, allow_python=False):
- if id(node) in self.function_kinds:
- return self.function_kinds[id(node)]
- if inspect.isfunction(self.original) and node is self.tree.body[0]:
- kind = protocol.function_kind(self.original)
- if kind is not None:
- options = {
- **kind.metadata.get("defaults", {}),
- **getattr(self.original, "__tvm_function_options__", {}),
- }
- mapping = kind.metadata.get("option_map", {})
- result = (
- kind,
- {
- mapping.get(key, key): value
- for key, value in options.items()
- if key != "check_well_formed"
- },
- )
- self.function_kinds[id(node)] = result
- return result
- for decorator in node.decorator_list:
- target = decorator.func if isinstance(decorator, ast.Call) else decorator
- value = _resolve(target, env, self.filename)
- kind = protocol.function_kind(value)
- if kind is not None:
- options = dict(kind.metadata.get("defaults", {}))
- if isinstance(decorator, ast.Call):
- if decorator.args:
- raise SyntaxError("Function decorators accept keyword options only")
- for item in decorator.keywords:
- if item.arg is None:
- options.update(_resolve(item.value, env, self.filename))
- else:
- options[item.arg] = _resolve(item.value, env, self.filename)
- mapping = kind.metadata.get("option_map", {})
- options = {
- mapping.get(key, key): value
- for key, value in options.items()
- if key != "check_well_formed"
- }
- self.function_kinds[id(node)] = (kind, options)
- return kind, options
- if allow_python:
- return protocol.FunctionKind(None, {"python": True}), {}
- raise SyntaxError(f"Function {node.name!r} has no registered construction kind")
-
- def declare(self, node, env, *, local=False):
- kind, options = self.function_kind(node, env)
- scope = AnnotationScope(env, kind.builder, self.filename, self.span)
- original = (
- getattr(self.original, node.name, None)
- if inspect.isclass(self.original)
- else self.original
- if inspect.isfunction(self.original) and node is self.tree.body[0]
- else None
+ def transformer(self, builder_name=None, **options):
+ """Create a syntax transformer sharing this unit's name allocator.
+
+ builder_name optionally selects an injected builder alias; options are
+ syntax-handler configuration passed to IRBuilderTranspiler. Returns a new
+ transformer with function-local binding/export analysis and no IR state.
+ """
+ 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,
+ **options,
)
- eager = inspect.isfunction(original) and not (
- original.__code__.co_flags & __future__.annotations.compiler_flag
- )
- eager_scope = getattr(original, "__tvm_eager_annotations__", None)
- if eager_scope is not None:
- scope.symbols.update(eager_scope.symbols)
- scope.env.update(eager_scope.symbols)
- spec = Signature(node, kind, options, scope)
- if kind.metadata.get("python"):
- return spec
- X = kind.builder
- mode = {"local": True} if local else {}
- with X.decl_function(**options, **mode, span=self.span(node)) as frame:
- X.func_name(node.name)
- scope.prepare_type_params(getattr(node, "type_params", []))
- arguments = node.args
- if eager:
- arguments = copy.copy(node.args)
- arguments.args = [
- argument
- for argument in node.args.args
- if isinstance(original.__annotations__.get(argument.arg), str)
- ]
- scope.prepare_parameters(arguments)
- if node.args.posonlyargs or node.args.kwonlyargs or node.args.vararg or node.args.kwarg:
- raise SyntaxError("IR signatures require ordinary named parameters")
- for argument in node.args.args:
- if argument.annotation is None:
- raise SyntaxError(f"Parameter {argument.arg!r} requires an annotation")
- annotation = (
- original.__annotations__[argument.arg]
- if eager
- else scope.evaluate(argument.annotation)
- )
- if eager and isinstance(annotation, str):
- annotation = scope.evaluate(argument.annotation)
- value = X.arg(
- argument.arg,
- scope.symbols.get(argument.arg, annotation),
- span=self.span(argument),
- )
- spec.params[argument.arg] = scope.env[argument.arg] = value
- if node.returns is not None:
- spec.result_type = (
- original.__annotations__["return"]
- if eager
- else scope.evaluate(node.returns, introduce=False)
- )
- if eager and isinstance(spec.result_type, str):
- spec.result_type = scope.evaluate(node.returns, introduce=False)
- X.func_ret_type(spec.result_type)
- spec.reference = frame.reference
- scope.env[node.name] = spec.reference
- return spec
-
- def define(self, spec, env, *, local=False):
- X = spec.kind.builder
- mode = {"local": True, "reference": spec.reference} if local else {}
- with X.function(**spec.options, **mode, span=self.span(spec.node)) as frame:
- X.func_name(spec.node.name)
- for name, value in spec.params.items():
- X.arg(name, value)
- if spec.result_type is not protocol.MISSING:
- X.func_ret_type(spec.result_type)
- scope_env = {**env, **spec.scope.env, **spec.params}
- self.run_statements(
- spec.node.body,
- X,
- scope_env,
- set(spec.params) | set(spec.scope.symbols),
- scope=spec.scope,
- )
- result = frame.function
- result.__name__ = spec.node.name
- return result
- def run_statements(self, body, builder, env, bound_names, *, scope=None, preserve_return=False):
- from .transform import Transformer
+ def function_kind(self, node, env=None, *, allow_python=False):
+ """Return registered syntax metadata and unevaluated option AST.
+ node is a FunctionDef; env is accepted for source-compatibility and unused.
+ allow_python permits ordinary nested helpers. No annotation or option
+ expression is evaluated by this lookup.
+ """
+ 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 the same AST-only lowering.
+
+ body is an original statement list, builder its registered namespace,
+ env an opaque execution namespace, bound_names supplied Python parameters,
+ and preserve_return whether helper returns retain ordinary Python control.
+ Returns the generated helper's opaque result; execution errors propagate.
+ """
namespace = dict(env)
- namespace.update(
- {
- self.builder_name: builder,
- self.infrastructure_name: protocol,
- self.span_name: self.spans,
- }
- )
- nested = {}
- nested_name = self.fresh("nested")
-
- def nested_statement(node):
- kind, _ = self.function_kind(node, namespace, allow_python=True)
- if kind.metadata.get("python"):
- return [copy.deepcopy(node)]
- index = len(nested)
- nested[index] = node
- value = ast.Call(
- ast.Name(nested_name, ast.Load()),
- [
- ast.Constant(index),
- ast.Call(
- ast.Attribute(
- ast.Name(self.infrastructure_name, ast.Load()), "locals", ast.Load()
- ),
- [],
- [],
- ),
- ],
- [],
- )
- return [ast.copy_location(ast.Assign([ast.Name(node.name, ast.Store())], value), node)]
-
- def build_nested(index, values):
- node = nested[index]
- local_env = {**namespace, **values}
- group = FunctionGroup(self, [node], local_env, local=True)
- return group.define(node.name)
-
- namespace[nested_name] = build_nested
- injected = {}
-
- def rewrite_expression(node):
- node = scope.rewrite(node) if scope is not None else copy.deepcopy(node)
- method = None
- values = []
- if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
- method, values = "logical_not", [node.operand]
- if method is not None:
- node = ast.copy_location(
- ast.Call(
- ast.Attribute(ast.Name(self.builder_name, ast.Load()), method, ast.Load()),
- values,
- [],
- ),
- node,
- )
- return node
-
- def rewrite_iterable(node):
- node = copy.deepcopy(node)
- if isinstance(node, ast.Call):
- target = scope._resolve(node.func) if scope is not None else None
- try:
- replacement = getattr(builder, "__tvm_call_overrides__", {}).get(target)
- except TypeError:
- replacement = None
- if replacement is not None:
- name = self.fresh("call")
- injected[name] = replacement
- node.func = ast.copy_location(ast.Name(name, ast.Load()), node.func)
- return node
-
- signature_values = {}
- if scope is not None:
- for name, value in scope.symbols.items():
- alias = self.fresh("symbol")
- namespace[alias] = value
- signature_values[name] = ast.Name(alias, ast.Load())
-
- transformer = Transformer(
- filename=self.filename,
- environment=namespace,
- builder_name=self.builder_name,
- infrastructure_name=self.infrastructure_name,
- span=self.span_ast,
- fresh=self.fresh,
- signature_names=set(bound_names),
- signature_values=signature_values,
- expression_rewriter=rewrite_expression,
- iterable_rewriter=rewrite_iterable,
- nested_function=nested_statement,
- preserve_return=preserve_return,
- )
- statements = transformer.transform_statements(copy.deepcopy(body))
- namespace.update(injected)
- if scope is not None:
- namespace.update(scope.env)
- bound_names = set(bound_names) | set(scope.symbols)
+ namespace.update({self.builder_name: builder, self.infrastructure_name: protocol})
+ 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("body")
- helper = ast.FunctionDef(
- name=helper_name,
- args=ast.arguments(
- posonlyargs=[],
- args=[ast.arg(name) for name in names],
- kwonlyargs=[],
- kw_defaults=[],
- defaults=[],
+ 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=statements or [ast.Pass()],
- decorator_list=[],
+ body[0],
)
- ast.copy_location(helper, 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),
@@ -518,98 +363,72 @@ class Compiler:
return namespace[helper_name](*(namespace[name] for name in names))
def build(self):
- nodes = self.tree.body
- env = dict(self.env)
- if nodes and isinstance(nodes[-1], ast.FunctionDef | ast.ClassDef):
- prefix, nodes = nodes[:-1], nodes[-1:]
- if any(isinstance(node, ast.FunctionDef | ast.ClassDef) for node in prefix):
- raise SyntaxError("Source must contain one function or module class")
- if prefix:
- setup = ast.fix_missing_locations(ast.Module(copy.deepcopy(prefix), []))
- exec(compile(setup, self.filename, "exec"), env)
- if len(nodes) == 1 and isinstance(nodes[0], ast.ClassDef):
- root = nodes[0]
- statements = root.body
- elif len(nodes) == 1 and isinstance(nodes[0], ast.FunctionDef):
- root = None
- statements = nodes
- else:
- raise SyntaxError("Source must contain one function or module class")
- functions = [node for node in statements if isinstance(node, ast.FunctionDef)]
- python_functions = []
- with IRBuilder() as builder:
- with I.ir_module():
- references = {node.name: I.reserve_function(node.name) for node in functions}
- env.update(references)
- if root is not None:
- env[root.name] = SimpleNamespace(**references)
- for statement in statements:
- if not isinstance(statement, ast.FunctionDef):
- # Eager signatures and the module share the original metadata
- # objects, including identity-sensitive device meshes.
- infos = getattr(self.original, "__tvm_script_global_infos__", None)
- if (
- infos is not None
- and isinstance(statement, ast.Expr)
- and isinstance(statement.value, ast.Call)
- and _resolve(statement.value.func, env, self.filename)
- is I.module_global_infos
- ):
- I.module_global_infos(infos)
- continue
- exec(
- compile(
- ast.fix_missing_locations(
- ast.Module([copy.deepcopy(statement)], [])
- ),
- self.filename,
- "exec",
- ),
- env,
- )
- if isinstance(statement, ast.Assign | ast.AnnAssign):
- targets = (
- statement.targets
- if isinstance(statement, ast.Assign)
- else [statement.target]
- )
- for target in targets:
- if isinstance(target, ast.Name):
- value = env[target.id]
- if isinstance(value, ir.BaseFunc):
- reference = I.decl_function(target.id, value)
- I.def_function(target.id, value)
- env[target.id] = reference
- setattr(env[root.name], target.id, env[target.id])
- ir_functions = []
- for node in functions:
- if is_python_function(self, node, env):
- original = self.env.get(node.name)
- original = (
- original
- if getattr(original, "__tvm_python_function__", False)
- else None
- )
- record = declare_python(self, node, env, original=original)
- python_functions.append(record)
- references[node.name] = env[node.name] = record.reference
- else:
- ir_functions.append(node)
- group = FunctionGroup(self, ir_functions, env)
- results = group.define_all()
- module = builder.get()
- if python_functions:
- attach_python(module, python_functions)
- if root is not None:
- module.__name__ = root.name
- original = self.original if inspect.isclass(self.original) else None
- bases = tuple(_resolve(base, env, self.filename) for base in root.bases)
- return adapt_python_module(module, original=original, bases=bases)
- return results[functions[0].name]
+ """Compile original-location builder AST and return its opaque result.
+
+ Compilation uses the original filename and original AST coordinates,
+ never ast.unparse/reparse. copy_location preserves all four location
+ fields on replaced nodes; fix_missing_locations fills only absent fields
+ on generated helpers. Python exceptions therefore name the source file
+ and original expression range (column detail requires Python 3.11+).
+ """
+ # 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 protocol.function_kind(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: protocol,
+ }
+ 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, **options):
- """Construct from source using only entry-module registered construction policies."""
+ """Transpile and execute a source string, Python function, or Python class.
+
+ source is the original object/text; extra_vars optionally supplies lexical
+ bindings overriding captured values. filename optionally overrides the source
+ filename. Remaining options are accepted for entry-point compatibility; actual
+ construction policy comes from registered decorators in the source.
+ Returns the generated builder program's opaque result. Source/host/builder
+ failures become DiagnosticError with original-source ranges; an existing
+ DiagnosticError is preserved. Source acquisition errors propagate directly.
+ Builder declaration/definition frames are entered only by generated execution.
+ """
env = {} if isinstance(source, str) else _capture(source)
env.update(extra_vars or {})
compiler = Compiler(source, env, filename)
@@ -622,7 +441,14 @@ def parse(source, extra_vars=None, *, filename=None, **options):
def ir_module(module=None, **options):
- """Build a module after its registered function signatures have been declared."""
+ """Decorate a Python class with two-phase module construction.
+
+ module is a class or None; options are forwarded to parse. With a class,
+ returns the generated module result immediately; with None, returns a class
+ decorator. Class host bindings are captured before transpilation. Generated
+ execution declares all registered signatures before defining their bodies.
+ Parsing and builder errors follow parse() diagnostics and frame cleanup.
+ """
def apply(module):
return parse(module, _capture(module), **options)
diff --git a/python/tvm/script/parser/functions.py b/python/tvm/script/parser/functions.py
deleted file mode 100644
index cce0dbfa19..0000000000
--- a/python/tvm/script/parser/functions.py
+++ /dev/null
@@ -1,237 +0,0 @@
-# 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-group construction and opaque Python member registration."""
-
-import ast
-import copy
-import dis
-import inspect
-import linecache
-from dataclasses import dataclass
-from types import CodeType
-
-from tvm.script.ir_builder import ir as I
-
-_OPAQUE_FACTORY = None
-_MODULE_ADAPTER = None
-
-
-def register_opaque_factory(factory, *, module_adapter=None):
- """Register the owner-provided constructor for a module's opaque function slot.
-
- The factory receives ``(name, python_callable, source_text, span)`` and must
- return a concrete BaseFunc. The optional adapter receives a completed module,
- its original class (if available), and its resolved base classes. Registration
- never executes the Python body.
- """
- global _OPAQUE_FACTORY, _MODULE_ADAPTER
- if not callable(factory):
- raise TypeError("An opaque function factory must be callable")
- if module_adapter is not None and not callable(module_adapter):
- raise TypeError("An opaque module adapter must be callable")
- _OPAQUE_FACTORY = factory
- _MODULE_ADAPTER = module_adapter
-
-
-def is_python_function(compiler, node, env):
- """Identify a Python function through its registered function-kind metadata."""
- kind, _ = compiler.function_kind(node, env)
- return bool(kind.metadata.get("python"))
-
-
-def _error(compiler, node, message):
- raise SyntaxError(
- message,
- (
- compiler.filename,
- node.lineno,
- node.col_offset + 1,
- linecache.getline(compiler.filename, node.lineno),
- ),
- )
-
-
-def materialize_python(compiler, node, env, *, original=None):
- """Retain an existing Python callable, or execute its unchanged definition.
-
- Source-only module members share ``env``, preserving ordinary global lookup.
- Nested Python definitions should remain directly in the construction helper's
- AST, so Python itself creates their lexical closure cells.
- """
- if original is None and not isinstance(compiler.original, str):
- candidate = env.get(node.name)
- if getattr(candidate, "__tvm_python_function__", False):
- original = candidate
- if original is not None:
- if not callable(original):
- _error(compiler, node, "A Python function definition must retain a callable")
- return original
- module = ast.fix_missing_locations(ast.Module([copy.deepcopy(node)], []))
- exec(
- compile(
- module,
- compiler.filename,
- "exec",
- flags=getattr(compiler, "compile_flags", 0),
- dont_inherit=True,
- ),
- env,
- )
- return env[node.name]
-
-
-@dataclass(frozen=True)
-class PythonFunction:
- """Original Python callable and its separate opaque module reference."""
-
- name: str
- function: object
- reference: object
-
-
-def declare_python(compiler, node, env, *, original=None):
- """Construct a real opaque module entry and retain its Python implementation."""
- if _OPAQUE_FACTORY is None:
- _error(compiler, node, "No opaque Python function constructor has been registered")
- function = materialize_python(compiler, node, env, original=original)
- try:
- source = inspect.getsource(function)
- except (OSError, TypeError):
- source = ast.unparse(node)
- opaque = _OPAQUE_FACTORY(node.name, function, source, compiler.span(node))
- reference = I.decl_function(node.name, opaque)
- I.def_function(node.name, opaque)
- return PythonFunction(node.name, function, reference)
-
-
-def attach_python(module, functions):
- """Attach Python callables to a completed module with opaque entries.
-
- This preserves construction and Python execution. Device conversion, method
- binding, and runtime registration remain responsibilities of the runtime's
- module adapter; attaching callables alone does not establish those bridges.
- """
- existing = dict(getattr(module, "pyfuncs", {}))
- for function in functions:
- existing[function.name] = function.function
- module.pyfuncs = existing
- return module
-
-
-def adapt_python_module(module, *, original=None, bases=()):
- """Let the registered owner supply an executable module wrapper if needed."""
- if _MODULE_ADAPTER is None:
- return module
- return _MODULE_ADAPTER(module, original, bases)
-
-
-def _global_loads(node, filename):
- """Find free sibling uses without confusing attributes or shadowed locals."""
- node = copy.deepcopy(node)
- node.decorator_list = []
- node.returns = None
- for argument in (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs):
- argument.annotation = None
- module = ast.fix_missing_locations(ast.Module([node], []))
- code = compile(module, filename, "exec", dont_inherit=True)
- names = set()
-
- def visit(current):
- for instruction in dis.get_instructions(current):
- if instruction.opname in ("LOAD_GLOBAL", "LOAD_NAME"):
- names.add(instruction.argval)
- for value in current.co_consts:
- if isinstance(value, CodeType):
- visit(value)
-
- # Skip the module's decoration/default expressions: only body references
- # determine whether a local definition uses a still-undefined peer.
- for value in code.co_consts:
- if isinstance(value, CodeType) and value.co_name == node.name:
- visit(value)
- return names
-
-
-class FunctionGroup:
- """Declare sibling IR signatures before defining any of their bodies.
-
- An optional ``reserve(name)`` callback allocates stable identities before
- annotation evaluation, when the enclosing module or builder supports it.
- Local definitions retain source ordering and reject forward sibling uses;
- self-recursion uses the function's own completed declaration.
- """
-
- def __init__(self, compiler, nodes, env, *, local=False, reserve=None):
- self.compiler = compiler
- self.local = local
- self.env = dict(env)
- self.signatures = {}
- self.references = {}
- self.results = {}
- nodes = list(nodes)
- names = {node.name for node in nodes}
- if len(names) != len(nodes):
- duplicate = next(
- node
- for index, node in enumerate(nodes)
- if node.name in {previous.name for previous in nodes[:index]}
- )
- _error(compiler, duplicate, f"Duplicate function declaration {duplicate.name!r}")
- if reserve is not None:
- self.references.update((node.name, reserve(node.name)) for node in nodes)
- self.env.update(self.references)
- for node in nodes:
- signature = compiler.declare(node, self.env, local=local)
- self.signatures[node.name] = signature
- self.references[node.name] = signature.reference
- self.env[node.name] = signature.reference
- for signature in self.signatures.values():
- shadowed = signature.params.keys() | signature.scope.symbols.keys()
- signature.scope.env.update(
- (name, reference)
- for name, reference in self.references.items()
- if name not in shadowed
- )
-
- def define(self, name, env=None):
- """Define one declared body, returning the cached construction reference."""
- signature = self.signatures[name]
- if name in self.results:
- _error(self.compiler, signature.node, f"Function {name!r} is already defined")
- if self.local:
- unresolved = (
- (_global_loads(signature.node, self.compiler.filename) & self.references.keys())
- - self.results.keys()
- - {name}
- )
- if unresolved:
- _error(
- self.compiler,
- signature.node,
- f"Local function {name!r} refers to undefined sibling "
- f"{sorted(unresolved)[0]!r}; "
- "mutually recursive local definitions are unsupported",
- )
- visible = {**self.env, **(env or {}), **self.references}
- self.results[name] = self.compiler.define(signature, visible, local=self.local)
- return self.references[name]
-
- def define_all(self):
- """Define bodies in source order after every signature has been registered."""
- for name in self.signatures:
- self.define(name)
- return self.results
diff --git a/python/tvm/script/parser/transform.py b/python/tvm/script/parser/transform.py
deleted file mode 100644
index b3270c45fc..0000000000
--- a/python/tvm/script/parser/transform.py
+++ /dev/null
@@ -1,782 +0,0 @@
-# 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 copy
-import inspect
-
-
-class Transformer(ast.NodeTransformer):
- """Lower statements without importing or discovering construction namespaces."""
-
- def __init__(
- self,
- filename,
- environment,
- builder_name,
- infrastructure_name,
- span,
- fresh,
- signature_names=None,
- expression_rewriter=None,
- nested_function=None,
- preserve_return=False,
- signature_values=None,
- iterable_rewriter=None,
- ):
- self.filename = filename
- self.environment = dict(environment)
- self.builder_name = builder_name
- self.infrastructure_name = infrastructure_name
- self.span = span
- self.fresh = fresh
- self.signature_names = set(signature_names or ())
- self.signature_values = dict(signature_values or {})
- self.bound = set(self.signature_names)
- self.optional = {}
- self.expression_rewriter = expression_rewriter
- self.iterable_rewriter = iterable_rewriter
- self.nested_function = nested_function
- self.preserve_return = preserve_return
-
- def transform_statements(self, body):
- """Transform a copy, retaining the caller's original source tree."""
- result = []
- for statement in copy.deepcopy(body):
- translated = self.visit(statement)
- if translated is not None:
- block = translated if isinstance(translated, list) else [translated]
- context = self._call(
- self.infrastructure_name, "span_context", [self.span(statement)], statement
- )
- result.append(self._with(context, block, statement))
- 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):
- 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.builder_name, 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.environment.get(node.id)
- 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 _is_declaration(self, node):
- if not isinstance(node, ast.Call):
- return False
- constructor = self._resolve(node.func)
- policy = getattr(constructor, "__tvm_declaration_args__", None)
- if policy is None or any(isinstance(arg, ast.Starred) for arg in node.args):
- return False
- if any(keyword.arg is None for keyword in node.keywords):
- return False
- try:
- arguments = inspect.signature(constructor).bind_partial(
- *node.args, **{keyword.arg: keyword.value for keyword in node.keywords}
- )
- except (TypeError, ValueError):
- return False
- value = arguments.arguments.get(policy.value_parameter)
- return value is None or (isinstance(value, ast.Constant) and value.value is None)
-
- def _declaration_pattern(self, node):
- if isinstance(node, ast.Tuple | ast.List):
- return [self._declaration_pattern(item) for item in node.elts]
- return self._is_declaration(node)
-
- def _expression(self, original, *, attach_span=True, declarations=False):
- 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
- )
- if self.expression_rewriter is not None:
- node = self.expression_rewriter(node)
- # Policy-generated scaffolding inherits the replaced expression range;
- # original descendants retain their more precise source coordinates.
- if not hasattr(node, "lineno"):
- ast.copy_location(node, original)
- ast.fix_missing_locations(node)
- if isinstance(node, ast.Await | ast.Yield | ast.YieldFrom | ast.NamedExpr):
- self._error(original, f"Unsupported expression: {type(node).__name__}")
- if isinstance(node, ast.BoolOp):
- method = "logical_and" if isinstance(node.op, ast.And) else "logical_or"
- operands = [self._lambda([], self._expression(value), value) for value in node.values]
- node = self._call(
- self.infrastructure_name,
- "logical_chain",
- [
- self._attribute(self.builder_name, method, node),
- ast.Tuple(operands, ast.Load()),
- ast.Constant(isinstance(node.op, ast.Or)),
- ],
- node,
- )
- elif isinstance(node, ast.IfExp):
- node = self._call(
- self.infrastructure_name,
- "select_lazy",
- [
- self._attribute(self.builder_name, "select", node),
- self._expression(node.test),
- self._lambda([], self._expression(node.body), node.body),
- self._lambda([], self._expression(node.orelse), node.orelse),
- ],
- node,
- )
- elif isinstance(node, ast.Compare) and len(node.ops) > 1:
- node = self._compare_chain(node)
- 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, declarations)
- if isinstance(node, ast.Call):
- node.func = self._call(
- self.infrastructure_name,
- "callee",
- [self._name(self.builder_name, original), node.func],
- original,
- )
- 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):
- operands = [
- self._lambda([], self._expression(value), value)
- for value in [node.left, *node.comparators]
- ]
- comparisons = []
- for operation in node.ops:
- left, right = self.fresh("left"), self.fresh("right")
- comparison = self._located(
- ast.Compare(self._name(left, node), [operation], [self._name(right, node)]), node
- )
- comparisons.append(self._lambda([left, right], comparison, node))
- return self._call(
- self.infrastructure_name,
- "compare_chain",
- [
- self._attribute(self.builder_name, "logical_and", node),
- ast.Tuple(operands, ast.Load()),
- ast.Tuple(comparisons, ast.Load()),
- ],
- 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, declarations=False):
- 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):
- policy = (
- declarations[index]
- if field == "elts" and isinstance(declarations, list)
- else False
- )
- value[index] = self._expression(
- item, attach_span=policy is not True, declarations=policy
- )
- 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, declaration=False, frame_value=False):
- 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 declaration is True:
- keywords["declaration"] = ast.Constant(True)
- if target.id in self.signature_names:
- keywords["previous"] = copy.deepcopy(
- self.signature_values.get(target.id, self._name(target.id, target))
- )
- elif 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
- )
- ]
- if isinstance(target, ast.Attribute):
- return [
- self._statement(
- self._operation(
- "setattr",
- [self._expression(target.value), ast.Constant(target.attr), value],
- statement,
- ),
- statement,
- )
- ]
- 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.builder_name, "unpack", [value], target)
- assignment = self._located(
- ast.Assign([ast.Tuple(pattern, ast.Store())], unpack), target
- )
- result = [assignment]
- star = next(
- (index for index, item in enumerate(target.elts) if isinstance(item, ast.Starred)),
- None,
- )
- for index, (item, name) in enumerate(zip(target.elts, names)):
- policy = False
- if isinstance(declaration, list):
- source_index = index
- if star is not None and index > star:
- source_index += len(declaration) - len(target.elts)
- if index != star and 0 <= source_index < len(declaration):
- policy = declaration[source_index]
- item = item.value if isinstance(item, ast.Starred) else item
- result.extend(
- self._bind(
- item,
- self._name(name, item),
- statement,
- declaration=policy,
- frame_value=frame_value,
- )
- )
- return result
- self._error(target, f"Unsupported assignment target: {type(target).__name__}")
-
- def visit_Assign(self, node):
- # Cache first: stores evaluate RHS before target base/index, and chained
- # assignments share exactly one RHS evaluation.
- declaration = self._declaration_pattern(node.value)
- cache, value = self._cache(
- self._expression(
- node.value, attach_span=declaration is not True, declarations=declaration
- ),
- node.value,
- )
- result = [cache]
- resolved = self._resolve(node.value)
- for target in node.targets:
- result.extend(self._bind(target, copy.deepcopy(value), node, declaration=declaration))
- if isinstance(target, ast.Name):
- if resolved is not None:
- self.environment[target.id] = resolved
- else:
- self.environment.pop(target.id, None)
- return result
-
- def visit_AnnAssign(self, node):
- 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, attach_span=not self._is_declaration(node.value)),
- node.value,
- )
- result.append(cache)
- result.extend(
- self._bind(
- node.target,
- value,
- node,
- self._expression(node.annotation),
- self._is_declaration(node.value),
- )
- )
- return result
-
- def visit_AugAssign(self, node):
- 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):
- return self._statement(self._operation("emit_", [self._expression(node.value)], node), node)
-
- def visit_Return(self, node):
- value = self._expression(node.value) if node.value is not None else ast.Constant(None)
- if self.preserve_return:
- return self._located(ast.Return(value), node)
- return self._statement(self._operation("return_", [value], node), node)
-
- def visit_Break(self, node):
- return self._statement(self._operation("break_", [], node), node)
-
- def visit_Continue(self, node):
- return self._statement(self._operation("continue_", [], node), node)
-
- def visit_Assert(self, node):
- 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):
- 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.environment, 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.environment = 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.environment, 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, mapping=None):
- result = []
- if mapping is None:
- mapping_stmt, mapping = self._cache(
- self._call(
- self.infrastructure_name,
- "frame_result",
- [self._name(frame, original)],
- original,
- ),
- original,
- "exports",
- )
- result.append(mapping_stmt)
- for name in sorted(candidates):
- key = ast.Constant(name)
- condition = self._located(
- ast.Compare(copy.deepcopy(key), [ast.In()], [copy.deepcopy(mapping)]), original
- )
- value = self._located(ast.Subscript(copy.deepcopy(mapping), key, ast.Load()), original)
- result.append(
- self._located(
- ast.If(condition, [self._assign(name, value, original)], []), original
- )
- )
- for name in candidates - self.bound:
- self.optional[name] = self._located(
- ast.Call(
- ast.Attribute(copy.deepcopy(mapping), "get", ast.Load()),
- [
- ast.Constant(name),
- self._attribute(self.infrastructure_name, "MISSING", original),
- ],
- [],
- ),
- original,
- )
- 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):
- 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._call(
- self.infrastructure_name, "frame_result", [self._name(frame, node)], node
- ),
- node,
- )
- )
- tail = self._exports(frame, candidates, node, self._name(mapping_name, 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):
- if node.orelse:
- self._error(node, "A construction loop does not support an else clause")
- frame, values = self.fresh("loop"), self.fresh("indices")
- if self.iterable_rewriter is not None:
- node.iter = self.iterable_rewriter(node.iter)
- context = self._assign(
- frame, self._operation("For", [self._expression(node.iter)], node), node
- )
- body = self._scope(
- node.body,
- node,
- "body",
- lambda: self._bind_entered(node.target, self._name(values, node.target), node),
- )
- region = self._with(self._name(frame, node), body, node, self._name(values, node, True))
- return [context, region, *self._exports(frame, self._assigned_names(node.body), node)]
-
- def visit_While(self, node):
- 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):
- 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):
- 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_kind__", 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):
- 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)
diff --git a/python/tvm/script/parser/transpile.py b/python/tvm/script/parser/transpile.py
new file mode 100644
index 0000000000..ea15baea82
--- /dev/null
+++ b/python/tvm/script/parser/transpile.py
@@ -0,0 +1,1334 @@
+# 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 tvm.script.ir_builder import protocol
+
+from .expression import rewrite_expression
+
+
+class NameCollector(ast.NodeVisitor):
+ """Reserve source identifiers in a shared name-to-counter dictionary.
+
+ names is owned by the enclosing transpilation unit and populated in place;
+ function/class/import/parameter names and parseable expression strings are
+ reserved before any generated binding is allocated. No source name changes.
+ """
+
+ 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)
+ self.generic_visit(node)
+
+ visit_AsyncFunctionDef = visit_FunctionDef
+ visit_ClassDef = visit_FunctionDef
+
+ 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[str, object]
+ Recognized host namespace/callable bindings for static metadata lookup.
+ No expression evaluation or concrete symbol construction occurs here.
+ builder_name, infrastructure_name : str
+ Collision-free injected aliases for current dialect and shared protocol.
+ span : callable
+ Maps an original AST node to an AST expression containing location data.
+ fresh : callable
+ Shared whole-unit allocator; accepts an optional hint and returns a unique
+ generated identifier. Source identifiers are never renamed.
+ signature_names : iterable[str] or None
+ Names initially bound as function parameters; default empty.
+ nested_function : callable or None
+ Optional syntax-only function handler, default None. Whole-program
+ lowering temporarily installs its nested declaration/definition emitter.
+ preserve_return : bool
+ Retain host Python return semantics in construction helpers, default False.
+ name_map : dict[str, int] or None
+ Shared name reservations/counters; initialized from environment if absent.
+ NameCollector reserves all source names before whole-program lowering.
+
+ State
+ -----
+ filename, span/fresh callbacks and infrastructure_name belong to one complete
+ source unit. namespace_bindings maps visible source spellings to registered Python
+ namespaces/callables; assignments and parameters shadow entries. dialect_prefix
+ selects the generated current builder alias and is restored after functions.
+ bound holds only Python identifier strings. optional maps
+ conditionally exported names to AST reads of builder result dictionaries;
+ for example a name defined only by one branch is checked when read. Each
+ lexical helper copies/restores bound/namespace_bindings/optional; no concrete values
+ or IR identity are stored. preserve_return and syntax hooks configure one
+ body pass. The whole-unit allocator is shared, never reset by nested scopes.
+
+ Pattern summary
+ ---------------
+ ``v = e`` -> ``v = X.bind_(e, name="v", span=location)``.
+ ``v: ty`` uses the builder MISSING initializer; ``v: ty = e`` also passes ty.
+ ``return e`` -> ``X.return_(e)``; bare return supplies no argument.
+ ``for i in range(n)`` -> ``with X.for_(X.range_(n)) as i``.
+ ``for i, j in T.grid(m,n)`` retains grid and the tuple target with X.for_.
+ Function factories emit declaration contexts first, followed by body
+ definition callbacks; builders retain parameter and symbol identities.
+ Marked expression strings become resolve_type_var calls; unmarked dtype
+ literals stay strings. 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,
+ ):
+ self.filename = filename
+ 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"):
+ """Return one unused generated name from the shared whole-unit allocator.
+
+ prefix is an optional identifier prefix, default '_t'. name_map reserves
+ original/generated names with value 0 and stores the next integer counter
+ for each prefix. All nested functions share this dictionary; no source
+ name is renamed and no scope resets the allocator. No builder effects.
+ """
+ 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):
+ """Return translated statements from an original ast.stmt sequence.
+
+ body is copied before visiting; source nodes are never mutated. Each
+ translated statement is wrapped in a location context inheriting its
+ original full range. Syntactic bound/optional-name state is updated for
+ subsequent statements in this body. Unsupported syntax raises SyntaxError;
+ this method never executes callbacks or enters construction frames.
+ """
+ result = []
+ for statement in copy.deepcopy(body):
+ translated = self.visit(statement)
+ if translated is not None:
+ block = translated if isinstance(translated, list) else [translated]
+ context = self._call(
+ self.infrastructure_name, "span_context", [self.span(statement)], statement
+ )
+ result.append(self._with(context, block, statement))
+ 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):
+ 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.logical_not(e).
+ if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
+ node = self._call(
+ self.dialect_prefix, "logical_not", [self._expression(node.operand)], node
+ )
+ # Pattern: a and/or b -> lazy host/builder logical_chain callbacks.
+ elif isinstance(node, ast.BoolOp):
+ method = "logical_and" if isinstance(node.op, ast.And) else "logical_or"
+ operands = [self._lambda([], self._expression(value), value) for value in node.values]
+ node = self._call(
+ self.infrastructure_name,
+ "logical_chain",
+ [
+ self._attribute(self.dialect_prefix, method, node),
+ ast.Tuple(operands, ast.Load()),
+ ast.Constant(isinstance(node.op, ast.Or)),
+ ],
+ node,
+ )
+ # Pattern: a if c else b -> lazy select callbacks, preserving host effects.
+ elif isinstance(node, ast.IfExp):
+ node = self._call(
+ self.infrastructure_name,
+ "select_lazy",
+ [
+ self._attribute(self.dialect_prefix, "select", node),
+ self._expression(node.test),
+ self._lambda([], self._expression(node.body), node.body),
+ self._lambda([], self._expression(node.orelse), node.orelse),
+ ],
+ node,
+ )
+ # Pattern: a < b < c -> compare_chain, evaluating each operand once.
+ 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)
+ if isinstance(node, ast.Call):
+ node.func = self._call(
+ self.infrastructure_name,
+ "callee",
+ [self._name(self.dialect_prefix, original), node.func],
+ original,
+ )
+ 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):
+ operands = [
+ self._lambda([], self._expression(value), value)
+ for value in [node.left, *node.comparators]
+ ]
+ comparisons = []
+ for operation in node.ops:
+ left, right = self.fresh("left"), self.fresh("right")
+ comparison = self._located(
+ ast.Compare(self._name(left, node), [operation], [self._name(right, node)]), node
+ )
+ comparisons.append(self._lambda([left, right], comparison, node))
+ return self._call(
+ self.infrastructure_name,
+ "compare_chain",
+ [
+ self._attribute(self.dialect_prefix, "logical_and", node),
+ ast.Tuple(operands, ast.Load()),
+ ast.Tuple(comparisons, ast.Load()),
+ ],
+ 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, mapping=None):
+ result = []
+ if mapping is None:
+ mapping_stmt, mapping = self._cache(
+ self._call(
+ self.infrastructure_name,
+ "frame_result",
+ [self._name(frame, original)],
+ original,
+ ),
+ original,
+ "exports",
+ )
+ result.append(mapping_stmt)
+ for name in sorted(candidates):
+ key = ast.Constant(name)
+ condition = self._located(
+ ast.Compare(copy.deepcopy(key), [ast.In()], [copy.deepcopy(mapping)]), original
+ )
+ value = self._located(ast.Subscript(copy.deepcopy(mapping), key, ast.Load()), original)
+ result.append(
+ self._located(
+ ast.If(condition, [self._assign(name, value, original)], []), original
+ )
+ )
+ for name in candidates - self.bound:
+ self.optional[name] = self._located(
+ ast.Call(
+ ast.Attribute(copy.deepcopy(mapping), "get", ast.Load()),
+ [
+ ast.Constant(name),
+ self._attribute(self.infrastructure_name, "MISSING", original),
+ ],
+ [],
+ ),
+ original,
+ )
+ 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._call(
+ self.infrastructure_name, "frame_result", [self._name(frame, node)], node
+ ),
+ node,
+ )
+ )
+ tail = self._exports(frame, candidates, node, self._name(mapping_name, 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_kind__", 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 registered decorator metadata without evaluating annotations.
+
+ node is an original FunctionDef; allow_python permits unregistered nested
+ Python helpers. Returns (FunctionKind, options AST). Raises SyntaxError
+ for an undecorated IR function or positional construction options.
+ """
+ for decorator in node.decorator_list:
+ target = decorator.func if isinstance(decorator, ast.Call) else decorator
+ kind = protocol.function_kind(self._resolve(target))
+ if kind is None:
+ continue
+ mapping = kind.metadata.get("option_map", {})
+ values = {
+ key: ast.Constant(value) for key, value in kind.metadata.get("defaults", {}).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.FunctionKind(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 one declaration factory and deferred definition callback.
+
+ node is an original FunctionDef with a registered decorator. runtime is
+ the generated builder-runtime alias. bindings is the caller-owned dict
+ receiving injected namespace aliases. local defaults False; True asks
+ the builder for a nested-function reference. Returns (statement list,
+ record-name, body-name), or None for registered ordinary Python functions.
+ This method allocates names and restores body-local syntax state after
+ recursive translation; it neither executes annotations nor creates IR.
+ Unsupported signatures and decorators raise SyntaxError. Generated code
+ enters the builder's declaration frame, then defines 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.metadata.get("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,
+ )
+ )
+ for parameter in node.args.args:
+ annotation = parameter.annotation
+ if annotation is None:
+ continue
+ constructor = self._resolve(
+ annotation.func if isinstance(annotation, ast.Call) else annotation
+ )
+ if (
+ getattr(constructor, "__tvm_declaration_args__", None) is not None
+ or getattr(constructor, "__tvm_parameter_dtype__", None) is not None
+ ):
+ declaration.append(
+ self._statement(
+ self._call(
+ record,
+ "predeclare",
+ [
+ ast.Constant(parameter.arg),
+ copy.deepcopy(annotation),
+ self.span(parameter),
+ ],
+ parameter,
+ ),
+ parameter,
+ )
+ )
+ 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
+ )
+ 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.metadata.get("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/function builder program from original AST.
+
+ tree is a copied ast.Module; runtime/original_name identify injected
+ runtime and original-object bindings. bindings receives generated builder
+ namespace aliases only. Returns (module AST, generated result binding).
+ Every generated wrapper inherits the source construct's range; original
+ descendants retain their own finer ranges. fix_missing_locations fills
+ missing fields, never recreates coordinates by rendering/reparsing 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.metadata.get("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/__init__.py b/python/tvm/tirx/script/builder/__init__.py
index 4069fed4fa..57b8c8b98a 100644
--- a/python/tvm/tirx/script/builder/__init__.py
+++ b/python/tvm/tirx/script/builder/__init__.py
@@ -18,6 +18,8 @@
import builtins as _python
import sys as _sys
+from dataclasses import dataclass as _dataclass
+from dataclasses import field as _field
from functools import partial as _partial
from functools import wraps as _wraps
@@ -33,7 +35,11 @@ from tvm.script.ir_builder.protocol import at as _at
from tvm.script.ir_builder.protocol import expression_args as _expression_args
from tvm.script.ir_builder.protocol import register_call_kind as _register_call_kind
from tvm.script.ir_builder.protocol import register_declaration as _register_declaration
+from tvm.script.ir_builder.protocol import source_span as _source_span
from tvm.script.ir_builder.protocol import span_context as _span_context
+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.tirx.lang.alloc_pool import SMEMPool, TMEMPool
from . import _ffi_api
@@ -49,8 +55,16 @@ is_type_var = _ir.is_prim_var
def type_var(name, *, dtype=None, span=None):
- """Construct a signature symbol; shape symbols default to int64."""
- return _ir.Var(name, "int64" if dtype is None else dtype, span)
+ """Construct an explicit standalone primitive symbol.
+
+ name is the source-name string; dtype is a dtype string or None (the legacy
+ explicit-construction default int64). span accepts None, an IR span or an
+ original-source location tuple. Returns a new primitive Var, entering no
+ frame and retaining no symbol state. Invalid types propagate constructor
+ errors. Generated expression strings use resolve_type_var instead, which
+ shares the active function frame and defaults newly introduced names to i32.
+ """
+ return _ir.Var(name, "int64" if dtype is None else dtype, _source_span(span))
@_expression_args(
@@ -61,6 +75,7 @@ def type_var(name, *, dtype=None, span=None):
introduce=True,
implicit_dtype="int32",
compound_declarations=True,
+ as_type=True,
)
def Buffer(
shape,
@@ -78,7 +93,23 @@ def Buffer(
*,
span=None,
):
- """Construct a concrete buffer from resolved shape expressions."""
+ """Construct a concrete buffer annotation from resolved expressions.
+
+ ``shape`` is a sequence of concrete dimensions; ``dtype`` defaults to float32.
+ ``data`` optionally supplies the underlying pointer. ``strides`` optionally
+ supplies element strides. ``elem_offset``/``byte_offset`` select the offset;
+ omitted values use native defaults. ``scope`` defaults to global, ``align``
+ and ``offset_factor`` to zero. ``layout`` defaults to the dialect's default
+ layout; ``allocated_addr`` optionally supplies allocation metadata.
+ ``buffer_name`` defaults to an empty name. ``span`` is None, ir.Span, or a
+ source-location tuple. Returns a concrete buffer Var without entering a
+ construction frame, temporarily applying the source span. Native shape,
+ layout and dtype validation errors propagate.
+
+ The protocol wrapper marks only shape/stride/offset fields: unresolved
+ eager annotation expressions return MissingType; dtype/scope/name strings
+ remain literals. Handwritten calls in builders require concrete expressions.
+ """
with _span_context(span):
return _at(
span,
@@ -106,7 +137,7 @@ def Ptr(dtype, storage_scope="global", *, span=None):
"""Construct a concrete pointer variable usable as a function annotation."""
if callable(dtype) and not isinstance(dtype, _ir.Expr):
dtype = dtype()
- if isinstance(dtype, _ir.Expr):
+ if isinstance(dtype, _ir.Expr | _TypeVarDecl):
dtype = dtype.ty
if isinstance(dtype, _ir.PrimType):
dtype = dtype.dtype
@@ -118,6 +149,10 @@ 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 = {}
@@ -156,6 +191,8 @@ def arg(name, annotation, *, span=None):
"""Use the same concrete parameter object in declaration and definition."""
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 _span_context(span):
@@ -184,7 +221,7 @@ def func_ret_type(annotation, *, span=None):
"""Set the signature's concrete return type."""
if callable(annotation) and not isinstance(annotation, _ir.Expr):
annotation = annotation()
- if isinstance(annotation, _ir.Expr):
+ if isinstance(annotation, _ir.Expr | _TypeVarDecl):
annotation = annotation.ty
with _span_context(span):
return _native.func_ret(annotation)
@@ -225,8 +262,37 @@ def bind_(
declaration=False,
frame_value=False,
):
- """Bind values, or name a frame-owned value without introducing new storage."""
+ """Construct a named binding using the active primitive function's policy.
+
+ ``value`` is a concrete expression/value, native frame, TypeVarDecl, or the
+ MISSING default for annotation-only bindings. ``ty`` is an optional concrete
+ type or zero-argument type constructor. ``name`` optionally names the result.
+ ``span`` and ``name_span`` accept None, ir.Span, or source-location tuples;
+ name_span defaults to span. ``previous`` optionally supplies the existing
+ binding for reassignment, otherwise MISSING. ``declaration`` (default False)
+ checks/reuses an existing primitive symbol. A TypeVarDecl always resolves a
+ compatible symbol through the nearest TypeVarFrame. ``frame_value`` (default
+ False) names an existing context-manager result without allocating storage.
+
+ Returns the constructed binding, symbol, or unchanged Python value. Scalar
+ annotations/assignments may allocate mutable scalar storage and emit stores;
+ native frames may be entered with deferred finalization callbacks. TypeVarDecl
+ creates no storage and retains canonical identity across declaration/body.
+ Missing initializers, illegal buffer/pointer/axis rebinding, incompatible
+ declaration types and unsupported annotations raise ValueError/TypeError;
+ other native construction errors propagate. Requires the relevant native
+ function frame and, for TypeVarDecl, a function symbol frame.
+ """
name_span = span if name_span is None else name_span
+ if isinstance(value, _TypeVarDecl):
+ 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")
+ return _TypeVarFrame.current().resolve(name, value.ty, span=name_span)
with _span_context(span):
if frame_value:
if isinstance(value, _frame.SBlockFrame):
@@ -269,7 +335,9 @@ def bind_(
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) else annotation
+ 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:
@@ -299,7 +367,9 @@ def bind_(
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) else annotation
+ 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
@@ -388,7 +458,7 @@ def setattr(target, name, value, *, span=None):
_python.setattr(target, name, value)
-def return_(value, *, span=None):
+def return_(value=None, *, span=None):
"""Emit a native return; subsequent unreachable statements remain in the IR."""
if value is None:
raise TypeError("A primitive function return requires an expression")
@@ -450,13 +520,136 @@ def Else(*, span=None):
return _Frame(_native.Else(), span)
-def For(iterable, *, span=None):
- """Adapt a native loop frame, or a Python range, to construction scope."""
- 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 a native loop frame or range")
- return _Frame(iterable, 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 a Cartesian iteration without entering or constructing loop frames.
+
+ ``extents`` are concrete integer/primitive expressions or (start, extent)
+ pairs; ``dtype`` optionally selects int32/uint32 loop variables, otherwise
+ each extent determines its variable's type. Returns an iteration descriptor
+ with immutable bounds for ``for_``; native validation occurs on consumption.
+ For handwritten compatibility, ``with grid(...) as variables`` delegates to
+ for_ and yields the native variable(s), finalizing the same frames on exit.
+ Sequential reuse creates fresh frames; nested direct reuse of one descriptor
+ raises ValueError. Expression strings are not supported by handwritten builders.
+ """
+ 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 loop scope from an iteration specification.
+
+ ``iterable`` accepts grid()/range_() descriptors, a native ForFrame from
+ serial/parallel/vectorized/thread_binding, or Python range. ``names`` is
+ None (leave native names), one source target string, or a tuple of target
+ strings optionally containing one '*name' unpack marker. A single string
+ names grouped loop variables name_0, name_1, etc.; starred groups use the
+ same suffix convention. Mismatched unpack arity raises ValueError on entry. ``span`` is None,
+ ir.Span, or a source-location tuple. Returns a context manager which pushes
+ native loop frames and yields one variable for serial range, or the grid's
+ variable sequence. Exit finalizes those frames. Invalid iteration inputs
+ raise TypeError; native bound/dtype errors propagate during construction.
+ """
+ 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 _span_context(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):
@@ -509,18 +702,26 @@ del _constructor
def range_(*args):
- """Construct a serial loop from the source builtin range arguments."""
- if len(args) in (1, 2):
- return _native.serial(*args)
- if len(args) != 3:
+ """Describe a serial iteration using Python's one/two/three argument forms.
+
+ ``args`` are concrete integer or primitive expressions: stop; start, stop;
+ or start, stop, step. Defaults are start=0 and step=1. Returns an immutable
+ descriptor consumed by for_ and creates no native frame or variable.
+ Direct ``with range_(...) as variable`` delegates entry/finalization to for_;
+ nested direct reuse of one descriptor raises ValueError. An omitted step
+ remains absent in native IR (whose semantic default is one).
+ Raises TypeError unless one to three arguments are supplied, ValueError for
+ a literal zero step. Other bound/step errors are checked by the native loop.
+ """
+ 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")
- start, stop, step = args
- if isinstance(step, _python.int) and step == 0:
+ if isinstance(args[2], _python.int) and args[2] == 0:
raise ValueError("range step cannot be zero")
- return _native.serial(start, stop, step=step)
-
-
-__tvm_call_overrides__ = {_python.range: range_}
+ return _IterationSpec("range", args)
def logical_and(*values):
@@ -578,8 +779,3 @@ def _global_callee(function):
_register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee)
-
-
-from tvm.script.parser.annotations import enable_eager_constructors as _enable_eager_constructors
-
-_enable_eager_constructors(_sys.modules[__name__], classes=("Buffer",))
diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py
index 3f4e2b600f..71590fcaa9 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -474,12 +474,20 @@ def func_ret(ret_type: Type | None) -> Type:
def Tuple(*fields: Type) -> Type: # pylint: disable=invalid-name
- """Construct a tuple type for a TIRx function or binding annotation."""
+ """Construct a tuple type from concrete types or scalar type declarations.
+
+ Fields accept types (including MissingType and nested TupleType), scalar
+ expressions, TypeVarDecl descriptors, and zero-argument constructors of
+ those values. Descriptors contribute their type without resolving a symbol
+ or entering a frame. Native TupleType validates the resulting field types.
+ """
+ from tvm.script.ir_builder.type_var_frame import TypeVarDecl
+
normalized_fields = []
for field in fields:
if callable(field) and not isinstance(field, Expr):
field = field()
- if isinstance(field, Expr):
+ if isinstance(field, Expr | TypeVarDecl):
field = field.ty
normalized_fields.append(field)
return ir.TupleType(normalized_fields)
@@ -1540,7 +1548,9 @@ def Bind( # pylint: disable=invalid-name
# already a resolved type annotation, rather than a deferred annotation factory.
if callable(type_annotation) and not isinstance(type_annotation, Expr):
type_annotation = type_annotation()
- if isinstance(type_annotation, ir.Var):
+ from tvm.script.ir_builder.type_var_frame import TypeVarDecl
+
+ if isinstance(type_annotation, ir.Var | TypeVarDecl):
type_annotation = type_annotation.ty
return _ffi_api.Bind(value, type_annotation, var) # type: ignore[attr-defined] # pylint: disable=no-member
@@ -1575,21 +1585,30 @@ class LetAnnotation:
return LetAnnotation(item)
def as_var(self, rhs_dtype=None):
- """Resolve to a tir.Var."""
- if self.type_spec is not None:
- if isinstance(self.type_spec, ir.Var):
- return self.type_spec # Already a Var (e.g. T.handle(...))
- elif callable(self.type_spec):
- return self.type_spec() # e.g. T.int32() -> Var
- elif isinstance(self.type_spec, Type):
- return Var("", self.type_spec)
- else:
- raise TypeError(f"Invalid type for T.let: {self.type_spec}")
- elif rhs_dtype is not None:
- rhs_ty = rhs_dtype if isinstance(rhs_dtype, Type) else ir.PrimType(rhs_dtype)
- return Var("", rhs_ty)
- else:
- raise TypeError("T.let requires either a type or an RHS value")
+ """Return an unnamed concrete variable for this immutable binding.
+
+ ``rhs_dtype`` is an optional ir.Type or dtype string used when no explicit
+ type_spec was supplied. Explicit type_spec accepts a concrete Var, Type,
+ TypeVarDecl, or zero-argument constructor of these. Returns a concrete
+ Var (preserving an explicitly supplied Var), without entering frames or
+ registering a TypeVar: let-bindings are ordinary lexical values. Raises
+ TypeError if neither source supplies a valid type; constructor errors
+ propagate.
+ """
+ from tvm.script.ir_builder.type_var_frame import TypeVarDecl
+
+ annotation = self.type_spec
+ if callable(annotation) and not isinstance(annotation, Expr):
+ annotation = annotation()
+ if isinstance(annotation, ir.Var):
+ return annotation
+ if isinstance(annotation, TypeVarDecl):
+ annotation = annotation.ty
+ if annotation is None and rhs_dtype is not None:
+ annotation = rhs_dtype if isinstance(rhs_dtype, Type) else ir.PrimType(rhs_dtype)
+ if isinstance(annotation, Type):
+ return Var("", annotation)
+ raise TypeError("T.let requires a concrete type or an RHS value")
let = LetAnnotation() # Singleton for T.let (no subscript)
@@ -1615,13 +1634,16 @@ class DtypeConstructor:
Replaces the plain functions previously returned by ``func_gen``.
- * ``T.float32()`` — same FFI call as before (returns ``Var``).
+ * ``T.float32()`` — TypeVarDecl in a function symbol frame, otherwise Var.
* ``T.float32[N]`` — returns ``LocalVectorAnnotation("float32", (N,))``.
* ``T.float32[M, N]`` — returns ``LocalVectorAnnotation("float32", (M, N))``.
- * ``x: T.float32`` — parser calls this object, gets a ``Var``.
+ * ``x: T.float32`` — builder consumes its primitive type.
"""
def __init__(self, ffi_name: str, dtype_str: str):
+ # Immutable constructor metadata: native entry-point name and dtype
+ # spelling. These live with the module-level callable, across functions;
+ # neither stores per-call values or symbol state.
self._ffi_name = ffi_name
self._dtype_str = dtype_str
@@ -1629,6 +1651,24 @@ class DtypeConstructor:
self,
expr: "Expr | Literal['inf', '-inf', 'nan'] | int | float | None" = None,
) -> "Expr":
+ """Construct/cast a scalar, or describe a named symbolic declaration.
+
+ ``expr`` is a concrete expression/number, a floating literal string
+ (inf, -inf, nan), or None (default). Concrete values use native casting.
+ None returns TypeVarDecl when a TypeVarFrame is active, for bind_ to
+ create/reuse a compatible named symbol; without such a frame it returns
+ an anonymous native Var for ordinary handwritten-builder compatibility.
+ Calls enter no frames. Invalid casts/dtypes propagate native errors;
+ unrecognized floating strings raise ValueError.
+ """
+ if expr is None:
+ from tvm.script.ir_builder.base import IRBuilder
+ from tvm.script.ir_builder.type_var_frame import TypeVarDecl, TypeVarFrame
+
+ if IRBuilder.is_in_scope() and any(
+ isinstance(frame, TypeVarFrame) for frame in IRBuilder.current().frames
+ ):
+ return TypeVarDecl(self._dtype_str)
if isinstance(expr, str):
expr = float(expr)
return getattr(_ffi_api, self._ffi_name)(expr)
diff --git a/src/relax/analysis/type_analysis.cc b/src/relax/analysis/type_analysis.cc
index f2a9506bed..e2ccd07e7e 100644
--- a/src/relax/analysis/type_analysis.cc
+++ b/src/relax/analysis/type_analysis.cc
@@ -205,11 +205,15 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase {
PrimExpr value = ret.value().as_or_throw<PrimExpr>();
if (value->IsInstance<IntImmNode>()) {
- return ffi::Any(tvm::prim::cast(PrimType::Int(64), value));
+ // Preserve the replaced symbol's arithmetic width inside compound
+ // dimensions; shape literals are normalized separately by ShapeType.
+ return ffi::Any(tvm::prim::cast(var->ty.as_or_throw<PrimType>(), value));
}
- TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64))
- << "Can only provide i64 expressions in shape";
- return ffi::Any(value);
+ TVM_FFI_ICHECK(value.ty().IsScalar() &&
+ (value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 32) ||
+ value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64)))
+ << "Can only provide scalar i32 or i64 expressions in shape";
+ return ffi::Any(tvm::prim::cast(var->ty.as_or_throw<PrimType>(), value));
};
// A well-definedness map may intentionally be self-referential (for example m -> m + 1).
// Post-order preserves the old substitution behavior by not re-entering that replacement.
diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index 0469dff885..ca75a9053e 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -139,7 +139,8 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor {
Var var = ffi::GetRef<Var>(op);
if (collect_scalar_ && !var.as<DataflowVarNode>()) {
if (auto prim_var = var.as<PrimVar>();
- prim_var && prim_var.value().ty()->dtype == DLDataType{kDLInt, 64, 1}) {
+ prim_var && (prim_var.value().ty()->dtype == DLDataType{kDLInt, 32, 1} ||
+ prim_var.value().ty()->dtype == DLDataType{kDLInt, 64, 1})) {
HandlePrimExpr(prim_var.value());
}
}
@@ -716,8 +717,12 @@ class VMShapeLowerMutator
ffi::Map<tirx::Var, PrimExpr> var_map;
for (const auto& [expr, slot] : slot_map_) {
if (auto var = expr.as<tirx::Var>()) {
- var_map.Set(var.value(),
- tirx::BufferLoad(buffer, {IntImm(tvm::PrimType(ShapeDType()), slot->index)}));
+ // Heap slots always store i64, while source dimensions retain their
+ // declared width. Narrow loads before substitution so i32 arithmetic
+ // keeps its semantics instead of silently becoming i64 arithmetic.
+ PrimExpr load =
+ tirx::BufferLoad(buffer, {IntImm(tvm::PrimType(ShapeDType()), slot->index)});
+ var_map.Set(var.value(), tvm::prim::cast(var.value()->ty.as_or_throw<PrimType>(), load));
}
}
auto f_substitute =
@@ -732,8 +737,10 @@ class VMShapeLowerMutator
slot->value_computed = true;
PrimExpr value = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(slot->expr, f_substitute)
.as_or_throw<PrimExpr>();
- seq.push_back(
- tirx::BufferStore(buffer, value, {IntImm(tvm::PrimType(ShapeDType()), slot->index)}));
+ // Widen only at the storage boundary; the shape heap/runtime ABI stays
+ // i64 even when the completed dimension expression is computed in i32.
+ seq.push_back(tirx::BufferStore(buffer, tvm::prim::cast(tvm::PrimType(ShapeDType()), value),
+ {IntImm(tvm::PrimType(ShapeDType()), slot->index)}));
}
tirx::Stmt body = tirx::SeqStmt::Flatten(seq);
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index 3b7ed49017..3e035e3483 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -219,8 +219,13 @@ ShapeType::ShapeType(ffi::Array<PrimExpr> values, Span span) : Type(ffi::UnsafeI
if (value->IsInstance<IntImmNode>()) {
return tvm::prim::cast(PrimType::Int(64), value);
}
- TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64))
- << "the value in ShapeType can only have dtype of int64";
+ // Keep symbolic dimensions as direct Vars/expressions: wrapping an i32 Var
+ // in an i64 cast here would hide its definition from shape-scope analysis.
+ // Only the runtime shape heap has a fixed i64 representation.
+ TVM_FFI_ICHECK(value.ty().IsScalar() &&
+ (value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 32) ||
+ value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64)))
+ << "the value in ShapeType must have scalar dtype int32 or int64";
return value;
});
n->span = span;
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 5c89eec025..3da7934ea9 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -416,8 +416,13 @@ ShapeExpr::ShapeExpr(ffi::Array<PrimExpr> values, Span span) {
if (value->IsInstance<IntImmNode>()) {
return tvm::prim::cast(PrimType::Int(64), value);
}
- TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64))
- << "the value in ShapeType can only have dtype of int64";
+ // Keep symbolic dimensions as direct Vars/expressions: wrapping an i32 Var
+ // in an i64 cast here would hide its definition from shape-scope analysis.
+ // Only the runtime shape heap has a fixed i64 representation.
+ TVM_FFI_ICHECK(value.ty().IsScalar() &&
+ (value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 32) ||
+ value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64)))
+ << "the value in ShapeType must have scalar dtype int32 or int64";
return value;
});
n->span = span;
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/relax/script/printer/tir.cc b/src/relax/script/printer/tir.cc
index 8508aef4de..125f0f54ad 100644
--- a/src/relax/script/printer/tir.cc
+++ b/src/relax/script/printer/tir.cc
@@ -65,7 +65,11 @@ Doc PrintCanonicalVar(Var n, AccessPath n_p, IRDocsifier d) {
f->type_vars->insert(n.get());
}
}
- IdDoc var = d->Define(n, ffi::GetRef<Frame>(f), n->name.empty() ? "v" : n->name);
+ // Signature symbols have declarations hoisted above the module. Keep their
+ // names reserved across sibling functions, including mixed i32/i64 symbols.
+ bool signature_symbol = f->type_vars && f->type_vars->count(n.get());
+ Frame name_scope = signature_symbol ? d->frames.front() : ffi::GetRef<Frame>(f);
+ IdDoc var = d->Define(n, name_scope, n->name.empty() ? "v" : n->name);
var->source_paths.push_back(n_p);
if (!f->func_vars || f->prim_params->count(n.get()) || !f->type_vars->count(n.get())) {
f->stmts.push_back(AssignDoc(var, PrintVarCreation(prim_var, n_p, d), std::nullopt));
diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc
index 1871a905a6..53871df111 100644
--- a/src/relax/transform/static_plan_block_memory.cc
+++ b/src/relax/transform/static_plan_block_memory.cc
@@ -173,7 +173,7 @@ class StorageToken : public ffi::ObjectRef {
if (const IntImmNode* const_dim_len = dim_len.as<IntImmNode>()) {
const_coeff *= const_dim_len->value;
} else {
- size *= dim_len;
+ size *= tvm::prim::cast(PrimType::Int(64), dim_len);
}
}
}
@@ -491,7 +491,18 @@ ffi::Array<PrimExpr> GetUpperBoundShape(ffi::Array<PrimExpr> shape, sym::Analyze
upper_bounded_shape.reserve(shape.size());
for (const PrimExpr& dim_len : shape) {
int64_t max_bound = ana->const_int_bound(dim_len)->max_value;
- if (max_bound == std::numeric_limits<int64_t>::max()) {
+ // A narrow integer's representable range is not an annotated allocation
+ // bound. Keep unbounded i32 dimensions symbolic, just as for i64, instead
+ // of reserving INT32_MAX elements (or overflowing a product of such bounds).
+ bool has_unbounded_i32 = false;
+ for (const VarNode* var : GetUsedVars(dim_len)) {
+ if (var->ty.as_or_throw<PrimType>().MatchesElementType(DLDataTypeCode::kDLInt, 32) &&
+ !dom_map.count(ffi::GetRef<Var>(var))) {
+ has_unbounded_i32 = true;
+ break;
+ }
+ }
+ if (has_unbounded_i32 || max_bound == std::numeric_limits<int64_t>::max()) {
sym::IntSet int_set = ana->int_set(dim_len, dom_map);
if (int_set.HasUpperBound()) {
upper_bounded_shape.push_back(int_set.max());
@@ -992,9 +1003,10 @@ class StorageAllocationRewriter : public ExprMutator {
TVM_FFI_ICHECK(!ty->IsUnknownDtype());
TVM_FFI_ICHECK_EQ(ty->dtype.value()->dtype,
call->args[1].as_or_throw<GenericConst>()->value.cast<DLDataType>());
- PrimExpr bytes = upper_bounded_shape[0];
+ // Allocation sizes use the runtime i64 ABI even when dimensions are i32.
+ PrimExpr bytes = tvm::prim::cast(PrimType::Int(64), upper_bounded_shape[0]);
for (int i = 1; i < static_cast<int>(upper_bounded_shape.size()); ++i) {
- bytes *= upper_bounded_shape[i];
+ bytes *= tvm::prim::cast(PrimType::Int(64), upper_bounded_shape[i]);
}
DLDataType dtype = ty->dtype.value()->dtype;
PrimType dtype_ty(dtype);
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/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index 11d4811cf8..b87b83fb14 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -76,7 +76,8 @@ ffi::Optional<ffi::String> GetTypeVarDeclarationName(const StmtDoc& stmt) {
return std::nullopt;
}
const auto* callee = call->callee.as<IdDocNode>();
- if (callee == nullptr || callee->name != "TypeVar") {
+ const auto* attribute = call->callee.as<AttrAccessDocNode>();
+ if ((callee == nullptr || callee->name != "TypeVar") && attribute == nullptr) {
return std::nullopt;
}
return lhs->name;
diff --git a/src/script/printer/utils.h b/src/script/printer/utils.h
index 1ec51e3a43..37ef740146 100644
--- a/src/script/printer/utils.h
+++ b/src/script/printer/utils.h
@@ -32,6 +32,7 @@
#include <algorithm>
#include <sstream>
#include <string>
+#include <tuple>
#include <unordered_set>
#include <utility>
#include <vector>
@@ -157,22 +158,27 @@ inline Doc HeaderWrapper(const IRDocsifier& d, const Doc& doc) {
return doc;
}
-inline std::vector<std::pair<std::string, ExprDoc>> DefineTypeVarDocs(
+inline std::vector<std::tuple<std::string, ExprDoc, DLDataType>> DefineTypeVarDocs(
const std::unordered_set<const VarNode*>& type_vars, const Frame& frame, const IRDocsifier& d) {
- std::vector<std::pair<std::string, ExprDoc>> type_var_docs;
+ std::vector<std::tuple<std::string, ExprDoc, DLDataType>> type_var_docs;
type_var_docs.reserve(type_vars.size());
for (const VarNode* var_node : type_vars) {
Var var = ffi::GetRef<Var>(var_node);
ffi::Optional<ExprDoc> existing_doc = d->GetVarDoc(var);
ExprDoc var_doc = existing_doc.has_value()
? existing_doc.value()
- : d->Define(var, frame, var->name.empty() ? "v" : var->name);
+ // Symbol declarations may be hoisted above a module.
+ // Reserve their names for the whole print unit, not
+ // just one function: two distinct n symbols (including
+ // i32/i64 pairs) must not collapse into one declaration.
+ : d->Define(var, d->frames.empty() ? frame : d->frames.front(),
+ var->name.empty() ? "v" : var->name);
const auto* id_doc = var_doc.as<IdDocNode>();
TVM_FFI_ICHECK(id_doc != nullptr);
- type_var_docs.emplace_back(id_doc->name, var_doc);
+ type_var_docs.emplace_back(id_doc->name, var_doc, var->ty.as<PrimTypeNode>()->dtype);
}
std::sort(type_var_docs.begin(), type_var_docs.end(),
- [](const auto& lhs, const auto& rhs) { return lhs.first < rhs.first; });
+ [](const auto& lhs, const auto& rhs) { return std::get<0>(lhs) < std::get<0>(rhs); });
return type_var_docs;
}
@@ -183,24 +189,36 @@ inline bool UsePEP695TypeVars(const IRDocsifier& d) {
inline Doc WrapFunctionDocWithTypeVars(
const IRDocsifier& d, FunctionDoc function_doc,
- const std::vector<std::pair<std::string, ExprDoc>>& type_var_docs) {
+ const std::vector<std::tuple<std::string, ExprDoc, DLDataType>>& type_var_docs) {
if (type_var_docs.empty()) {
return HeaderWrapper(d, function_doc);
}
- if (UsePEP695TypeVars(d)) {
+ // PEP 695 scalar parameters use the language default i32. Fall back to
+ // explicit declarations when a symbol has another dtype; a bare type
+ // parameter would otherwise silently lose its type on roundtrip.
+ const bool default_types =
+ std::all_of(type_var_docs.begin(), type_var_docs.end(), [](const auto& entry) {
+ const auto dtype = std::get<2>(entry);
+ return dtype.code == kDLInt && dtype.bits == 32 && dtype.lanes == 1;
+ });
+ if (UsePEP695TypeVars(d) && default_types) {
d->ir_usage.insert("future_annotations");
for (const auto& type_var_doc : type_var_docs) {
- function_doc->type_params.push_back(type_var_doc.second);
+ function_doc->type_params.push_back(std::get<1>(type_var_doc));
}
return HeaderWrapper(d, function_doc);
}
d->ir_usage.insert("type_var");
ffi::Array<StmtDoc> stmts;
- for (const auto& [name, var_doc] : type_var_docs) {
- stmts.push_back(AssignDoc(
- var_doc, IdDoc("TypeVar")->Call({LiteralDoc::Str(name, ffi::Optional<AccessPath>())}),
- std::nullopt));
+ for (const auto& [name, var_doc, dtype] : type_var_docs) {
+ // Expression-string symbols default to i32. Preserve every other explicit
+ // dtype in source, so reparsing cannot silently change symbol semantics.
+ ExprDoc declaration =
+ dtype.code == kDLInt && dtype.bits == 32 && dtype.lanes == 1
+ ? IdDoc("TypeVar")->Call({LiteralDoc::Str(name, ffi::Optional<AccessPath>())})
+ : TIR(d, ffi::String(ffi::DLDataTypeToString(dtype)))->Call({});
+ stmts.push_back(AssignDoc(var_doc, declaration, std::nullopt));
}
stmts.push_back(function_doc);
return HeaderWrapper(d, StmtBlockDoc(stmts));
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;
diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py
index 23752550c9..b2c833b9fa 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -187,7 +187,7 @@ def test_symbolic_var_across_functions():
def test_symbolic_var_invalid_type():
- with pytest.raises(RuntimeError, match="the value in ShapeType can only have dtype of int64"):
+ with pytest.raises(RuntimeError, match="scalar dtype int32 or int64"):
dim = tirx.Var("dim", "float32")
y = rx.Var("y", R.Tensor([dim], "float32"))
gv0 = rx.Var("gv0", R.Tensor([dim], "float32"))
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py
index c9f4fda65e..1ef80f0f1e 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -73,6 +73,9 @@ def test_dispatch_scanop_cuda():
lowered to the packed func `"gpu_2d_continuous_cumsum"`.
"""
+ # Preserve the explicitly i64 shapes used by the reference lowering.
+ m = T.int64()
+
@I.ir_module
class Before:
I.module_global_infos({"vdevice": [I.vdevice("cuda", 0)]})
@@ -119,6 +122,9 @@ def test_dispatch_scanop_cuda():
def test_dispatch_sort():
+ # Preserve the explicitly i64 shapes used by the reference lowering.
+ m = T.int64()
+
@I.ir_module
class Before:
I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
@@ -216,6 +222,9 @@ def test_dispatch_sort_cuda():
def test_dispatch_argsort():
+ # Preserve the explicitly i64 shapes used by the reference lowering.
+ m = T.int64()
+
@I.ir_module
class Before:
I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
@@ -309,6 +318,9 @@ def test_dispatch_argsort_cuda():
def test_dispatch_topk():
+ # Preserve the explicitly i64 shapes used by the reference lowering.
+ m = T.int64()
+
@I.ir_module
class Before:
I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
diff --git a/tests/python/relax/test_backend_transform_shape_lower.py b/tests/python/relax/test_backend_transform_shape_lower.py
index 746edf9ace..b6dc0bb1cc 100644
--- a/tests/python/relax/test_backend_transform_shape_lower.py
+++ b/tests/python/relax/test_backend_transform_shape_lower.py
@@ -173,6 +173,9 @@ def test_symbolic_compute():
MS = MatchShapeCode
MK = MakeShapeCode
+ # These fixtures exercise explicit i64 arithmetic, independent of the script default.
+ n, m = T.int64(), T.int64()
+
@tvm.script.ir_module
class Before:
@R.function
@@ -461,6 +464,9 @@ def test_return_match_check_with_new_expr():
"""
MS = MatchShapeCode
+ # These fixtures exercise explicit i64 arithmetic, independent of the script default.
+ n = T.int64()
+
@tvm.script.ir_module
class Before:
@R.function
@@ -540,6 +546,9 @@ def test_symbolic_shape_multiple_function():
MS = MatchShapeCode
MK = MakeShapeCode
+ # These fixtures exercise explicit i64 arithmetic, independent of the script default.
+ m, n, m2, n2 = (T.int64() for _ in range(4))
+
@I.ir_module
class Before:
@R.function
@@ -550,10 +559,10 @@ def test_symbolic_shape_multiple_function():
return A
@R.function
- def fn2(A: R.Tensor(("n", "m"), dtype="float32")):
+ def fn2(A: R.Tensor(("n2", "m2"), dtype="float32")):
R.func_attr({"relax.force_pure": True})
- n = T.int64()
- m = T.int64()
+ n2 = T.int64()
+ m2 = T.int64()
return A
# slot assignment:
@@ -601,10 +610,12 @@ def test_symbolic_shape_multiple_function():
return A
@R.function
- def fn2(A: R.Tensor(("n", "m"), dtype="float32")) -> R.Tensor(("n", "m"), dtype="float32"):
+ def fn2(A: R.Tensor(("n2", "m2"), dtype="float32")) -> R.Tensor(
+ ("n2", "m2"), dtype="float32"
+ ):
R.func_attr({"relax.force_pure": True})
- n = T.int64()
- m = T.int64()
+ n2 = T.int64()
+ m2 = T.int64()
shape_heap: R.Tensor(dtype="int64", ndim=1) = R.call_builtin_with_ctx(
"vm.builtin.alloc_shape_heap",
(R.prim_value(2),),
@@ -730,6 +741,9 @@ def test_check_lifted_weights():
def test_check_weights_with_dynamic_shape():
MS = MatchShapeCode
+ # These fixtures exercise explicit i64 arithmetic, independent of the script default.
+ n = T.int64()
+
@I.ir_module
class Before:
@R.function
diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py
index 212b9d4388..f163ed1e8c 100644
--- a/tests/python/relax/test_blockbuilder_core.py
+++ b/tests/python/relax/test_blockbuilder_core.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F401, F841
+# ruff: noqa: F841
"""Block builder unit test"""
# The test here do not depend on tvmscript to cover most basic features
@@ -660,6 +660,7 @@ def test_emit_nested_tuple(emit_nested_tuple):
return bb.finalize()["func"]
def make_expected(emit_nested_tuple: bool):
+ n, m = T.int64(), T.int64()
if emit_nested_tuple:
@R.function
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index 93d554424e..b6f6e7c27b 100644
--- a/tests/python/relax/test_expr.py
+++ b/tests/python/relax/test_expr.py
@@ -292,8 +292,9 @@ def test_shape_expr():
tvm.ir.assert_structural_equal(x.ty.shape.ty, R.Shape((10, 20)))
m = tirx.Var("m", "int32")
- with pytest.raises(RuntimeError, match="the value in ShapeType can only have dtype of int64"):
- rx.ShapeExpr([m, 3])
+ assert rx.ShapeExpr([m, 3]).values[0].same_as(m)
+ with pytest.raises(RuntimeError, match="scalar dtype int32 or int64"):
+ rx.ShapeExpr([tirx.Var("invalid", "int16"), 3])
def test_prim_value_helper_from_python_scalar():
diff --git a/tests/python/relax/test_shape_i32.py b/tests/python/relax/test_shape_i32.py
new file mode 100644
index 0000000000..dc4ddf50be
--- /dev/null
+++ b/tests/python/relax/test_shape_i32.py
@@ -0,0 +1,147 @@
+# 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.
+"""Signed shape symbols retain their width until the runtime heap boundary."""
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import ir, relax, tirx
+from tvm.testing import env
+
+
+@pytest.mark.parametrize("dtype", ["int32", "int64"])
+def test_shape_symbol_identity(dtype):
+ symbol = ir.Var("n", dtype)
+ shape = relax.ShapeExpr([symbol])
+ shape_type = relax.ShapeType([symbol])
+ tensor_type = relax.TensorType([symbol], "float32")
+ assert shape[0].same_as(symbol)
+ assert shape_type.values[0].same_as(symbol)
+ assert tensor_type.shape[0].same_as(symbol)
+
+
+@pytest.mark.parametrize("dtype", ["int16", "uint32", "float32", "int32x4"])
+def test_shape_rejects_unsupported_symbol_type(dtype):
+ symbol = ir.Var("n", dtype)
+ for constructor in (relax.ShapeExpr, relax.ShapeType):
+ with pytest.raises(tvm.error.InternalError, match="scalar dtype int32 or int64"):
+ constructor([symbol])
+
+
+def _shape_computation(dtype, scalar_parameter=False):
+ symbol = ir.Var("n", dtype)
+ parameter = symbol if scalar_parameter else ir.Var("x", relax.TensorType([symbol], "float32"))
+ builder = relax.BlockBuilder()
+ with builder.function("main", [parameter], attrs={"relax.force_pure": True}):
+ output = builder.emit(relax.ShapeExpr([symbol + 1]))
+ builder.emit_func_output(output)
+ return builder.get()
+
+
+@pytest.mark.parametrize("dtype", ["int32", "int64"])
+def test_shape_heap_preserves_arithmetic_width(dtype):
+ lowered = relax.transform.VMShapeLower()(_shape_computation(dtype))
+ functions = [value for value in lowered.functions.values() if isinstance(value, tirx.PrimFunc)]
+ assert len(functions) == 1
+ stores = []
+
+ def collect_store(node, visitor):
+ stores.append(node)
+ visitor.default_visit(node)
+
+ tvm_ffi.structural_visit(functions[0].body, [(tirx.BufferStore, collect_store)])
+ assert len(stores) == 1
+ store = stores[0]
+ assert store.buffer.ty.dtype == "int64"
+ assert store.value.ty == "int64"
+ value = store.value.value if dtype == "int32" else store.value
+ assert value.ty == dtype
+ if dtype == "int32":
+ assert isinstance(store.value, tirx.Cast)
+ assert isinstance(value.a, tirx.Cast)
+ assert value.a.value.ty == "int64"
+ else:
+ assert not isinstance(store.value, tirx.Cast)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="LLVM support required")
+@pytest.mark.parametrize("dtype", ["int32", "int64"])
+@pytest.mark.parametrize("scalar_parameter", [False, True], ids=["tensor", "scalar"])
+def test_shape_vm_execution(dtype, scalar_parameter):
+ module = _shape_computation(dtype, scalar_parameter=scalar_parameter)
+ executable = relax.build(module, target="llvm")
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+ argument = 7 if scalar_parameter else tvm.runtime.tensor(np.empty((7,), dtype="float32"))
+ assert tuple(vm["main"](argument)) == (8,)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="LLVM support required")
+@pytest.mark.parametrize("dtype", ["int32", "int64"])
+@pytest.mark.parametrize("bounded", [False, True], ids=["dynamic", "bounded"])
+def test_dynamic_matmul_vm_execution(dtype, bounded):
+ # This crosses legalization/call_tir, allocation planning, and VM shape
+ # lowering. Dtype bounds alone must not reserve INT32_MAX-sized tensors;
+ # explicit function bounds remain available to the memory planner.
+ m, n, k = [ir.Var(name, dtype) for name in ("m", "n", "k")]
+ lhs = ir.Var("lhs", relax.TensorType([m, k], "float32"))
+ rhs = ir.Var("rhs", relax.TensorType([k, n], "float32"))
+ attributes = {"relax.force_pure": True, "relax.memory_plan_dynamic_func_output": True}
+ if bounded:
+ attributes["tir_var_upper_bound"] = {"m": 8, "n": 8, "k": 8}
+ builder = relax.BlockBuilder()
+ with builder.function("main", [lhs, rhs], attrs=attributes):
+ product = builder.emit(relax.op.matmul(lhs, rhs))
+ # Keep an intermediate allocation eligible for memory reuse planning.
+ result = builder.emit(relax.op.add(product, product))
+ builder.emit_func_output(result)
+ module = builder.get()
+ planned = tvm.transform.Sequential(
+ [
+ relax.transform.LegalizeOps(),
+ relax.transform.CallTIRRewrite(),
+ relax.transform.StaticPlanBlockMemory(),
+ ]
+ )(module)
+ allocation_sizes = []
+
+ def collect_allocation(node, visitor):
+ if isinstance(node.op, ir.Op) and node.op.name == "relax.memory.alloc_storage":
+ allocation_sizes.append(node.args[0][0])
+ visitor.default_visit(node)
+
+ tvm_ffi.structural_visit(planned["main"].body, [(ir.Call, collect_allocation)])
+ assert allocation_sizes
+ assert all(size.ty == "int64" for size in allocation_sizes)
+ if bounded:
+ # The explicit bound, unlike the dtype maximum, must plan fixed storage.
+ assert all(int(size) == 8 * 8 * 4 for size in allocation_sizes)
+ else:
+ assert all(not isinstance(size, tirx.IntImm) for size in allocation_sizes)
+
+ executable = relax.build(module, target="llvm")
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+ lhs_array = np.arange(6, dtype="float32").reshape(2, 3)
+ rhs_array = np.arange(12, dtype="float32").reshape(3, 4)
+ actual = vm["main"](tvm.runtime.tensor(lhs_array), tvm.runtime.tensor(rhs_array))
+ tvm.testing.assert_allclose(actual.numpy(), 2 * (lhs_array @ rhs_array))
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_normalize.py b/tests/python/relax/test_transform_normalize.py
index d136c80df3..6100f1201d 100644
--- a/tests/python/relax/test_transform_normalize.py
+++ b/tests/python/relax/test_transform_normalize.py
@@ -105,6 +105,8 @@ def test_normalize_if():
def test_normalize_no_op():
# the normalize pass should be no-op for IR in ANF
+ m, n = T.int64(), T.int64()
+
@tvm.script.ir_module
class ANFMod1:
@R.function
diff --git a/tests/python/relax/test_tvmscript_parser.py b/tests/python/relax/test_tvmscript_parser.py
index 9cda2610d3..d52ac94b90 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -414,15 +414,15 @@ def test_relax_shape_to_tensor():
def test_symbolic_shape():
@R.function
def foo(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
- m = T.int64()
- n = T.int64()
+ m = T.int32()
+ n = T.int32()
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
return gv0
@R.function
def bar(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
- m = T.int64()
- n = T.int64()
+ m = T.int32()
+ n = T.int32()
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
return gv0
@@ -430,13 +430,13 @@ def test_symbolic_shape():
@R.function
def mismatch_dtype(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2):
- m = T.int64()
- n = T.int32() # The shape dtype should be int64
+ m = T.int32()
+ n = T.int64() # Conflicts with the newly introduced i32 shape symbol.
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
return gv0
def _expected(name: str):
- n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
+ n, m = tirx.Var("n", "int32"), tirx.Var("m", "int32")
x = relax.Var("x", R.Tensor([m, n], "float32"))
bb = relax.BlockBuilder()
with bb.function(name, (x,)):
@@ -888,7 +888,7 @@ def test_annotation():
y: R.Tensor(("m",), "float32"),
r: R.Tensor(dtype="int64"),
) -> R.Any:
- m = T.int64()
+ m = T.int32()
z: R.Tensor((32, m), "float32") = R.multiply(x, y)
w: R.Tensor(ndim=2) = R.multiply(z, z)
q: R.Tensor = R.add(w, w)
@@ -980,13 +980,13 @@ def test_call_tir_with_tir_var():
def main(
dumb_param: R.Tensor(("n",), "float32"), x: R.Tensor(("n * 2",), "float32")
) -> R.Tensor(("n * 2",), "float32"):
- n = T.int64()
+ n = T.int32()
cls = Module
y = R.call_tir(cls.copy, (x, n), R.Tensor((n * 2,), dtype="float32"))
return y
@T.prim_func(s_tir=True)
- def copy(var_x: T.handle, n: T.int64, var_y: T.handle):
+ def copy(var_x: T.handle, n: T.int32, var_y: T.handle):
X = T.match_buffer(var_x, (n * 2,), dtype="float32")
Y = T.match_buffer(var_y, (n * 2,), dtype="float32")
for i in T.grid(n * 2):
@@ -1344,7 +1344,7 @@ def test_computed_prim_value_as_branch_condition():
@R.function
def func(x: R.Tensor(["N"], "float32")):
- N = T.int64()
+ N = T.int32()
if R.prim_value(N % 16 == 0):
out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
else:
@@ -1363,7 +1363,7 @@ def test_tir_expr_as_branch_condition():
@R.function(private=True)
def sugared(x: R.Tensor(["N"], "float32")):
- N = T.int64()
+ N = T.int32()
if N % 16 == 0:
out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
else:
@@ -1372,7 +1372,7 @@ def test_tir_expr_as_branch_condition():
@R.function(private=True)
def unsugared(x: R.Tensor(["N"], "float32")):
- N = T.int64()
+ N = T.int32()
if R.prim_value(N % 16 == 0):
out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
else:
@@ -1417,7 +1417,7 @@ def test_computed_prim_value_as_assert_condition():
@R.function(pure=False)
def func(x: R.Tensor(["N"], "float32")):
- N = T.int64()
+ N = T.int32()
_ = R.assert_op(R.prim_value(N % 16 == 0))
out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
return out
@@ -1435,14 +1435,14 @@ def test_tir_expr_as_assert_condition():
@R.function(pure=False, private=True)
def sugared(x: R.Tensor(["N"], "float32")):
- N = T.int64()
+ N = T.int32()
_ = R.assert_op(N % 16 == 0)
out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
return out
@R.function(pure=False, private=True)
def unsugared(x: R.Tensor(["N"], "float32")):
- N = T.int64()
+ N = T.int32()
_ = R.assert_op(R.prim_value(N % 16 == 0))
out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
return out
@@ -1468,7 +1468,7 @@ def test_erase_to_well_defined_keeps_variables_exposed_by_tensor_shape():
@R.function
def foo(x: R.Tensor(["m", "n"])):
q = x
- m, n = T.int64(), T.int64()
+ m, n = T.int32(), T.int32()
z = R.match_cast(q, R.Tensor((m, n)))
w = z
return w
@@ -1481,7 +1481,7 @@ def test_erase_to_well_defined_keeps_variants_exposed_by_shape_expr():
@R.function
def foo(x: R.Tensor, _: R.Shape(["m", "n"])):
q = x
- m, n = T.int64(), T.int64()
+ m, n = T.int32(), T.int32()
z = R.match_cast(q, R.Tensor((m, n)))
w = z
return w
@@ -1497,7 +1497,7 @@ def test_erase_to_well_defined_infers_from_shape_expr():
@R.function
def subroutine(x: R.Tensor, _: R.Shape(["m", "n"])) -> R.Tensor(["m", "n"]):
q = x
- m, n = T.int64(), T.int64()
+ m, n = T.int32(), T.int32()
z = R.match_cast(q, R.Tensor((m, n)))
w = z
return w
@@ -1538,9 +1538,9 @@ def test_symbolic_vars_in_tensor_shape_with_usage_first():
z = R.add(x, y)
return z
- m = tirx.Var("m", "int64")
- x = relax.Var("x", relax.TensorType([m + 1], "float32"))
- y = relax.Var("y", relax.TensorType([m, 1], "float32"))
+ m = tirx.Var("m", "int32")
+ x = relax.Var("x", R.Tensor([m + 1], "float32"))
+ y = relax.Var("y", R.Tensor([m, 1], "float32"))
bb = relax.BlockBuilder()
with bb.function("foo", (x, y)):
z = bb.emit(relax.op.add(x, y))
@@ -1556,13 +1556,13 @@ def test_symbolic_vars_in_tensor_shape_with_definition_first():
def bar(x: R.Tensor(("m",), "float32"), y: R.Tensor(("T.max(m, 20)",), "float32")) -> R.Tensor(
("T.max(m, 20) + 1",), "float32"
):
- m = T.int64()
+ m = T.int32()
z = R.call_dps_packed("test_intrin", (x, y), R.Tensor((T.max(m, 20) + 1,), dtype="float32"))
return z
- m = tirx.Var("m", "int64")
- x = relax.Var("x", relax.TensorType([m], "float32"))
- y = relax.Var("y", relax.TensorType([tirx.max(m, 20)], "float32"))
+ m = tirx.Var("m", "int32")
+ x = relax.Var("x", R.Tensor([m], "float32"))
+ y = relax.Var("y", R.Tensor([tirx.max(m, 20)], "float32"))
bb = relax.BlockBuilder()
with bb.function("bar", (x, y)):
z = bb.emit(
@@ -1619,26 +1619,29 @@ class Module:
_check(mod)
-def test_later_prim_param_not_adopted_by_usage_first_symbol():
- with pytest.raises(tvm.error.DiagnosticError):
- tvm.script.from_source(
- """
+def test_later_explicit_prim_param_reserves_symbol_type():
+ func = tvm.script.from_source(
+ """
@R.function
def main(x: R.Tensor(["n"], "float32"), n: R.Prim("int64")):
return x
"""
- )
+ )
+ assert func.params[0].ty.shape[0].same_as(func.params[1])
-def test_non_int64_prim_param_rejected_in_shape_annotation():
- with pytest.raises(tvm.error.DiagnosticError):
- tvm.script.from_source(
- """
+def test_i32_prim_param_preserves_shape_identity():
+ func = tvm.script.from_source(
+ """
@R.function
def main(n: R.Prim("int32"), x: R.Tensor(["n"], "float32")):
return x
"""
- )
+ )
+ dimension = func.params[1].ty.shape[0]
+ assert dimension.ty.dtype == "int32"
+ assert dimension.same_as(func.params[0])
+ assert func.params[0].ty.dtype == "int32"
def test_recursive_local_function_reuses_earlier_prim_param_in_signature():
@@ -1674,11 +1677,11 @@ def test_symbolic_vars_in_shape():
@R.function
def baz(x: R.Shape(("m",)), y: R.Tensor(("m * 2",), "float32")):
- m = T.int64()
+ m = T.int32()
z = R.call_dps_packed("test_intrin", y, R.Tensor((m * 2,), dtype="float32"))
return z
- m = tirx.Var("m", "int64")
+ m = tirx.Var("m", "int32")
x = relax.Var("x", relax.ShapeType([m]))
y = relax.Var("y", relax.TensorType([m * 2], "float32"))
bb = relax.BlockBuilder()
@@ -1689,21 +1692,18 @@ def test_symbolic_vars_in_shape():
_check(baz, bb.get()["baz"])
-def test_undefined_symbolic_var_raises_error():
- """An undefined symbolic variable in an error
+def test_compound_expression_introduces_i32_symbol():
+ """Every symbol in an expression string resolves through the function frame."""
- A symbolic variables is defined at the first site where it appears
- as a shape parameter without any modification. TVMScript does not
- support solving for a symbolic variable in terms of the argument
- shape. That is, this test case raises an error, and will not
- attempt to define `m` as either `x.shape[0]-1` or `x.shape[1]//2`.
- """
- with pytest.raises(tvm.error.DiagnosticError):
+ @R.function
+ def foo(x: R.Tensor(("m + 1", "m * 2"), "float32")):
+ z = R.add(x, x)
+ return z
- @R.function
- def foo(x: R.Tensor(("m + 1", "m * 2"), "float32")): # name 'm' is not defined
- z = R.add(x, x)
- return z
+ first, second = foo.params[0].ty.shape
+ assert first.a.same_as(second.a)
+ assert first.a.ty.dtype == "int32"
+ _check(foo)
def test_arith_operators():
@@ -1728,8 +1728,8 @@ def test_arith_operators():
t2 = tuple_expr[0][0] # <= Will normalize to two bindings
return (a0, a1, a2, a3, a4, a5, a6, c0, c1, c2, c3, t0, t1, t2)
- m = tirx.Var("m", "int64")
- n = tirx.Var("n", "int64")
+ m = tirx.Var("m", "int32")
+ n = tirx.Var("n", "int32")
x = relax.Var("x", relax.TensorType([m, n], "float32"))
y = relax.Var("y", relax.TensorType([m, n], "float32"))
bb = relax.BlockBuilder()
@@ -1760,8 +1760,8 @@ def test_arith_operators():
def test_memory_ops():
@R.function
def foo(x: R.Tensor(("m", "n"), dtype="float32")):
- m = T.int64()
- n = T.int64()
+ m = T.int32()
+ n = T.int32()
storage = R.memory.alloc_storage(
R.shape([4 * m * n]), virtual_device_index=0, storage_scope="global", dtype="float32"
)
@@ -1776,8 +1776,8 @@ def test_memory_ops():
def test_vm_ops():
@R.function(pure=False)
def foo(x: R.Tensor(("m", "n"), dtype="float32")):
- m = T.int64()
- n = T.int64()
+ m = T.int32()
+ n = T.int32()
storage = R.vm.alloc_storage(R.shape([4 * m * n]), runtime_device_index=0, dtype="uint8")
alloc = R.vm.alloc_tensor(storage, offset=0, shape=R.shape([m, n]), dtype="float32")
tensor = R.builtin.alloc_tensor(R.shape([m, n]), dtype="float32", runtime_device_index=0)
@@ -2367,7 +2367,7 @@ def test_function_attributes_are_defined():
@R.function
def subroutine(x: R.Tensor, _: R.Shape(["m", "n"])) -> R.Tensor(["m", "n"]):
q = x
- m, n = T.int64(), T.int64()
+ m, n = T.int32(), T.int32()
z = R.match_cast(q, R.Tensor((m, n)))
w = z
return w
@@ -2472,7 +2472,7 @@ def test_shared_meta_var_skips_relax_bindings():
@R.function(private=True)
def func(A: R.Tensor(["N"], "float32")):
- N: R.Prim("int64") = T.int64()
+ N: R.Prim("int32") = T.int32()
via_i = I.meta_var(N)
via_t = T.meta_var(via_i)
output = R.reshape(A, R.shape([via_t]))
@@ -2540,7 +2540,7 @@ def test_conditional_may_use_symbolic_variables_from_function_scope():
B: R.Tensor(["N"], "float32"),
cond: R.Prim("bool"),
) -> R.Tensor(["N"], "float32"):
- N = T.int64()
+ N = T.int32()
if cond:
out: R.Tensor([N], "float32") = A + B
@@ -2555,7 +2555,7 @@ def test_conditional_may_use_symbolic_variables_from_function_scope():
B: R.Tensor(["N"], "float32"),
cond: R.Prim("bool"),
):
- N = T.int64()
+ N = T.int32()
if cond:
out = A + B
else:
diff --git a/tests/python/relax/test_tvmscript_type_vars.py b/tests/python/relax/test_tvmscript_type_vars.py
index 990a9582f7..bb0fe575d0 100644
--- a/tests/python/relax/test_tvmscript_type_vars.py
+++ b/tests/python/relax/test_tvmscript_type_vars.py
@@ -21,6 +21,7 @@ from typing import TypeVar
import tvm
import tvm.testing
from tvm.script import relax as R
+from tvm.script import tirx as T
M = TypeVar("M")
UNUSED_GENERIC = TypeVar("UNUSED_GENERIC", bound=int)
@@ -64,5 +65,44 @@ def func[M: int](x: R.Tensor((M, M * 2), "float32")):
tvm.ir.assert_structural_equal(func, tvm.script.from_source(portable))
+def test_explicit_dtype_roundtrip_and_local_shadowing():
+ n = T.int64()
+
+ @R.function(private=True)
+ def func(x: R.Tensor(("n",), "float32")):
+ return x
+
+ assert func.params[0].ty.shape[0].same_as(n)
+ for pep695 in [False, True]:
+ script = func.script(extra_config={"script.use_pep695": pep695})
+ assert "n = T.int64()" in script
+ tvm.ir.assert_structural_equal(func, tvm.script.from_source(script))
+
+ @T.prim_func
+ def local():
+ n = T.int32()
+ T.evaluate(n)
+
+ assert str(local.body.value.ty.dtype) == "int32"
+
+
+def test_module_preserves_distinct_symbol_types():
+ functions = {}
+ for dtype in ["int32", "int64"]:
+ n = tvm.tirx.Var("n", dtype)
+ x = tvm.relax.Var("x", tvm.relax.TensorType([n], "float32"))
+ builder = tvm.relax.BlockBuilder()
+ with builder.function(dtype, [x]):
+ builder.emit_func_output(x)
+ functions[dtype] = builder.get()[dtype]
+ module = tvm.IRModule(functions)
+ for pep695 in [False, True]:
+ script = module.script(extra_config={"script.use_pep695": pep695})
+ # PEP695 requires a supporting host; portable syntax is always exercised.
+ if pep695 and sys.version_info < (3, 12):
+ continue
+ tvm.ir.assert_structural_equal(module, tvm.script.from_source(script))
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py
index 7a10f81538..4fbd7221ad 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -212,6 +212,8 @@ def test_vm_compile_e2e(exec_mode):
def test_vm_compile_e2e_func_param_with_shape(exec_mode):
+ m, n, k = T.int64(), T.int64(), T.int64()
+
@tvm.script.ir_module
class TestVMCompileE2E2:
@T.prim_func(s_tir=True)
@@ -560,6 +562,8 @@ def test_vm_relax_symbolic_shape(exec_mode):
def test_vm_relax_symbolic_shape_tuple(exec_mode):
+ m, n = T.int64(), T.int64()
+
@I.ir_module(s_tir=True)
class mod:
@R.function
@@ -927,6 +931,7 @@ class TestVMSetInput:
def test_multi_systemlib(exec_mode):
+ m, n = T.int64(), T.int64()
pytest.importorskip("cloudpickle") # needed by popen_pool.PopenWorker
@tvm.script.ir_module
diff --git a/tests/python/tvmscript/test_parser.py b/tests/python/tvmscript/test_parser.py
index 2b23e7e4e9..61195c72f2 100644
--- a/tests/python/tvmscript/test_parser.py
+++ b/tests/python/tvmscript/test_parser.py
@@ -30,10 +30,11 @@ from types import ModuleType, SimpleNamespace
import pytest
from tvm import ir
+from tvm.ir.prim import Cast
from tvm.script import parser
from tvm.script.ir_builder import IRBuilder, protocol
from tvm.script.parser.frontend import Compiler, make_decorator
-from tvm.script.parser.transform import Transformer
+from tvm.script.parser.transpile import IRBuilderTranspiler
class _Recorder:
@@ -93,7 +94,7 @@ def test_source_translation():
compiler, function, _ = _registered(source)
original = ast.dump(function, include_attributes=True)
counter = itertools.count()
- transformer = Transformer(
+ transformer = IRBuilderTranspiler(
compiler.filename,
compiler.env,
"X",
@@ -207,7 +208,7 @@ def test_callable_metadata_and_missing_initializer():
)
_run(compiler, function, builder)
_, metadata = builder.bindings["n"]
- assert metadata["declaration"] is True
+ assert "declaration" not in metadata # Assignment lowering never inspects runtime kinds.
assert metadata["previous"] is previous
assert builder.bindings["absent"][0] is protocol.MISSING
assert builder.bindings["explicit"][0] is None
@@ -289,6 +290,12 @@ def test_shared_parser_dependency_direction():
values = [node.attr]
elif isinstance(node, ast.Constant) and isinstance(node.value, str):
values = [node.value]
+ if isinstance(node, ast.ImportFrom):
+ assert node.module not in {"tvm.ir", "tvm.ir.prim"}, path
+ if node.module == "tvm":
+ assert all(alias.name != "ir" for alias in node.names), path
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+ assert node.func.id != "eval", path
if any(
forbidden.intersection(re.findall(r"[A-Za-z_]\w*", value.lower()))
for value in values
@@ -356,16 +363,26 @@ def test_annotation_identity_effects_and_recovery(postponed, monkeypatch):
assert len(module.errors) == 2
assert "annotation argument failed" in str(module.errors[0])
assert "unknown dtype" in str(module.errors[1]).lower()
- assert module.calls == ["x", "x", "x", "return", "x", "return", "object", "object"]
+ # Source annotations are safe to re-evaluate, including ordinary eager mode;
+ # the implementation deliberately gives no exactly-once side-effect contract.
+ assert module.calls.count("x") >= 4
+ assert module.calls.count("return") >= 2
+ assert module.calls.count("object") >= 2
+ module.captures[:] = [value for value in module.captures if not value.is_missing()]
+
+ def symbol(dimension):
+ # Explicit casts preserve the same underlying canonical symbol.
+ return dimension.value if isinstance(dimension, Cast) else dimension
+
for first, second in zip(module.captures, module.captures[1:]):
- assert not first.shape[0].same_as(second.shape[0])
+ assert not symbol(first.shape[0]).same_as(symbol(second.shape[0]))
first, second = module.functions
for function in module.functions:
for argument_dim, return_dim in zip(function.params[0].ty.shape, function.ret_ty.shape):
- assert argument_dim.same_as(return_dim)
- assert str(argument_dim.ty.dtype) == "int64"
+ assert symbol(argument_dim).same_as(symbol(return_dim))
+ assert str(symbol(argument_dim).ty.dtype) == "int32"
for first_dim, second_dim in zip(first.params[0].ty.shape, second.params[0].ty.shape):
- assert not first_dim.same_as(second_dim)
+ assert not symbol(first_dim).same_as(symbol(second_dim))
assert module.quoted.params[0].ty == module.quoted.ret_ty
@@ -405,3 +422,42 @@ def test_annotation_module_context(postponed, monkeypatch):
for annotation in [function.params[0].ty, function.ret_ty]:
actual = getattr(annotation, field)
assert actual.__chandle__() == module.global_infos[key][0].__chandle__()
+
+
+@pytest.mark.parametrize("postponed", [False, True], ids=["eager", "postponed"])
+def test_composite_missing_annotations(postponed, monkeypatch):
+ module = _execute_annotations(
+ """
+ from tvm.script import relax as R
+ eager_type = R.Tuple(R.Tensor(("n",), "float32"), R.Object())
+ @R.function
+ def f(x: R.Tuple(R.Tensor(("n",), "float32"), R.Object())):
+ return x
+ """,
+ postponed,
+ monkeypatch,
+ )
+ assert module.eager_type.fields[0].is_missing()
+ assert not module.eager_type.fields[1].is_missing()
+ actual = module.f.params[0].ty.fields[0]
+ assert str(actual.shape[0].ty.dtype) == "int32"
+ assert module.f.ret_ty == module.f.params[0].ty
+
+
+def test_shared_allocator_preserves_source_names():
+ compiler, function, builder = _registered(
+ """
+ @D.function
+ def f():
+ _t0 = 1
+ value0 = 2
+ _t1 = _t0 + value0
+ return _t1
+ """
+ )
+ _run(compiler, function, builder)
+ assert builder.returned == [3]
+ assert {"_t0", "value0", "_t1"}.issubset(builder.bindings)
+ transformer = compiler.transformer()
+ assert transformer.fresh_unique_name("value") != "value0"
+ assert transformer.fresh_unique_name() not in {"_t0", "_t1"}