This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 900bc3477 [REFACTOR] Bind parser source locations to individual names
(#103)
900bc3477 is described below
commit 900bc34770e5efa3f5b87e3144e7d63b04ab4f98
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 27 21:23:34 2026 -0400
[REFACTOR] Bind parser source locations to individual names (#103)
Bind generated parser source locations to individual fresh names instead
of indexing a shared location list. Reuse each binding for repeated
source coordinates and preserve collision avoidance through the existing
name allocator.
Update parser comments and examples to match the generated names.
---
python/tvm_tirx/script/ir_builder/parser_protocol.py | 4 ++--
python/tvm_tirx/script/parser/entry.py | 20 +++++++-------------
python/tvm_tirx/script/parser/transpile.py | 20 ++++++++++----------
3 files changed, 19 insertions(+), 25 deletions(-)
diff --git a/python/tvm_tirx/script/ir_builder/parser_protocol.py
b/python/tvm_tirx/script/ir_builder/parser_protocol.py
index 902803ab8..4dc342501 100644
--- a/python/tvm_tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm_tirx/script/ir_builder/parser_protocol.py
@@ -427,7 +427,7 @@ def bind_(
# Source
x = value
# Generated builder
- x = X.bind_(value, name="x", loc=_L[0], value_loc=_L[1])
+ x = X.bind_(value, name="x", loc=_L_0, value_loc=_L_1)
# TIRx concise scope entry
tid = T.launch_thread("threadIdx.x", 128)
@@ -596,7 +596,7 @@ def emit_(value: Any, *, loc: _Loc = None) -> None:
# Source
T.evaluate(1)
# Generated builder
- X.emit_(X.evaluate(1), loc=_L[0])
+ X.emit_(X.evaluate(1), loc=_L_0)
"""
raise NotImplementedError
diff --git a/python/tvm_tirx/script/parser/entry.py
b/python/tvm_tirx/script/parser/entry.py
index 2aab23341..4517a1b79 100644
--- a/python/tvm_tirx/script/parser/entry.py
+++ b/python/tvm_tirx/script/parser/entry.py
@@ -604,17 +604,14 @@ def _prepare_transpiler(
builder_name, ir_prefix = make_fresh_name("_X"), make_fresh_name("_I")
namespace[ir_prefix] = builder_ir
- loc_table_name = make_fresh_name("_L") if track_loc else None
if track_loc:
# Entries contain fixed native metadata only. The existing rewrite
creates
# them on demand; there is no location collection pass or retained AST.
source_name = SourceName(filename)
- loc_entries: list[LocationEntry] = []
- loc_indices: dict[tuple[int, int, int, int], int] = {}
- namespace[loc_table_name] = loc_entries
+ loc_names: dict[tuple[int, int, int, int], str] = {}
def make_loc_expr(node: ast.AST) -> ast.expr:
- """Materialize a needed location and emit its injected table
reference."""
+ """Materialize a needed location and emit its injected name."""
if not track_loc:
return ast.copy_location(ast.Constant(None), node)
coordinates = (
@@ -623,14 +620,11 @@ def _prepare_transpiler(
node.end_lineno,
node.end_col_offset + 1,
)
- index = loc_indices.get(coordinates)
- if index is None:
- index = len(loc_entries)
- loc_indices[coordinates] = index
- loc_entries.append(LocationEntry(SourceLoc(source_name,
*coordinates)))
- location = ast.Subscript(
- ast.Name(loc_table_name, ast.Load()), ast.Constant(index),
ast.Load()
- )
+ name = loc_names.get(coordinates)
+ if name is None:
+ name = loc_names[coordinates] = make_fresh_name("_L_")
+ namespace[name] = LocationEntry(SourceLoc(source_name,
*coordinates))
+ location = ast.Name(name, ast.Load())
return ast.copy_location(location, node)
context = ModuleContext(
diff --git a/python/tvm_tirx/script/parser/transpile.py
b/python/tvm_tirx/script/parser/transpile.py
index 1735b8937..432d8fd27 100644
--- a/python/tvm_tirx/script/parser/transpile.py
+++ b/python/tvm_tirx/script/parser/transpile.py
@@ -174,7 +174,7 @@ class ModuleContext:
self.environment = environment
# Generated name of the shared IR namespace used for module frames and
MISSING.
self.ir_prefix = ir_prefix
- # Entry callback creates/reuses a LocationEntry and emits its AST
table reference
+ # Entry callback creates/reuses a LocationEntry and emits its AST name
reference
# for _call_dialect, binding locations and expression instrumentation.
self.make_loc_expr = make_loc_expr
# _call and _attach_loc omit native source instrumentation when false.
@@ -344,7 +344,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
or node is self.binding_expression
):
return value
- # _L[i](value)
+ # _L_i(value)
return ast.copy_location(ast.Call(self.module.make_loc_expr(node),
[value], []), node)
@staticmethod
@@ -462,7 +462,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# Module.f
#
# Builder:
- # _L[i](Module.f)
+ # _L_i(Module.f)
# -------------------------------------------------
# Keeping the module owner prevents a local f from shadowing its
GlobalVar.
root = node
@@ -496,7 +496,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# [a, b]
#
# Builder:
- # _L[i]([a, b])
+ # _L_i([a, b])
# -------------------------------------------------
result = self.generic_visit(node)
return result if self.bypass_ast_rewrite else self._attach_loc(result,
node)
@@ -527,7 +527,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# buffer[index]
#
# Builder:
- # _L[i](buffer[index])
+ # _L_i(buffer[index])
# -------------------------------------------------
node.value = self.visit(node.value)
node.slice = self._rewrite_index(node.slice)
@@ -613,11 +613,11 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# f(a)
#
# Builder:
- # _L[i].ctx(lambda: f(a))
+ # _L_i.ctx(lambda: f(a))
# -------------------------------------------------
# Calls retain their construction context; bind_ owns ordinary RHS
attribution.
if self.module.track_loc and callee is None:
- # _L[i].ctx(lambda: callee(*args, **keywords))
+ # _L_i.ctx(lambda: callee(*args, **keywords))
return ast.copy_location(
ast.Call(
ast.Attribute(self.module.make_loc_expr(node), "ctx",
ast.Load()),
@@ -650,7 +650,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# a + b
#
# Builder:
- # _L[i](a + b)
+ # _L_i(a + b)
# -------------------------------------------------
return (
self._attach_loc(self.generic_visit(node), node)
@@ -1164,7 +1164,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# f()
#
# Builder:
- # X.emit_(_L[i].ctx(lambda: f()), loc=_L[i])
+ # X.emit_(_L_i.ctx(lambda: f()), loc=_L_i)
# -------------------------------------------------
if self.bypass_ast_rewrite:
return self.generic_visit(node)
@@ -1878,7 +1878,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
[translated, ast.Name(const_args, ast.Load())],
[],
)
- # x = X.arg_("x", annotation, loc=_L[i])
+ # x = X.arg_("x", annotation, loc=_L_i)
fallback = self._call_dialect("arg_", [ast.Constant(name),
translated], parameter)
# Constexprs have no ABI slot. Optional None parameters retain a
slot
# while binding Python None for later annotations and body guards.
