Skip to content

Commit

Permalink
Switch from pybind11 to nanobind
Browse files Browse the repository at this point in the history
  • Loading branch information
francesco-ballarin committed Nov 7, 2023
1 parent a50b164 commit bd5c79f
Show file tree
Hide file tree
Showing 8 changed files with 74 additions and 62 deletions.
6 changes: 2 additions & 4 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -36,14 +36,14 @@ jobs:
echo "OMPI_ALLOW_RUN_AS_ROOT=1" >> $GITHUB_ENV
echo "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1" >> $GITHUB_ENV
rm /usr/local/lib/python3.11/site-packages/petsc4py/py.typed
python3 -m pip install pybind11 scikit-build-core[pyproject]
python3 -m pip install nanobind scikit-build-core[pyproject]
- backend: none-complex
container: numericalpdes/base_images:slepc4py-complex
setup_container: |
echo "OMPI_ALLOW_RUN_AS_ROOT=1" >> $GITHUB_ENV
echo "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1" >> $GITHUB_ENV
rm /usr/local/lib/python3.11/site-packages/petsc4py/py.typed
python3 -m pip install pybind11 scikit-build-core[pyproject]
python3 -m pip install nanobind scikit-build-core[pyproject]
- backend: dolfinx-real
container: dolfinx/dolfinx:nightly
setup_container: |
Expand All @@ -53,7 +53,6 @@ jobs:
echo "PYTHONPATH=$PYTHONPATH" >> $GITHUB_ENV
echo "PKG_CONFIG_PATH=$PKG_CONFIG_PATH" >> $GITHUB_ENV
rm /usr/local/lib/python3.10/dist-packages/petsc4py/py.typed
python3 -m pip install scikit-build-core[pyproject]
- backend: dolfinx-complex
container: dolfinx/dolfinx:nightly
setup_container: |
Expand All @@ -63,7 +62,6 @@ jobs:
echo "PYTHONPATH=$PYTHONPATH" >> $GITHUB_ENV
echo "PKG_CONFIG_PATH=$PKG_CONFIG_PATH" >> $GITHUB_ENV
rm /usr/local/lib/python3.10/dist-packages/petsc4py/py.typed
python3 -m pip install scikit-build-core[pyproject]
fail-fast: false
container:
image: ${{ matrix.container }}
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
[build-system]
requires = [
"nanobind >= 1.8.0",
"petsc4py",
"pybind11 >= 2.9.1",
"scikit-build-core[pyproject]"
]
build-backend = "scikit_build_core.build"
Expand Down
1 change: 1 addition & 0 deletions rbnicsx/_cpp/.gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,5 @@ CMakeCache.txt
CMakeFiles
Makefile
cmake_install.cmake
libnanobind-static.a
rbnicsx_cpp.*.so
33 changes: 24 additions & 9 deletions rbnicsx/_cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,30 @@ set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

# Set default build type
if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
endif()

# Find python
find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
find_package(Python COMPONENTS Interpreter Development REQUIRED)

