[REFACTOR][TVMScript] Simplify annotation capture and transpiler contexts

docs/arch/tvmscript.rst

@@ -45,12 +45,37 @@ function decorators retain definitions until module construction, which can decl
 signatures before building bodies. Captured Python values belong to the source
 context; symbolic IR values are created and resolved by builders.
 
-Use ``I.dynamic("n")`` or the identical ``T.dynamic`` constructor to create a fresh
-primitive symbol, defaulting to int64. Reuse that object in ordinary Python shape,
+Function construction, JIT and macro decorators require ``@`` application at the
+function definition site. Later application to an existing function is rejected.
+Use a registered namespace, such as ``@T.prim_func`` or ``@Ts.prim_func(private=True)``.
+Python definitions support namespace aliases such as ``Alias = T``; bare callable
+aliases and preconfigured decorator aliases are unsupported. Source strings resolve
+namespace aliases from imports or ``extra_vars``, not executable prefix assignments.
+
+``GeneratedBuilder`` records connect generated body helpers to original functions and
+identify bindings preserved during recomposition. Original-name wrapper parameters
+capture definition values as defaults. Execution globals remain separate from the
+locals that evaluate those defaults. Known unshadowed signature values use these
+parameters; lazy snapshots remain for missing or shadowed names and body annotations.
+Class setup and declaration snapshots execute in source order before function bodies.
+Body globals, closures and local shadowing retain their Python scope. Missing values
+are read only when needed, including conditional and optional annotations.
+
+Annotations read concrete values from their definition scope, preserving missing-name
+errors when a value is used. Create external symbols with ``n = I.dynamic("n")``
+or the identical ``T.dynamic`` and ``Ts.dynamic`` constructors. Each call creates a
+fresh native variable, defaulting to int64. Reuse that object in ordinary Python shape,
 stride and offset expressions to share identity; strings in those fields are not
-parsed as expressions. Whole quoted Python annotations remain supported. On Python
-3.12+, function parameters such as ``def f[n, k: T.int32](...)`` introduce local
-symbols with default int64 and explicit int32 dtypes.
+parsed as expressions. On Python 3.12+, explicit headers such as
+``def f[n, k: T.int32](...)`` declare local symbols; ``n: int`` retains the int64
+default. Quote the whole annotation or use ``from __future__ import annotations``
+to defer eager Python evaluation of header symbols. Captured runtime ``typing.TypeVar``
+objects are not script symbols; ordinary Python typing uses remain unaffected.
+
+An explicit scalar annotation ``n: n`` preserves a captured native symbol's identity.
+An independently typed parameter such as ``n: T.int32`` and ordinary body locals
+shadow definition captures normally. Annotation classes, Python unions and deferred
+return-constructor evaluation retain their builder behavior.
 
 Syntax and construction protocol
 --------------------------------
@@ -59,14 +84,23 @@ The syntax transpiler rewrites a fresh Python AST using scope and declaration fa
 from a prescan. A registered decorator selects the construction namespace. Assignments
 call binding hooks, expression statements call emission hooks, and control flow opens
 builder frames. The namespace owns the meaning of these operations and the supported
-IR constructs.
+IR constructs. Standalone decorators pass already-evaluated ``root_function_kwargs``
+so option expressions execute once. Nested functions and module members use their own
+decorator options. Builder hooks own option defaults; ``check_well_formed`` remains a
+separate parser setting.
 
 ``tvm.script.parser.protocol_registry`` records syntax policies under registered
-namespace paths. These policies identify declarations and scalar annotation dtypes.
-Source aliases resolve to those paths;
+namespace paths. These policies identify scalar annotations, mutable declarations and
+result span handling. Symbolic shapes use concrete expressions. Source aliases resolve to those paths;
 ordinary Python calls remain calls in the generated program. Explicit ``constexpr``
 markers select host control flow during construction.
 
+JIT supplies ``const_args``, a mapping from parameter names to fixed values. Explicit
+``None`` values mark absent optional arguments and skip their annotation evaluation.
+Generated code reads the mapping through a fresh ``_const_args`` name. An empty map
+still selects root JIT construction; an absent map selects ordinary parsing. Syntax
+translation does not inspect these values to select a branch.
+
 Generated operations retain source locations. Syntax restrictions raise source-located
 ``SyntaxError`` exceptions; builder and Python helper errors retain their original
 exception types. Temporary parse state is released when construction finishes or fails.

docs/reference/api/python/script/parser.rst

@@ -24,6 +24,14 @@ tvm.script.parser
    :members:
    :imported-members:
 
