Support ABI3

Signed-off-by: take-cheeze <[email protected]>

Author: take-cheeze

.github/workflows/build-and-test.yml

@@ -46,6 +46,11 @@ jobs:
         submodules: recursive
     - name: Build wheels
       uses: pypa/[email protected]
+    - name: Check ABI3
+      run: |
+        pip install abi3audit
+        abi3audit --strict -v --report ./wheelhouse/*.whl
+      if: ${{ contains(fromJSON('["cp312", "cp313"]'), matrix.python) }}
     - uses: actions/upload-artifact@v4
       with:
         name: artifact-${{ matrix.os }}-${{ matrix.python }}

.gitmodules

@@ -4,3 +4,6 @@
 [submodule "third_party/protobuf"]
 	path = third_party/protobuf
 	url = https://github.com/protocolbuffers/protobuf.git
+[submodule "third_party/nanobind"]
+	path = third_party/nanobind
+	url = https://github.com/wjakob/nanobind.git

CMakeLists.txt

@@ -67,9 +67,12 @@ target_include_directories(onnx_optimizer_c_api PUBLIC
     )
 
 if(ONNX_BUILD_PYTHON)
-  find_package(Python3 REQUIRED COMPONENTS Interpreter Development.Module)
-  find_package(pybind11)
-  pybind11_add_module(onnx_opt_cpp2py_export "onnxoptimizer/cpp2py_export.cc")
+  find_package(
+    Python 3
+    REQUIRED COMPONENTS Interpreter Development.Module
+    OPTIONAL_COMPONENTS Development.SABIModule)
+  add_subdirectory(third_party/nanobind EXCLUDE_FROM_ALL)
+  nanobind_add_module(onnx_opt_cpp2py_export "onnxoptimizer/cpp2py_export.cc" STABLE_ABI)
   target_link_libraries(onnx_opt_cpp2py_export PRIVATE onnx_optimizer)
 endif()
 

cmake/utils.cmake

@@ -3,7 +3,7 @@ include(${PROJECT_SOURCE_DIR}/third_party/onnx/cmake/Utils.cmake)
 # Poor man's FetchContent
 function(add_subdirectory_if_no_target dir target)
   if (NOT TARGET ${target})
-    add_subdirectory(${dir})
+    add_subdirectory(${dir} EXCLUDE_FROM_ALL)
   endif()
 endfunction()
 

onnxoptimizer/cpp2py_export.cc

@@ -2,39 +2,50 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/vector.h>
 
-#include "onnx/py_utils.h"
 #include "onnxoptimizer/model_util.h"
 #include "onnxoptimizer/optimize.h"
 
 namespace ONNX_NAMESPACE {
-namespace py = pybind11;
-using namespace pybind11::literals;
-PYBIND11_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
+namespace nb = nanobind;
+using namespace nanobind::literals;
+
+template <typename Proto>
+bool ParseProtoFromPyBytes(Proto* proto, const nb::bytes& bytes) {
+  // Get the buffer from Python bytes object
+  char* buffer = nullptr;
+  Py_ssize_t length = 0;
+  PyBytes_AsStringAndSize(bytes.ptr(), &buffer, &length);
+
+  return ParseProtoFromBytes(proto, buffer, length);
+}
+
+NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
   onnx_opt_cpp2py_export.doc() = "ONNX Optimizer";
 
   onnx_opt_cpp2py_export.def(
       "optimize",
-      [](const py::bytes& bytes, const std::vector<std::string>& names) {
+      [](const nb::bytes& bytes, const std::vector<std::string>& names) {
         ModelProto proto{};
         ParseProtoFromPyBytes(&proto, bytes);
         auto const result = optimization::Optimize(proto, names);
         std::string out;
         result.SerializeToString(&out);
-        return py::bytes(out);
+        return nb::bytes(out.data(), out.size());
       });
 
   onnx_opt_cpp2py_export.def(
       "optimize_fixedpoint",
-      [](const py::bytes& bytes, const std::vector<std::string>& names) {
+      [](const nb::bytes& bytes, const std::vector<std::string>& names) {
         ModelProto proto{};
         ParseProtoFromPyBytes(&proto, bytes);
         auto const result = optimization::OptimizeFixed(proto, names);
         std::string out;
         result.SerializeToString(&out);
-        return py::bytes(out);
+        return nb::bytes(out.data(), out.size());
       });
 
   onnx_opt_cpp2py_export.def(

setup.py

@@ -288,10 +288,21 @@ def run(self):
 # Extensions
 ################################################################################
 
+py_limited_api = sys.version_info[0] >= 3 and sys.version_info[1] >= 12
+if py_limited_api:
+    setup_opts = {
+        'bdist_wheel': {
+            'py_limited_api': 'cp312'
+        },
+    }
+else:
+    setup_opts = {}
+
 ext_modules = [
     setuptools.Extension(
         name=str('onnxoptimizer.onnx_opt_cpp2py_export'),
-        sources=[])
+        sources=[],
+        py_limited_api=py_limited_api)
 ]
 
 ################################################################################
@@ -348,4 +359,5 @@ def run(self):
             'onnxoptimizer=onnxoptimizer:main',
         ],
     },
+    options=setup_opts,
 )