[Python] Add Stable ABI (abi3) support for Python 3.12+ (#23600)

Build Python bindings against the Limited C API (`Py_LIMITED_API`, PEP
384) so that a single wheel built with Python 3.12 works on all future
CPython versions. This reduces the total release artifact size.

The Limited API forbids direct access to CPython struct layouts and most
convenience macros. The main impact is replacing the NumPy C API (whose
headers use non-limited internals) with Python-level equivalents like
`numpy.frombuffer()` and `numpy.dtype()`, and using nanobind's
`type_slots()` to register buffer/sequence/mapping protocols instead of
casting to `PyHeapTypeObject`.

Limitations:
- Only CPython 3.12+. Per-version wheels are still built for 3.10-3.11.
- Free-threaded builds (3.13t) are excluded: the free-threaded ABI is
not yet stable.
- `Development.SABIModule` requires CMake 3.26+ (IREE minimum is 3.21),
so `IREE_ENABLE_PYTHON_STABLE_ABI` is OFF by default. setup.py
auto-enables it when building with CPython 3.12+. May want to bump it in
the future.

The existing release workflow already builds with Python 3.12 on all
platforms, so abi3 wheels are produced with no workflow changes. pkgci
only builds cp311 by default (no abi3).

The LLVM submodule includes cherry-picks of llvm/llvm-project PRs for
MLIR-side abi3 support.

Assisted-by: claude

Author: kuhar

CMakeLists.txt

@@ -726,6 +726,13 @@ include(iree_setup_toolchain)
 # Otherwise, for features that just require the interpreter, find that alone.
 #-------------------------------------------------------------------------------
 
+# Build Python bindings with Python Stable ABI (abi3) for Python 3.12+.
+# When enabled, a single extension module can be used across Python 3.12+.
+# nanobind handles the interaction with FREE_THREADED: if the interpreter
+# is free-threaded, STABLE_ABI is ignored and the free-threaded ABI is used.
+# Declared here (before first use) and also cascaded to MLIR below.
+option(IREE_ENABLE_PYTHON_STABLE_ABI "Build Python bindings with Stable ABI (abi3) for Python 3.12+" OFF)
+
 if(IREE_BUILD_PYTHON_BINDINGS)
   # After CMake 3.18, we are able to limit the scope of the search to just
   # Development.Module. Searching for Development will fail in situations where
@@ -739,13 +746,19 @@ if(IREE_BUILD_PYTHON_BINDINGS)
   # See: https://reviews.llvm.org/D118148
   # If building Python packages, we have a hard requirement on 3.10+.
   find_package(Python3 3.10 COMPONENTS Interpreter Development NumPy)
-  find_package(Python3 3.10 COMPONENTS Interpreter Development.Module NumPy REQUIRED)
+  # Development.SABIModule is needed for Stable ABI (abi3) builds. CMake 3.26+
+  # provides it; on older CMake it is silently ignored if not found.
+  set(_PYTHON_SABI_COMPONENT "")
+  if(IREE_ENABLE_PYTHON_STABLE_ABI)
+    set(_PYTHON_SABI_COMPONENT Development.SABIModule)
+  endif()
+  find_package(Python3 3.10 COMPONENTS Interpreter Development.Module NumPy ${_PYTHON_SABI_COMPONENT} REQUIRED)
   # Some parts of the build use FindPython instead of FindPython3. Why? No
   # one knows, but they are different. So make sure to bootstrap this one too.
   # Not doing this here risks them diverging, which on multi-Python systems,
   # can be troublesome. Note that nanobind requires FindPython.
   set(Python_EXECUTABLE "${Python3_EXECUTABLE}")
-  find_package(Python 3.10 COMPONENTS Interpreter Development.Module NumPy REQUIRED)
+  find_package(Python 3.10 COMPONENTS Interpreter Development.Module NumPy ${_PYTHON_SABI_COMPONENT} REQUIRED)
 elseif(IREE_BUILD_COMPILER OR IREE_BUILD_TESTS)
   find_package(Python3 COMPONENTS Interpreter REQUIRED)
   set(Python_EXECUTABLE "${Python3_EXECUTABLE}")
@@ -870,6 +883,11 @@ endif()
 # MLIR/LLVM Dependency
 #-------------------------------------------------------------------------------
 
+# Cascade IREE_ENABLE_PYTHON_STABLE_ABI to MLIR.
+if(IREE_ENABLE_PYTHON_STABLE_ABI)
+  set(MLIR_ENABLE_PYTHON_STABLE_ABI ON CACHE BOOL "" FORCE)
+endif()
+
 # Both the IREE and MLIR Python bindings require nanobind. We initialize it here
 # at the top level so that everything uses ours consistently.
 if(IREE_BUILD_PYTHON_BINDINGS)

build_tools/python_deploy/build_linux_packages.sh

@@ -65,7 +65,9 @@ this_dir="$(cd $(dirname $0) && pwd)"
 script_name="$(basename $0)"
 repo_root=$(cd "${this_dir}" && find_git_dir_parent)
 manylinux_docker_image="${manylinux_docker_image:-$(uname -m | awk '{print ($1 == "aarch64") ? "quay.io/pypa/manylinux_2_28_aarch64" : "ghcr.io/iree-org/manylinux_x86_64@sha256:2e0246137819cf10ed84240a971f9dd75cc3eb62dc6907dfd2080ee966b3c9f4" }')}"
-python_versions="${override_python_versions:-cp310-cp310 cp311-cp311 cp312-cp312 cp313-cp313 cp313-cp313t}"
+# Python versions to build. cp312 is the abi3 build (produces a wheel
+# compatible with 3.12+). Per-version builds are used for <3.12 and 3.13t.
+python_versions="${override_python_versions:-cp310-cp310 cp311-cp311 cp312-cp312 cp313-cp313t}"
 output_dir="${output_dir:-${this_dir}/wheelhouse}"
 packages="${packages:-iree-base-runtime iree-base-compiler}"
 package_suffix="${package_suffix:-}"
@@ -173,8 +175,14 @@ function build_iree_compiler() {
 function run_audit_wheel() {
   local wheel_basename="$1"
   local python_version="$2"
+  # For abi3 builds (cp312 non-free-threaded), the wheel tag is cp312-abi3
+  # instead of cp312-cp312.
+  local wheel_tag="${python_version}"
+  if [[ "${python_version}" == "cp312-cp312" ]]; then
+    wheel_tag="cp312-abi3"
+  fi
   # Force wildcard expansion here
-  generic_wheel="$(echo "${output_dir}/${wheel_basename}-"*"-${python_version}-linux_$(uname -m).whl")"
+  generic_wheel="$(echo "${output_dir}/${wheel_basename}-"*"-${wheel_tag}-linux_$(uname -m).whl")"
   ls "${generic_wheel}"
   echo ":::: Auditwheel ${generic_wheel}"
   auditwheel repair -w "${output_dir}" "${generic_wheel}"
@@ -186,6 +194,10 @@ function clean_wheels() {
   local python_version="$2"
   echo ":::: Clean wheels ${wheel_basename} ${python_version}"
   rm -f -v "${output_dir}/${wheel_basename}-"*"-${python_version}-"*".whl"
+  # Also clean abi3 wheels for cp312.
+  if [[ "${python_version}" == "cp312-cp312" ]]; then
+    rm -f -v "${output_dir}/${wheel_basename}-"*"-cp312-abi3-"*".whl"
+  fi
 }
 
 function prepare_python() {

build_tools/python_deploy/build_macos_packages.sh

@@ -21,7 +21,9 @@ set -eu -o errtrace
 
 this_dir="$(cd $(dirname $0) && pwd)"
 repo_root="$(cd $this_dir/../../ && pwd)"
-python_versions="${override_python_versions:-3.11}"
+# Python versions to build. 3.12 produces an abi3 wheel (compatible with 3.12+).
+# Per-version builds are used for <3.12 and free-threaded builds.
+python_versions="${override_python_versions:-3.11 3.12}"
 output_dir="${output_dir:-${this_dir}/wheelhouse}"
 packages="${packages:-iree-base-runtime iree-base-compiler}"
 

build_tools/python_deploy/build_windows_packages.ps1

@@ -9,7 +9,9 @@
 
 # Configure settings with script parameters.
 param(
-    [array]$python_versions=@("3.11"),
+    # Python versions to build. 3.12 produces an abi3 wheel (compatible with 3.12+).
+    # Per-version builds are used for <3.12 and free-threaded builds.
+    [array]$python_versions=@("3.11", "3.12"),
     [array]$packages=@("iree-base-runtime", "iree-base-compiler"),
     [System.String]$output_dir
 )

compiler/setup.py

@@ -41,6 +41,12 @@
 from setuptools.command.build_py import build_py as _build_py
 from setuptools.command.egg_info import egg_info
 
+# Detect whether we should build an abi3 (Stable ABI) wheel.
+# This applies to CPython 3.12+ when not in free-threaded mode.
+_is_abi3_build = sys.version_info >= (3, 12) and not sysconfig.get_config_var(
+    "Py_GIL_DISABLED"
+)
+
 
 def check_pip_version():
     from packaging import version
@@ -269,6 +275,8 @@ def prepare_installation():
             get_env_cmake_option("IREE_TARGET_BACKEND_CUDA", "OFF"),
             get_env_cmake_option("IREE_ENABLE_LLD", "OFF"),
         ]
+        if _is_abi3_build:
+            cmake_args.append("-DIREE_ENABLE_PYTHON_STABLE_ABI=ON")
         cmake_args.extend(get_cmake_version_info_args())
 
         # These usually flow through the environment, but we add them explicitly
@@ -382,6 +390,24 @@ def run(self):
         egg_info.run(self)
 
 
+# Override bdist_wheel to produce abi3 wheel tags when applicable.
+_bdist_wheel_cmdclass = {}
+try:
+    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+
+    class bdist_wheel(_bdist_wheel):
+        def get_tag(self):
+            python, abi, plat = _bdist_wheel.get_tag(self)
+            if _is_abi3_build:
+                python, abi = "cp312", "abi3"
+            return python, abi, plat
+
+    _bdist_wheel_cmdclass = {"bdist_wheel": bdist_wheel}
+except ImportError:
+    # wheel package not available (e.g., during sdist). Not an error.
+    pass
+
+
 def generate_version_py():
     return f"""# Auto-generated version info.
 PACKAGE_SUFFIX = "{PACKAGE_SUFFIX}"
@@ -485,12 +511,15 @@ def find_git_submodule_revision(submodule_path):
         CMakeExtension("iree.compiler._mlir_libs._mlirGPUPasses"),
         CMakeExtension("iree.compiler._mlir_libs._site_initialize_0"),
     ],
-    cmdclass={
-        "build": CustomBuild,
-        "built_ext": NoopBuildExtension,
-        "build_py": CMakeBuildPy,
-        "egg_info": CleanEggInfo,
-    },
+    cmdclass=dict(
+        {
+            "build": CustomBuild,
+            "built_ext": NoopBuildExtension,
+            "build_py": CMakeBuildPy,
+            "egg_info": CleanEggInfo,
+        },
+        **_bdist_wheel_cmdclass,
+    ),
     zip_safe=False,
     package_dir={
         # Note: Must be relative path, so we line this up with the absolute

runtime/bindings/python/CMakeLists.txt

@@ -42,8 +42,13 @@ iree_select_compiler_opts(_RTTI_AND_EXCEPTION_COPTS
     "/GR"
 )
 
+set(_STABLE_ABI_FLAG "")
+if(IREE_ENABLE_PYTHON_STABLE_ABI)
+  set(_STABLE_ABI_FLAG STABLE_ABI)
+endif()
+
 nanobind_add_module(iree_runtime_bindings_python_PyExtRt
-  NB_STATIC LTO FREE_THREADED
+  NB_STATIC LTO FREE_THREADED ${_STABLE_ABI_FLAG}
   "binding.h"
   "initialize_module.cc"
   "invoke.h"
@@ -87,8 +92,6 @@ target_link_libraries(iree_runtime_bindings_python_PyExtRt
   iree::tooling::modules
   iree::vm
   iree::vm::bytecode::module
-
-  Python::NumPy
 )
 
 target_compile_options(iree_runtime_bindings_python_PyExtRt

runtime/bindings/python/binding.h

@@ -58,7 +58,15 @@ class ApiRefCounted {
   }
   void operator=(const ApiRefCounted&) = delete;
 
-  ~ApiRefCounted() { Release(); }
+  ~ApiRefCounted() {
+    // In stable ABI (abi3) mode, types are heap-allocated via PyType_FromSpec
+    // and instances may be destroyed during Py_FinalizeEx after IREE's type
+    // registry is no longer valid. Skip release when nanobind's internals
+    // are being torn down since the process is exiting anyway.
+    if (instance_ && py::is_alive()) {
+      Release();
+    }
+  }
 
   // Steals the reference to the object referenced by the given raw pointer and
   // returns a wrapper (transfers ownership).
@@ -130,37 +138,39 @@ inline py::object create_empty_tuple() {
   return py::steal(py::handle(PyTuple_New(0)));
 }
 
-// For a bound class, binds the buffer protocol. This will result in a call
-// to on the CppType:
-//   HandleBufferProtocol(Py_buffer *view, int flags)
-// This is a low level callback and must not raise any exceptions. If
-// error conditions are warranted the usual PyErr_SetString approach must be
-// used (and -1 returned). Return 0 on success.
+// Returns a nanobind::type_slots() descriptor for the buffer protocol on
+// CppType. CppType must implement HandleBufferProtocol(Py_buffer*, int).
+// The slots are only read during py::class_ construction, so the static local
+// array inside is fine (no persistent global state).
 template <typename CppType>
-void BindBufferProtocol(py::handle clazz) {
-  PyBufferProcs buffer_procs;
-  memset(&buffer_procs, 0, sizeof(buffer_procs));
-  buffer_procs.bf_getbuffer =
-      // It is not legal to raise exceptions from these callbacks.
-      +[](PyObject* raw_self, Py_buffer* view, int flags) -> int {
-    if (view == NULL) {
-      PyErr_SetString(PyExc_ValueError, "NULL view in getbuffer");
-      return -1;
-    }
-
-    // Cast must succeed due to invariants.
-    auto self = py::cast<CppType*>(py::handle(raw_self));
-
-    Py_INCREF(raw_self);
-    view->obj = raw_self;
-    return self->HandleBufferProtocol(view, flags);
-  };
-  buffer_procs.bf_releasebuffer =
-      +[](PyObject* raw_self, Py_buffer* view) -> void {};
-  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_buffer = buffer_procs;
+py::type_slots buffer_protocol_slots() {
+  // It is not legal to raise exceptions from buffer protocol callbacks.
+  static const PyType_Slot slots[] = {
+      {Py_bf_getbuffer,
+       reinterpret_cast<void*>(
+           +[](PyObject* raw_self, Py_buffer* view, int flags) -> int {
+             if (!view) {
+               PyErr_SetString(PyExc_ValueError, "NULL view in getbuffer");
+               return -1;
+             }
+
+             // Cast must succeed due to invariants.
+             auto self = py::cast<CppType*>(py::handle(raw_self));
+
+             int rc = self->HandleBufferProtocol(view, flags);
+             if (rc == 0) {
+               Py_INCREF(raw_self);
+               view->obj = raw_self;
+             }
+             return rc;
+           })},
+      // No-op: PyBuffer_Release handles Py_DECREF(view->obj) after calling
+      // this callback, so the Py_INCREF in getbuffer is already balanced.
+      {Py_bf_releasebuffer,
+       reinterpret_cast<void*>(
+           +[](PyObject* raw_self, Py_buffer* view) -> void {})},
+      {0, nullptr}};
+  return py::type_slots(slots);
 }
 
 // Nanobind 2.0 had a backwards compatibility bug where it left out the

runtime/bindings/python/hal.cc

@@ -1238,6 +1238,24 @@ int HalModuleDebugSinkTpClear(PyObject* self) {
   return 0;
 }
 
+//------------------------------------------------------------------------------
+// HalMappedMemory buffer protocol
+//------------------------------------------------------------------------------
+
+int HalMappedMemory::HandleBufferProtocol(Py_buffer* view, int flags) {
+  view->buf = mapped_memory_.contents.data;
+  view->len = mapped_memory_.contents.data_length;
+  view->readonly = 1;  // Mapped with IREE_HAL_MEMORY_ACCESS_READ
+  view->itemsize = 1;
+  view->format = (char*)"B";
+  view->ndim = 1;
+  view->shape = nullptr;
+  view->strides = nullptr;
+  view->suboffsets = nullptr;
+  view->internal = nullptr;
+  return 0;
+}
+
 //------------------------------------------------------------------------------
 // Bindings
 //------------------------------------------------------------------------------
@@ -1344,9 +1362,7 @@ void SetupHalBindings(nanobind::module_ m) {
   hal_element_type
       .def_static("map_to_dtype",
                   [](iree_hal_element_type_t element_type) {
-                    int typenum = numpy::ConvertHalElementTypeToNumPyTypeNum(
-                        element_type);
-                    return numpy::DescrNewFromType(typenum);
+                    return numpy::DescrNewFromType(element_type);
                   })
       .def_static("is_byte_aligned",
                   [](iree_hal_element_type_t element_type) {
@@ -1856,7 +1872,8 @@ void SetupHalBindings(nanobind::module_ m) {
           py::arg("timeout") = py::none(), py::arg("deadline") = py::none(),
           kHalWait);
 
-  py::class_<HalMappedMemory>(m, "MappedMemory")
+  py::class_<HalMappedMemory>(m, "MappedMemory",
+                              buffer_protocol_slots<HalMappedMemory>())
       .def(
           "asarray",
           [](HalMappedMemory* self, py::handle shape, py::object dtype_descr) {
@@ -1867,8 +1884,7 @@ void SetupHalBindings(nanobind::module_ m) {
             for (size_t i = 0; i < rank; ++i) {
               dims[i] = py::cast<intptr_t>(shape[i]);
             }
-            int typenum = numpy::TypenumFromDescr(dtype_descr);
-            return numpy::SimpleNewFromData(rank, dims, typenum,
+            return numpy::SimpleNewFromData(rank, dims, dtype_descr,
                                             self->mapped_memory().contents.data,
                                             py_mapped_memory);
           },

runtime/bindings/python/hal.h

@@ -286,6 +286,9 @@ class HalMappedMemory {
 
   iree_hal_buffer_mapping_t& mapped_memory() { return mapped_memory_; }
 
+  // Buffer protocol for PyMemoryView_FromObject (used by SimpleNewFromData).
+  int HandleBufferProtocol(Py_buffer* view, int flags);
+
  private:
   iree_hal_buffer_mapping_t mapped_memory_ = {{0}};
   iree_hal_buffer_t* buffer_ = nullptr;

runtime/bindings/python/initialize_module.cc

@@ -14,7 +14,6 @@
 #include "./invoke.h"
 #include "./io.h"
 #include "./loop.h"
-#include "./numpy_interop.h"
 #include "./py_module.h"
 #include "./status_utils.h"
 #include "./vm.h"
@@ -33,7 +32,6 @@ namespace iree {
 namespace python {
 
 NB_MODULE(_runtime, m) {
-  numpy::InitializeNumPyInterop();
   IREE_TRACE_APP_ENTER();
 
   IREE_CHECK_OK(iree_hal_register_all_available_drivers(