[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@anthropic.com>
This commit is contained in:
Alisson Azzolini
GitHub
parent
2a3667bbb5
commit
d8c14330dc
5 changed files with 272 additions and 7 deletions
|
|
@ -129,6 +129,20 @@ public:
|
|||
|
||||
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 @@ public:
|
|||
|
||||
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;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ from ._mlir_libs._mlir import (
|
|||
register_type_caster,
|
||||
register_value_caster,
|
||||
globals as _globals,
|
||||
OnExplicitAction,
|
||||
CurrentLocAction,
|
||||
)
|
||||
from ._mlir_libs import (
|
||||
get_dialect_registry,
|
||||
|
|
@ -71,7 +73,12 @@ def get_ops_of_type(
|
|||
|
||||
|
||||
@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.
|
||||
|
|
|
|||
|
|
@ -95,3 +95,111 @@ def testInferLocations():
|
|||
_cext.globals.set_loc_tracebacks_frame_limit(0)
|
||||
# CHECK: loc(unknown)
|
||||
bar1()
|
||||
|
||||
|
||||
# CHECK-LABEL: TEST: testNamelocWrap
|
||||
@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
|
||||
@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
|
||||
@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
|
||||
@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
|
||||
@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
|
||||
@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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue