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
The following commit(s) were added to refs/heads/script/canonical-parser-df by
this push:
new 7505c9a88d [REFACTOR][TVMSCRIPT] Let dialects resolve global metadata
7505c9a88d is described below
commit 7505c9a88d43b97931cfbb254b0b42cefa09ca44
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 24 13:15:14 2026 +0000
[REFACTOR][TVMSCRIPT] Let dialects resolve global metadata
Global metadata selectors belong to the language that defines their meaning.
Dispatch registered global-info arguments through the active parser
protocol,
retain Relax selector behavior, and make unsupported TIRx lookup explicit.
Introduce a modular, customizable parser infrastructure for TVMScript
through transpilation. The key idea is to deterministically transpile language
syntax into a corresponding sequence of builder calls while preserving source
locations and providing span context.
Key design elements include:
- Acquire source and its definition scope, collect syntax facts once,
rewrite the owned AST, and execute the resulting builder program with the
original lexical bindings.
- Each language variant customizes behavior through a `parser_protocol.py`
file, which defines special handling for syntax such as `for` and `if`.
- Create reusable span entries during rewriting. Attach locations to
meaningful expressions and pass spans directly to builder operations that own
them.
- Reuse existing builder and IR APIs, so new types and functions can be
added without changing the parser, provided they have well-defined builder
behavior.
- Generate specialization handling only when specialization is requested.
For example, this kernel:
```python
@T.prim_func
def add_one(A: T.Buffer((128,), "float32"), B: T.Buffer((16, 8),
"float32")):
for i in range(128):
B[i // 8, i % 8] = A[i] + 1
```
is lowered to the following builder calls. This presentation wraps the
emitted program in `build_add_one` and gives its local scaffolding readable
aliases. `I` supplies shared builders, `T` is the dialect namespace, `_S` holds
source-span entries, and `_host0` and `_host1` are captured annotation helpers.
```python
def build_add_one():
with I.IRBuilder() as builder:
with T.function_(span=_S[6]) as fn:
T.func_name('add_one')
_parameter0 = T.arg('A', _host0(_S[1](T.Buffer(_S[0]((128,)),
'float32'))), span=_S[2])
_parameter1 = T.arg('B', _host1(_S[4](T.Buffer(_S[3]((16, 8)),
'float32'))), span=_S[5])
def build_body():
global T
A, B = fn.params
with T.for_(T.range_(128), span=_S[7], names='i') as i:
_value0 = _S[9](_S[8](A[i]) + 1)
T.setitem_(B, (_S[10](i // 8), _S[11](i % 8)), _value0,
span=_S[12])
build_body()
result = fn.function
result.__name__ = 'add_one'
T.check_well_formed_(result)
return fn.function
```
---
python/tvm/relax/script/builder/__init__.py | 7 +--
python/tvm/relax/script/builder/parser_protocol.py | 60 ++++++++++++++++++++++
python/tvm/script/ir_builder/ir/__init__.py | 2 +-
python/tvm/script/ir_builder/ir/parser_protocol.py | 53 ++++---------------
python/tvm/script/parser/transpile.py | 6 +--
python/tvm/tirx/script/builder/__init__.py | 1 +
python/tvm/tirx/script/builder/parser_protocol.py | 8 +++
tests/python/script/minilang.py | 2 +-
.../python/script/test_special_parser_protocol.py | 53 +++++++++++++++++--
9 files changed, 138 insertions(+), 54 deletions(-)
diff --git a/python/tvm/relax/script/builder/__init__.py
b/python/tvm/relax/script/builder/__init__.py
index 6d3f1eed09..3d8dbc1d41 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/relax/script/builder/__init__.py
@@ -30,7 +30,6 @@ 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 resolve_global_info_ as
_lookup_global_info
from tvm.script.parser.protocol_registry import ARGS_POLICIES as _ARGS_POLICIES
from tvm.script.parser.protocol_registry import args_policy as _args_policy
from tvm.script.parser.protocol_registry import constexpr as constexpr
@@ -68,6 +67,7 @@ from .parser_protocol import (
not_,
or_,
range_,
+ resolve_global_info_,
resolve_type_var_,
return_,
set_mutable_cell_,
@@ -92,7 +92,7 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *,
span=None):
# Python evaluates annotations before the module decorator opens its
# frame. The generated declaration resolves this module-owned
reference.
return _ir.Type.missing()
- vdevice = _lookup_global_info(vdevice)
+ vdevice = resolve_global_info_(vdevice)
return _relax.TensorType(shape, dtype, vdevice, ndim, _source_span(span))
@@ -106,7 +106,7 @@ def DTensor(shape=None, dtype=None, device_mesh=None,
placement="", *, ndim=-1,
if device_mesh is None:
device_mesh = _DeviceMesh([], _ir.Range(0, 1))
else:
- device_mesh = _lookup_global_info(device_mesh)
+ device_mesh = resolve_global_info_(device_mesh)
if isinstance(placement, _python.str):
placement = _Placement.from_text(placement)
return _DTensorType(Tensor(shape, dtype, ndim=ndim), device_mesh,
placement, _source_span(span))
@@ -232,6 +232,7 @@ __all__ = [
"if_",
"check_well_formed_",
"resolve_type_var_",
+ "resolve_global_info_",
"call_global_var_",
"decl_mutable_cell_",
"set_mutable_cell_",
diff --git a/python/tvm/relax/script/builder/parser_protocol.py
b/python/tvm/relax/script/builder/parser_protocol.py
index 2a88efc25e..9a27a03a77 100644
--- a/python/tvm/relax/script/builder/parser_protocol.py
+++ b/python/tvm/relax/script/builder/parser_protocol.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import builtins as _python
import numbers as _numbers
+import re as _re
from collections.abc import Sequence
from typing import Any, NoReturn
@@ -31,6 +32,7 @@ from tvm import tirx as _tir
from tvm.ir.prim import _ffi_api as _prim_ffi
from tvm.script.ir_builder import IRBuilder as _IRBuilder
from tvm.script.ir_builder import base as _base
+from tvm.script.ir_builder.ir.frame import IRModuleFrame as _IRModuleFrame
from .. import builder as _builder
from . import _ffi_api
@@ -40,6 +42,64 @@ from . import ir as _native
_Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int]
| None
+def resolve_global_info_(content: Any) -> Any:
+ """Resolve a module-owned selector or retain a concrete object.
+
+ Parameters
+ ----------
+ content : str or Any
+ Original global-info selector or concrete value. "mesh[0]" indexes a
named list;
+ "cuda:1" selects the second CUDA vdevice; "vdevice:0" selects by
absolute index.
+ A trailing memory-scope suffix is accepted without changing device
selection.
+
+ Returns
+ -------
+ Any
+ The exact registered global-info object, or the unchanged non-string
input.
+
+ Notes
+ -----
+ String lookup requires the nearest active native module frame and creates
no metadata.
+ Missing context, malformed selectors or unmatched devices raise
ValueError; missing map
+ entries or out-of-range indices propagate KeyError/IndexError. Non-string
values require
+ no frame. No source span is attached to an existing metadata object.
+
+ .. code:: python
+
+ # Source
+ R.Tensor((n,), "float32", vdevice="cuda:0")
+ # Generated builder
+ R.Tensor((n,), "float32", vdevice=R.resolve_global_info_("cuda:0"))
+ """
+ if not isinstance(content, str):
+ return content
+ 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 bind_(
value: Any = _base.MISSING,
*,
diff --git a/python/tvm/script/ir_builder/ir/__init__.py
b/python/tvm/script/ir_builder/ir/__init__.py
index f38bc481e7..1fd0646e98 100644
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ b/python/tvm/script/ir_builder/ir/__init__.py
@@ -21,7 +21,7 @@ from tvm.ir import GenericConst, Range, StringImm, StringType
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_
+from .parser_protocol import check_well_formed_, module_member_
from .frame import IRModuleFrame
from .ir import (
decl_function,
diff --git a/python/tvm/script/ir_builder/ir/parser_protocol.py
b/python/tvm/script/ir_builder/ir/parser_protocol.py
index 39779390f7..f781895ff1 100644
--- a/python/tvm/script/ir_builder/ir/parser_protocol.py
+++ b/python/tvm/script/ir_builder/ir/parser_protocol.py
@@ -81,74 +81,43 @@ They specify a structural contract, without a base class or
runtime dispatcher.
from __future__ import annotations
-import re
from collections.abc import Sequence
from typing import Any
from tvm import ir as _ir
-from ..base import MISSING, IRBuilder, IRBuilderFrame, SpanEntry
-from .frame import IRModuleFrame
+from ..base import MISSING, IRBuilderFrame, SpanEntry
_Span = SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int] | None
def resolve_global_info_(content: Any) -> Any:
- """Resolve a module-owned selector or retain a concrete object.
+ """Look up language-owned global metadata selected by a source argument.
Parameters
----------
content : str or Any
- Original global-info selector or concrete value. "mesh[0]" indexes a
named list;
- "cuda:1" selects the second CUDA vdevice; "vdevice:0" selects by
absolute index.
- A trailing memory-scope suffix is accepted without changing device
selection.
+ A selector or concrete value interpreted by the active language.
Returns
-------
Any
- The exact registered global-info object, or the unchanged non-string
input.
+ The language's metadata value. Selector syntax, lookup context and
+ unsupported-value errors belong to the dialect implementation.
Notes
-----
- String lookup requires the nearest active native module frame and creates
no metadata.
- Missing context, malformed selectors or unmatched devices raise
ValueError; missing map
- entries or out-of-range indices propagate KeyError/IndexError. Non-string
values require
- no frame. No source span is attached to an existing metadata object.
+ A string argument marked ``global_info`` is evaluated once and passed to
+ this hook without attaching a new span to the existing metadata object.
.. code:: python
- # Source
- R.Tensor((n,), "float32", vdevice="cuda:0")
+ # Source: an argument policy marks metadata as global_info.
+ X.CustomType(metadata="mesh[0]")
# Generated builder
- X.Tensor((n,), "float32", vdevice=I.resolve_global_info_("cuda:0"))
+ X.CustomType(metadata=X.resolve_global_info_("mesh[0]"))
"""
- if not isinstance(content, str):
- return content
- 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}")
+ raise NotImplementedError
def resolve_type_var_(
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index a68f286714..b71a67e95c 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -616,11 +616,11 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# X.Tensor(vdevice="cuda:1")
#
# Builder:
- # X.Tensor(vdevice=I.resolve_global_info_("cuda:1"))
+ # X.Tensor(vdevice=X.resolve_global_info_("cuda:1"))
# -------------------------------------------------
value = self.visit(node)
return self._call(
- self.module.infrastructure_name, "resolve_global_info_",
[value], node
+ self.function.dialect_prefix, "resolve_global_info_",
[value], node
)
return self.visit(node)
@@ -649,7 +649,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# X.Tensor(("n",), vdevice="cuda:0")
#
# Builder:
- # X.Tensor((X.resolve_type_var_("n"),),
vdevice=I.resolve_global_info_("cuda:0"))
+ # X.Tensor((X.resolve_type_var_("n"),),
vdevice=X.resolve_global_info_("cuda:0"))
# -------------------------------------------------
binding_value = node is self.binding_expression
marker = self._read_constexpr_operand(node)
diff --git a/python/tvm/tirx/script/builder/__init__.py
b/python/tvm/tirx/script/builder/__init__.py
index df83450d41..cb510ad7f3 100644
--- a/python/tvm/tirx/script/builder/__init__.py
+++ b/python/tvm/tirx/script/builder/__init__.py
@@ -69,6 +69,7 @@ from .parser_protocol import (
not_,
or_,
range_,
+ resolve_global_info_,
resolve_type_var_,
return_,
set_mutable_cell_,
diff --git a/python/tvm/tirx/script/builder/parser_protocol.py
b/python/tvm/tirx/script/builder/parser_protocol.py
index e17e19b402..0d6fbc4d7b 100644
--- a/python/tvm/tirx/script/builder/parser_protocol.py
+++ b/python/tvm/tirx/script/builder/parser_protocol.py
@@ -37,6 +37,14 @@ from . import ir as _native
_Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int]
| None
+def resolve_global_info_(content: Any) -> Any:
+ """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.resolve_global_info_`.
+
+ TIRx does not define global-info selectors.
+ """
+ raise NotImplementedError("TIRx does not support global-info lookup")
+
+
def _name(value: Any, name: str | None, span: _Span) -> Any:
if name is not None:
_IRBuilder.name(name, value)
diff --git a/tests/python/script/minilang.py b/tests/python/script/minilang.py
index 8ae88c245d..68cc30a5bf 100644
--- a/tests/python/script/minilang.py
+++ b/tests/python/script/minilang.py
@@ -157,7 +157,6 @@ class Language:
ir_module=lambda: Frame(self, "module"),
at_=self.at,
with_at_group_=self.with_at_group,
- resolve_global_info_=self.resolve_global_info,
module_member_=lambda name, value: value,
require_defined=self.require_defined,
annotation_value_=lambda name, value: value,
@@ -182,6 +181,7 @@ class Language:
break_=lambda **kwargs: self.statement("break"),
continue_=lambda **kwargs: self.statement("continue"),
resolve_type_var_=self.resolve_type_var,
+ resolve_global_info_=self.resolve_global_info,
bind_=self.bind,
check_well_formed_=lambda result: None,
emit_=self.emit,
diff --git a/tests/python/script/test_special_parser_protocol.py
b/tests/python/script/test_special_parser_protocol.py
index abb0d8ee1f..ed14fe7eb1 100644
--- a/tests/python/script/test_special_parser_protocol.py
+++ b/tests/python/script/test_special_parser_protocol.py
@@ -18,15 +18,20 @@
A failed constructor must not change the next function's interpretation of
literals.
"""
+import inspect
+import traceback
+
import pytest
+from tvm.script import tirx as T
from tvm.script.parser import protocol_registry as registry
def test_constructor_policy_survives_a_failed_definition(language):
- # A failing definition must not erase registered argument semantics for
the next one.
+ # A registered argument uses the dialect lookup, even after a constructor
failure.
M = language.M
calls = []
+ dialect_infos = {}
first, invalid, second = object(), object(), object()
failure = RuntimeError("constructor failed")
@@ -37,6 +42,10 @@ def
test_constructor_policy_survives_a_failed_definition(language):
raise failure
return device
+ def resolve(name):
+ return dialect_infos[name]
+
+ M.resolve_global_info_ = resolve
M.constructor = constructor
def build():
@@ -47,13 +56,13 @@ def
test_constructor_policy_survives_a_failed_definition(language):
return main
- language.global_infos["mesh"] = first
+ dialect_infos["mesh"] = first
assert build().body == [("emit", first), ("emit", "mesh")]
- language.global_infos["mesh"] = invalid
+ dialect_infos["mesh"] = invalid
with pytest.raises(RuntimeError, match="constructor failed") as caught:
build()
assert caught.value is failure
- language.global_infos["mesh"] = second
+ dialect_infos["mesh"] = second
assert build().body == [("emit", second), ("emit", "mesh")]
assert calls == [first, "mesh", invalid, second, "mesh"]
@@ -77,3 +86,39 @@ def
test_argument_policy_preserves_expression_dtype(language):
assert increment.op == "add"
assert increment.args[0] is n
assert increment.args[1] == 1
+
+
+def test_tirx_rejects_global_info_at_the_source_argument(monkeypatch):
+ # A custom global-info argument must report TIRx's unsupported policy at
its source.
+ @registry.args_policy("T.global_annotation", {"device": "global_info"})
+ def global_annotation(device):
+ return T.int32
+
+ monkeypatch.setattr(T, "global_annotation", global_annotation,
raising=False)
+ with pytest.raises(
+ NotImplementedError, match="TIRx does not support global-info lookup"
+ ) as caught:
+
+ @T.prim_func
+ def main(value: T.global_annotation(device="cuda:0")):
+ T.evaluate(value)
+
+ lines, first =
inspect.getsourcelines(test_tirx_rejects_global_info_at_the_source_argument)
+ index, line = next((i, line) for i, line in enumerate(lines) if "def
main(value:" in line)
+ location = first + index
+ frames = traceback.extract_tb(caught.value.__traceback__)
+ source_frames = [
+ frame for frame in frames if frame.filename == __file__ and
frame.lineno == location
+ ]
+ assert source_frames
+ if getattr(source_frames[-1], "colno", None) is not None:
+ column = line.index('"cuda:0"')
+ assert (
+ source_frames[-1].colno,
+ source_frames[-1].end_lineno,
+ source_frames[-1].end_colno,
+ ) == (
+ column,
+ location,
+ column + len('"cuda:0"'),
+ )