+tvm.script.parser.jit_support
+*****************************
+These helpers carry ``const_args``, the selected fixed argument map, into generated builder
+execution. Parameter validation and caching remain with the JIT entry point.
+
+.. automodule:: tvm.script.parser.jit_support
+   :members: use_specialization, read_specialization_bindings
+
 tvm.script.parser.protocol_registry
 ***********************************
 Language variants register source syntax policies independently of their

python/tvm/relax/script/ir_builder/__init__.py

@@ -28,7 +28,6 @@
 from tvm.relax.distributed import Placement as _Placement
 from tvm.relax.distributed import device_mesh as device_mesh
 from tvm.script.ir_builder import resolve_global_info_args as _resolve_global_info_args
-from tvm.script.ir_builder.base import annotation_constructor as _annotation_constructor
 from tvm.script.ir_builder.base import at as _at
 from tvm.script.ir_builder.base import source_span as _source_span
 from tvm.script.parser.protocol_registry import constexpr as constexpr
@@ -83,7 +82,6 @@
 
 
 @_resolve_global_info_args("vdevice", resolver=resolve_global_info_)
-@_annotation_constructor("shape")
 def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
     """Construct a Relax tensor type.
 
@@ -108,8 +106,8 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
 
     Returns
     -------
-    result : TensorType or Type
-        The tensor type, or a missing type for an unresolved eager shape annotation.
+    result : TensorType
+        The constructed tensor type.
         String selectors outside an active module always raise ValueError.
     """
     if isinstance(shape, _python.str) and dtype is None:
@@ -118,7 +116,6 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
 
 
 @_resolve_global_info_args("device_mesh", resolver=resolve_global_info_)
-@_annotation_constructor("shape")
 def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1, span=None):
     """Construct a Relax distributed tensor type.
 
@@ -143,8 +140,8 @@ def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1,
 
     Returns
     -------
-    result : DTensorType or Type
-        The distributed type, or a missing type for an unresolved eager shape annotation.
+    result : DTensorType
+        The constructed distributed type.
         String selectors outside an active module always raise ValueError.
     """
     if device_mesh is None:
@@ -154,7 +151,7 @@ def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1,
     return _DTensorType(Tensor(shape, dtype, ndim=ndim), device_mesh, placement, _source_span(span))
 
 
-# The distributed source spelling shares concrete constructors and argument policy.
+# The distributed source spelling shares the decorated concrete constructor.
 dist.DTensor = DTensor
 dist.device_mesh = device_mesh
 
@@ -164,7 +161,6 @@ def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1,
 __tvm_value_if__ = True
 
 
