[PyCDE] Set MLIR op location based on Python call site

[teqdruid]

Right now, everything created in PyCDE has the operation location set to 'Unknown'. This makes debugging obscenely difficult. What I'd prefer is a mechanism to set the default location (DefaultingPyLocation) to call a function (either C++ or Python) to compute the location. In that function, we could get a stack trace, find the user code which is calling into PyCDE and creating the operation, then return a location which corresponds to that Python code.

@stellaraccident Is it possible to do this now? I haven't found a way. I also think it would just be a matter of modifying PyLocation to accept a callback rather than just a static location, yes? Alternatively, what do you think about making get() virtual and allowing users to override it from Python (a la https://pybind11.readthedocs.io/en/stable/advanced/classes.html#overriding-virtual-functions-in-python)?

[stellaraccident]

Yeah, been on my list, but most of my use cases all carry explicit locations so has been low priority.

The most straight forward thing to do is have another level of fallback when resolving a default location. This can use the traceback facility to get the current stack.

There is subtlety, though, because a lot of times, this will point to some library code vs end user entry point. For libraries such as yours, you probably want something like Location.current_python(stack_offset: int = 0).

You can then create a decorator for your library entry point functions like @library_entry, which gets the location at stack offset -1 (should be user code) and sets it as the default in a with block.

Based on experience, this will work but require some fiddling to get right (ie. Nested library calls can mess it up). The fully generic way of doing this will also cost you some efficiency, which is one reason I avoided it (for non user facing things).

The alternative is to do what Tensorflow does, and preserve the whole call stack at each op (as nested locations). I personally detest the sloppiness of that though. I wouldn't call Tensorflow a paragon of debuggability and i would advise against emulating this design point.

[teqdruid]

Based on experience, this will work but require some fiddling to get right (ie. Nested library calls can mess it up). The fully generic way of doing this will also cost you some efficiency, which is one reason I avoided it (for non user facing things).

Yeah, I agree so that's why I mentioned the callback approach -- so different Python bindings clients could choose how they dissect the stack trace. If it's a callback, it doesn't even need to be based on a stack trace! For instance, I was going to start by unwinding the stack until I get to a file which isn't part of my library then report that.

[teqdruid]

To be more specific about my proposal: On the Python API side, I'd add a Location.callback (alongside Location.file) which takes just a Python function which would be called whenever the C++ side needs a default location. On the C++ Pybind11 side, that Location.callback function would create a PyLocation with an Unknown loc field, and keep the callback py::function in another field alongside loc. Whenever get() is called: if the callback field is non-null it calls the callback, if it's null get() proceeds as normal (returns loc).

Does that make sense? Would it work?