[Mlir-commits] [mlir] [mlir][python] Add location source composition to loc_tracebacks() (PR #192310)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Apr 15 11:42:11 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Alisson Azzolini (aazzolini)
<details>
<summary>Changes</summary>
## Summary
Add two new kwargs to `loc_tracebacks()` controlling how the three location sources (explicit `loc=`, traceback, `Location.current`) compose:
- **`on_explicit`**: `"use_explicit"` (default) | `"use_traceback"` — what to do when explicit `loc=` is passed
- **`current_loc`**: `"fallback"` (default) | `"nameloc_wrap"` — how to compose `Location.current` NameLoc scopes on top
The two flags are orthogonal and only take effect when `loc_tracebacks()` is active.
## Use cases
1. **DSL profiling attribution**: `loc_tracebacks(current_loc="nameloc_wrap", on_explicit="use_traceback")` — push NameLoc scopes for profiling and have them survive traceback generation
2. **Override explicit locs with tracebacks**: `loc_tracebacks(on_explicit="use_traceback")` — get full Python stack traces without removing every `loc=` argument
## Test plan
- [x] All 102 existing Python binding tests pass (5 unsupported unchanged)
- [x] 6 new tests covering: nameloc_wrap, use_explicit default, use_traceback, DSL profiling combo, orthogonality (use_explicit + nameloc_wrap), tracebacks-disabled fallback
---
Full diff: https://github.com/llvm/llvm-project/pull/192310.diff
5 Files Affected:
- (modified) mlir/include/mlir/Bindings/Python/Globals.h (+22)
- (modified) mlir/lib/Bindings/Python/Globals.cpp (+24)
- (modified) mlir/lib/Bindings/Python/IRCore.cpp (+79-6)
- (modified) mlir/python/mlir/ir.py (+38-1)
- (modified) mlir/test/python/ir/auto_location.py (+102)
``````````diff
diff --git a/mlir/include/mlir/Bindings/Python/Globals.h b/mlir/include/mlir/Bindings/Python/Globals.h
index b2aa169744c97..6282f37f885b3 100644
--- a/mlir/include/mlir/Bindings/Python/Globals.h
+++ b/mlir/include/mlir/Bindings/Python/Globals.h
@@ -129,6 +129,18 @@ class MLIR_PYTHON_API_EXPORTED PyGlobals {
class MLIR_PYTHON_API_EXPORTED TracebackLoc {
public:
+ /// Policy for handling explicit loc= when loc_tracebacks() is active.
+ enum class OnExplicitAction : uint8_t {
+ UseExplicit, // use loc= as base (default)
+ UseTraceback, // discard loc=, generate traceback
+ };
+
+ /// Policy for composing Location.current with the computed location.
+ enum class CurrentLocAction : uint8_t {
+ Fallback, // use Location.current only as fallback (default)
+ NamelocWrap, // extract NameLoc names, wrap computed loc
+ };
+
bool locTracebacksEnabled();
void setLocTracebacksEnabled(bool value);
@@ -143,12 +155,22 @@ class MLIR_PYTHON_API_EXPORTED PyGlobals {
bool isUserTracebackFilename(std::string_view file);
+ OnExplicitAction tracebackActionOnExplicitLoc();
+
+ void setTracebackActionOnExplicitLoc(OnExplicitAction action);
+
+ CurrentLocAction tracebackActionOnCurrentLoc();
+
+ void setTracebackActionOnCurrentLoc(CurrentLocAction action);
+
static constexpr size_t kMaxFrames = 512;
private:
nanobind::ft_mutex mutex;
bool locTracebackEnabled_ = false;
size_t locTracebackFramesLimit_ = 10;
+ OnExplicitAction onExplicitAction_ = OnExplicitAction::UseExplicit;
+ CurrentLocAction currentLocAction_ = CurrentLocAction::Fallback;
std::unordered_set<std::string> userTracebackIncludeFiles;
std::unordered_set<std::string> userTracebackExcludeFiles;
std::regex userTracebackIncludeRegex;
diff --git a/mlir/lib/Bindings/Python/Globals.cpp b/mlir/lib/Bindings/Python/Globals.cpp
index 1e48eac27dd83..dcc14367c42f1 100644
--- a/mlir/lib/Bindings/Python/Globals.cpp
+++ b/mlir/lib/Bindings/Python/Globals.cpp
@@ -285,6 +285,30 @@ void PyGlobals::TracebackLoc::setLocTracebackFramesLimit(size_t value) {
locTracebackFramesLimit_ = std::min(value, kMaxFrames);
}
+PyGlobals::TracebackLoc::OnExplicitAction
+PyGlobals::TracebackLoc::tracebackActionOnExplicitLoc() {
+ nanobind::ft_lock_guard lock(mutex);
+ return onExplicitAction_;
+}
+
+void PyGlobals::TracebackLoc::setTracebackActionOnExplicitLoc(
+ OnExplicitAction action) {
+ nanobind::ft_lock_guard lock(mutex);
+ onExplicitAction_ = action;
+}
+
+PyGlobals::TracebackLoc::CurrentLocAction
+PyGlobals::TracebackLoc::tracebackActionOnCurrentLoc() {
+ nanobind::ft_lock_guard lock(mutex);
+ return currentLocAction_;
+}
+
+void PyGlobals::TracebackLoc::setTracebackActionOnCurrentLoc(
+ CurrentLocAction action) {
+ nanobind::ft_lock_guard lock(mutex);
+ currentLocAction_ = action;
+}
+
void PyGlobals::TracebackLoc::registerTracebackFileInclusion(
const std::string &file) {
nanobind::ft_lock_guard lock(mutex);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index b52328a37e5f1..fe303c72dbbe0 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2734,17 +2734,68 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
#endif
}
+/// Apply currentLocAction: wrap or fuse Location.current onto baseLoc.
+static MlirLocation applyCurrentLocAction(
+ MlirContext ctx, MlirLocation baseLoc,
+ PyGlobals::TracebackLoc::CurrentLocAction action) {
+ using Action = PyGlobals::TracebackLoc::CurrentLocAction;
+ if (action == Action::Fallback)
+ return baseLoc;
+
+ auto *currentLoc = PyThreadContextEntry::getDefaultLocation();
+ if (!currentLoc)
+ return baseLoc;
+
+ // NamelocWrap: walk the NameLoc chain on Location.current, collect scope
+ // names, wrap baseLoc innermost-first so result is Outer(Inner(baseLoc)).
+ // If Location.current is not a NameLoc, scopeNames is empty and baseLoc
+ // is returned unchanged (nameloc_wrap is a no-op for non-NameLoc contexts).
+ thread_local std::vector<MlirStringRef> scopeNames;
+ scopeNames.clear();
+ MlirLocation walk = currentLoc->get();
+ while (mlirLocationIsAName(walk)) {
+ scopeNames.push_back(mlirIdentifierStr(mlirLocationNameGetName(walk)));
+ walk = mlirLocationNameGetChildLoc(walk);
+ }
+ for (auto it = scopeNames.rbegin(); it != scopeNames.rend(); ++it)
+ baseLoc = mlirLocationNameGet(ctx, *it, baseLoc);
+ return baseLoc;
+}
+
PyLocation
maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
- if (location.has_value())
- return location.value();
- if (!PyGlobals::get().getTracebackLoc().locTracebacksEnabled())
- return DefaultingPyLocation::resolve();
+ auto &tbl = PyGlobals::get().getTracebackLoc();
+
+ // Tracebacks not enabled — return explicit loc or fall back to Location.current.
+ if (!tbl.locTracebacksEnabled())
+ return location.has_value() ? location.value()
+ : DefaultingPyLocation::resolve();
+ // From here: tracebacks are enabled.
+ using OnExplicit = PyGlobals::TracebackLoc::OnExplicitAction;
PyMlirContext &ctx = DefaultingPyMlirContext::resolve();
- MlirLocation mlirLoc = tracebackToLocation(ctx.get());
+ MlirLocation baseLoc;
+
+ // Step 1: on_explicit — resolve explicit loc= vs traceback.
+ if (location.has_value()) {
+ switch (tbl.tracebackActionOnExplicitLoc()) {
+ case OnExplicit::UseExplicit:
+ baseLoc = location->get();
+ break;
+ case OnExplicit::UseTraceback:
+ baseLoc = tracebackToLocation(ctx.get());
+ break;
+ }
+ } else {
+ baseLoc = tracebackToLocation(ctx.get());
+ }
+
+ // Step 2: current_loc — compose with Location.current.
+ baseLoc = applyCurrentLocAction(ctx.get(), baseLoc,
+ tbl.tracebackActionOnCurrentLoc());
+
PyMlirContextRef ref = PyMlirContext::forContext(ctx.get());
- return {ref, mlirLoc};
+ return {ref, baseLoc};
}
} // namespace
@@ -2869,6 +2920,28 @@ void populateRoot(nb::module_ &m) {
.def("register_traceback_file_exclusion",
[](PyGlobals &self, const std::string &filename) {
self.getTracebackLoc().registerTracebackFileExclusion(filename);
+ })
+ .def("traceback_action_on_explicit_loc",
+ [](PyGlobals &self) {
+ return static_cast<int>(
+ self.getTracebackLoc().tracebackActionOnExplicitLoc());
+ })
+ .def("set_traceback_action_on_explicit_loc",
+ [](PyGlobals &self, int action) {
+ self.getTracebackLoc().setTracebackActionOnExplicitLoc(
+ static_cast<PyGlobals::TracebackLoc::OnExplicitAction>(
+ action));
+ })
+ .def("traceback_action_on_current_loc",
+ [](PyGlobals &self) {
+ return static_cast<int>(
+ self.getTracebackLoc().tracebackActionOnCurrentLoc());
+ })
+ .def("set_traceback_action_on_current_loc",
+ [](PyGlobals &self, int action) {
+ self.getTracebackLoc().setTracebackActionOnCurrentLoc(
+ static_cast<PyGlobals::TracebackLoc::CurrentLocAction>(
+ action));
});
// Aside from making the globals accessible to python, having python manage
diff --git a/mlir/python/mlir/ir.py b/mlir/python/mlir/ir.py
index c5b00a561831b..b62fe2c116257 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -70,8 +70,18 @@ def collect_ops(op: Operation):
return ops
+# Enum values matching C++ OnExplicitAction / CurrentLocAction.
+_ON_EXPLICIT_MAP = {"use_explicit": 0, "use_traceback": 1}
+_CURRENT_LOC_MAP = {"fallback": 0, "nameloc_wrap": 1}
+
+
@contextmanager
-def loc_tracebacks(*, max_depth: int | None = None) -> Generator[None, None, None]:
+def loc_tracebacks(
+ *,
+ max_depth: int | None = None,
+ on_explicit: str = "use_explicit",
+ current_loc: str = "fallback",
+) -> Generator[None, None, None]:
"""Enables automatic traceback-based locations for MLIR operations.
Operations created within this context will have their location
@@ -80,12 +90,37 @@ def loc_tracebacks(*, max_depth: int | None = None) -> Generator[None, None, Non
Args:
max_depth: Maximum number of frames to include in the location.
If None, the default limit is used.
+ on_explicit: Policy when an explicit loc= is passed to an op constructor.
+ "use_explicit" (default) — use loc= as base, skip traceback.
+ "use_traceback" — discard loc=, generate traceback instead.
+ current_loc: Policy for composing Location.current with the result.
+ "fallback" (default) — use Location.current only as fallback.
+ "nameloc_wrap" — extract NameLoc names from Location.current,
+ wrap the computed location with them.
"""
+ if on_explicit not in _ON_EXPLICIT_MAP:
+ raise ValueError(
+ f"on_explicit must be one of {list(_ON_EXPLICIT_MAP)}, "
+ f"got {on_explicit!r}"
+ )
+ if current_loc not in _CURRENT_LOC_MAP:
+ raise ValueError(
+ f"current_loc must be one of {list(_CURRENT_LOC_MAP)}, "
+ f"got {current_loc!r}"
+ )
old_enabled = _globals.loc_tracebacks_enabled()
old_limit = _globals.loc_tracebacks_frame_limit()
+ old_on_explicit = _globals.traceback_action_on_explicit_loc()
+ old_current_loc = _globals.traceback_action_on_current_loc()
max_depth = old_limit if max_depth is None else max_depth
try:
_globals.set_loc_tracebacks_frame_limit(max_depth)
+ _globals.set_traceback_action_on_explicit_loc(
+ _ON_EXPLICIT_MAP[on_explicit]
+ )
+ _globals.set_traceback_action_on_current_loc(
+ _CURRENT_LOC_MAP[current_loc]
+ )
if not old_enabled:
_globals.set_loc_tracebacks_enabled(True)
yield
@@ -93,6 +128,8 @@ def loc_tracebacks(*, max_depth: int | None = None) -> Generator[None, None, Non
if not old_enabled:
_globals.set_loc_tracebacks_enabled(False)
_globals.set_loc_tracebacks_frame_limit(old_limit)
+ _globals.set_traceback_action_on_explicit_loc(old_on_explicit)
+ _globals.set_traceback_action_on_current_loc(old_current_loc)
# Convenience decorator for registering user-friendly Attribute builders.
diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index ba0fa9b7e8e2e..140b06b9a736e 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -95,3 +95,105 @@ def bar3():
_cext.globals.set_loc_tracebacks_frame_limit(0)
# CHECK: loc(unknown)
bar1()
+
+
+# CHECK-LABEL: TEST: testNamelocWrap
+ at run
+def testNamelocWrap():
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ with loc_tracebacks(current_loc="nameloc_wrap"):
+ # Build nested NameLoc: TaskA(ResourceB(unknown))
+ inner = Location.name("ResourceB", childLoc=Location.unknown())
+ outer = Location.name("TaskA", childLoc=inner)
+ with outer:
+ op = Operation.create("custom.op1")
+ # fmt: off
+ # CHECK: loc("TaskA"("ResourceB"(callsite(
+ # fmt: on
+ print(op.location)
+
+
+# CHECK-LABEL: TEST: testOnExplicitDefault
+ at run
+def testOnExplicitDefault():
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ with loc_tracebacks():
+ explicit = Location.file("explicit.py", 1, 1)
+ op = Operation.create("custom.op1", loc=explicit)
+ # CHECK: loc("explicit.py":1:1)
+ print(op.location)
+
+
+# CHECK-LABEL: TEST: testOnExplicitUseTraceback
+ at run
+def testOnExplicitUseTraceback():
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ with loc_tracebacks(on_explicit="use_traceback"):
+ explicit = Location.file("explicit.py", 1, 1)
+ op = Operation.create("custom.op1", loc=explicit)
+ # fmt: off
+ # CHECK: loc(callsite(
+ # fmt: on
+ print(op.location)
+
+
+# CHECK-LABEL: TEST: testDslProfilingUseCase
+ at run
+def testDslProfilingUseCase():
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ with loc_tracebacks(
+ current_loc="nameloc_wrap", on_explicit="use_traceback"
+ ):
+ task_loc = Location.name("Task", childLoc=Location.unknown())
+ with task_loc:
+ op1 = Operation.create(
+ "custom.op1", loc=Location.file("x.py", 1, 1)
+ )
+ # fmt: off
+ # CHECK: loc("Task"(callsite(
+ # fmt: on
+ print(op1.location)
+
+ op2 = Operation.create("custom.op2")
+ # fmt: off
+ # CHECK: loc("Task"(callsite(
+ # fmt: on
+ print(op2.location)
+
+
+# CHECK-LABEL: TEST: testOnExplicitFallbackWhenTracebacksDisabled
+ at run
+def testOnExplicitFallbackWhenTracebacksDisabled():
+ """When on_explicit != use_explicit but tracebacks are disabled,
+ the explicit loc should be returned (not Location.current)."""
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ # Set on_explicit to use_traceback WITHOUT enabling tracebacks.
+ _cext.globals.set_traceback_action_on_explicit_loc(1) # UseTraceback
+ try:
+ explicit = Location.file("keep_me.py", 42, 1)
+ op = Operation.create("custom.op1", loc=explicit)
+ # CHECK: loc("keep_me.py":42:1)
+ print(op.location)
+ finally:
+ _cext.globals.set_traceback_action_on_explicit_loc(0) # UseExplicit
+
+
+# CHECK-LABEL: TEST: testUseExplicitWithNamelocWrap
+ at run
+def testUseExplicitWithNamelocWrap():
+ """on_explicit=use_explicit + current_loc=nameloc_wrap should wrap
+ the explicit loc with the NameLoc chain (flags are orthogonal)."""
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ with loc_tracebacks(on_explicit="use_explicit", current_loc="nameloc_wrap"):
+ task_loc = Location.name("Task", childLoc=Location.unknown())
+ with task_loc:
+ explicit = Location.file("framework.py", 10, 1)
+ op = Operation.create("custom.op1", loc=explicit)
+ # CHECK: loc("Task"("framework.py":10:1))
+ print(op.location)
``````````
</details>
https://github.com/llvm/llvm-project/pull/192310
More information about the Mlir-commits
mailing list