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 <[email protected]> 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 [email protected]("protocol", [P, TP, RP]) [email protected]("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
