[Mlir-commits] [mlir] [MLIR] Auto generate source location for python bindings (PR #112923)

[Maksim Levental]

makslevental wrote:

I have another idea: https://docs.python.org/3/library/sys.html#sys.setprofile. Here's an example snippet:

```python
def my_profiler(frame, event, arg):
    if event == "call":
        frame_info = inspect.getframeinfo(frame)
        print(
            f"CALL: Function '{frame.f_code.co_name}' called in '{frame.f_code.co_filename}' at line {frame.f_lineno} at col: {frame_info.positions.col_offset}"
        )

import sys
sys.setprofile(my_profiler)
```

I dropped this into a random script and got prints like

```
CALL: Function 'yield_' called in '/home/mlevental/dev_projects/mlir-python-extras/mlir/extras/dialects/ext/scf.py' at line 368 at col: 0
CALL: Function 'yield_' called in '/home/mlevental/miniconda3/envs/mlir-python-extras-py.313/lib/python3.12/site-packages/mlir/dialects/_scf_ops_gen.py' at line 532 at col: 0
CALL: Function '__init__' called in '/home/mlevental/miniconda3/envs/mlir-python-extras-py.313/lib/python3.12/site-packages/mlir/dialects/_scf_ops_gen.py' at line 517 at col: 0
```

So the idea would be to basically keep a window of these locations:

```python

prev_locations = deque()

def location_tracker(frame, event, arg):
    if event == "call":
        frame_info = inspect.getframeinfo(frame)
        prev_locations.append((frame.f_code.co_name, frame.f_code.co_filename, frame.f_lineno, frame_info.positions.col_offset))
        if len(prev_locations) > 3:
            prev_locations.popleft()

import sys
sys.setprofile(location_tracker)
```

and then query them in the `__init__`s/builders:

```python
if loc is None:
    # find the most recent entry in prev_locations that is not within `_op_gen.py`
```

Note, obvious using `deque` is silly here (we can use `prev_loc1`, `prev_loc2`, etc and just update them round-robin).
Note, also, https://docs.python.org/3/library/threading.html#threading.setprofile also exists and is probably the better approach given we're moving to a "free threading" world soon-ish.

The advantage of this approach is that it does not force walk the frame stack each time we want to get a location. Naturally the downside is we're recording a window of the frame stack on every single call and that's high overhead.

**Here's the good part**: lucky for us, as of py3.12, there's a much faster way to do this kind of monitoring: https://peps.python.org/pep-0669/. 

So this works just like `my_profiler` above:

```python
sys.monitoring.use_tool_id(sys.monitoring.PROFILER_ID, "location_tracker")

def my_profiler_fast(code: CodeType, instruction_offset: int, callable: object, arg0: object):
    print(code.co_name, code.co_filename, code.co_firstlineno)

sys.monitoring.set_events(sys.monitoring.PROFILER_ID, sys.monitoring.events.CALL)
sys.monitoring.register_callback(
    sys.monitoring.PROFILER_ID, sys.monitoring.events.CALL, my_profiler
)
```
...
```
yield_ /home/mlevental/dev_projects/mlir-python-extras/mlir/extras/dialects/ext/scf.py 368
yield_ /home/mlevental/miniconda3/envs/mlir-python-extras-py.313/lib/python3.12/site-packages/mlir/dialects/_scf_ops_gen.py 532
__init__ /home/mlevental/miniconda3/envs/mlir-python-extras-py.313/lib/python3.12/site-packages/mlir/dialects/_scf_ops_gen.py 517
```

but is much faster (very evident from trying both).

So we could implement both of these (use a branch like `if sys.version_info.minor < 12: register_slow else: register_fast`), have them both off by default but togglable and then eventually (couple of years?) remove the slow path.

What do you guys think?

https://github.com/llvm/llvm-project/pull/112923