Skip to content

Commit

Permalink
Migrate from pybind11 to nanobind
Browse files Browse the repository at this point in the history
  • Loading branch information
calcmogul committed Jul 11, 2024
1 parent b6ffa2d commit c5da740
Show file tree
Hide file tree
Showing 28 changed files with 1,062 additions and 758 deletions.
3 changes: 3 additions & 0 deletions .github/workflows/build-python.yml
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@ jobs:
with:
python-version: ${{ matrix.version }}

- run: pip install typing_extensions
if: matrix.version == '3.9' || matrix.version == '3.10'

- run: python3 ./tools/update_version.py
- run: pip3 install build pytest
- run: ${{ matrix.cmake-env }} python3 -m build --wheel
Expand Down
3 changes: 1 addition & 2 deletions .github/workflows/website.yml
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,7 @@ jobs:
- run: pip3 install \
build \
py-build-cmake \
pybind11 \
pybind11-stubgen \
nanobind \
pybind11-mkdoc \
clang

Expand Down
2 changes: 1 addition & 1 deletion .styleguide
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,6 @@ modifiableFileExclude {
includeOtherLibs {
^Eigen/
^catch2/
^pybind11/
^nanobind/
^sleipnir/
}
77 changes: 50 additions & 27 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ include(FetchContent)

# Options for using a package manager (e.g., vcpkg) for certain dependencies
option(USE_SYSTEM_EIGEN "Use system eigen" OFF)
option(USE_SYSTEM_PYBIND "Use system pybind" OFF)
option(USE_SYSTEM_NANOBIND "Use system nanobind" OFF)

# Required for std::async()
find_package(Threads REQUIRED)
Expand Down Expand Up @@ -276,23 +276,29 @@ if(BUILD_PYTHON)
set(PY_DEST lib/python${Python_VERSION_MAJOR}.${Python_VERSION_MINOR})
endif()

# pybind11 dependency
# nanobind dependency
if(NOT USE_SYSTEM_PYBIND)
fetchcontent_declare(
pybind11
GIT_REPOSITORY https://github.com/pybind/pybind11.git
GIT_TAG v2.13.1
nanobind
GIT_REPOSITORY https://github.com/wjakob/nanobind.git
# master on 2024-07-06
GIT_TAG b0136fe6ac1967cb2399456adc346a1af06a3b88
PATCH_COMMAND
git apply
${CMAKE_CURRENT_SOURCE_DIR}/cmake/0001-Replace-zero-size-array.patch
${CMAKE_CURRENT_SOURCE_DIR}/cmake/0002-Give-anonymous-union-a-name.patch
UPDATE_DISCONNECTED 1
)
fetchcontent_makeavailable(pybind11)
fetchcontent_makeavailable(nanobind)
else()
find_package(pybind11 CONFIG REQUIRED)
find_package(nanobind CONFIG REQUIRED)
endif()

file(GLOB_RECURSE jormungandr_src jormungandr/cpp/*.cpp)

# Build Sleipnir dependency directly into the wheel to avoid having to
# configure RPATHs
pybind11_add_module(_jormungandr ${jormungandr_src} ${Sleipnir_src})
nanobind_add_module(_jormungandr ${jormungandr_src} ${Sleipnir_src})
compiler_flags(_jormungandr)
target_compile_definitions(_jormungandr PRIVATE JORMUNGANDR=1)
target_include_directories(
Expand All @@ -302,10 +308,12 @@ if(BUILD_PYTHON)
${CMAKE_CURRENT_SOURCE_DIR}/include
${CMAKE_CURRENT_SOURCE_DIR}/jormungandr/cpp
)
target_link_libraries(
_jormungandr
PUBLIC pybind11::module Threads::Threads Eigen3::Eigen
)
target_link_libraries(_jormungandr PUBLIC Threads::Threads Eigen3::Eigen)

# Suppress compiler warnings in nanobind
if(${CMAKE_CXX_COMPILER_ID} STREQUAL "AppleClang")
target_compile_options(nanobind-static PRIVATE "-Wno-array-bounds")
endif()

install(
TARGETS _jormungandr
Expand All @@ -314,22 +322,37 @@ if(BUILD_PYTHON)
DESTINATION ${PY_DEST}
)

# pybind11-stubgen and pybind11_mkdoc don't support Windows
if(NOT WIN32 AND NOT CMAKE_CROSSCOMPILING)
# pybind11-stubgen dependency
fetchcontent_declare(
pybind11-stubgen
GIT_REPOSITORY https://github.com/sizmailov/pybind11-stubgen.git
GIT_TAG v2.5.1
GIT_SUBMODULES ""
)
fetchcontent_makeavailable(pybind11-stubgen)

# Generate stubs for the Python module
include(cmake/modules/Pybind11Stubgen.cmake)
pybind11_stubgen(_jormungandr)
pybind11_stubgen_install(_jormungandr ${PY_DEST})
nanobind_add_stub(
_jormungandr_stub
INSTALL_TIME
MARKER_FILE jormungandr/py.typed
MODULE _jormungandr
OUTPUT jormungandr/__init__.pyi
PYTHON_PATH $<TARGET_FILE_DIR:_jormungandr>
DEPENDS _jormungandr
COMPONENT python_modules
)
nanobind_add_stub(
_jormungandr_autodiff_stub
INSTALL_TIME
MODULE _jormungandr.autodiff
OUTPUT jormungandr/autodiff.pyi
PYTHON_PATH $<TARGET_FILE_DIR:_jormungandr>
DEPENDS _jormungandr
COMPONENT python_modules
)
nanobind_add_stub(
_jormungandr_optimization_stub
INSTALL_TIME
MODULE _jormungandr.optimization
OUTPUT jormungandr/optimization.pyi
PYTHON_PATH $<TARGET_FILE_DIR:_jormungandr>
DEPENDS _jormungandr
COMPONENT python_modules
)

# pybind11_mkdoc doesn't support Windows
if(NOT WIN32 AND NOT CMAKE_CROSSCOMPILING)
# pybind11_mkdoc dependency
fetchcontent_declare(
pybind11_mkdoc
Expand Down
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ See the [examples](https://github.com/SleipnirGroup/Sleipnir/tree/main/examples)
* On Linux, install via `sudo apt install python`
* On macOS, install via `brew install python`
* [Eigen](https://gitlab.com/libeigen/eigen)
* [pybind11](https://github.com/pybind/pybind11) (build only)
* [nanobind](https://github.com/wjakob/nanobind) (build only)
* [Catch2](https://github.com/catchorg/Catch2) (tests only)

Library dependencies which aren't installed locally will be automatically downloaded and built by CMake.
Expand Down
27 changes: 27 additions & 0 deletions cmake/0001-Replace-zero-size-array.patch
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Tyler Veness <[email protected]>
Date: Sat, 6 Jul 2024 14:27:32 -0700
Subject: [PATCH 1/2] Replace zero-size array

---
include/nanobind/nb_attr.h | 6 ------
1 file changed, 6 deletions(-)

diff --git a/include/nanobind/nb_attr.h b/include/nanobind/nb_attr.h
index 43aec3f196c3847cf72abce29248bd4a2a4bcd6b..b6a9c17f914f9bc86ca507f036c59f9ddb85a890 100644
--- a/include/nanobind/nb_attr.h
+++ b/include/nanobind/nb_attr.h
@@ -204,13 +204,7 @@ template <size_t Size> struct func_data_prelim {
//
// As a workaround it is likely possible to restrict the scope of style
// checks to particular C++ namespaces or source code locations.
-#if defined(_MSC_VER)
- // MSVC doesn't support zero-length arrays
arg_data args[Size == 0 ? 1 : Size];
-#else
- // GCC and Clang do.
- arg_data args[Size];
-#endif
};

template <typename F>
210 changes: 210 additions & 0 deletions cmake/0002-Give-anonymous-union-a-name.patch
Original file line number Diff line number Diff line change
@@ -0,0 +1,210 @@
From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Tyler Veness <[email protected]>
Date: Sat, 6 Jul 2024 14:31:24 -0700
Subject: [PATCH 2/2] Give anonymous union a name

---
include/nanobind/nb_class.h | 2 +-
src/implicit.cpp | 24 ++++++++++++------------
src/nb_enum.cpp | 18 +++++++++---------
src/nb_type.cpp | 18 +++++++++---------
4 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/include/nanobind/nb_class.h b/include/nanobind/nb_class.h
index 05ffad9c7aee250c8a2c128df1777e48423218f5..5fd9d82da5b4e61d2072aff7f2b0ceb5e4e73651 100644
--- a/include/nanobind/nb_class.h
+++ b/include/nanobind/nb_class.h
@@ -118,7 +118,7 @@ struct type_data {
void *fwd;
void *rev;
} enum_tbl;
- };
+ } u;
void (*set_self_py)(void *, PyObject *) noexcept;
bool (*keep_shared_from_this_alive)(PyObject *) noexcept;
#if defined(Py_LIMITED_API)
diff --git a/src/implicit.cpp b/src/implicit.cpp
index 721ddbc80088d4ff928ce3a01add9e90c446ca03..b1999f596d3a62ef8d9a776d9148996a0f2cdc32 100644
--- a/src/implicit.cpp
+++ b/src/implicit.cpp
@@ -22,22 +22,22 @@ void implicitly_convertible(const std::type_info *src,
size_t size = 0;

if (t->flags & (uint32_t) type_flags::has_implicit_conversions) {
- while (t->implicit.cpp && t->implicit.cpp[size])
+ while (t->u.implicit.cpp && t->u.implicit.cpp[size])
size++;
} else {
- t->implicit.cpp = nullptr;
- t->implicit.py = nullptr;
+ t->u.implicit.cpp = nullptr;
+ t->u.implicit.py = nullptr;
t->flags |= (uint32_t) type_flags::has_implicit_conversions;
}

void **data = (void **) malloc(sizeof(void *) * (size + 2));

if (size)
- memcpy(data, t->implicit.cpp, size * sizeof(void *));
+ memcpy(data, t->u.implicit.cpp, size * sizeof(void *));
data[size] = (void *) src;
data[size + 1] = nullptr;
- free(t->implicit.cpp);
- t->implicit.cpp = (decltype(t->implicit.cpp)) data;
+ free(t->u.implicit.cpp);
+ t->u.implicit.cpp = (decltype(t->u.implicit.cpp)) data;
}

void implicitly_convertible(bool (*predicate)(PyTypeObject *, PyObject *,
@@ -50,21 +50,21 @@ void implicitly_convertible(bool (*predicate)(PyTypeObject *, PyObject *,
size_t size = 0;

if (t->flags & (uint32_t) type_flags::has_implicit_conversions) {
- while (t->implicit.py && t->implicit.py[size])
+ while (t->u.implicit.py && t->u.implicit.py[size])
size++;
} else {
- t->implicit.cpp = nullptr;
- t->implicit.py = nullptr;
+ t->u.implicit.cpp = nullptr;
+ t->u.implicit.py = nullptr;
t->flags |= (uint32_t) type_flags::has_implicit_conversions;
}

void **data = (void **) malloc(sizeof(void *) * (size + 2));
if (size)
- memcpy(data, t->implicit.py, size * sizeof(void *));
+ memcpy(data, t->u.implicit.py, size * sizeof(void *));
data[size] = (void *) predicate;
data[size + 1] = nullptr;
- free(t->implicit.py);
- t->implicit.py = (decltype(t->implicit.py)) data;
+ free(t->u.implicit.py);
+ t->u.implicit.py = (decltype(t->u.implicit.py)) data;
}

NAMESPACE_END(detail)
diff --git a/src/nb_enum.cpp b/src/nb_enum.cpp
index d81d415ddf9ed4516eb4b5ab3b6933b5a7b0ec1e..25f0b11a6248cdd5a82b5f8d4fb5f22da28e814b 100644
--- a/src/nb_enum.cpp
+++ b/src/nb_enum.cpp
@@ -65,8 +65,8 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
t->type = ed->type;
t->type_py = (PyTypeObject *) result.ptr();
t->flags = ed->flags;
- t->enum_tbl.fwd = new enum_map();
- t->enum_tbl.rev = new enum_map();
+ t->u.enum_tbl.fwd = new enum_map();
+ t->u.enum_tbl.rev = new enum_map();
t->scope = ed->scope;

it.value() = t;
@@ -76,8 +76,8 @@ PyObject *enum_create(enum_init_data *ed) noexcept {

result.attr("__nb_enum__") = capsule(t, [](void *p) noexcept {
type_init_data *t = (type_init_data *) p;
- delete (enum_map *) t->enum_tbl.fwd;
- delete (enum_map *) t->enum_tbl.rev;
+ delete (enum_map *) t->u.enum_tbl.fwd;
+ delete (enum_map *) t->u.enum_tbl.rev;
nb_type_unregister(t);
free((char*)t->name);
delete t;
@@ -138,10 +138,10 @@ void enum_append(PyObject *tp_, const char *name_, int64_t value_,

member_map[name] = el;

- enum_map *fwd = (enum_map *) t->enum_tbl.fwd;
+ enum_map *fwd = (enum_map *) t->u.enum_tbl.fwd;
fwd->emplace(value_, (int64_t) (uintptr_t) el.ptr());

- enum_map *rev = (enum_map *) t->enum_tbl.rev;
+ enum_map *rev = (enum_map *) t->u.enum_tbl.rev;
rev->emplace((int64_t) (uintptr_t) el.ptr(), value_);
}

@@ -150,7 +150,7 @@ bool enum_from_python(const std::type_info *tp, PyObject *o, int64_t *out, uint8
if (!t)
return false;

- enum_map *rev = (enum_map *) t->enum_tbl.rev;
+ enum_map *rev = (enum_map *) t->u.enum_tbl.rev;
enum_map::iterator it = rev->find((int64_t) (uintptr_t) o);

if (it != rev->end()) {
@@ -159,7 +159,7 @@ bool enum_from_python(const std::type_info *tp, PyObject *o, int64_t *out, uint8
}

if (flags & (uint8_t) cast_flags::convert) {
- enum_map *fwd = (enum_map *) t->enum_tbl.fwd;
+ enum_map *fwd = (enum_map *) t->u.enum_tbl.fwd;

if (t->flags & (uint32_t) type_flags::is_signed) {
long long value = PyLong_AsLongLong(o);
@@ -195,7 +195,7 @@ PyObject *enum_from_cpp(const std::type_info *tp, int64_t key) noexcept {
if (!t)
return nullptr;

- enum_map *fwd = (enum_map *) t->enum_tbl.fwd;
+ enum_map *fwd = (enum_map *) t->u.enum_tbl.fwd;

enum_map::iterator it = fwd->find(key);
if (it != fwd->end()) {
diff --git a/src/nb_type.cpp b/src/nb_type.cpp
index c9f76821bcf6b31be791819344ae89f6e90efa57..88c5654662053e7d4a924f52c24bfbdb59049c0a 100644
--- a/src/nb_type.cpp
+++ b/src/nb_type.cpp
@@ -371,8 +371,8 @@ static void nb_type_dealloc(PyObject *o) {
nb_type_unregister(t);

if (t->flags & (uint32_t) type_flags::has_implicit_conversions) {
- free(t->implicit.cpp);
- free(t->implicit.py);
+ free(t->u.implicit.cpp);
+ free(t->u.implicit.py);
}

free((char *) t->name);
@@ -420,8 +420,8 @@ static int nb_type_init(PyObject *self, PyObject *args, PyObject *kwds) {
t->name = strdup_check(PyUnicode_AsUTF8AndSize(name, nullptr));
Py_DECREF(name);
t->type_py = (PyTypeObject *) self;
- t->implicit.cpp = nullptr;
- t->implicit.py = nullptr;
+ t->u.implicit.cpp = nullptr;
+ t->u.implicit.py = nullptr;
t->alias_chain = nullptr;

return 0;
@@ -1146,8 +1146,8 @@ static NB_NOINLINE bool nb_type_get_implicit(PyObject *src,
const type_data *dst_type,
nb_internals *internals_,
cleanup_list *cleanup, void **out) noexcept {
- if (dst_type->implicit.cpp && cpp_type_src) {
- const std::type_info **it = dst_type->implicit.cpp;
+ if (dst_type->u.implicit.cpp && cpp_type_src) {
+ const std::type_info **it = dst_type->u.implicit.cpp;
const std::type_info *v;

while ((v = *it++)) {
@@ -1155,7 +1155,7 @@ static NB_NOINLINE bool nb_type_get_implicit(PyObject *src,
goto found;
}

- it = dst_type->implicit.cpp;
+ it = dst_type->u.implicit.cpp;
while ((v = *it++)) {
const type_data *d = nb_type_c2p(internals_, v);
if (d && PyType_IsSubtype(Py_TYPE(src), d->type_py))
@@ -1163,9 +1163,9 @@ static NB_NOINLINE bool nb_type_get_implicit(PyObject *src,
}
}

- if (dst_type->implicit.py) {
+ if (dst_type->u.implicit.py) {
bool (**it)(PyTypeObject *, PyObject *, cleanup_list *) noexcept =
- dst_type->implicit.py;
+ dst_type->u.implicit.py;
bool (*v2)(PyTypeObject *, PyObject *, cleanup_list *) noexcept;

while ((v2 = *it++)) {
Loading

0 comments on commit c5da740

Please sign in to comment.