This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 64721bde7e [TVMScript] Fill AST locations at compilation and simplify
span inputs (#20450)
64721bde7e is described below
commit 64721bde7ea22d2712eb5f8f3ecd729ee2a53839
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Sep 25 11:21:13 2026 -0400
[TVMScript] Fill AST locations at compilation and simplify span inputs
(#20450)
Remove the redundant AST location fill after rewriting; fill missing
locations at each compilation boundary after the generated builder has
been assembled.
Restrict script builder span inputs to native spans, span entries, or
`None`. Unwrap entries at native boundaries and remove tuple/string
source-range normalization while preserving span attachment, context
handling, and emission receipts. Place the receipt class before the
span-entry class.
---
python/tvm/relax/script/ir_builder/__init__.py | 56 ++++++++-----
.../tvm/relax/script/ir_builder/parser_protocol.py | 30 +++----
python/tvm/script/ir_builder/base.py | 94 +++++++++-------------
python/tvm/script/ir_builder/ir.py | 10 +--
python/tvm/script/ir_builder/parser_protocol.py | 58 ++++++-------
python/tvm/script/parser/transpile.py | 2 +-
python/tvm/tirx/script/ir_builder/__init__.py | 5 +-
.../tvm/tirx/script/ir_builder/parser_protocol.py | 14 ++--
tests/python/script/test_source_locations.py | 2 +-
9 files changed, 138 insertions(+), 133 deletions(-)
diff --git a/python/tvm/relax/script/ir_builder/__init__.py
b/python/tvm/relax/script/ir_builder/__init__.py
index a9919cf4d2..2ca8e53e25 100644
--- a/python/tvm/relax/script/ir_builder/__init__.py
+++ b/python/tvm/relax/script/ir_builder/__init__.py
@@ -28,8 +28,8 @@ from tvm.relax.distributed import DTensorType as _DTensorType
from tvm.relax.distributed import Placement as _Placement
from tvm.relax.distributed import device_mesh as device_mesh
from tvm.script.ir_builder import resolve_global_info_args as
_resolve_global_info_args
+from tvm.script.ir_builder.base import SpanEntry as _SpanEntry
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.parser.protocol_registry import constexpr as constexpr
from . import distributed as dist
@@ -101,7 +101,7 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1,
*, span=None):
ndim : int, optional
Rank when shape is unknown; -1 means unknown rank. Do not supply
an explicit rank together with a known shape.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -112,7 +112,9 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1,
*, span=None):
"""
if isinstance(shape, _python.str) and dtype is None:
dtype, shape = shape, None
- return _relax.TensorType(shape, dtype, vdevice, ndim, _source_span(span))
+ return _relax.TensorType(
+ shape, dtype, vdevice, ndim, span.span if isinstance(span, _SpanEntry)
else span
+ )
@_resolve_global_info_args("device_mesh", resolver=resolve_global_info_)
@@ -135,7 +137,7 @@ def DTensor(shape=None, dtype=None, device_mesh=None,
placement="", *, ndim=-1,
parsed with Placement.from_text.
ndim : int, optional
Global rank when shape is unknown; -1 means unknown rank.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -148,7 +150,12 @@ def DTensor(shape=None, dtype=None, device_mesh=None,
placement="", *, ndim=-1,
device_mesh = _DeviceMesh([], _ir.Range(0, 1))
if isinstance(placement, _python.str):
placement = _Placement.from_text(placement)
- return _DTensorType(Tensor(shape, dtype, ndim=ndim), device_mesh,
placement, _source_span(span))
+ return _DTensorType(
+ Tensor(shape, dtype, ndim=ndim),
+ device_mesh,
+ placement,
+ span.span if isinstance(span, _SpanEntry) else span,
+ )
# The distributed source spelling shares the decorated concrete constructor.
@@ -172,7 +179,7 @@ def Shape(values=None, ndim=-1, *, span=None):
ndim : int, optional
Number of dimensions when values is None; -1 leaves it unknown.
Do not supply an explicit count together with known values.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -180,7 +187,7 @@ def Shape(values=None, ndim=-1, *, span=None):
result : ShapeType
The constructed shape type.
"""
- return _relax.ShapeType(values, ndim, _source_span(span))
+ return _relax.ShapeType(values, ndim, span.span if isinstance(span,
_SpanEntry) else span)
def _type(value):
@@ -212,7 +219,7 @@ def Callable(params=None, ret=None, purity=None,
derive_func=None, *, span=None)
derive_func : str or EnvFunc, optional
Custom result-type derivation for an opaque callable. It is not
accepted when params supplies a concrete parameter list.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -233,14 +240,17 @@ def Callable(params=None, ret=None, purity=None,
derive_func=None, *, span=None)
ret=None if ret is None else _type(ret),
derive_func=derive_func,
purity=purity,
- span=_source_span(span),
+ span=span.span if isinstance(span, _SpanEntry) else span,
)
if derive_func is not None:
raise ValueError("A derivation function requires an opaque callable")
if not isinstance(params, list | _python.tuple):
params = [params]
return _relax.FuncType(
- [_type(param) for param in params], _type(ret), purity,
_source_span(span)
+ [_type(param) for param in params],
+ _type(ret),
+ purity,
+ span.span if isinstance(span, _SpanEntry) else span,
)
@@ -252,7 +262,7 @@ def Tuple(*fields, span=None):
fields : Type or callable
Field annotations as positional arguments, or one list or tuple.
Each annotation is normalized to a type; None denotes an empty tuple.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -262,7 +272,9 @@ def Tuple(*fields, span=None):
"""
if len(fields) == 1 and isinstance(fields[0], list | _python.tuple):
fields = fields[0]
- return _ir.TupleType([_type(field) for field in fields],
_source_span(span))
+ return _ir.TupleType(
+ [_type(field) for field in fields], span.span if isinstance(span,
_SpanEntry) else span
+ )
def Object(*, span=None):
@@ -270,7 +282,7 @@ def Object(*, span=None):
Parameters
----------
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -278,7 +290,7 @@ def Object(*, span=None):
result : AnyType
A type accepting any Relax value.
"""
- return _relax.AnyType(_source_span(span))
+ return _relax.AnyType(span.span if isinstance(span, _SpanEntry) else span)
Any = Object
@@ -296,7 +308,7 @@ def type_var(name, *, dtype=None, span=None):
Name of the symbol.
dtype : str or PrimType, optional
Primitive type of the symbol; None selects "int64".
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -309,7 +321,11 @@ def type_var(name, *, dtype=None, span=None):
This constructor creates a new symbol on each call. Use the language
variant
resolver for symbols shared by name within a function signature.
"""
- return _ir.Var(name, "int64" if dtype is None else dtype,
_source_span(span))
+ return _ir.Var(
+ name,
+ "int64" if dtype is None else dtype,
+ span.span if isinstance(span, _SpanEntry) else span,
+ )
func_ret_ty = func_ret_type
@@ -320,7 +336,7 @@ def dataflow(*, span=None):
Parameters
----------
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR.
Returns
@@ -351,7 +367,7 @@ def match_cast(value, ty, *, span=None):
are converted to Relax expressions; None is not accepted.
ty : Type or callable
Asserted type, or a zero-argument factory producing its annotation.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -363,7 +379,9 @@ def match_cast(value, ty, *, span=None):
if value is None:
raise ValueError("The match-cast value cannot be None")
ty = _type(ty)
- return _relax.MatchCast(_ir.Var("", ty), _value(value), ty,
_source_span(span))
+ return _relax.MatchCast(
+ _ir.Var("", ty), _value(value), ty, span.span if isinstance(span,
_SpanEntry) else span
+ )
__all__ = [
diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py
b/python/tvm/relax/script/ir_builder/parser_protocol.py
index f7ba61e7c7..85ac05e23e 100644
--- a/python/tvm/relax/script/ir_builder/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -45,7 +45,7 @@ from . import _ffi_api
from . import frame as _frame
from . import ir as _native
-_Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int]
| None
+_Span = _base.SpanEntry | _ir.Span | None
# --------------------------------------
@@ -172,7 +172,7 @@ def lt_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
return _base.at_(span, _relax.op.less(lhs, rhs))
- return _prim_ffi._OpLT(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpLT(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def le_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
@@ -181,7 +181,7 @@ def le_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
return _base.at_(span, _relax.op.less_equal(lhs, rhs))
- return _prim_ffi._OpLE(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpLE(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def gt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
@@ -190,7 +190,7 @@ def gt_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
return _base.at_(span, _relax.op.greater(lhs, rhs))
- return _prim_ffi._OpGT(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpGT(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def ge_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
@@ -199,7 +199,7 @@ def ge_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
return _base.at_(span, _relax.op.greater_equal(lhs, rhs))
- return _prim_ffi._OpGE(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpGE(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def eq_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
@@ -208,7 +208,7 @@ def eq_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
return _base.at_(span, _relax.op.equal(lhs, rhs))
- return _prim_ffi._OpEQ(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpEQ(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def ne_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
@@ -217,7 +217,7 @@ def ne_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
return _base.at_(span, _relax.op.not_equal(lhs, rhs))
- return _prim_ffi._OpNE(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpNE(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
# --------------------------------------
@@ -350,7 +350,7 @@ def arg(name: str, ty: Any, *, span: _Span = None) ->
_ir.Var:
:func:`tvm.script.ir_builder.parser_protocol.arg` contract. Primitive
annotations resolve through the active signature symbol context; an
existing variable retains its identity.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR; None leaves it
unspecified.
Returns
@@ -446,14 +446,18 @@ def bind_(
Relax values emit a binding or MatchCast; value_span belongs to the actual
RHS and span to the binding target. Metadata passes through unchanged.
"""
- name_span = _base.source_span(span if name_span is None else name_span)
+ if isinstance(span, _base.SpanEntry):
+ span = span.span
+ name_span = span if name_span is None else name_span
+ if isinstance(name_span, _base.SpanEntry):
+ name_span = name_span.span
if frame_value:
if isinstance(value, _python.list | _python.tuple | _ir.Array):
for index, item in enumerate(value):
bind_(
item,
name=None if name is None else f"{name}_{index}",
- span=_base.source_span(span),
+ span=span,
name_span=name_span,
frame_value=True,
)
@@ -472,12 +476,10 @@ def bind_(
_base.at_(value_span, value.value)
if ty is not None and not _ffi.structural_equal(ty, value.ty):
raise TypeError("The binding annotation differs from the
match-cast type")
- result = _ffi_api.EmitMatchCastWithSpan(
- value.value, value.ty, name_span, _base.source_span(span)
- )
+ result = _ffi_api.EmitMatchCastWithSpan(value.value, value.ty,
name_span, span)
elif isinstance(value, _relax.Expr):
_base.at_(value_span, value)
- result = _ffi_api.EmitWithSpan(value, ty, name_span,
_base.source_span(span))
+ result = _ffi_api.EmitWithSpan(value, ty, name_span, span)
else:
return value
if name is not None:
diff --git a/python/tvm/script/ir_builder/base.py
b/python/tvm/script/ir_builder/base.py
index bba7ea9471..0543fd8613 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -274,6 +274,29 @@ class IRBuilder(_Object):
_T = TypeVar("_T")
+class AlreadyEmitted(Generic[_T]):
+ """Hold the exact emitted value; location handling preserves this receipt.
+
+ Parameters
+ ----------
+ value : Any
+ Object already emitted by a language variant's builder. The receipt
+ retains this object without copying it. Binding or emitting the receipt
+ must not emit the object again.
+
+ Attributes
+ ----------
+ value : Any
+ The same emitted object, available for identity checks and source-span
+ attachment.
+ """
+
+ __slots__ = ("value",)
+
+ def __init__(self, value: _T) -> None:
+ self.value = value
+
+
class SpanEntry:
"""A materialized source range shared by generated builder operations.
@@ -282,7 +305,7 @@ class SpanEntry:
during evaluation. ``ctx(thunk, attach_result=False)`` supplies only the
evaluation context, leaving result attachment to the binding operation.
Both compose the active caller context at invocation.
- Builders accepting an explicit span normalize an entry with
``source_span``.
+ Builders accepting an explicit span unwrap the entry at native boundaries.
"""
__slots__ = ("span",)
@@ -303,54 +326,14 @@ class SpanEntry:
return with_at_group_(self.span, thunk, attach_result=attach_result)
-def source_span(
- location: SpanEntry | ir.Span | tuple[str | ir.SourceName, int, int, int,
int] | None,
-) -> ir.Span | None:
- """Normalize an entry or materialize a range without retaining source-unit
state."""
- if isinstance(location, SpanEntry):
- return location.span
- if location is None or isinstance(location, ir.Span | ir.SequentialSpan):
- return location
- source_name, line, end_line, column, end_column = location
- if isinstance(source_name, str):
- source_name = ir.SourceName(source_name)
- return ir.Span(source_name, line, end_line, column, end_column)
-
-
-class AlreadyEmitted(Generic[_T]):
- """Hold the exact emitted value; location handling preserves this receipt.
-
- Parameters
- ----------
- value : Any
- Object already emitted by a language variant's builder. The receipt
- retains this object without copying it. Binding or emitting the receipt
- must not emit the object again.
-
- Attributes
- ----------
- value : Any
- The same emitted object, available for identity checks and source-span
- attachment.
- """
-
- __slots__ = ("value",)
-
- def __init__(self, value: _T) -> None:
- self.value = value
-
-
-def at(
- span: SpanEntry | ir.Span | tuple[str | ir.SourceName, int, int, int, int]
| None, value: _T
-) -> _T:
+def at(span: SpanEntry | ir.Span | None, value: _T) -> _T:
"""Attach source context to the same IR node, emission receipt, or frame.
Parameters
----------
- span : SpanEntry, Span, tuple or None
- Source location to compose with the active construction context. A
tuple
- contains ``(source_name, line, end_line, column, end_column)``; its
source
- name may be a string or SourceName. None leaves the value unchanged.
+ span : SpanEntry, Span or None
+ Source location to compose with the active construction context.
+ None leaves the value unchanged.
value : Any
Native object, :class:`AlreadyEmitted` receipt, or list/tuple of native
objects to annotate. A receipt's contained object receives the span.
@@ -369,16 +352,18 @@ def at(
"""
if span is None or not IRBuilder.is_in_scope():
return value
+ if isinstance(span, SpanEntry):
+ span = span.span
target = value.value if isinstance(value, AlreadyEmitted) else value
targets = target if isinstance(target, list | tuple) else (target,)
for item in targets:
if isinstance(item, _Object):
- _ffi_api.IRBuilderSetSourceSpan(IRBuilder.current(), item,
source_span(span))
+ _ffi_api.IRBuilderSetSourceSpan(IRBuilder.current(), item, span)
return value
def with_at_group_(
- location: SpanEntry | ir.Span | tuple[str | ir.SourceName, int, int, int,
int] | None,
+ location: SpanEntry | ir.Span | None,
thunk: Callable[[], _T],
*,
attach_result: bool = True,
@@ -387,10 +372,9 @@ def with_at_group_(
Parameters
----------
- location : SpanEntry, Span, tuple or None
- Source context for the call. A tuple contains
- ``(source_name, line, end_line, column, end_column)``. None, or the
- absence of an active builder, leaves construction context unchanged.
+ location : SpanEntry, Span or None
+ Source context for the call. None, or the absence of an active builder,
+ leaves construction context unchanged.
thunk : Callable[[], Any]
Zero-argument callable evaluated exactly once inside that context.
attach_result : bool, optional
@@ -409,7 +393,7 @@ def with_at_group_(
exception propagates unchanged. Frames created during the call retain their
construction spans regardless of ``attach_result``.
"""
- span = source_span(location)
+ span = location.span if isinstance(location, SpanEntry) else location
context = (
IRBuilder.current().with_source_span(span)
if span is not None and IRBuilder.is_in_scope()
@@ -417,7 +401,7 @@ def with_at_group_(
)
with context:
value = thunk()
- return at(location, value) if attach_result else value
+ return at(span, value) if attach_result else value
at_ = at
@@ -435,7 +419,9 @@ def _resolve_type_var(frame, ffi_resolver, name,
dtype=None, *, value=None, span
raise TypeError("A symbolic variable requires a primitive type")
if value is not None and not ir.is_prim_var(value):
raise TypeError("A symbolic binding requires a primitive Var")
- return ffi_resolver(frame, name, dtype, value, source_span(span))
+ return ffi_resolver(
+ frame, name, dtype, value, span.span if isinstance(span, SpanEntry)
else span
+ )
def _current_function_frame():
diff --git a/python/tvm/script/ir_builder/ir.py
b/python/tvm/script/ir_builder/ir.py
index 7230bdba71..7f46a436f1 100644
--- a/python/tvm/script/ir_builder/ir.py
+++ b/python/tvm/script/ir_builder/ir.py
@@ -19,17 +19,17 @@
import inspect
from typing import TypeVar
-from tvm.ir import BaseFunc, GlobalInfo, GlobalVar, Var
+from tvm.ir import BaseFunc, GlobalInfo, GlobalVar, Span, Var
from tvm.runtime import Object as tvm_Object
from . import _ffi_api
-from .base import IRBuilder, source_span
+from .base import IRBuilder, SpanEntry
from .frame import IRModuleFrame
T = TypeVar("T")
-def dynamic(name: str, dtype: str = "int64", *, span=None) -> Var:
+def dynamic(name: str, dtype: str = "int64", *, span: SpanEntry | Span | None
= None) -> Var:
"""Create a fresh primitive symbolic variable, independently of builder
scope.
Parameters
@@ -38,7 +38,7 @@ def dynamic(name: str, dtype: str = "int64", *, span=None) ->
Var:
The symbol's display name. Repeated names do not share identity.
dtype : str
Primitive dtype, defaulting to int64.
- span : Optional[Span]
+ span : SpanEntry, Span or None
Source location of the symbol.
Returns
@@ -47,7 +47,7 @@ def dynamic(name: str, dtype: str = "int64", *, span=None) ->
Var:
A fresh symbol. Reuse this object to share dimensions across
annotations
and function bodies, including outside an ``I.ir_module`` definition.
"""
- return Var(name, dtype, source_span(span))
+ return Var(name, dtype, span.span if isinstance(span, SpanEntry) else span)
def meta_var(value: T) -> T:
diff --git a/python/tvm/script/ir_builder/parser_protocol.py
b/python/tvm/script/ir_builder/parser_protocol.py
index 1d09c15eaa..cc28473634 100644
--- a/python/tvm/script/ir_builder/parser_protocol.py
+++ b/python/tvm/script/ir_builder/parser_protocol.py
@@ -35,7 +35,7 @@ from tvm import ir as _ir
from .base import MISSING, AlreadyEmitted, IRBuilderFrame, SpanEntry
-_Span = SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int] | None
+_Span = SpanEntry | _ir.Span | None
# --------------------------------------
@@ -60,7 +60,7 @@ def if_(condition: Any, *, span: _Span = None) ->
IRBuilderFrame:
condition : Expr or bool
Already-evaluated predicate; both branch bodies construct IR without
testing it
in Python.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -99,7 +99,7 @@ def then_(*, span: _Span = None) -> IRBuilderFrame:
Parameters
----------
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -137,7 +137,7 @@ def else_(*, span: _Span = None) -> IRBuilderFrame:
Parameters
----------
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -184,7 +184,7 @@ def for_(
Source target names, including at most one ``*starred`` group. None
(the default)
retains constructor defaults. Native configuration expands/validates
them before
entry; names never control the entry return shape.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -222,7 +222,7 @@ def while_(condition: Any, *, span: _Span = None) ->
IRBuilderFrame:
----------
condition : Expr or bool
Loop predicate expression, constructed once and evaluated by the IR at
runtime.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -288,7 +288,7 @@ def break_(*, span: _Span = None) -> AlreadyEmitted[Any]:
Parameters
----------
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -319,7 +319,7 @@ def continue_(*, span: _Span = None) -> AlreadyEmitted[Any]:
Parameters
----------
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -495,7 +495,7 @@ def lt_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
Already-evaluated left operand.
rhs : Expr or scalar
Already-evaluated right operand.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -531,7 +531,7 @@ def le_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
Already-evaluated left operand.
rhs : Expr or scalar
Already-evaluated right operand.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -567,7 +567,7 @@ def gt_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
Already-evaluated left operand.
rhs : Expr or scalar
Already-evaluated right operand.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -603,7 +603,7 @@ def ge_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
Already-evaluated left operand.
rhs : Expr or scalar
Already-evaluated right operand.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -639,7 +639,7 @@ def eq_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
Already-evaluated left operand.
rhs : Expr or scalar
Already-evaluated right operand.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -675,7 +675,7 @@ def ne_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
Already-evaluated left operand.
rhs : Expr or scalar
Already-evaluated right operand.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -764,7 +764,7 @@ def resolve_type_var_(
value : Var, optional
Existing primitive variable to register without replacement. None
creates a
variable only if the name is not already registered.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -890,7 +890,7 @@ def function_(*, decl: bool = False, span: _Span = None,
**options: Any) -> IRBu
decl : bool, optional
False (default) constructs a complete function on one entry. True
collects a
signature on the first entry and retains this frame for body re-entry.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -940,7 +940,7 @@ def arg(name: str, annotation: Any, *, span: _Span = None)
-> _ir.Var:
annotation : Type, Var, Buffer or callable
Concrete rewritten annotation or existing native parameter. A callable
annotation is evaluated; an existing variable retains identity.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -1007,7 +1007,7 @@ def func_ret_type(annotation: Any, *, span: _Span = None)
-> None:
annotation : Type, Expr or callable
Rewritten return annotation; expression annotations supply their type.
None
denotes a void/empty tuple return as supported by the language variant.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -1109,15 +1109,15 @@ def bind_(
the binding type.
name : str, optional
Source target name. None (the default) requests no source-derived name.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Binding-target location for a newly constructed binding. None (the
default)
leaves it unspecified; this is separate from the RHS location.
- value_span : SpanEntry, Span or source-location tuple, optional
+ value_span : SpanEntry, Span or None, optional
RHS source location, passed separately without first stamping the
returned value.
None (the default) leaves explicit RHS attribution unspecified. The
language variant
applies it only when binding/conversion requires value attribution;
TIRx
variable and metadata passthrough retain producer names and spans.
- name_span : SpanEntry, Span or source-location tuple, optional
+ name_span : SpanEntry, Span or None, optional
Location of the target identifier. None (the default) uses span; it
can differ
from the emitted statement location.
frame_value : bool, optional
@@ -1182,11 +1182,11 @@ def decl_mutable_cell_(
created storage handle.
name : str, optional
Source target name. None (the default) requests no source-derived name.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
- name_span : SpanEntry, Span or source-location tuple, optional
+ name_span : SpanEntry, Span or None, optional
Location of the target identifier. None (the default) uses span; it
can differ
from the emitted statement location.
@@ -1224,7 +1224,7 @@ def set_mutable_cell_(target: Any, value: Any, *, span:
_Span = None) -> Already
value : Expr or scalar convertible to Expr
Once-evaluated value to store. Native store checking validates its
type and
indices against the target.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location of the emitted store. None (the default) leaves explicit
location
unspecified; existing source-call provenance is retained.
@@ -1305,7 +1305,7 @@ def emit_(value: Any, *, span: _Span = None) -> None:
value : Any
Once-evaluated expression result. AlreadyEmitted receipts and None
produce no
additional emission.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location of the emitted statement and expression, or the existing node
in an
AlreadyEmitted receipt. None leaves explicit attribution unspecified.
Active
caller provenance is composed without adding a construction context;
native
@@ -1345,7 +1345,7 @@ def return_(value: Any = None, *, span: _Span = None) ->
AlreadyEmitted[Any] | N
value : Any, optional
Return operand. None (the default) means an empty tuple in Relax; TIRx
requires
an expression.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -1384,7 +1384,7 @@ def setitem_(target: Any, key: Any, value: Any, *, span:
_Span = None) -> Alread
Indices in written order; native buffer-store rules validate supported
forms.
value : Expr or scalar
Once-evaluated stored value.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -1424,7 +1424,7 @@ def setattr_(
Attribute identifier, evaluated by source syntax before this hook.
value : Any
Once-evaluated replacement or stored value.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
@@ -1467,7 +1467,7 @@ def assert_(
message : str or assertion metadata, optional
Empty text by default. Relax requires construction-time text. TIRx
also accepts
message parts or an (error_kind, parts) pair.
- span : SpanEntry, Span or source-location tuple, optional
+ span : SpanEntry, Span or None, optional
Location for the constructed result. None (the default) leaves
explicit location
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index b62f9f6ec2..1a41a587cc 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -2412,4 +2412,4 @@ class IRBuilderTranspiler(ast.NodeTransformer):
root, result, python_functions,
check_well_formed=check_well_formed
)
)
- return ast.fix_missing_locations(ast.Module(translated, [])), result
+ return ast.Module(translated, []), result
diff --git a/python/tvm/tirx/script/ir_builder/__init__.py
b/python/tvm/tirx/script/ir_builder/__init__.py
index 6f2e2ed550..3677ba9b02 100644
--- a/python/tvm/tirx/script/ir_builder/__init__.py
+++ b/python/tvm/tirx/script/ir_builder/__init__.py
@@ -26,7 +26,6 @@ from tvm import tirx as _tir
from tvm.script.ir_builder import dynamic as dynamic
from tvm.script.ir_builder.base import annotation_constructor as
_annotation_constructor
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.parser.protocol_registry import constexpr as constexpr
from tvm.script.parser.protocol_registry import (
mutable_cell_decl as _mutable_cell_decl,
@@ -155,7 +154,7 @@ def Buffer(
buffer_name : str
The name of the buffer.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR.
Returns
@@ -196,7 +195,7 @@ def Ptr(dtype, storage_scope="global", *, span=None):
storage_scope : str
The storage scope of the pointer.
- span : Span or source location, optional
+ span : SpanEntry, Span or None, optional
Source location attached to the constructed IR.
Returns
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 0d170dbc54..47de5704bf 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -40,7 +40,7 @@ from . import _ffi_api
from . import frame as _frame
from . import ir as _native
-_Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int]
| None
+_Span = _base.SpanEntry | _ir.Span | None
# --------------------------------------
@@ -160,32 +160,32 @@ def not_(value: Any) -> Any:
def lt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.lt_`."""
- return _prim_ffi._OpLT(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpLT(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def le_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.le_`."""
- return _prim_ffi._OpLE(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpLE(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def gt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.gt_`."""
- return _prim_ffi._OpGT(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpGT(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def ge_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.ge_`."""
- return _prim_ffi._OpGE(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpGE(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def eq_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.eq_`."""
- return _prim_ffi._OpEQ(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpEQ(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
def ne_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.ne_`."""
- return _prim_ffi._OpNE(lhs, rhs, _base.source_span(span))
+ return _prim_ffi._OpNE(lhs, rhs, span.span if isinstance(span,
_base.SpanEntry) else span)
# --------------------------------------
diff --git a/tests/python/script/test_source_locations.py
b/tests/python/script/test_source_locations.py
index e54dea0156..4ec3cfc97d 100644
--- a/tests/python/script/test_source_locations.py
+++ b/tests/python/script/test_source_locations.py
@@ -336,7 +336,7 @@ def
test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
captured = T.Buffer((4, 4), "float32")
original = type(captured).view
seen, produced, observed = [], [], []
- span = base.source_span(("producer.py", 7, 7, 2, 19))
+ span = ir.Span(ir.SourceName("producer.py"), 7, 7, 2, 19)
@wraps(original)
def view(buffer, *args):
