This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tvmscript-ast-only-transpiler
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/tvmscript-ast-only-transpiler
by this push:
new c6ba25159f Resolve named outputs through explicit completed builder
frames
c6ba25159f is described below
commit c6ba25159f25ea039172859667f5de2fbb78fe0e
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Sep 21 17:59:38 2026 +0000
Resolve named outputs through explicit completed builder frames
---
python/tvm/relax/script/builder/__init__.py | 20 +++++++++++
python/tvm/script/ir_builder/protocol.py | 16 +++++----
python/tvm/script/parser/transpile.py | 53 +++++++++++++----------------
python/tvm/tirx/script/builder/__init__.py | 18 ++++++++++
tests/python/tvmscript/test_parser.py | 36 ++++++++++++++++++++
5 files changed, 108 insertions(+), 35 deletions(-)
diff --git a/python/tvm/relax/script/builder/__init__.py
b/python/tvm/relax/script/builder/__init__.py
index 3a5b83058e..7607bebffa 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/relax/script/builder/__init__.py
@@ -223,6 +223,25 @@ class _Frame:
return False
+def frame_result(completed_frame, name):
+ """Return one supported output of an explicitly supplied completed region.
+
+ completed_frame is a completed Relax builder context or a host Python
branch result
+ dictionary; name is the original source variable-name string. A conditional
+ exports its native designated result, whose final binding name matches in
+ both branches. Dataflow frames export their declared outputs; local
function
+ frames export their function reference. Internal branch bindings and names
+ owned by other frames are not exposed. Host dictionaries contain only the
+ lexical outputs explicitly supplied by their executed branch.
+
+ Returns the exported value, including None, or protocol.MISSING when
absent.
+ A non-string name raises TypeError; custom frame/result access errors
+ propagate. The operation is read-only, enters no frames, and retains no
+ implicit completed-frame state, so nested frame results remain independent.
+ """
+ return _protocol._frame_result(completed_frame, name)
+
+
def function(is_pure=True, is_private=False, *, local=False, reference=None,
span=None):
"""Enter a definition using the native Relax function frame."""
if local:
@@ -447,6 +466,7 @@ __all__ = [
"DTensor",
"For",
"for_",
+ "frame_result",
"resolve_type_var",
"Object",
"Prim",
diff --git a/python/tvm/script/ir_builder/protocol.py
b/python/tvm/script/ir_builder/protocol.py
index ff70228b19..ebfd9d110f 100644
--- a/python/tvm/script/ir_builder/protocol.py
+++ b/python/tvm/script/ir_builder/protocol.py
@@ -200,14 +200,18 @@ def at(span, value):
_at = at
-def frame_result(frame):
- """Return frame.result, or an empty mapping when no result is exposed.
+def _frame_result(frame, name):
+ """Read one explicit export without consulting ambient construction state.
- frame is a finalized builder context object. The result maps source names
- to exported values and is owned by the frame. This read-only operation
- enters no frames; custom attribute-access errors propagate.
+ Dialects pass a completed frame whose result is its finalized export map,
+ or a Python-branch dictionary containing host lexical results. The map
stays
+ private to its owner; a missing name returns MISSING instead of leaking a
+ branch-local binding. No frame is entered and no result is cached globally.
"""
- return getattr(frame, "result", {})
+ if not isinstance(name, str):
+ raise TypeError("A frame result name must be a string")
+ exports = frame if isinstance(frame, dict) else getattr(frame, "result",
{})
+ return exports.get(name, MISSING)
def require_defined(value, name):
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index e78d9fbec7..01394c8c59 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -148,7 +148,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
namespaces/callables; assignments and parameters shadow entries.
dialect_prefix
selects the generated current builder alias and is restored after
functions.
bound holds only Python identifier strings. optional maps
- conditionally exported names to AST reads of builder result dictionaries;
+ conditionally exported names to AST reads of cached named builder outputs;
for example a name defined only by one branch is checked when read. Each
lexical helper copies/restores bound/namespace_bindings/optional; no
concrete values
or IR identity are stored. preserve_return and syntax hooks configure one
@@ -726,43 +726,40 @@ class IRBuilderTranspiler(ast.NodeTransformer):
invocation = self._located(ast.Call(self._name(helper, original), [],
[]), original)
return [definition, self._statement(invocation, original)]
- def _exports(self, frame, candidates, original, mapping=None):
+ def _exports(self, frame, candidates, original):
+ # Pattern: names assigned in a completed region -> one named builder
+ # request per candidate, e.g. y = X.frame_result(frame, "y"). Builders
+ # decide which names are designated outputs; absent results remain
+ # MISSING so reads cannot expose branch-local bindings accidentally.
result = []
- if mapping is None:
- mapping_stmt, mapping = self._cache(
+ for name in sorted(candidates):
+ assignment, value = self._cache(
self._call(
- self.infrastructure_name,
+ self.dialect_prefix,
"frame_result",
- [self._name(frame, original)],
+ [self._name(frame, original), ast.Constant(name)],
original,
),
original,
- "exports",
+ "output",
)
- result.append(mapping_stmt)
- for name in sorted(candidates):
- key = ast.Constant(name)
+ result.append(assignment)
condition = self._located(
- ast.Compare(copy.deepcopy(key), [ast.In()],
[copy.deepcopy(mapping)]), original
+ ast.Compare(
+ copy.deepcopy(value),
+ [ast.IsNot()],
+ [self._attribute(self.infrastructure_name, "MISSING",
original)],
+ ),
+ original,
)
- value = self._located(ast.Subscript(copy.deepcopy(mapping), key,
ast.Load()), original)
result.append(
self._located(
- ast.If(condition, [self._assign(name, value, original)],
[]), original
+ ast.If(condition, [self._assign(name,
copy.deepcopy(value), original)], []),
+ original,
)
)
- for name in candidates - self.bound:
- self.optional[name] = self._located(
- ast.Call(
- ast.Attribute(copy.deepcopy(mapping), "get", ast.Load()),
- [
- ast.Constant(name),
- self._attribute(self.infrastructure_name, "MISSING",
original),
- ],
- [],
- ),
- original,
- )
+ if name not in self.bound:
+ self.optional[name] = copy.deepcopy(value)
return result
def _with(self, context, body, original, target=None):
@@ -814,13 +811,11 @@ class IRBuilderTranspiler(ast.NodeTransformer):
native.append(
self._assign(
mapping_name,
- self._call(
- self.infrastructure_name, "frame_result",
[self._name(frame, node)], node
- ),
+ self._name(frame, node),
node,
)
)
- tail = self._exports(frame, candidates, node,
self._name(mapping_name, node))
+ tail = self._exports(mapping_name, candidates, node)
dispatch = self._located(
ast.If(
self._call(
diff --git a/python/tvm/tirx/script/builder/__init__.py
b/python/tvm/tirx/script/builder/__init__.py
index 10537d2d08..fb332d95c1 100644
--- a/python/tvm/tirx/script/builder/__init__.py
+++ b/python/tvm/tirx/script/builder/__init__.py
@@ -32,6 +32,7 @@ from tvm.script.ir_builder import ir as _I
from tvm.script.ir_builder.base import IRBuilderFrame as _NativeFrame
from tvm.script.ir_builder.protocol import MISSING as _MISSING
from tvm.script.ir_builder.protocol import _construction_span
+from tvm.script.ir_builder.protocol import _frame_result as _named_frame_result
from tvm.script.ir_builder.protocol import at as _at
from tvm.script.ir_builder.protocol import register_call_kind as
_register_call_kind
from tvm.script.ir_builder.protocol import source_span as _source_span
@@ -175,6 +176,23 @@ class _Frame:
return getattr(self.native, name)
+def frame_result(completed_frame, name):
+ """Return one named export from an explicitly supplied completed region.
+
+ completed_frame is a completed builder context or the result dictionary
from a host
+ Python branch; name is the original source variable-name string. Returns
+ its exported value, including None, or protocol.MISSING when not exported.
+ TIR statement frames have no implicit lexical exports, so their internal
+ bindings remain unavailable through this operation. Host branch
dictionaries
+ expose only the names explicitly supplied by that branch.
+
+ A non-string name raises TypeError; custom frame/result access errors
+ propagate. This read-only operation enters no frames, modifies no result
+ maps, and retains no last-completed-frame or cross-region state.
+ """
+ return _named_frame_result(completed_frame, name)
+
+
def function(*, private=False, s_tir=False, persistent=False, span=None):
"""Enter a native primitive-function definition frame."""
with _construction_span(span):
diff --git a/tests/python/tvmscript/test_parser.py
b/tests/python/tvmscript/test_parser.py
index a029cadc52..46667cac02 100644
--- a/tests/python/tvmscript/test_parser.py
+++ b/tests/python/tvmscript/test_parser.py
@@ -45,6 +45,7 @@ class _Recorder:
def __init__(self):
self.bindings = {}
self.returned = []
+ self.frame_requests = []
def bind_(self, value=protocol.MISSING, **metadata):
self.bindings[metadata.get("name")] = (value, metadata)
@@ -65,6 +66,12 @@ class _Recorder:
def unpack(self, value):
return value
+ def frame_result(self, frame, name):
+ self.frame_requests.append((frame, name))
+ return getattr(frame, "result", frame if isinstance(frame, dict) else
{}).get(
+ name, protocol.MISSING
+ )
+
def _registered(source, env=None, **options):
recorder = _Recorder()
@@ -655,3 +662,32 @@ def f():
]
assert chains == [(1, 2, 3)]
assert builder.returned == [(False, True, 10, True, True)]
+
+
+def test_named_results_keep_nested_frame_identity():
+ class Scope:
+ def __init__(self, value):
+ self.value = value
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ self.result = {"value": self.value}
+
+ outer, inner = Scope(30), Scope(20)
+ compiler, function, builder = _registered(
+ """
+ @D.function
+ def f():
+ with outer:
+ with inner:
+ value = 1
+ value = value + 1
+ return value
+ """,
+ {"outer": outer, "inner": inner},
+ )
+ _run(compiler, function, builder)
+ assert builder.frame_requests == [(inner, "value"), (outer, "value")]
+ assert builder.returned == [30]
