| # 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) |
