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 544ac6a818 [REFACTOR][IR] Consolidate builder location resolution and
attachment (#20723)
544ac6a818 is described below
commit 544ac6a818b6dea0f3b193a4268ef37438ff19e8
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 11:56:58 2026 -0400
[REFACTOR][IR] Consolidate builder location resolution and attachment
(#20723)
Rename IRBuilder location resolution to `ResolveLoc` and consolidate
node attachment under `AttachLoc` with an optional explicit location.
Preserve full-stack normalization, caller composition, and object
identity while updating native consumers and Python FFI helpers.
---
include/tvm/script/ir_builder/base.h | 10 ++-
python/tvm/script/ir_builder/base.py | 11 +--
src/relax/script/ir_builder/ir.cc | 2 +-
src/script/ir_builder/base.cc | 15 ++--
src/script/ir_builder/utils.h | 2 +-
tests/python/script/test_ir_builder.py | 79 ++++++++++++++++++++++
.../script/test_tirx_script_source_locations.py | 20 ++++++
7 files changed, 116 insertions(+), 23 deletions(-)
diff --git a/include/tvm/script/ir_builder/base.h
b/include/tvm/script/ir_builder/base.h
index b5e5c75a6d..e3bfd29d89 100644
--- a/include/tvm/script/ir_builder/base.h
+++ b/include/tvm/script/ir_builder/base.h
@@ -179,12 +179,10 @@ class IRBuilderObj : public ffi::Object {
void PushLoc(Location loc);
/*! \brief Pop the innermost frontend source location. */
void PopLoc();
- /*! \brief Return the normalized active source loc, including expansion
history. */
- Location GetCurrentLoc(Location location = UnknownLoc()) const;
- /*! \brief Compose active source context onto a supported node or
construction frame. */
- ffi::ObjectRef SetCurrentLoc(ffi::ObjectRef obj) const;
- /*! \brief Attach an explicit location composed with the active source-call
context. */
- ffi::ObjectRef SetLoc(ffi::ObjectRef obj, Location loc) const;
+ /*! \brief Normalize the full active source stack with an optional explicit
location. */
+ Location ResolveLoc(Location location = UnknownLoc()) const;
+ /*! \brief Compose active and explicit source context onto the same
supported node or frame. */
+ ffi::ObjectRef AttachLoc(ffi::ObjectRef obj, Location loc = UnknownLoc())
const;
};
/*!
diff --git a/python/tvm/script/ir_builder/base.py
b/python/tvm/script/ir_builder/base.py
index 803bbe729a..6593e5baf0 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -225,11 +225,12 @@ class IRBuilder(_Object):
self
)
- def _set_current_loc(self, value):
- """Compose the active source location onto the same supported IR node
or frame."""
- return _ffi_api.IRBuilderSetCurrentLoc( # type: ignore[attr-defined]
# pylint: disable=no-member
- self, value
+ def _attach_loc(self, value, loc=None):
+ """Compose active and explicit source context onto the same supported
IR node or frame."""
+ _ffi_api.IRBuilderAttachLoc( # type: ignore[attr-defined] # pylint:
disable=no-member
+ self, value, UnknownLoc() if loc is None else loc
)
+ return value
@staticmethod
def name(s: str, v: Any) -> Any:
@@ -366,7 +367,7 @@ def at(loc: LocationEntry | ir.Location | None, value: _T)
-> _T:
targets = target if isinstance(target, list | tuple) else (target,)
for item in targets:
if isinstance(item, _Object):
- _ffi_api.IRBuilderSetLoc(IRBuilder.current(), item, loc)
+ IRBuilder.current()._attach_loc(item, loc)
return value
diff --git a/src/relax/script/ir_builder/ir.cc
b/src/relax/script/ir_builder/ir.cc
index fffaa51814..b1b8dba9d4 100644
--- a/src/relax/script/ir_builder/ir.cc
+++ b/src/relax/script/ir_builder/ir.cc
@@ -307,7 +307,7 @@ tvm::Var EmitVarBinding(const tvm::relax::VarBinding&
binding) {
namespace {
tvm::Var RecordBindingLoc(tvm::Var var, const Location& name_loc, const
Location& statement_loc) {
- Location loc = IRBuilder::Current()->GetCurrentLoc(statement_loc);
+ Location loc = IRBuilder::Current()->ResolveLoc(statement_loc);
if (!loc.as<UnknownLocNode>()) {
CheckBindingBlockFrameExistAndUnended()->binding_locs.Set(var, loc);
}
diff --git a/src/script/ir_builder/base.cc b/src/script/ir_builder/base.cc
index 669618ddd4..d874c60fb0 100644
--- a/src/script/ir_builder/base.cc
+++ b/src/script/ir_builder/base.cc
@@ -124,7 +124,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
IRBuilderFrameObj::IRBuilderFrameObj() {
- if (IRBuilder::IsInScope()) loc = IRBuilder::Current()->GetCurrentLoc();
+ if (IRBuilder::IsInScope()) loc = IRBuilder::Current()->ResolveLoc();
}
void IRBuilderFrameObj::EnterWithScope() {
@@ -162,7 +162,7 @@ void IRBuilderObj::PopLoc() {
locs.pop_back();
}
-Location IRBuilderObj::GetCurrentLoc(Location location) const {
+Location IRBuilderObj::ResolveLoc(Location location) const {
std::vector<Location> normalized;
normalized.reserve(locs.size());
for (const Location& loc : locs) {
@@ -172,12 +172,8 @@ Location IRBuilderObj::GetCurrentLoc(Location location)
const {
return NormalizedLoc(normalized);
}
-ffi::ObjectRef IRBuilderObj::SetCurrentLoc(ffi::ObjectRef obj) const {
- return SetLoc(std::move(obj), UnknownLoc());
-}
-
-ffi::ObjectRef IRBuilderObj::SetLoc(ffi::ObjectRef obj, Location loc) const {
- loc = ComposeLoc(GetCurrentLoc(), loc);
+ffi::ObjectRef IRBuilderObj::AttachLoc(ffi::ObjectRef obj, Location loc) const
{
+ loc = ComposeLoc(ResolveLoc(), loc);
if (!loc.as<UnknownLocNode>() && obj.defined()) {
if (Location* target = details::LocationAccessor::vtable()(obj)) {
*target = ComposeLoc(loc, *target);
@@ -266,8 +262,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_method("script.ir_builder.IRBuilderGet",
&IRBuilderObj::Get<ffi::ObjectRef>)
.def_method("script.ir_builder.IRBuilderPushLoc", &IRBuilderObj::PushLoc)
.def_method("script.ir_builder.IRBuilderPopLoc", &IRBuilderObj::PopLoc)
- .def_method("script.ir_builder.IRBuilderSetCurrentLoc",
&IRBuilderObj::SetCurrentLoc)
- .def_method("script.ir_builder.IRBuilderSetLoc", &IRBuilderObj::SetLoc)
+ .def_method("script.ir_builder.IRBuilderAttachLoc",
&IRBuilderObj::AttachLoc)
.def("script.ir_builder.IRBuilderName", IRBuilder::Name<ffi::ObjectRef>);
}
diff --git a/src/script/ir_builder/utils.h b/src/script/ir_builder/utils.h
index c283a201b7..7c060da297 100644
--- a/src/script/ir_builder/utils.h
+++ b/src/script/ir_builder/utils.h
@@ -44,7 +44,7 @@ inline IRModuleFrame FindModuleFrame(const ffi::String&
method) {
/*! \brief Add an eager statement under the current source-call context. */
inline void AddToParent(tvm::Stmt stmt) {
// Some builder paths use an undefined statement as an omitted branch.
- if (stmt.defined()) IRBuilder::Current()->SetCurrentLoc(stmt);
+ if (stmt.defined()) IRBuilder::Current()->AttachLoc(stmt);
AddToParent(std::move(stmt), UnknownLoc());
}
diff --git a/tests/python/script/test_ir_builder.py
b/tests/python/script/test_ir_builder.py
index 0680ea5b4d..edf6f9ba55 100644
--- a/tests/python/script/test_ir_builder.py
+++ b/tests/python/script/test_ir_builder.py
@@ -49,3 +49,82 @@ def test_ir_builder_irmodule():
ir_module_expected = ir.IRModule(None, None)
assert_structural_equal(ir_module_actual, ir_module_expected,
map_free_vars=True)
+
+
+def _source_loc(name, line, column=1, end_line=None, end_column=20):
+ return ir.SourceLoc(ir.SourceName(name), line, column, end_line or line,
end_column)
+
+
+def _source_chain(loc):
+ if isinstance(loc, ir.CallSiteLoc):
+ return [*_source_chain(loc.caller), *_source_chain(loc.callee)]
+ return [loc]
+
+
+def test_attach_loc_normalizes_full_stack_and_preserves_identity():
+ caller = _source_loc("caller.py", 2)
+ outer = _source_loc("definition.py", 1, end_line=10)
+ inner = _source_loc("definition.py", 3)
+ leaf = _source_loc("leaf.py", 5)
+ node = ir.prim.IntImm("int32", 1)
+ with IRBuilder() as ib:
+ with ib.with_loc(caller), ib.with_loc(outer),
ib.with_loc(ir.UnknownLoc()):
+ with ib.with_loc(ir.CallSiteLoc(inner, caller)):
+ assert ib._attach_loc(node, leaf) is node
+ assert_structural_equal(node.loc, ir.CallSiteLoc(leaf,
ir.CallSiteLoc(inner, caller)))
+
+
+def test_attach_loc_preserves_shared_prefix_and_receipt_identity():
+ from tvm.script.ir_builder.base import AlreadyEmitted, at
+
+ caller = _source_loc("caller.py", 2)
+ callee = _source_loc("definition.py", 3)
+ node = ir.prim.IntImm("int32", 1)
+ receipt = AlreadyEmitted(node)
+ container = [node]
+ with IRBuilder() as ib:
+ with ib.with_loc(caller):
+ assert ib._attach_loc(node) is node
+ with ib.with_loc(callee):
+ assert at(caller, receipt) is receipt
+ assert receipt.value is node
+ assert at(caller, container) is container
+ assert_structural_equal(node.loc, ir.CallSiteLoc(callee, caller))
+
+
+def test_unknown_attachment_and_scope_exception_restore_context():
+ from tvm.script.ir_builder.base import at
+
+ caller = _source_loc("caller.py", 2)
+ callee = _source_loc("definition.py", 3)
+ node = ir.prim.IntImm("int32", 1)
+ with IRBuilder() as ib:
+ assert ib._attach_loc(node) is node
+ assert isinstance(node.loc, ir.UnknownLoc)
+ with ib.with_loc(caller):
+ assert at(ir.UnknownLoc(), node) is node
+ assert isinstance(node.loc, ir.UnknownLoc)
+ with pytest.raises(RuntimeError, match="scope failure"):
+ with ib.with_loc(callee):
+ raise RuntimeError("scope failure")
+ ib._attach_loc(node)
+ assert node.loc.same_as(caller)
+ untouched = ir.prim.IntImm("int32", 2)
+ ib._attach_loc(untouched)
+ assert isinstance(untouched.loc, ir.UnknownLoc)
+ assert not IRBuilder.is_in_scope()
+
+
+def test_frame_captures_full_location_stack_before_entering():
+ caller = _source_loc("caller.py", 2)
+ callee = _source_loc("definition.py", 3)
+ with IRBuilder() as ib:
+ with ib.with_loc(caller), ib.with_loc(callee):
+ frame = I.ir_module()
+ assert_structural_equal(frame.loc, ir.CallSiteLoc(callee, caller))
+ with frame:
+ pass
+ assert [loc.source_name.name for loc in _source_chain(frame.loc)] == [
+ "caller.py",
+ "definition.py",
+ ]
diff --git a/tests/python/tirx/script/test_tirx_script_source_locations.py
b/tests/python/tirx/script/test_tirx_script_source_locations.py
index 13e2c66948..58c20e915d 100644
--- a/tests/python/tirx/script/test_tirx_script_source_locations.py
+++ b/tests/python/tirx/script/test_tirx_script_source_locations.py
@@ -425,3 +425,23 @@ def _position(test, statement, expression):
def _loc_position(loc):
return loc.start_line, loc.start_column, loc.end_line, loc.end_column
+
+
+def test_eager_and_deferred_nodes_keep_nested_locations_and_receipts():
+ from tvm.script.ir_builder.base import at
+
+ caller = ir.SourceLoc(ir.SourceName("caller.py"), 2, 1, 2, 20)
+ callee = ir.SourceLoc(ir.SourceName("definition.py"), 3, 1, 3, 20)
+ explicit = ir.SourceLoc(ir.SourceName("definition.py"), 3, 5, 3, 10)
+ with I.IRBuilder() as ib:
+ with ib.with_loc(caller), ib.with_loc(callee):
+ loop = T.serial(2)
+ with loop:
+ with ib.with_loc(caller), ib.with_loc(callee):
+ receipt = T.evaluate(1)
+ assert at(explicit, receipt) is receipt
+ result = ib.get()
+ assert isinstance(result, ir.For)
+ assert_structural_equal(result.loc, ir.CallSiteLoc(callee, caller))
+ assert receipt.value.same_as(result.body[0])
+ assert_structural_equal(receipt.value.loc, ir.CallSiteLoc(explicit,
caller))