-@_annotation_constructor("values")
 def Shape(values=None, ndim=-1, *, span=None):
     """Construct a Relax shape type.
 

python/tvm/script/ir_builder/__init__.py

@@ -26,9 +26,7 @@
     MISSING,
     AlreadyEmitted,
     IRBuilder,
-    annotation_value_,
     at_,
-    require_defined,
     resolve_global_info_args,
     with_at_group_,
 )
@@ -57,7 +55,6 @@
     "Range",
     "StringImm",
     "StringType",
-    "annotation_value_",
     "at_",
     "check_well_formed_",
     "constexpr",
@@ -72,7 +69,6 @@
     "module_global_infos",
     "module_member_",
     "module_set_attr",
-    "require_defined",
     "resolve_global_info_args",
     "with_at_group_",
 ]

python/tvm/script/ir_builder/base.py

@@ -23,7 +23,7 @@
 from typing import Any, Generic, TypeVar
 
 from tvm_ffi import register_object as _register_object
-from tvm_ffi.dataclasses import MISSING
+from tvm_ffi.dataclasses import MISSING as MISSING
 
 from tvm import ir
 from tvm.runtime import Object as _Object
@@ -377,32 +377,6 @@ def at(
     return value
 
 
-def require_defined(value, name):
-    """Report a source name whose designated region output was not produced.
-
-    Parameters
-    ----------
-    value : Any
-        Candidate binding. Only the canonical ``MISSING`` singleton denotes an
-        absent value; explicit None is a defined value.
-    name : str
-        Source identifier to include in the missing-name diagnostic.
-
-    Returns
-    -------
-    Any
-        The exact ``value`` when it is defined.
-
-    Raises
-    ------
-    NameError
-        If ``value`` is ``MISSING``.
-    """
-    if value is MISSING:
-        raise NameError(f"name {name!r} is not defined")
-    return value
-
-
 def with_at_group_(
     location: SpanEntry | ir.Span | tuple[str | ir.SourceName, int, int, int, int] | None,
     thunk: Callable[[], _T],
@@ -465,116 +439,34 @@ def _resolve_type_var(frame, ffi_resolver, name, dtype=None, *, value=None, span
 
 
 def _current_function_frame():
-    """Find the nearest function for eager shared annotation constructors."""
+    """Find the nearest native function frame for explicit symbol declarations."""
     if IRBuilder.is_in_scope():
         for frame in reversed(IRBuilder.current().frames):
             if callable(getattr(frame, "resolve_type_var", None)):
                 return frame
     raise ValueError("Symbol resolution requires an active function frame")
 
 
-def annotation_constructor(*fields: str, as_type: bool = False):
-    """Adapt eager Python type parameters on concrete annotation constructors.
-
-    Unresolved ``typing.TypeVar`` values defer annotations outside a builder;
-    an active native function owns their resolution. Ordinary values, including
-    strings, are passed unchanged to the concrete API.
-    """
-
-    def decorate(constructor):
-        call_signature = signature(constructor)
-
-        def unresolved(value):
-            if isinstance(value, TypeVar):
-                return True
-            if isinstance(value, tuple | list):
-                return any(unresolved(item) for item in value)
-            return False
-
-        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 _current_function_frame().resolve_type_var(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
-
-        @wraps(constructor)
-        def invoke(*args, **kwargs):
-            bound = call_signature.bind(*args, **kwargs)
-            for field in fields:
-                if field not in bound.arguments:
-                    continue
-                value = bound.arguments[field]
-                if IRBuilder.is_in_scope():
-                    bound.arguments[field] = resolve(value)
-                elif unresolved(value):
-                    return ir.Type.missing()
-            return constructor(*bound.args, **bound.kwargs)
-
-        if not as_type:
-            return invoke
-        # A real annotation class supports Python unions while constructing
-        # ordinary native types, with no proxy values or parser policy state.
-        return type(
-            constructor.__name__,
-            (),
-            {
-                "__new__": lambda cls, *args, **kwargs: invoke(*args, **kwargs),
-                "__signature__": call_signature,
-                "__doc__": constructor.__doc__,
-                "__module__": constructor.__module__,
-            },
-        )
-
-    return decorate
-
-
 def _return_annotation(annotation):
-    """Evaluate a return annotation without introducing return-only symbols."""
-    if not callable(annotation) or isinstance(annotation, ir.Expr | ir.Type):
-        return annotation
-    frame = _current_function_frame()
-    declared = set(frame.type_var_map)
-    annotation = annotation()
-    introduced = set(frame.type_var_map) - declared
-    if introduced:
-        raise ValueError(f"Return annotation introduces unbound symbol {sorted(introduced)[0]!r}")
+    """Evaluate the deferred return expression before normalizing its annotation value."""
+    if callable(annotation) and not isinstance(annotation, ir.Expr | ir.Type):
+        return annotation()
     return annotation
 
 
-def annotation_value_(name, value):
-    """Adapt a real definition-context symbol using the native function map.
-
-    Parameters
-    ----------
-    name : str
-        Source spelling used to resolve the symbol in the nearest active
-        function frame.
-    value : TypeVar, Var or Any
-        Captured definition-context value. An unconstrained ``typing.TypeVar``
-        resolves a symbolic variable; a primitive IR Var supplies its existing
-        value to the resolver. Other values pass through unchanged.
+def annotation_constructor(constructor):
+    """Expose a constructor as a real annotation class supporting Python unions.
 
-    Returns
-    -------
-    Any
-        The function frame's resolved symbol, or the unchanged nonsymbolic value.
-
-    Raises
-    ------
-    TypeError
-        If a TypeVar has a bound or constraints.
-    ValueError
-        If a symbolic value requires resolution without an active function frame.
+    Calls construct ordinary native values directly. The class preserves the
+    constructor's signature and documentation without adapting its arguments.
     """
-    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 _current_function_frame().resolve_type_var(name)
-    if ir.is_prim_var(value):
-        return _current_function_frame().resolve_type_var(name, value=value)
-    return value
+    return type(
+        constructor.__name__,
+        (),
+        {
+            "__new__": lambda cls, *args, **kwargs: constructor(*args, **kwargs),
+            "__signature__": signature(constructor),
+            "__doc__": constructor.__doc__,
+            "__module__": constructor.__module__,
+        },
+    )

python/tvm/script/parser/__init__.py

@@ -20,7 +20,7 @@
 
 import importlib
 from collections.abc import Callable
-from typing import Any, TypeVar
+from typing import Any
 
 _NAMESPACES: dict[str, object] = {}
 _NAMESPACE_INITIALIZERS: list[Callable[[], None]] = []
@@ -104,7 +104,6 @@ def _initialize() -> None:
     _initializing = True
     try:
         importlib.import_module(__name__ + ".ir")
-        register_namespace("TypeVar", TypeVar)
         for initialize in _NAMESPACE_INITIALIZERS:
             initialize()
         _initialized = True