[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

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

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

mlir/lib/Bindings/Python/IRModule.h

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 ```

**Context:** Update llvm, stablehlo and enzyme, 2025 Q4 The latest pair of good versions, indicated by stablehlo, is openxla/stablehlo@0a4440a ``` stablehlo=0a4440a5c8de45c4f9649bf3eb4913bf3f97da0d llvm=113f01aa82d055410f22a9d03b3468fa68600589 ``` For Enzyme, we go to the latest release https://github.com/EnzymeAD/Enzyme/releases/tag/v0.0.203 ``` enzyme=v0.0.203 ``` with commit `476c8e3193a8577ba24ff845ae2294109225f83a` **Description of the Change:** - `llvm-project/mlir/include/mlir/InitAllPasses.h` header no longer includes individual pass headers. To incorporate specific passes, you now need to include either `mlir/Conversion/Passes.h` or the header for each dialect pass. This change stems from the upstream cleanup of `mlir/InitAllDialects.h`, which previously handled these includes. Following the decision from the owner (as seen in this pull request comment: https://github.com/ftynse/water/pull/16/files#r2259526343), we will specifically include each required header, faster compilation (no unnecessary headers). For further context, refer to this PR: llvm/llvm-project#151150 - To bufferize custom tensors into custom buffers. The upstream PR llvm/llvm-project#144867 changed the the return type of `bufferization::detail::defaultGetBufferType` to be `BufferLikeType` instead of `BaseMemRefType`. We aligned the behaviour with the upstream PR, return the BufferLikeType from our `getBufferType()` implementation as well. - `elemwise_unary` and `elemwise_binary` of `linalg` have been deprecated from the upstream, replaced with `elementwise`. llvm/llvm-project#147082. Follow this PR, using `linalg.add` directly. register_all_dialects/passes/extension` should now be configured as a static library in CMake, refer to this PR: llvm/llvm-project#151150 - An opt-in feature (`{op}::create` method) be added to support the same behaviour of `rewriter.create` method. But it provides a meaningful message when using this new opt-in feature compared to `rewrite.create`. Since the new generated `create` function will call the `build` function, it requires us to define the instance of the builder (`GraidentOps.td`). Refer to the PR: llvm/llvm-project#147168. **Take `linalg.abs` as an example:** ```cpp void AbsOp::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ValueRange inputs, ValueRange outputs, ArrayRef<NamedAttribute> attributes) { buildStructuredOp(odsBuilder, odsState, std::nullopt, inputs, outputs, attributes, AbsOp::getRegionBuilder()); } AbsOp AbsOp::create(::mlir::OpBuilder &builder, ::mlir::Location location, ValueRange inputs, ValueRange outputs, ArrayRef<NamedAttribute> attributes) { ::mlir::OperationState __state__(location, getOperationName()); build(builder, __state__, inputs, outputs, attributes); auto __res__ = ::llvm::dyn_cast<AbsOp>(builder.create(__state__)); assert(__res__ && "builder didn't return the right type"); return __res__; } AbsOp AbsOp::create(::mlir::ImplicitLocOpBuilder &builder, ValueRange inputs, ValueRange outputs, ArrayRef<NamedAttribute> attributes) { return create(builder, builder.getLoc(), inputs, outputs, attributes); } ``` - Changed `EnzymeStatic-21` to `EnzymeStatic-22` - Mock `_ods_cext.globals.register_traceback_file_exclusion` due to API conflicts between Catalyst's MLIR version and the MLIR version used by JAX. The current JAX version we used has not yet updated to the latest MLIR, causing compatibility issues. This workaround will be removed once JAX updates to a compatible MLIR version. llvm/llvm-project#151246 TODO: - [x] Update .dep_version - [ ] Change all `rewrite.create` to `{op}::create` Even it’s compatible to the current `rewrite.create` method, we don’t need to change all `rewrite.create` to `{op}::create` right away, but this new feature provides more friendly development experience. Just for an example: llvm/llvm-project#147311 **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** [sc-100911] [sc-100912] --------- Co-authored-by: David Ittah <[email protected]>

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 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 jpienaar and rolfmorel August 12, 2025 05:24

makslevental commented Aug 12, 2025

View reviewed changes

mlir/lib/Bindings/Python/IRModule.h Outdated

Comment on lines -195 to -204

Copy link

Contributor Author

makslevental Aug 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

drive-by DCE (unused/no defn)

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

rniczh mentioned this pull request Oct 29, 2025

Update LLVM version, 2025 Q4 PennyLaneAI/catalyst#2122

Merged

2 tasks

[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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

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