[Mlir-commits] [mlir] d8c1433 - [mlir][python] Add location source composition to loc_tracebacks() (#192310)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Apr 19 01:08:01 PDT 2026
Author: Alisson Azzolini
Date: 2026-04-19T10:07:56+02:00
New Revision: d8c14330dcc17610f634b45d1730b7a6e46d91cd
URL: https://github.com/llvm/llvm-project/commit/d8c14330dcc17610f634b45d1730b7a6e46d91cd
DIFF: https://github.com/llvm/llvm-project/commit/d8c14330dcc17610f634b45d1730b7a6e46d91cd.diff
LOG: [mlir][python] Add location source composition to loc_tracebacks() (#192310)
## Summary
Add two new kwargs to `loc_tracebacks()` controlling how the three
location sources (explicit `loc=`, traceback, `Location.current`)
compose:
- **`on_explicit`**: `OnExplicitAction.USE_EXPLICIT` (default) |
`OnExplicitAction.USE_TRACEBACK` — what to do when explicit `loc=` is
passed
- **`current_loc`**: `CurrentLocAction.FALLBACK` (default) |
`CurrentLocAction.NAMELOC_WRAP` — how to compose `Location.current`
NameLoc scopes on top
Both enums are exposed to Python via `nb::enum_<>` bindings, following
the `PassDisplayMode` pattern in `Pass.cpp`.
The two flags are orthogonal and only take effect when
`loc_tracebacks()` is active.
## Use cases
1. **DSL profiling attribution**:
`loc_tracebacks(current_loc=CurrentLocAction.NAMELOC_WRAP,
on_explicit=OnExplicitAction.USE_TRACEBACK)` — push NameLoc scopes for
profiling and have them survive traceback generation
2. **Override explicit locs with tracebacks**:
`loc_tracebacks(on_explicit=OnExplicitAction.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, on_explicit default,
on_explicit=use_traceback, DSL profiling combo, orthogonality
(use_explicit + nameloc_wrap), tracebacks-disabled fallback
Co-authored by Claude Code ;)
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/include/mlir/Bindings/Python/Globals.h
mlir/lib/Bindings/Python/Globals.cpp
mlir/lib/Bindings/Python/IRCore.cpp
mlir/python/mlir/ir.py
mlir/test/python/ir/auto_location.py
Removed:
################################################################################
diff --git a/mlir/include/mlir/Bindings/Python/Globals.h b/mlir/include/mlir/Bindings/Python/Globals.h
index b2aa169744c97..87f6357939285 100644
--- a/mlir/include/mlir/Bindings/Python/Globals.h
+++ b/mlir/include/mlir/Bindings/Python/Globals.h
@@ -129,6 +129,20 @@ 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.
+ /// TODO: possibly add CallSiteLoc wrap and a generic Fuse option
+ /// (`fused[Location.current, baseLoc]`) for non-NameLoc cases.
+ 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 +157,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..73e08c6a8b1c1 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..0fa84a11f2f35 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -20,6 +20,7 @@
#include "mlir-c/Support.h"
#include <array>
+#include <cassert>
#include <functional>
#include <optional>
#include <string>
@@ -2734,17 +2735,71 @@ 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;
+ assert(mlirLocationGetContext(currentLoc->get()).ptr == ctx.ptr &&
+ "Location.current must belong to the current MLIR context");
+
+ // 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
@@ -2825,6 +2880,19 @@ void populateRoot(nb::module_ &m) {
m.attr("T") = nb::type_var("T");
m.attr("U") = nb::type_var("U");
+ // Policies for how loc_tracebacks() composes the three location sources
+ // (explicit loc=, generated traceback, Location.current).
+ nb::enum_<PyGlobals::TracebackLoc::OnExplicitAction>(m, "OnExplicitAction")
+ .value("USE_EXPLICIT",
+ PyGlobals::TracebackLoc::OnExplicitAction::UseExplicit)
+ .value("USE_TRACEBACK",
+ PyGlobals::TracebackLoc::OnExplicitAction::UseTraceback);
+
+ nb::enum_<PyGlobals::TracebackLoc::CurrentLocAction>(m, "CurrentLocAction")
+ .value("FALLBACK", PyGlobals::TracebackLoc::CurrentLocAction::Fallback)
+ .value("NAMELOC_WRAP",
+ PyGlobals::TracebackLoc::CurrentLocAction::NamelocWrap);
+
nb::class_<PyGlobals>(m, "_Globals")
.def_prop_rw("dialect_search_modules",
&PyGlobals::getDialectSearchPrefixes,
@@ -2869,6 +2937,24 @@ 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 self.getTracebackLoc().tracebackActionOnExplicitLoc();
+ })
+ .def("set_traceback_action_on_explicit_loc",
+ [](PyGlobals &self,
+ PyGlobals::TracebackLoc::OnExplicitAction action) {
+ self.getTracebackLoc().setTracebackActionOnExplicitLoc(action);
+ })
+ .def("traceback_action_on_current_loc",
+ [](PyGlobals &self) {
+ return self.getTracebackLoc().tracebackActionOnCurrentLoc();
+ })
+ .def("set_traceback_action_on_current_loc",
+ [](PyGlobals &self,
+ PyGlobals::TracebackLoc::CurrentLocAction action) {
+ self.getTracebackLoc().setTracebackActionOnCurrentLoc(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..011ecfffd355f 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -14,6 +14,8 @@
register_type_caster,
register_value_caster,
globals as _globals,
+ OnExplicitAction,
+ CurrentLocAction,
)
from ._mlir_libs import (
get_dialect_registry,
@@ -71,7 +73,12 @@ def collect_ops(op: Operation):
@contextmanager
-def loc_tracebacks(*, max_depth: int | None = None) -> Generator[None, None, None]:
+def loc_tracebacks(
+ *,
+ max_depth: int | None = None,
+ on_explicit_actn: OnExplicitAction = OnExplicitAction.USE_EXPLICIT,
+ current_loc_actn: CurrentLocAction = CurrentLocAction.FALLBACK,
+) -> Generator[None, None, None]:
"""Enables automatic traceback-based locations for MLIR operations.
Operations created within this context will have their location
@@ -80,12 +87,26 @@ 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_actn: Policy when an explicit loc= is passed to an op
+ constructor.
+ OnExplicitAction.USE_EXPLICIT (default) — use loc= as base, skip
+ traceback.
+ OnExplicitAction.USE_TRACEBACK — discard loc=, generate traceback.
+ current_loc_actn: Policy for composing Location.current with the result.
+ CurrentLocAction.FALLBACK (default) — use Location.current only as
+ fallback.
+ CurrentLocAction.NAMELOC_WRAP — extract NameLoc names from
+ Location.current and wrap the computed location with them.
"""
old_enabled = _globals.loc_tracebacks_enabled()
old_limit = _globals.loc_tracebacks_frame_limit()
+ old_on_explicit_actn = _globals.traceback_action_on_explicit_loc()
+ old_current_loc_actn = _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_actn)
+ _globals.set_traceback_action_on_current_loc(current_loc_actn)
if not old_enabled:
_globals.set_loc_tracebacks_enabled(True)
yield
@@ -93,6 +114,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_actn)
+ _globals.set_traceback_action_on_current_loc(old_current_loc_actn)
# 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..3af20739b0c46 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -95,3 +95,111 @@ 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_actn=CurrentLocAction.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_actn=OnExplicitAction.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_actn=CurrentLocAction.NAMELOC_WRAP,
+ on_explicit_actn=OnExplicitAction.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(
+ OnExplicitAction.USE_TRACEBACK
+ )
+ 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(
+ OnExplicitAction.USE_EXPLICIT
+ )
+
+
+# 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_actn=OnExplicitAction.USE_EXPLICIT,
+ current_loc_actn=CurrentLocAction.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)
More information about the Mlir-commits
mailing list