Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .github/workflows/build-and-test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,11 @@ jobs:
submodules: recursive
- name: Build wheels
uses: pypa/cibuildwheel@v3.2.0
- 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 }}
Expand Down
3 changes: 3 additions & 0 deletions .gitmodules
Original file line number Diff line number Diff line change
Expand Up @@ -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
9 changes: 6 additions & 3 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
2 changes: 1 addition & 1 deletion cmake/utils.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
31 changes: 21 additions & 10 deletions onnxoptimizer/cpp2py_export.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
14 changes: 13 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
]

################################################################################
Expand Down Expand Up @@ -348,4 +359,5 @@ def run(self):
'onnxoptimizer=onnxoptimizer:main',
],
},
options=setup_opts,
)
1 change: 1 addition & 0 deletions third_party/nanobind
Submodule nanobind added at ab3456