102 lines · python
1# RUN: %PYTHON %s | FileCheck %s2# REQUIRES: python-ge-3113import gc4from contextlib import contextmanager5 6from mlir.ir import *7from mlir.dialects._ods_common import _cext8from mlir.dialects import arith, _arith_ops_gen9 10 11def run(f):12 print("\nTEST:", f.__name__)13 f()14 gc.collect()15 assert Context._get_live_count() == 016 17 18@contextmanager19def with_infer_location():20 _cext.globals.set_loc_tracebacks_enabled(True)21 yield22 _cext.globals.set_loc_tracebacks_enabled(False)23 24 25# CHECK-LABEL: TEST: testInferLocations26@run27def testInferLocations():28 with Context() as ctx, with_infer_location():29 ctx.allow_unregistered_dialects = True30 31 op = Operation.create("custom.op1")32 one = arith.constant(IndexType.get(), 1)33 _cext.globals.register_traceback_file_exclusion(arith.__file__)34 two = arith.constant(IndexType.get(), 2)35 36 # fmt: off37 # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP:[/\\]+]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":31:13 to :43) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))38 # fmt: on39 print(op.location)40 41 # fmt: off42 # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":65:12 to :76) at callsite("constant"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":110:40 to :81) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":32:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))))43 # fmt: on44 print(one.location)45 46 # fmt: off47 # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":34:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))48 # fmt: on49 print(two.location)50 51 _cext.globals.register_traceback_file_inclusion(_arith_ops_gen.__file__)52 three = arith.constant(IndexType.get(), 3)53 # fmt: off54 # 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":52:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))55 # fmt: on56 print(three.location)57 58 def foo():59 four = arith.constant(IndexType.get(), 4)60 print(four.location)61 62 # fmt: off63 # 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":59:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":65:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))))64 # fmt: on65 foo()66 67 _cext.globals.register_traceback_file_exclusion(__file__)68 69 # fmt: off70 # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235))71 # fmt: on72 foo()73 74 def bar1():75 def bar2():76 def bar3():77 five = arith.constant(IndexType.get(), 5)78 print(five.location)79 80 bar3()81 82 bar2()83 84 _cext.globals.register_traceback_file_inclusion(__file__)85 _cext.globals.register_traceback_file_exclusion(_arith_ops_gen.__file__)86 87 _cext.globals.set_loc_tracebacks_frame_limit(2)88 # fmt: off89 # CHECK: loc(callsite("testInferLocations.<locals>.bar1.<locals>.bar2.<locals>.bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":77:27 to :61) at "testInferLocations.<locals>.bar1.<locals>.bar2"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":80:16 to :22)))90 # fmt: on91 bar1()92 93 _cext.globals.set_loc_tracebacks_frame_limit(1)94 # fmt: off95 # CHECK: loc("testInferLocations.<locals>.bar1.<locals>.bar2.<locals>.bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":77:27 to :61))96 # fmt: on97 bar1()98 99 _cext.globals.set_loc_tracebacks_frame_limit(0)100 # CHECK: loc(unknown)101 bar1()102