[Mlir-commits] [mlir] [mlir][python] Add stable ABI (abi3) support (PR #183856)
Jakub Kuderski
llvmlistbot at llvm.org
Sat Feb 28 11:13:46 PST 2026
https://github.com/kuhar updated https://github.com/llvm/llvm-project/pull/183856
>From a282f07653beac6ef51e49885142142f17d9d588 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 16:52:10 -0500
Subject: [PATCH 01/10] [mlir][python] Add stable ABI (abi3) support in
bindings
Add `MLIR_ENABLE_PYTHON_STABLE_ABI` cmake flag to build
bindings against the Python limited/stable API (abi3 / PEP 384).
This allow for compatibility across different >=3.12 versions with a
single .so / wheel.
The stable ABI restricts usage to a subset of the CPython C API: frame
and code object structs are opaque, so introspection APIs like
`PyCode_Addr2Location`, `PyFrame_GetLasti`, and `PyFrame_GetCode` are
unavailable. The traceback-based location logic is reworked to use
`sys._getframe()` and Python-level attribute access instead. Since column
information is not available through the limited API, stable ABI builds
emit `:0` for columns in traceback locations (hence a separate test
gated by `REQUIRES`).
Assisted-by: claude
---
mlir/CMakeLists.txt | 8 ++
mlir/cmake/modules/AddMLIRPython.cmake | 39 +++++++-
mlir/cmake/modules/MLIRDetectPythonEnv.cmake | 6 ++
mlir/include/mlir/Bindings/Python/IRCore.h | 4 +-
.../mlir/Bindings/Python/NanobindAdaptors.h | 9 +-
.../mlir/Bindings/Python/NanobindUtils.h | 91 +++++++++--------
mlir/lib/Bindings/Python/IRAttributes.cpp | 25 ++++-
mlir/lib/Bindings/Python/IRCore.cpp | 99 ++++++++++---------
mlir/test/CMakeLists.txt | 1 +
mlir/test/lit.cfg.py | 3 +
mlir/test/lit.site.cfg.py.in | 1 +
mlir/test/python/ir/auto_location.py | 3 +-
.../python/ir/auto_location_stable_abi.py | 46 +++++++++
13 files changed, 230 insertions(+), 105 deletions(-)
create mode 100644 mlir/test/python/ir/auto_location_stable_abi.py
diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
index 9e1e9314511e3..be15007a50c95 100644
--- a/mlir/CMakeLists.txt
+++ b/mlir/CMakeLists.txt
@@ -200,6 +200,14 @@ set(MLIR_PYTHON_PACKAGE_PREFIX "mlir"
embedded in a relocatable way).")
set(MLIR_ENABLE_BINDINGS_PYTHON 0 CACHE BOOL
"Enables building of Python bindings.")
+set(MLIR_ENABLE_PYTHON_STABLE_ABI 0 CACHE BOOL
+ "Build Python bindings against the stable ABI (abi3, PEP 384) for \
+ cross-version compatibility. Requires Python 3.12+.")
+if(MLIR_ENABLE_PYTHON_STABLE_ABI AND CMAKE_VERSION VERSION_LESS "3.26.0")
+ message(FATAL_ERROR
+ "MLIR_ENABLE_PYTHON_STABLE_ABI requires CMake >= 3.26 "
+ "(Python::SABIModule not available in ${CMAKE_VERSION}).")
+endif()
set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/mlir_core/mlir" CACHE STRING
"Prefix under install directory to place python bindings")
set(MLIR_DETECT_PYTHON_ENV_PRIME_SEARCH 1 CACHE BOOL
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 54c59f41404b7..1821cfbf35d2a 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -311,18 +311,37 @@ function(build_nanobind_lib)
# Only build in free-threaded mode if the Python ABI supports it.
# See https://github.com/wjakob/nanobind/blob/4ba51fcf795971c5d603d875ae4184bc0c9bd8e6/cmake/nanobind-config.cmake#L363-L371.
- if (NB_ABI MATCHES "[0-9]t")
+ if(NB_ABI MATCHES "[0-9]t")
set(_ft "-ft")
+ set(_abi3 "")
+ else()
+ set(_ft "")
+ # Match nanobind_add_module's naming: with STABLE_ABI, the shared library
+ # name includes "-abi3" (e.g., "nanobind-abi3-mlir").
+ if(MLIR_ENABLE_PYTHON_STABLE_ABI)
+ set(_abi3 "-abi3")
+ else()
+ set(_abi3 "")
+ endif()
endif()
# nanobind does a string match on the suffix to figure out whether to build
# the lib with free threading...
- set(NB_LIBRARY_TARGET_NAME "nanobind${_ft}-${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}")
+ set(NB_LIBRARY_TARGET_NAME "nanobind${_ft}${_abi3}-${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}")
set(NB_LIBRARY_TARGET_NAME "${NB_LIBRARY_TARGET_NAME}" PARENT_SCOPE)
nanobind_build_library(${NB_LIBRARY_TARGET_NAME} AS_SYSINCLUDE)
target_compile_definitions(${NB_LIBRARY_TARGET_NAME}
PRIVATE
NB_DOMAIN=${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
)
+ # Propagate stable ABI to the shared nanobind library. nanobind internally
+ # skips this when the interpreter is free-threaded.
+ if(MLIR_ENABLE_PYTHON_STABLE_ABI AND NOT (NB_ABI MATCHES "[0-9]t"))
+ target_compile_definitions(${NB_LIBRARY_TARGET_NAME}
+ PUBLIC
+ Py_LIMITED_API=0x030C0000
+ )
+ target_link_libraries(${NB_LIBRARY_TARGET_NAME} PRIVATE Python::SABIModule)
+ endif()
if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
# nanobind handles this correctly for MacOS by explicitly setting -U for all the necessary Python symbols
# (see https://github.com/wjakob/nanobind/blob/master/cmake/darwin-ld-cpython.sym)
@@ -913,9 +932,15 @@ function(add_mlir_python_extension libname extname nb_library_target_name)
set_property(TARGET ${libname} PROPERTY WINDOWS_EXPORT_ALL_SYMBOLS ON)
endif()
else()
+ if(MLIR_ENABLE_PYTHON_STABLE_ABI)
+ set(_stable_abi_flag STABLE_ABI)
+ else()
+ set(_stable_abi_flag "")
+ endif()
nanobind_add_module(${libname}
NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
FREE_THREADED
+ ${_stable_abi_flag}
NB_SHARED
${ARG_SOURCES}
)
@@ -973,9 +998,15 @@ function(add_mlir_python_extension libname extname nb_library_target_name)
# Same for the rest.
target_link_options(${libname} PUBLIC
"LINKER:-U,_PyClassMethod_New"
- "LINKER:-U,_PyCode_Addr2Location"
- "LINKER:-U,_PyFrame_GetLasti"
)
+ if(NOT MLIR_ENABLE_PYTHON_STABLE_ABI)
+ # PyCode_Addr2Location and PyFrame_GetLasti are not part of the stable
+ # ABI and are not referenced when Py_LIMITED_API is defined.
+ target_link_options(${libname} PUBLIC
+ "LINKER:-U,_PyCode_Addr2Location"
+ "LINKER:-U,_PyFrame_GetLasti"
+ )
+ endif()
endif()
target_compile_options(${libname} PRIVATE ${eh_rtti_enable})
diff --git a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
index 01dd7437cc8fc..068d6712393c2 100644
--- a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
+++ b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
@@ -19,6 +19,9 @@ macro(mlir_configure_python_dev_packages)
# Development.Module.
# See https://pybind11.readthedocs.io/en/stable/compiling.html#findpython-mode
set(_python_development_component Development.Module)
+ if(MLIR_ENABLE_PYTHON_STABLE_ABI)
+ list(APPEND _python_development_component Development.SABIModule)
+ endif()
find_package(Python3 ${MLIR_MINIMUM_PYTHON_VERSION}
COMPONENTS Interpreter ${_python_development_component} REQUIRED)
@@ -44,6 +47,9 @@ macro(mlir_configure_python_dev_packages)
COMPONENTS Interpreter ${_python_development_component} REQUIRED)
unset(_python_development_component)
+ if(MLIR_ENABLE_PYTHON_STABLE_ABI)
+ message(STATUS "MLIR Python stable ABI (abi3) enabled")
+ endif()
message(STATUS "Found python include dirs: ${Python3_INCLUDE_DIRS}")
message(STATUS "Found python libraries: ${Python3_LIBRARIES}")
message(STATUS "Found numpy v${Python3_NumPy_VERSION}: ${Python3_NumPy_INCLUDE_DIRS}")
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index e9669c4b2726d..cdd99b52f5200 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -4,7 +4,6 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//===----------------------------------------------------------------------===//
#ifndef MLIR_BINDINGS_PYTHON_IRCORE_H
@@ -1877,7 +1876,8 @@ MLIR_PYTHON_API_EXPORTED void populateRoot(nanobind::module_ &m);
template <class Func, typename... Args>
inline nanobind::object classmethod(Func f, Args... args) {
nanobind::object cf = nanobind::cpp_function(f, args...);
- return nanobind::borrow<nanobind::object>((PyClassMethod_New(cf.ptr())));
+ nanobind::object builtins = nanobind::module_::import_("builtins");
+ return builtins.attr("classmethod")(cf);
}
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 9b54baa5bd4b1..95151bdaf8ca6 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -186,10 +186,11 @@ struct type_caster<MlirContext> {
// If there is no context, including thread-bound, emit a warning (since
// this function is not allowed to throw) and fail to cast.
if (src.is_none()) {
- PyErr_Warn(
+ PyErr_WarnEx(
PyExc_RuntimeWarning,
"Passing None as MLIR Context is only allowed inside "
- "the " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " context manager.");
+ "the " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " context manager.",
+ /*stacklevel=*/1);
return false;
}
if (std::optional<nanobind::object> capsule = mlirApiObjectToCapsule(src)) {
@@ -495,8 +496,8 @@ class pure_subclass {
std::forward<Func>(f),
nanobind::name(name), // nanobind::scope(thisClass),
extra...);
- thisClass.attr(name) =
- nanobind::borrow<nanobind::object>(PyClassMethod_New(cf.ptr()));
+ nanobind::object builtins = nanobind::module_::import_("builtins");
+ thisClass.attr(name) = builtins.attr("classmethod")(cf);
return *this;
}
diff --git a/mlir/include/mlir/Bindings/Python/NanobindUtils.h b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
index 55d183ab2df42..53c8b823b4d86 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindUtils.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
@@ -383,8 +383,55 @@ class Sliceable {
return elements;
}
+ // Manually implement the sequence protocol via the C API. We do this
+ // because it is approx 4x faster than via nanobind, largely because that
+ // formulation requires a C++ exception to be thrown to detect end of
+ // sequence.
+ // Since we are in a C-context, any C++ exception that happens here
+ // will terminate the program. There is nothing in this implementation
+ // that should throw in a non-terminal way, so we forgo further
+ // exception marshalling.
+ // See: https://github.com/pybind/nanobind/issues/2842
+ //
/// Binds the indexing and length methods in the Python class.
static void bind(nanobind::module_ &m) {
+ // These slots are passed via nanobind::type_slots() at class creation
+ // time, which is compatible with both the full and limited (stable ABI)
+ // Python APIs.
+ static PyType_Slot sequenceSlots[] = {
+ {Py_sq_length, (void *)(+[](PyObject *rawSelf) -> Py_ssize_t {
+ auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
+ return self->length;
+ })},
+ // sq_item is called as part of the sequence protocol for iteration,
+ // list construction, etc.
+ {Py_sq_item,
+ (void *)(+[](PyObject *rawSelf, Py_ssize_t index) -> PyObject * {
+ auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
+ return self->getItem(index).release().ptr();
+ })},
+ // mp_subscript is used for both slices and integer lookups.
+ {Py_mp_subscript,
+ (void *)(+[](PyObject *rawSelf,
+ PyObject *rawSubscript) -> PyObject * {
+ auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
+ Py_ssize_t index =
+ PyNumber_AsSsize_t(rawSubscript, PyExc_IndexError);
+ if (!PyErr_Occurred()) {
+ // Integer indexing.
+ return self->getItem(index).release().ptr();
+ }
+ PyErr_Clear();
+
+ // Assume slice-based indexing.
+ if (PySlice_Check(rawSubscript)) {
+ return self->getItemSlice(rawSubscript).release().ptr();
+ }
+
+ PyErr_SetString(PyExc_ValueError, "expected integer or slice");
+ return nullptr;
+ })},
+ {0, nullptr}};
const std::type_info &elemTy = typeid(ElementTy);
PyObject *elemTyInfo = nanobind::detail::nb_type_lookup(&elemTy);
assert(elemTyInfo &&
@@ -394,52 +441,10 @@ class Sliceable {
"(collections.abc.Sequence[" +
nanobind::cast<std::string>(elemTyName) + "])";
auto clazz = nanobind::class_<Derived>(m, Derived::pyClassName,
+ nanobind::type_slots(sequenceSlots),
nanobind::sig(sig.c_str()))
.def("__add__", &Sliceable::dunderAdd);
Derived::bindDerived(clazz);
-
- // Manually implement the sequence protocol via the C API. We do this
- // because it is approx 4x faster than via nanobind, largely because that
- // formulation requires a C++ exception to be thrown to detect end of
- // sequence.
- // Since we are in a C-context, any C++ exception that happens here
- // will terminate the program. There is nothing in this implementation
- // that should throw in a non-terminal way, so we forgo further
- // exception marshalling.
- // See: https://github.com/pybind/nanobind/issues/2842
- auto heap_type = reinterpret_cast<PyHeapTypeObject *>(clazz.ptr());
- assert(heap_type->ht_type.tp_flags & Py_TPFLAGS_HEAPTYPE &&
- "must be heap type");
- heap_type->as_sequence.sq_length = +[](PyObject *rawSelf) -> Py_ssize_t {
- auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
- return self->length;
- };
- // sq_item is called as part of the sequence protocol for iteration,
- // list construction, etc.
- heap_type->as_sequence.sq_item =
- +[](PyObject *rawSelf, Py_ssize_t index) -> PyObject * {
- auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
- return self->getItem(index).release().ptr();
- };
- // mp_subscript is used for both slices and integer lookups.
- heap_type->as_mapping.mp_subscript =
- +[](PyObject *rawSelf, PyObject *rawSubscript) -> PyObject * {
- auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
- Py_ssize_t index = PyNumber_AsSsize_t(rawSubscript, PyExc_IndexError);
- if (!PyErr_Occurred()) {
- // Integer indexing.
- return self->getItem(index).release().ptr();
- }
- PyErr_Clear();
-
- // Assume slice-based indexing.
- if (PySlice_Check(rawSubscript)) {
- return self->getItemSlice(rawSubscript).release().ptr();
- }
-
- PyErr_SetString(PyExc_ValueError, "expected integer or slice");
- return nullptr;
- };
}
/// Hook for derived classes willing to bind more methods.
diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index c0c555ac51aa1..59eb9b8e81cf0 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1039,9 +1039,28 @@ nb::int_ PyDenseIntElementsAttribute::dunderGetItem(intptr_t pos) const {
void PyDenseIntElementsAttribute::bindDerived(ClassTy &c) {
c.def("__getitem__", &PyDenseIntElementsAttribute::dunderGetItem);
}
-// Check if the python version is less than 3.13. Py_IsFinalizing is a part
-// of stable ABI since 3.13 and before it was available as _Py_IsFinalizing.
-#if PY_VERSION_HEX < 0x030d0000
+
+// Py_IsFinalizing is part of the stable ABI since 3.13. Before that, it was
+// available as the private _Py_IsFinalizing, which is not part of the limited
+// API.
+#if defined(Py_LIMITED_API) && Py_LIMITED_API < 0x030d0000
+// Under limited API targeting < 3.13, use sys.is_finalizing() via C API.
+// PySys_GetObject avoids import machinery (safe during finalization).
+static int Py_IsFinalizing(void) {
+ // PySys_GetObject returns a borrowed reference; no Py_DECREF needed.
+ PyObject *fn = PySys_GetObject("is_finalizing");
+ if (!fn)
+ return 0;
+ PyObject *result = PyObject_CallNoArgs(fn);
+ if (!result) {
+ PyErr_Clear();
+ return 0;
+ }
+ int val = PyObject_IsTrue(result);
+ Py_DECREF(result);
+ return val > 0 ? 1 : 0;
+}
+#elif PY_VERSION_HEX < 0x030d0000
#define Py_IsFinalizing _Py_IsFinalizing
#endif
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 88d890f36b811..f9594f38297fd 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -803,7 +803,7 @@ nb::tuple PyDiagnostic::getNotes() {
for (intptr_t i = 0; i < numNotes; ++i) {
MlirDiagnostic noteDiag = mlirDiagnosticGetNote(diagnostic, i);
nb::object diagnostic = nb::cast(PyDiagnostic(noteDiag));
- PyTuple_SET_ITEM(notes.ptr(), i, diagnostic.release().ptr());
+ PyTuple_SetItem(notes.ptr(), i, diagnostic.release().ptr());
}
materializedNotes = std::move(notes);
@@ -2672,47 +2672,6 @@ void PyDynamicOpTraits::NoTerminator::bind(nb::module_ &m) {
} // namespace mlir
namespace {
-// see
-// https://raw.githubusercontent.com/python/pythoncapi_compat/master/pythoncapi_compat.h
-
-#ifndef _Py_CAST
-#define _Py_CAST(type, expr) ((type)(expr))
-#endif
-
-// Static inline functions should use _Py_NULL rather than using directly NULL
-// to prevent C++ compiler warnings. On C23 and newer and on C++11 and newer,
-// _Py_NULL is defined as nullptr.
-#ifndef _Py_NULL
-#if (defined(__STDC_VERSION__) && __STDC_VERSION__ > 201710L) || \
- (defined(__cplusplus) && __cplusplus >= 201103)
-#define _Py_NULL nullptr
-#else
-#define _Py_NULL NULL
-#endif
-#endif
-
-// Python 3.10.0a3
-#if PY_VERSION_HEX < 0x030A00A3
-
-// bpo-42262 added Py_XNewRef()
-#if !defined(Py_XNewRef)
-[[maybe_unused]] PyObject *_Py_XNewRef(PyObject *obj) {
- Py_XINCREF(obj);
- return obj;
-}
-#define Py_XNewRef(obj) _Py_XNewRef(_PyObject_CAST(obj))
-#endif
-
-// bpo-42262 added Py_NewRef()
-#if !defined(Py_NewRef)
-[[maybe_unused]] PyObject *_Py_NewRef(PyObject *obj) {
- Py_INCREF(obj);
- return obj;
-}
-#define Py_NewRef(obj) _Py_NewRef(_PyObject_CAST(obj))
-#endif
-
-#endif // Python 3.10.0a3
using namespace mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN;
@@ -2725,6 +2684,43 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
size_t count = 0;
nb::gil_scoped_acquire acquire;
+
+#if defined(Py_LIMITED_API)
+ // Under the limited API (targeting Python 3.12+), most frame introspection
+ // C APIs (PyCode_Addr2Location, PyFrame_GetLasti, etc.) are not available.
+ // Use Python-level sys._getframe() and attribute access instead. Note:
+ // co_qualname (used below) requires Python 3.11+.
+ nb::object sys_mod = nb::module_::import_("sys");
+ nb::object frameObj;
+ try {
+ // _getframe(0) from C++ returns the topmost Python frame, equivalent
+ // to PyThreadState_GetFrame() in the non-limited-API path.
+ frameObj = sys_mod.attr("_getframe")(0);
+ } catch (nb::python_error &) {
+ return mlirLocationUnknownGet(ctx);
+ }
+
+ while (!frameObj.is_none() && frameObj.ptr() != nullptr &&
+ count < framesLimit) {
+ nb::object codeObj = frameObj.attr("f_code");
+ auto fileNameStr = nb::cast<std::string>(codeObj.attr("co_filename"));
+ std::string_view fileName(fileNameStr);
+ if (PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName)) {
+ std::string name = nb::cast<std::string>(codeObj.attr("co_qualname"));
+ std::string_view funcName(name);
+ int startLine = nb::cast<int>(frameObj.attr("f_lineno"));
+ MlirLocation loc = mlirLocationFileLineColGet(
+ ctx, mlirStringRefCreate(fileName.data(), fileName.size()), startLine,
+ 0);
+ frames[count] = mlirLocationNameGet(
+ ctx, mlirStringRefCreate(funcName.data(), funcName.size()), loc);
+ ++count;
+ }
+ frameObj = frameObj.attr("f_back");
+ if (frameObj.is_none())
+ break;
+ }
+#else
PyThreadState *tstate = PyThreadState_GET();
PyFrameObject *next;
PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate);
@@ -2736,24 +2732,30 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
for (; pyFrame != nullptr && count < framesLimit;
next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) {
PyCodeObject *code = PyFrame_GetCode(pyFrame);
- auto fileNameStr =
- nb::cast<std::string>(nb::borrow<nb::str>(code->co_filename));
+ // Use attribute access instead of direct struct member access for
+ // forward compatibility with the limited (stable) API where
+ // PyCodeObject is opaque.
+ nb::object fileNameObj =
+ nb::steal(PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_filename"));
+ auto fileNameStr = nb::cast<std::string>(fileNameObj);
std::string_view fileName(fileNameStr);
if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName))
continue;
// co_qualname and PyCode_Addr2Location added in py3.11
#if PY_VERSION_HEX < 0x030B00F0
- std::string name =
- nb::cast<std::string>(nb::borrow<nb::str>(code->co_name));
+ nb::object nameObj =
+ nb::steal(PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_name"));
+ std::string name = nb::cast<std::string>(nameObj);
std::string_view funcName(name);
int startLine = PyFrame_GetLineNumber(pyFrame);
MlirLocation loc = mlirLocationFileLineColGet(
ctx, mlirStringRefCreate(fileName.data(), fileName.size()), startLine,
0);
#else
- std::string name =
- nb::cast<std::string>(nb::borrow<nb::str>(code->co_qualname));
+ nb::object nameObj =
+ nb::steal(PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_qualname"));
+ std::string name = nb::cast<std::string>(nameObj);
std::string_view funcName(name);
int startLine, startCol, endLine, endCol;
int lasti = PyFrame_GetLasti(pyFrame);
@@ -2773,6 +2775,7 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
// When the loop breaks (after the last iter), current frame (if non-null)
// is leaked without this.
Py_XDECREF(pyFrame);
+#endif
if (count == 0)
return mlirLocationUnknownGet(ctx);
diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index 55534b9910503..a5492d69b5ee1 100644
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -70,6 +70,7 @@ llvm_canonicalize_cmake_booleans(
LLVM_HAS_NVPTX_TARGET
LLVM_INCLUDE_SPIRV_TOOLS_TESTS
MLIR_ENABLE_BINDINGS_PYTHON
+ MLIR_ENABLE_PYTHON_STABLE_ABI
MLIR_ENABLE_CUDA_RUNNER
MLIR_ENABLE_ROCM_CONVERSIONS
MLIR_ENABLE_ROCM_RUNNER
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index 174623eb347b9..b20bfe5a4147d 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -392,3 +392,6 @@ def have_host_jit_feature_support(feature_name):
if sys.version_info >= (3, 11):
config.available_features.add("python-ge-311")
+
+if config.enable_python_stable_abi:
+ config.available_features.add("python-stable-abi")
diff --git a/mlir/test/lit.site.cfg.py.in b/mlir/test/lit.site.cfg.py.in
index b14a11163c107..0f3a014875946 100644
--- a/mlir/test/lit.site.cfg.py.in
+++ b/mlir/test/lit.site.cfg.py.in
@@ -44,6 +44,7 @@ config.enable_levelzero_runner = @MLIR_ENABLE_LEVELZERO_RUNNER@
config.enable_spirv_cpu_runner = @MLIR_ENABLE_SPIRV_CPU_RUNNER@
config.enable_vulkan_runner = @MLIR_ENABLE_VULKAN_RUNNER@
config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@
+config.enable_python_stable_abi = @MLIR_ENABLE_PYTHON_STABLE_ABI@
config.intel_sde_executable = "@INTEL_SDE_EXECUTABLE@"
config.mlir_run_amx_tests = @MLIR_RUN_AMX_TESTS@
config.mlir_run_arm_sve_tests = @MLIR_RUN_ARM_SVE_TESTS@
diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 6448a88dc1775..ba0fa9b7e8e2e 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -1,5 +1,6 @@
# RUN: %PYTHON %s | FileCheck %s
# REQUIRES: python-ge-311
+# UNSUPPORTED: python-stable-abi
import gc
from contextlib import contextmanager
@@ -27,7 +28,7 @@ def testInferLocations():
two = arith.constant(IndexType.get(), 2)
# fmt: off
- # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP:[/\\]+]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}: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":{{[0-9]+}}:1 to :4))))
+ # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP:[/\\]+]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:13 to :43) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":14:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))
# fmt: on
print(op.location)
diff --git a/mlir/test/python/ir/auto_location_stable_abi.py b/mlir/test/python/ir/auto_location_stable_abi.py
new file mode 100644
index 0000000000000..4d413a8d457ca
--- /dev/null
+++ b/mlir/test/python/ir/auto_location_stable_abi.py
@@ -0,0 +1,46 @@
+# RUN: %PYTHON %s | FileCheck %s
+# REQUIRES: python-stable-abi
+#
+# Verify that traceback-based locations work under the stable ABI (abi3).
+# The limited API path uses sys._getframe() and cannot provide column info,
+# so locations have :0 for columns. This test checks that function names,
+# file paths, callsite nesting, and frame limiting all work correctly.
+
+import gc
+from mlir.ir import *
+from mlir.dialects import arith
+
+
+def run(f):
+ print("\nTEST:", f.__name__)
+ f()
+ gc.collect()
+ assert Context._get_live_count() == 0
+
+
+# CHECK-LABEL: TEST: testStableABILocations
+ at run
+def testStableABILocations():
+ with Context() as ctx, loc_tracebacks():
+ ctx.allow_unregistered_dialects = True
+
+ # Basic op creation produces a callsite location with function names.
+ op = Operation.create("custom.op1")
+ # CHECK: loc(callsite("testStableABILocations"({{.*}}auto_location_stable_abi.py":{{[0-9]+}}:0)
+ print(op.location)
+
+ # Nested function calls produce nested callsite locations.
+ def inner():
+ return arith.constant(IndexType.get(), 1)
+
+ val = inner()
+ # CHECK: loc(callsite({{.*}} at callsite("testStableABILocations.<locals>.inner"({{.*}}auto_location_stable_abi.py":{{[0-9]+}}:0)
+ print(val.location)
+
+ # Frame limit of 0 produces unknown location.
+ from mlir.dialects._ods_common import _cext
+
+ _cext.globals.set_loc_tracebacks_frame_limit(0)
+ val2 = arith.constant(IndexType.get(), 2)
+ # CHECK: loc(unknown)
+ print(val2.location)
>From 4eaaa4f64bba64df44cfbc1332b900013b617f1b Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 17:54:26 -0500
Subject: [PATCH 02/10] Format
---
mlir/include/mlir/Bindings/Python/NanobindUtils.h | 3 +--
mlir/lib/Bindings/Python/IRCore.cpp | 12 ++++++------
2 files changed, 7 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Bindings/Python/NanobindUtils.h b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
index 53c8b823b4d86..15ed8c73190eb 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindUtils.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
@@ -412,8 +412,7 @@ class Sliceable {
})},
// mp_subscript is used for both slices and integer lookups.
{Py_mp_subscript,
- (void *)(+[](PyObject *rawSelf,
- PyObject *rawSubscript) -> PyObject * {
+ (void *)(+[](PyObject *rawSelf, PyObject *rawSubscript) -> PyObject * {
auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
Py_ssize_t index =
PyNumber_AsSsize_t(rawSubscript, PyExc_IndexError);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index f9594f38297fd..48dd5a7f55802 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2735,8 +2735,8 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
// Use attribute access instead of direct struct member access for
// forward compatibility with the limited (stable) API where
// PyCodeObject is opaque.
- nb::object fileNameObj =
- nb::steal(PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_filename"));
+ nb::object fileNameObj = nb::steal(PyObject_GetAttrString(
+ reinterpret_cast<PyObject *>(code), "co_filename"));
auto fileNameStr = nb::cast<std::string>(fileNameObj);
std::string_view fileName(fileNameStr);
if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName))
@@ -2744,8 +2744,8 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
// co_qualname and PyCode_Addr2Location added in py3.11
#if PY_VERSION_HEX < 0x030B00F0
- nb::object nameObj =
- nb::steal(PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_name"));
+ nb::object nameObj = nb::steal(
+ PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_name"));
std::string name = nb::cast<std::string>(nameObj);
std::string_view funcName(name);
int startLine = PyFrame_GetLineNumber(pyFrame);
@@ -2753,8 +2753,8 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
ctx, mlirStringRefCreate(fileName.data(), fileName.size()), startLine,
0);
#else
- nb::object nameObj =
- nb::steal(PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_qualname"));
+ nb::object nameObj = nb::steal(PyObject_GetAttrString(
+ reinterpret_cast<PyObject *>(code), "co_qualname"));
std::string name = nb::cast<std::string>(nameObj);
std::string_view funcName(name);
int startLine, startCol, endLine, endCol;
>From 3d70cc3a25481b9de9e96b3fb170a4406166306e Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 18:01:12 -0500
Subject: [PATCH 03/10] Style
---
mlir/test/lit.cfg.py | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index b20bfe5a4147d..6e922a0954870 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -354,6 +354,9 @@ def find_real_python_interpreter():
else:
config.available_features.add("noasserts")
+if config.enable_python_stable_abi:
+ config.available_features.add("python-stable-abi")
+
if config.expensive_checks:
config.available_features.add("expensive_checks")
@@ -393,5 +396,3 @@ def have_host_jit_feature_support(feature_name):
if sys.version_info >= (3, 11):
config.available_features.add("python-ge-311")
-if config.enable_python_stable_abi:
- config.available_features.add("python-stable-abi")
>From fdd8c64089c8e1a83fde5bfa5558d2e11e34ee37 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 18:07:02 -0500
Subject: [PATCH 04/10] Avoid casts
---
mlir/lib/Bindings/Python/IRCore.cpp | 18 +++---------------
1 file changed, 3 insertions(+), 15 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 48dd5a7f55802..cb87dc9a16b1c 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2732,31 +2732,19 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
for (; pyFrame != nullptr && count < framesLimit;
next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) {
PyCodeObject *code = PyFrame_GetCode(pyFrame);
- // Use attribute access instead of direct struct member access for
- // forward compatibility with the limited (stable) API where
- // PyCodeObject is opaque.
- nb::object fileNameObj = nb::steal(PyObject_GetAttrString(
- reinterpret_cast<PyObject *>(code), "co_filename"));
- auto fileNameStr = nb::cast<std::string>(fileNameObj);
- std::string_view fileName(fileNameStr);
+ std::string_view fileName(PyUnicode_AsUTF8(code->co_filename));
if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName))
continue;
// co_qualname and PyCode_Addr2Location added in py3.11
#if PY_VERSION_HEX < 0x030B00F0
- nb::object nameObj = nb::steal(
- PyObject_GetAttrString(reinterpret_cast<PyObject *>(code), "co_name"));
- std::string name = nb::cast<std::string>(nameObj);
- std::string_view funcName(name);
+ std::string_view funcName(PyUnicode_AsUTF8(code->co_name));
int startLine = PyFrame_GetLineNumber(pyFrame);
MlirLocation loc = mlirLocationFileLineColGet(
ctx, mlirStringRefCreate(fileName.data(), fileName.size()), startLine,
0);
#else
- nb::object nameObj = nb::steal(PyObject_GetAttrString(
- reinterpret_cast<PyObject *>(code), "co_qualname"));
- std::string name = nb::cast<std::string>(nameObj);
- std::string_view funcName(name);
+ std::string_view funcName(PyUnicode_AsUTF8(code->co_qualname));
int startLine, startCol, endLine, endCol;
int lasti = PyFrame_GetLasti(pyFrame);
if (!PyCode_Addr2Location(code, lasti, &startLine, &startCol, &endLine,
>From b6b9dff9590f0ed7ce7825b3d4917fa1032be54e Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 18:11:04 -0500
Subject: [PATCH 05/10] Also check size for possible null chars
---
mlir/lib/Bindings/Python/IRCore.cpp | 13 ++++++++++---
1 file changed, 10 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index cb87dc9a16b1c..6cfb1cad1129d 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2732,19 +2732,26 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
for (; pyFrame != nullptr && count < framesLimit;
next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) {
PyCodeObject *code = PyFrame_GetCode(pyFrame);
- std::string_view fileName(PyUnicode_AsUTF8(code->co_filename));
+ Py_ssize_t fileNameLen;
+ const char *fileNamePtr = PyUnicode_AsUTF8AndSize(code->co_filename,
+ &fileNameLen);
+ std::string_view fileName(fileNamePtr, fileNameLen);
if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName))
continue;
// co_qualname and PyCode_Addr2Location added in py3.11
#if PY_VERSION_HEX < 0x030B00F0
- std::string_view funcName(PyUnicode_AsUTF8(code->co_name));
+ Py_ssize_t nameLen;
+ const char *namePtr = PyUnicode_AsUTF8AndSize(code->co_name, &nameLen);
+ std::string_view funcName(namePtr, nameLen);
int startLine = PyFrame_GetLineNumber(pyFrame);
MlirLocation loc = mlirLocationFileLineColGet(
ctx, mlirStringRefCreate(fileName.data(), fileName.size()), startLine,
0);
#else
- std::string_view funcName(PyUnicode_AsUTF8(code->co_qualname));
+ Py_ssize_t nameLen;
+ const char *namePtr = PyUnicode_AsUTF8AndSize(code->co_qualname, &nameLen);
+ std::string_view funcName(namePtr, nameLen);
int startLine, startCol, endLine, endCol;
int lasti = PyFrame_GetLasti(pyFrame);
if (!PyCode_Addr2Location(code, lasti, &startLine, &startCol, &endLine,
>From 163ecdfb263bc8954ad83fb57a19e4fa767cef96 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 18:15:37 -0500
Subject: [PATCH 06/10] lint
---
mlir/test/lit.cfg.py | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index 6e922a0954870..a716ba0adb480 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -395,4 +395,3 @@ def have_host_jit_feature_support(feature_name):
if sys.version_info >= (3, 11):
config.available_features.add("python-ge-311")
-
>From 23ad2454777d04729e0a415baff750c3791e4602 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 18:18:29 -0500
Subject: [PATCH 07/10] lint
---
mlir/lib/Bindings/Python/IRCore.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 6cfb1cad1129d..eb8b057a6b84c 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2733,8 +2733,8 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) {
PyCodeObject *code = PyFrame_GetCode(pyFrame);
Py_ssize_t fileNameLen;
- const char *fileNamePtr = PyUnicode_AsUTF8AndSize(code->co_filename,
- &fileNameLen);
+ const char *fileNamePtr =
+ PyUnicode_AsUTF8AndSize(code->co_filename, &fileNameLen);
std::string_view fileName(fileNamePtr, fileNameLen);
if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName))
continue;
>From 4581c681600b30d9505a0ab57c5274a8ca8b6c4b Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 27 Feb 2026 18:25:40 -0500
Subject: [PATCH 08/10] Add line number checks
---
mlir/test/python/ir/auto_location_stable_abi.py | 10 ++++++++--
1 file changed, 8 insertions(+), 2 deletions(-)
diff --git a/mlir/test/python/ir/auto_location_stable_abi.py b/mlir/test/python/ir/auto_location_stable_abi.py
index 4d413a8d457ca..5d43c41697e70 100644
--- a/mlir/test/python/ir/auto_location_stable_abi.py
+++ b/mlir/test/python/ir/auto_location_stable_abi.py
@@ -26,7 +26,9 @@ def testStableABILocations():
# Basic op creation produces a callsite location with function names.
op = Operation.create("custom.op1")
- # CHECK: loc(callsite("testStableABILocations"({{.*}}auto_location_stable_abi.py":{{[0-9]+}}:0)
+ # CHECK: loc(callsite("testStableABILocations"({{.*}}auto_location_stable_abi.py":28:0)
+ # CHECK-SAME: at callsite("run"({{.*}}auto_location_stable_abi.py":16:0)
+ # CHECK-SAME: at "<module>"({{.*}}auto_location_stable_abi.py":22:0))))
print(op.location)
# Nested function calls produce nested callsite locations.
@@ -34,7 +36,11 @@ def inner():
return arith.constant(IndexType.get(), 1)
val = inner()
- # CHECK: loc(callsite({{.*}} at callsite("testStableABILocations.<locals>.inner"({{.*}}auto_location_stable_abi.py":{{[0-9]+}}:0)
+ # CHECK: loc(callsite(
+ # CHECK-SAME: "testStableABILocations.<locals>.inner"({{.*}}auto_location_stable_abi.py":36:0)
+ # CHECK-SAME: at callsite("testStableABILocations"({{.*}}auto_location_stable_abi.py":38:0)
+ # CHECK-SAME: at callsite("run"({{.*}}auto_location_stable_abi.py":16:0)
+ # CHECK-SAME: at "<module>"({{.*}}auto_location_stable_abi.py":22:0)))))))
print(val.location)
# Frame limit of 0 produces unknown location.
>From 28e306068321878616ab1ee6949bc5423f759200 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Sat, 28 Feb 2026 08:23:07 -0500
Subject: [PATCH 09/10] Disable autolocation under stable abi
---
mlir/lib/Bindings/Python/IRCore.cpp | 43 +++----------------
.../python/ir/auto_location_stable_abi.py | 37 +++-------------
2 files changed, 11 insertions(+), 69 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index eb8b057a6b84c..b0acda6001455 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2676,6 +2676,11 @@ namespace {
using namespace mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN;
MlirLocation tracebackToLocation(MlirContext ctx) {
+#if defined(Py_LIMITED_API)
+ // Frame introspection C APIs are not available under the limited API.
+ // Traceback-based auto-location is not supported; return unknown.
+ return mlirLocationUnknownGet(ctx);
+#else
size_t framesLimit =
PyGlobals::get().getTracebackLoc().locTracebackFramesLimit();
// Use a thread_local here to avoid requiring a large amount of space.
@@ -2685,42 +2690,6 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
nb::gil_scoped_acquire acquire;
-#if defined(Py_LIMITED_API)
- // Under the limited API (targeting Python 3.12+), most frame introspection
- // C APIs (PyCode_Addr2Location, PyFrame_GetLasti, etc.) are not available.
- // Use Python-level sys._getframe() and attribute access instead. Note:
- // co_qualname (used below) requires Python 3.11+.
- nb::object sys_mod = nb::module_::import_("sys");
- nb::object frameObj;
- try {
- // _getframe(0) from C++ returns the topmost Python frame, equivalent
- // to PyThreadState_GetFrame() in the non-limited-API path.
- frameObj = sys_mod.attr("_getframe")(0);
- } catch (nb::python_error &) {
- return mlirLocationUnknownGet(ctx);
- }
-
- while (!frameObj.is_none() && frameObj.ptr() != nullptr &&
- count < framesLimit) {
- nb::object codeObj = frameObj.attr("f_code");
- auto fileNameStr = nb::cast<std::string>(codeObj.attr("co_filename"));
- std::string_view fileName(fileNameStr);
- if (PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName)) {
- std::string name = nb::cast<std::string>(codeObj.attr("co_qualname"));
- std::string_view funcName(name);
- int startLine = nb::cast<int>(frameObj.attr("f_lineno"));
- MlirLocation loc = mlirLocationFileLineColGet(
- ctx, mlirStringRefCreate(fileName.data(), fileName.size()), startLine,
- 0);
- frames[count] = mlirLocationNameGet(
- ctx, mlirStringRefCreate(funcName.data(), funcName.size()), loc);
- ++count;
- }
- frameObj = frameObj.attr("f_back");
- if (frameObj.is_none())
- break;
- }
-#else
PyThreadState *tstate = PyThreadState_GET();
PyFrameObject *next;
PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate);
@@ -2770,7 +2739,6 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
// When the loop breaks (after the last iter), current frame (if non-null)
// is leaked without this.
Py_XDECREF(pyFrame);
-#endif
if (count == 0)
return mlirLocationUnknownGet(ctx);
@@ -2786,6 +2754,7 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
caller = mlirLocationCallSiteGet(frames[i], caller);
return mlirLocationCallSiteGet(callee, caller);
+#endif
}
PyLocation
diff --git a/mlir/test/python/ir/auto_location_stable_abi.py b/mlir/test/python/ir/auto_location_stable_abi.py
index 5d43c41697e70..a55b9ffe68c0f 100644
--- a/mlir/test/python/ir/auto_location_stable_abi.py
+++ b/mlir/test/python/ir/auto_location_stable_abi.py
@@ -1,14 +1,11 @@
# RUN: %PYTHON %s | FileCheck %s
# REQUIRES: python-stable-abi
#
-# Verify that traceback-based locations work under the stable ABI (abi3).
-# The limited API path uses sys._getframe() and cannot provide column info,
-# so locations have :0 for columns. This test checks that function names,
-# file paths, callsite nesting, and frame limiting all work correctly.
+# Verify that traceback-based auto-location is not supported under the stable
+# ABI (abi3) and always produces unknown locations.
import gc
from mlir.ir import *
-from mlir.dialects import arith
def run(f):
@@ -18,35 +15,11 @@ def run(f):
assert Context._get_live_count() == 0
-# CHECK-LABEL: TEST: testStableABILocations
+# CHECK-LABEL: TEST: testAutoLocationIsUnknown
@run
-def testStableABILocations():
+def testAutoLocationIsUnknown():
with Context() as ctx, loc_tracebacks():
ctx.allow_unregistered_dialects = True
-
- # Basic op creation produces a callsite location with function names.
op = Operation.create("custom.op1")
- # CHECK: loc(callsite("testStableABILocations"({{.*}}auto_location_stable_abi.py":28:0)
- # CHECK-SAME: at callsite("run"({{.*}}auto_location_stable_abi.py":16:0)
- # CHECK-SAME: at "<module>"({{.*}}auto_location_stable_abi.py":22:0))))
- print(op.location)
-
- # Nested function calls produce nested callsite locations.
- def inner():
- return arith.constant(IndexType.get(), 1)
-
- val = inner()
- # CHECK: loc(callsite(
- # CHECK-SAME: "testStableABILocations.<locals>.inner"({{.*}}auto_location_stable_abi.py":36:0)
- # CHECK-SAME: at callsite("testStableABILocations"({{.*}}auto_location_stable_abi.py":38:0)
- # CHECK-SAME: at callsite("run"({{.*}}auto_location_stable_abi.py":16:0)
- # CHECK-SAME: at "<module>"({{.*}}auto_location_stable_abi.py":22:0)))))))
- print(val.location)
-
- # Frame limit of 0 produces unknown location.
- from mlir.dialects._ods_common import _cext
-
- _cext.globals.set_loc_tracebacks_frame_limit(0)
- val2 = arith.constant(IndexType.get(), 2)
# CHECK: loc(unknown)
- print(val2.location)
+ print(op.location)
>From c782c259b26c191f53a96c89806d8593342ec8de Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Sat, 28 Feb 2026 14:13:32 -0500
Subject: [PATCH 10/10] Cache imports
---
mlir/include/mlir/Bindings/Python/IRCore.h | 7 ++-
.../mlir/Bindings/Python/NanobindAdaptors.h | 43 +++----------------
.../mlir/Bindings/Python/NanobindUtils.h | 37 ++++++++++++++++
3 files changed, 47 insertions(+), 40 deletions(-)
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index cdd99b52f5200..5953f26d07370 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1876,8 +1876,11 @@ MLIR_PYTHON_API_EXPORTED void populateRoot(nanobind::module_ &m);
template <class Func, typename... Args>
inline nanobind::object classmethod(Func f, Args... args) {
nanobind::object cf = nanobind::cpp_function(f, args...);
- nanobind::object builtins = nanobind::module_::import_("builtins");
- return builtins.attr("classmethod")(cf);
+ static SafeInit<nanobind::object> classmethodFn([]() {
+ return std::make_unique<nanobind::object>(
+ nanobind::module_::import_("builtins").attr("classmethod"));
+ });
+ return classmethodFn.get()(cf);
}
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 95151bdaf8ca6..6669433550a00 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -19,7 +19,6 @@
#ifndef MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
#define MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
-#include <atomic>
#include <cstdint>
#include <memory>
#include <optional>
@@ -36,41 +35,6 @@ namespace mlir {
namespace python {
namespace {
-// Safely calls Python initialization code on first use, avoiding deadlocks.
-template <typename T>
-class SafeInit {
-public:
- typedef std::unique_ptr<T> (*F)();
-
- explicit SafeInit(F init_fn) : initFn(init_fn) {}
-
- T &get() {
- if (T *result = output.load()) {
- return *result;
- }
-
- // Note: init_fn() may be called multiple times if, for example, the GIL is
- // released during its execution. The intended use case is for module
- // imports which are safe to perform multiple times. We are careful not to
- // hold a lock across init_fn() to avoid lock ordering problems.
- std::unique_ptr<T> m = initFn();
- {
- nanobind::ft_lock_guard lock(mu);
- if (T *result = output.load()) {
- return *result;
- }
- T *p = m.release();
- output.store(p);
- return *p;
- }
- }
-
-private:
- nanobind::ft_mutex mu;
- std::atomic<T *> output{nullptr};
- F initFn;
-};
-
nanobind::module_ &irModule() {
static SafeInit<nanobind::module_> init([]() {
return std::make_unique<nanobind::module_>(
@@ -496,8 +460,11 @@ class pure_subclass {
std::forward<Func>(f),
nanobind::name(name), // nanobind::scope(thisClass),
extra...);
- nanobind::object builtins = nanobind::module_::import_("builtins");
- thisClass.attr(name) = builtins.attr("classmethod")(cf);
+ static SafeInit<nanobind::object> classmethodFn([]() {
+ return std::make_unique<nanobind::object>(
+ nanobind::module_::import_("builtins").attr("classmethod"));
+ });
+ thisClass.attr(name) = classmethodFn.get()(cf);
return *this;
}
diff --git a/mlir/include/mlir/Bindings/Python/NanobindUtils.h b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
index 15ed8c73190eb..5bffc2b5983d4 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindUtils.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
@@ -13,7 +13,9 @@
#include "mlir-c/Support.h"
#include "mlir/Bindings/Python/Nanobind.h"
+#include <atomic>
#include <fstream>
+#include <memory>
#include <sstream>
#include <string>
#include <string_view>
@@ -33,6 +35,41 @@ struct std::iterator_traits<nanobind::detail::fast_iterator> {
namespace mlir {
namespace python {
+/// Safely calls Python initialization code on first use, avoiding deadlocks.
+template <typename T>
+class SafeInit {
+public:
+ typedef std::unique_ptr<T> (*F)();
+
+ explicit SafeInit(F init_fn) : initFn(init_fn) {}
+
+ T &get() {
+ if (T *result = output.load()) {
+ return *result;
+ }
+
+ // Note: init_fn() may be called multiple times if, for example, the GIL is
+ // released during its execution. The intended use case is for module
+ // imports which are safe to perform multiple times. We are careful not to
+ // hold a lock across init_fn() to avoid lock ordering problems.
+ std::unique_ptr<T> m = initFn();
+ {
+ nanobind::ft_lock_guard lock(mu);
+ if (T *result = output.load()) {
+ return *result;
+ }
+ T *p = m.release();
+ output.store(p);
+ return *p;
+ }
+ }
+
+private:
+ nanobind::ft_mutex mu;
+ std::atomic<T *> output{nullptr};
+ F initFn;
+};
+
struct MlirTypeIDHash {
size_t operator()(MlirTypeID typeID) const {
return mlirTypeIDHashValue(typeID);
More information about the Mlir-commits
mailing list