[mlir][python] automatic location inference #151246

makslevental · 2025-07-29T22:44:31Z

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.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 (cc @hawkinsp 😄).

github-actions · 2025-07-29T23:01:54Z

✅ With the latest revision this PR passed the Python code formatter.

makslevental · 2025-08-12T05:19:13Z

mlir/python/mlir/dialects/_ods_common.py

+        if _cext.ir.Location.current:
+            return _cext.ir.Location.current.context
+        return None


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.

@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 · 2025-08-12T05:25:46Z

mlir/lib/Bindings/Python/IRCore.cpp

+          [](nb::object & /*class*/) -> std::optional<PyLocation *> {
            auto *loc = PyThreadContextEntry::getDefaultLocation();
            if (!loc)
-              throw nb::value_error("No current Location");
+              return std::nullopt;


change Location.current to return None instead of throwing - see above

makslevental · 2025-08-12T05:27:47Z

mlir/lib/Bindings/Python/IRModule.h

-  /// 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();
-


drive-by DCE (unused/no defn)

…d fix warning unused Introduced (but omitted from this CMake) in #151246.

…153861) Introduced (but omitted from this CMake) in #151246.

…y nanobind (#153861) Introduced (but omitted from this CMake) in llvm/llvm-project#151246.

…(#153861) Introduced (but omitted from this CMake) in llvm/llvm-project#151246.

hanhanW · 2025-08-19T14:28:23Z

Hi, I'm seeing the below error on macos build:

ld: warning: ignoring duplicate libraries: '-lm'
ld: Undefined symbols:
  _PyCode_Addr2Location, referenced from:
      (anonymous namespace)::maybeGetTracebackLocation(std::__1::optional<mlir::python::PyLocation> const&) in IRCore.cpp.o
  _PyFrame_GetLasti, referenced from:
      (anonymous namespace)::maybeGetTracebackLocation(std::__1::optional<mlir::python::PyLocation> const&) in IRCore.cpp.o
clang: error: linker command failed with exit code 1 (use -v to see invocation)

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.

hanhanW · 2025-08-19T14:29:55Z

Oh, 6fc1deb seems to fix the issue. Let me give it a shot.

hanhanW · 2025-08-19T14:34:31Z

Yes, confirmed that it is fixed by the commit. Thanks

anmyachev · 2025-09-03T12:34:12Z

mlir/test/python/ir/auto_location.py

+        two = arith.constant(IndexType.get(), 2)
+
+        # fmt: off
+        # 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))))


Hi @makslevental

I recently faced the following error: ******************** TEST 'MLIR :: python/ir/auto_location.py' FAILED ******************** when LLVM was compiled for windows in OpenAI Triton project: https://github.com/triton-lang/triton/actions/runs/17419518063/job/49454994952.

More logs:

******************** TEST 'MLIR :: python/ir/auto_location.py' FAILED ******************** Exit Code: 1 Command Output (stdout): -- # RUN: at line 1 "C:/hostedtoolcache/windows/Python/3.11.9/x64/python3.exe" D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py | d:\a\triton\triton\llvm-project\build\bin\filecheck.exe D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py # executed command: C:/hostedtoolcache/windows/Python/3.11.9/x64/python3.exe 'D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py' # executed command: 'd:\a\triton\triton\llvm-project\build\bin\filecheck.exe' 'D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py' # .---command stderr------------ # | D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py:37:11: error: CHECK: expected string not found in input # | # 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)))) # | ^ # | <stdin>:2:25: note: scanning from here # | TEST: testInferLocations # | ^ # | # | Input file: <stdin> # | Check file: D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py # | # | -dump-input=help explains the following input dump. # | # | Input was: # | <<<<<< # | 1: # | 2: TEST: testInferLocations # | check:37 X error: no match found # | 3: loc(callsite("testInferLocations"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":31:13 to :43) at callsite("run"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":13:4 to :7) at "<module>"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":26:1 to :4)))) # | check:37

Probably this line should be changed [[SEP:[/\\]]] -> [[SEP]]?

Um no [[SEP:[/\\]]] defines SEP as the regex [/\\] so that's correct so I'm not sure what's going on there.

Thanks for the fast response! I'll try to debug it further.

That pattern should work and we do have windows tests here but actually I think they're not running because the CI (here, on Windows) doesn't have the minimum version of Python. I didn't have a Windows box handy when I put this together so it's true I wasn't able to test this thoroughly.

I managed to fix it with the following regexp: [[SEP:(/|\\\\)]]

In the previous approach, only one \ character was written to SEP.

I checked it with lit and the following file:

// RUN: cat %s | FileCheck %s # 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)))) 1: 2: TEST: testInferLocations 3: loc(callsite("testInferLocations"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":31:13 to :43) at callsite("run"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":13:4 to :7) at "<module>"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":26:1 to :4))))

Ok that makes sense but just for sanity check can you try SEP:[/\\]+

But btw if you send a PR I'll happily approve it (otherwise I can send one soon).

SEP:[/\]+

yes, it also works. Do you prefer this way?

But btw if you send a PR I'll happily approve it (otherwise I can send one soon).

I'll open a PR for it.

Initially found in: #151246 (comment) To fix: ```txt ******************** TEST 'MLIR :: python/ir/auto_location.py' FAILED ******************** Exit Code: 1 Command Output (stdout): -- # RUN: at line 1 "C:/hostedtoolcache/windows/Python/3.11.9/x64/python3.exe" D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py | d:\a\triton\triton\llvm-project\build\bin\filecheck.exe D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py # executed command: C:/hostedtoolcache/windows/Python/3.11.9/x64/python3.exe 'D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py' # executed command: 'd:\a\triton\triton\llvm-project\build\bin\filecheck.exe' 'D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py' # .---command stderr------------ # | D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py:37:11: error: CHECK: expected string not found in input # | # 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)))) # | ^ # | <stdin>:2:25: note: scanning from here # | TEST: testInferLocations # | ^ # | # | Input file: <stdin> # | Check file: D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py # | # | -dump-input=help explains the following input dump. # | # | Input was: # | <<<<<< # | 1: # | 2: TEST: testInferLocations # | check:37 X error: no match found # | 3: loc(callsite("testInferLocations"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":31:13 to :43) at callsite("run"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":13:4 to :7) at "<module>"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":26:1 to :4)))) # | check:37 ```

Initially found in: llvm/llvm-project#151246 (comment) To fix: ```txt ******************** TEST 'MLIR :: python/ir/auto_location.py' FAILED ******************** Exit Code: 1 Command Output (stdout): -- # RUN: at line 1 "C:/hostedtoolcache/windows/Python/3.11.9/x64/python3.exe" D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py | d:\a\triton\triton\llvm-project\build\bin\filecheck.exe D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py # executed command: C:/hostedtoolcache/windows/Python/3.11.9/x64/python3.exe 'D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py' # executed command: 'd:\a\triton\triton\llvm-project\build\bin\filecheck.exe' 'D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py' # .---command stderr------------ # | D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py:37:11: error: CHECK: expected string not found in input # | # 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)))) # | ^ # | <stdin>:2:25: note: scanning from here # | TEST: testInferLocations # | ^ # | # | Input file: <stdin> # | Check file: D:\a\triton\triton\llvm-project\mlir\test\python\ir\auto_location.py # | # | -dump-input=help explains the following input dump. # | # | Input was: # | <<<<<< # | 1: # | 2: TEST: testInferLocations # | check:37 X error: no match found # | 3: loc(callsite("testInferLocations"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":31:13 to :43) at callsite("run"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":13:4 to :7) at "<module>"("D:\\a\\triton\\triton\\llvm-project\\mlir\\test\\python\\ir\\auto_location.py":26:1 to :4)))) # | check:37 ```

makslevental force-pushed the makslevental/auto-locs branch 3 times, most recently from 61b0669 to 9ea8751 Compare July 29, 2025 22:59

makslevental force-pushed the makslevental/auto-locs branch 11 times, most recently from 34d5f14 to 7c53bf1 Compare July 30, 2025 03:25

makslevental mentioned this pull request Jul 30, 2025

[MLIR] Auto generate source location for python bindings #112923

Open

makslevental force-pushed the makslevental/auto-locs branch 14 times, most recently from 389db39 to f63eaf2 Compare July 30, 2025 05:05

makslevental force-pushed the makslevental/auto-locs branch from 3ab9de1 to 2297849 Compare August 12, 2025 05:12

fix live contexts

a4bc05a

makslevental force-pushed the makslevental/auto-locs branch from 2297849 to a4bc05a Compare August 12, 2025 05:13

fix location is None test

21d68e6

makslevental commented Aug 12, 2025

View reviewed changes

makslevental requested review from rolfmorel and jpienaar August 12, 2025 05:24

makslevental commented Aug 12, 2025

View reviewed changes

upper bound kMaxFrames

70244f3

makslevental force-pushed the makslevental/auto-locs branch from b0ff949 to 70244f3 Compare August 12, 2025 05:35

jpienaar approved these changes Aug 12, 2025

View reviewed changes

add comment

2e63eee

makslevental merged commit a40f47c into llvm:main Aug 12, 2025
9 checks passed

makslevental deleted the makslevental/auto-locs branch August 12, 2025 22:00

joker-eph mentioned this pull request Aug 13, 2025

[MLIR] Update Maintainers.md: add Python bindings under "Core" #153410

Closed

makslevental mentioned this pull request Aug 15, 2025

[mlir][python] handle more undefined symbols not covered by nanobind #153861

Merged

makslevental added a commit that referenced this pull request Aug 15, 2025

[mlir][python] handle more undefined symbols not covered by pybind an…

632c7a7

…d fix warning unused Introduced (but omitted from this CMake) in #151246.

makslevental added a commit that referenced this pull request Aug 16, 2025

[mlir][python] handle more undefined symbols not covered by nanobind (#…

6fc1deb

…153861) 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

Automerge: [mlir][python] handle more undefined symbols not covered b…

45253fa

…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

[mlir][python] handle more undefined symbols not covered by nanobind …

556e703

…(#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

[mlir][python] handle more undefined symbols not covered by nanobind …

b3563ff

…(#153861) Introduced (but omitted from this CMake) in llvm/llvm-project#151246.

anmyachev reviewed Sep 3, 2025

View reviewed changes

anmyachev mentioned this pull request Sep 3, 2025

Fix python/ir/auto_location.py test on Windows #156711

Merged

[mlir][python] automatic location inference #151246

[mlir][python] automatic location inference #151246

Uh oh!

Conversation

makslevental commented Jul 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jul 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

makslevental Aug 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

makslevental Aug 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hanhanW commented Aug 19, 2025

Uh oh!

hanhanW commented Aug 19, 2025

Uh oh!

hanhanW commented Aug 19, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

makslevental commented Jul 29, 2025 •

edited

Loading

github-actions bot commented Jul 29, 2025 •

edited

Loading

makslevental Aug 12, 2025 •

edited

Loading

makslevental Aug 12, 2025 •

edited

Loading