You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/09/23 18:36:28 UTC
(tvm) 04/04: [Script] Keep syntax registration in the parser registry
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch script/canonical-parser-df
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit eee9ba3f29c23ed3c7fe3bc82d3d8076452af435
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Sep 23 18:35:57 2026 +0000
[Script] Keep syntax registration in the parser registry
---
docs/arch/tvmscript.rst | 12 +-
python/tvm/relax/script/builder/__init__.py | 4 +-
python/tvm/script/ir_builder/ir/__init__.py | 18 +-
python/tvm/script/ir_builder/ir/ir.py | 2 +-
python/tvm/script/ir_builder/ir/parser_protocol.py | 824 +++----------------
python/tvm/script/parser/call_args_policy.py | 2 +-
python/tvm/script/parser/entry.py | 4 +-
python/tvm/script/parser/prescan.py | 9 +-
python/tvm/script/parser/protocol_registry.py | 868 +++++++++++++++++++++
python/tvm/script/parser/transpile.py | 3 +-
python/tvm/tirx/script/builder/__init__.py | 6 +-
python/tvm/tirx/script/builder/ir.py | 2 +-
python/tvm/tirx/script/builder/parser_protocol.py | 2 +-
python/tvm/tirx/script/jit.py | 4 +-
tests/python/script/dummy_builder.py | 12 +-
.../script/test_builder_protocol_contract.py | 156 +++-
tests/python/script/test_parser.py | 32 +-
tests/python/script/test_parser_latest_design.py | 14 +-
.../python/script/test_parser_shared_resolvers.py | 4 +-
.../python/tirx/script/test_parser_buffer_views.py | 6 +-
.../tirx/script/test_parser_tirx_resolvers.py | 4 +-
.../script/test_tirx_builder_protocol_contract.py | 41 +-
.../tirx/transform/test_transform_lower_tirx.py | 2 +-
23 files changed, 1210 insertions(+), 821 deletions(-)
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index df5d045485..4a8655637e 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -172,7 +172,7 @@ then lowers to operations supplied by that namespace: for example, unmarked ``if
TIR control-flow frames in a primitive function and Relax control-flow frames in a Relax
function. The transpiler emits these operations without implementing dialect IR semantics.
-``ir_builder/ir/parser_protocol.py`` records callable syntax metadata. Argument policies such
+``parser/protocol_registry.py`` owns callable syntax metadata and registration. Argument policies such
as ``expr_str`` translate symbolic strings written directly in source expressions, while ``global_info``
preserves a module reference for builder-side lookup. Dtype and placement strings remain
literal. Captured or computed symbolic shapes must already contain explicit IR variables;
@@ -246,9 +246,13 @@ Generated programs use the shared ``I.at_`` and ``I.with_at_group_`` helpers to
locations and preserve caller/definition provenance while evaluating source calls once.
Their implementation lives in ``ir_builder.base``. These helpers retain returned object
identity; builders handle results that already emitted a statement or introduced a binding.
-The shared ``ir_builder.ir.parser_protocol`` module owns syntax registration, metadata
-and the documented generated builder contract. Dialect ``parser_protocol`` modules
-register their syntax and implement their construction hooks.
+The parser-owned ``parser.protocol_registry`` module defines registration APIs, metadata and
+persistent registration state. Dialect ``parser_protocol`` modules import these APIs
+directly, register their syntax and implement their construction hooks. The shared
+``ir_builder.ir.parser_protocol`` file-level documentation explains the complete
+registration and generated-builder contract; that module implements builder hooks
+without re-exporting registration APIs. For example, dialect registration imports
+``args_policy`` and ``direct_call`` from ``tvm.script.parser.protocol_registry``.
Calls registered with ``direct_call`` keep their ordinary results without automatic
binding, emission or result-span attachment. Their arguments still undergo normal
diff --git a/python/tvm/relax/script/builder/__init__.py b/python/tvm/relax/script/builder/__init__.py
index 826e78d9e3..76df38ebc1 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/relax/script/builder/__init__.py
@@ -30,9 +30,9 @@ from tvm.relax.distributed import device_mesh as device_mesh
from tvm.script.ir_builder import IRBuilder as _IRBuilder
from tvm.script.ir_builder.base import at as _at
from tvm.script.ir_builder.base import source_span as _source_span
-from tvm.script.ir_builder.ir.parser_protocol import args_policy as _args_policy
-from tvm.script.ir_builder.ir.parser_protocol import constexpr as constexpr
from tvm.script.ir_builder.ir.parser_protocol import resolve_global_info_ as _lookup_global_info
+from tvm.script.parser.protocol_registry import args_policy as _args_policy
+from tvm.script.parser.protocol_registry import constexpr as constexpr
from . import distributed as dist
from . import ir as _native
diff --git a/python/tvm/script/ir_builder/ir/__init__.py b/python/tvm/script/ir_builder/ir/__init__.py
index 7f768a3f8d..1087cca36e 100644
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ b/python/tvm/script/ir_builder/ir/__init__.py
@@ -18,23 +18,7 @@
"""Package tvm.script.ir_builder.ir"""
from tvm.ir import GenericConst, Range, StringImm, StringType
-from .parser_protocol import (
- args_policy,
- constexpr,
- direct_call,
- function_info,
- get_args_policy,
- get_result_members,
- is_direct_call,
- is_mutable_var_decl,
- is_scope_var_query_or_decl,
- register_binding_decl,
- register_function,
- register_mutable_var_decl,
- register_result_members,
- register_scope_var_query_or_decl,
- register_type_var_decl,
-)
+from tvm.script.parser.protocol_registry import constexpr
from ..base import IRBuilder, MISSING, annotation_value_, at_, with_at_group_, require_defined
from .parser_protocol import check_well_formed_, module_member_, resolve_global_info_
diff --git a/python/tvm/script/ir_builder/ir/ir.py b/python/tvm/script/ir_builder/ir/ir.py
index ef73ca6551..e30119e5db 100644
--- a/python/tvm/script/ir_builder/ir/ir.py
+++ b/python/tvm/script/ir_builder/ir/ir.py
@@ -21,11 +21,11 @@ from typing import TYPE_CHECKING, TypeVar
from tvm.ir import BaseFunc, GlobalInfo, GlobalVar
from tvm.runtime import Object as tvm_Object
+from tvm.script.parser.protocol_registry import direct_call
from ..base import IRBuilder
from . import _ffi_api
from .frame import IRModuleFrame
-from .parser_protocol import direct_call
if TYPE_CHECKING:
from tvm.relax import DummyGlobalInfo, VDevice
diff --git a/python/tvm/script/ir_builder/ir/parser_protocol.py b/python/tvm/script/ir_builder/ir/parser_protocol.py
index a680c9ca4d..ed4a233b07 100644
--- a/python/tvm/script/ir_builder/ir/parser_protocol.py
+++ b/python/tvm/script/ir_builder/ir/parser_protocol.py
@@ -25,19 +25,118 @@ references and results. Ordinary functions can use one ``X.function()`` entry.
When forward references are needed, ``X.function(decl=True)`` reserves each
signature before body entry, allowing sibling calls without another ownership
record. Re-entering that same frame completes its body.
-This module also owns syntax metadata and registration: args_policy selects
-argument rewriting, function registration selects a builder namespace, and
-registered declaration categories choose assignment policy without inspecting IR.
-Registries hold callable identities and flat syntax facts only, never native
-frames, constructed values or per-parse environments. The parser consumes this
-contract; it is not imported here. Module setup precedes all signatures, which
-precede all bodies. Completed results are validated only after frame exit.
+Syntax registration is owned by ``tvm.script.parser.protocol_registry``. Dialects import
+its APIs directly; this file documents the complete customization contract but
+does not import or re-export those registration APIs. Registration state lives
+across parses, while PrescanContext, ModuleContext and FunctionContext only
+consume it. Callable markers stay on the registered callables. The argument
+policy table shares one immutable policy between original and adapted callable
+identities. No registry stores active frames, construction results or per-parse
+environments. Module setup precedes all signatures, which precede all bodies.
+Completed results are validated only after frame exit.
+
+A dialect selects source syntax through these registration APIs:
+
+* ``args_policy(fields, scalar_strings=True, dtype=None, as_type=False)`` marks
+ individual parameters as ``expr_str`` or ``global_info``. Positional and
+ keyword arguments use the same signature-based policy. Expression strings
+ lower through ``X.resolve_type_var_``; global-info strings use
+ ``I.resolve_global_info_``. Unmarked strings remain literal. The optional
+ opaque dtype reaches symbol resolution unchanged. Concrete arguments still
+ execute the constructor; eager expression/type adapters retain the builder's
+ MissingType and active-function behavior. Unknown policies or parameter names
+ are rejected at registration.
+* ``register_type_var_decl`` marks an omitted-value scalar declaration and its
+ optional dtype, selecting ``X.resolve_type_var_``. ``register_binding_decl``
+ marks an explicit ordinary binding that shadows outer mutable storage and
+ still uses ``X.bind_``. ``register_mutable_var_decl`` advertises call,
+ annotation or parameter positions for ``X.decl_mutable_var_``; later stores
+ use ``X.set_mutable_var_``.
+* ``register_scope_var_query_or_decl`` marks calls whose producer already owns
+ the returned variables. This declaration category precedes mutable stores;
+ ``X.scope_var_query_or_decl_`` preserves identity and performs dialect naming
+ and validation without another IR binding.
+* ``direct_call`` returns the registered callable unchanged. Source calls omit
+ automatic binding, emission, result-span attachment and the outer source-call
+ wrapper. Callee and arguments still evaluate once in order, with normal child
+ rewriting. The callable owns its effects. This differs from AlreadyEmitted,
+ whose wrapped result retains normal source-result handling.
+* ``register_result_members`` associates a producer or annotation with an
+ opaque namespace of actual member callables/descriptors. Prescan facts and
+ chained producer syntax can expose those members without evaluating a value
+ or inspecting its IR type. Each member still needs its own declaration or
+ direct-call registration; a method name alone never selects a policy.
+* ``register_function`` attaches an opaque builder namespace, copied option
+ mapping, defaults and an optional Python-body flag to a source decorator.
+ ``function_info`` reads that ``FunctionDecoratorInfo`` record. The flag keeps
+ original Python callables in module ``__pyfuncs__``; other functions lower
+ through the registered dialect's function/signature/body hooks.
+
+``get_args_policy`` reads ``ArgsPolicy`` and its ``ExprStrPolicy`` record;
+``get_type_var_decl`` reads ``DeclarationArguments``. ``is_binding_decl``,
+``is_mutable_var_decl``, ``is_scope_var_query_or_decl`` and ``is_direct_call``
+read declaration/call markers. ``copy_function_info`` shares the exact registered
+decorator record with a source function, leaving definition scope and applied
+options to the entry point. ``get_result_members`` reads member namespaces,
+normalizing bound methods and property getters. Argument, scope-query and
+direct-call readers also recognize Python bound-method identities where
+applicable. ``DeclarationArguments`` carries omitted-value parameter and dtype
+facts. These readers neither call constructors nor own native state. The shared
+``constexpr`` marker is recognized by identity for host control expressions and
+specialization annotations; dialect semantic exports refer to that same marker.
+
+For example, a dialect registers producers directly with the parser owner:
+
+.. code:: python
+
+ from tvm.script.parser.protocol_registry import (
+ args_policy,
+ direct_call,
+ register_binding_decl,
+ register_function,
+ register_mutable_var_decl,
+ register_result_members,
+ register_scope_var_query_or_decl,
+ register_type_var_decl,
+ )
+
+ register_type_var_decl(X.int32, dtype="int32")
+ register_mutable_var_decl(X.local_scalar, syntax="call")
+ register_binding_decl(X.alloc_buffer)
+ register_scope_var_query_or_decl(X.bind)
+ register_function(prim_func, X, option_map={"private": "private"})
+ register_result_members(X.Buffer, BufferMethods)
+ direct_call(BufferMethods.permute)
+
+ @args_policy({"shape": "expr_str", "device": "global_info"})
+ def tensor(shape, device=None):
+ return make_tensor(shape, device)
+
+ @direct_call
+ def identity(value):
+ return value
+
+ # Source
+ value = tensor(("n",), device="cuda:0")
+ kept = identity(value)
+ # Generated builder
+ value = X.bind_(
+ tensor((X.resolve_type_var_("n"),), device=I.resolve_global_info_("cuda:0")),
+ name="value",
+ )
+ kept = identity(value)
+
+Importing the registration module does not initialize parser entry points or
+concrete builders. Applying an expression/type argument policy requests the
+shared eager annotation adapter when needed; native frames still own all IR
+construction, symbol identity and validation. Definitions and detailed API
+docstrings remain beside the parser-owned implementation.
Shared source support remains in ir_builder.base and is exported through I:
I.at_(location, value) annotates that same object (or AlreadyEmitted.value),
records a returned frame's deferred location, and preserves identity;
I.with_at_group_(location, thunk) evaluates one source call inside restored
-caller/definition provenance and applies at_ before returning. I.annotation_value_
+caller/definition provenance and applies ``at_`` before returning. ``I.annotation_value_``
converts a rewritten annotation without creating a parser owner. Locations are
Span objects or (SourceName, start_line, start_column, end_line, end_column)
tuples; None omits explicit attribution. Track-span disabling omits source
@@ -105,10 +204,8 @@ They specify a structural contract, without a base class or runtime dispatcher.
from __future__ import annotations
import re
-from collections.abc import Callable, Mapping, Sequence
-from inspect import signature
-from types import MappingProxyType, MethodType
-from typing import Any, NamedTuple, NoReturn, TypeVar
+from collections.abc import Sequence
+from typing import Any
from tvm import ir as _ir
@@ -539,628 +636,6 @@ def check_well_formed_(module: _ir.IRModule) -> None:
raise ValueError(f"{message}\n{error}") from error
-_Callable = TypeVar("_Callable", bound=Callable[..., Any])
-
-
-def constexpr(value: object) -> NoReturn:
- """Mark a host control expression or a JIT specialization annotation.
-
- Parameters
- ----------
- value : object
- Source operand evaluated as ordinary Python after rewriting; never passed to
- this marker at runtime.
-
- Returns
- -------
- NoReturn
- The marker is removed from generated execution; direct Python invocation raises
- TypeError.
-
- Notes
- -----
- No builder context or IR effects. Recognition uses callable identity, including aliases,
- rather than name spelling or return values. The same marker can annotate a
- specialization parameter. Host operators/control flow retain Python behavior while their
- source calls still receive ordinary source-call instrumentation.
-
- .. code:: python
-
- # Source
- if I.constexpr(enabled):
- T.evaluate(1)
- # Generated builder
- if enabled:
- X.emit_(X.evaluate(1))
- """
- raise TypeError("constexpr is a parser syntax marker, not a runtime operation")
-
-
-class ExprStrPolicy(NamedTuple):
- """Immutable syntax policy for registered constructor arguments.
-
- Parameters
- ----------
- fields : tuple of str
- Parameter names whose string values represent expressions.
- dtype : object, optional
- Dtype spelling passed to builder symbol resolution. Default is None.
- scalar_strings : bool, optional
- Interpret bare strings as expressions. Default is True. Nested
- strings in marked fields are always treated as expressions.
-
- Notes
- -----
- This tuple record lives with its registered callable across transpilation
- passes. It stores no symbols, eager results, or function-local state.
- Dtype spellings are forwarded as builder arguments; the transpiler never
- constructs or validates symbols. Builders retain already-declared symbol
- identity and type. Construction enters no frame.
- """
-
- fields: tuple[str, ...]
- dtype: object = None
- scalar_strings: bool = True
-
-
-# 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.
-_ARGS_POLICIES: dict[object, ArgsPolicy] = {}
-
-
-class ArgsPolicy(NamedTuple):
- """Immutable per-parameter policies shared by a callable and its adapter.
-
- Parameters
- ----------
- fields : Mapping[str, str]
- Immutable parameter-to-policy mapping containing expr_str/global_info.
- expression : ExprStrPolicy
- Expression fields, scalar-string behavior and opaque symbolic dtype.
-
- Notes
- -----
- This process-wide syntax record stores no evaluated arguments, scopes,
- frames or results. Registration copies its mappings; parses only read it.
- """
-
- fields: Mapping[str, str]
- expression: ExprStrPolicy
-
-
-def get_args_policy(constructor: object) -> ArgsPolicy | None:
- """Read argument policy metadata without evaluating a constructor.
-
- Parameters
- ----------
- constructor : object
- Resolved host callable or bound method; bound methods use their underlying
- function identity.
-
- Returns
- -------
- ArgsPolicy | None
- Shared immutable policy, or None for unregistered/unhashable values.
-
- Notes
- -----
- No builder context, IR construction or source-span changes occur. The returned policy
- belongs to process-wide registration, not a particular parse. Custom hashing errors
- other than TypeError propagate.
-
- .. code:: python
-
- # Source
- tensor(("n",))
- # Rewrite-time lookup; generated code receives decoded arguments
- policy = I.get_args_policy(tensor)
- tensor((X.resolve_type_var_("n"),))
- """
- if isinstance(constructor, MethodType):
- constructor = constructor.__func__
- try:
- return _ARGS_POLICIES.get(constructor)
- except TypeError:
- return None
-
-
-def args_policy(
- fields: Mapping[str, str],
- *,
- scalar_strings: bool = True,
- dtype: object = None,
- as_type: bool = False,
-) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
- """Register per-argument source policies without storing construction results.
-
- Parameters
- ----------
- fields : mapping of str to str
- Parameter names mapped to expr_str or global_info. Copied into immutable
- registration data; positional and keyword calls use the same policy.
- scalar_strings : bool, optional
- True by default. False preserves a bare string in an expression field; nested
- strings are still expression syntax.
- dtype : object, optional
- Opaque dtype passed to symbol resolution for expression strings. None (default)
- leaves the builder default.
- as_type : bool, optional
- False by default. True preserves the annotation-class surface, including Python
- type unions, while construction returns the underlying result.
-
- Returns
- -------
- Callable[[Callable[..., Any]], Callable[..., Any]]
- Decorator returning the callable or its eager annotation adapter. Original and
- adapted identities share one argument policy.
-
- Notes
- -----
- Registration enters no frame, constructs no IR and retains no per-function values.
- Concrete arguments call the original constructor. Outside construction, unresolved
- expression annotations yield MissingType; inside construction, unresolved expression
- strings raise TypeError and TypeVar values use the nearest function resolver. Unknown
- policy kinds/parameter names raise ValueError; signature/argument-binding errors
- propagate. The adapter belongs to the builder, not parser execution.
-
- .. code:: python
-
- # Source
- @I.args_policy({"shape": "expr_str", "device": "global_info"})
- def tensor(shape, device=None):
- ...
- # Source
- tensor(("n",), device="cuda:0")
- # Generated builder
- tensor((X.resolve_type_var_("n"),), device=I.resolve_global_info_("cuda:0"))
- """
- fields = dict(fields)
- unsupported = set(fields.values()).difference(("expr_str", "global_info"))
- if unsupported:
- raise ValueError(f"Unknown argument policies: {sorted(unsupported)}")
-
- def decorate(constructor: Callable[..., Any]) -> Callable[..., Any]:
- call_signature = signature(constructor)
- unknown = set(fields).difference(call_signature.parameters)
- if unknown:
- raise ValueError(f"Unknown argument policy fields: {sorted(unknown)}")
- expression_fields = tuple(name for name, kind in fields.items() if kind == "expr_str")
- expression = ExprStrPolicy(expression_fields, dtype, bool(scalar_strings))
- if expression_fields or as_type:
- # Builders own eager construction and active-frame/MissingType decisions.
- from tvm.script.ir_builder.base import wrap_expression_constructor
-
- result = wrap_expression_constructor(
- constructor, call_signature, expression, as_type=as_type
- )
- else:
- result = constructor
- policy = ArgsPolicy(MappingProxyType(fields.copy()), expression)
- _ARGS_POLICIES[constructor] = policy
- _ARGS_POLICIES[result] = policy
- return result
-
- return decorate
-
-
-class DeclarationArguments(NamedTuple):
- """Immutable metadata for scalar-constructor predeclarations.
-
- Parameters
- ----------
- value_parameter : str
- Optional value parameter whose absence denotes a declaration.
- dtype : object, optional
- Explicit primitive dtype metadata. Default is None.
-
- Notes
- -----
- `register_type_var_decl` attaches this record to its callable. Aliases
- and transpilation passes share it for the callable's lifetime; it stores
- no symbols or construction results and enters no builder frame.
- """
-
- value_parameter: str
- dtype: object = None
-
-
-def register_type_var_decl(
- constructor: _Callable, *, value_parameter: str = "expr", dtype: object = None
-) -> _Callable:
- """Mark a constructor whose omitted value argument declares a symbolic variable.
-
- Parameters
- ----------
- constructor : callable
- Constructor supporting syntax attribute assignment. Identity is retained.
- value_parameter : str, optional
- Value parameter whose absence means declaration; defaults to expr.
- dtype : object, optional
- Opaque primitive dtype, default None. A string dtype also allows prescan to
- recognize direct zero-argument body declarations.
-
- Returns
- -------
- _Callable
- The original callable with DeclarationArguments metadata.
-
- Notes
- -----
- No constructor executes and no frame is entered. Syntax predeclaration precedes
- dependent signature shapes while original body ordering remains. Only direct declaration
- forms are collected, not nested/effectful expressions. Aliases share callable metadata;
- unsupported attribute assignment raises AttributeError.
-
- .. code:: python
-
- # Source
- I.register_type_var_decl(T.int32, dtype="int32")
- # Source
- n = T.int32()
- # Generated builder
- n = X.resolve_type_var_("n", dtype="int32")
- """
- constructor.__tvm_type_var_decl__ = DeclarationArguments(value_parameter, dtype)
- return constructor
-
-
-def register_binding_decl(constructor: _Callable) -> _Callable:
- """Mark a call that explicitly introduces an ordinary source binding.
-
- Parameters
- ----------
- constructor : callable
- Producer supporting syntax attribute assignment.
-
- Returns
- -------
- _Callable
- The unchanged producer with its declaration marker.
-
- Notes
- -----
- No evaluation, frame entry or source attachment occurs. Its assignment targets shadow an
- outer mutable declaration instead of storing through it; ordinary dialect ``bind_`` still
- owns result construction. Scalar and unpacked targets follow the same precedence.
- Attribute errors propagate.
-
- .. code:: python
-
- # Source
- I.register_binding_decl(make_value)
- # Source, even if an outer x is mutable
- x = make_value()
- # Generated builder
- x = X.bind_(make_value(), name="x")
- """
- constructor.__tvm_binding_decl__ = True
- return constructor
-
-
-def register_mutable_var_decl(constructor: _Callable, *, syntax: str = "call") -> _Callable:
- """Advertise explicit mutable storage declaration syntax.
-
- Parameters
- ----------
- constructor : callable
- Storage/annotation constructor supporting attribute assignment.
- syntax : str, optional
- One of call (default), annotation or parameter. Repeated registration adds to
- its immutable syntax set.
-
- Returns
- -------
- _Callable
- The same callable with its mutable syntax metadata.
-
- Notes
- -----
- No value is constructed, cached or inspected. Builders own storage and stores; prescan
- reads only syntax category. Unknown syntax raises ValueError and unsupported attribute
- assignment raises AttributeError. Declaration patterns take precedence over a prior
- target binding.
-
- .. code:: python
-
- # Source
- I.register_mutable_var_decl(T.int32, syntax="annotation")
- # Source
- x: T.int32 = 0
- # Generated builder
- x = X.decl_mutable_var_(0, ty=X.int32, name="x")
- """
- if syntax not in ("call", "annotation", "parameter"):
- raise ValueError("Mutable declaration syntax must be call, annotation or parameter")
- kinds: frozenset[str] = getattr(constructor, "__tvm_mutable_var_decl__", frozenset())
- constructor.__tvm_mutable_var_decl__ = kinds | frozenset((syntax,))
- return constructor
-
-
-def is_mutable_var_decl(constructor: object, *, syntax: str) -> bool:
- """Read whether a callable advertises a mutable declaration position.
-
- Parameters
- ----------
- constructor : object
- Resolved host constructor or annotation value; it is never called.
- syntax : str
- Source position to test, normally call, annotation or parameter.
-
- Returns
- -------
- bool
- True exactly when the position is registered; False for absent metadata.
-
- Notes
- -----
- No frame, IR effect or source attachment. This reads a callable attribute; custom
- attribute-access errors propagate. It does not inspect constructed storage or infer
- mutability from IR types.
-
- .. code:: python
-
- # Source
- x = T.local_scalar("int32")
- # Rewrite-time decision and generated operation
- I.is_mutable_var_decl(T.local_scalar, syntax="call")
- x = X.decl_mutable_var_(X.local_scalar("int32"), name="x")
- """
- return syntax in getattr(constructor, "__tvm_mutable_var_decl__", ())
-
-
-class FunctionDecoratorInfo(NamedTuple):
- """Flat syntax registration for a source function decorator.
-
- Parameters
- ----------
- builder : object or None
- Opaque construction namespace, or None for ordinary Python functions.
- option_map : dict of str to str, optional
- Public option names mapped to builder option names. Default is None,
- interpreted as empty. `register_function` copies supplied mappings.
- defaults : dict of str to object, optional
- Default builder keyword values. Default is None, interpreted as
- empty. `register_function` copies supplied mappings.
- python : bool, optional
- Preserve the source body as ordinary Python. Default is False.
-
- Notes
- -----
- These are the complete supported fields; no generic metadata dictionary
- is retained. Records live with decorators across compilations. Consumers
- must treat mappings as read-only. Records retain no frame, function
- result, annotation result, or symbol state and enter no builder frames.
- """
-
- builder: object
- option_map: dict[str, str] | None = None
- defaults: dict[str, Any] | None = None
- python: bool = False
-
-
-def register_function(
- decorator: _Callable,
- builder: object,
- *,
- option_map: Mapping[str, str] | None = None,
- defaults: Mapping[str, Any] | None = None,
- python: bool = False,
-) -> _Callable:
- """Register a source decorator and its explicit builder options.
-
- Parameters
- ----------
- decorator : callable
- Function decorator supporting attribute assignment; retained unchanged.
- builder : object or None
- Opaque dialect builder namespace, or None for ordinary Python functions.
- option_map : mapping of str to str, optional
- Source decorator keyword names mapped to builder keyword names. None (default)
- means empty; values are copied.
- defaults : mapping of str to Any, optional
- Default builder keyword values. None (default) means empty; values are copied.
- python : bool, optional
- False (default) requests script lowering. True preserves the original Python
- callable/body for module.__pyfuncs__.
-
- Returns
- -------
- _Callable
- The original decorator with its flat FunctionDecoratorInfo record.
-
- Notes
- -----
- No annotation, function or builder executes and no frame/span is retained. Registration
- replaces callable metadata for its lifetime; mappings are read-only by convention.
- Unsupported attributes or invalid mappings raise AttributeError/TypeError/ValueError.
- These records contain syntax options, never per-parse results.
-
- .. code:: python
-
- # Source
- I.register_function(prim_func, X, option_map={"private": "private"})
- # Source
- @prim_func(private=True)
- def f():
- pass
- # Generated builder entry
- with X.function(private=True):
- pass
- """
- decorator.__tvm_function_info__ = FunctionDecoratorInfo(
- builder, dict(option_map or {}), dict(defaults or {}), bool(python)
- )
- return decorator
-
-
-def function_info(decorator: object) -> FunctionDecoratorInfo | None:
- """Read registered function-decorator metadata.
-
- Parameters
- ----------
- decorator : object
- Resolved source decorator; it is not called by this lookup.
-
- Returns
- -------
- FunctionDecoratorInfo | None
- Shared flat record, or None for an unregistered decorator.
-
- Notes
- -----
- No context, construction effect or source attachment. Option mappings belong to
- registration and must not be mutated by a parse. Custom attribute-access errors
- propagate; unknown decorators are diagnosed by the parser consumer.
-
- .. code:: python
-
- # Source
- @T.prim_func
- def f():
- pass
- # Rewrite-time lookup and generated builder
- I.function_info(T.prim_func)
- with X.function():
- pass
- """
- return getattr(decorator, "__tvm_function_info__", None)
-
-
-def register_scope_var_query_or_decl(constructor: _Callable) -> _Callable:
- """Mark a call that already owns its scope variable or immutable declaration.
-
- Parameters
- ----------
- constructor : callable
- Scope query/declaration producer, for example a block-axis constructor, scope-ID
- helper or T.bind.
-
- Returns
- -------
- _Callable
- The exact producer with scope_var_query_or_decl syntax metadata.
-
- Notes
- -----
- Registration neither executes the producer nor stores an IR result/frame. The RHS
- executes once; generated assignment preserves returned identities and delegates source
- naming/validation to the dialect hook. This category precedes ordinary binding or a
- store to an outer mutable name. Callable attribute errors propagate.
-
- .. code:: python
-
- # Source
- I.register_scope_var_query_or_decl(T.bind)
- # Source
- x = T.bind(value)
- # Generated builder
- x = X.scope_var_query_or_decl_(X.bind(value), name="x")
- """
- constructor.__tvm_scope_var_query_or_decl__ = True
- return constructor
-
-
-def is_scope_var_query_or_decl(constructor: object) -> bool:
- """Read scope-query/declaration metadata from a resolved callable.
-
- Parameters
- ----------
- constructor : object
- Resolved producer or bound method; bound methods read the underlying function.
-
- Returns
- -------
- bool
- True for the registered syntax category, otherwise False.
-
- Notes
- -----
- No builder context/evaluation or span handling occurs. This is syntax recognition, not a
- runtime variable-type test; custom attribute-access errors propagate.
-
- .. code:: python
-
- # Source
- x = T.bind(value)
- # Rewrite-time decision and generated operation
- I.is_scope_var_query_or_decl(T.bind)
- x = X.scope_var_query_or_decl_(X.bind(value), name="x")
- """
- if isinstance(constructor, MethodType):
- constructor = constructor.__func__
- return bool(getattr(constructor, "__tvm_scope_var_query_or_decl__", False))
-
-
-def direct_call(constructor: _Callable) -> _Callable:
- """Mark a call that owns its result without automatic source-result handling.
-
- Parameters
- ----------
- constructor : callable
- Producer, method or class supporting attribute assignment. Callable identity and
- ordinary return values are retained.
-
- Returns
- -------
- _Callable
- The same callable; no forwarding callable or result wrapper is created.
-
- Notes
- -----
- No builder context or constructor evaluation. Registered source calls omit ``bind_``, ``emit_``,
- result ``at_`` and the outer ``with_at_group_`` wrapper. Callee/arguments still evaluate once in
- order, and argument expressions follow their own syntax rules. The callable owns any IR
- effects or source attribution it needs. This is separate from AlreadyEmitted, which
- retains ordinary result-span attachment. Attribute errors propagate.
-
- .. code:: python
-
- # Source
- @I.direct_call
- def identity(value):
- return value
- # Source and generated Python
- x = identity(value)
- identity(value)
- """
- constructor.__tvm_direct_call__ = True
- return constructor
-
-
-def is_direct_call(constructor: object) -> bool:
- """Read direct-call syntax metadata without invoking the callable.
-
- Parameters
- ----------
- constructor : object
- Resolved callable, class or bound method; methods read their underlying function
- metadata.
-
- Returns
- -------
- bool
- True if direct_call is registered, otherwise False.
-
- Notes
- -----
- No frame, IR effect, naming or span attachment. Attribute-access errors propagate.
- Recognition is by resolved callable metadata, never inferred from a constructed return
- type or an arbitrary method spelling.
-
- .. code:: python
-
- # Source
- x = I.meta_var(value)
- # Rewrite-time decision; generated execution remains a direct call
- I.is_direct_call(I.meta_var)
- x = I.meta_var(value)
- """
- if isinstance(constructor, MethodType):
- constructor = constructor.__func__
- return bool(getattr(constructor, "__tvm_direct_call__", False))
-
-
def scope_var_query_or_decl_(
value: Any, *, name: str | None = None, span: _Span = None, name_span: _Span = None
) -> Any:
@@ -2176,84 +1651,3 @@ def ne(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
X.ne(lhs, rhs)
"""
raise NotImplementedError
-
-
-def register_result_members(constructor: _Callable, members: object) -> _Callable:
- """Advertise a producer's static member namespace for syntax recognition.
-
- Parameters
- ----------
- constructor : callable
- Registered annotation or value producer. Bound methods attach metadata
- to their underlying function; the supplied callable is returned unchanged.
- members : object
- Opaque namespace of actual member callables/descriptors. It describes
- syntax available on the produced value, without describing an IR type.
-
- Returns
- -------
- _Callable
- The original producer with its member-namespace attribute.
-
- Notes
- -----
- No constructor executes and no result, frame or per-parse binding is stored.
- Parameter annotation and unambiguous declaration syntax can supply this
- namespace from existing prescan facts. A registered member producer may
- supply another namespace for chained views. Unknown or ambiguous aliases
- retain their ordinary policy; spelling alone never marks a method direct.
- Attribute-assignment errors propagate. Registration performs no source
- attachment and does not itself mark a call as direct_call.
-
- .. code:: python
-
- # Registration: BufferMethods.permute is the actual registered method.
- I.register_result_members(X.Buffer, BufferMethods)
- I.direct_call(BufferMethods.permute)
- # Source
- view = A.permute(1, 0) # A has source annotation X.Buffer(...).
- # Generated Python preserves the producer's ordinary result.
- view = A.permute(1, 0)
- """
- target = constructor.__func__ if isinstance(constructor, MethodType) else constructor
- target.__tvm_result_members__ = members
- return constructor
-
-
-def get_result_members(constructor: object) -> object | None:
- """Read a producer's opaque syntax member namespace without evaluation.
-
- Parameters
- ----------
- constructor : object
- Resolved annotation/producer, bound method or property descriptor.
- Methods use their function and properties use their getter metadata.
-
- Returns
- -------
- object or None
- Registered namespace, or None when static member information is absent.
-
- Notes
- -----
- This reads syntax metadata only. It enters no frame, calls no producer and
- inspects no constructed IR result. Metadata has callable lifetime; no
- function-local lookup state is cached. Custom attribute-access errors
- propagate. Member lookup still checks the actual member's direct_call or
- declaration metadata; a namespace does not make all its methods direct.
-
- .. code:: python
-
- # Source
- view = A.permute(1, 0)
- # Rewrite-time lookup for A's registered source annotation.
- members = I.get_result_members(X.Buffer)
- I.is_direct_call(members.permute)
- # Generated Python
- view = A.permute(1, 0)
- """
- if isinstance(constructor, property):
- constructor = constructor.fget
- if isinstance(constructor, MethodType):
- constructor = constructor.__func__
- return getattr(constructor, "__tvm_result_members__", None)
diff --git a/python/tvm/script/parser/call_args_policy.py b/python/tvm/script/parser/call_args_policy.py
index 196b2d1aa6..159728f8f3 100644
--- a/python/tvm/script/parser/call_args_policy.py
+++ b/python/tvm/script/parser/call_args_policy.py
@@ -31,7 +31,7 @@ import linecache
import re
from collections.abc import Callable
-from tvm.script.ir_builder.ir import parser_protocol as protocol
+from . import protocol_registry as protocol
class _LiteralParser:
diff --git a/python/tvm/script/parser/entry.py b/python/tvm/script/parser/entry.py
index 7fef2e1973..d59d425615 100644
--- a/python/tvm/script/parser/entry.py
+++ b/python/tvm/script/parser/entry.py
@@ -33,9 +33,9 @@ from typing import TYPE_CHECKING, Any, TypeVar
from tvm.ir import SourceName
from tvm.script.ir_builder import base
from tvm.script.ir_builder import ir as builder_ir
-from tvm.script.ir_builder.ir import parser_protocol as syntax_protocol
from . import jit_support
+from . import protocol_registry as syntax_protocol
from .inspect_source import (
acquire_source,
capture_annotation_bindings,
@@ -311,7 +311,7 @@ def make_decorator(
definition_scope = {} if deferred else capture_definition_scope(frame)
finally:
del frame
- function.__tvm_function_info__ = decorator.__tvm_function_info__
+ syntax_protocol.copy_function_info(decorator, function)
function.__tvm_function_options__ = options
if deferred:
return function
diff --git a/python/tvm/script/parser/prescan.py b/python/tvm/script/parser/prescan.py
index 80a7f1475a..61a34c7f64 100644
--- a/python/tvm/script/parser/prescan.py
+++ b/python/tvm/script/parser/prescan.py
@@ -25,8 +25,7 @@ from collections.abc import Mapping, Sequence
from types import GetSetDescriptorType, MemberDescriptorType
from typing import NamedTuple, NoReturn
-from tvm.script.ir_builder.ir import parser_protocol as protocol
-
+from . import protocol_registry as protocol
from .call_args_policy import parse_annotation
@@ -446,7 +445,7 @@ class PrescanCollector(ast.NodeVisitor):
self.environment,
)
dtype = getattr(constructor, "__tvm_parameter_dtype__", None)
- declaration = getattr(constructor, "__tvm_type_var_decl__", None)
+ declaration = protocol.get_type_var_decl(constructor)
self._record_binding(
arg.arg,
arg,
@@ -537,7 +536,7 @@ class PrescanCollector(ast.NodeVisitor):
kind: str = "ordinary",
) -> None:
constructor = resolve_constructor(value, self.environment, self.bindings[self.scope])
- if getattr(constructor, "__tvm_binding_decl__", False):
+ if protocol.is_binding_decl(constructor):
binding_declaration = value.func
while isinstance(binding_declaration, ast.Attribute):
binding_declaration = binding_declaration.value
@@ -556,7 +555,7 @@ class PrescanCollector(ast.NodeVisitor):
elif protocol.is_direct_call(constructor):
kind = "direct_call"
if isinstance(target, ast.Name):
- declaration = getattr(constructor, "__tvm_type_var_decl__", None)
+ declaration = protocol.get_type_var_decl(constructor)
if declaration is not None and not value.args and not value.keywords:
self._record_binding(target.id, target, "symbol", value, declaration.dtype)
elif getattr(self.builder, "supports_mutable_declarations", True) and (
diff --git a/python/tvm/script/parser/protocol_registry.py b/python/tvm/script/parser/protocol_registry.py
new file mode 100644
index 0000000000..2e69437dbe
--- /dev/null
+++ b/python/tvm/script/parser/protocol_registry.py
@@ -0,0 +1,868 @@
+# 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 registration, metadata and lookup helpers.
+
+Dialect builders register their callables here, and prescan, argument rewriting
+and transpilation read the resulting metadata. The argument-policy table lives
+across parses; other markers live on their callables. Aliases and bound methods
+share the existing callable identities. Records contain syntax facts and opaque
+builder namespaces, never active frames, constructed results or parse contexts.
+
+Importing this module needs only the standard library. Registration does not
+initialize parser entry points or concrete dialects. Expression-argument policies
+request the shared builder's eager annotation adapter only when applied; that
+adapter retains native construction and MissingType decisions.
+
+The complete customization overview and generated-operation contract are in
+``tvm.script.ir_builder.ir.parser_protocol``. Its documentation shows direct
+imports from this module; it does not own or re-export the registry.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from inspect import signature
+from types import MappingProxyType, MethodType
+from typing import Any, NamedTuple, NoReturn, TypeVar
+
+_Callable = TypeVar("_Callable", bound=Callable[..., Any])
+
+
+def constexpr(value: object) -> NoReturn:
+ """Mark a host control expression or a JIT specialization annotation.
+
+ Parameters
+ ----------
+ value : object
+ Source operand evaluated as ordinary Python after rewriting; never passed to
+ this marker at runtime.
+
+ Returns
+ -------
+ NoReturn
+ The marker is removed from generated execution; direct Python invocation raises
+ TypeError.
+
+ Notes
+ -----
+ No builder context or IR effects. Recognition uses callable identity, including aliases,
+ rather than name spelling or return values. The same marker can annotate a
+ specialization parameter. Host operators/control flow retain Python behavior while their
+ source calls still receive ordinary source-call instrumentation.
+
+ .. code:: python
+
+ # Source
+ if I.constexpr(enabled):
+ T.evaluate(1)
+ # Generated builder
+ if enabled:
+ X.emit_(X.evaluate(1))
+ """
+ raise TypeError("constexpr is a parser syntax marker, not a runtime operation")
+
+
+class ExprStrPolicy(NamedTuple):
+ """Immutable syntax policy for registered constructor arguments.
+
+ Parameters
+ ----------
+ fields : tuple of str
+ Parameter names whose string values represent expressions.
+ dtype : object, optional
+ Dtype spelling passed to builder symbol resolution. Default is None.
+ scalar_strings : bool, optional
+ Interpret bare strings as expressions. Default is True. Nested
+ strings in marked fields are always treated as expressions.
+
+ Notes
+ -----
+ This tuple record lives with its registered callable across transpilation
+ passes. It stores no symbols, eager results, or function-local state.
+ Dtype spellings are forwarded as builder arguments; the transpiler never
+ constructs or validates symbols. Builders retain already-declared symbol
+ identity and type. Construction enters no frame.
+ """
+
+ fields: tuple[str, ...]
+ dtype: object = None
+ scalar_strings: bool = True
+
+
+# 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.
+_ARGS_POLICIES: dict[object, ArgsPolicy] = {}
+
+
+class ArgsPolicy(NamedTuple):
+ """Immutable per-parameter policies shared by a callable and its adapter.
+
+ Parameters
+ ----------
+ fields : Mapping[str, str]
+ Immutable parameter-to-policy mapping containing expr_str/global_info.
+ expression : ExprStrPolicy
+ Expression fields, scalar-string behavior and opaque symbolic dtype.
+
+ Notes
+ -----
+ This process-wide syntax record stores no evaluated arguments, scopes,
+ frames or results. Registration copies its mappings; parses only read it.
+ """
+
+ fields: Mapping[str, str]
+ expression: ExprStrPolicy
+
+
+def get_args_policy(constructor: object) -> ArgsPolicy | None:
+ """Read argument policy metadata without evaluating a constructor.
+
+ Parameters
+ ----------
+ constructor : object
+ Resolved host callable or bound method; bound methods use their underlying
+ function identity.
+
+ Returns
+ -------
+ ArgsPolicy | None
+ Shared immutable policy, or None for unregistered/unhashable values.
+
+ Notes
+ -----
+ No builder context, IR construction or source-span changes occur. The returned policy
+ belongs to process-wide registration, not a particular parse. Custom hashing errors
+ other than TypeError propagate.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import get_args_policy
+
+ # Source
+ tensor(("n",))
+ # Rewrite-time lookup; generated code receives decoded arguments
+ policy = get_args_policy(tensor)
+ tensor((X.resolve_type_var_("n"),))
+ """
+ if isinstance(constructor, MethodType):
+ constructor = constructor.__func__
+ try:
+ return _ARGS_POLICIES.get(constructor)
+ except TypeError:
+ return None
+
+
+def args_policy(
+ fields: Mapping[str, str],
+ *,
+ scalar_strings: bool = True,
+ dtype: object = None,
+ as_type: bool = False,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+ """Register per-argument source policies without storing construction results.
+
+ Parameters
+ ----------
+ fields : mapping of str to str
+ Parameter names mapped to expr_str or global_info. Copied into immutable
+ registration data; positional and keyword calls use the same policy.
+ scalar_strings : bool, optional
+ True by default. False preserves a bare string in an expression field; nested
+ strings are still expression syntax.
+ dtype : object, optional
+ Opaque dtype passed to symbol resolution for expression strings. None (default)
+ leaves the builder default.
+ as_type : bool, optional
+ False by default. True preserves the annotation-class surface, including Python
+ type unions, while construction returns the underlying result.
+
+ Returns
+ -------
+ Callable[[Callable[..., Any]], Callable[..., Any]]
+ Decorator returning the callable or its eager annotation adapter. Original and
+ adapted identities share one argument policy.
+
+ Notes
+ -----
+ Registration enters no frame, constructs no IR and retains no per-function values.
+ Concrete arguments call the original constructor. Outside construction, unresolved
+ expression annotations yield MissingType; inside construction, unresolved expression
+ strings raise TypeError and TypeVar values use the nearest function resolver. Unknown
+ policy kinds/parameter names raise ValueError; signature/argument-binding errors
+ propagate. The adapter belongs to the builder, not parser execution.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import args_policy
+
+ # Source
+ @args_policy({"shape": "expr_str", "device": "global_info"})
+ def tensor(shape, device=None):
+ ...
+ # Source
+ tensor(("n",), device="cuda:0")
+ # Generated builder
+ tensor((X.resolve_type_var_("n"),), device=I.resolve_global_info_("cuda:0"))
+ """
+ fields = dict(fields)
+ unsupported = set(fields.values()).difference(("expr_str", "global_info"))
+ if unsupported:
+ raise ValueError(f"Unknown argument policies: {sorted(unsupported)}")
+
+ def decorate(constructor: Callable[..., Any]) -> Callable[..., Any]:
+ call_signature = signature(constructor)
+ unknown = set(fields).difference(call_signature.parameters)
+ if unknown:
+ raise ValueError(f"Unknown argument policy fields: {sorted(unknown)}")
+ expression_fields = tuple(name for name, kind in fields.items() if kind == "expr_str")
+ expression = ExprStrPolicy(expression_fields, dtype, bool(scalar_strings))
+ if expression_fields or as_type:
+ # Builders own eager construction and active-frame/MissingType decisions.
+ from tvm.script.ir_builder.base import wrap_expression_constructor
+
+ result = wrap_expression_constructor(
+ constructor, call_signature, expression, as_type=as_type
+ )
+ else:
+ result = constructor
+ policy = ArgsPolicy(MappingProxyType(fields.copy()), expression)
+ _ARGS_POLICIES[constructor] = policy
+ _ARGS_POLICIES[result] = policy
+ return result
+
+ return decorate
+
+
+class DeclarationArguments(NamedTuple):
+ """Immutable metadata for scalar-constructor predeclarations.
+
+ Parameters
+ ----------
+ value_parameter : str
+ Optional value parameter whose absence denotes a declaration.
+ dtype : object, optional
+ Explicit primitive dtype metadata. Default is None.
+
+ Notes
+ -----
+ `register_type_var_decl` attaches this record to its callable. Aliases
+ and transpilation passes share it for the callable's lifetime; it stores
+ no symbols or construction results and enters no builder frame.
+ """
+
+ value_parameter: str
+ dtype: object = None
+
+
+def register_type_var_decl(
+ constructor: _Callable, *, value_parameter: str = "expr", dtype: object = None
+) -> _Callable:
+ """Mark a constructor whose omitted value argument declares a symbolic variable.
+
+ Parameters
+ ----------
+ constructor : callable
+ Constructor supporting syntax attribute assignment. Identity is retained.
+ value_parameter : str, optional
+ Value parameter whose absence means declaration; defaults to expr.
+ dtype : object, optional
+ Opaque primitive dtype, default None. A string dtype also allows prescan to
+ recognize direct zero-argument body declarations.
+
+ Returns
+ -------
+ _Callable
+ The original callable with DeclarationArguments metadata.
+
+ Notes
+ -----
+ No constructor executes and no frame is entered. Syntax predeclaration precedes
+ dependent signature shapes while original body ordering remains. Only direct declaration
+ forms are collected, not nested/effectful expressions. Aliases share callable metadata;
+ unsupported attribute assignment raises AttributeError.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import register_type_var_decl
+
+ # Source
+ register_type_var_decl(T.int32, dtype="int32")
+ # Source
+ n = T.int32()
+ # Generated builder
+ n = X.resolve_type_var_("n", dtype="int32")
+ """
+ constructor.__tvm_type_var_decl__ = DeclarationArguments(value_parameter, dtype)
+ return constructor
+
+
+def register_binding_decl(constructor: _Callable) -> _Callable:
+ """Mark a call that explicitly introduces an ordinary source binding.
+
+ Parameters
+ ----------
+ constructor : callable
+ Producer supporting syntax attribute assignment.
+
+ Returns
+ -------
+ _Callable
+ The unchanged producer with its declaration marker.
+
+ Notes
+ -----
+ No evaluation, frame entry or source attachment occurs. Its assignment targets shadow an
+ outer mutable declaration instead of storing through it; ordinary dialect ``bind_`` still
+ owns result construction. Scalar and unpacked targets follow the same precedence.
+ Attribute errors propagate.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import register_binding_decl
+
+ # Source
+ register_binding_decl(make_value)
+ # Source, even if an outer x is mutable
+ x = make_value()
+ # Generated builder
+ x = X.bind_(make_value(), name="x")
+ """
+ constructor.__tvm_binding_decl__ = True
+ return constructor
+
+
+def get_type_var_decl(constructor: object) -> DeclarationArguments | None:
+ """Read scalar-declaration metadata without evaluating the constructor.
+
+ Parameters
+ ----------
+ constructor : object
+ Resolved source constructor or annotation.
+
+ Returns
+ -------
+ DeclarationArguments or None
+ The attached record, or None when no declaration is registered.
+
+ Notes
+ -----
+ Ordinary attribute lookup preserves callable aliases and propagates custom
+ attribute-access errors. No frame, symbol or per-parse state is created.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import get_type_var_decl
+
+ # Source
+ n = X.int32()
+ # Rewrite-time lookup and generated builder
+ declaration = get_type_var_decl(X.int32)
+ n = X.resolve_type_var_("n", dtype=declaration.dtype)
+ """
+ return getattr(constructor, "__tvm_type_var_decl__", None)
+
+
+def is_binding_decl(constructor: object) -> bool:
+ """Read whether a producer explicitly introduces an ordinary binding.
+
+ Parameters
+ ----------
+ constructor : object
+ Resolved source callable; it is not invoked.
+
+ Returns
+ -------
+ bool
+ The truth value of its binding marker, or False when absent.
+
+ Notes
+ -----
+ The marker is read and tested once. Attribute or truth-conversion errors
+ propagate unchanged. No native value, frame or source location is inspected.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import is_binding_decl
+
+ # Source
+ x = make_value()
+ # Rewrite-time decision and generated builder
+ is_binding_decl(make_value)
+ x = X.bind_(make_value(), name="x")
+ """
+ return bool(getattr(constructor, "__tvm_binding_decl__", False))
+
+
+def register_mutable_var_decl(constructor: _Callable, *, syntax: str = "call") -> _Callable:
+ """Advertise explicit mutable storage declaration syntax.
+
+ Parameters
+ ----------
+ constructor : callable
+ Storage/annotation constructor supporting attribute assignment.
+ syntax : str, optional
+ One of call (default), annotation or parameter. Repeated registration adds to
+ its immutable syntax set.
+
+ Returns
+ -------
+ _Callable
+ The same callable with its mutable syntax metadata.
+
+ Notes
+ -----
+ No value is constructed, cached or inspected. Builders own storage and stores; prescan
+ reads only syntax category. Unknown syntax raises ValueError and unsupported attribute
+ assignment raises AttributeError. Declaration patterns take precedence over a prior
+ target binding.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import register_mutable_var_decl
+
+ # Source
+ register_mutable_var_decl(T.int32, syntax="annotation")
+ # Source
+ x: T.int32 = 0
+ # Generated builder
+ x = X.decl_mutable_var_(0, ty=X.int32, name="x")
+ """
+ if syntax not in ("call", "annotation", "parameter"):
+ raise ValueError("Mutable declaration syntax must be call, annotation or parameter")
+ kinds: frozenset[str] = getattr(constructor, "__tvm_mutable_var_decl__", frozenset())
+ constructor.__tvm_mutable_var_decl__ = kinds | frozenset((syntax,))
+ return constructor
+
+
+def is_mutable_var_decl(constructor: object, *, syntax: str) -> bool:
+ """Read whether a callable advertises a mutable declaration position.
+
+ Parameters
+ ----------
+ constructor : object
+ Resolved host constructor or annotation value; it is never called.
+ syntax : str
+ Source position to test, normally call, annotation or parameter.
+
+ Returns
+ -------
+ bool
+ True exactly when the position is registered; False for absent metadata.
+
+ Notes
+ -----
+ No frame, IR effect or source attachment. This reads a callable attribute; custom
+ attribute-access errors propagate. It does not inspect constructed storage or infer
+ mutability from IR types.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import is_mutable_var_decl
+
+ # Source
+ x = T.local_scalar("int32")
+ # Rewrite-time decision and generated operation
+ is_mutable_var_decl(T.local_scalar, syntax="call")
+ x = X.decl_mutable_var_(X.local_scalar("int32"), name="x")
+ """
+ return syntax in getattr(constructor, "__tvm_mutable_var_decl__", ())
+
+
+class FunctionDecoratorInfo(NamedTuple):
+ """Flat syntax registration for a source function decorator.
+
+ Parameters
+ ----------
+ builder : object or None
+ Opaque construction namespace, or None for ordinary Python functions.
+ option_map : dict of str to str, optional
+ Public option names mapped to builder option names. Default is None,
+ interpreted as empty. `register_function` copies supplied mappings.
+ defaults : dict of str to object, optional
+ Default builder keyword values. Default is None, interpreted as
+ empty. `register_function` copies supplied mappings.
+ python : bool, optional
+ Preserve the source body as ordinary Python. Default is False.
+
+ Notes
+ -----
+ These are the complete supported fields; no generic metadata dictionary
+ is retained. Records live with decorators across compilations. Consumers
+ must treat mappings as read-only. Records retain no frame, function
+ result, annotation result, or symbol state and enter no builder frames.
+ """
+
+ builder: object
+ option_map: dict[str, str] | None = None
+ defaults: dict[str, Any] | None = None
+ python: bool = False
+
+
+def register_function(
+ decorator: _Callable,
+ builder: object,
+ *,
+ option_map: Mapping[str, str] | None = None,
+ defaults: Mapping[str, Any] | None = None,
+ python: bool = False,
+) -> _Callable:
+ """Register a source decorator and its explicit builder options.
+
+ Parameters
+ ----------
+ decorator : callable
+ Function decorator supporting attribute assignment; retained unchanged.
+ builder : object or None
+ Opaque dialect builder namespace, or None for ordinary Python functions.
+ option_map : mapping of str to str, optional
+ Source decorator keyword names mapped to builder keyword names. None (default)
+ means empty; values are copied.
+ defaults : mapping of str to Any, optional
+ Default builder keyword values. None (default) means empty; values are copied.
+ python : bool, optional
+ False (default) requests script lowering. True preserves the original Python
+ callable/body for module.__pyfuncs__.
+
+ Returns
+ -------
+ _Callable
+ The original decorator with its flat FunctionDecoratorInfo record.
+
+ Notes
+ -----
+ No annotation, function or builder executes and no frame/span is retained. Registration
+ replaces callable metadata for its lifetime; mappings are read-only by convention.
+ Unsupported attributes or invalid mappings raise AttributeError/TypeError/ValueError.
+ These records contain syntax options, never per-parse results.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import register_function
+
+ # Source
+ register_function(prim_func, X, option_map={"private": "private"})
+ # Source
+ @prim_func(private=True)
+ def f():
+ pass
+ # Generated builder entry
+ with X.function(private=True):
+ pass
+ """
+ decorator.__tvm_function_info__ = FunctionDecoratorInfo(
+ builder, dict(option_map or {}), dict(defaults or {}), bool(python)
+ )
+ return decorator
+
+
+def function_info(decorator: object) -> FunctionDecoratorInfo | None:
+ """Read registered function-decorator metadata.
+
+ Parameters
+ ----------
+ decorator : object
+ Resolved source decorator; it is not called by this lookup.
+
+ Returns
+ -------
+ FunctionDecoratorInfo | None
+ Shared flat record, or None for an unregistered decorator.
+
+ Notes
+ -----
+ No context, construction effect or source attachment. Option mappings belong to
+ registration and must not be mutated by a parse. Custom attribute-access errors
+ propagate; unknown decorators are diagnosed by the parser consumer.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import function_info
+
+ # Source
+ @T.prim_func
+ def f():
+ pass
+ # Rewrite-time lookup and generated builder
+ function_info(T.prim_func)
+ with X.function():
+ pass
+ """
+ return getattr(decorator, "__tvm_function_info__", None)
+
+
+def copy_function_info(source: Callable[..., Any], target: Callable[..., Any]) -> None:
+ """Attach an existing decorator registration to its source function.
+
+ Parameters
+ ----------
+ source : callable
+ Registered decorator whose function metadata is required.
+ target : callable
+ Source function that will share the exact registration record.
+
+ Returns
+ -------
+ None
+ The target is updated in place without constructing another record.
+
+ Notes
+ -----
+ Missing registration or unsupported attribute assignment raises AttributeError.
+ The source attribute is read once and assigned unchanged. Definition scope,
+ applied decorator options and native construction remain entry-point concerns;
+ no parse environment or result is retained by this operation.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import copy_function_info
+
+ # Source
+ @prim_func
+ def f():
+ pass
+ # Decorator preparation before the generated builder executes
+ copy_function_info(prim_func, f)
+ """
+ target.__tvm_function_info__ = source.__tvm_function_info__
+
+
+def register_scope_var_query_or_decl(constructor: _Callable) -> _Callable:
+ """Mark a call that already owns its scope variable or immutable declaration.
+
+ Parameters
+ ----------
+ constructor : callable
+ Scope query/declaration producer, for example a block-axis constructor, scope-ID
+ helper or T.bind.
+
+ Returns
+ -------
+ _Callable
+ The exact producer with scope_var_query_or_decl syntax metadata.
+
+ Notes
+ -----
+ Registration neither executes the producer nor stores an IR result/frame. The RHS
+ executes once; generated assignment preserves returned identities and delegates source
+ naming/validation to the dialect hook. This category precedes ordinary binding or a
+ store to an outer mutable name. Callable attribute errors propagate.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import register_scope_var_query_or_decl
+
+ # Source
+ register_scope_var_query_or_decl(T.bind)
+ # Source
+ x = T.bind(value)
+ # Generated builder
+ x = X.scope_var_query_or_decl_(X.bind(value), name="x")
+ """
+ constructor.__tvm_scope_var_query_or_decl__ = True
+ return constructor
+
+
+def is_scope_var_query_or_decl(constructor: object) -> bool:
+ """Read scope-query/declaration metadata from a resolved callable.
+
+ Parameters
+ ----------
+ constructor : object
+ Resolved producer or bound method; bound methods read the underlying function.
+
+ Returns
+ -------
+ bool
+ True for the registered syntax category, otherwise False.
+
+ Notes
+ -----
+ No builder context/evaluation or span handling occurs. This is syntax recognition, not a
+ runtime variable-type test; custom attribute-access errors propagate.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import is_scope_var_query_or_decl
+
+ # Source
+ x = T.bind(value)
+ # Rewrite-time decision and generated operation
+ is_scope_var_query_or_decl(T.bind)
+ x = X.scope_var_query_or_decl_(X.bind(value), name="x")
+ """
+ if isinstance(constructor, MethodType):
+ constructor = constructor.__func__
+ return bool(getattr(constructor, "__tvm_scope_var_query_or_decl__", False))
+
+
+def direct_call(constructor: _Callable) -> _Callable:
+ """Mark a call that owns its result without automatic source-result handling.
+
+ Parameters
+ ----------
+ constructor : callable
+ Producer, method or class supporting attribute assignment. Callable identity and
+ ordinary return values are retained.
+
+ Returns
+ -------
+ _Callable
+ The same callable; no forwarding callable or result wrapper is created.
+
+ Notes
+ -----
+ No builder context or constructor evaluation. Registered source calls omit ``bind_``, ``emit_``,
+ result ``at_`` and the outer ``with_at_group_`` wrapper. Callee/arguments still evaluate once in
+ order, and argument expressions follow their own syntax rules. The callable owns any IR
+ effects or source attribution it needs. This is separate from AlreadyEmitted, which
+ retains ordinary result-span attachment. Attribute errors propagate.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import direct_call
+
+ # Source
+ @direct_call
+ def identity(value):
+ return value
+ # Source and generated Python
+ x = identity(value)
+ identity(value)
+ """
+ constructor.__tvm_direct_call__ = True
+ return constructor
+
+
+def is_direct_call(constructor: object) -> bool:
+ """Read direct-call syntax metadata without invoking the callable.
+
+ Parameters
+ ----------
+ constructor : object
+ Resolved callable, class or bound method; methods read their underlying function
+ metadata.
+
+ Returns
+ -------
+ bool
+ True if direct_call is registered, otherwise False.
+
+ Notes
+ -----
+ No frame, IR effect, naming or span attachment. Attribute-access errors propagate.
+ Recognition is by resolved callable metadata, never inferred from a constructed return
+ type or an arbitrary method spelling.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import is_direct_call
+
+ # Source
+ x = I.meta_var(value)
+ # Rewrite-time decision; generated execution remains a direct call
+ is_direct_call(I.meta_var)
+ x = I.meta_var(value)
+ """
+ if isinstance(constructor, MethodType):
+ constructor = constructor.__func__
+ return bool(getattr(constructor, "__tvm_direct_call__", False))
+
+
+def register_result_members(constructor: _Callable, members: object) -> _Callable:
+ """Advertise a producer's static member namespace for syntax recognition.
+
+ Parameters
+ ----------
+ constructor : callable
+ Registered annotation or value producer. Bound methods attach metadata
+ to their underlying function; the supplied callable is returned unchanged.
+ members : object
+ Opaque namespace of actual member callables/descriptors. It describes
+ syntax available on the produced value, without describing an IR type.
+
+ Returns
+ -------
+ _Callable
+ The original producer with its member-namespace attribute.
+
+ Notes
+ -----
+ No constructor executes and no result, frame or per-parse binding is stored.
+ Parameter annotation and unambiguous declaration syntax can supply this
+ namespace from existing prescan facts. A registered member producer may
+ supply another namespace for chained views. Unknown or ambiguous aliases
+ retain their ordinary policy; spelling alone never marks a method direct.
+ Attribute-assignment errors propagate. Registration performs no source
+ attachment and does not itself mark a call as direct_call.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import direct_call, register_result_members
+
+ # Registration: BufferMethods.permute is the actual registered method.
+ register_result_members(X.Buffer, BufferMethods)
+ direct_call(BufferMethods.permute)
+ # Source
+ view = A.permute(1, 0) # A has source annotation X.Buffer(...).
+ # Generated Python preserves the producer's ordinary result.
+ view = A.permute(1, 0)
+ """
+ target = constructor.__func__ if isinstance(constructor, MethodType) else constructor
+ target.__tvm_result_members__ = members
+ return constructor
+
+
+def get_result_members(constructor: object) -> object | None:
+ """Read a producer's opaque syntax member namespace without evaluation.
+
+ Parameters
+ ----------
+ constructor : object
+ Resolved annotation/producer, bound method or property descriptor.
+ Methods use their function and properties use their getter metadata.
+
+ Returns
+ -------
+ object or None
+ Registered namespace, or None when static member information is absent.
+
+ Notes
+ -----
+ This reads syntax metadata only. It enters no frame, calls no producer and
+ inspects no constructed IR result. Metadata has callable lifetime; no
+ function-local lookup state is cached. Custom attribute-access errors
+ propagate. Member lookup still checks the actual member's direct_call or
+ declaration metadata; a namespace does not make all its methods direct.
+
+ .. code:: python
+
+ from tvm.script.parser.protocol_registry import get_result_members, is_direct_call
+
+ # Source
+ view = A.permute(1, 0)
+ # Rewrite-time lookup for A's registered source annotation.
+ members = get_result_members(X.Buffer)
+ is_direct_call(members.permute)
+ # Generated Python
+ view = A.permute(1, 0)
+ """
+ if isinstance(constructor, property):
+ constructor = constructor.fget
+ if isinstance(constructor, MethodType):
+ constructor = constructor.__func__
+ return getattr(constructor, "__tvm_result_members__", None)
diff --git a/python/tvm/script/parser/transpile.py b/python/tvm/script/parser/transpile.py
index 53975dc20d..6dc89da6db 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -43,8 +43,7 @@ from contextlib import contextmanager
from types import FunctionType
from typing import Any, NoReturn, TypeVar
-from tvm.script.ir_builder.ir import parser_protocol as protocol
-
+from . import protocol_registry as protocol
from .call_args_policy import handle_call_args_policy, parse_annotation, parse_expression_string
from .prescan import Binding, PrescanContext, resolve_constructor, resolve_syntax
diff --git a/python/tvm/tirx/script/builder/__init__.py b/python/tvm/tirx/script/builder/__init__.py
index ef17cf7ddc..a8d93720c3 100644
--- a/python/tvm/tirx/script/builder/__init__.py
+++ b/python/tvm/tirx/script/builder/__init__.py
@@ -24,9 +24,9 @@ from tvm import ir as _ir
from tvm import tirx as _tir
from tvm.script.ir_builder.base import at as _at
from tvm.script.ir_builder.base import source_span as _source_span
-from tvm.script.ir_builder.ir.parser_protocol import args_policy as _args_policy
-from tvm.script.ir_builder.ir.parser_protocol import constexpr as constexpr
-from tvm.script.ir_builder.ir.parser_protocol import (
+from tvm.script.parser.protocol_registry import args_policy as _args_policy
+from tvm.script.parser.protocol_registry import constexpr as constexpr
+from tvm.script.parser.protocol_registry import (
register_type_var_decl as _register_type_var_decl,
)
from tvm.tirx.lang.alloc_pool import SMEMPool as SMEMPool
diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py
index 1b4d6016a4..bc03a505e3 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -2890,7 +2890,7 @@ if TYPE_CHECKING:
else:
def _install_meta_class(cls):
- from tvm.script.ir_builder.ir.parser_protocol import direct_call
+ from tvm.script.parser.protocol_registry import direct_call
direct_call(cls)
if cls.__dict__.get("_tirx_meta_class_installed", False):
diff --git a/python/tvm/tirx/script/builder/parser_protocol.py b/python/tvm/tirx/script/builder/parser_protocol.py
index 6d74d26eb3..6b311b01f6 100644
--- a/python/tvm/tirx/script/builder/parser_protocol.py
+++ b/python/tvm/tirx/script/builder/parser_protocol.py
@@ -427,7 +427,7 @@ def set_mutable_var_(
def _register_declarations() -> None:
- from tvm.script.ir_builder.ir.parser_protocol import (
+ from tvm.script.parser.protocol_registry import (
direct_call,
register_mutable_var_decl,
register_result_members,
diff --git a/python/tvm/tirx/script/jit.py b/python/tvm/tirx/script/jit.py
index 280481f111..a06160087f 100644
--- a/python/tvm/tirx/script/jit.py
+++ b/python/tvm/tirx/script/jit.py
@@ -24,7 +24,7 @@ from collections.abc import Callable, Mapping
from types import FunctionType
from typing import Any
-from tvm.script.ir_builder.ir import parser_protocol as protocol
+from tvm.script.parser import protocol_registry as protocol
from tvm.tirx import PrimFunc
@@ -100,7 +100,7 @@ def make_jit(builder: object) -> Callable[..., Any]:
definition_scope = capture_definition_scope(frame)
finally:
del frame
- function.__tvm_function_info__ = jit.__tvm_function_info__
+ protocol.copy_function_info(jit, function)
function.__tvm_function_options__ = {
"private": private,
"s_tir": is_stir,
diff --git a/tests/python/script/dummy_builder.py b/tests/python/script/dummy_builder.py
index 6ab2d926fc..f093e23ced 100644
--- a/tests/python/script/dummy_builder.py
+++ b/tests/python/script/dummy_builder.py
@@ -29,8 +29,8 @@ from dataclasses import dataclass, field
from types import SimpleNamespace
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder.ir import parser_protocol as protocol
from tvm.script.parser import entry
+from tvm.script.parser import protocol_registry as registry
@dataclass(eq=False)
@@ -144,7 +144,7 @@ class Language:
annotation_value_=lambda name, value: value,
MISSING=self.missing,
check_well_formed_=lambda result: None,
- constexpr=protocol.constexpr,
+ constexpr=registry.constexpr,
)
self.X = SimpleNamespace(
supports_mutable_declarations=True,
@@ -174,7 +174,7 @@ class Language:
and_=lambda *args, **kwargs: Value("and", args),
or_=lambda *args: Value("or", args),
not_=lambda value: Value("not", (value,)),
- constexpr=protocol.constexpr,
+ constexpr=registry.constexpr,
value=lambda *args: Value("value", args),
record=self.record,
)
@@ -187,7 +187,7 @@ class Language:
setattr(self.X, name + "_", operation)
self.X.script = entry.make_decorator(self.X)
- @protocol.args_policy({"shape": "expr_str", "device": "global_info"}, scalar_strings=False)
+ @registry.args_policy({"shape": "expr_str", "device": "global_info"}, scalar_strings=False)
def tensor(shape=None, dtype="float32", device=None, placement="S[0]"):
return Value("tensor", (shape, dtype, device, placement))
@@ -198,8 +198,8 @@ class Language:
return Value("cell", (value,))
self.X.tensor = tensor
- self.X.symbol = protocol.register_type_var_decl(symbol, dtype="int64")
- self.X.cell = protocol.register_mutable_var_decl(cell)
+ self.X.symbol = registry.register_type_var_decl(symbol, dtype="int64")
+ self.X.cell = registry.register_mutable_var_decl(cell)
@contextmanager
def context(self):
diff --git a/tests/python/script/test_builder_protocol_contract.py b/tests/python/script/test_builder_protocol_contract.py
index 9bbd1c2688..138ab6e713 100644
--- a/tests/python/script/test_builder_protocol_contract.py
+++ b/tests/python/script/test_builder_protocol_contract.py
@@ -18,6 +18,9 @@
import inspect
import re
+import subprocess
+import sys
+import textwrap
import typing
import pytest
@@ -26,10 +29,11 @@ from tvm import ir
from tvm.relax.script.builder import parser_protocol as RP
from tvm.script.ir_builder import ir as IB
from tvm.script.ir_builder.ir import parser_protocol as P
+from tvm.script.parser import protocol_registry as registry
from tvm.tirx.script.builder import parser_protocol as TP
-@pytest.mark.parametrize("protocol", [P, TP, RP])
+@pytest.mark.parametrize("protocol", [P, TP, RP, registry])
def test_protocol_signatures_and_local_examples_are_complete(protocol):
for name, function in vars(protocol).items():
if name.startswith("_") or not inspect.isfunction(function):
@@ -53,27 +57,27 @@ def test_registration_keeps_callable_identity_and_bound_method_metadata():
calls = []
class Factory:
- @P.direct_call
+ @registry.direct_call
def make(self, value):
calls.append(value)
return value
factory = Factory()
function = Factory.make
- assert P.direct_call(function) is function
- assert P.is_direct_call(factory.make)
+ assert registry.direct_call(function) is function
+ assert registry.is_direct_call(factory.make)
assert calls == []
token = object()
assert factory.make(token) is token
assert calls == [token]
- assert not P.is_direct_call(object())
+ assert not registry.is_direct_call(object())
def scope_value():
return token
- assert P.register_scope_var_query_or_decl(scope_value) is scope_value
- assert P.is_scope_var_query_or_decl(scope_value)
- assert not P.is_mutable_var_decl(scope_value, syntax="call")
+ assert registry.register_scope_var_query_or_decl(scope_value) is scope_value
+ assert registry.is_scope_var_query_or_decl(scope_value)
+ assert not registry.is_mutable_var_decl(scope_value, syntax="call")
assert scope_value() is token
@@ -86,4 +90,138 @@ def test_meta_var_retains_exact_payload_and_existing_span():
assert IB.meta_var(pair) is pair
left, right = IB.meta_var(pair)
assert left is value and right is pair[1]
- assert P.is_direct_call(IB.meta_var)
+ assert registry.is_direct_call(IB.meta_var)
+
+
+def test_registration_has_one_parser_owner_and_no_builder_reexports():
+ for name in (
+ "ExprStrPolicy",
+ "ArgsPolicy",
+ "DeclarationArguments",
+ "FunctionDecoratorInfo",
+ "args_policy",
+ "get_args_policy",
+ "direct_call",
+ "is_direct_call",
+ "register_type_var_decl",
+ "get_type_var_decl",
+ "register_binding_decl",
+ "is_binding_decl",
+ "register_mutable_var_decl",
+ "is_mutable_var_decl",
+ "register_scope_var_query_or_decl",
+ "is_scope_var_query_or_decl",
+ "register_function",
+ "function_info",
+ "copy_function_info",
+ "register_result_members",
+ "get_result_members",
+ ):
+ assert getattr(registry, name).__module__ == registry.__name__, name
+ assert name not in vars(P), name
+ assert name not in vars(IB), name
+ # constexpr is a real source marker, shared by builder namespaces by identity.
+ assert IB.constexpr is registry.constexpr
+
+ def source():
+ raise AssertionError("metadata lookup must not invoke its callable")
+
+ def target():
+ raise AssertionError("metadata transfer must not invoke its callable")
+
+ assert registry.get_type_var_decl(source) is None
+ assert not registry.is_binding_decl(source)
+ assert registry.function_info(target) is None
+ with pytest.raises(AttributeError, match="__tvm_function_info__"):
+ registry.copy_function_info(source, target)
+ assert registry.function_info(target) is None
+
+ dtype = object()
+ assert registry.register_type_var_decl(source, dtype=dtype) is source
+ declaration = registry.get_type_var_decl(source)
+ assert declaration.dtype is dtype
+ assert registry.get_type_var_decl(source) is declaration
+ assert registry.register_binding_decl(source) is source
+ assert registry.is_binding_decl(source)
+ builder = object()
+ assert registry.register_function(source, builder) is source
+ metadata = registry.function_info(source)
+ assert metadata.builder is builder
+ assert registry.copy_function_info(source, target) is None
+ assert registry.function_info(target) is metadata
+ assert registry.function_info(source) is metadata
+
+
+def test_registry_import_and_metadata_registration_need_no_tvm_initialization():
+ # Load the actual module in a fresh interpreter with every TVM import blocked.
+ script = textwrap.dedent(
+ """
+ import importlib.abc
+ import importlib.util
+ import sys
+
+ class BlockTVM(importlib.abc.MetaPathFinder):
+ def find_spec(self, fullname, path=None, target=None):
+ if fullname == "tvm" or fullname.startswith("tvm."):
+ raise AssertionError("Registry initialized TVM: " + fullname)
+ return None
+
+ sys.meta_path.insert(0, BlockTVM())
+ spec = importlib.util.spec_from_file_location("isolated_registry", sys.argv[1])
+ registry = importlib.util.module_from_spec(spec)
+ sys.modules[spec.name] = registry
+ spec.loader.exec_module(registry)
+
+ def constructor(device):
+ return device
+
+ assert registry.args_policy({"device": "global_info"})(constructor) is constructor
+ policy = registry.get_args_policy(constructor)
+ assert dict(policy.fields) == {"device": "global_info"}
+ assert registry.direct_call(constructor) is constructor
+ assert registry.is_direct_call(constructor)
+ assert registry.register_function(constructor, object()) is constructor
+ assert registry.function_info(constructor) is not None
+ assert not any(name == "tvm" or name.startswith("tvm.") for name in sys.modules)
+ """
+ )
+ result = subprocess.run(
+ [sys.executable, "-c", script, registry.__file__],
+ capture_output=True,
+ text=True,
+ check=False,
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_registration_persists_across_successful_and_failed_parse_contexts(language):
+ calls = []
+ first, invalid, second = object(), object(), object()
+ failure = RuntimeError("constructor failed")
+
+ @registry.args_policy({"device": "global_info"})
+ @registry.direct_call
+ def constructor(device):
+ calls.append(device)
+ if device is invalid:
+ raise failure
+ return device
+
+ policy = registry.get_args_policy(constructor)
+ function_info = registry.function_info(language.X.script)
+ source = '@X.script\ndef main():\n X.record(constructor(device="mesh"))\n'
+ language.global_infos["mesh"] = first
+ assert language.parse(source, constructor=constructor).body == [("emit", first)]
+ language.global_infos["mesh"] = invalid
+ with pytest.raises(RuntimeError, match="constructor failed") as caught:
+ language.parse(source, constructor=constructor)
+ assert caught.value is failure
+ assert registry.get_args_policy(constructor) is policy
+ assert registry.function_info(language.X.script) is function_info
+ assert registry.is_direct_call(constructor)
+ language.global_infos["mesh"] = second
+ assert language.parse(source, constructor=constructor).body == [("emit", second)]
+ assert calls == [first, invalid, second]
+ assert registry.get_args_policy(constructor) is policy
+ assert registry.function_info(language.X.script) is function_info
+ assert registry.is_direct_call(constructor)
diff --git a/tests/python/script/test_parser.py b/tests/python/script/test_parser.py
index 9dd0212945..8954527149 100644
--- a/tests/python/script/test_parser.py
+++ b/tests/python/script/test_parser.py
@@ -24,8 +24,8 @@ from types import SimpleNamespace
import pytest
from dummy_builder import Value
-from tvm.script.ir_builder.ir import parser_protocol as protocol
from tvm.script.parser import entry
+from tvm.script.parser import protocol_registry as registry
@pytest.mark.parametrize(
@@ -131,27 +131,27 @@ def test_policy_registration_is_immutable_and_aliases_keep_identity():
# Before: alias = constructor; alias(shape=("n",))
# Expected builder program: alias(shape=(X.resolve_type_var_("n"),))
fields = {"shape": "expr_str", "device": "global_info"}
- decorate = protocol.args_policy(fields)
+ decorate = registry.args_policy(fields)
fields["shape"] = "global_info"
def constructor(shape, device):
return shape, device
alias = decorate(constructor)
- policy = protocol.get_args_policy(alias)
- assert policy is protocol.get_args_policy(constructor)
+ policy = registry.get_args_policy(alias)
+ assert policy is registry.get_args_policy(constructor)
assert dict(policy.fields) == {"shape": "expr_str", "device": "global_info"}
with pytest.raises(TypeError):
policy.fields["shape"] = "global_info"
- assert protocol.get_args_policy([]) is None
- assert protocol.get_args_policy(None) is None
+ assert registry.get_args_policy([]) is None
+ assert registry.get_args_policy(None) is None
- @protocol.args_policy({"values": "expr_str"}, scalar_strings=False)
+ @registry.args_policy({"values": "expr_str"}, scalar_strings=False)
def shorthand(values):
return values
- assert dict(protocol.get_args_policy(shorthand).fields) == {"values": "expr_str"}
- assert protocol.get_args_policy(shorthand).expression.scalar_strings is False
+ assert dict(registry.get_args_policy(shorthand).fields) == {"values": "expr_str"}
+ assert registry.get_args_policy(shorthand).expression.scalar_strings is False
@pytest.mark.parametrize(
@@ -165,7 +165,7 @@ def test_invalid_policy_registration_is_rejected(fields, message):
# Before: register an unknown policy or a nonexistent constructor parameter.
# Expected builder program: registration raises ValueError before parsing.
with pytest.raises(ValueError, match=message):
- protocol.args_policy(fields)(lambda shape: shape)
+ registry.args_policy(fields)(lambda shape: shape)
@pytest.mark.parametrize("marker", ["I.constexpr", "X.constexpr"])
@@ -331,7 +331,7 @@ def test_binding_declarations_override_mutable_targets_and_unpack_once(language,
}[target]
calls = []
- @protocol.register_binding_decl
+ @registry.register_binding_decl
def axes():
calls.append("axes")
return returned
@@ -403,7 +403,7 @@ def test_shadowed_binding_declaration_alias_uses_ordinary_assignment(language, s
# An ambient declaration policy cannot override a lexical callable binding.
marker, calls = object(), []
- @protocol.register_binding_decl
+ @registry.register_binding_decl
def axis_alias():
raise AssertionError("The shadowed ambient declaration must not run")
@@ -530,7 +530,7 @@ def test_policy_resolution_preserves_instance_and_namespace_callables(language):
seen = []
class Owner:
- @protocol.args_policy({"shape": "expr_str"})
+ @registry.args_policy({"shape": "expr_str"})
def constructor(self, shape):
seen.append(shape)
return shape
@@ -649,7 +649,7 @@ def main():
if marker(True):
X.record(1)
""",
- marker=protocol.constexpr,
+ marker=registry.constexpr,
)
@@ -811,7 +811,7 @@ def test_non_call_expression_reads_keep_their_source_range(language, expression)
@pytest.mark.parametrize("dtype", [None, "int32", "int64"])
def test_argument_policy_preserves_expression_dtype(language, dtype):
- @protocol.args_policy({"values": "expr_str"}, dtype=dtype)
+ @registry.args_policy({"values": "expr_str"}, dtype=dtype)
def shape(values):
return values
@@ -826,4 +826,4 @@ def main():
assert increment.op == "add"
assert increment.args[0] is n
assert increment.args[1] == 1
- assert protocol.get_args_policy(shape).expression.dtype == dtype
+ assert registry.get_args_policy(shape).expression.dtype == dtype
diff --git a/tests/python/script/test_parser_latest_design.py b/tests/python/script/test_parser_latest_design.py
index f8cea3847c..5c258f3e11 100644
--- a/tests/python/script/test_parser_latest_design.py
+++ b/tests/python/script/test_parser_latest_design.py
@@ -23,8 +23,8 @@ import ast
import pytest
from tvm.script import ir as I
-from tvm.script.ir_builder.ir import parser_protocol as protocol
from tvm.script.parser import entry
+from tvm.script.parser import protocol_registry as registry
@pytest.mark.parametrize("module", [False, True])
@@ -153,12 +153,12 @@ def test_direct_and_scope_declaration_calls_preserve_identity_once(language, cou
values = tuple(object() for _ in range(count))
calls = []
- @protocol.register_scope_var_query_or_decl
+ @registry.register_scope_var_query_or_decl
def declared():
calls.append("declared")
return values[0] if count == 1 else values
- @protocol.direct_call
+ @registry.direct_call
def direct(value):
calls.append("direct")
return value
@@ -190,9 +190,9 @@ def test_registered_global_callee_does_not_override_lexical_shadow(language, cat
pytest.fail("shadowed global callable executed")
if category == "direct":
- protocol.direct_call(operation)
+ registry.direct_call(operation)
else:
- protocol.register_scope_var_query_or_decl(operation)
+ registry.register_scope_var_query_or_decl(operation)
result = language.parse(
"@X.script\ndef main():\n operation = ordinary\n"
" value = operation()\n X.record(value)\n",
@@ -213,7 +213,7 @@ def test_direct_call_keeps_nested_expression_rewrites_and_call_scopes(language):
calls.append((value, len(language.source_stack)))
return value
- @protocol.direct_call
+ @registry.direct_call
def direct(value):
calls.append(("direct", len(language.source_stack)))
return value
@@ -258,7 +258,7 @@ def test_scope_declaration_reuses_symbol_spelling_but_plain_write_is_rejected(
# Expected: declaration precedence applies; a subsequent ordinary write remains illegal.
value = object()
- @protocol.register_scope_var_query_or_decl
+ @registry.register_scope_var_query_or_decl
def declared():
return value
diff --git a/tests/python/script/test_parser_shared_resolvers.py b/tests/python/script/test_parser_shared_resolvers.py
index 407311e6ca..a6e704e85b 100644
--- a/tests/python/script/test_parser_shared_resolvers.py
+++ b/tests/python/script/test_parser_shared_resolvers.py
@@ -24,7 +24,7 @@ from tvm.script import parser
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import ir as I
-from tvm.script.ir_builder.ir import parser_protocol as protocol
+from tvm.script.parser import protocol_registry as registry
def test_global_info_selectors_use_module_map():
@@ -43,7 +43,7 @@ def test_global_info_selectors_use_module_map():
def test_annotation_class_supports_union_and_missing_type():
# Before: X.Annotation("n") | None
# Expected builder program: X.Annotation((X.resolve_type_var_("n"))) in a function frame.
- @protocol.args_policy({"shape": "expr_str"}, as_type=True)
+ @registry.args_policy({"shape": "expr_str"}, as_type=True)
def annotation(shape):
return shape
diff --git a/tests/python/tirx/script/test_parser_buffer_views.py b/tests/python/tirx/script/test_parser_buffer_views.py
index 6c59cfb613..d4f58f09fc 100644
--- a/tests/python/tirx/script/test_parser_buffer_views.py
+++ b/tests/python/tirx/script/test_parser_buffer_views.py
@@ -25,7 +25,7 @@ from tvm import ir, tirx
from tvm.script import parser
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import ir as I
-from tvm.script.ir_builder.ir import parser_protocol
+from tvm.script.parser import protocol_registry as registry
from tvm.tirx import script as T
from tvm.tirx.layout import TCol, TLane
from tvm.tirx.script import builder as B
@@ -251,14 +251,14 @@ def test_direct_view_keeps_producer_identity_name_span_and_evaluation(
seen.append(name)
return value
- @I.direct_call
+ @registry.direct_call
def observe(view):
observed.append((view, view.span))
# A real bound method exercises registration lookup through __func__.
monkeypatch.setattr(type(captured), "operation", operation, raising=False)
bound = MethodType(operation, captured)
- assert parser_protocol.is_direct_call(bound)
+ assert registry.is_direct_call(bound)
result = parse(
f"""
@T.prim_func
diff --git a/tests/python/tirx/script/test_parser_tirx_resolvers.py b/tests/python/tirx/script/test_parser_tirx_resolvers.py
index cbd238b489..f56f0263dc 100644
--- a/tests/python/tirx/script/test_parser_tirx_resolvers.py
+++ b/tests/python/tirx/script/test_parser_tirx_resolvers.py
@@ -24,7 +24,7 @@ from tvm import ir
from tvm.relax.script import builder as R
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import ir as I
-from tvm.script.ir_builder.ir import parser_protocol as protocol
+from tvm.script.parser import protocol_registry as registry
from tvm.tirx.script import builder as T
@@ -68,7 +68,7 @@ def test_symbol_lookup_requires_function_and_reuses_native_identity(dialect):
def test_eager_annotations_preserve_missing_type_and_concrete_construction():
# Before: X.Tensor(("n",)); captured_shape = (n,); X.Tensor(captured_shape)
# Expected builder program: X.Tensor((X.resolve_type_var_("n"),)); X.Tensor(captured_shape)
- @protocol.args_policy({"shape": "expr_str"}, scalar_strings=False)
+ @registry.args_policy({"shape": "expr_str"}, scalar_strings=False)
def tensor(shape):
return shape
diff --git a/tests/python/tirx/script/test_tirx_builder_protocol_contract.py b/tests/python/tirx/script/test_tirx_builder_protocol_contract.py
index d3b04d12a0..ae389cf207 100644
--- a/tests/python/tirx/script/test_tirx_builder_protocol_contract.py
+++ b/tests/python/tirx/script/test_tirx_builder_protocol_contract.py
@@ -28,6 +28,7 @@ from tvm.script import tirx as T
from tvm.script.ir_builder import AlreadyEmitted, IRBuilder
from tvm.script.ir_builder import ir as IB
from tvm.script.ir_builder.ir import parser_protocol as P
+from tvm.script.parser import protocol_registry as registry
from tvm.tirx.script import builder as TB
from tvm.tirx.script.builder import parser_protocol as TP
@@ -35,12 +36,14 @@ from tvm.tirx.script.builder import parser_protocol as TP
def test_documented_two_function_program_builds_the_same_ir():
# The contract's complete source and builder examples must stay executable.
blocks = re.findall(r"\.\. code:: python\n\n((?: .*\n|\n)+)", P.__doc__)
- assert len(blocks) == 2
+ source_blocks = [block for block in blocks if "@I.ir_module" in block]
+ builder_blocks = [block for block in blocks if "with IRBuilder()" in block]
+ assert len(source_blocks) == len(builder_blocks) == 1
source = {"I": I, "T": T}
# Text parse supplies retrievable source coordinates for this documentation.
- module = tvm.script.from_source(textwrap.dedent(blocks[0]), extra_vars=source)
+ module = tvm.script.from_source(textwrap.dedent(source_blocks[0]), extra_vars=source)
generated = {"I": IB, "X": TB, "IRBuilder": IRBuilder}
- exec(textwrap.dedent(blocks[1]), generated)
+ exec(textwrap.dedent(builder_blocks[0]), generated)
tvm.ir.assert_structural_equal(module, generated["result"], map_free_vars=True)
assert not module["first"].params[0].ty.shape[0].same_as(module["second"].params[0].ty.shape[0])
@@ -59,7 +62,7 @@ def test_scope_category_keeps_explicit_bind_and_sequence_identity():
function = builder.get()
assert isinstance(function.body, tirx.Bind)
assert function.body.var.same_as(value)
- assert P.is_scope_var_query_or_decl(TB.bind)
+ assert registry.is_scope_var_query_or_decl(TB.bind)
def test_scope_category_preserves_pointer_bind_identity_and_ir():
@@ -117,7 +120,7 @@ def test_direct_object_producers_keep_constructor_names_and_identities():
produced = []
objects = []
- @P.direct_call
+ @registry.direct_call
def record(value):
objects.append(value)
return value
@@ -142,10 +145,10 @@ def test_direct_object_producers_keep_constructor_names_and_identities():
assert len(objects) == 2
assert objects[1].var.name == "explicit"
assert produced[0].buffer.name == "owned"
- assert P.is_direct_call(Resources)
- assert P.is_direct_call(TB.TileLayout)
- assert P.is_direct_call(TB.iter_var)
- assert P.is_direct_call(TB.TileLayout(TB.S[1]).canonicalize)
+ assert registry.is_direct_call(Resources)
+ assert registry.is_direct_call(TB.TileLayout)
+ assert registry.is_direct_call(TB.iter_var)
+ assert registry.is_direct_call(TB.TileLayout(TB.S[1]).canonicalize)
assert function is not None
@@ -154,7 +157,7 @@ def test_source_meta_var_keeps_existing_expression_span_without_a_binding():
existing = tirx.IntImm("int32", 7, span=span)
observed = []
- @P.direct_call
+ @registry.direct_call
def observe(value):
observed.append(value.span)
@@ -192,18 +195,18 @@ def test_result_member_registration_uses_real_callable_and_descriptor_identities
from tvm.tirx import _buffer_view
from tvm.tirx.buffer import _BufferMethods
- assert P.get_result_members(TB.Buffer) is _BufferMethods
- assert P.get_result_members(TB.alloc_buffer) is _BufferMethods
- assert P.get_result_members(_BufferMethods.view) is _BufferMethods
- assert P.get_result_members(_BufferMethods.sub) is _buffer_view.SubIndexer
- assert P.get_result_members(_BufferMethods.tile) is _buffer_view.TileIndexer
+ assert registry.get_result_members(TB.Buffer) is _BufferMethods
+ assert registry.get_result_members(TB.alloc_buffer) is _BufferMethods
+ assert registry.get_result_members(_BufferMethods.view) is _BufferMethods
+ assert registry.get_result_members(_BufferMethods.sub) is _buffer_view.SubIndexer
+ assert registry.get_result_members(_BufferMethods.tile) is _buffer_view.TileIndexer
for indexer in (_buffer_view.SubIndexer, _buffer_view.TileIndexer, _buffer_view.ChunkIndexer):
- assert P.is_direct_call(indexer.__getitem__)
- assert P.get_result_members(indexer.__getitem__) is _BufferMethods
+ assert registry.is_direct_call(indexer.__getitem__)
+ assert registry.get_result_members(indexer.__getitem__) is _BufferMethods
class Unrelated:
def view(self):
return self
- assert P.get_result_members(Unrelated.view) is None
- assert not P.is_direct_call(Unrelated.view)
+ assert registry.get_result_members(Unrelated.view) is None
+ assert not registry.is_direct_call(Unrelated.view)
diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py b/tests/python/tirx/transform/test_transform_lower_tirx.py
index 8adf7d5d16..3ec575800c 100644
--- a/tests/python/tirx/transform/test_transform_lower_tirx.py
+++ b/tests/python/tirx/transform/test_transform_lower_tirx.py
@@ -21,7 +21,7 @@ import tvm_ffi
import tvm
import tvm.testing
from tvm.script import tirx as T
-from tvm.script.ir_builder.ir.parser_protocol import register_mutable_var_decl
+from tvm.script.parser.protocol_registry import register_mutable_var_decl
from tvm.script.tirx import tile as Tx
from tvm.tirx.function import PrimFunc
from tvm.tirx.layout import laneid, warpid, wg_local_layout