[mlir][python] automatic location inference #151246
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
makslevental
force-pushed
the
makslevental/auto-locs
branch
3 times, most recently
from
61b0669 to
9ea8751
Compare
|
✅ With the latest revision this PR passed the Python code formatter. |
makslevental
force-pushed
the
makslevental/auto-locs
branch
11 times, most recently
from
34d5f14 to
7c53bf1
Compare
makslevental
force-pushed
the
makslevental/auto-locs
branch
14 times, most recently
from
389db39 to
f63eaf2
Compare
makslevental
force-pushed
the
makslevental/auto-locs
branch
5 times, most recently
from
3ab9de1 to
2297849
Compare
makslevental
force-pushed
the
makslevental/auto-locs
branch
from
2297849 to
a4bc05a
Compare
makslevental
commented
Aug 12, 2025
Comment on lines
+84
to
+86
| if _cext.ir.Location.current: | ||
| return _cext.ir.Location.current.context | ||
| return None |
There was a problem hiding this comment.
I realized since we're now supporting loc inference we shouldn't require a location to be in the context; specifically Location.current shouldn't throw but just return None. See tests in context_manager.py below.
There was a problem hiding this comment.
@jpienaar @rolfmorel i'm asking for a re-review because I made this change which wasn't in the original PR and slightly changes semantics; see also the change to Location.current here below.
makslevental
commented
Aug 12, 2025
Comment on lines
+3130
to
+3132
| [](nb::object & /*class*/) -> std::optional<PyLocation *> { | ||
| auto *loc = PyThreadContextEntry::getDefaultLocation(); | ||
| if (!loc) | ||
| throw nb::value_error("No current Location"); | ||
| return std::nullopt; |
There was a problem hiding this comment.
change Location.current to return None instead of throwing - see above
makslevental
commented
Aug 12, 2025
Comment on lines
-195
to
-204
| /// For the case of a python __init__ (nanobind::init) method, pybind11 is | ||
| /// quite strict about needing to return a pointer that is not yet associated | ||
| /// to an nanobind::object. Since the forContext() method acts like a pool, | ||
| /// possibly returning a recycled context, it does not satisfy this need. The | ||
| /// usual way in python to accomplish such a thing is to override __new__, but | ||
| /// that is also not supported by pybind11. Instead, we use this entry | ||
| /// point which always constructs a fresh context (which cannot alias an | ||
| /// existing one because it is fresh). | ||
| static PyMlirContext *createNewContextForInit(); | ||
|
|
There was a problem hiding this comment.
drive-by DCE (unused/no defn)
makslevental
force-pushed
the
makslevental/auto-locs
branch
from
b0ff949 to
70244f3
Compare
jpienaar
approved these changes
Aug 12, 2025
makslevental
added a commit
that referenced
this pull request
Aug 15, 2025
Introduced (but omitted from this CMake) in #151246.
makslevental
added a commit
that referenced
this pull request
Aug 15, 2025
…d fix warning unused Introduced (but omitted from this CMake) in #151246.
llvm-sync bot
pushed a commit
to arm/arm-toolchain
that referenced
this pull request
Aug 16, 2025
…y nanobind (#153861) Introduced (but omitted from this CMake) in llvm/llvm-project#151246.
bump-llvm bot
pushed a commit
to makslevental/python_bindings_fork
that referenced
this pull request
Aug 17, 2025
…(#153861) Introduced (but omitted from this CMake) in llvm/llvm-project#151246.
bump-llvm bot
pushed a commit
to makslevental/python_bindings_fork
that referenced
this pull request
Aug 17, 2025
…(#153861) Introduced (but omitted from this CMake) in llvm/llvm-project#151246.
|
Hi, I'm seeing the below error on macos build: I haven't tracked the change yet, but I'll try to provide a patch. Currently, we can only revert three commits together, otherwsie there are conflicts. |
|
Oh, 6fc1deb seems to fix the issue. Let me give it a shot. |
|
Yes, confirmed that it is fixed by the commit. Thanks |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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 adapted from JAX's implementation 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.pyfiles.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 theLocationusing specifically aCallSiteLoc. This can be useful for profiling tools (flamegraphs etc.).Shoutout to the folks at JAX for coming up with a good system (cc @hawkinsp 😄).