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 cfda25480ded71e7a927c253cb74714e97ef2a85 Author: Tianqi Chen <[email protected]> AuthorDate: Wed Sep 23 04:12:51 2026 +0000 Keep native module aliases and shared resolution direct --- python/tvm/relax/script/builder/__init__.py | 3 ++ python/tvm/relax/script/builder/protocol.py | 3 ++ python/tvm/script/ir_builder/ir/frame.py | 32 +++++++--------------- python/tvm/script/ir_builder/ir/protocol.py | 41 +++++++++++++++++++++++----- python/tvm/tirx/script/builder/__init__.py | 3 ++ python/tvm/tirx/script/builder/protocol.py | 8 ++++-- tests/python/unittest/test_parser_support.py | 7 ++--- 7 files changed, 62 insertions(+), 35 deletions(-) diff --git a/python/tvm/relax/script/builder/__init__.py b/python/tvm/relax/script/builder/__init__.py index 820e393921..a3ee4b15ac 100644 --- a/python/tvm/relax/script/builder/__init__.py +++ b/python/tvm/relax/script/builder/__init__.py @@ -59,6 +59,9 @@ from .protocol import emit_ as emit_ from .protocol import resolve_type_var_ as resolve_type_var_ from .protocol import set_mutable_var_ as set_mutable_var_ +# Syntax capability: mutable declaration policies apply only in this dialect. +supports_mutable_declarations = False + @_args_policy({"shape": "expr_str", "vdevice": "global_info"}, scalar_strings=False) def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None): diff --git a/python/tvm/relax/script/builder/protocol.py b/python/tvm/relax/script/builder/protocol.py index 8aa078b293..fe7ea7cf20 100644 --- a/python/tvm/relax/script/builder/protocol.py +++ b/python/tvm/relax/script/builder/protocol.py @@ -46,6 +46,9 @@ def bind_( """Emit a Relax binding or retain a named frame-owned value.""" if isinstance(value, _BypassBind): return value.value + if isinstance(value, _I.IRModuleFrame): + # A source class alias retains its active module; it is not a new scope. + return value name_span = _source_span(span if name_span is None else name_span) if frame_value: if isinstance(value, _python.list | _python.tuple | _ir.Array): diff --git a/python/tvm/script/ir_builder/ir/frame.py b/python/tvm/script/ir_builder/ir/frame.py index 87204fb398..346ec00c62 100644 --- a/python/tvm/script/ir_builder/ir/frame.py +++ b/python/tvm/script/ir_builder/ir/frame.py @@ -16,8 +16,6 @@ # under the License. """Package tvm.script.ir_builder.ir.frame""" -import re - from tvm_ffi import register_object as _register_object from ..base import IRBuilderFrame @@ -25,23 +23,13 @@ from ..base import IRBuilderFrame @_register_object("script.ir_builder.IRModuleFrame") class IRModuleFrame(IRBuilderFrame): - def resolve_global_info(self, content): - """Resolve a map index or target-kind device ordinal in this module.""" - if not isinstance(content, str): - return content - match = re.fullmatch(r"([^\[\]]+)\[(\d+)\]", content) - if match: - name, index = match.groups() - return self.global_infos[name][int(index)] - selector = re.fullmatch(r"([^:\[\]]+)(?::(\d+)(?::([^:]+))?)?", content) - if selector is None: - raise ValueError(f"Invalid global-info reference: {content!r}") - target, index, _scope = selector.groups() - ordinal = int(index) if index is not None else 0 - devices = self.global_infos.get("vdevice", ()) - if target == "vdevice": - return devices[ordinal] - devices = [device for device in devices if device.target.kind.name == target] - if ordinal >= len(devices): - raise ValueError(f"Global-info device reference was not found: {content!r}") - return devices[ordinal] + def __getattr__(self, name): + """Expose a declared function through the native module's reference map. + + Reflected native fields and Python methods resolve before this fallback. + Source aliases therefore retain this frame and its plain GlobalVars. + """ + try: + return self.global_vars[name] + except KeyError: + raise AttributeError(f"IRModuleFrame has no attribute {name!r}") from None diff --git a/python/tvm/script/ir_builder/ir/protocol.py b/python/tvm/script/ir_builder/ir/protocol.py index 8baa44a7bf..61b161193b 100644 --- a/python/tvm/script/ir_builder/ir/protocol.py +++ b/python/tvm/script/ir_builder/ir/protocol.py @@ -21,9 +21,14 @@ ordinary Python. Its state describes syntax, never IR values. Generated code uses ``I`` for shared module and source-location support and ``X`` for the current dialect's builders and protocol hooks. Native function frames own parameters, symbols and completed functions; module and region frames own their -references and results. A declaration entry reserves each function signature -before any body entry, allowing sibling calls without another ownership record. +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. Syntax policy registration belongs to ``tvm.script.parser.protocol``. +The dialect namespace advertises ``supports_mutable_declarations`` so a +primitive annotation imported from another dialect does not grant mutable +storage syntax to a dialect whose bindings are immutable. For example, source functions can share a symbolic shape spelling while each function retains its own symbol identity: @@ -75,6 +80,8 @@ The function stubs below document dialect hooks, implemented in each dialect's runtime dispatcher. ``resolve_global_info`` is the real shared ``I`` operation. """ +import re + from ..base import MISSING, IRBuilder from .frame import IRModuleFrame @@ -88,11 +95,31 @@ def resolve_global_info(content): """ if not isinstance(content, str): return content - if IRBuilder.is_in_scope(): - for frame in reversed(IRBuilder.current().frames): - if isinstance(frame, IRModuleFrame): - return frame.resolve_global_info(content) - raise ValueError("Global-info lookup requires an enclosing module frame") + if not IRBuilder.is_in_scope(): + raise ValueError("Global-info lookup requires an enclosing module frame") + for frame in reversed(IRBuilder.current().frames): + if isinstance(frame, IRModuleFrame): + break + else: + raise ValueError("Global-info lookup requires an enclosing module frame") + match = re.fullmatch(r"([^\[\]]+)\[(\d+)\]", content) + if match: + name, index = match.groups() + return frame.global_infos[name][int(index)] + selector = re.fullmatch(r"([^:\[\]]+)(?::(\d+)(?::([^:]+))?)?", content) + if selector is None: + raise ValueError(f"Invalid global-info reference: {content!r}") + target, index, _scope = selector.groups() + ordinal = int(index) if index is not None else 0 + devices = frame.global_infos.get("vdevice", ()) + if target == "vdevice": + return devices[ordinal] + for device in devices: + if device.target.kind.name == target: + if ordinal == 0: + return device + ordinal -= 1 + raise ValueError(f"Global-info device reference was not found: {content!r}") def resolve_type_var_(name, dtype=None, *, value=None, span=None): diff --git a/python/tvm/tirx/script/builder/__init__.py b/python/tvm/tirx/script/builder/__init__.py index 985df14a63..ff13d25ad5 100644 --- a/python/tvm/tirx/script/builder/__init__.py +++ b/python/tvm/tirx/script/builder/__init__.py @@ -65,6 +65,9 @@ from .utils import buffer_proxy as buffer_proxy from .utils import frame_scope as frame_scope from .utils import seq_scope as seq_scope +# Syntax capability: mutable declaration policies apply only in this dialect. +supports_mutable_declarations = True + is_type_var = _ir.is_prim_var diff --git a/python/tvm/tirx/script/builder/protocol.py b/python/tvm/tirx/script/builder/protocol.py index dc79b0f33a..aa9b3beecf 100644 --- a/python/tvm/tirx/script/builder/protocol.py +++ b/python/tvm/tirx/script/builder/protocol.py @@ -62,10 +62,14 @@ def bind_( if not value.value.name and name is not None: _IRBuilder.name(name, value.value) return value.value + if isinstance(value, _I.IRModuleFrame): + # A source class alias retains its active module; it is not a new scope. + return value name_span = span if name_span is None else name_span # Axis and environment-thread variables already belong to native frames; - # naming an assignment must preserve their registration identities. - if not frame_value and _ir.is_prim_var(value): + # naming an assignment must preserve their registration identities. An + # explicit T.let annotation instead requests a distinct immutable binding. + if not frame_value and not isinstance(ty, _native.LetAnnotation) and _ir.is_prim_var(value): for frame in reversed(_IRBuilder.current().frames): if isinstance(frame, _frame.SBlockFrame) and _python.any( axis.var.same_as(value) for axis in frame.iter_vars diff --git a/tests/python/unittest/test_parser_support.py b/tests/python/unittest/test_parser_support.py index eb0b07d99d..3eef773a72 100644 --- a/tests/python/unittest/test_parser_support.py +++ b/tests/python/unittest/test_parser_support.py @@ -52,10 +52,9 @@ def test_global_info_selectors_use_module_map(): devices = [I.vdevice("llvm"), I.vdevice("cuda"), I.vdevice("cuda", 1)] I.module_global_infos({"vdevice": devices, "other": [I.dummy_global_info()]}) for spelling, index in (("cuda:1", 2), ("vdevice[1]", 1), ("cuda", 1)): - assert I.resolve_global_info(spelling).__chandle__() == devices[index].__chandle__() - assert ( - module.resolve_global_info(spelling).__chandle__() == devices[index].__chandle__() - ) + resolved = I.resolve_global_info(spelling) + assert resolved.__chandle__() == devices[index].__chandle__() + assert resolved.__chandle__() == module.global_infos["vdevice"][index].__chandle__() assert I.resolve_global_info("other[0]") is not None