# Find pybind11
find_package(
pybind11 2.9.1 REQUIRED
CONFIG HINTS ${PYBIND11_DIR} ${PYBIND11_ROOT} $ENV{PYBIND11_DIR} $ENV{PYBIND11_ROOT}
# Find nanobind
execute_process(
COMMAND ${Python_EXECUTABLE} -m nanobind --cmake_dir
OUTPUT_VARIABLE NANOBIND_CMAKE_DIR
RESULT_VARIABLE NANOBIND_CMAKE_DIR_COMMAND_RESULT
ERROR_VARIABLE NANOBIND_CMAKE_DIR_COMMAND_ERROR
OUTPUT_STRIP_TRAILING_WHITESPACE
)
if(NOT NANOBIND_CMAKE_DIR_COMMAND_RESULT)
list(APPEND CMAKE_PREFIX_PATH "${NANOBIND_CMAKE_DIR}")
find_package(nanobind CONFIG REQUIRED)
message(STATUS "Found nanobind python wrappers at ${NANOBIND_CMAKE_DIR}")
else()
message(FATAL_ERROR "nanobind could not be found.")
endif()

# Check for PETSc
find_package(PkgConfig REQUIRED)
Expand All @@ -25,7 +41,7 @@ pkg_search_module(PETSC REQUIRED IMPORTED_TARGET PETSc>=3.15 petsc>=3.15)

# Check for petsc4py
execute_process(
COMMAND ${Python3_EXECUTABLE} -c "import petsc4py; print(petsc4py.get_include())"
COMMAND ${Python_EXECUTABLE} -c "import petsc4py; print(petsc4py.get_include())"
OUTPUT_VARIABLE PETSC4PY_INCLUDE_DIR
RESULT_VARIABLE PETSC4PY_INCLUDE_COMMAND_RESULT
ERROR_VARIABLE PETSC4PY_INCLUDE_COMMAND_ERROR
Expand All @@ -38,13 +54,12 @@ else()
message(FATAL_ERROR "petsc4py could not be found.")
endif()

# Compile rbnicsx C++ backend and pybind11 wrappers
pybind11_add_module(
# Compile rbnicsx C++ backend and nanobind wrappers
nanobind_add_module(
rbnicsx_cpp
MODULE
rbnicsx/_backends/frobenius_inner_product.cpp
rbnicsx/_backends/petsc_error.cpp
rbnicsx/_backends/petsc_casters.cpp
rbnicsx/wrappers/_backends.cpp
rbnicsx/wrappers/rbnicsx.cpp
)
Expand Down
7 changes: 0 additions & 7 deletions rbnicsx/_cpp/rbnicsx/_backends/petsc_casters.cpp

This file was deleted.

70 changes: 38 additions & 32 deletions rbnicsx/_cpp/rbnicsx/_backends/petsc_casters.h
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright (C) 2017-2021 Chris Richardson and Garth N. Wells
// Copyright (C) 2017-2023 Chris Richardson and Garth N. Wells
//
// This file is part of RBniCSx (copied from DOLFINx).
//
Expand All @@ -8,25 +8,28 @@

#pragma once

#include <nanobind/nanobind.h>
#include <petsc4py/petsc4py.h>
#include <petscdm.h>
#include <petscis.h>
#include <petscksp.h>
#include <petscmat.h>
#include <petscsnes.h>
#include <petscvec.h>
#include <pybind11/pybind11.h>

// pybind11 casters for PETSc/petsc4py objects
// nanobind casters for PETSc/petsc4py objects

namespace py = pybind11;
namespace nb = nanobind;

// Import petsc4py on demand
#define VERIFY_PETSC4PY(func) \
#define VERIFY_PETSC4PY_FROMPY(func) \
if (!func) \
{ \
if (import_petsc4py() != 0) \
throw std::runtime_error("Error when importing petsc4py"); \
return false; \
}

#define VERIFY_PETSC4PY_FROMCPP(func) \
if (!func) \
{ \
if (import_petsc4py() != 0) \
return {}; \
}

// Macro for casting between PETSc and petsc4py objects
Expand All @@ -35,42 +38,45 @@ namespace py = pybind11;
class type_caster<_p_##TYPE> \
{ \
public: \
PYBIND11_TYPE_CASTER(TYPE, _(#NAME)); \
bool load(handle src, bool) \
NB_TYPE_CASTER(TYPE, const_name(#NAME)) \
bool from_python(handle src, uint8_t, cleanup_list*) noexcept \
{ \
VERIFY_PETSC4PY(PyPetsc##P4PYTYPE##_Get); \
VERIFY_PETSC4PY_FROMPY(PyPetsc##P4PYTYPE##_Get); \
if (PyObject_TypeCheck(src.ptr(), &PyPetsc##P4PYTYPE##_Type) != 0) \
{ \
value = PyPetsc##P4PYTYPE##_Get(src.ptr()); \
return true; \
} \
\
return false; \
else \
return false; \
} \
\
static handle cast(TYPE src, py::return_value_policy policy, \
handle parent) \
static handle from_cpp(TYPE src, rv_policy policy, \
cleanup_list* /*cleanup*/) noexcept \
{ \
VERIFY_PETSC4PY(PyPetsc##P4PYTYPE##_New); \
auto obj = PyPetsc##P4PYTYPE##_New(src); \
if (policy == py::return_value_policy::take_ownership) \
VERIFY_PETSC4PY_FROMCPP(PyPetsc##P4PYTYPE##_New); \
if (policy == rv_policy::take_ownership) \
{ \
PyObject* obj = PyPetsc##P4PYTYPE##_New(src); \
PetscObjectDereference((PetscObject)src); \
else if (policy == py::return_value_policy::reference_internal) \
keep_alive_impl(obj, parent); \
return py::handle(obj); \
return nb::handle(obj); \
} \
else if (policy == rv_policy::automatic_reference \
or policy == rv_policy::reference) \
{ \
PyObject* obj = PyPetsc##P4PYTYPE##_New(src); \
return nb::handle(obj); \
} \
else \
{ \
return {}; \
} \
} \
\
operator TYPE() { return value; } \
}

namespace pybind11::detail
namespace nanobind::detail
{
PETSC_CASTER_MACRO(DM, DM, dm);
PETSC_CASTER_MACRO(IS, IS, is);
PETSC_CASTER_MACRO(KSP, KSP, ksp);
PETSC_CASTER_MACRO(Mat, Mat, mat);
PETSC_CASTER_MACRO(SNES, SNES, snes);
PETSC_CASTER_MACRO(Vec, Vec, vec);
} // namespace pybind11::detail

#undef PETSC_CASTER_MACRO
} // namespace nanobind::detail
7 changes: 3 additions & 4 deletions rbnicsx/_cpp/rbnicsx/wrappers/_backends.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,16 @@
//
// SPDX-License-Identifier: LGPL-3.0-or-later

#include <pybind11/complex.h>
#include <pybind11/pybind11.h>
#include <nanobind/nanobind.h>

#include <rbnicsx/_backends/petsc_casters.h>
#include <rbnicsx/_backends/frobenius_inner_product.h>

namespace py = pybind11;
namespace nb = nanobind;

namespace rbnicsx_wrappers
{
void _backends(py::module& m)
void _backends(nb::module_& m)
{
m.def("frobenius_inner_product", &rbnicsx::_backends::frobenius_inner_product,
"Frobenius inner product between PETSc Mat objects.");
Expand Down
10 changes: 5 additions & 5 deletions rbnicsx/_cpp/rbnicsx/wrappers/rbnicsx.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,21 @@
//
// SPDX-License-Identifier: LGPL-3.0-or-later

#include <pybind11/pybind11.h>
#include <nanobind/nanobind.h>

namespace py = pybind11;
namespace nb = nanobind;

namespace rbnicsx_wrappers
{
void _backends(py::module& m);
void _backends(nb::module_& m);
}

PYBIND11_MODULE(rbnicsx_cpp, m)
NB_MODULE(rbnicsx_cpp, m)
{
// Create module for C++ wrappers
m.doc() = "RBniCSx Python interface";

// Create internal backends submodule
py::module _backends = m.def_submodule("_backends", "Internal backends module");
nb::module_ _backends = m.def_submodule("_backends", "Internal backends module");
rbnicsx_wrappers::_backends(_backends);
}

0 comments on commit bd5c79f

Please sign in to comment.