This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 974a099a75 [REFACTOR][TVMScript] Fix syntax roots and simplify
annotation scopes (#20446)
974a099a75 is described below
commit 974a099a75deecb4d90e8f1caa5e08e4dee27f69
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Sep 25 10:09:36 2026 -0400
[REFACTOR][TVMScript] Fix syntax roots and simplify annotation scopes
(#20446)
TVMScript selects special syntax from fixed dialect roots and reserved
`range`/`int` names. Validate invalid bindings during prescan, then
match direct `X.member` calls without reconstructing Python shadowing or
following callable aliases.
Annotation helpers capture values under their original names so Python
handles parameter, lambda, comprehension, and body scopes. Keep
declaration snapshots separate from body closures, preserve lazy
missing-name checks, ordinary calls and symbolic range lowering, and use
ordinary call contexts for nested namespace members.
Explicit symbols remove the need for unconditional-declaration tracking
and expression-string namespace context. Remove unused module-decorator
metadata while retaining definition-site decorator options, lazy
captures and source diagnostics.
---
docs/arch/tvmscript.rst | 5 +-
docs/reference/api/python/script/parser.rst | 2 +-
python/tvm/backend/trn/script.py | 4 +-
python/tvm/relax/script/__init__.py | 6 -
python/tvm/relax/script/ir_builder/ir.py | 5 -
python/tvm/script/ir_builder/parser_protocol.py | 1 -
python/tvm/script/parser/entry.py | 32 +-
python/tvm/script/parser/prescan.py | 231 +++++---
python/tvm/script/parser/protocol_registry.py | 37 +-
python/tvm/script/parser/transpile.py | 637 ++++++++-------------
python/tvm/tirx/script/jit.py | 20 +-
tests/python/script/test_basic_usage.py | 8 +-
tests/python/script/test_error_handling.py | 94 +++
tests/python/script/test_jit_specialization.py | 52 ++
tests/python/script/test_meta_programming.py | 68 +++
tests/python/script/test_parser_entry.py | 16 +-
.../python/script/test_special_parser_protocol.py | 34 ++
17 files changed, 702 insertions(+), 550 deletions(-)
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index 5740a4f4dc..612527f591 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -55,8 +55,9 @@ namespace aliases from imports or ``extra_vars``, not
executable prefix assignme
``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.
+locals that evaluate those defaults. Declaration and annotation helpers bind
snapshot
+values under their original names, so Python handles parameter and local
shadowing.
+Missing snapshot values raise only when their annotation expressions read them.
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.
diff --git a/docs/reference/api/python/script/parser.rst
b/docs/reference/api/python/script/parser.rst
index 8f41fa584f..70a7cc989a 100644
--- a/docs/reference/api/python/script/parser.rst
+++ b/docs/reference/api/python/script/parser.rst
@@ -40,7 +40,7 @@ registered through
:func:`tvm.script.parser.register_namespace` and
:func:`tvm.script.parser.register_namespace_initializer`.
.. automodule:: tvm.script.parser.protocol_registry
- :members: constexpr, register_scalar_annotation, mutable_cell_decl,
result_span, module_decorator, declaration_kind
+ :members: constexpr, register_scalar_annotation, mutable_cell_decl,
result_span, declaration_kind
The language variant aliases below share the public construction namespaces
documented
in :doc:`script`. Parser entry points above use the canonical frontend.
diff --git a/python/tvm/backend/trn/script.py b/python/tvm/backend/trn/script.py
index 701523f273..03d134d582 100644
--- a/python/tvm/backend/trn/script.py
+++ b/python/tvm/backend/trn/script.py
@@ -21,8 +21,6 @@ from __future__ import annotations
from collections.abc import Callable
from typing import Any
-from tvm.script.parser.protocol_registry import result_span as _result_span
-
from . import op as _trn_op
OpWrapper = Callable[[Callable[..., Any]], Callable[..., Any]]
@@ -33,7 +31,7 @@ def _default_op_wrapper(func: Callable[..., Any]) ->
Callable[..., Any]:
from tvm.tirx.script.ir_builder.ir import _op_wrapper # pylint:
disable=import-outside-toplevel
# NKI producers return a single call_intrin node and never emit separately.
- return
_result_span(f"T.nki.{func.__name__.removeprefix('nki_')}")(_op_wrapper(func))
+ return _op_wrapper(func)
class NKINamespace:
diff --git a/python/tvm/relax/script/__init__.py
b/python/tvm/relax/script/__init__.py
index b1bdfe93e8..fd4d9b9686 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -26,8 +26,6 @@ from collections.abc import Callable as _Callable
from typing import TYPE_CHECKING
from typing import Any as _Any
-from tvm.script.parser.protocol_registry import module_decorator as
_module_decorator
-
if TYPE_CHECKING:
from tvm.ir import IRModule
from tvm.relax.base_py_module import BasePyModule
@@ -212,7 +210,3 @@ def py_module(
return result
return apply(module) if module is not None else apply
-
-
-# Member decorators may defer to this shared-parser entry without importing
Relax.
-_module_decorator("R.py_module")(py_module)
diff --git a/python/tvm/relax/script/ir_builder/ir.py
b/python/tvm/relax/script/ir_builder/ir.py
index 4996b000e6..58f46e5813 100644
--- a/python/tvm/relax/script/ir_builder/ir.py
+++ b/python/tvm/relax/script/ir_builder/ir.py
@@ -209,7 +209,6 @@ from tvm.runtime._tensor import (
)
from tvm.script.ir_builder import IRBuilder, decl_function
from tvm.script.ir_builder.ir import _class_global_infos
-from tvm.script.parser.protocol_registry import module_decorator as
_module_decorator
from . import _ffi_api, frame
@@ -474,10 +473,6 @@ def rewriter(rewriter_mod: IRModule | type) ->
PatternMatchingRewriter:
return PatternMatchingRewriter.from_module(rewriter_mod)
-# Member decorators defer to the same shared module construction boundary.
-_module_decorator("R.rewriter")(rewriter)
-
-
############################# BindingBlock ##############################
diff --git a/python/tvm/script/ir_builder/parser_protocol.py
b/python/tvm/script/ir_builder/parser_protocol.py
index 258d5b9b3b..1d09c15eaa 100644
--- a/python/tvm/script/ir_builder/parser_protocol.py
+++ b/python/tvm/script/ir_builder/parser_protocol.py
@@ -874,7 +874,6 @@ def module_member_(name: str, value: Any) -> Any:
# ``register_scalar_annotation(path, constructor, dtype=...)`` describes
scalar annotations.
# ``mutable_cell_decl(path)`` marks mutable storage declarations.
# ``result_span(path)`` permits attaching a call's result span without a call
context.
-# ``module_decorator(path)`` marks module declaration decorators.
# ``declaration_kind(path, kind)`` marks function and helper declarations.
# ``with X.function_(...) as fn:`` builds a function frame.
# ``a = X.arg("a", ty)`` declares a parameter.
diff --git a/python/tvm/script/parser/entry.py
b/python/tvm/script/parser/entry.py
index ceb9d94d0c..a5c26857bb 100644
--- a/python/tvm/script/parser/entry.py
+++ b/python/tvm/script/parser/entry.py
@@ -503,7 +503,9 @@ def make_macro_decorator(
{**environment, **bound.arguments},
set(bound.arguments),
preserve_return=preserve_return,
- definition_scope=definition_scope,
+ definition_scope=(
+ definition_scope if options.get("hygienic", True) else
environment
+ ),
)
return invoke
@@ -672,7 +674,16 @@ def _run_statements(
)
namespace[transformer.function.dialect_prefix] = builder
node = tree.body[-1]
- statements = transformer.transform_statements(node.body)
+ # Macros enter through their body instead of function declaration lowering.
+ # Snapshot annotations before defining the helper, in the definition scope.
+ # Bound arguments remain original-name parameters of the macro body.
+ _, _, annotation_names = transformer._read_function_annotations(
+ node, [], transformer.module.prescan_ctx.bindings.get(node, ())
+ )
+ captures = transformer.module.make_fresh_name("_definition")
+ transformer.function.definition_captures = captures
+ statements = transformer._create_definition_bindings(node,
annotation_names, captures=captures)
+ body = transformer.transform_statements(node.body)
names = sorted(name for name in bound_names if name in namespace)
helper_name = transformer.module.make_fresh_name("_macro")
helper = ast.copy_location(
@@ -685,7 +696,7 @@ def _run_statements(
kw_defaults=[],
defaults=[],
),
- statements or [ast.Pass()],
+ body or [ast.Pass()],
[],
None,
),
@@ -693,17 +704,23 @@ def _run_statements(
)
if "type_params" in ast.FunctionDef._fields:
helper.type_params = []
+ transformer.module.generated_builders.append(
+ GeneratedBuilder(
+ helper, node.name, set(names) |
transformer.module.prescan_ctx.namespaces.keys()
+ )
+ )
runnable = _recompose_builder(
- ast.Module([helper], []),
+ ast.Module([*statements, helper], []),
source_fn=source,
definition_scope=transformer.module.definition_scope,
generated_builders=transformer.module.generated_builders,
filename=filename,
flags=flags,
- name=helper_name,
+ name=transformer.module.make_fresh_name("_builder"),
make_fresh_name=transformer.module.make_fresh_name,
environment=namespace,
- )
+ result=helper_name,
+ )()
return runnable(*(namespace[name] for name in names))
@@ -894,7 +911,4 @@ def ir_module(module: type | None = None, **options: Any)
-> IRModule | Callable
return apply(module) if module is not None else apply
-syntax_protocol.module_decorator("I.ir_module")(ir_module)
-syntax_protocol.module_decorator("script.ir_module")(ir_module)
-
from_source = parse
diff --git a/python/tvm/script/parser/prescan.py
b/python/tvm/script/parser/prescan.py
index 2cc344a582..f5aea4b0bb 100644
--- a/python/tvm/script/parser/prescan.py
+++ b/python/tvm/script/parser/prescan.py
@@ -19,8 +19,9 @@
from __future__ import annotations
import ast
+import builtins
import inspect
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
from types import ModuleType
from typing import NamedTuple, NoReturn
@@ -28,24 +29,23 @@ from . import protocol_registry as protocol
from .annotation import parse_annotation
-def collect_annotation_free_names(
- node: ast.expr, bound: set[str] | None = None
-) -> dict[str, ast.Name]:
- """Find annotation reads, respecting Python lambda/comprehension binders.
+def collect_annotation_free_names(node: ast.expr) -> dict[str, ast.Name]:
+ """Find annotation captures and their source introduction locations."""
+ return {reference.id: reference for reference in
collect_annotation_free_reads(node)}
- The returned nodes retain the introduction locations for diagnostics and
- supply the names needed by temporary/deferred definition capture.
- """
+
+def collect_annotation_free_reads(node: ast.expr, bound: set[str] | None =
None) -> list[ast.Name]:
+ """Collect original free reads for captures and lazy missing-name
checks."""
bound = set() if bound is None else bound
if isinstance(node, ast.Name):
- return {node.id: node} if isinstance(node.ctx, ast.Load) and node.id
not in bound else {}
+ return [node] if isinstance(node.ctx, ast.Load) and node.id not in
bound else []
if isinstance(node, ast.Lambda):
# -------------------- Pattern --------------------
# Python source:
# lambda n=outer: shape(n)
#
# Builder:
- # lambda n=captured_outer: shape(n)
+ # lambda n=outer: shape(n)
# -------------------------------------------------
# Only outer is free; n belongs to the lambda.
arguments = [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]
@@ -53,12 +53,12 @@ def collect_annotation_free_names(
arguments.append(node.args.vararg)
if node.args.kwarg:
arguments.append(node.args.kwarg)
- result: dict[str, ast.Name] = {}
+ result: list[ast.Name] = []
for value in [*node.args.defaults, *node.args.kw_defaults]:
if value is not None:
- result.update(collect_annotation_free_names(value, bound))
- result.update(
- collect_annotation_free_names(node.body, bound | {arg.arg for arg
in arguments})
+ result.extend(collect_annotation_free_reads(value, bound))
+ result.extend(
+ collect_annotation_free_reads(node.body, bound | {arg.arg for arg
in arguments})
)
return result
if isinstance(node, ast.ListComp | ast.SetComp | ast.DictComp |
ast.GeneratorExp):
@@ -67,35 +67,34 @@ def collect_annotation_free_names(
# tuple(n for n in shape)
#
# Builder:
- # tuple(n for n in captured_shape)
+ # tuple(n for n in shape)
# -------------------------------------------------
# Only shape is free; n belongs to the comprehension.
- result = {}
+ result = []
local = set(bound)
for generator in node.generators:
- result.update(collect_annotation_free_names(generator.iter, local))
+ result.extend(collect_annotation_free_reads(generator.iter, local))
local.update(
item.id for item in ast.walk(generator.target) if
isinstance(item, ast.Name)
)
for condition in generator.ifs:
- result.update(collect_annotation_free_names(condition, local))
+ result.extend(collect_annotation_free_reads(condition, local))
values = [node.key, node.value] if isinstance(node, ast.DictComp) else
[node.elt]
for value in values:
- result.update(collect_annotation_free_names(value, local))
+ result.extend(collect_annotation_free_reads(value, local))
return result
- result = {}
+ result = []
for child in ast.iter_child_nodes(node):
if isinstance(child, ast.expr):
- result.update(collect_annotation_free_names(child, bound))
+ result.extend(collect_annotation_free_reads(child, bound))
elif isinstance(child, ast.keyword):
- result.update(collect_annotation_free_names(child.value, bound))
+ result.extend(collect_annotation_free_reads(child.value, bound))
return result
def resolve_namespace_key(
node: ast.AST | None,
environment: Mapping[str, object],
- bindings: Sequence[Binding] = (),
) -> str | None:
"""Normalize a fixed namespace alias, without resolving its members or
receivers."""
from . import _NAMESPACES
@@ -104,7 +103,7 @@ def resolve_namespace_key(
while isinstance(node, ast.Attribute):
parts.append(node.attr)
node = node.value
- if not isinstance(node, ast.Name) or any(item.name == node.id for item in
bindings):
+ if not isinstance(node, ast.Name):
return None
owner = environment.get(node.id)
parts.reverse()
@@ -139,15 +138,22 @@ def resolve_namespace_value(node: ast.AST | None,
environment: Mapping[str, obje
return value
-def resolve_constructor(
- node: ast.AST | None, environment: Mapping[str, object], bindings:
Sequence[Binding] = ()
-) -> str | None:
- """Read only a fixed-namespace source call's canonical policy key."""
- return (
- resolve_namespace_key(node.func, environment, bindings)
- if isinstance(node, ast.Call)
- else None
- )
+# Design principle: Once builtin names and language namespaces are selected
+# for script syntax, they cannot be shadowed. Prescan enforces this rule and
+# reports clear errors at the offending bindings. This lets the transpiler
+# rely on fixed bindings without tracking shadowing.
+#
+# Imports may establish an initial dialect root;
+# aliases created inside a function cannot rename it. Other names and calls
keep
+# ordinary Python scope and bind_ behavior. Annotation captures and IRModule
aliases
+# serve separate purposes and retain their own declaration facts.
+def _match_special_func(node: ast.AST | None, namespaces: Mapping[str, str])
-> str | None:
+ """Match direct X.member syntax without evaluating a callee or receiver."""
+ if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
+ root = namespaces.get(node.value.id)
+ if root is not None:
+ return f"{root}.{node.attr}"
+ return None
class Binding(NamedTuple):
@@ -158,7 +164,6 @@ class Binding(NamedTuple):
kind: str
annotation: ast.expr | None = None
dtype: object = None
- direct: bool = False
class PrescanContext:
@@ -179,7 +184,7 @@ class PrescanContext:
bindings: dict[ast.AST | None, list[Binding]],
sites: dict[ast.AST, Binding],
conditional_outputs: dict[ast.stmt, str],
- namespaces: set[str],
+ namespaces: dict[str, str],
with_outputs: dict[ast.With, list[str]],
recursive_functions: set[ast.FunctionDef | ast.AsyncFunctionDef],
) -> None:
@@ -198,7 +203,7 @@ class PrescanContext:
}
# Collected branch-ending names select conditional results without
later updates.
self.conditional_outputs = conditional_outputs
- # Collected entry namespaces forbid source rebinding; the set stays
read-only.
+ # Fixed source roots map to canonical dialect names; the map stays
read-only.
self.namespaces = namespaces
# Collected explicit output names select exports after each
with-region exits;
# rewriting reads the original lists without adding inferred outputs.
@@ -207,6 +212,10 @@ class PrescanContext:
# lowering; rewriting does not add or remove functions from this set.
self.recursive_functions = recursive_functions
+ def _match_special_func(self, node: ast.AST | None) -> str | None:
+ """Match direct source syntax against the roots fixed by this
prescan."""
+ return _match_special_func(node, self.namespaces)
+
class PrescanCollector(ast.NodeVisitor):
"""Collect binding syntax once, then discard all traversal accumulators."""
@@ -220,7 +229,7 @@ class PrescanCollector(ast.NodeVisitor):
self.bindings: dict[ast.AST | None, list[Binding]] = {}
self.sites: dict[ast.AST, Binding] = {}
self.outputs: dict[ast.stmt, str] = {}
- self.namespaces: set[str] = set()
+ self.namespaces: dict[str, str] = {}
self.exports: dict[ast.With, list[str]] = {}
self.recursive: set[ast.FunctionDef | ast.AsyncFunctionDef] = set()
# Active source declarations, used only to recognize self references.
@@ -228,11 +237,9 @@ class PrescanCollector(ast.NodeVisitor):
self.module_name: str | None = None
# Temporary lexical with-stack routes explicit output calls, then
resets.
self.regions: list[ast.With] = []
- # Lexical scope and language variant restore on function/class exit.
direct marks
- # unconditional function-body declarations eligible before signatures.
+ # Lexical scope and language variant restore on function/class exit.
self.scope: ast.Module | ast.ClassDef | ast.FunctionDef |
ast.AsyncFunctionDef | None = None
self.builder: object = None
- self.direct: bool = False
def collect(self, tree: ast.Module) -> PrescanContext:
"""Collect reserved names, scoped declarations and region-result
syntax.
@@ -254,12 +261,21 @@ class PrescanCollector(ast.NodeVisitor):
SyntaxError
If source bindings or declarations violate a parser restriction.
"""
- # Only registered namespace objects establish fixed source aliases.
- self.namespaces.update(
- name
- for name in self.environment
- if resolve_namespace_key(ast.Name(name, ast.Load()),
self.environment) is not None
- )
+ from . import _NAMESPACES
+
+ # Resolve roots once. Qualified receivers remain ordinary Python
expressions.
+ for name, value in self.environment.items():
+ for alias, namespace in _NAMESPACES.items():
+ if value is namespace:
+ self.namespaces[name] = alias
+ break
+ # Entry executes only these direct source-prefix imports before
prescan.
+ self.initial_imports = {
+ alias
+ for statement in tree.body[:-1]
+ if isinstance(statement, ast.Import | ast.ImportFrom)
+ for alias in statement.names
+ }
self.scope = tree
self.bindings[tree] = []
self.visit(tree)
@@ -288,6 +304,14 @@ class PrescanCollector(ast.NodeVisitor):
),
)
+ def _check_reserved(self, name: str, node: ast.AST) -> None:
+ if name in self.namespaces:
+ self._raise_error(node, f"Script namespace {name!r} cannot be
rebound or shadowed")
+ if name in ("range", "int"):
+ self._raise_error(
+ node, f"Name {name!r} is reserved for script syntax and cannot
be rebound"
+ )
+
def _record_binding(
self,
name: str,
@@ -296,10 +320,9 @@ class PrescanCollector(ast.NodeVisitor):
annotation: ast.expr | None = None,
dtype: object = None,
) -> None:
- if name in self.namespaces:
- self._raise_error(node, f"Script namespace {name!r} cannot be
rebound or shadowed")
+ self._check_reserved(name, node)
self.names.add(name)
- item = Binding(name, node, kind, annotation, dtype, self.direct)
+ item = Binding(name, node, kind, annotation, dtype)
self.bindings[self.scope].append(item)
self.sites[node] = item
@@ -312,8 +335,14 @@ class PrescanCollector(ast.NodeVisitor):
# value = X.bind_(f(), name="value")
# -------------------------------------------------
# Reserve source names and recognize self references before lowering.
- if isinstance(node.ctx, ast.Store) and node.id in self.namespaces:
- self._raise_error(node, f"Script namespace {node.id!r} cannot be
rebound or shadowed")
+ if isinstance(node.ctx, ast.Store | ast.Del):
+ self._check_reserved(node.id, node)
+ elif node.id in ("range", "int") and self.environment.get(
+ node.id, getattr(builtins, node.id)
+ ) is not getattr(builtins, node.id):
+ self._raise_error(
+ node, f"Name {node.id!r} is reserved for script syntax and
cannot be shadowed"
+ )
self.names.add(node.id)
if isinstance(node.ctx, ast.Load):
for function in reversed(self.functions):
@@ -331,8 +360,7 @@ class PrescanCollector(ast.NodeVisitor):
# x = X.arg("x", ty)
# -------------------------------------------------
# Parameter names cannot shadow registered namespaces.
- if node.arg in self.namespaces:
- self._raise_error(node, f"Script namespace {node.arg!r} cannot be
rebound or shadowed")
+ self._check_reserved(node.arg, node)
self.names.add(node.arg)
self.generic_visit(node)
@@ -346,10 +374,40 @@ class PrescanCollector(ast.NodeVisitor):
# -------------------------------------------------
# Reserve the imported Python name.
name = node.asname or node.name.split(".")[0]
- if isinstance(self.scope, ast.FunctionDef) and name in self.namespaces:
- self._raise_error(node, f"Script namespace {name!r} cannot be
rebound or shadowed")
+ if node not in self.initial_imports:
+ self._check_reserved(name, node)
self.names.add(name)
+ def visit_TypeVar(self, node: ast.AST) -> None:
+ self._check_reserved(node.name, node)
+ self.generic_visit(node)
+
+ visit_ParamSpec = visit_TypeVar
+ visit_TypeVarTuple = visit_TypeVar
+
+ def visit_ExceptHandler(self, node: ast.ExceptHandler) -> None:
+ if node.name:
+ self._check_reserved(node.name, node)
+ self.generic_visit(node)
+
+ def visit_MatchAs(self, node: ast.MatchAs) -> None:
+ if node.name:
+ self._check_reserved(node.name, node)
+ self.generic_visit(node)
+
+ def visit_MatchStar(self, node: ast.MatchStar) -> None:
+ if node.name:
+ self._check_reserved(node.name, node)
+
+ def visit_MatchMapping(self, node: ast.MatchMapping) -> None:
+ if node.rest:
+ self._check_reserved(node.rest, node)
+ self.generic_visit(node)
+
+ def visit_NamedExpr(self, node: ast.NamedExpr) -> None:
+ self._collect_target(node.target, node.value)
+ self.visit(node.value)
+
def visit_ClassDef(self, node: ast.ClassDef) -> None:
# -------------------- Pattern --------------------
# Python source:
@@ -364,6 +422,7 @@ class PrescanCollector(ast.NodeVisitor):
# X.func_name("f")
# -------------------------------------------------
# Class host bindings and members share one lexical scope.
+ self._check_reserved(node.name, node)
self.names.add(node.name)
old, old_module = self.scope, self.module_name
self.scope, self.module_name = node, node.name
@@ -384,14 +443,14 @@ class PrescanCollector(ast.NodeVisitor):
# -------------------------------------------------
# Collect this function separately from its enclosing scope.
self._record_binding(node.name, node, "function")
- old_scope, old_builder, old_direct = self.scope, self.builder,
self.direct
- self.scope, self.direct = node, True
+ old_scope, old_builder = self.scope, self.builder
+ self.scope = node
self.functions.append(node)
self.bindings[node] = []
for decorator in node.decorator_list:
target = decorator.func if isinstance(decorator, ast.Call) else
decorator
if isinstance(target, ast.Attribute) and
protocol.DECLARATION_KIND.get(
- resolve_namespace_key(target, self.environment)
+ _match_special_func(target, self.namespaces)
) in ("function", "helper"):
namespace = resolve_namespace_value(target.value,
self.environment)
if namespace is not None:
@@ -411,7 +470,7 @@ class PrescanCollector(ast.NodeVisitor):
"int64"
if bound is None or (isinstance(bound, ast.Name) and bound.id
== "int")
else protocol.SCALAR_ANNOTATION_DTYPE.get(
- resolve_namespace_key(bound, self.environment)
+ _match_special_func(bound, self.namespaces)
)
)
self._record_binding(parameter.name, parameter, "symbol",
dtype=dtype)
@@ -429,9 +488,9 @@ class PrescanCollector(ast.NodeVisitor):
# without evaluating types.
annotation = parse_annotation(arg.annotation, self.filename) if
arg.annotation else None
arg.annotation = annotation
- constructor = resolve_namespace_key(
+ constructor = _match_special_func(
annotation.func if isinstance(annotation, ast.Call) else
annotation,
- self.environment,
+ self.namespaces,
)
self._record_binding(
arg.arg,
@@ -443,10 +502,13 @@ class PrescanCollector(ast.NodeVisitor):
else "parameter",
arg.annotation,
)
+ # Parameters, defaults and annotation-local binders are part of the
same
+ # fixed-name invariant, even when signature lowering handles them
separately.
+ self.visit(node.args)
+ for parameter in getattr(node, "type_params", ()):
+ self.visit(parameter)
for statement in node.body:
self.visit(statement)
- if isinstance(statement, ast.Return | ast.Raise):
- self.direct = False
for decorator in node.decorator_list:
self.visit(decorator)
if node.returns:
@@ -454,7 +516,7 @@ class PrescanCollector(ast.NodeVisitor):
self.visit(node.returns)
self._validate_symbols(node)
self.functions.pop()
- self.scope, self.builder, self.direct = old_scope, old_builder,
old_direct
+ self.scope, self.builder = old_scope, old_builder
visit_AsyncFunctionDef = visit_FunctionDef
@@ -500,7 +562,21 @@ class PrescanCollector(ast.NodeVisitor):
*,
kind: str = "ordinary",
) -> None:
- constructor = resolve_constructor(value, self.environment,
self.bindings[self.scope])
+ if self.functions:
+ values = [value]
+ while values:
+ item = values.pop()
+ if isinstance(item, ast.Tuple | ast.List):
+ values.extend(item.elts)
+ elif isinstance(item, ast.Starred):
+ values.append(item.value)
+ elif isinstance(item, ast.Name) and item.id in self.namespaces:
+ self._raise_error(
+ item, f"Script namespace {item.id!r} cannot be renamed
through an alias"
+ )
+ constructor = _match_special_func(
+ value.func if isinstance(value, ast.Call) else None,
self.namespaces
+ )
if isinstance(target, ast.Name):
if getattr(self.builder, "supports_mutable_declarations", True)
and (
"call" in protocol.MUTABLE_CELL_DECL.get(constructor, ())
@@ -508,11 +584,11 @@ class PrescanCollector(ast.NodeVisitor):
annotation is not None
and "annotation"
in protocol.MUTABLE_CELL_DECL.get(
- resolve_namespace_key(
+ _match_special_func(
annotation.value
if isinstance(annotation, ast.Subscript)
else annotation,
- self.environment,
+ self.namespaces,
),
(),
)
@@ -606,7 +682,6 @@ class PrescanCollector(ast.NodeVisitor):
# y = frame.var
# -------------------------------------------------
# Record matching branch outputs; lexical helpers are assembled during
rewriting.
- old_direct, self.direct = self.direct, False
self.generic_visit(node)
marker = (
isinstance(node.test, ast.Call)
@@ -643,7 +718,6 @@ class PrescanCollector(ast.NodeVisitor):
location, "IR conditional branches must end with the
same named output"
)
self.outputs[node] = then
- self.direct = old_direct
def visit_For(self, node: ast.For) -> None:
# -------------------- Pattern --------------------
@@ -656,27 +730,10 @@ class PrescanCollector(ast.NodeVisitor):
# X.emit_(body(i))
# -------------------------------------------------
# Loop binders remain assignable and are never signature declarations.
- old_direct, self.direct = self.direct, False
self._collect_target(node.target, kind="loop")
self.visit(node.iter)
for statement in node.body + node.orelse:
self.visit(statement)
- self.direct = old_direct
-
- def visit_While(self, node: ast.While) -> None:
- # -------------------- Pattern --------------------
- # Python source:
- # while condition:
- # body()
- #
- # Builder:
- # with X.While(condition):
- # X.emit_(body())
- # -------------------------------------------------
- # Declarations inside this region are conditional.
- old_direct, self.direct = self.direct, False
- self.generic_visit(node)
- self.direct = old_direct
def visit_With(self, node: ast.With) -> None:
# -------------------- Pattern --------------------
@@ -690,7 +747,6 @@ class PrescanCollector(ast.NodeVisitor):
# X.emit_(body(value))
# -------------------------------------------------
# Collect targets and explicit region outputs.
- old_direct, self.direct = self.direct, False
self.regions.append(node)
for item in node.items:
self.visit(item.context_expr)
@@ -699,7 +755,6 @@ class PrescanCollector(ast.NodeVisitor):
for statement in node.body:
self.visit(statement)
self.regions.pop()
- self.direct = old_direct
def visit_Call(self, node: ast.Call) -> None:
# -------------------- Pattern --------------------
diff --git a/python/tvm/script/parser/protocol_registry.py
b/python/tvm/script/parser/protocol_registry.py
index f0119cf662..17f16137f0 100644
--- a/python/tvm/script/parser/protocol_registry.py
+++ b/python/tvm/script/parser/protocol_registry.py
@@ -19,7 +19,7 @@
Each language variant registers its namespace through ``register_namespace``
and supplies
explicit canonical keys such as ``T.int32`` or ``R.Tensor`` at its marker
sites.
The first registered alias names that namespace. Source aliases normalize to
this
-same root; nested namespace members append their attribute names. Consumers
read
+same root; only direct root members select special syntax. Consumers read
the public dictionaries directly. Ordinary captured callables, local aliases,
instance methods and descriptors do not acquire syntax policies through
identity
or receiver inference.
@@ -42,7 +42,6 @@ _Callable = TypeVar("_Callable", bound=Callable[..., Any])
SCALAR_ANNOTATION_DTYPE: dict[str, object] = {}
MUTABLE_CELL_DECL: dict[str, frozenset[str]] = {}
RESULT_SPAN: dict[str, bool] = {}
-MODULE_DECORATOR: dict[str, bool] = {}
DECLARATION_KIND: dict[str, Literal["function", "helper"]] = {}
@@ -204,40 +203,6 @@ def result_span(namespace_path: str) ->
Callable[[_Callable], _Callable]:
return decorate
-def module_decorator(namespace_path: str) -> Callable[[_Callable], _Callable]:
- """Mark a module entry point and return the same callable.
-
- Parameters
- ----------
- namespace_path : str
- Canonical registered namespace alias and exported module-decorator
path,
- such as ``"I.ir_module"``.
-
- Returns
- -------
- Callable
- Registration decorator that returns the module entry point unchanged.
-
- Notes
- -----
- Member function decorators use this fixed namespace syntax to defer parsing
- until the complete enclosing module is available. No source-function
records,
- captured scopes or constructed module results are retained in the table.
-
- .. code:: python
-
- @module_decorator("I.ir_module")
- def ir_module(source):
- return parse(source)
- """
-
- def decorate(decorator: _Callable) -> _Callable:
- MODULE_DECORATOR[namespace_path] = True
- return decorator
-
- return decorate
-
-
def declaration_kind(
namespace_path: str, kind: Literal["function", "helper"]
) -> Callable[[_Callable], _Callable]:
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index fbe3eb141a..b62f9f6ec2 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -38,7 +38,6 @@ from __future__ import annotations
import ast
import builtins
-import copy
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from types import FunctionType
@@ -51,8 +50,7 @@ from .prescan import (
Binding,
PrescanContext,
collect_annotation_free_names,
- resolve_namespace_key,
- resolve_namespace_value,
+ collect_annotation_free_reads,
)
_Node = TypeVar("_Node", bound=ast.AST)
@@ -87,6 +85,15 @@ def _require_constexpr_arg(value: _Value, name: str) ->
_Value:
return value
+def _require_annotation_value(value: _Value, name: str) -> _Value:
+ """Keep captured values unchanged, reporting missing names only when
read."""
+ from tvm.script.ir_builder.base import MISSING
+
+ if value is MISSING:
+ raise NameError(f"name {name!r} is not defined")
+ return value
+
+
def _unwrap_optional_annotation(
annotation: Any, const_args: Mapping[str, Any] | None = None
) -> Any:
@@ -172,7 +179,7 @@ class ModuleContext:
) -> None:
# Source filename used by _raise_error and parse_annotation
diagnostics.
self.filename = filename
- # Definition/lexical lookup layers used by _resolve to identify
namespaces.
+ # Definition/lexical lookup layers used to identify construction
namespaces.
# These preserve source meanings independently of injected execution
globals.
self.environment = environment
# Generated name of the shared IR namespace used for module frames and
MISSING.
@@ -226,22 +233,20 @@ class FunctionContext:
"""
def __init__(self, current_scope: ast.AST | None, dialect_prefix: str) ->
None:
- # Active source function (or None before root lowering);
_binding_facts and
- # declaration rewriting use it to select lexical prescan facts.
+ # Active source function (or None before root lowering); declaration
rewriting
+ # uses it to select lexical prescan facts.
self.current_scope = current_scope
# Namespace identifier used by _call_dialect for this function's
builder operations.
self.dialect_prefix = dialect_prefix
- # Annotation reads start with definition captures and grow in signature
- # order. Annotation rewriting reads this one map; lexical masks
temporarily
- # replace it and restore it on exit. Ordinary body lookup does not use
it.
- self.annotation_bindings: dict[str, ast.expr] = {}
+ # One snapshot supplies original-name helper defaults at the
definition site.
+ self.definition_captures: str | None = None
class IRBuilderTranspiler(ast.NodeTransformer):
"""A single statement/expression visitor over the entry-owned AST.
The module context shares source inputs, name allocation and injected
values.
- The active function context selects lexical facts and annotation
substitutions.
+ The active function context selects lexical facts and definition captures.
Function entry and temporary expression modes save and restore their state
at
lexical boundaries, including failures. Generated AST calls construct
native
frames at execution time; this visitor does not own their IR state.
@@ -258,9 +263,8 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# Bypass syntax-to-builder lowering.
# Still traverse children and instrument source calls with spans.
self.bypass_ast_rewrite = False
- # Only annotation syntax consults the function's one substitution map;
- # ordinary body globals keep Python lookup even when names coincide.
- self.annotation_expression = False
+ # Only free reads of the active annotation need lazy missing-name
checks.
+ self.annotation_reads: set[ast.Name] = set()
# Only this expression's result is already located by its enclosing
emit_.
# Child operations retain their locations; restore the borrowed node
on exit.
self.emitted_expression: ast.expr | None = None
@@ -415,55 +419,14 @@ class IRBuilderTranspiler(ast.NodeTransformer):
self.bypass_ast_rewrite = old
@contextmanager
- def _use_aliases(self, mapping: dict[str, ast.expr]) -> Iterator[None]:
- """Restore lexical annotation substitutions even when a visitor
fails."""
- # Mask the active context's one map, retaining the enclosing map by
identity.
- old = self.function.annotation_bindings
- self.function.annotation_bindings = mapping
+ def _rewrite_annotation(self, node: ast.expr) -> Iterator[None]:
+ """Check missing values only at free reads, without renaming Python
bindings."""
+ previous_reads = self.annotation_reads
+ self.annotation_reads = set(collect_annotation_free_reads(node))
try:
yield
finally:
- # Nested annotation scopes cannot leak substitutions into their
caller.
- self.function.annotation_bindings = old
-
- @contextmanager
- def _rewrite_annotation(self) -> Iterator[None]:
- """Use definition substitutions only while rewriting annotation
syntax."""
- previous, self.annotation_expression = self.annotation_expression, True
- try:
- yield
- finally:
- self.annotation_expression = previous
-
- def _binding_facts(self) -> list[Binding]:
- """Read existing lexical bindings from the innermost source scope
outward."""
- scope = self.function.current_scope
- facts: list[Binding] = []
- while isinstance(scope, ast.FunctionDef):
- current = self.module.prescan_ctx.bindings.get(scope, ())
- hidden = {item.name for item in facts}
- facts.extend(item for item in current if item.name not in hidden)
- scope = next(
- (
- parent
- for parent, items in
self.module.prescan_ctx.bindings.items()
- if any(item.node is scope and item.kind == "function" for
item in items)
- ),
- None,
- )
- return facts
-
- def _resolve(self, node: ast.AST | None) -> str | None:
- """Recognize fixed namespace syntax, leaving all other Python values
opaque."""
- return resolve_namespace_key(node, self.module.environment,
self._binding_facts())
-
- def _is_builtin(self, node: ast.AST, value: object) -> bool:
- """Respect ordinary lexical/global shadowing of a builtin spelling."""
- return (
- isinstance(node, ast.Name)
- and not any(item.name == node.id for item in self._binding_facts())
- and self.module.environment.get(node.id, getattr(builtins,
node.id, None)) is value
- )
+ self.annotation_reads = previous_reads
def _read_constexpr_operand(self, node: ast.expr) -> ast.expr | None:
# -------------------- Pattern --------------------
@@ -474,11 +437,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# expr
# -------------------------------------------------
# The same visitor keeps Python operators and still instruments nested
source calls.
- if not isinstance(node, ast.Call) or not isinstance(node.func,
ast.Attribute):
- return None
- if node.func.attr != "constexpr" or not isinstance(node.func.value,
ast.Name):
- return None
- if node.func.value.id not in self.module.prescan_ctx.namespaces:
+ if not isinstance(node, ast.Call) or not
self._is_constexpr_annotation(node.func):
return None
if len(node.args) != 1 or node.keywords or isinstance(node.args[0],
ast.Starred):
self._raise_error(node, "constexpr expects exactly one controlling
value")
@@ -490,26 +449,18 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# value: X.Tensor((n,))
#
# Builder:
- # value_type = X.Tensor((_definition["n"],))
+ # value_type = X.Tensor((_require_annotation_value(n, "n"),))
# -------------------------------------------------
- # Definition substitutions apply only within annotation syntax.
- # A preceding body target is already a Python binding (including
symbols).
- # Consult its existing prescan location, without a second body-alias
map.
- if self.annotation_expression:
- if any(
- item.name == node.id
- and isinstance(item.node, ast.Name)
- and (item.node.lineno, item.node.col_offset) < (node.lineno,
node.col_offset)
- for item in
self.module.prescan_ctx.bindings.get(self.function.current_scope, ())
- ):
+ # Helpers bind captured values under their original names. Python
handles
+ # lambda/comprehension locals and preceding signature parameters
directly.
+ if node in self.annotation_reads:
+ if node.id in self.module.prescan_ctx.namespaces:
return node
- if (
- self.annotation_expression
- and isinstance(node.ctx, ast.Load)
- and node.id in self.function.annotation_bindings
- ):
return ast.copy_location(
- copy.deepcopy(self.function.annotation_bindings[node.id]), node
+ ast.Call(
+ self._inject(_require_annotation_value), [node,
ast.Constant(node.id)], []
+ ),
+ node,
)
return node
@@ -537,60 +488,6 @@ class IRBuilderTranspiler(ast.NodeTransformer):
else result
)
- def visit_Lambda(self, node: ast.Lambda) -> ast.Lambda:
- # -------------------- Pattern --------------------
- # Python source:
- # lambda n: n + outer
- #
- # Builder:
- # lambda n: n + captured_outer
- # -------------------------------------------------
- # Inside annotations, lambda-local binders mask definition
substitutions.
- node.args.defaults = [self.visit(value) for value in
node.args.defaults]
- node.args.kw_defaults = [
- self.visit(value) if value is not None else None for value in
node.args.kw_defaults
- ]
- arguments = [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]
- arguments += [arg for arg in (node.args.vararg, node.args.kwarg) if
arg]
- local_names = {argument.arg for argument in arguments}
- with self._use_aliases(
- {
- name: alias
- for name, alias in self.function.annotation_bindings.items()
- if name not in local_names
- }
- ):
- node.body = self.visit(node.body)
- return node
-
- def visit_ListComp(
- self, node: ast.ListComp | ast.SetComp | ast.DictComp |
ast.GeneratorExp
- ) -> ast.ListComp | ast.SetComp | ast.DictComp | ast.GeneratorExp:
- # -------------------- Pattern --------------------
- # Python source:
- # [f(n) for n in values]
- #
- # Builder:
- # [_S[i].ctx(lambda: f(n)) for n in values]
- # -------------------------------------------------
- # Comprehension binders mask annotation substitutions in Python
evaluation order.
- with self._use_aliases(dict(self.function.annotation_bindings)):
- for generator in node.generators:
- generator.iter = self.visit(generator.iter)
- for target in ast.walk(generator.target):
- if isinstance(target, ast.Name):
- self.function.annotation_bindings.pop(target.id, None)
- generator.ifs = [self.visit(value) for value in generator.ifs]
- if isinstance(node, ast.DictComp):
- node.key, node.value = self.visit(node.key),
self.visit(node.value)
- else:
- node.elt = self.visit(node.elt)
- return node
-
- visit_SetComp = visit_ListComp
- visit_DictComp = visit_ListComp
- visit_GeneratorExp = visit_ListComp
-
def visit_Constant(self, node: ast.Constant) -> ast.expr:
# -------------------- Pattern --------------------
# Python source:
@@ -662,7 +559,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
def _visit_direct_operand(self, node: ast.expr) -> ast.expr:
"""Preserve an existing payload's span without bypassing child
operations."""
if isinstance(node, ast.Name):
- return self.visit(node) if self.annotation_expression else node
+ return self.visit(node)
if isinstance(node, ast.Attribute):
node.value = self._visit_direct_operand(node.value)
return node
@@ -687,7 +584,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
if marker is not None:
with self._bypass_rewrite():
return self.visit(marker)
- constructor = self._resolve(node.func)
+ constructor = self.module.prescan_ctx._match_special_func(node.func)
global_call = (
isinstance(node.func, ast.Name) and node.func.id in
self.module.global_func_names
) or (
@@ -1210,8 +1107,12 @@ class IRBuilderTranspiler(ast.NodeTransformer):
if node.value
else ast.Attribute(ast.Name(self.module.ir_prefix, ast.Load()),
"MISSING", ast.Load())
)
- with self._rewrite_annotation():
- annotation = self.visit(parse_annotation(node.annotation,
self.module.filename))
+ annotation = parse_annotation(node.annotation, self.module.filename)
+ names = set(collect_annotation_free_names(annotation))
+ with self._rewrite_annotation(annotation):
+ helper = self._create_lambda([], self.visit(annotation))
+ self._capture_defaults(helper, names, node.annotation,
body_locals=True)
+ annotation = ast.copy_location(ast.Call(helper, [], []),
node.annotation)
return self._bind(node.target, value, node, ty=annotation,
value_span=value_span)
def visit_AugAssign(self, node: ast.AugAssign) -> ast.AugAssign |
list[ast.stmt]:
@@ -1471,9 +1372,13 @@ class IRBuilderTranspiler(ast.NodeTransformer):
return self.generic_visit(node)
if node.orelse:
self._raise_error(node, "A construction loop does not support an
else clause")
- if isinstance(node.iter, ast.Call) and
self._is_builtin(node.iter.func, range):
- # Normalize only the known builtin; a lexical range binding follows
- # the ordinary source-call path. Arguments keep that one call
scope.
+ if (
+ isinstance(node.iter, ast.Call)
+ and isinstance(node.iter.func, ast.Name)
+ and node.iter.func.id == "range"
+ ):
+ # Prescan reserves range, so symbolic bounds need no per-use scope
lookup.
+ # Arguments keep the original call's ordinary evaluation order.
iterable = self.visit_Call(
node.iter,
callee=ast.Attribute(
@@ -1664,7 +1569,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
"""Read the construction namespace from a qualified source
decorator."""
for decorator in node.decorator_list:
target = decorator.func if isinstance(decorator, ast.Call) else
decorator
- key = self._resolve(target)
+ key = self.module.prescan_ctx._match_special_func(target)
if key is None:
continue
kind = protocol.DECLARATION_KIND.get(key)
@@ -1672,13 +1577,9 @@ class IRBuilderTranspiler(ast.NodeTransformer):
return None, ast.Dict([], [])
if kind != "function":
continue
- namespace = (
- resolve_namespace_value(target.value, self.module.environment)
- if isinstance(target, ast.Attribute)
- else None
- )
- if namespace is None:
- continue
+ # A matched decorator is a direct member of a fixed root.
Construction
+ # needs that root's actual builder, not another syntax/value
resolution.
+ namespace = self.module.environment[target.value.id]
if (
self.module.root_function_kwargs is not None
and self.module.module_name is None
@@ -1728,13 +1629,8 @@ class IRBuilderTranspiler(ast.NodeTransformer):
return self._inject(namespace, "_X").id
def _read_function_annotations(
- self,
- node: ast.FunctionDef,
- parameters: list[ast.arg],
- facts: list[Binding],
- *,
- captures: str,
- ) -> tuple[list[ast.expr | None], ast.expr | None, dict[str, ast.expr]]:
+ self, node: ast.FunctionDef, parameters: list[ast.arg], facts:
list[Binding]
+ ) -> tuple[list[ast.expr | None], ast.expr | None, set[str]]:
"""Find definition-scope names needed by signatures and body
annotations."""
declared_names = {item.name for item in getattr(node, "type_params",
())}
annotations = [
@@ -1745,10 +1641,10 @@ class IRBuilderTranspiler(ast.NodeTransformer):
]
returns = parse_annotation(node.returns, self.module.filename) if
node.returns else None
annotation_names = {
- item.id
+ name
for annotation in [*annotations, returns]
if annotation is not None
- for item in collect_annotation_free_names(annotation).values()
+ for name in collect_annotation_free_names(annotation)
}
body_annotations = [
parse_annotation(item.annotation, self.module.filename)
@@ -1756,61 +1652,31 @@ class IRBuilderTranspiler(ast.NodeTransformer):
if item.annotation is not None
and item.kind not in ("parameter", "mutable_parameter", "symbol")
]
- body_annotation_names = {
- item.id
+ annotation_names.update(
+ name
for annotation in body_annotations
- for item in collect_annotation_free_names(annotation).values()
- }
- annotation_names.update(body_annotation_names)
- # Known, unshadowed signature inputs use original-name wrapper
defaults.
- # Body annotations and potentially absent/rebound names retain lazy
- # declaration-point snapshots: source body assignments must not
redirect them.
- available = _AnnotationScope(annotation_names,
self.module.definition_scope)
- rebound = {
- item.name for scope in self.module.prescan_ctx.bindings.values()
for item in scope
- }
- direct = available.keys() - body_annotation_names - rebound
- bindings = {
- name: (
- ast.Name(name, ast.Load())
- if name in direct
- or (
- name in self.module.prescan_ctx.namespaces
- and name in self.module.exec_globals
- and self.module.environment.get(name) is
self.module.exec_globals[name]
- )
- else ast.Subscript(ast.Name(captures, ast.Load()),
ast.Constant(name), ast.Load())
- )
- for name in sorted(annotation_names - declared_names)
- }
- return annotations, returns, bindings
+ for name in collect_annotation_free_names(annotation)
+ )
+ return annotations, returns, annotation_names - declared_names
def _create_definition_bindings(
self,
node: ast.FunctionDef,
- bindings: dict[str, ast.expr],
+ names: set[str],
*,
captures: str,
) -> list[ast.stmt]:
"""Snapshot only lexical annotation/constexpr names at their
definition site."""
- names = {
- name for name, binding in bindings.items() if isinstance(binding,
ast.Subscript)
- } | {
+ names = names | {
parameter.arg
for parameter in [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]
if self._is_constexpr_annotation(parameter.annotation)
}
- if not names:
- return []
# Inject builtin operations so same-named source bindings cannot
replace them.
# Acquire each scope once; the snapshot retains selected values, never
frames.
scopes: list[ast.expr] = [ast.Call(self._inject(locals), [], [])]
scopes.append(ast.Call(self._inject(globals), [], []))
- defaults = {
- name: getattr(builtins, name)
- for name in names
- if name in bindings and hasattr(builtins, name)
- }
+ defaults = {name: getattr(builtins, name) for name in names if
hasattr(builtins, name)}
if defaults:
scopes.append(self._inject(defaults))
captured = ast.Call(
@@ -1820,6 +1686,47 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
return [self._assign(captures, captured, node)]
+ def _capture_defaults(
+ self,
+ helper: ast.FunctionDef | ast.Lambda,
+ names: set[str],
+ node: ast.AST,
+ *,
+ body_locals: bool = False,
+ ) -> None:
+ """Give a helper original-name captures; Python owns all lexical
shadowing."""
+ for name in sorted(names):
+ value = self._call(
+ self.function.definition_captures,
+ "get",
+ [
+ ast.Constant(name),
+ ast.Attribute(
+ ast.Name(self.module.ir_prefix, ast.Load()),
"MISSING", ast.Load()
+ ),
+ ],
+ node,
+ )
+ if body_locals and any(
+ item.name == name
+ and (item.node.lineno, item.node.col_offset) < (node.lineno,
node.col_offset)
+ for item in
self.module.prescan_ctx.bindings.get(self.function.current_scope, ())
+ ):
+ # Keep a real name read so Python captures enclosing helper
cells.
+ # A prior conditional binding may not have executed: snapshot
its
+ # absence without reading it until the annotation selects that
name.
+ value = ast.IfExp(
+ ast.Compare(
+ ast.Constant(name), [ast.In()],
[ast.Call(self._inject(locals), [], [])]
+ ),
+ ast.Name(name, ast.Load()),
+ ast.Attribute(
+ ast.Name(self.module.ir_prefix, ast.Load()),
"MISSING", ast.Load()
+ ),
+ )
+ helper.args.kwonlyargs.append(ast.arg(name))
+ helper.args.kw_defaults.append(value)
+
def _create_const_args(self, node: ast.FunctionDef, *, const_args: str) ->
list[ast.stmt]:
"""Read root JIT inputs; nested functions keep ordinary runtime
parameters."""
from . import jit_support
@@ -1832,12 +1739,9 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
return [self._assign(const_args, const_args_expr, node)]
- def _create_symbol_declarations(
- self, node: ast.FunctionDef
- ) -> tuple[list[ast.stmt], dict[str, ast.expr]]:
+ def _create_symbol_declarations(self, node: ast.FunctionDef) ->
list[ast.stmt]:
"""Bind explicit signature type parameters with their declared
dtypes."""
declaration: list[ast.stmt] = []
- symbol_aliases: dict[str, ast.expr] = {}
# -------------------- Pattern --------------------
# Python source:
# def f[n]():
@@ -1851,19 +1755,22 @@ class IRBuilderTranspiler(ast.NodeTransformer):
self._raise_error(parameter, "Only scalar type parameters are
supported")
bound = getattr(parameter, "bound", None)
dtype = self.module.prescan_ctx.sites[parameter].dtype
- if bound is not None and not self._is_builtin(bound, int):
- if protocol.SCALAR_ANNOTATION_DTYPE.get(self._resolve(bound))
is None:
+ if bound is not None and not (isinstance(bound, ast.Name) and
bound.id == "int"):
+ if (
+ protocol.SCALAR_ANNOTATION_DTYPE.get(
+ self.module.prescan_ctx._match_special_func(bound)
+ )
+ is None
+ ):
self._raise_error(
parameter,
"A symbolic type parameter bound must be int or a
registered scalar dtype",
)
if getattr(parameter, "default_value", None) is not None:
self._raise_error(parameter, "A symbolic type parameter cannot
have a default")
- alias = self.module.make_fresh_name("_symbol")
- symbol_aliases[parameter.name] = ast.Name(alias, ast.Load())
declaration.append(
self._assign(
- alias,
+ parameter.name,
self._call_dialect(
"resolve_type_var_",
[ast.Constant(parameter.name)],
@@ -1873,7 +1780,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
parameter,
)
)
- return declaration, symbol_aliases
+ return declaration
@staticmethod
def _select_specialized_value(
@@ -1898,8 +1805,9 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
def _is_constexpr_annotation(self, annotation: ast.expr | None) -> bool:
- """Recognize the registered marker identity, including ordinary
aliases."""
- return resolve_namespace_value(annotation, self.module.environment) is
protocol.constexpr
+ """Recognize constexpr syntax through a canonical registered
root/member key."""
+ key = self.module.prescan_ctx._match_special_func(annotation)
+ return key is not None and key.endswith(".constexpr")
def _rewrite_parameters(
self,
@@ -1910,7 +1818,6 @@ class IRBuilderTranspiler(ast.NodeTransformer):
const_args: str | None,
) -> tuple[list[ast.stmt], dict[str, str]]:
"""Declare signature parameters, making constexpr values available
first."""
-
declaration: list[ast.stmt] = []
constexpr_aliases: dict[str, str] = {}
constexpr_params: list[tuple[ast.arg, ast.expr | None, bool]] = []
@@ -1921,15 +1828,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
is_constexpr = self._is_constexpr_annotation(annotation)
group = constexpr_params if is_constexpr else other_params
group.append((parameter, annotation, is_constexpr))
- constexpr_reads = {
- parameter.arg
- for parameter, _, _ in constexpr_params
- if parameter.arg in self.function.annotation_bindings
- }
- constexpr_scope = self.module.make_fresh_name("_constexpr") if
constexpr_reads else None
- for index, (parameter, annotation, is_constexpr) in enumerate(
- [*constexpr_params, *other_params]
- ):
+ for parameter, annotation, is_constexpr in [*constexpr_params,
*other_params]:
if annotation is None:
self._raise_error(parameter, f"Parameter {parameter.arg!r}
requires an annotation")
name = parameter.arg
@@ -1960,7 +1859,8 @@ class IRBuilderTranspiler(ast.NodeTransformer):
[],
)
else:
- translated = self.visit(annotation)
+ with self._rewrite_annotation(annotation):
+ translated = self.visit(annotation)
if const_args is not None:
# Optional annotations are unwrapped only for
specialization;
# ordinary annotation validation belongs to the language
variant arg.
@@ -1981,30 +1881,8 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
)
declaration.append(self._assign(alias, value, parameter))
- self.function.annotation_bindings[name] = (
- ast.Subscript(ast.Name(constexpr_scope, ast.Load()),
ast.Constant(name), ast.Load())
- if is_constexpr and name in constexpr_reads
- else ast.Name(alias, ast.Load())
- )
- if constexpr_scope is not None and index + 1 ==
len(constexpr_params):
- # Selected host values may be MISSING. Preserve lazy
missing-name
- # errors without changing definition captures or Python body
locals.
- names = sorted(constexpr_reads)
- values = ast.Dict(
- [ast.Constant(name) for name in names],
- [ast.Name(constexpr_aliases[name], ast.Load()) for name in
names],
- )
- declaration.append(
- self._assign(
- constexpr_scope,
- ast.Call(
- self._inject(_AnnotationScope),
- [ast.Tuple([ast.Constant(name) for name in names],
ast.Load()), values],
- [],
- ),
- parameter,
- )
- )
+ if alias != name:
+ declaration.append(self._assign(name, ast.Name(alias,
ast.Load()), parameter))
return declaration, constexpr_aliases
def _create_function_frame(
@@ -2145,66 +2023,38 @@ class IRBuilderTranspiler(ast.NodeTransformer):
self,
node: ast.FunctionDef,
frame: str,
- body_name: str,
definition: ast.FunctionDef,
frame_declaration: list[ast.stmt],
+ capture_names: set[str],
+ constexpr_aliases: dict[str, str],
) -> list[ast.stmt]:
- """Give each signature a lexical scope and retain its body closure."""
- # Each signature needs its own Python parameter scope. Keep
- # definition captures outside, at the original declaration
- # point, and retain the body helper's signature closures.
- # -------------------- Pattern --------------------
- # Python source:
- # def f(x: annotation):
- # body(x)
- #
- # Builder:
- # def declare(frame):
- # with frame:
- # x = X.arg("x", annotation)
- # def build():
- # x, = frame.params
- # body(x)
- # return frame, build
- # frame, build = declare(X.function_(decl=True))
- # f = frame.global_var
- # with frame:
- # build()
- # -------------------------------------------------
+ """Keep signature snapshots separate from the body's enclosing Python
scope."""
+ # The declaration helper returns only the frame and selected host
values.
+ # Runtime parameters and explicit symbols are recovered from the frame
by
+ # the sibling body helper, whose closures still observe source class
setup.
declare_name = self.module.make_fresh_name("_declare")
signature = frame_declaration[0]
- # Decorator options belong to the outer definition scope;
- # a same-named parameter must not hide them in the helper.
+ # Decorator options evaluate outside the signature's same-named
parameters.
constructor = signature.items[0].context_expr
signature.items[0].context_expr = ast.copy_location(ast.Name(frame,
ast.Load()), node)
+ outputs = [frame, *constexpr_aliases.values()]
returned = ast.copy_location(
- ast.Return(
- ast.Tuple(
- [ast.Name(frame, ast.Load()), ast.Name(body_name,
ast.Load())],
- ast.Load(),
- )
- ),
+ ast.Return(ast.Tuple([ast.Name(name, ast.Load()) for name in
outputs], ast.Load())),
node,
)
- declare = self._create_definition(declare_name, [signature,
definition, returned], node)
+ declare = self._create_definition(declare_name, [signature, returned],
node)
declare.args.args.append(ast.arg(frame))
+ self._capture_defaults(declare, capture_names, node)
return [
declare,
ast.copy_location(
ast.Assign(
- [
- ast.Tuple(
- [
- ast.Name(frame, ast.Store()),
- ast.Name(body_name, ast.Store()),
- ],
- ast.Store(),
- )
- ],
+ [ast.Tuple([ast.Name(name, ast.Store()) for name in
outputs], ast.Store())],
ast.Call(ast.Name(declare_name, ast.Load()),
[constructor], []),
),
node,
),
+ definition,
frame_declaration[1],
]
@@ -2213,10 +2063,9 @@ class IRBuilderTranspiler(ast.NodeTransformer):
) -> tuple[list[ast.stmt], str, ast.With]:
"""Declare a native frame and emit a lexical body helper inside its
scope.
- Definition captures retain outer annotation values. Signature bindings
add
- declared symbols and each preceding parameter; constexpr aliases retain
- compile-time values for the body. One substitution map serves
annotation syntax;
- ordinary body reads retain Python globals/closures. None owns native
IR.
+ Original-name helper defaults retain definition values. Subsequent
Python
+ assignments expose each signature parameter to later annotations. Body
+ helpers preserve source globals/closures, independently of annotation
captures.
"""
kind, options = self.read_function_metadata(node)
builder = self._function_namespace(node, kind)
@@ -2226,111 +2075,133 @@ class IRBuilderTranspiler(ast.NodeTransformer):
old = self.function
self.function = FunctionContext(node, builder)
try:
- with self._use_aliases({}):
- parameters = [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]
- if node.args.vararg or node.args.kwarg:
- self._raise_error(node, "IR signatures require ordinary
named parameters")
- facts = self.module.prescan_ctx.bindings.get(node, [])
- captures = self.module.make_fresh_name("_definition")
- annotations, returns, definition_aliases =
self._read_function_annotations(
- node, parameters, facts, captures=captures
- )
- self.function.annotation_bindings = definition_aliases
- statements = self._create_definition_bindings(
- node, definition_aliases, captures=captures
+ parameters = [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]
+ if node.args.vararg or node.args.kwarg:
+ self._raise_error(node, "IR signatures require ordinary named
parameters")
+ facts = self.module.prescan_ctx.bindings.get(node, [])
+ annotations, returns, capture_names =
self._read_function_annotations(
+ node, parameters, facts
+ )
+ captures = self.module.make_fresh_name("_definition")
+ self.function.definition_captures = captures
+ statements = self._create_definition_bindings(node, capture_names,
captures=captures)
+ const_args = None
+ if self.module.enable_jit_map and not local_function:
+ const_args = self.module.make_fresh_name("_const_args")
+ statements.extend(self._create_const_args(node,
const_args=const_args))
+ declaration: list[ast.stmt] = [
+ ast.copy_location(
+ ast.Expr(self._call(builder, "func_name",
[ast.Constant(node.name)], node)),
+ node,
)
- const_args = None
- if self.module.enable_jit_map and not local_function:
- const_args = self.module.make_fresh_name("_const_args")
- statements.extend(self._create_const_args(node,
const_args=const_args))
- declaration: list[ast.stmt] = [
+ ]
+ declaration.extend(self._create_symbol_declarations(node))
+ arguments, constexpr_aliases = self._rewrite_parameters(
+ parameters, annotations, captures=captures,
const_args=const_args
+ )
+ declaration.extend(arguments)
+ if returns is not None:
+ with self._rewrite_annotation(returns):
+ translated_return = self.visit(returns)
+ declaration.append(
ast.copy_location(
- ast.Expr(self._call(builder, "func_name",
[ast.Constant(node.name)], node)),
- node,
- )
- ]
- symbols, symbol_aliases =
self._create_symbol_declarations(node)
- declaration.extend(symbols)
- with (
- self._use_aliases({**definition_aliases,
**symbol_aliases}),
- self._rewrite_annotation(),
- ):
- arguments, constexpr_aliases = self._rewrite_parameters(
- parameters, annotations, captures=captures,
const_args=const_args
- )
- declaration.extend(arguments)
- if returns is not None:
- declaration.append(
- ast.copy_location(
- ast.Expr(
- self._call(
- builder,
- "func_ret_type",
- [self._create_lambda([],
self.visit(returns))],
- returns,
- )
- ),
+ ast.Expr(
+ self._call(
+ builder,
+ "func_ret_type",
+ [self._create_lambda([], translated_return)],
returns,
)
- )
-
definition_aliases.update(self.function.annotation_bindings)
- with self._bypass_rewrite():
- options = self.visit(options)
- frame_declaration, body_entry = self._create_function_frame(
- node,
- options,
- declaration,
- frame=frame,
- local_function=local_function,
- split_declare=split_declare,
+ ),
+ returns,
+ )
+ )
+ with self._bypass_rewrite():
+ options = self.visit(options)
+ frame_declaration, body_entry = self._create_function_frame(
+ node,
+ options,
+ declaration,
+ frame=frame,
+ local_function=local_function,
+ split_declare=split_declare,
+ )
+ body = self._create_body_parameters(
+ node, parameters, constexpr_aliases, frame=frame,
const_args=const_args
+ )
+ body.extend(self.transform_statements(node.body))
+ definition = self._create_definition(body_name, body, node)
+ if not local_function:
+ protected = (
+ {item.name for item in getattr(node, "type_params", ())}
+ | {node.name}
+ | set(self.module.global_func_names)
+ | self.module.prescan_ctx.namespaces.keys()
)
- body = self._create_body_parameters(
- node, parameters, constexpr_aliases, frame=frame,
const_args=const_args
+ if self.module.module_name:
+ protected.add(self.module.module_name)
+ self.module.generated_builders.append(
+ GeneratedBuilder(definition, node.name, protected)
)
- body.extend(self.transform_statements(node.body))
- definition = self._create_definition(body_name, body, node)
- if not local_function:
- protected = (
- {item.name for item in getattr(node, "type_params",
())}
- | {node.name}
- | set(self.module.global_func_names)
- | self.module.prescan_ctx.namespaces
- )
- if self.module.module_name:
- protected.add(self.module.module_name)
- self.module.generated_builders.append(
- GeneratedBuilder(definition, node.name, protected)
+ # -------------------- Pattern --------------------
+ # Python source:
+ # def f():
+ # body()
+ #
+ # Builder:
+ # with frame:
+ # def build():
+ # X.emit_(body())
+ # build()
+ # -------------------------------------------------
+ # The helper owns lexical scope; the surrounding with owns frame
+ # entry/exit and error unwinding.
+ invocation = ast.copy_location(
+ ast.Expr(ast.Call(ast.Name(body_name, ast.Load()), [], [])),
node
+ )
+ if split_declare:
+ statements.extend(
+ self._create_split_declaration(
+ node, frame, definition, frame_declaration,
capture_names, constexpr_aliases
)
- # -------------------- Pattern --------------------
- # Python source:
- # def f():
- # body()
- #
- # Builder:
- # with frame:
- # def build():
- # X.emit_(body())
- # build()
- # -------------------------------------------------
- # The helper owns lexical scope; the surrounding with owns
frame
- # entry/exit and error unwinding.
- invocation = ast.copy_location(
- ast.Expr(ast.Call(ast.Name(body_name, ast.Load()), [],
[])), node
)
- if split_declare:
- statements.extend(
- self._create_split_declaration(
- node, frame, body_name, definition,
frame_declaration
+ else:
+ outputs = list(constexpr_aliases.values())
+ if outputs:
+ declaration.append(
+ ast.copy_location(
+ ast.Return(
+ ast.Tuple(
+ [ast.Name(name, ast.Load()) for name in
outputs], ast.Load()
+ )
+ ),
+ node,
)
)
- resumed = ast.copy_location(
- ast.With(
- [body_entry],
- [invocation] if split_declare else [*declaration,
definition, invocation],
- ),
+ signature = self._create_definition(
+ self.module.make_fresh_name("_declare"), declaration, node
+ )
+ self._capture_defaults(signature, capture_names, node)
+ call = ast.Call(ast.Name(signature.name, ast.Load()), [], [])
+ signature_call = ast.copy_location(
+ ast.Assign(
+ [ast.Tuple([ast.Name(name, ast.Store()) for name in
outputs], ast.Store())],
+ call,
+ )
+ if outputs
+ else ast.Expr(call),
node,
)
- return statements, frame, resumed
+ resumed = ast.copy_location(
+ ast.With(
+ [body_entry],
+ [invocation]
+ if split_declare
+ else [signature, signature_call, definition, invocation],
+ ),
+ node,
+ )
+ return statements, frame, resumed
finally:
# Restore the enclosing function before returning fragments or
propagating
# an error; partially rewritten nested functions cannot leave
active state.
diff --git a/python/tvm/tirx/script/jit.py b/python/tvm/tirx/script/jit.py
index 7ca5f43e68..544b1c89e9 100644
--- a/python/tvm/tirx/script/jit.py
+++ b/python/tvm/tirx/script/jit.py
@@ -236,7 +236,7 @@ class TIRJit:
raw_anns = getattr(func, "__annotations__", {}) or {}
annotation_scope = {**func.__globals__, **self._closure_vars,
**self._definition_scope}
- from tvm.script.parser.prescan import resolve_namespace_value
+ from tvm.script.parser.prescan import resolve_namespace_key,
resolve_namespace_value
sig = inspect.signature(func)
constexpr_names: set[str] = set()
@@ -250,11 +250,21 @@ class TIRJit:
node = ast.parse(ann, mode="eval").body
if isinstance(node, ast.Constant) and
isinstance(node.value, str):
node = ast.parse(node.value, mode="eval").body
- ann = resolve_namespace_value(
- node.func if isinstance(node, ast.Call) else node,
annotation_scope
+ key = (
+ resolve_namespace_key(node, annotation_scope)
+ if isinstance(node, ast.Attribute) and
isinstance(node.value, ast.Name)
+ else None
)
- if isinstance(node, ast.Call) and ann is not
OptionalAnnotation:
- ann = None
+ if key is not None and key.endswith(".constexpr"):
+ ann = protocol.constexpr
+ else:
+ ann = resolve_namespace_value(
+ node.func if isinstance(node, ast.Call) else node,
annotation_scope
+ )
+ if ann is protocol.constexpr or (
+ isinstance(node, ast.Call) and ann is not
OptionalAnnotation
+ ):
+ ann = None
except SyntaxError:
ann = None
if ann is protocol.constexpr:
diff --git a/tests/python/script/test_basic_usage.py
b/tests/python/script/test_basic_usage.py
index 6d05987a2c..5068c9aa8b 100644
--- a/tests/python/script/test_basic_usage.py
+++ b/tests/python/script/test_basic_usage.py
@@ -407,8 +407,8 @@ def
test_void_branch_statements_need_no_synthetic_named_output(language):
assert [value for kind, value in main.body] == [None, None, 9]
-def test_lexical_range_binding_calls_the_custom_iterator_once(language):
- # A lexical range alias must call the captured iterator once instead of
the builtin range.
+def test_ordinary_iterator_binding_calls_the_custom_iterator_once(language):
+ # An ordinary iterator alias calls its captured implementation exactly
once.
M = language.M
calls = []
@@ -418,8 +418,8 @@ def
test_lexical_range_binding_calls_the_custom_iterator_once(language):
@M.function
def main():
- range = custom_range
- for i in range(4):
+ iterator = custom_range
+ for i in iterator(4):
M.record(i)
assert calls == [4]
diff --git a/tests/python/script/test_error_handling.py
b/tests/python/script/test_error_handling.py
index 61d7f29b8d..cb7764d68b 100644
--- a/tests/python/script/test_error_handling.py
+++ b/tests/python/script/test_error_handling.py
@@ -28,6 +28,7 @@ import pytest
from tvm import ir
from tvm.ir import prim
from tvm.script import tirx as T
+from tvm.script.parser import entry
def
test_nested_function_failure_preserves_error_and_recovers(spanned_language):
@@ -190,6 +191,99 @@ def
test_namespace_rebinding_reports_source_location(language):
assert not language.functions
[email protected](
+ "name, statement",
+ [
+ ("range", "range = 1"),
+ ("int", "int: object = 1"),
+ ("M", "M += 1"),
+ ("range", "(range, other) = (1, 2)"),
+ ("int", "for int in M.grid(2):\n pass"),
+ ("M", "with M.grid(2) as M:\n pass"),
+ ("range", "def range():\n pass"),
+ ("int", "class int:\n pass"),
+ ("M", "import math as M"),
+ ("range", "from math import floor as range"),
+ ("int", "values = [0 for int in (1,)]"),
+ ("M", "value = lambda M: 0"),
+ ("range", "value = (range := 1)"),
+ ("int", "del int"),
+ ("M", "try:\n pass\nexcept ValueError as M:\n pass"),
+ ("range", "match 1:\n case range:\n pass"),
+ ("int", "match [1]:\n case [*int]:\n pass"),
+ ("M", "match {}:\n case {**M}:\n pass"),
+ ],
+)
+def test_reserved_bindings_fail_before_construction(language, name, statement):
+ source = "@M.function\ndef main():\n effect()\n" + "\n".join(
+ " " + line for line in statement.splitlines()
+ )
+ effects = []
+ with pytest.raises(SyntaxError, match=repr(name)) as caught:
+ entry.parse(
+ source, {"M": language.M, "effect": lambda: effects.append(1)},
filename="reserved.py"
+ )
+ error = caught.value
+ assert error.filename == "reserved.py"
+ assert error.lineno >= 4
+ assert name in source.splitlines()[error.lineno - 1]
+ assert error.offset > 0
+ assert not language.functions
+ assert not effects
+
+
[email protected](
+ "name, parameter", [("M", "M: object"), ("range", "*range"), ("int",
"**int")]
+)
+def test_reserved_parameters_fail_before_construction(language, name,
parameter):
+ source = "@M.function\ndef main(" + parameter.format(name=name) + "):\n
pass"
+ with pytest.raises(SyntaxError, match=repr(name)) as caught:
+ entry.parse(source, {"M": language.M}, filename="parameter.py")
+ assert (caught.value.filename, caught.value.lineno) == ("parameter.py", 2)
+ assert source.splitlines()[1][caught.value.offset - 1 :].startswith(name)
+ assert not language.functions
+
+
[email protected]("name", ["range", "int"])
+def test_reserved_captured_builtin_reports_its_use(language, name):
+ source = f"@M.function\ndef main():\n {name}(2)"
+ with pytest.raises(SyntaxError, match=repr(name)) as caught:
+ entry.parse(source, {"M": language.M, name: lambda value: value},
filename="capture.py")
+ assert (caught.value.filename, caught.value.lineno, caught.value.offset)
== ("capture.py", 3, 5)
+ assert not language.functions
+
+
[email protected]("binding", ["alias, *other = M, 1, 2", "value =
(alias := M)"])
+def test_local_namespace_alias_is_rejected(language, binding):
+ source = "@M.function\ndef main():\n " + binding
+ with pytest.raises(SyntaxError, match="namespace 'M'.*alias") as caught:
+ entry.parse(source, {"M": language.M}, filename="alias.py")
+ assert (caught.value.filename, caught.value.lineno) == ("alias.py", 3)
+ assert source.splitlines()[2][caught.value.offset - 1] == "M"
+ assert not language.functions
+
+
[email protected](
+ "prefix",
+ ["class Module:", "if True:"],
+)
+def test_host_import_cannot_replace_reserved_names(language, prefix):
+ effects = []
+ source = prefix + "\n from math import floor as range\n effect()\n"
+ source += (
+ " @M.function\n def main():\n pass"
+ if prefix.startswith("class")
+ else "@M.function\ndef main():\n pass"
+ )
+ with pytest.raises(SyntaxError, match="'range'.*reserved") as caught:
+ entry.parse(
+ source, {"M": language.M, "effect": lambda: effects.append(1)},
filename="import.py"
+ )
+ assert (caught.value.filename, caught.value.lineno) == ("import.py", 2)
+ assert not effects
+ assert not language.functions
+
+
def test_undefined_name_reports_the_original_source(language):
# A missing body value must remain a Python NameError at the user's
identifier.
M = language.M
diff --git a/tests/python/script/test_jit_specialization.py
b/tests/python/script/test_jit_specialization.py
index b5b1966f70..f2fce055e7 100644
--- a/tests/python/script/test_jit_specialization.py
+++ b/tests/python/script/test_jit_specialization.py
@@ -24,6 +24,7 @@ from __future__ import annotations
import gc
import weakref
+from types import SimpleNamespace
import pytest
@@ -257,3 +258,54 @@ def
test_jit_does_not_retain_ordinary_decorator_owners(jit_language):
function, reference = make()
assert reference() is None
assert function.specialize().body == []
+
+
[email protected]("postponed", [False, True])
+def test_jit_initial_alias_preserves_annotation_forms(tmp_path, postponed):
+ source = """@Script.jit(private=True)
+def kernel(output: Script.Buffer((1,), "int32"), *, value: Script.constexpr,
+ optional: Script.Optional(Script.Buffer((1,), "int32"))):
+ output[0] = value
+"""
+ if postponed:
+ source = "from __future__ import annotations\n" + source.replace(
+ "value: Script.constexpr", "value: 'Script.constexpr'"
+ )
+ path = tmp_path / "aliased_jit.py"
+ path.write_text(source)
+ namespace = {"Script": T}
+ exec(compile(source, str(path), "exec", dont_inherit=True), namespace)
+ kernel = namespace["kernel"]
+ assert kernel.constexpr_names == {"value"}
+ assert kernel.optional_names == {"optional"}
+ result = kernel.specialize(value=3, optional=None)
+ assert len(result.params) == 1
+ assert int(result.body.value) == 3
+
+
+def test_postponed_jit_uses_registered_constexpr_syntax(jit_language,
monkeypatch):
+ M = jit_language.M
+ Alias = M
+ # Syntax classification must not inspect the value behind a registered
marker name.
+ monkeypatch.setattr(M, "constexpr", object())
+ Ordinary = SimpleNamespace(constexpr=I.constexpr, Optional=T.Optional)
+ M.marker_alias = I.constexpr
+
+ def unavailable():
+ pytest.fail("Unselected annotations must remain unevaluated")
+
+ @M.jit
+ def specialized(value: Alias.constexpr):
+ M.record(value)
+
+ assert specialized.specialize(value=4).body == [("emit", 4)]
+
+ @M.jit
+ def ordinary(
+ value: Ordinary.constexpr,
+ optional: Ordinary.Optional(unavailable()),
+ alias: M.marker_alias,
+ ):
+ M.record(value)
+
+ assert ordinary.constexpr_names == ordinary.optional_names == set()
diff --git a/tests/python/script/test_meta_programming.py
b/tests/python/script/test_meta_programming.py
index 513241e695..868bfae8c5 100644
--- a/tests/python/script/test_meta_programming.py
+++ b/tests/python/script/test_meta_programming.py
@@ -28,10 +28,13 @@ import gc
import weakref
from contextlib import contextmanager
from types import SimpleNamespace
+from typing import TypeVar
import pytest
+from tvm import ir
from tvm.script import ir as I
+from tvm.script import relax as R
from tvm.script.parser import entry, protocol_registry
EXTENT = 11
@@ -101,6 +104,27 @@ def
test_module_parameter_does_not_replace_another_functions_capture(language):
assert function.body == [("emit", function.params[0])]
+def test_body_local_shadows_capture_for_later_annotation(language):
+ M = language.M
+ extent = 7
+ seen = []
+
+ def annotation(shape):
+ seen.append(shape)
+ return M.Tensor(shape)
+
+ @M.function
+ def main(x: M.Tensor((extent,))):
+ extent = 3
+ value: annotation((extent,)) = x
+ M.record(value)
+
+ assert main.params[0].args[0].args[0] == (7,)
+ assert seen == [(3,)]
+ assert extent == 7
+ assert main.body == [("emit", main.params[0])]
+
+
def
test_unrelated_same_file_caller_does_not_supply_annotation_locals(language):
# Before: unrelated caller locals shadow both annotation and body global
names.
# Expected builder: X.arg receives global 11; X.record receives global 1.
@@ -240,6 +264,19 @@ def
test_local_annotation_preserves_lambda_and_comprehension_bindings(language):
assert result.args == (function.params[0], function.params[0])
+def test_return_annotation_keeps_local_symbols_and_unused_captures():
+ n = ir.Var("n", "int64")
+ unused = TypeVar("unused", bound=int)
+
+ @R.function
+ def main(x: R.Tensor((n,), "float32")) -> R.Tensor(
+ ((lambda local: local if I.constexpr(True) else unused)(n),), "float32"
+ ):
+ return x
+
+ assert main.ret_ty.shape[0].same_as(n)
+
+
def test_class_annotation_scope_keeps_distinct_method_closure(language):
# Before: an enclosing extent=7 and class extent=3 share a method spelling.
# Expected builder program: the annotation reads class 3, the body closes
over 7.
@@ -375,6 +412,37 @@ def
test_nonlocal_declaration_preserves_captured_values(language):
assert dtype == "float32" and value == 3
[email protected]("hygienic", [True, False])
+def test_macro_local_annotation_captures_definition_and_argument_names(
+ language, hygienic, monkeypatch
+):
+ M = language.M
+ M.macro = protocol_registry.declaration_kind("M.macro",
"helper")(entry.make_macro_decorator(M))
+ MACRO_VALUE = 7 # noqa: F841 — captured only by the postponed local
annotation.
+ observed = []
+
+ def annotation(shape):
+ observed.append(shape)
+ return M.Tensor(shape, "float32")
+
+ monkeypatch.setitem(globals(), "annotation", annotation)
+
+ @M.macro(hygienic=hygienic)
+ def typed_local(size):
+ value: annotation((MACRO_VALUE, size)) = M.value(1, 2)
+ M.record(value)
+
+ monkeypatch.setitem(globals(), "MACRO_VALUE", 9)
+
+ @M.function
+ def function():
+ typed_local(3)
+
+ assert observed == [(7 if hygienic else 9, 3)]
+ assert function.body[0][0] == "emit"
+ assert function.body[0][1].op == "value"
+
+
def test_macro_capture_policy_remains_explicit(language, monkeypatch):
# A changed global must distinguish captured macro scope from caller
scope: values 2 then 1.
M = language.M
diff --git a/tests/python/script/test_parser_entry.py
b/tests/python/script/test_parser_entry.py
index 4679d25fa9..f2e1f3189c 100644
--- a/tests/python/script/test_parser_entry.py
+++ b/tests/python/script/test_parser_entry.py
@@ -23,7 +23,6 @@ from __future__ import annotations
# Keep the single generated-program golden exactly as ast.unparse emits it.
import ast
-import sys
from textwrap import dedent
from tvm.script.parser import entry
@@ -74,9 +73,14 @@ def
test_callable_entry_emits_the_expected_builder_program(language, monkeypatch
expected = dedent(
"""
with _I0.IRBuilder() as _builder0:
+ _definition0 = _host2(('M',), _host0(), _host1())
with M.function_() as _fn0:
- M.func_name('identity')
- x = M.arg('x', M.Tensor((4,)))
+
+ def _declare0(*, M=_definition0.get('M', _I0.MISSING)):
+ M.func_name('identity')
+ x = M.arg('x', M.Tensor((4,)))
+
+ _declare0()
def _build0():
x, = _fn0.params
@@ -87,10 +91,8 @@ def
test_callable_entry_emits_the_expected_builder_program(language, monkeypatch
M.check_well_formed_(_result0)
"""
).strip()
- if sys.version_info[:2] == (3, 10):
- # Python 3.10's ast.unparse parenthesizes this tuple assignment target.
- expected = expected.replace("x, = _fn0.params", "(x,) = _fn0.params")
- assert programs == [expected]
+ assert len(programs) == 1
+ assert ast.dump(ast.parse(programs[0])) == ast.dump(ast.parse(expected))
assert result.params[0].args[0].args[0] == (4,)
assert result.body == [("emit", result.params[0])]
diff --git a/tests/python/script/test_special_parser_protocol.py
b/tests/python/script/test_special_parser_protocol.py
index 6cb7c63bf8..2cdf1fdcb8 100644
--- a/tests/python/script/test_special_parser_protocol.py
+++ b/tests/python/script/test_special_parser_protocol.py
@@ -20,11 +20,13 @@ A failed constructor must not change the next function's
interpretation of liter
import inspect
import traceback
+from types import SimpleNamespace
import pytest
from tvm.script import tirx as T
from tvm.script.ir_builder import resolve_global_info_args
+from tvm.script.parser import entry
def test_constructor_policy_survives_a_failed_definition(language):
@@ -121,3 +123,35 @@ def
test_tirx_rejects_global_info_at_the_call_site(monkeypatch):
location,
column + len('T.global_annotation(device="cuda:0")'),
)
+
+
+def test_ordinary_calls_preserve_literal_arguments(language):
+ M = language.M
+ Alias = M
+ seen = []
+
+ M.special = lambda value: seen.append(value)
+ M.nested = SimpleNamespace(special=M.special)
+ special = M.special
+
+ @M.function
+ def main():
+ Alias.special("n")
+ M.nested.special("n")
+ special("n")
+ M.record(int("3"))
+
+ assert seen == ["n", "n", "n"]
+ assert main.body[-1] == ("emit", 3)
+
+
+def test_initial_import_alias_and_symbolic_range():
+ source = """from tvm.script import tirx as Script
[email protected]_func
+def main(n: Script.int32):
+ for i in range(n):
+ Script.evaluate(i)
+"""
+ function = entry.parse(source)
+ assert function.body.extent.same_as(function.params[0])
+ assert function.body.body.value.same_as(function.body.loop_var)
