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 8ff466b351 feat(tvmscript): add scoped Python line tracing for
IRBuilder (#20731)
8ff466b351 is described below
commit 8ff466b3517e38a3b445662822db97d77886a783
Author: Bohan Hou <[email protected]>
AuthorDate: Wed Oct 14 17:39:28 2026 -0400
feat(tvmscript): add scoped Python line tracing for IRBuilder (#20731)
Add `IRBuilder.line_trace_scope()` for frontends that execute ordinary
Python to construct IR. Traced operations retain their source lines and
helper call sites, with optional implementation-path exclusions and
restoration of prior tracing state on exit.
Use the owned `PushLoc` / `UpdateLoc` location slots from #20727 so
advancing a traced line preserves nested explicit `with_loc` scopes and
deferred frame locations. Cover same-line helper returns, nested
builders, existing trace callbacks, and exception cleanup. Initialize
CUDA in void-emission tests through the idempotent backend loader so
repeated setup preserves canonical script namespaces.
Void emission uses the native script API boundary independently of
tracing. Raw operators construct Calls; script-facing void operators
emit whenever an IRBuilder is active. Apply the existing native
decorator to the CUDA tile `apply` and `ewise` exports, and test both
with tracing enabled and disabled. Document tracing and native Tensor
parameter contracts, including compiled host checks for shared
dimensions.
Downstream TIRx-lite migration:
https://github.com/mlc-ai/TIRx-kernels/pull/11 (merge this TVM change
first).
Validation:
- Native build and changed-file pre-commit passed.
- 365 script, namespace, and tracing tests passed (2 skipped, 1
xfailed), including the combined test order that previously replaced
canonical CUDA namespaces.
- The new documentation sections parse without warnings.
- 167 downstream lite API, entry, IR-contract, descriptor, registry, and
compiler tests passed.
- Bench-suite imports passed for 93 kernels; strict registry discovery
passed for 105 kernels.
- Full TIRx suite on all seven Blackwell GPUs: 3,994 passed, 300
skipped, and 5 xfailed in 561.08 seconds. Eleven cases hit GPU-memory
contention while allocating test inputs; all eleven passed when rerun on
physical GPU 6.
---------
Co-authored-by: Tianqi Chen <[email protected]>
---
docs/tirx/api/script.rst | 57 ++++
include/tvm/script/ir_builder/base.h | 14 +-
python/tvm/backend/cuda/script.py | 2 +
python/tvm/script/ir_builder/base.py | 15 +
python/tvm/script/ir_builder/tracing.py | 163 +++++++++++
src/script/ir_builder/base.cc | 12 +-
tests/python/script/test_builder_void_emission.py | 7 +-
.../python/tirx/script/test_tirx_script_tracing.py | 312 +++++++++++++++++++++
8 files changed, 575 insertions(+), 7 deletions(-)
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index ce3c2f6e36..196d407a9c 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -57,6 +57,63 @@ thread-placement and tensor alias policies remain in the
TIRx/S-TIR builders.
Marker used by ``Tx.let`` and ``Tx.let[dtype]`` annotations to construct an
explicit ``LetStmt``.
+Tracing ordinary Python
+-----------------------
+
+The native builders also support frontends that execute ordinary Python to
+construct IR. ``IRBuilder.line_trace_scope()`` records file-backed Python
source
+lines. Script-facing void calls automatically emit into an active IRBuilder,
+with or without line tracing::
+
+ from tvm.script import tirx as Tx
+ from tvm.script.ir_builder import IRBuilder
+ from tvm.tirx.script import ir_builder as I
+
+ with IRBuilder() as builder:
+ with I.function_():
+ I.func_name_("synchronize")
+ with builder.line_trace_scope():
+ with Tx.device_entry(
+ launch=Tx.cuda.LaunchConfig(grid=1, block=32)
+ ):
+ Tx.ptx["bar.sync"](0)
+ function = builder.get()
+
+Auto emission applies to script-facing calls through ``Tx``, including PTX
+and backend namespaces. A void call returns an ``AlreadyEmitted`` receipt;
+passing that receipt to ``evaluate`` or ``emit_`` does not emit it again.
+Calls producing values remain expressions. Raw IR constructors such as
+``tvm.ir.Call``, ``tvm.tirx.op`` and raw backend APIs remain pure even inside
+an IRBuilder. Line tracing only supplies source locations; it neither enables
+nor disables emission. The trace scope restores its previous trace and location
+state on exit, including exceptional exits, and applies only to its owning
builder.
+Explicit ``with_loc`` and ``at_`` locations compose with the traced lines.
+
+Line tracing excludes TVM and standard-library implementation files. Frontends
+can supply additional directories using ``line_trace_scope(exclude_paths=(...
,))``.
+User helpers retain their own source line and the caller's location through
+``CallSiteLoc``. These locations cover whole lines; the TVMScript parser
records
+expression ranges. Python loops and branches execute during construction; use
+the native statement frames to construct runtime control flow.
+
+Tensor parameters and shared dimensions
+---------------------------------------
+
+Traced frontends can bind native types directly with ``I.arg_``. A shared
symbolic
+dimension imposes the same size constraint on multiple tensor arguments::
+
+ n = Tx.Var("n", Tx.int64)
+ # Inside an I.function_() frame:
+ x = I.arg_("x", Tx.Tensor((n,), Tx.float32))
+ y = I.arg_("y", Tx.Tensor((n,), Tx.float32))
+
+The compiled host entry accepts tensor arguments, including PyTorch tensors,
+and checks their dtype, rank and shared dimensions. Use independent symbolic
+dimensions for unrelated sizes. If the same dimension is also an explicit
scalar
+parameter, ``I.arg_("n", n)`` binds that existing Var. ``x.data`` and
+``x.ptr_to(indices)`` provide pointers inside the function. ``Tx.Ptr`` retains
+its raw-pointer ABI and does not imply tensor argument conversion.
+
Explicit register fragments
---------------------------
diff --git a/include/tvm/script/ir_builder/base.h
b/include/tvm/script/ir_builder/base.h
index 7add479039..bccff9eec2 100644
--- a/include/tvm/script/ir_builder/base.h
+++ b/include/tvm/script/ir_builder/base.h
@@ -24,6 +24,7 @@
#include <tvm/ir/function.h>
#include <tvm/ir/object_functor.h>
+#include <cstdint>
#include <vector>
namespace tvm {
@@ -175,8 +176,17 @@ class IRBuilderObj : public ffi::Object {
*/
template <typename TObjectRef>
inline TObjectRef Get() const;
- /*! \brief Push a frontend source location for IR constructed in the nested
scope. */
- TVM_TIRX_DLL void PushLoc(Location loc);
+ /*!
+ * \brief Push a frontend source location for IR constructed in the nested
scope.
+ * \return The index of the pushed scope, valid until it is popped.
+ */
+ TVM_TIRX_DLL int64_t PushLoc(Location loc);
+ /*!
+ * \brief Update a source location without changing its nested scopes.
+ * \param index The index returned by PushLoc for a scope that has not been
popped.
+ * \param loc The replacement location for that scope.
+ */
+ TVM_TIRX_DLL void UpdateLoc(int64_t index, Location loc);
/*! \brief Pop the innermost frontend source location. */
TVM_TIRX_DLL void PopLoc();
/*! \brief Normalize the full active source stack with an optional explicit
location. */
diff --git a/python/tvm/backend/cuda/script.py
b/python/tvm/backend/cuda/script.py
index 871da7af12..4b6cdedec8 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -159,6 +159,8 @@ __activemask = _auto_emit_void_call(_cuda_op.__activemask)
tile = _make_tensor_namespace()
+tile.apply = _auto_emit_void_call(tile.apply)
+tile.ewise = _auto_emit_void_call(tile.ewise)
_ir.op._init_op_api("tirx.cuda", __name__, wrapper=_auto_emit_void_call)
diff --git a/python/tvm/script/ir_builder/base.py
b/python/tvm/script/ir_builder/base.py
index 52ac16dcbc..d87d18afff 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -203,6 +203,21 @@ class IRBuilder(_Object):
"""Get the constructed IR."""
return _ffi_api.IRBuilderGet(self) # type: ignore[attr-defined] #
pylint: disable=no-member
+ def line_trace_scope(self, *, exclude_paths=()):
+ """Track Python source lines while executing builder code.
+
+ TVM and Python implementation frames are excluded automatically.
+ ``exclude_paths`` supplies additional implementation directories.
+ User helpers retain their own line and the caller's location. Nested
+ scopes and exceptions restore the preceding trace and location state.
+ Locations cover whole lines; use the parser for expression ranges.
+ Tracing only supplies locations: script-facing void calls emit under
+ any active builder, independently of this scope.
+ """
+ from .tracing import LineTraceScope # pylint:
disable=import-outside-toplevel
+
+ return LineTraceScope(self, exclude_paths=exclude_paths)
+
@contextmanager
def with_loc(self, loc):
"""Attach ``loc`` to IR nodes constructed in the nested scope.
diff --git a/python/tvm/script/ir_builder/tracing.py
b/python/tvm/script/ir_builder/tracing.py
new file mode 100644
index 0000000000..c6ee9d69db
--- /dev/null
+++ b/python/tvm/script/ir_builder/tracing.py
@@ -0,0 +1,163 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Executed Python source locations for native IR builders."""
+
+import linecache
+import sys
+import sysconfig
+from pathlib import Path
+
+import tvm_ffi
+
+from tvm import ir
+
+from . import _ffi_api
+from .base import IRBuilder
+
+
+class LineTraceScope:
+ """Track file-backed Python lines under one builder, composing helper
locations."""
+
+ def __init__(self, builder, *, exclude_paths=()):
+ self.builder = builder
+ self.excluded = (
+ Path(__file__).resolve().parents[2],
+ Path(tvm_ffi.__file__).resolve().parent,
+ *(Path(path).resolve() for path in exclude_paths),
+ )
+ self.stdlib = Path(sysconfig.get_paths()["stdlib"]).resolve()
+ self.site_roots = tuple(
+ Path(sysconfig.get_paths()[key]).resolve() for key in ("purelib",
"platlib")
+ )
+ self.files = {}
+ self.locations = {}
+ self.frames = {}
+ self.previous_callbacks = {}
+ self.slot = None
+ self.saved_traces = {}
+ self.active_sources = ()
+ self.previous_trace = None
+
+ def _is_user_source(self, filename):
+ if filename not in self.files:
+ path = Path(filename).resolve()
+ excluded = any(path.is_relative_to(root) for root in self.excluded)
+ stdlib = path.is_relative_to(self.stdlib) and not any(
+ path.is_relative_to(root) for root in self.site_roots
+ )
+ self.files[filename] = not filename.startswith("<") and not
excluded and not stdlib
+ return self.files[filename]
+
+ def source_loc(self, filename, line):
+ """Return a whole-line location, or None for excluded/generated
source."""
+ if not self._is_user_source(filename):
+ return None
+ key = (filename, line)
+ if key not in self.locations:
+ text = linecache.getline(filename, line).rstrip("\r\n")
+ self.locations[key] = ir.SourceLoc(
+ ir.SourceName(filename), line, 1, line, max(2, len(text) + 1)
+ )
+ return self.locations[key]
+
+ def _activate(self, frame):
+ if frame not in self.saved_traces:
+ self.saved_traces[frame] = (frame.f_trace, frame.f_trace_lines,
frame.f_trace_opcodes)
+ frame.f_trace = self._trace
+ frame.f_trace_lines = True
+
+ def __enter__(self):
+ self.previous_trace = sys.gettrace()
+ self.root_frame = sys._getframe(1)
+ # Own one stable slot below any explicit with_loc scopes in the body.
+ self.slot = _ffi_api.IRBuilderPushLoc(self.builder, ir.UnknownLoc())
+ try:
+ if self._is_user_source(self.root_frame.f_code.co_filename):
+ if self.root_frame.f_trace is not None:
+ self.previous_callbacks[self.root_frame] =
self.root_frame.f_trace
+ self._activate(self.root_frame)
+ self.frames[self.root_frame] = self.source_loc(
+ self.root_frame.f_code.co_filename,
self.root_frame.f_lineno
+ )
+ self._restore(self.root_frame)
+ sys.settrace(self._trace)
+ except BaseException:
+ self.__exit__(*sys.exc_info())
+ raise
+ return self
+
+ def __exit__(self, *_):
+ sys.settrace(self.previous_trace)
+ for frame, (trace, lines, opcodes) in self.saved_traces.items():
+ frame.f_trace = trace
+ frame.f_trace_lines = lines
+ frame.f_trace_opcodes = opcodes
+ self.saved_traces.clear()
+ _ffi_api.IRBuilderPopLoc(self.builder)
+ self.slot = None
+ self.active_sources = ()
+ self.frames.clear()
+ self.previous_callbacks.clear()
+ self.root_frame = None
+
+ def _restore(self, frame):
+ locs = []
+ while frame is not None:
+ if frame in self.frames:
+ locs.append(self.frames[frame])
+ frame = frame.f_back
+ sources = tuple(id(loc) for loc in locs)
+ if sources == self.active_sources:
+ return
+ self.active_sources = sources
+ loc = ir.UnknownLoc()
+ if locs:
+ loc = locs[-1]
+ for callee in reversed(locs[:-1]):
+ loc = ir.CallSiteLoc(callee, loc)
+ _ffi_api.IRBuilderUpdateLoc(self.builder, self.slot, loc)
+
+ def _trace(self, frame, event, arg):
+ callback = self.previous_trace if event == "call" else
self.previous_callbacks.get(frame)
+ if callback is not None:
+ next_callback = callback(frame, event, arg)
+ if next_callback is not None:
+ self.previous_callbacks[frame] = next_callback
+ else:
+ self.previous_callbacks.pop(frame, None)
+ user = (
+ self._is_user_source(frame.f_code.co_filename)
+ and IRBuilder.is_in_scope()
+ and IRBuilder.current().same_as(self.builder)
+ )
+ if user and event == "call":
+ self._activate(frame)
+ elif user and event == "line":
+ self.frames[frame] = self.source_loc(frame.f_code.co_filename,
frame.f_lineno)
+ self._restore(frame)
+ elif event == "return":
+ previous = self.previous_callbacks.pop(frame, None)
+ if frame in self.frames:
+ self.frames.pop(frame)
+ self._restore(frame.f_back)
+ if frame in self.saved_traces and frame is not self.root_frame:
+ trace, lines, opcodes = self.saved_traces.pop(frame)
+ frame.f_trace = previous if previous is not None else trace
+ frame.f_trace_lines = lines
+ frame.f_trace_opcodes = opcodes
+ return frame.f_trace
+ return self._trace if user or frame in self.previous_callbacks else
None
diff --git a/src/script/ir_builder/base.cc b/src/script/ir_builder/base.cc
index d874c60fb0..6a2d83ebcf 100644
--- a/src/script/ir_builder/base.cc
+++ b/src/script/ir_builder/base.cc
@@ -154,7 +154,16 @@ IRBuilder::IRBuilder() {
data_ = n;
}
-void IRBuilderObj::PushLoc(Location loc) { locs.push_back(std::move(loc)); }
+int64_t IRBuilderObj::PushLoc(Location loc) {
+ locs.push_back(std::move(loc));
+ return static_cast<int64_t>(locs.size()) - 1;
+}
+
+void IRBuilderObj::UpdateLoc(int64_t index, Location loc) {
+ TVM_FFI_CHECK(index >= 0 && static_cast<size_t>(index) < locs.size(),
ValueError)
+ << "Invalid source location scope";
+ locs[index] = std::move(loc);
+}
void IRBuilderObj::PopLoc() {
TVM_FFI_CHECK(!locs.empty(), ValueError)
@@ -261,6 +270,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.IRBuilderIsInScope", IRBuilder::IsInScope)
.def_method("script.ir_builder.IRBuilderGet",
&IRBuilderObj::Get<ffi::ObjectRef>)
.def_method("script.ir_builder.IRBuilderPushLoc", &IRBuilderObj::PushLoc)
+ .def_method("script.ir_builder.IRBuilderUpdateLoc",
&IRBuilderObj::UpdateLoc)
.def_method("script.ir_builder.IRBuilderPopLoc", &IRBuilderObj::PopLoc)
.def_method("script.ir_builder.IRBuilderAttachLoc",
&IRBuilderObj::AttachLoc)
.def("script.ir_builder.IRBuilderName", IRBuilder::Name<ffi::ObjectRef>);
diff --git a/tests/python/script/test_builder_void_emission.py
b/tests/python/script/test_builder_void_emission.py
index 446e3c3059..68d732d7e5 100644
--- a/tests/python/script/test_builder_void_emission.py
+++ b/tests/python/script/test_builder_void_emission.py
@@ -19,16 +19,15 @@
import pytest
import tvm_ffi
-from tvm import ir, tirx
-from tvm.backend import cuda
+from tvm import backend, ir, tirx
from tvm.script import ir_builder as I
from tvm.script import tirx as T
from tvm.script.ir_builder.base import AlreadyEmitted
@pytest.fixture(scope="module", autouse=True)
-def register_cuda():
- cuda.register_backend()
+def load_cuda():
+ backend.load("cuda")
@pytest.mark.parametrize("ty", [ir.PrimType("void"), ir.TupleType([])])
diff --git a/tests/python/tirx/script/test_tirx_script_tracing.py
b/tests/python/tirx/script/test_tirx_script_tracing.py
new file mode 100644
index 0000000000..3b22f46cba
--- /dev/null
+++ b/tests/python/tirx/script/test_tirx_script_tracing.py
@@ -0,0 +1,312 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, missing-docstring
+"""Native Python tracing, emission and parameter contracts."""
+
+import linecache
+import sys
+from contextlib import nullcontext
+
+import pytest
+from tvm_ffi import structural_walk
+
+import tvm
+from tvm import ir
+from tvm.script import tirx as T
+from tvm.script.ir_builder import AlreadyEmitted, IRBuilder
+from tvm.tirx.script import ir_builder as I
+
+
+def _calls(function):
+ calls = []
+ structural_walk(function, (ir.Call, calls.append))
+ return calls
+
+
+def _leaf(loc):
+ while isinstance(loc, ir.CallSiteLoc):
+ loc = loc.callee
+ return loc
+
+
[email protected]("trace", [False, True])
+def test_void_calls_emit_once_and_values_remain_expressions(trace):
+ with IRBuilder() as builder:
+ with I.function_():
+ with builder.line_trace_scope() if trace else nullcontext():
+ receipt = T.ptx["bar.sync"](0)
+ assert isinstance(receipt, AlreadyEmitted)
+ assert T.evaluate(receipt) is receipt
+ I.emit_(receipt)
+ T.gpu_storage_sync("shared")
+ T.cuda.func_call("side_effect", "__device__ void side_effect()
{}")
+ value = T.cuda.elect_sync()
+ assert isinstance(value, ir.Call)
+ raw = ir.Call("tirx.gpu_storage_sync",
[ir.StringImm("shared")])
+ assert isinstance(raw, ir.Call)
+ T.evaluate(raw)
+ names = [call.op.name for call in _calls(builder.get())]
+ assert names.count("tirx.ptx.bar_sync") == 1
+ assert names.count("tirx.gpu_storage_sync") == 2
+ assert "tirx.cuda.elect_sync" not in names
+
+
+def test_auto_emit_uses_current_builder_without_tracing():
+ with IRBuilder() as outer:
+ with I.function_():
+ T.ptx["bar.sync"](0)
+ with pytest.raises(RuntimeError, match="failed"):
+ with IRBuilder() as inner:
+ with I.function_():
+ receipt = T.ptx["bar.sync"](1)
+ assert isinstance(receipt, AlreadyEmitted)
+ raise RuntimeError("failed")
+ T.ptx["bar.sync"](2)
+ assert isinstance(T.ptx["bar.sync"](3), ir.Call)
+ assert len(_calls(outer.get())) == 2
+ assert inner.result is None # Aborted functions are not finalized.
+
+
[email protected]("trace", [False, True])
+def test_backend_tile_namespace_emits_once(trace):
+ from tvm.backend.cuda.tile.patterns import apply as raw_apply
+ from tvm.backend.cuda.tile.patterns.ewise import ewise as raw_ewise
+
+ @T.function(private=True)
+ def identity(value: T.float32) -> T.float32:
+ return value
+
+ with IRBuilder() as builder:
+ with I.function_():
+ source = I.arg_("source", T.Tensor((1, 4), T.float32))
+ dest = I.bind(T.alloc_tensor(ty_args=[T.Tensor((1, 4), T.float32,
scope="local")]))
+ with builder.line_trace_scope() if trace else nullcontext():
+ receipt = T.cuda.tile.apply(T.ptx["ld.global.v4.f32"], dest,
source)
+ assert isinstance(receipt, AlreadyEmitted)
+ T.evaluate(receipt)
+ receipt = T.cuda.tile.ewise(identity, dest, dest)
+ assert isinstance(receipt, AlreadyEmitted)
+ T.evaluate(receipt)
+ assert isinstance(raw_apply(T.ptx["ld.global.v4.f32"], dest,
source), ir.Call)
+ assert isinstance(raw_ewise(identity, dest, dest), ir.Call)
+ names = [call.op.name for call in _calls(builder.get())]
+ assert names.count("tirx.cuda.tile.apply") == 1
+ assert names.count("tirx.cuda.tile.ewise") == 1
+
+
+def test_ptx_failed_and_ambiguous_candidates_emit_nothing():
+ with IRBuilder() as builder:
+ with I.function_():
+ with pytest.raises(ValueError):
+ T.ptx["bar.sync"]()
+ chain = T.ptx["bar.sync"]
+ ambiguous = type(chain)(chain._cands * 2)
+ with pytest.raises(AssertionError, match="ambiguous"):
+ ambiguous(0)
+ T.ptx["bar.sync"](1)
+ assert len(_calls(builder.get())) == 1
+
+
+def test_line_trace_tracks_direct_calls_helpers_and_frames():
+ def helper():
+ T.ptx["bar.sync"](1)
+
+ with IRBuilder() as builder:
+ with I.function_():
+ with builder.line_trace_scope():
+ T.ptx["bar.sync"](0)
+ helper()
+ with T.serial(2):
+ T.ptx["bar.sync"](2)
+ calls = _calls(builder.get())
+ assert len(calls) == 3
+ for call, fragment in zip(calls, ("(0)", "(1)", "(2)")):
+ loc = _leaf(call.loc)
+ assert loc.source_name.name == __file__
+ assert fragment in linecache.getline(__file__, loc.start_line)
+ assert isinstance(calls[1].loc, ir.CallSiteLoc)
+ assert "helper()" in linecache.getline(__file__,
_leaf(calls[1].loc.caller).start_line)
+
+
+def test_line_trace_preserves_existing_callback_and_restores_on_exception():
+ previous = sys.gettrace()
+ events = []
+
+ def existing(frame, event, arg):
+ if frame.f_code.co_name == "helper":
+ events.append(event)
+ return existing
+
+ def helper():
+ T.ptx["bar.sync"](0)
+ raise RuntimeError("failed")
+
+ frame = sys._getframe()
+ saved = (frame.f_trace, frame.f_trace_lines, frame.f_trace_opcodes)
+ sys.settrace(existing)
+ try:
+ frame.f_trace = existing
+ frame.f_trace_lines = False
+ frame.f_trace_opcodes = True
+ with IRBuilder() as builder:
+ with I.function_():
+ with pytest.raises(RuntimeError, match="failed"):
+ with builder.line_trace_scope():
+ helper()
+ assert sys.gettrace() is existing
+ assert {"call", "line", "return"} <= set(events)
+ assert (frame.f_trace, frame.f_trace_lines, frame.f_trace_opcodes) == (
+ existing,
+ False,
+ True,
+ )
+ finally:
+ sys.settrace(previous)
+ frame.f_trace, frame.f_trace_lines, frame.f_trace_opcodes = saved
+
+
+def test_nested_line_trace_restores_outer_builder():
+ with IRBuilder() as outer:
+ with I.function_():
+ with outer.line_trace_scope():
+ T.ptx["bar.sync"](0)
+ with IRBuilder() as inner:
+ with I.function_():
+ with inner.line_trace_scope():
+ T.ptx["bar.sync"](1)
+ T.ptx["bar.sync"](2)
+ for function, fragments in ((outer.get(), ("(0)", "(2)")), (inner.get(),
("(1)",))):
+ calls = _calls(function)
+ assert len(calls) == len(fragments)
+ for call, fragment in zip(calls, fragments):
+ assert fragment in linecache.getline(__file__,
_leaf(call.loc).start_line)
+
+
+def _source_chain(loc):
+ if isinstance(loc, ir.CallSiteLoc):
+ return [*_source_chain(loc.caller), *_source_chain(loc.callee)]
+ return [loc]
+
+
+def test_trace_preserves_explicit_scopes_and_deferred_frames():
+ from tvm.script.ir_builder import at_
+
+ caller = ir.SourceLoc(ir.SourceName("caller.py"), 2, 1, 2, 20)
+ explicit = ir.SourceLoc(ir.SourceName("definition.py"), 3, 1, 3, 20)
+ attached = ir.SourceLoc(ir.SourceName("definition.py"), 3, 5, 3, 10)
+
+ def helper():
+ return T.evaluate(1)
+
+ with IRBuilder() as builder:
+ with builder.with_loc(caller):
+ with builder.line_trace_scope():
+ loop_line = sys._getframe().f_lineno + 1
+ loop = T.serial(2)
+ with loop, builder.line_trace_scope():
+ same_line = sys._getframe().f_lineno + 1
+ first, second = helper(), T.evaluate(2)
+ with builder.with_loc(explicit):
+ third = T.evaluate(3)
+ fourth = at_(attached, T.evaluate(4))
+ last = T.evaluate(5)
+ result = builder.get()
+ assert isinstance(result, ir.For)
+ assert result.loc.caller.same_as(caller)
+ assert result.loc.callee.start_line == loop_line
+ assert len(result.body.seq) == 5
+ for receipt, statement, fragment, extra in zip(
+ (first, second, third, fourth, last),
+ result.body.seq,
+ (
+ "return T.evaluate(1)",
+ "helper(), T.evaluate(2)",
+ "T.evaluate(3)",
+ "T.evaluate(4)",
+ "T.evaluate(5)",
+ ),
+ (None, None, explicit, attached, None),
+ ):
+ assert receipt.value.same_as(statement)
+ chain = _source_chain(statement.loc)
+ assert chain[0].same_as(caller)
+ if extra is not None:
+ assert chain.pop().same_as(extra)
+ assert chain[-1].source_name.name == __file__
+ assert fragment in linecache.getline(__file__, chain[-1].start_line)
+ # Returning from a helper restores the caller even on the same Python line.
+ assert _leaf(second.value.loc).start_line == same_line
+ assert _source_chain(first.value.loc)[-2].start_line == same_line
+ with builder:
+ assert isinstance(T.evaluate(0).value.loc, ir.UnknownLoc)
+
+
+def test_nested_trace_callback_failure_restores_outer_scope(monkeypatch):
+ from tvm.script.ir_builder import _ffi_api
+
+ original = _ffi_api.IRBuilderUpdateLoc
+ fail = [False]
+
+ def update(*args):
+ if fail[0]:
+ fail[0] = False
+ raise RuntimeError("trace callback failed")
+ return original(*args)
+
+ monkeypatch.setattr(_ffi_api, "IRBuilderUpdateLoc", update)
+ previous = sys.gettrace()
+ frame = sys._getframe()
+ saved = (frame.f_trace, frame.f_trace_lines, frame.f_trace_opcodes)
+ with IRBuilder() as builder:
+ with I.function_():
+ with builder.line_trace_scope():
+ callback = sys.gettrace()
+ with pytest.raises(RuntimeError, match="trace callback
failed"):
+ with builder.line_trace_scope():
+ fail[0] = True
+ T.evaluate(0)
+ assert sys.gettrace() == callback
+ receipt = T.evaluate(5)
+ assert sys.gettrace() is previous
+ assert (frame.f_trace, frame.f_trace_lines, frame.f_trace_opcodes) == saved
+ assert "receipt = T.evaluate(5)" in linecache.getline(
+ __file__, _leaf(receipt.value.loc).start_line
+ )
+ assert len(builder.get().body.seq) == 1
+
+
+def test_shared_tensor_shapes_are_checked_by_compiled_host():
+ import torch
+
+ n = T.Var("n", T.int64)
+ with IRBuilder() as builder:
+ with I.function_():
+ I.func_name_("check_shape")
+ I.func_attr({"global_symbol": "check_shape"})
+ I.func_ret_type_(T.int64)
+ I.arg_("x", T.Tensor((n,), T.float32))
+ I.arg_("y", T.Tensor((n,), T.float32))
+ I.return_(n)
+ compiled = tvm.compile(builder.get(), target="llvm", tir_pipeline="tirx")
+ assert compiled(torch.zeros(3), torch.zeros(3)) == 3
+ for other, error, message in (
+ (torch.zeros(4), ValueError, "expected to match"),
+ (torch.zeros(3, dtype=torch.int32), TypeError, "expected float32"),
+ (torch.zeros(1, 3), ValueError, "expected 1"),
+ ):
+ with pytest.raises(error, match=message):
+ compiled(torch.zeros(3), other)
