[Mlir-commits] [mlir] [mlir][python] Add location source composition to loc_tracebacks() (PR #192310)
Alisson Azzolini
llvmlistbot at llvm.org
Fri Apr 17 09:56:06 PDT 2026
https://github.com/aazzolini updated https://github.com/llvm/llvm-project/pull/192310
>From c61354acef28d000f6ca7aa7e22d93caef4c988d Mon Sep 17 00:00:00 2001
From: Alisson Gusatti Azzolini <aazzolini at nvidia.com>
Date: Tue, 14 Apr 2026 16:45:20 -0700
Subject: [PATCH 1/4] [mlir][python] Add location source composition to
loc_tracebacks()
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"
Controls what happens when an explicit loc= is passed to an op
constructor while loc_tracebacks() is active.
- current_loc: "fallback" (default) | "nameloc_wrap"
Controls how Location.current (the thread-local context from
Location.name() context managers) composes with the computed loc.
"nameloc_wrap" extracts NameLoc names and wraps the result.
The two flags are orthogonal: on_explicit resolves the base location,
then current_loc wraps on top. Both only take effect when
loc_tracebacks() is active (kill switch).
This enables DSL frameworks to:
1. Push NameLoc scopes (via Location.name()) for profiling attribution
and have them survive traceback generation (current_loc="nameloc_wrap")
2. Override pre-existing explicit loc= arguments with full tracebacks
without modifying every callsite (on_explicit="use_traceback")
Backward compatible: default values match existing behavior exactly.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply at anthropic.com>
---
mlir/include/mlir/Bindings/Python/Globals.h | 22 +++++
mlir/lib/Bindings/Python/Globals.cpp | 24 +++++
mlir/lib/Bindings/Python/IRCore.cpp | 85 ++++++++++++++--
mlir/python/mlir/ir.py | 39 +++++++-
mlir/test/python/ir/auto_location.py | 102 ++++++++++++++++++++
5 files changed, 265 insertions(+), 7 deletions(-)
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)
>From 0682b805f96831b33d464d97c2fa2ac0c9f1d51e Mon Sep 17 00:00:00 2001
From: Alisson Gusatti Azzolini <aazzolini at nvidia.com>
Date: Thu, 16 Apr 2026 14:54:51 -0700
Subject: [PATCH 2/4] [mlir][python] Expose OnExplicitAction/CurrentLocAction
as nanobind enums
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Address review feedback on #192310: bind the two policy enums directly
via nanobind (following the PassDisplayMode pattern in Pass.cpp) rather
than exposing them as opaque ints wrapped in Python-side string→int maps.
- IRCore.cpp: `nb::enum_<>` bindings for OnExplicitAction and
CurrentLocAction, exposed as `mlir.OnExplicitAction` and
`mlir.CurrentLocAction`. The getter/setter lambdas now take the typed
enum directly — no more int round-trips.
- ir.py: `loc_tracebacks()` accepts either the typed enum or the
convenience string alias (`"use_explicit"`, `"nameloc_wrap"`, etc.).
The maps now hold enum values; the int magic numbers are gone.
Also apply clang-format and black fixups flagged by upstream CI.
---
mlir/lib/Bindings/Python/IRCore.cpp | 42 +++++++++++-------
mlir/python/mlir/ir.py | 65 +++++++++++++++++-----------
mlir/test/python/ir/auto_location.py | 8 +---
3 files changed, 67 insertions(+), 48 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index fe303c72dbbe0..8c27e2d832fcb 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2735,9 +2735,9 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
}
/// Apply currentLocAction: wrap or fuse Location.current onto baseLoc.
-static MlirLocation applyCurrentLocAction(
- MlirContext ctx, MlirLocation baseLoc,
- PyGlobals::TracebackLoc::CurrentLocAction action) {
+static MlirLocation
+applyCurrentLocAction(MlirContext ctx, MlirLocation baseLoc,
+ PyGlobals::TracebackLoc::CurrentLocAction action) {
using Action = PyGlobals::TracebackLoc::CurrentLocAction;
if (action == Action::Fallback)
return baseLoc;
@@ -2766,7 +2766,8 @@ PyLocation
maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
auto &tbl = PyGlobals::get().getTracebackLoc();
- // Tracebacks not enabled — return explicit loc or fall back to Location.current.
+ // Tracebacks not enabled — return explicit loc or fall back to
+ // Location.current.
if (!tbl.locTracebacksEnabled())
return location.has_value() ? location.value()
: DefaultingPyLocation::resolve();
@@ -2876,6 +2877,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,
@@ -2923,25 +2937,21 @@ void populateRoot(nb::module_ &m) {
})
.def("traceback_action_on_explicit_loc",
[](PyGlobals &self) {
- return static_cast<int>(
- self.getTracebackLoc().tracebackActionOnExplicitLoc());
+ return self.getTracebackLoc().tracebackActionOnExplicitLoc();
})
.def("set_traceback_action_on_explicit_loc",
- [](PyGlobals &self, int action) {
- self.getTracebackLoc().setTracebackActionOnExplicitLoc(
- static_cast<PyGlobals::TracebackLoc::OnExplicitAction>(
- action));
+ [](PyGlobals &self,
+ PyGlobals::TracebackLoc::OnExplicitAction action) {
+ self.getTracebackLoc().setTracebackActionOnExplicitLoc(action);
})
.def("traceback_action_on_current_loc",
[](PyGlobals &self) {
- return static_cast<int>(
- self.getTracebackLoc().tracebackActionOnCurrentLoc());
+ return self.getTracebackLoc().tracebackActionOnCurrentLoc();
})
.def("set_traceback_action_on_current_loc",
- [](PyGlobals &self, int action) {
- self.getTracebackLoc().setTracebackActionOnCurrentLoc(
- static_cast<PyGlobals::TracebackLoc::CurrentLocAction>(
- action));
+ [](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 b62fe2c116257..846885f3b588f 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,
@@ -70,17 +72,24 @@ 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}
+# String aliases for the typed enum values exposed from C++. Accepting
+# strings is convenient for callers; internally we always hold the enum.
+_ON_EXPLICIT_MAP = {
+ "use_explicit": OnExplicitAction.USE_EXPLICIT,
+ "use_traceback": OnExplicitAction.USE_TRACEBACK,
+}
+_CURRENT_LOC_MAP = {
+ "fallback": CurrentLocAction.FALLBACK,
+ "nameloc_wrap": CurrentLocAction.NAMELOC_WRAP,
+}
@contextmanager
def loc_tracebacks(
*,
max_depth: int | None = None,
- on_explicit: str = "use_explicit",
- current_loc: str = "fallback",
+ on_explicit: str | OnExplicitAction = "use_explicit",
+ current_loc: str | CurrentLocAction = "fallback",
) -> Generator[None, None, None]:
"""Enables automatic traceback-based locations for MLIR operations.
@@ -91,23 +100,31 @@ def loc_tracebacks(
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.
+ OnExplicitAction.USE_EXPLICIT / "use_explicit" (default) — use loc=
+ as base, skip traceback.
+ OnExplicitAction.USE_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.
+ CurrentLocAction.FALLBACK / "fallback" (default) — use
+ Location.current only as fallback.
+ CurrentLocAction.NAMELOC_WRAP / "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}"
- )
+ if isinstance(on_explicit, str):
+ 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}"
+ )
+ on_explicit = _ON_EXPLICIT_MAP[on_explicit]
+ if isinstance(current_loc, str):
+ 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}"
+ )
+ current_loc = _CURRENT_LOC_MAP[current_loc]
+
old_enabled = _globals.loc_tracebacks_enabled()
old_limit = _globals.loc_tracebacks_frame_limit()
old_on_explicit = _globals.traceback_action_on_explicit_loc()
@@ -115,12 +132,8 @@ def loc_tracebacks(
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]
- )
+ _globals.set_traceback_action_on_explicit_loc(on_explicit)
+ _globals.set_traceback_action_on_current_loc(current_loc)
if not old_enabled:
_globals.set_loc_tracebacks_enabled(True)
yield
diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 140b06b9a736e..04a45269a5f61 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -145,14 +145,10 @@ def testOnExplicitUseTraceback():
def testDslProfilingUseCase():
with Context() as ctx, Location.unknown():
ctx.allow_unregistered_dialects = True
- with loc_tracebacks(
- current_loc="nameloc_wrap", on_explicit="use_traceback"
- ):
+ 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)
- )
+ op1 = Operation.create("custom.op1", loc=Location.file("x.py", 1, 1))
# fmt: off
# CHECK: loc("Task"(callsite(
# fmt: on
>From 2c09fb539443b96be506829ff1ca123829c28891 Mon Sep 17 00:00:00 2001
From: Alisson Gusatti Azzolini <aazzolini at nvidia.com>
Date: Thu, 16 Apr 2026 15:04:32 -0700
Subject: [PATCH 3/4] Drop string aliases; loc_tracebacks() takes the typed
enums directly
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Follow-up on the nb::enum_ refactor: now that OnExplicitAction and
CurrentLocAction are exposed to Python as typed enums, there's no
reason to keep the parallel string API — it duplicates the surface
area for no real benefit and doesn't match the upstream style
(cf. PassDisplayMode, which has no string shim).
- ir.py: drop _ON_EXPLICIT_MAP / _CURRENT_LOC_MAP, the isinstance(str)
branches, and the `str | OnExplicitAction` union types. The kwargs
now take the enum directly; defaults use the enum values.
- auto_location.py: update call sites to pass the typed enum. Also
replace the `set_*(0/1)` magic numbers in the
tracebacks-disabled edge-case test with enum values.
---
mlir/python/mlir/ir.py | 46 ++++++----------------------
mlir/test/python/ir/auto_location.py | 24 ++++++++++-----
2 files changed, 26 insertions(+), 44 deletions(-)
diff --git a/mlir/python/mlir/ir.py b/mlir/python/mlir/ir.py
index 846885f3b588f..f98cff8aed3a8 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -72,24 +72,12 @@ def collect_ops(op: Operation):
return ops
-# String aliases for the typed enum values exposed from C++. Accepting
-# strings is convenient for callers; internally we always hold the enum.
-_ON_EXPLICIT_MAP = {
- "use_explicit": OnExplicitAction.USE_EXPLICIT,
- "use_traceback": OnExplicitAction.USE_TRACEBACK,
-}
-_CURRENT_LOC_MAP = {
- "fallback": CurrentLocAction.FALLBACK,
- "nameloc_wrap": CurrentLocAction.NAMELOC_WRAP,
-}
-
-
@contextmanager
def loc_tracebacks(
*,
max_depth: int | None = None,
- on_explicit: str | OnExplicitAction = "use_explicit",
- current_loc: str | CurrentLocAction = "fallback",
+ on_explicit: OnExplicitAction = OnExplicitAction.USE_EXPLICIT,
+ current_loc: CurrentLocAction = CurrentLocAction.FALLBACK,
) -> Generator[None, None, None]:
"""Enables automatic traceback-based locations for MLIR operations.
@@ -100,31 +88,15 @@ def loc_tracebacks(
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.
- OnExplicitAction.USE_EXPLICIT / "use_explicit" (default) — use loc=
- as base, skip traceback.
- OnExplicitAction.USE_TRACEBACK / "use_traceback" — discard loc=,
- generate traceback instead.
+ OnExplicitAction.USE_EXPLICIT (default) — use loc= as base, skip
+ traceback.
+ OnExplicitAction.USE_TRACEBACK — discard loc=, generate traceback.
current_loc: Policy for composing Location.current with the result.
- CurrentLocAction.FALLBACK / "fallback" (default) — use
- Location.current only as fallback.
- CurrentLocAction.NAMELOC_WRAP / "nameloc_wrap" — extract NameLoc
- names from Location.current, wrap the computed location with them.
+ 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.
"""
- if isinstance(on_explicit, str):
- 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}"
- )
- on_explicit = _ON_EXPLICIT_MAP[on_explicit]
- if isinstance(current_loc, str):
- 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}"
- )
- current_loc = _CURRENT_LOC_MAP[current_loc]
-
old_enabled = _globals.loc_tracebacks_enabled()
old_limit = _globals.loc_tracebacks_frame_limit()
old_on_explicit = _globals.traceback_action_on_explicit_loc()
diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 04a45269a5f61..fe199cda07feb 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -102,7 +102,7 @@ def bar3():
def testNamelocWrap():
with Context() as ctx, Location.unknown():
ctx.allow_unregistered_dialects = True
- with loc_tracebacks(current_loc="nameloc_wrap"):
+ with loc_tracebacks(current_loc=CurrentLocAction.NAMELOC_WRAP):
# Build nested NameLoc: TaskA(ResourceB(unknown))
inner = Location.name("ResourceB", childLoc=Location.unknown())
outer = Location.name("TaskA", childLoc=inner)
@@ -131,7 +131,7 @@ def testOnExplicitDefault():
def testOnExplicitUseTraceback():
with Context() as ctx, Location.unknown():
ctx.allow_unregistered_dialects = True
- with loc_tracebacks(on_explicit="use_traceback"):
+ with loc_tracebacks(on_explicit=OnExplicitAction.USE_TRACEBACK):
explicit = Location.file("explicit.py", 1, 1)
op = Operation.create("custom.op1", loc=explicit)
# fmt: off
@@ -145,7 +145,10 @@ def testOnExplicitUseTraceback():
def testDslProfilingUseCase():
with Context() as ctx, Location.unknown():
ctx.allow_unregistered_dialects = True
- with loc_tracebacks(current_loc="nameloc_wrap", on_explicit="use_traceback"):
+ with loc_tracebacks(
+ current_loc=CurrentLocAction.NAMELOC_WRAP,
+ on_explicit=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))
@@ -168,15 +171,19 @@ def testOnExplicitFallbackWhenTracebacksDisabled():
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
+ # 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(0) # UseExplicit
+ _cext.globals.set_traceback_action_on_explicit_loc(
+ OnExplicitAction.USE_EXPLICIT
+ )
# CHECK-LABEL: TEST: testUseExplicitWithNamelocWrap
@@ -186,7 +193,10 @@ def testUseExplicitWithNamelocWrap():
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"):
+ with loc_tracebacks(
+ on_explicit=OnExplicitAction.USE_EXPLICIT,
+ current_loc=CurrentLocAction.NAMELOC_WRAP,
+ ):
task_loc = Location.name("Task", childLoc=Location.unknown())
with task_loc:
explicit = Location.file("framework.py", 10, 1)
>From 29f577839bdc26434d37e7ec97ec4c3228b99745 Mon Sep 17 00:00:00 2001
From: Alisson Gusatti Azzolini <aazzolini at nvidia.com>
Date: Fri, 17 Apr 2026 09:55:43 -0700
Subject: [PATCH 4/4] Assert Location.current belongs to the current MLIR
context
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Per @grypp review comment on #192310. Catches the case where Python
code pushes a Location from context A while Context.current is ctx B —
reading identifier strings out of a mismatched-context location is
undefined (identifier storage is owned by its context). The assert
fires at the cheapest point instead of letting the misuse manifest
later as a cryptic lifetime bug.
---
mlir/lib/Bindings/Python/IRCore.cpp | 3 +++
1 file changed, 3 insertions(+)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8c27e2d832fcb..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>
@@ -2745,6 +2746,8 @@ applyCurrentLocAction(MlirContext ctx, MlirLocation 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)).
More information about the Mlir-commits
mailing list