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 615d506e0e [REFACTOR][Script] Simplify annotation scopes and signature
variables (#20456)
615d506e0e is described below
commit 615d506e0eac781284f472bbf3c185235dc1c76e
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Sep 25 15:18:57 2026 -0400
[REFACTOR][Script] Simplify annotation scopes and signature variables
(#20456)
---
include/tvm/relax/script/ir_builder/frame.h | 3 -
include/tvm/script/ir_builder/base.h | 5 -
include/tvm/tirx/script/ir_builder/frame.h | 3 -
python/tvm/relax/script/ir_builder/frame.py | 10 +-
python/tvm/relax/script/ir_builder/ir.py | 4 +-
.../tvm/relax/script/ir_builder/parser_protocol.py | 18 +-
python/tvm/script/ir_builder/base.py | 26 --
python/tvm/script/ir_builder/parser_protocol.py | 52 +---
python/tvm/script/parser/entry.py | 7 +-
python/tvm/script/parser/inspect_source.py | 6 +-
python/tvm/script/parser/prescan.py | 95 +++---
python/tvm/script/parser/protocol_registry.py | 2 +-
python/tvm/script/parser/transpile.py | 336 +++++++++++----------
python/tvm/tirx/script/ir_builder/frame.py | 8 +-
.../tvm/tirx/script/ir_builder/parser_protocol.py | 16 +-
src/relax/analysis/well_formed.cc | 14 +-
src/relax/script/ir_builder/ir.cc | 4 -
src/script/ir_builder/base.cc | 13 -
src/tirx/script/ir_builder/ir.cc | 4 -
tests/python/relax/test_tvmscript_parser.py | 24 +-
tests/python/script/minilang.py | 21 +-
tests/python/script/test_parser_entry.py | 3 +-
.../python/tvmscript/test_tvmscript_parser_tir.py | 14 +-
23 files changed, 292 insertions(+), 396 deletions(-)
diff --git a/include/tvm/relax/script/ir_builder/frame.h
b/include/tvm/relax/script/ir_builder/frame.h
index 8ddba3659d..ad63178dbd 100644
--- a/include/tvm/relax/script/ir_builder/frame.h
+++ b/include/tvm/relax/script/ir_builder/frame.h
@@ -97,8 +97,6 @@ class FunctionFrameNode : public SeqExprFrameNode {
* However, we must specify the name by `R.func_name` before exit this
frame.
*/
ffi::Optional<ffi::String> name;
- /*! \brief Function-local symbols retained across declaration and body
entry. */
- ffi::Map<ffi::String, tvm::Var> type_var_map;
/*! \brief The function params. */
ffi::Array<tvm::Var> params;
/*!
@@ -130,7 +128,6 @@ class FunctionFrameNode : public SeqExprFrameNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FunctionFrameNode>()
- .def_ro("type_var_map", &FunctionFrameNode::type_var_map)
.def_ro("name", &FunctionFrameNode::name)
.def_ro("params", &FunctionFrameNode::params)
.def_ro("ret_ty", &FunctionFrameNode::ret_ty)
diff --git a/include/tvm/script/ir_builder/base.h
b/include/tvm/script/ir_builder/base.h
index ff9c798618..ee7cd9d72b 100644
--- a/include/tvm/script/ir_builder/base.h
+++ b/include/tvm/script/ir_builder/base.h
@@ -111,11 +111,6 @@ class IRBuilderFrame : public ffi::ObjectRef {
}
};
-/*! \brief Resolve one symbol in a native function frame's retained map. */
-TVM_DLL tvm::Var ResolveTypeVar(ffi::Map<ffi::String, tvm::Var>* symbols,
const ffi::String& name,
- ffi::Optional<PrimType> dtype,
ffi::Optional<tvm::Var> value,
- Span span);
-
////////////////////////////// IRBuilder //////////////////////////////
/*!
diff --git a/include/tvm/tirx/script/ir_builder/frame.h
b/include/tvm/tirx/script/ir_builder/frame.h
index 7836276b56..72defaff79 100644
--- a/include/tvm/tirx/script/ir_builder/frame.h
+++ b/include/tvm/tirx/script/ir_builder/frame.h
@@ -74,8 +74,6 @@ class TIRFrame : public IRBuilderFrame {
*/
class PrimFuncFrameNode : public TIRFrameNode {
public:
- /*! \brief Function-local symbols retained across declaration and body
entry. */
- ffi::Map<ffi::String, tvm::Var> type_var_map;
/*! \brief The name of the block. */
ffi::Optional<ffi::String> name;
/*! \brief Function parameters. */
@@ -101,7 +99,6 @@ class PrimFuncFrameNode : public TIRFrameNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PrimFuncFrameNode>()
- .def_ro("type_var_map", &PrimFuncFrameNode::type_var_map)
.def_ro("name", &PrimFuncFrameNode::name)
.def_ro("args", &PrimFuncFrameNode::args)
.def_ro("is_private", &PrimFuncFrameNode::is_private)
diff --git a/python/tvm/relax/script/ir_builder/frame.py
b/python/tvm/relax/script/ir_builder/frame.py
index e81d11808b..3ffc0bc2aa 100644
--- a/python/tvm/relax/script/ir_builder/frame.py
+++ b/python/tvm/relax/script/ir_builder/frame.py
@@ -18,9 +18,7 @@
from tvm_ffi import register_object as _register_object
-from tvm.script.ir_builder.base import IRBuilderFrame, _resolve_type_var
-
-from . import _ffi_api
+from tvm.script.ir_builder.base import IRBuilderFrame
@_register_object("script.ir_builder.relax.RelaxFrame")
@@ -34,11 +32,7 @@ class SeqExprFrame(RelaxFrame): ...
@_register_object("script.ir_builder.relax.FunctionFrame")
class FunctionFrame(SeqExprFrame):
- """Native function frame retaining signature, symbols and finalized
results."""
-
- def resolve_type_var(self, name, dtype=None, *, value=None, span=None):
- """Resolve a primitive symbol in this function's native map."""
- return _resolve_type_var(self, _ffi_api.ResolveTypeVar, name, dtype,
value=value, span=span)
+ """Native function frame retaining signature and finalized results."""
@_register_object("script.ir_builder.relax.BindingBlockFrame")
diff --git a/python/tvm/relax/script/ir_builder/ir.py
b/python/tvm/relax/script/ir_builder/ir.py
index c39005fa36..481fc74afc 100644
--- a/python/tvm/relax/script/ir_builder/ir.py
+++ b/python/tvm/relax/script/ir_builder/ir.py
@@ -526,8 +526,8 @@ def type_var(name, *, dtype=None, span=None):
Notes
-----
- This constructor creates a new symbol on each call. Use the language
variant
- resolver for symbols shared by name within a function signature.
+ This constructor creates a new symbol on each call. Reuse the returned
+ variable to share its identity across annotations and function parameters.
"""
return _ir.Var(
name,
diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py
b/python/tvm/relax/script/ir_builder/parser_protocol.py
index 1513cae10a..d23f748a11 100644
--- a/python/tvm/relax/script/ir_builder/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -106,8 +106,8 @@ def arg_(name: str, ty: Any, *, span: _Span = None) ->
_ir.Var:
ty : Type, Var, or callable
Parameter annotation; corresponds to ``annotation`` in the shared
:func:`tvm.script.ir_builder.parser_protocol.arg_` contract. Primitive
- annotations resolve through the active signature symbol context; an
- existing variable retains its identity.
+ annotations create a fresh parameter; an existing variable retains
+ its identity.
span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
@@ -118,8 +118,6 @@ def arg_(name: str, ty: Any, *, span: _Span = None) ->
_ir.Var:
"""
if not isinstance(ty, _ir.Var):
ty = _native._type(ty)
- if isinstance(ty, _ir.PrimType) or _ir.is_prim_var(ty):
- ty = resolve_type_var_(name, ty, span=span)
if isinstance(ty, _ir.Var):
return _ffi_api.ArgVar(name, ty)
return _base.at_(span, _ffi_api.Arg(name, _native._type(ty)))
@@ -195,17 +193,6 @@ def _check_module_well_formed(module: _ir.IRModule) ->
None:
# --------------------------------------
-def resolve_type_var_(
- name: str,
- dtype: str | _ir.Type | _ir.Var | None = None,
- *,
- value: _ir.Var | None = None,
- span: _Span = None,
-) -> _ir.Var:
- """Implements
:func:`tvm.script.ir_builder.parser_protocol.resolve_type_var_`."""
- return _base._current_function_frame().resolve_type_var(name, dtype,
value=value, span=span)
-
-
def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
"""Implements
:func:`tvm.script.ir_builder.parser_protocol.call_global_var_`."""
return _relax.Call(function, [_relax.utils.convert_to_expr(value) for
value in args])
@@ -605,7 +592,6 @@ __all__ = [
"output",
"range_",
"resolve_global_info_",
- "resolve_type_var_",
"return_",
"seq_expr",
"set_mutable_cell_",
diff --git a/python/tvm/script/ir_builder/base.py
b/python/tvm/script/ir_builder/base.py
index 0543fd8613..3234dc0960 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -407,32 +407,6 @@ def with_at_group_(
at_ = at
-def _resolve_type_var(frame, ffi_resolver, name, dtype=None, *, value=None,
span=None):
- """Validate a native function resolver's inputs, leaving its map owned by
C++."""
- if not isinstance(name, str) or not name:
- raise ValueError("A symbolic variable requires a nonempty string name")
- if isinstance(dtype, ir.Var):
- value, dtype = dtype, dtype.ty
- if isinstance(dtype, str):
- dtype = ir.PrimType(dtype)
- if dtype is not None and not isinstance(dtype, ir.PrimType):
- raise TypeError("A symbolic variable requires a primitive type")
- if value is not None and not ir.is_prim_var(value):
- raise TypeError("A symbolic binding requires a primitive Var")
- return ffi_resolver(
- frame, name, dtype, value, span.span if isinstance(span, SpanEntry)
else span
- )
-
-
-def _current_function_frame():
- """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 _return_annotation(annotation):
"""Evaluate the deferred return expression before normalizing its
annotation value."""
if callable(annotation) and not isinstance(annotation, ir.Expr | ir.Type):
diff --git a/python/tvm/script/ir_builder/parser_protocol.py
b/python/tvm/script/ir_builder/parser_protocol.py
index df66a680c4..6a16ef8146 100644
--- a/python/tvm/script/ir_builder/parser_protocol.py
+++ b/python/tvm/script/ir_builder/parser_protocol.py
@@ -162,10 +162,10 @@ def arg_(name: str, annotation: Any, *, span: _Span =
None) -> _ir.Var:
Notes
-----
- Requires an active function signature frame. Primitive symbols use that
frame's
- resolver, preserving annotation/body identity. TIRx handles buffer layout
according to
- its function policy; Relax converts its type annotation. Invalid
annotations or context
- raise TypeError/native builder errors.
+ Requires an active function signature frame. Existing variables retain
identity;
+ type annotations create fresh variables. TIRx handles buffer layout
according to
+ its function policy; Relax converts its type annotation. Invalid
annotations or
+ context raise TypeError/native builder errors.
.. code:: python
@@ -333,50 +333,6 @@ def resolve_global_info_(content: str) -> _ir.GlobalInfo:
return _global_infos()[name][int(index)]
-def resolve_type_var_(
- name: str,
- dtype: str | _ir.Type | _ir.Var | None = None,
- *,
- value: _ir.Var | None = None,
- span: _Span = None,
-) -> _ir.Var:
- """Resolve or declare a symbolic variable in the nearest native function.
-
- Parameters
- ----------
- name : str
- Function-local lookup key for explicit header parameters and captured
symbols.
- dtype : str, Type or Var, optional
- Explicit primitive type or supplied variable. None defaults new
symbols to
- int64; existing symbols must agree with an explicit dtype.
- value : Var, optional
- Existing primitive variable to register without replacement. None
creates a
- variable only if the name is not already registered.
- span : SpanEntry, Span or None, optional
- Location for the constructed result. None (the default) leaves
explicit location
- unspecified; active source-call provenance is composed by the builder.
Frames
- retain their location until finalization.
-
- Returns
- -------
- _ir.Var
- The exact existing or newly registered variable.
-
- Notes
- -----
- Requires an active native function frame. The signature and resumed body
use the same
- symbol map; nested functions use separate maps. Conflicting dtype/variable
declarations
- or missing context raise a builder error. No statement is emitted.
-
- .. code:: python
-
- # Source: def f[n: T.int32](...)
- # Generated builder
- n = X.resolve_type_var_("n", dtype="int32")
- """
- raise NotImplementedError
-
-
def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
"""Construct a call to a declared module function.
diff --git a/python/tvm/script/parser/entry.py
b/python/tvm/script/parser/entry.py
index bae99d679d..537a22abaf 100644
--- a/python/tvm/script/parser/entry.py
+++ b/python/tvm/script/parser/entry.py
@@ -688,13 +688,12 @@ def _run_statements(
# 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)
+ statements = transformer._create_definition_bindings(
+ node, transformer.function.capture_names, captures=captures
+ )
names = sorted(name for name in bound_names if name in namespace)
helper_name = transformer.module.make_fresh_name("_macro")
helper = ast.copy_location(
diff --git a/python/tvm/script/parser/inspect_source.py
b/python/tvm/script/parser/inspect_source.py
index 8ced5a8bba..8ec89704f1 100644
--- a/python/tvm/script/parser/inspect_source.py
+++ b/python/tvm/script/parser/inspect_source.py
@@ -40,7 +40,7 @@ from tvm_ffi.dataclasses import MISSING
from tvm.ir import SourceName, Span
from . import protocol_registry
-from .prescan import collect_annotation_free_names, resolve_namespace_key,
resolve_namespace_value
+from .prescan import collect_annotation_free_reads, resolve_namespace_key,
resolve_namespace_value
class _AnnotationScope(dict):
@@ -238,7 +238,7 @@ def capture_annotation_bindings(
protocol_registry.DefinitionKind.MACRO,
protocol_registry.DefinitionKind.PYTHON,
):
- names.update(collect_annotation_free_names(target.value))
+ names.update(read.id for read in
collect_annotation_free_reads(target.value))
for node in ast.walk(tree):
annotation = (
node.annotation
@@ -248,7 +248,7 @@ def capture_annotation_bindings(
else None
)
if annotation is not None:
- names.update(collect_annotation_free_names(annotation))
+ names.update(read.id for read in
collect_annotation_free_reads(annotation))
return {name: definition_scope[name] for name in names if name in
definition_scope}
diff --git a/python/tvm/script/parser/prescan.py
b/python/tvm/script/parser/prescan.py
index f10344e471..1d7b43f8dd 100644
--- a/python/tvm/script/parser/prescan.py
+++ b/python/tvm/script/parser/prescan.py
@@ -22,19 +22,15 @@ import ast
import builtins
import inspect
from collections.abc import Mapping
+from enum import IntEnum, auto
from types import ModuleType
from typing import NamedTuple, NoReturn
from . import protocol_registry
-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)}
-
-
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."""
+ """Collect lexical free reads for definition capture and shadowing
checks."""
bound = set() if bound is None else bound
if isinstance(node, ast.Name):
return [node] if isinstance(node.ctx, ast.Load) and node.id not in
bound else []
@@ -155,14 +151,26 @@ def _match_special_func(node: ast.AST | None, namespaces:
Mapping[str, str]) ->
return None
+class BindingKind(IntEnum):
+ """Syntax categories used by binding lowering and explicit-symbol
checks."""
+
+ ORDINARY = auto()
+ PARAMETER = auto()
+ FORWARD_PARAMETER = auto()
+ MUTABLE_PARAMETER = auto()
+ SYMBOL = auto()
+ MUTABLE = auto()
+ LOOP = auto()
+ MODULE_ALIAS = auto()
+ MUTABLE_UPDATE = auto()
+
+
class Binding(NamedTuple):
"""A source binding site, retained through rewriting without value
state."""
name: str
node: ast.AST
- kind: str
- annotation: ast.expr | None = None
- dtype: object = None
+ kind: BindingKind
class PrescanContext:
@@ -197,7 +205,11 @@ class PrescanContext:
# Derive mutable-name sets once for assignment dispatch, then read
them only;
# these are syntax categories, not independently updated runtime value
state.
self.mutable_names = {
- scope: {item.name for item in items if item.kind in ("mutable",
"mutable_parameter")}
+ scope: {
+ item.name
+ for item in items
+ if item.kind in (BindingKind.MUTABLE,
BindingKind.MUTABLE_PARAMETER)
+ }
for scope, items in bindings.items()
}
# Collected branch-ending names select conditional results without
later updates.
@@ -343,13 +355,11 @@ class PrescanCollector(ast.NodeVisitor):
self,
name: str,
node: ast.AST,
- kind: str = "ordinary",
- annotation: ast.expr | None = None,
- dtype: object = None,
+ kind: BindingKind = BindingKind.ORDINARY,
) -> None:
self._check_reserved(name, node)
self.names.add(name)
- item = Binding(name, node, kind, annotation, dtype)
+ item = Binding(name, node, kind)
self.bindings[self.scope].append(item)
self.sites[node] = item
@@ -470,7 +480,7 @@ class PrescanCollector(ast.NodeVisitor):
# x = X.arg_("x", ty)
# -------------------------------------------------
# Collect this function separately from its enclosing scope.
- self._record_binding(node.name, node, "function")
+ self._record_binding(node.name, node)
old_scope, old_builder = self.scope, self.builder
self.scope = node
self.functions.append(node)
@@ -495,17 +505,11 @@ class PrescanCollector(ast.NodeVisitor):
# body(n)
#
# Builder:
- # n = X.resolve_type_var_("n")
+ # n = I.dynamic("n")
# -------------------------------------------------
- bound = getattr(parameter, "bound", None)
- dtype = (
- "int64"
- if bound is None or (isinstance(bound, ast.Name) and bound.id
== "int")
- else protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
- _match_special_func(bound, self.namespaces)
- )
- )
- self._record_binding(parameter.name, parameter, "symbol",
dtype=dtype)
+ self._record_binding(parameter.name, parameter, BindingKind.SYMBOL)
+ earlier_reads: set[str] = set()
+ header_names = {parameter.name for parameter in getattr(node,
"type_params", ())}
for arg in [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]:
# -------------------- Pattern --------------------
# Python source:
@@ -526,13 +530,21 @@ class PrescanCollector(ast.NodeVisitor):
self._record_binding(
arg.arg,
arg,
- "mutable_parameter"
+ BindingKind.MUTABLE_PARAMETER
if inspect.getattr_static(self.builder,
"supports_mutable_declarations", True)
is True
and "parameter" in
protocol_registry.MUTABLE_CELL_DECL.get(constructor, ())
- else "parameter",
- arg.annotation,
+ else BindingKind.FORWARD_PARAMETER
+ if arg.arg in earlier_reads
+ and arg.arg not in header_names
+ and protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
+ _match_special_func(annotation, self.namespaces)
+ )
+ is not None
+ else BindingKind.PARAMETER,
)
+ if annotation is not None:
+ earlier_reads.update(read.id for read in
collect_annotation_free_reads(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)
@@ -560,19 +572,26 @@ class PrescanCollector(ast.NodeVisitor):
assignable = {
item.name
for item in facts
- if item.kind in ("parameter", "mutable_parameter", "mutable",
"loop")
+ if item.kind
+ in (
+ BindingKind.PARAMETER,
+ BindingKind.FORWARD_PARAMETER,
+ BindingKind.MUTABLE_PARAMETER,
+ BindingKind.MUTABLE,
+ BindingKind.LOOP,
+ )
}
# This temporary diagnostic index is derived from existing source
facts;
# it is not retained in the prescan result or used as value state.
origins: dict[str, ast.AST] = {}
for item in facts:
- if item.kind == "symbol":
+ if item.kind == BindingKind.SYMBOL:
origin = origins.get(item.name)
if origin is None or item.node.lineno < origin.lineno:
origins[item.name] = item.node
for item in facts:
# Explicit declarations and mutable updates are not symbol
rebindings.
- if item.kind == "symbol" or item.name in assignable:
+ if item.kind == BindingKind.SYMBOL or item.name in assignable:
continue
origin = origins.get(item.name)
if origin is not None and (item.node.lineno, item.node.col_offset)
> (
@@ -591,7 +610,7 @@ class PrescanCollector(ast.NodeVisitor):
value: ast.expr | None = None,
annotation: ast.expr | None = None,
*,
- kind: str = "ordinary",
+ kind: BindingKind = BindingKind.ORDINARY,
) -> None:
if self.functions:
values = [value]
@@ -625,11 +644,11 @@ class PrescanCollector(ast.NodeVisitor):
)
)
):
- self._record_binding(target.id, target, "mutable", annotation)
+ self._record_binding(target.id, target, BindingKind.MUTABLE)
elif isinstance(value, ast.Name) and value.id == self.module_name:
- self._record_binding(target.id, target, "module_alias")
+ self._record_binding(target.id, target,
BindingKind.MODULE_ALIAS)
elif isinstance(value, ast.Name) and any(
- item.name == value.id and item.kind == "module_alias"
+ item.name == value.id and item.kind == BindingKind.MODULE_ALIAS
for item in self.bindings[self.scope]
):
# -------------------- Pattern --------------------
@@ -640,9 +659,9 @@ class PrescanCollector(ast.NodeVisitor):
# second_alias = first_alias
# -------------------------------------------------
# Preserve the active module frame identity.
- self._record_binding(target.id, target, "module_alias")
+ self._record_binding(target.id, target,
BindingKind.MODULE_ALIAS)
else:
- self._record_binding(target.id, target, kind,
annotation=annotation)
+ self._record_binding(target.id, target, kind)
elif isinstance(target, ast.Tuple | ast.List):
values = (
value.elts
@@ -760,7 +779,7 @@ class PrescanCollector(ast.NodeVisitor):
# X.emit_(body(i))
# -------------------------------------------------
# Loop binders remain assignable and are never signature declarations.
- self._collect_target(node.target, kind="loop")
+ self._collect_target(node.target, kind=BindingKind.LOOP)
self.visit(node.iter)
for statement in node.body + node.orelse:
self.visit(statement)
diff --git a/python/tvm/script/parser/protocol_registry.py
b/python/tvm/script/parser/protocol_registry.py
index b8f745f419..67cfce82be 100644
--- a/python/tvm/script/parser/protocol_registry.py
+++ b/python/tvm/script/parser/protocol_registry.py
@@ -89,7 +89,7 @@ def register_scalar_annotation(
register_scalar_annotation("tirx.int32", T.int32, dtype="int32")
# Source: def f[n: T.int32](...):
- # Builder: n = X.resolve_type_var_("n", dtype="int32")
+ # Builder: n = I.dynamic("n", "int32")
"""
SCALAR_ANNOTATION_DTYPE[namespace_path] = dtype
return constructor
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index 50a7fc9440..1f5afefe20 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -30,8 +30,8 @@ entry-owned definition scope and hand generated helpers to
private recomposition
in one direction, without back-references to the rewriter. Entry executes the
recomposed builder and releases temporary captures. Short-lived frame names,
statement lists and assembly results stay local to the methods that need them.
-Native frames own symbols, declarations, parameters, region results and final
-IR; the Python contexts do not mirror that construction state.
+Native frames own declarations, parameters, region results and final IR;
+the Python contexts do not mirror that construction state.
"""
from __future__ import annotations
@@ -46,9 +46,8 @@ from typing import Any, NamedTuple, NoReturn, TypeVar
from . import protocol_registry
from .inspect_source import _AnnotationScope
from .prescan import (
- Binding,
+ BindingKind,
PrescanContext,
- collect_annotation_free_names,
collect_annotation_free_reads,
)
@@ -84,15 +83,6 @@ 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:
@@ -237,8 +227,9 @@ class FunctionContext:
self.current_scope = current_scope
# Namespace identifier used by _call_dialect for this function's
builder operations.
self.dialect_prefix = dialect_prefix
- # One snapshot supplies original-name helper defaults at the
definition site.
+ # Only reads that cross a Python scope boundary need definition
snapshots.
self.definition_captures: str | None = None
+ self.capture_names: set[str] = set()
class IRBuilderTranspiler(ast.NodeTransformer):
@@ -262,7 +253,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# Bypass syntax-to-builder lowering.
# Still traverse children and instrument source calls with spans.
self.bypass_ast_rewrite = False
- # Only free reads of the active annotation need lazy missing-name
checks.
+ # Only free reads that cross a scope boundary use the definition
snapshot.
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.
@@ -418,10 +409,11 @@ class IRBuilderTranspiler(ast.NodeTransformer):
self.bypass_ast_rewrite = old
@contextmanager
- def _rewrite_annotation(self, node: ast.expr) -> Iterator[None]:
- """Check missing values only at free reads, without renaming Python
bindings."""
+ def _rewrite_annotation(self, reads: set[ast.Name]) -> Iterator[None]:
+ """Keep Python reads direct unless their definition scope would be
shadowed."""
previous_reads = self.annotation_reads
- self.annotation_reads = set(collect_annotation_free_reads(node))
+ self.annotation_reads = reads
+ self.function.capture_names.update(read.id for read in
self.annotation_reads)
try:
yield
finally:
@@ -448,16 +440,16 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# value: X.Tensor((n,))
#
# Builder:
- # value_type = X.Tensor((_require_annotation_value(n, "n"),))
+ # value_type = X.Tensor((n,))
# -------------------------------------------------
- # Helpers bind captured values under their original names. Python
handles
- # lambda/comprehension locals and preceding signature parameters
directly.
+ # Python handles ordinary names, lambda/comprehension locals, and
earlier
+ # parameters. A selected cross-scope read keeps lazy NameError
semantics.
if node in self.annotation_reads:
- if node.id in self.module.prescan_ctx.namespaces:
- return node
return ast.copy_location(
- ast.Call(
- self._inject(_require_annotation_value), [node,
ast.Constant(node.id)], []
+ ast.Subscript(
+ ast.Name(self.function.definition_captures, ast.Load()),
+ ast.Constant(node.id),
+ ast.Load(),
),
node,
)
@@ -553,7 +545,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
for item in
self.module.prescan_ctx.bindings.get(self.function.current_scope, ())
if item.name == node.id
]
- return bool(records) and all(item.kind == "module_alias" for item in
records)
+ return bool(records) and all(item.kind == BindingKind.MODULE_ALIAS for
item in records)
def _visit_direct_operand(self, node: ast.expr) -> ast.expr:
"""Preserve an existing payload's span without bypassing child
operations."""
@@ -834,21 +826,21 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
return self.visit(node)
- def _binding_kind(self, target: ast.Name, *, frame_value: bool = False) ->
str:
+ def _binding_kind(self, target: ast.Name, *, frame_value: bool = False) ->
BindingKind:
"""Select the declaration/store/binding operation from existing syntax
facts."""
if frame_value:
- return "ordinary"
+ return BindingKind.ORDINARY
site = self.module.prescan_ctx.sites.get(target)
- kind = site.kind if site is not None else "ordinary"
- if kind in ("mutable", "module_alias"):
+ kind = site.kind if site is not None else BindingKind.ORDINARY
+ if kind in (BindingKind.MUTABLE, BindingKind.MODULE_ALIAS):
return kind
mutable =
self.module.prescan_ctx.mutable_names.get(self.function.current_scope, ())
- return "mutable_update" if target.id in mutable else "ordinary"
+ return BindingKind.MUTABLE_UPDATE if target.id in mutable else
BindingKind.ORDINARY
def _uses_ordinary_binding(self, target: ast.expr) -> bool:
"""Test assignment targets without evaluating or inspecting their RHS
values."""
if isinstance(target, ast.Name):
- return self._binding_kind(target) == "ordinary"
+ return self._binding_kind(target) == BindingKind.ORDINARY
if isinstance(target, ast.Starred):
return self._uses_ordinary_binding(target.value)
if isinstance(target, ast.Tuple | ast.List):
@@ -881,7 +873,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
keywords["ty"] = ty
if frame_value:
keywords["frame_value"] = ast.Constant(True)
- if kind == "mutable" and not frame_value:
+ if kind == BindingKind.MUTABLE and not frame_value:
# -------------------- Pattern --------------------
# Python source:
# x = X.local_scalar(initial)
@@ -896,7 +888,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
name_span=self.module.make_span_expr(target),
keywords=keywords,
)
- elif kind == "module_alias" and not frame_value:
+ elif kind == BindingKind.MODULE_ALIAS and not frame_value:
# -------------------- Pattern --------------------
# Python source:
# alias = Module
@@ -905,7 +897,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# alias = Module
# -------------------------------------------------
pass
- elif kind == "mutable_update":
+ elif kind == BindingKind.MUTABLE_UPDATE:
# -------------------- Pattern --------------------
# Python source:
# x = value
@@ -1103,12 +1095,8 @@ class IRBuilderTranspiler(ast.NodeTransformer):
if node.value
else ast.Attribute(ast.Name(self.module.ir_prefix, ast.Load()),
"MISSING", ast.Load())
)
- annotation = node.annotation
- 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)
+ with self._rewrite_annotation(self._body_annotation_captures(node)):
+ annotation = self.visit(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]:
@@ -1627,31 +1615,31 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# Conflicting definition and execution namespaces need an injected
name.
return self._inject(namespace, "_X").id
- def _read_function_annotations(
- 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 = [parameter.annotation for parameter in parameters]
- returns = node.returns
- annotation_names = {
- name
- for annotation in [*annotations, returns]
- if annotation is not None
- for name in collect_annotation_free_names(annotation)
- }
- body_annotations = [
- item.annotation
- for item in facts
- if item.annotation is not None
- and item.kind not in ("parameter", "mutable_parameter", "symbol")
- ]
- annotation_names.update(
- name
- for annotation in body_annotations
- for name in collect_annotation_free_names(annotation)
- )
- return annotations, returns, annotation_names - declared_names
+ def _body_annotation_captures(self, node: ast.AnnAssign) -> set[ast.Name]:
+ """Retain definition reads only where the body has a different Python
scope."""
+ facts =
self.module.prescan_ctx.bindings.get(self.function.current_scope, ())
+ captures = set()
+ for read in collect_annotation_free_reads(node.annotation):
+ name = read.id
+ if name in self.module.prescan_ctx.namespaces:
+ continue
+ bindings = [item for item in facts if item.name == name]
+ if any(
+ (item.node.lineno, item.node.col_offset)
+ < (node.annotation.lineno, node.annotation.col_offset)
+ for item in bindings
+ ):
+ continue
+ if (
+ bindings
+ or self.module.module_name is not None
+ or (
+ name in self.module.definition_scope
+ and self.module.definition_scope[name] is not
self.module.exec_globals.get(name)
+ )
+ ):
+ captures.add(read)
+ return captures
def _create_definition_bindings(
self,
@@ -1660,12 +1648,9 @@ class IRBuilderTranspiler(ast.NodeTransformer):
*,
captures: str,
) -> list[ast.stmt]:
- """Snapshot only lexical annotation/constexpr names at their
definition site."""
- names = names | {
- parameter.arg
- for parameter in [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]
- if self._is_constexpr_annotation(parameter.annotation)
- }
+ """Snapshot only consumed cross-scope reads at their definition
site."""
+ 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), [], [])]
@@ -1680,47 +1665,6 @@ 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
@@ -1733,22 +1677,31 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
return [self._assign(const_args, const_args_expr, node)]
- def _create_symbol_declarations(self, node: ast.FunctionDef) ->
list[ast.stmt]:
- """Bind explicit signature type parameters with their declared
dtypes."""
+ def _create_symbol_declarations(
+ self, node: ast.FunctionDef
+ ) -> tuple[list[ast.stmt], dict[str, str]]:
+ """Create explicit header symbols and carry their identities into the
body."""
declaration: list[ast.stmt] = []
+ aliases: dict[str, str] = {}
# -------------------- Pattern --------------------
# Python source:
# def f[n]():
# body(n)
#
# Builder:
- # n = X.resolve_type_var_("n")
+ # n = I.dynamic("n")
# -------------------------------------------------
for parameter in getattr(node, "type_params", ()):
if not isinstance(parameter, getattr(ast, "TypeVar", ())):
self._raise_error(parameter, "Only scalar type parameters are
supported")
bound = getattr(parameter, "bound", None)
- dtype = self.module.prescan_ctx.sites[parameter].dtype
+ dtype = (
+ "int64"
+ if bound is None or (isinstance(bound, ast.Name) and bound.id
== "int")
+ else protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
+ self.module.prescan_ctx._match_special_func(bound)
+ )
+ )
if bound is not None and not (isinstance(bound, ast.Name) and
bound.id == "int"):
if (
protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
@@ -1762,19 +1715,23 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
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")
+ aliases[parameter.name] = alias
declaration.append(
self._assign(
- parameter.name,
- self._call_dialect(
- "resolve_type_var_",
- [ast.Constant(parameter.name)],
+ alias,
+ self._call(
+ self.module.ir_prefix,
+ "dynamic",
+ [ast.Constant(parameter.name), ast.Constant(dtype)],
parameter,
- keywords={"dtype": ast.Constant(dtype)},
+ span=self.module.make_span_expr(parameter),
),
parameter,
)
)
- return declaration
+ declaration.append(self._assign(parameter.name, ast.Name(alias,
ast.Load()), parameter))
+ return declaration, aliases
@staticmethod
def _select_specialized_value(
@@ -1810,6 +1767,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
*,
captures: str,
const_args: str | None,
+ symbols: dict[str, str],
) -> tuple[list[ast.stmt], dict[str, str]]:
"""Declare signature parameters, making constexpr values available
first."""
declaration: list[ast.stmt] = []
@@ -1822,6 +1780,47 @@ 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))
+ forward: dict[str, str] = {}
+ for parameter, annotation, _ in other_params:
+ if self.module.prescan_ctx.sites[parameter].kind !=
BindingKind.FORWARD_PARAMETER:
+ continue
+ name = parameter.arg
+ self.function.capture_names.add(name)
+ alias = self.module.make_fresh_name("_parameter")
+ forward[name] = alias
+ dtype = protocol_registry.SCALAR_ANNOTATION_DTYPE[
+ self.module.prescan_ctx._match_special_func(annotation)
+ ]
+ created = self._call(
+ self.module.ir_prefix,
+ "dynamic",
+ [ast.Constant(name), ast.Constant(dtype)],
+ parameter,
+ span=self.module.make_span_expr(parameter),
+ )
+ if const_args is not None:
+ created = self._select_specialized_value(
+ name, created, parameter, const_args=const_args
+ )
+ # An outer dimension still belongs to earlier annotations. The
later
+ # scalar gets a fresh parameter unless this declaration supplied
it.
+ supplied = ast.Compare(ast.Constant(name), [ast.In()],
[ast.Name(captures, ast.Load())])
+ declaration.append(
+ self._assign(alias, ast.IfExp(supplied,
self.visit(annotation), created), parameter)
+ )
+ declaration.append(
+ self._assign(
+ name,
+ self._call(
+ captures,
+ "get",
+ [ast.Constant(name), ast.Name(alias, ast.Load())],
+ parameter,
+ ),
+ parameter,
+ )
+ )
+ unbound = {parameter.arg for parameter in parameters} - symbols.keys()
- forward.keys()
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")
@@ -1832,6 +1831,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# alias so that binding does not become a local `name = name`.
alias = self.module.make_fresh_name("_parameter")
constexpr_aliases[name] = alias
+ self.function.capture_names.add(name)
fallback = ast.Call(
self._inject(_require_constexpr_arg),
[
@@ -1853,8 +1853,24 @@ class IRBuilderTranspiler(ast.NodeTransformer):
[],
)
else:
- with self._rewrite_annotation(annotation):
- translated = self.visit(annotation)
+ if name in forward:
+ translated = ast.Name(forward[name], ast.Load())
+ elif (
+ name in symbols
+ and protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
+ self.module.prescan_ctx._match_special_func(annotation)
+ )
+ is not None
+ ):
+ translated = ast.Name(name, ast.Load())
+ else:
+ reads = {
+ read
+ for read in collect_annotation_free_reads(annotation)
+ if read.id in unbound
+ }
+ with self._rewrite_annotation(reads):
+ 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.
@@ -1875,6 +1891,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
)
declaration.append(self._assign(alias, value, parameter))
+ unbound.discard(name)
if alias != name:
declaration.append(self._assign(name, ast.Name(alias,
ast.Load()), parameter))
return declaration, constexpr_aliases
@@ -1943,6 +1960,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
node: ast.FunctionDef,
parameters: list[ast.arg],
constexpr_aliases: dict[str, str],
+ symbol_aliases: dict[str, str],
*,
frame: str,
const_args: str | None,
@@ -2001,16 +2019,8 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
)
body.append(self._assign(name, value, parameter))
- for parameter in getattr(node, "type_params", ()):
- body.append(
- self._assign(
- parameter.name,
- self._call_dialect(
- "resolve_type_var_", [ast.Constant(parameter.name)],
parameter
- ),
- parameter,
- )
- )
+ for name, alias in symbol_aliases.items():
+ body.append(self._assign(name, ast.Name(alias, ast.Load()), node))
return body
def _create_split_declaration(
@@ -2019,26 +2029,24 @@ class IRBuilderTranspiler(ast.NodeTransformer):
frame: str,
definition: ast.FunctionDef,
frame_declaration: list[ast.stmt],
- capture_names: set[str],
- constexpr_aliases: dict[str, str],
+ declaration_values: list[str],
) -> list[ast.stmt]:
"""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.
+ # The declaration helper returns its frame, explicit symbols and host
values.
+ # The sibling body reads runtime parameters from the frame and receives
+ # explicit symbol identities directly, without resolving names again.
declare_name = self.module.make_fresh_name("_declare")
signature = frame_declaration[0]
# 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()]
+ outputs = [frame, *declaration_values]
returned = ast.copy_location(
ast.Return(ast.Tuple([ast.Name(name, ast.Load()) for name in
outputs], ast.Load())),
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(
@@ -2057,9 +2065,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.
- 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.
+ Python closures supply annotation values; only shadowed reads need a
+ definition snapshot. Signature assignments expose each parameter to
later
+ annotations, while body helpers preserve source globals and closures.
"""
kind, options = self.read_function_metadata(node)
builder = self._function_namespace(node, kind)
@@ -2072,13 +2080,11 @@ class IRBuilderTranspiler(ast.NodeTransformer):
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
- )
+ annotations = [parameter.annotation for parameter in parameters]
+ returns = node.returns
captures = self.module.make_fresh_name("_definition")
self.function.definition_captures = captures
- statements = self._create_definition_bindings(node, capture_names,
captures=captures)
+ statements: list[ast.stmt] = []
const_args = None
if self.module.enable_jit_map and not local_function:
const_args = self.module.make_fresh_name("_const_args")
@@ -2089,14 +2095,18 @@ class IRBuilderTranspiler(ast.NodeTransformer):
node,
)
]
- declaration.extend(self._create_symbol_declarations(node))
+ symbols, symbol_aliases = self._create_symbol_declarations(node)
+ declaration.extend(symbols)
arguments, constexpr_aliases = self._rewrite_parameters(
- parameters, annotations, captures=captures,
const_args=const_args
+ parameters,
+ annotations,
+ captures=captures,
+ const_args=const_args,
+ symbols=symbol_aliases,
)
declaration.extend(arguments)
if returns is not None:
- with self._rewrite_annotation(returns):
- translated_return = self.visit(returns)
+ translated_return = self.visit(returns)
declaration.append(
ast.copy_location(
ast.Expr(
@@ -2121,7 +2131,12 @@ class IRBuilderTranspiler(ast.NodeTransformer):
split_declare=split_declare,
)
body = self._create_body_parameters(
- node, parameters, constexpr_aliases, frame=frame,
const_args=const_args
+ node,
+ parameters,
+ constexpr_aliases,
+ symbol_aliases,
+ frame=frame,
+ const_args=const_args,
)
body.extend(self.transform_statements(node.body))
definition = self._create_definition(body_name, body, node)
@@ -2153,14 +2168,15 @@ class IRBuilderTranspiler(ast.NodeTransformer):
invocation = ast.copy_location(
ast.Expr(ast.Call(ast.Name(body_name, ast.Load()), [], [])),
node
)
+ declaration_values = [*symbol_aliases.values(),
*constexpr_aliases.values()]
if split_declare:
statements.extend(
self._create_split_declaration(
- node, frame, definition, frame_declaration,
capture_names, constexpr_aliases
+ node, frame, definition, frame_declaration,
declaration_values
)
)
else:
- outputs = list(constexpr_aliases.values())
+ outputs = declaration_values
if outputs:
declaration.append(
ast.copy_location(
@@ -2175,7 +2191,6 @@ class IRBuilderTranspiler(ast.NodeTransformer):
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(
@@ -2195,7 +2210,10 @@ class IRBuilderTranspiler(ast.NodeTransformer):
),
node,
)
- return statements, frame, resumed
+ snapshots = self._create_definition_bindings(
+ node, self.function.capture_names, captures=captures
+ )
+ return [*snapshots, *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/ir_builder/frame.py
b/python/tvm/tirx/script/ir_builder/frame.py
index 870dd25d16..2aaedb4e71 100644
--- a/python/tvm/tirx/script/ir_builder/frame.py
+++ b/python/tvm/tirx/script/ir_builder/frame.py
@@ -21,7 +21,7 @@ from collections.abc import Sequence
from tvm_ffi import Array
from tvm_ffi import register_object as _register_object
-from tvm.script.ir_builder.base import IRBuilderFrame, _resolve_type_var
+from tvm.script.ir_builder.base import IRBuilderFrame
from tvm.tirx import Buffer, Var
from . import _ffi_api
@@ -33,7 +33,7 @@ class TIRFrame(IRBuilderFrame): ...
@_register_object("script.ir_builder.tirx.PrimFuncFrame")
class PrimFuncFrame(TIRFrame):
- """Native function frame retaining signature, symbols and finalized
results."""
+ """Native function frame retaining signature and finalized results."""
def default_buffer_layout(self, shape, scope):
"""Select the default layout for buffers constructed in this
function."""
@@ -46,10 +46,6 @@ class PrimFuncFrame(TIRFrame):
"""The native declared parameters, shared with the resumed body."""
return self.args
- def resolve_type_var(self, name, dtype=None, *, value=None, span=None):
- """Resolve a primitive symbol in this function's native map."""
- return _resolve_type_var(self, _ffi_api.ResolveTypeVar, name, dtype,
value=value, span=span)
-
@_register_object("script.ir_builder.tirx.ForFrame")
class ForFrame(TIRFrame):
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index aab9fa4147..3be37d2dbc 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -121,9 +121,7 @@ def arg_(name: str, annotation: Any, *, span: _Span = None)
-> _ir.Var:
raise TypeError("T.Optional is only supported by @T.jit")
if callable(annotation) and not isinstance(annotation, _ir.Expr):
annotation = annotation()
- if isinstance(annotation, _ir.PrimType) or _ir.is_prim_var(annotation):
- annotation = resolve_type_var_(name, annotation, span=span)
- elif isinstance(annotation, _ir.Type):
+ if isinstance(annotation, _ir.Type):
annotation = _ir.Var(name, annotation)
return _ffi_api.Arg(name, _base.at_(span, annotation))
@@ -209,17 +207,6 @@ def resolve_global_info_(content: Any) -> Any:
raise NotImplementedError("TIRx does not support global-info lookup")
-def resolve_type_var_(
- name: str,
- dtype: str | _ir.Type | _ir.Var | None = None,
- *,
- value: _ir.Var | None = None,
- span: _Span = None,
-) -> _ir.Var:
- """Implements
:func:`tvm.script.ir_builder.parser_protocol.resolve_type_var_`."""
- return _base._current_function_frame().resolve_type_var(name, dtype,
value=value, span=span)
-
-
def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
"""Implements
:func:`tvm.script.ir_builder.parser_protocol.call_global_var_`."""
return _op._call_global(function, *args)
@@ -1107,7 +1094,6 @@ __all__ = [
"prim_func",
"range_",
"resolve_global_info_",
- "resolve_type_var_",
"return_",
"serial",
"set_mutable_cell_",
diff --git a/src/relax/analysis/well_formed.cc
b/src/relax/analysis/well_formed.cc
index aa5b443b40..6146707eb8 100644
--- a/src/relax/analysis/well_formed.cc
+++ b/src/relax/analysis/well_formed.cc
@@ -295,13 +295,23 @@ class WellFormedChecker : public relax::ExprVisitor,
public relax::TypeVisitor {
}
}
- // Runtime parameters themselves are definitions only in source order.
+ // Primitive parameters may supply dimensions anywhere in the signature.
+ // Register each once so duplicate parameters remain an error.
+ for (Var param : op->params) {
+ if (GetType(param).as<PrimTypeNode>()) {
+ RegisterVarDefinition(param);
+ }
+ }
+
+ // Other runtime parameters retain source-order definition scope.
for (Var param : op->params) {
if (auto* dataflow_var = param.as<DataflowVarNode>()) {
TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef<DataflowVar>(dataflow_var))
<< "DataflowVar " << param << " is defined outside DataflowBlock.";
}
- RegisterVarDefinition(param);
+ if (!GetType(param).as<PrimTypeNode>()) {
+ RegisterVarDefinition(param);
+ }
auto it = param_var_func_map_.find(param);
if (it != param_var_func_map_.end() && it->second != cur_visited_func_) {
diff --git a/src/relax/script/ir_builder/ir.cc
b/src/relax/script/ir_builder/ir.cc
index 4104523ac9..4316b9937c 100644
--- a/src/relax/script/ir_builder/ir.cc
+++ b/src/relax/script/ir_builder/ir.cc
@@ -203,10 +203,6 @@ void FuncRetValue(const tvm::relax::Expr& value) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("script.ir_builder.relax.ResolveTypeVar",
- [](FunctionFrame frame, ffi::String name, ffi::Optional<PrimType>
dtype,
- ffi::Optional<tvm::Var> value,
- Span span) { return ResolveTypeVar(&frame->type_var_map, name,
dtype, value, span); })
.def("script.ir_builder.relax.Function", Function)
.def("script.ir_builder.relax.DeclFunction", DeclFunction)
.def("script.ir_builder.relax.LocalFunction", LocalFunction)
diff --git a/src/script/ir_builder/base.cc b/src/script/ir_builder/base.cc
index 6231a9581d..8a98ab2fb1 100644
--- a/src/script/ir_builder/base.cc
+++ b/src/script/ir_builder/base.cc
@@ -29,19 +29,6 @@ namespace tvm {
namespace script {
namespace ir_builder {
-tvm::Var ResolveTypeVar(ffi::Map<ffi::String, tvm::Var>* symbols, const
ffi::String& name,
- ffi::Optional<PrimType> dtype, ffi::Optional<tvm::Var>
value, Span span) {
- TVM_FFI_CHECK(!name.empty(), ValueError) << "A symbolic variable requires a
nonempty name";
- if (auto existing = symbols->Get(name)) return existing.value();
- tvm::Var symbol =
- value.has_value() ? value.value() : tvm::Var(name,
dtype.value_or(PrimType::Int(64)), span);
- TVM_FFI_CHECK(symbol->ty.as<PrimTypeNode>(), TypeError)
- << "A symbolic variable requires a primitive type";
- if (symbol->name.empty()) details::Namer::Name(symbol, name);
- symbols->Set(name, symbol);
- return symbol;
-}
-
namespace {
bool PositionLessEqual(int lhs_line, int lhs_column, int rhs_line, int
rhs_column) {
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index e569af59f4..fbe17af68f 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -726,10 +726,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("script.ir_builder.tirx.ResolveTypeVar",
- [](PrimFuncFrame frame, ffi::String name, ffi::Optional<PrimType>
dtype,
- ffi::Optional<tvm::Var> value,
- Span span) { return ResolveTypeVar(&frame->type_var_map, name,
dtype, value, span); })
.def("script.ir_builder.tirx.Buffer",
static_cast<BufferVar (*)(ffi::Array<PrimExpr>, PrimType,
ffi::String,
ffi::Optional<Expr>,
ffi::Optional<ffi::Array<PrimExpr>>,
diff --git a/tests/python/relax/test_tvmscript_parser.py
b/tests/python/relax/test_tvmscript_parser.py
index 3a88563bc4..488a3a5e75 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -1695,21 +1695,23 @@ class Module:
_check(mod)
-def test_later_prim_param_requires_external_shape_symbol():
- with pytest.raises(NameError):
- tvm.script.from_source(
- """
+def test_later_prim_param_reuses_shape_symbol():
+ function = tvm.script.from_source(
+ """
@R.function
def main(x: R.Tensor([n], "float32"), n: T.int64):
return x
""",
- extra_vars={
- "I": tvm.script.ir,
- "R": tvm.script.relax,
- "T": tvm.script.tirx,
- "Ts": tvm.script.s_tir,
- },
- )
+ extra_vars={
+ "I": tvm.script.ir,
+ "R": tvm.script.relax,
+ "T": tvm.script.tirx,
+ "Ts": tvm.script.s_tir,
+ },
+ )
+ x, n = function.params
+ assert x.ty.shape[0].same_as(n)
+ assert str(n.ty.dtype) == "int64"
def test_non_int64_prim_param_rejected_in_shape_annotation():
diff --git a/tests/python/script/minilang.py b/tests/python/script/minilang.py
index 6439b299c9..a0f2686a0e 100644
--- a/tests/python/script/minilang.py
+++ b/tests/python/script/minilang.py
@@ -78,7 +78,6 @@ class Frame:
)
self.function = Function() if kind == "function" else None
self.params = []
- self.type_var_map = {}
self.global_var = Value("global", (self,))
self.local_var = Value("local", (self,))
self.result = None
@@ -108,15 +107,6 @@ class Frame:
self.language.result = Module(self.language.functions)
return False
- def resolve_type_var(self, name, dtype=None, *, value=None, **kwargs):
- if name not in self.type_var_map:
- self.type_var_map[name] = (
- value
- if value is not None
- else Value("symbol", ("int64" if dtype is None else dtype,),
name)
- )
- return self.type_var_map[name]
-
def __getitem__(self, name):
return self.language.functions[name]
@@ -180,7 +170,6 @@ class Language:
),
break_=lambda **kwargs: self.statement("break"),
continue_=lambda **kwargs: self.statement("continue"),
- resolve_type_var_=self.resolve_type_var,
resolve_global_info_=self.resolve_global_info,
bind_=self.bind,
check_well_formed_=lambda result: None,
@@ -217,14 +206,14 @@ class Language:
def Tensor(shape=None, dtype="float32", device=None, placement="S[0]"):
return Value("tensor", (shape, dtype, device, placement))
- def dynamic(name, dtype="int64"):
- return Value("symbol", (dtype,), name)
+ def dynamic(name, dtype="int64", *, span=None):
+ return self.at(span, Value("symbol", (dtype,), name))
def cell(value=None):
return Value("cell", (value,))
self.M.Tensor = Tensor
- self.M.dynamic = dynamic
+ self.M.dynamic = self.I.dynamic = dynamic
self.M.int32 = protocol_registry.register_scalar_annotation(
"M.int32", lambda: None, dtype="int32"
)
@@ -262,10 +251,6 @@ class Language:
self.frame().function.body.append(statement)
return AlreadyEmitted(statement)
- def resolve_type_var(self, name, dtype=None, **kwargs):
- value = self.frame().resolve_type_var(name, dtype, **kwargs)
- return value
-
def resolve_global_info(self, name):
return self.global_infos[name]
diff --git a/tests/python/script/test_parser_entry.py
b/tests/python/script/test_parser_entry.py
index 931a4535d3..e0f4f44f34 100644
--- a/tests/python/script/test_parser_entry.py
+++ b/tests/python/script/test_parser_entry.py
@@ -73,10 +73,9 @@ 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:
- def _declare0(*, M=_definition0.get('M', _I0.MISSING)):
+ def _declare0():
M.func_name_('identity')
x = M.arg_('x', M.Tensor((4,)))
diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py
b/tests/python/tvmscript/test_tvmscript_parser_tir.py
index 59002c571e..14e5fa6415 100644
--- a/tests/python/tvmscript/test_tvmscript_parser_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py
@@ -172,11 +172,15 @@ class Module:
""",
],
)
-def test_tir_direct_later_prim_param_is_undefined(source):
- with pytest.raises(NameError):
- tvm.script.from_source(
- source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx,
"Ts": tvm.script.s_tir}
- )
+def test_tir_direct_later_prim_param_reuses_shape_symbol(source):
+ result = tvm.script.from_source(
+ source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx, "Ts":
tvm.script.s_tir}
+ )
+ function = result["main"] if isinstance(result, tvm.IRModule) else result
+ A, n = function.params
+ assert A.ty.shape[0].same_as(n)
+ assert str(n.ty.dtype) == "int32"
+ assert function.body.value.same_as(n)
def test_tir_return_annotation_does_not_define_symbolic_var():
