Doomhowl-Interactive/llvm-project
0
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity
Files
main
llvm-project/mlir/test/python/ir/auto_location.py
Alisson Azzolini d8c14330dc [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>
2026-04-19 10:07:56 +02:00

206 lines
9.0 KiB
Python
Raw Permalink Blame History

# RUN: %PYTHON %s | FileCheck %s
# REQUIRES: python-ge-311
# UNSUPPORTED: python-stable-abi
import gc
from contextlib import contextmanager
from mlir.ir import *
from mlir.dialects._ods_common import _cext
from mlir.dialects import arith, _arith_ops_gen
def run(f):
print("\nTEST:", f.__name__)
f()
gc.collect()
assert Context._get_live_count() == 0
# CHECK-LABEL: TEST: testInferLocations
@run
def testInferLocations():
with Context() as ctx, loc_tracebacks():
ctx.allow_unregistered_dialects = True
op = Operation.create("custom.op1")
one = arith.constant(IndexType.get(), 1)
_cext.globals.register_traceback_file_exclusion(arith.__file__)
two = arith.constant(IndexType.get(), 2)
# fmt: off
# CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP:[/\\]+]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:13 to :43) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":14:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))
# fmt: on
print(op.location)
# Test nesting of loc_tracebacks().
with loc_tracebacks():
# fmt: off
# CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":45:12 to :76) at callsite("constant"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":90:40 to :81) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))))
# fmt: on
print(one.location)
# fmt: off
# CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))
# fmt: on
print(two.location)
_cext.globals.register_traceback_file_inclusion(_arith_ops_gen.__file__)
three = arith.constant(IndexType.get(), 3)
# fmt: off
# CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4)))))
# fmt: on
print(three.location)
def foo():
four = arith.constant(IndexType.get(), 4)
print(four.location)
# fmt: off
# CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235) at callsite("testInferLocations.<locals>.foo"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))))
# fmt: on
foo()
_cext.globals.register_traceback_file_exclusion(__file__)
# fmt: off
# CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235))
# fmt: on
foo()
def bar1():
def bar2():
def bar3():
five = arith.constant(IndexType.get(), 5)
print(five.location)
bar3()
bar2()
_cext.globals.register_traceback_file_inclusion(__file__)
_cext.globals.register_traceback_file_exclusion(_arith_ops_gen.__file__)
_cext.globals.set_loc_tracebacks_frame_limit(2)
# fmt: off
# CHECK: loc(callsite("testInferLocations.<locals>.bar1.<locals>.bar2.<locals>.bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:27 to :61) at "testInferLocations.<locals>.bar1.<locals>.bar2"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:16 to :22)))
# fmt: on
bar1()
_cext.globals.set_loc_tracebacks_frame_limit(1)
# fmt: off
# CHECK: loc("testInferLocations.<locals>.bar1.<locals>.bar2.<locals>.bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:27 to :61))
# fmt: on
bar1()
_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)
Reference in New Issue View Git Blame Copy Permalink