a40f47c972
This PR implements "automatic" location inference in the bindings. The
way it works is it walks the frame stack collecting source locations
(Python captures these in the frame itself). It is inspired by JAX's
[implementation](523ddcfbca/jax/_src/interpreters/mlir.py (L462))
but moves the frame stack traversal into the bindings for better
performance.
The system supports registering "included" and "excluded" filenames;
frames originating from functions in included filenames **will not** be
filtered and frames originating from functions in excluded filenames
**will** be filtered (in that order). This allows excluding all the
generated `*_ops_gen.py` files.
The system is also "toggleable" and off by default to save people who
have their own systems (such as JAX) from the added cost.
Note, the system stores the entire stacktrace (subject to
`locTracebackFramesLimit`) in the `Location` using specifically a
`CallSiteLoc`. This can be useful for profiling tools (flamegraphs
etc.).
Shoutout to the folks at JAX for coming up with a good system.
---------
Co-authored-by: Jacques Pienaar <jpienaar@google.com>
91 lines
2.3 KiB
Python
91 lines
2.3 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
import gc
|
|
from mlir.ir import *
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
gc.collect()
|
|
assert Context._get_live_count() == 0
|
|
|
|
|
|
# CHECK-LABEL: TEST: testContextEnterExit
|
|
def testContextEnterExit():
|
|
with Context() as ctx:
|
|
assert Context.current is ctx
|
|
assert Context.current is None
|
|
|
|
|
|
run(testContextEnterExit)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testLocationEnterExit
|
|
def testLocationEnterExit():
|
|
ctx1 = Context()
|
|
with Location.unknown(ctx1) as loc1:
|
|
assert Context.current is ctx1
|
|
assert Location.current is loc1
|
|
|
|
# Re-asserting the same context should not change the location.
|
|
with ctx1:
|
|
assert Context.current is ctx1
|
|
assert Location.current is loc1
|
|
# Asserting a different context should clear it.
|
|
with Context() as ctx2:
|
|
assert Context.current is ctx2
|
|
assert Location.current is None
|
|
|
|
# And should restore.
|
|
assert Context.current is ctx1
|
|
assert Location.current is loc1
|
|
|
|
# All should clear.
|
|
assert Location.current is None
|
|
|
|
|
|
run(testLocationEnterExit)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testInsertionPointEnterExit
|
|
def testInsertionPointEnterExit():
|
|
ctx1 = Context()
|
|
m = Module.create(Location.unknown(ctx1))
|
|
ip = InsertionPoint(m.body)
|
|
|
|
with ip:
|
|
assert InsertionPoint.current is ip
|
|
# Asserting a location from the same context should preserve.
|
|
with Location.unknown(ctx1) as loc1:
|
|
assert InsertionPoint.current is ip
|
|
assert Location.current is loc1
|
|
# Location should clear.
|
|
assert Location.current is None
|
|
|
|
# Asserting the same Context should preserve.
|
|
with ctx1:
|
|
assert InsertionPoint.current is ip
|
|
|
|
# Asserting a different context should clear it.
|
|
with Context() as ctx2:
|
|
assert Context.current is ctx2
|
|
try:
|
|
_ = InsertionPoint.current
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
assert False, "Expected exception"
|
|
|
|
# All should clear.
|
|
try:
|
|
_ = InsertionPoint.current
|
|
except ValueError as e:
|
|
# CHECK: No current InsertionPoint
|
|
print(e)
|
|
else:
|
|
assert False, "Expected exception"
|
|
|
|
|
|
run(testInsertionPointEnterExit)
|