This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tvmscript-ast-only-transpiler
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/tvmscript-ast-only-transpiler
by this push:
new 5b7969b327 Own syntax registration in the parser protocol
5b7969b327 is described below
commit 5b7969b3279b8fda3cdefcc2521e2b5ece29ddd5
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Sep 21 17:16:36 2026 +0000
Own syntax registration in the parser protocol
Keep expression-string and function-decorator metadata in one parser-owned
registry, with explicit flat NamedTuple fields. Delegate concrete eager
annotation behavior to builders and load parser entry points lazily to preserve
the import boundary.
---
python/tvm/relax/script/builder/__init__.py | 7 +-
python/tvm/script/ir_builder/__init__.py | 6 +-
python/tvm/script/ir_builder/ir/__init__.py | 3 +-
python/tvm/script/ir_builder/protocol.py | 300 +++++++++-------------------
python/tvm/script/parser/__init__.py | 37 ++--
python/tvm/script/parser/expression.py | 2 +-
python/tvm/script/parser/frontend.py | 12 +-
python/tvm/script/parser/protocol.py | 192 ++++++++++++++++++
python/tvm/script/parser/transpile.py | 23 +--
python/tvm/tirx/script/builder/__init__.py | 2 +-
tests/python/tvmscript/test_parser.py | 33 +++
11 files changed, 371 insertions(+), 246 deletions(-)
diff --git a/python/tvm/relax/script/builder/__init__.py
b/python/tvm/relax/script/builder/__init__.py
index 16e6386287..ac22ef94df 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/relax/script/builder/__init__.py
@@ -37,6 +37,7 @@ from tvm.script.ir_builder import protocol as _protocol
from tvm.script.ir_builder.type_var_frame import TypeVarDecl as _TypeVarDecl
from tvm.script.ir_builder.type_var_frame import TypeVarFrame as _TypeVarFrame
from tvm.script.ir_builder.type_var_frame import resolve_type_var
+from tvm.script.parser.protocol import expr_str_args as _expr_str_args
from . import _ffi_api, distributed, frame, ir
from . import distributed as dist
@@ -46,7 +47,7 @@ from .distributed.ir import _lookup_device_mesh
from .ir import *
-@_protocol.expression_args("shape", introduce=True, dtype="int64",
scalar_strings=False)
+@_expr_str_args("shape", introduce=True, dtype="int64", scalar_strings=False)
def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
"""Construct a tensor type from resolved shape dimensions.
@@ -69,7 +70,7 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *,
span=None):
return _relax.TensorType(shape, dtype, vdevice, ndim,
_protocol.source_span(span))
-@_protocol.expression_args("shape", introduce=True, dtype="int64",
scalar_strings=False)
+@_expr_str_args("shape", introduce=True, dtype="int64", scalar_strings=False)
def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *,
ndim=-1, span=None):
"""Construct a concrete distributed tensor type with resolved shape
symbols."""
if device_mesh is None:
@@ -86,7 +87,7 @@ def DTensor(shape=None, dtype=None, device_mesh=None,
placement="", *, ndim=-1,
Range = _ir.Range
-@_protocol.expression_args("values", introduce=True, dtype="int64")
+@_expr_str_args("values", introduce=True, dtype="int64")
def Shape(values=None, ndim=-1, *, span=None):
"""Construct a shape type from concrete dimension expressions.
diff --git a/python/tvm/script/ir_builder/__init__.py
b/python/tvm/script/ir_builder/__init__.py
index 30f34d0a51..398759b181 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -35,11 +35,15 @@ import importlib
from typing import Any
from .base import IRBuilder
-from .protocol import MISSING, ExprStrPolicy, _at, expr_str_args,
expression_args
+from .protocol import MISSING, _at
from .type_var_frame import TypeVarDecl, TypeVarFrame
def __getattr__(name: str) -> Any:
+ if name in {"ExprStrPolicy", "expr_str_args", "expression_args"}:
+ from tvm.script.parser import protocol
+
+ return getattr(protocol, name)
# Lazy import to avoid loading tvm.script during dialect bootstrap.
from tvm.script import _DIALECT_REGISTRY # pylint:
disable=import-outside-toplevel
diff --git a/python/tvm/script/ir_builder/ir/__init__.py
b/python/tvm/script/ir_builder/ir/__init__.py
index 04a973c307..16608098a7 100644
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ b/python/tvm/script/ir_builder/ir/__init__.py
@@ -17,7 +17,8 @@
# under the License.
"""Package tvm.script.ir_builder.ir"""
-from ..protocol import MISSING, _at, expression_args
+from ..protocol import MISSING, _at
+from tvm.script.parser.protocol import expr_str_args as expression_args
from .frame import IRModuleFrame
from .ir import (
reserve_function,
diff --git a/python/tvm/script/ir_builder/protocol.py
b/python/tvm/script/ir_builder/protocol.py
index b6247e3760..a22153c176 100644
--- a/python/tvm/script/ir_builder/protocol.py
+++ b/python/tvm/script/ir_builder/protocol.py
@@ -14,20 +14,18 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Shared construction metadata for explicit TVMScript builders.
+"""Concrete runtime helpers for explicit TVMScript builders.
-Metadata describes source syntax. Decorated constructors return concrete values
-or a composable missing type during unresolved eager annotation evaluation.
-Dialects register their own policies here, so translation reads metadata
without
-importing dialects or evaluating annotations.
+Expression-string and function-decorator registration live in parser.protocol.
+This module implements their eager-call adapter, concrete source spans and
+construction-time control helpers. Unresolved eager annotations return a
+composable missing type; handwritten builder calls require concrete arguments.
"""
from builtins import locals as locals
from builtins import slice as slice
from contextlib import contextmanager, nullcontext
-from dataclasses import dataclass
from functools import wraps
-from inspect import signature
from typing import Any, NamedTuple, TypeVar
from tvm import ir
@@ -45,173 +43,77 @@ class _Missing:
MISSING = _Missing()
-class ExprStrPolicy(NamedTuple):
- """Immutable syntax policy for registered constructor arguments.
+def wrap_expression_constructor(constructor, call_signature, policy, *,
as_type=False):
+ """Create the concrete eager-call adapter for parser-owned syntax metadata.
- ``fields`` names parameters whose strings represent expressions.
- ``scalar_strings`` controls bare strings; nested strings are always marked.
- fields is a tuple of parameter-name strings; scalar_strings defaults True.
- introduce (False), dtype (None), implicit_dtype (None), and
- compound_declarations (False) retain legacy registration metadata only.
- They never drive transpiler symbol construction or type validation; those
- operations belong exclusively to builders. Construction returns this tuple
- record and has no builder effects; NamedTuple enforces field arity only.
- Instances live as long as their registered callable, across all
transpilation
- passes, and hold no symbols, eager results, or function-local state.
- """
-
- fields: tuple[str, ...]
- introduce: bool = False
- dtype: Any = None
- scalar_strings: bool = True
- implicit_dtype: Any = None
- compound_declarations: bool = False
-
-
-# Process-wide registry: callable identity -> immutable syntax policy. Dialect
-# imports register once; aliases share identities. No per-function entries or
-# evaluation results are cached, and transpilers only read this table.
-_EXPR_STR_POLICIES = {}
+ constructor is the registered callable; call_signature is its inspected
+ Python signature; policy is immutable metadata with marked field names and
+ scalar_strings. as_type defaults False; True preserves annotation-class and
+ type-union use. Returns a wrapped callable/class with the original
signature.
+ Calling it resolves host TypeVars in active function frames, returns a
+ composable MissingType for unresolved eager annotations, and rejects
strings
+ in handwritten builders. Constructor/type/binding errors propagate.
-
-def expr_str_policy(constructor):
- """Return a callable's ExprStrPolicy, or None when it is unregistered.
-
- ``constructor`` may be any Python value. This read-only lookup neither
- imports dialects nor evaluates the value and has no builder/frame effects.
- """
- try:
- return _EXPR_STR_POLICIES.get(constructor)
- except TypeError:
- return None
-
-
-def expr_str_args(
- *fields,
- introduce=False,
- dtype=None,
- scalar_strings=True,
- implicit_dtype=None,
- compound_declarations=False,
- as_type=False,
-):
- """Register expression-string fields and wrap eager constructor calls.
-
- Parameters
- ----------
- fields : str
- Named parameters whose string values denote source expressions. Unknown
- parameter names raise ValueError when applying the decorator.
- introduce : bool, optional
- Legacy syntax metadata, default False. Builders own symbol
introduction.
- dtype : object, optional
- Legacy explicit symbol-type metadata, default None.
- scalar_strings : bool, optional
- Whether bare strings are expressions, default True. False preserves
- literal shorthand such as ``Tensor("float32")``; nested strings in
- tuples/lists remain expressions.
- implicit_dtype : object, optional
- Legacy implicit-type metadata, default None; newly resolved symbols
- default to int32 in TypeVarFrame.
- compound_declarations : bool, optional
- Legacy expression metadata, default False.
- as_type : bool, optional
- Preserve use as an annotation class (including Python type unions),
- default False. True returns a class whose construction invokes the
- wrapped callable; its instances are still the callable's return values.
-
- Returns
- -------
- decorator : callable
- Registers and wraps a constructor. Unresolved strings or typing.TypeVar
- values in marked fields return a concrete ``ir.MissingType()`` during
- eager Python annotation evaluation. Concrete arguments call the
original
- constructor. Inside an active builder, unresolved inputs raise
TypeError:
- handwritten builders require concrete expressions. Calls are never
- cached and annotations must be safe to re-evaluate. Ordinary
call-binding
- errors propagate from inspect.Signature.bind.
+ This factory owns no registry. Its closure retains only callable/signature/
+ policy configuration for that callable's lifetime, never annotation results
+ or symbols. Each invocation uses the current native builder frame stack.
"""
-
- def decorate(constructor):
- call_signature = signature(constructor)
- unknown = set(fields).difference(call_signature.parameters)
- if unknown:
- raise ValueError(f"Unknown expression argument fields:
{sorted(unknown)}")
- policy = ExprStrPolicy(
- tuple(fields),
- bool(introduce),
- dtype,
- bool(scalar_strings),
- implicit_dtype,
- bool(compound_declarations),
- )
-
- def unresolved(value, nested=False):
- if isinstance(value, str):
- return nested or policy.scalar_strings
- if isinstance(value, TypeVar):
- return True
- if isinstance(value, tuple | list):
- return any(unresolved(item, True) for item in value)
- return False
-
- @wraps(constructor)
- def invoke(*args, **kwargs):
- bound = call_signature.bind(*args, **kwargs)
+ fields = policy.fields
+
+ def unresolved(value, nested=False):
+ if isinstance(value, str):
+ return nested or policy.scalar_strings
+ if isinstance(value, TypeVar):
+ return True
+ if isinstance(value, tuple | list):
+ return any(unresolved(item, True) for item in value)
+ return False
+
+ @wraps(constructor)
+ def invoke(*args, **kwargs):
+ bound = call_signature.bind(*args, **kwargs)
+ if IRBuilder.is_in_scope():
+ # typing.TypeVar is ordinary eager Python metadata. Resolve it
+ # here, never in the syntax-only transpiler.
+ from .type_var_frame import TypeVarFrame
+
+ def resolve(value):
+ if isinstance(value, TypeVar):
+ if value.__bound__ is not None or value.__constraints__:
+ raise TypeError("A symbolic TypeVar cannot have
constraints or a bound")
+ return TypeVarFrame.current().resolve(value.__name__)
+ if isinstance(value, tuple):
+ return tuple(resolve(item) for item in value)
+ if isinstance(value, list):
+ return [resolve(item) for item in value]
+ return value
+
+ for field in fields:
+ if field in bound.arguments:
+ bound.arguments[field] = resolve(bound.arguments[field])
+ if any(unresolved(bound.arguments[field]) for field in fields if field
in bound.arguments):
if IRBuilder.is_in_scope():
- # typing.TypeVar is ordinary eager Python metadata. Resolve it
- # here, never in the syntax-only transpiler.
- from .type_var_frame import TypeVarFrame
-
- def resolve(value):
- if isinstance(value, TypeVar):
- if value.__bound__ is not None or
value.__constraints__:
- raise TypeError("A symbolic TypeVar cannot have
constraints or a bound")
- return TypeVarFrame.current().resolve(value.__name__)
- if isinstance(value, tuple):
- return tuple(resolve(item) for item in value)
- if isinstance(value, list):
- return [resolve(item) for item in value]
- return value
-
- for field in fields:
- if field in bound.arguments:
- bound.arguments[field] =
resolve(bound.arguments[field])
- if any(
- unresolved(bound.arguments[field]) for field in fields if
field in bound.arguments
- ):
- if IRBuilder.is_in_scope():
- raise TypeError(
- "Builder expression arguments require concrete
symbols, not strings"
- )
- return ir.MissingType()
- return constructor(*bound.args, **bound.kwargs)
-
- result = invoke
- if as_type:
- # The class is an annotation surface, not an IR or proxy type.
- # __new__ returns the concrete construction result (or
MissingType).
- result = type(
- constructor.__name__,
- (),
- {
- "__new__": lambda cls, *args, **kwargs: invoke(*args,
**kwargs),
- "__signature__": call_signature,
- "__doc__": constructor.__doc__,
- "__module__": constructor.__module__,
- },
- )
- result.__tvm_expression_args__ = policy
- _EXPR_STR_POLICIES[result] = policy
- _EXPR_STR_POLICIES[constructor] = policy
- return result
-
- return decorate
-
-
-# Compatibility spellings share the registry and wrapper implementation.
-ExpressionArguments = ExprStrPolicy
-expression_args = expr_str_args
+ raise TypeError(
+ "Builder expression arguments require concrete symbols,
not strings"
+ )
+ return ir.MissingType()
+ return constructor(*bound.args, **bound.kwargs)
+
+ result = invoke
+ if as_type:
+ # The class is an annotation surface, not an IR or proxy type.
+ # __new__ returns the concrete construction result (or MissingType).
+ result = type(
+ constructor.__name__,
+ (),
+ {
+ "__new__": lambda cls, *args, **kwargs: invoke(*args,
**kwargs),
+ "__signature__": call_signature,
+ "__doc__": constructor.__doc__,
+ "__module__": constructor.__module__,
+ },
+ )
+ return result
class DeclarationArguments(NamedTuple):
@@ -241,44 +143,6 @@ def register_declaration(constructor, *,
value_parameter="expr", dtype=None):
return constructor
-@dataclass(frozen=True)
-class FunctionKind:
- """Builder namespace and syntax options attached to a function decorator.
-
- builder is the dialect construction namespace; metadata maps string policy
- names to values (for example option_map/defaults). Registration owns this
- record for the decorator's lifetime; transpilers read it across functions.
- No concrete function, frame, annotation result, or symbol is stored here.
- """
-
- builder: Any
- metadata: dict
-
-
-def register_function(decorator, builder, **metadata):
- """Associate a decorator with its builder namespace and syntax policies.
-
- ``decorator`` must support attribute assignment; ``builder`` is its dialect
- construction namespace and ``metadata`` contains named syntax policies.
- Returns the same decorator after replacing its FunctionKind record. No
- construction or annotation evaluation occurs. Unsupported attribute
- assignment raises the ordinary Python error; metadata is retained for the
- decorator's lifetime and must not be mutated by a transpilation pass.
- """
- decorator.__tvm_function_kind__ = FunctionKind(builder, metadata)
- return decorator
-
-
-def function_kind(decorator):
- """Return a decorator's FunctionKind record, or None when unregistered.
-
- ``decorator`` is any already resolved Python value. This lookup neither
- evaluates it nor enters frames; ordinary custom attribute-access errors
- propagate. Returned metadata is shared and must be treated as read-only.
- """
- return getattr(decorator, "__tvm_function_kind__", None)
-
-
def source_span(location):
"""Materialize an IR span from a source location, or preserve an existing
span.
@@ -458,3 +322,23 @@ def callee(builder, value):
if isinstance(value, value_type):
return adapter(value)
return value
+
+
+def __getattr__(name):
+ # Compatibility exports forward to the single parser-owned registry. Lazy
+ # lookup avoids importing parser entry points during builder
initialization.
+ aliases = {"function_kind": "function_info"}
+ exported = {
+ "ExprStrPolicy",
+ "ExpressionArguments",
+ "expr_str_args",
+ "expression_args",
+ "expr_str_policy",
+ "register_function",
+ "function_kind",
+ }
+ if name in exported:
+ from tvm.script.parser import protocol as syntax_protocol
+
+ return getattr(syntax_protocol, aliases.get(name, name))
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/python/tvm/script/parser/__init__.py
b/python/tvm/script/parser/__init__.py
index 6d9941886d..22e8585bb4 100644
--- a/python/tvm/script/parser/__init__.py
+++ b/python/tvm/script/parser/__init__.py
@@ -23,21 +23,30 @@ those of the original source. Entry modules register
policies; this parser never
imports their namespaces.
"""
-from .frontend import (
- _NAMESPACES,
- from_source,
- ir_module,
- make_decorator,
- make_helper,
- parse,
- pyfunc,
- register_namespace,
+import importlib
+import sys
+
+# Process-wide fixed export names, never per-function construction state. Lazy
+# entry loading permits builder imports to register syntax policy without a
+# parser/frontend/builder-construction import cycle.
+_FRONTEND_EXPORTS = (
+ "_NAMESPACES",
+ "from_source",
+ "ir_module",
+ "make_decorator",
+ "make_helper",
+ "parse",
+ "pyfunc",
+ "register_namespace",
)
+__all__ = [name for name in _FRONTEND_EXPORTS if not name.startswith("_")]
def __getattr__(name):
- # Entry modules register source aliases; parsing never imports their
owners.
- try:
- return _NAMESPACES[name]
- except KeyError:
- raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
from None
+ if name in _FRONTEND_EXPORTS:
+ frontend = importlib.import_module(f"{__name__}.frontend")
+ return getattr(frontend, name)
+ frontend = sys.modules.get(f"{__name__}.frontend")
+ if frontend is not None and name in getattr(frontend, "_NAMESPACES", {}):
+ return frontend._NAMESPACES[name]
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/python/tvm/script/parser/expression.py
b/python/tvm/script/parser/expression.py
index c8b31c0669..0efc15d966 100644
--- a/python/tvm/script/parser/expression.py
+++ b/python/tvm/script/parser/expression.py
@@ -29,7 +29,7 @@ import inspect
import linecache
import re
-from tvm.script.ir_builder import protocol
+from . import protocol
class _LiteralParser:
diff --git a/python/tvm/script/parser/frontend.py
b/python/tvm/script/parser/frontend.py
index c613835855..ea6516cff7 100644
--- a/python/tvm/script/parser/frontend.py
+++ b/python/tvm/script/parser/frontend.py
@@ -30,6 +30,7 @@ from typing import TypeVar
from tvm.error import DiagnosticError
from tvm.script.ir_builder import construction, protocol
+from . import protocol as syntax_protocol
from .diagnostics import diagnostic_error
from .transpile import IRBuilderTranspiler, NameCollector
@@ -117,7 +118,7 @@ def make_decorator(builder, *, option_map=None,
defaults=None):
raise ValueError("Construction decorators require a function or
keyword options")
def apply(function):
- function.__tvm_function_kind__ = decorator.__tvm_function_kind__
+ function.__tvm_function_info__ = decorator.__tvm_function_info__
function.__tvm_function_options__ = options
if _inside_class(function):
return function
@@ -125,7 +126,7 @@ def make_decorator(builder, *, option_map=None,
defaults=None):
return apply(function) if function is not None else apply
- return protocol.register_function(
+ return syntax_protocol.register_function(
decorator, builder, option_map=mapping, defaults=default_options
)
@@ -191,7 +192,7 @@ def pyfunc(function):
return function
-protocol.register_function(pyfunc, None, python=True)
+syntax_protocol.register_function(pyfunc, None, python=True)
class Compiler:
@@ -383,7 +384,10 @@ class Compiler:
exec(
compile(ast.Module(imports, []), self.filename, "exec",
dont_inherit=True), self.env
)
- if inspect.isfunction(self.original) and
protocol.function_kind(self.original) is not None:
+ if (
+ inspect.isfunction(self.original)
+ and syntax_protocol.function_info(self.original) is not None
+ ):
# Direct decorator application (T.prim_func(host_function)) has no
# decorator in source AST. Inject only its registered host
metadata;
# already evaluated option values stay opaque execution bindings.
diff --git a/python/tvm/script/parser/protocol.py
b/python/tvm/script/parser/protocol.py
new file mode 100644
index 0000000000..82f38d9eb9
--- /dev/null
+++ b/python/tvm/script/parser/protocol.py
@@ -0,0 +1,192 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Parser-owned syntax policies and function-decorator registration.
+
+This module stores only host callable identities and flat syntax metadata.
+Concrete eager annotation behavior is delegated to a builder-owned adapter;
+no IR definition, concrete annotation result, symbol or frame is stored here.
+"""
+
+from inspect import signature
+from typing import Any, NamedTuple
+
+
+class ExprStrPolicy(NamedTuple):
+ """Immutable syntax policy for registered constructor arguments.
+
+ ``fields`` names parameters whose strings represent expressions.
+ ``scalar_strings`` controls bare strings; nested strings are always marked.
+ fields is a tuple of parameter-name strings; scalar_strings defaults True.
+ introduce (False), dtype (None), implicit_dtype (None), and
+ compound_declarations (False) retain legacy registration metadata only.
+ They never drive transpiler symbol construction or type validation; those
+ operations belong exclusively to builders. Construction returns this tuple
+ record and has no builder effects; NamedTuple enforces field arity only.
+ Instances live as long as their registered callable, across all
transpilation
+ passes, and hold no symbols, eager results, or function-local state.
+ """
+
+ fields: tuple[str, ...]
+ introduce: bool = False
+ dtype: Any = None
+ scalar_strings: bool = True
+ implicit_dtype: Any = None
+ compound_declarations: bool = False
+
+
+# Process-wide registry: callable identity -> immutable syntax policy. Dialect
+# imports register once; aliases share identities. No per-function entries or
+# evaluation results are cached, and transpilers only read this table.
+_EXPR_STR_POLICIES = {}
+
+
+def expr_str_policy(constructor):
+ """Return a callable's ExprStrPolicy, or None when it is unregistered.
+
+ ``constructor`` may be any Python value. This read-only lookup neither
+ imports dialects nor evaluates the value and has no builder/frame effects.
+ """
+ try:
+ return _EXPR_STR_POLICIES.get(constructor)
+ except TypeError:
+ return None
+
+
+def expr_str_args(
+ *fields,
+ introduce=False,
+ dtype=None,
+ scalar_strings=True,
+ implicit_dtype=None,
+ compound_declarations=False,
+ as_type=False,
+):
+ """Register expression-string fields and wrap eager constructor calls.
+
+ Parameters
+ ----------
+ fields : str
+ Named parameters whose string values denote source expressions. Unknown
+ parameter names raise ValueError when applying the decorator.
+ introduce : bool, optional
+ Legacy syntax metadata, default False. Builders own symbol
introduction.
+ dtype : object, optional
+ Legacy explicit symbol-type metadata, default None.
+ scalar_strings : bool, optional
+ Whether bare strings are expressions, default True. False preserves
+ literal shorthand such as ``Tensor("float32")``; nested strings in
+ tuples/lists remain expressions.
+ implicit_dtype : object, optional
+ Legacy implicit-type metadata, default None; newly resolved symbols
+ default to int32 in TypeVarFrame.
+ compound_declarations : bool, optional
+ Legacy expression metadata, default False.
+ as_type : bool, optional
+ Preserve use as an annotation class (including Python type unions),
+ default False. True returns a class whose construction invokes the
+ wrapped callable; its instances are still the callable's return values.
+
+ Returns
+ -------
+ decorator : callable
+ Registers and wraps a constructor. Unresolved strings or typing.TypeVar
+ values in marked fields return a concrete ``ir.MissingType()`` during
+ eager Python annotation evaluation. Concrete arguments call the
original
+ constructor. Inside an active builder, unresolved inputs raise
TypeError:
+ handwritten builders require concrete expressions. Calls are never
+ cached and annotations must be safe to re-evaluate. Ordinary
call-binding
+ errors propagate from inspect.Signature.bind.
+ """
+
+ def decorate(constructor):
+ call_signature = signature(constructor)
+ unknown = set(fields).difference(call_signature.parameters)
+ if unknown:
+ raise ValueError(f"Unknown expression argument fields:
{sorted(unknown)}")
+ policy = ExprStrPolicy(
+ tuple(fields),
+ bool(introduce),
+ dtype,
+ bool(scalar_strings),
+ implicit_dtype,
+ bool(compound_declarations),
+ )
+
+ # Registration is syntax-only. Builders own eager construction and the
+ # active-frame/MissingType decisions behind this generic wrapper
factory.
+ from tvm.script.ir_builder.protocol import wrap_expression_constructor
+
+ result = wrap_expression_constructor(constructor, call_signature,
policy, as_type=as_type)
+ result.__tvm_expression_args__ = policy
+ _EXPR_STR_POLICIES[result] = policy
+ _EXPR_STR_POLICIES[constructor] = policy
+ return result
+
+ return decorate
+
+
+# Compatibility spellings share the registry and wrapper implementation.
+ExpressionArguments = ExprStrPolicy
+expression_args = expr_str_args
+
+
+class FunctionDecoratorInfo(NamedTuple):
+ """Flat syntax registration for a source function decorator.
+
+ builder is the opaque construction namespace, or None for ordinary Python.
+ option_map optionally maps public option names to builder option names;
+ defaults optionally maps builder option names to default Python constants.
+ Both default None (empty) and are copied by register_function. python
defaults
+ False; True preserves the source function body as ordinary Python. These
are
+ the complete supported fields: no generic metadata dictionary is retained.
+ Records live with their decorators across compilations; mappings are
read-only
+ to consumers and contain no frame, function result, annotation or symbol
state.
+ NamedTuple enforces field arity; construction enters no builder frames.
+ """
+
+ builder: Any
+ option_map: dict | None = None
+ defaults: dict | None = None
+ python: bool = False
+
+
+def register_function(decorator, builder, *, option_map=None, defaults=None,
python=False):
+ """Register and return a decorator with explicit supported syntax options.
+
+ decorator must allow attribute assignment. builder is an opaque
construction
+ namespace, or None with python=True. option_map/defaults are optional
mappings
+ copied into the flat FunctionDecoratorInfo; omitted mappings mean empty.
+ python defaults False and selects ordinary host-function preservation when
+ True. Registration replaces __tvm_function_info__ for the callable's
lifetime.
+ No constructor or annotation executes and no frame is entered. Invalid
+ mappings/attribute writes raise ordinary Python errors; unknown keywords
are
+ rejected by this explicit signature rather than retained as arbitrary
policy.
+ """
+ decorator.__tvm_function_info__ = FunctionDecoratorInfo(
+ builder, dict(option_map or {}), dict(defaults or {}), bool(python)
+ )
+ return decorator
+
+
+def function_info(decorator):
+ """Return FunctionDecoratorInfo for a resolved callable, or None.
+
+ decorator may be any already-resolved host value. This read-only lookup
+ neither evaluates it nor enters frames. Custom attribute errors propagate;
+ returned option mappings are shared registration state and must not be
edited.
+ """
+ return getattr(decorator, "__tvm_function_info__", None)
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index ea15baea82..b90adb965c 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -26,8 +26,7 @@ import builtins
import copy
import inspect
-from tvm.script.ir_builder import protocol
-
+from . import protocol
from .expression import rewrite_expression
@@ -884,7 +883,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# Ordinary construction helpers retain normal Python execution.
Registered
# function-kind entry points are resolved by the enclosing compiler
callback.
if any(
- getattr(self._resolve(decorator), "__tvm_function_kind__", None)
+ getattr(self._resolve(decorator), "__tvm_function_info__", None)
for decorator in node.decorator_list
):
self._error(node, "A registered nested function requires a
function compiler")
@@ -903,18 +902,16 @@ class IRBuilderTranspiler(ast.NodeTransformer):
"""Read registered decorator metadata without evaluating annotations.
node is an original FunctionDef; allow_python permits unregistered
nested
- Python helpers. Returns (FunctionKind, options AST). Raises SyntaxError
+ Python helpers. Returns (FunctionDecoratorInfo, options AST). Raises
SyntaxError
for an undecorated IR function or positional construction options.
"""
for decorator in node.decorator_list:
target = decorator.func if isinstance(decorator, ast.Call) else
decorator
- kind = protocol.function_kind(self._resolve(target))
+ kind = protocol.function_info(self._resolve(target))
if kind is None:
continue
- mapping = kind.metadata.get("option_map", {})
- values = {
- key: ast.Constant(value) for key, value in
kind.metadata.get("defaults", {}).items()
- }
+ mapping = kind.option_map or {}
+ values = {key: ast.Constant(value) for key, value in
(kind.defaults or {}).items()}
expansions = []
if isinstance(decorator, ast.Call):
if decorator.args:
@@ -930,7 +927,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
)
return kind, ast.copy_location(result, node)
if allow_python:
- return protocol.FunctionKind(None, {"python": True}), ast.Dict([],
[])
+ return protocol.FunctionDecoratorInfo(None, python=True),
ast.Dict([], [])
self._error(node, f"Function {node.name!r} has no registered
construction kind")
def function_program(self, node, runtime, bindings, *, local=False):
@@ -950,7 +947,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# A factory gives each function an ordinary Python closure: parameters
and
# symbols from sibling signatures never overwrite one another's
captures.
kind, options = self.function_metadata(node)
- if kind.metadata.get("python"):
+ if kind.python:
return None
builder = self.fresh("_t")
bindings[builder] = kind.builder
@@ -1099,7 +1096,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
def nested(child):
nested_kind, _ = self.function_metadata(child, allow_python=True)
- if nested_kind.metadata.get("python"):
+ if nested_kind.python:
return copy.deepcopy(child)
statements, child_record, child_body = self.function_program(
child, runtime, bindings, local=True
@@ -1287,7 +1284,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
records = []
for function in functions:
kind, _ = self.function_metadata(function)
- if kind.metadata.get("python"):
+ if kind.python:
# Pattern: @I.pyfunc def f(...): body -> unchanged host
function
# plus opaque builder registration.
host = copy.deepcopy(function)
diff --git a/python/tvm/tirx/script/builder/__init__.py
b/python/tvm/tirx/script/builder/__init__.py
index 57b8c8b98a..ae63a3535d 100644
--- a/python/tvm/tirx/script/builder/__init__.py
+++ b/python/tvm/tirx/script/builder/__init__.py
@@ -32,7 +32,6 @@ from tvm.script.ir_builder import ir as _I
from tvm.script.ir_builder.base import IRBuilderFrame as _NativeFrame
from tvm.script.ir_builder.protocol import MISSING as _MISSING
from tvm.script.ir_builder.protocol import at as _at
-from tvm.script.ir_builder.protocol import expression_args as _expression_args
from tvm.script.ir_builder.protocol import register_call_kind as
_register_call_kind
from tvm.script.ir_builder.protocol import register_declaration as
_register_declaration
from tvm.script.ir_builder.protocol import source_span as _source_span
@@ -40,6 +39,7 @@ from tvm.script.ir_builder.protocol import span_context as
_span_context
from tvm.script.ir_builder.type_var_frame import TypeVarDecl as _TypeVarDecl
from tvm.script.ir_builder.type_var_frame import TypeVarFrame as _TypeVarFrame
from tvm.script.ir_builder.type_var_frame import resolve_type_var
+from tvm.script.parser.protocol import expr_str_args as _expression_args
from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool
from . import _ffi_api
diff --git a/tests/python/tvmscript/test_parser.py
b/tests/python/tvmscript/test_parser.py
index 61195c72f2..c1ca3e4207 100644
--- a/tests/python/tvmscript/test_parser.py
+++ b/tests/python/tvmscript/test_parser.py
@@ -33,6 +33,7 @@ from tvm import ir
from tvm.ir.prim import Cast
from tvm.script import parser
from tvm.script.ir_builder import IRBuilder, protocol
+from tvm.script.parser import protocol as syntax_protocol
from tvm.script.parser.frontend import Compiler, make_decorator
from tvm.script.parser.transpile import IRBuilderTranspiler
@@ -461,3 +462,35 @@ def test_shared_allocator_preserves_source_names():
transformer = compiler.transformer()
assert transformer.fresh_unique_name("value") != "value0"
assert transformer.fresh_unique_name() not in {"_t0", "_t1"}
+
+
+def test_parser_metadata_registration_boundary():
+ options, defaults = {"private": "local"}, {"local": False}
+
+ def decorator():
+ pass
+
+ syntax_protocol.register_function(decorator, _Recorder(),
option_map=options, defaults=defaults)
+ info = syntax_protocol.function_info(decorator)
+ assert isinstance(info, tuple)
+ assert info._fields == ("builder", "option_map", "defaults", "python")
+ assert not hasattr(info, "metadata")
+ options["private"] = "changed"
+ assert info.option_map == {"private": "local"}
+ with pytest.raises(TypeError, match="unexpected keyword"):
+ syntax_protocol.register_function(decorator, None, unknown=True)
+
+ @syntax_protocol.expr_str_args("shape")
+ def constructor(shape):
+ return shape
+
+ policy = syntax_protocol.expr_str_policy(constructor)
+ assert isinstance(policy, tuple)
+ assert policy.fields == ("shape",)
+ assert policy.__class__.__module__ == "tvm.script.parser.protocol"
+ assert constructor((4,)) == (4,)
+ assert constructor(("n",)).is_missing()
+ # Existing metadata imports forward to the same parser-owned registration.
+ assert protocol.expr_str_policy(constructor) is policy
+ with IRBuilder(), pytest.raises(TypeError, match="concrete symbols"):
+ constructor(("n",))
