Add Cython bindings for the PETSc API

.github/workflows/core.yml

@@ -130,6 +130,15 @@ jobs:
           . venv-petsctools/bin/activate
           pytest petsctools-repo
 
+      - name: Run Cython demo
+        if: success() || steps.install-petsc4py.conclusion == 'success'
+        run: |
+          . venv-petsctools/bin/activate
+          cd petsctools-repo/docs/source/_static/cython_demo
+          CC=mpicc python setup.py build_ext --inplace
+          python -c "import slow"
+          python -c "import fast"
+
       - name: Build documentation
         id: build_docs
         if: success() || steps.install.conclusion == 'success'

.gitignore

@@ -1,7 +1,14 @@
 petsctools/config.ini
 
-*/__pycache__
+**/__pycache__
 *.egg-info
 
+build/
+
 docs/build
 docs/source/generated
+
+docs/source/_static/cython_demo/build
+docs/source/_static/cython_demo/*.c
+docs/source/_static/cython_demo/*.so
+docs/source/_static/cython_demo/*.html

docs/source/_static/cython_demo/fast.pyx

@@ -0,0 +1,36 @@
+import time
+
+import cython
+from petsc4py import PETSc
+
+from petsctools cimport cpetsc
+
+
+def medium():
+    N: cython.int = int(1e8)
+    section: PETSc.Section = PETSc.Section().create()
+    section.setChart(0, N)
+
+    start = time.time()
+    i: cython.int
+    for i in range(N):
+        if i % 2 == 0:
+            section.setDof(i, 1)
+    print(f"Time elapsed: {time.time() - start}")
+
+
+def fast():
+    N: cython.int = int(1e8)
+    section: cpetsc.PetscSection_py = PETSc.Section().create()
+    section.setChart(0, N)
+
+    start = time.time()
+    i: cython.int
+    for i in range(N):
+        if i % 2 == 0:
+            cpetsc.CHKERR(cpetsc.PetscSectionSetDof(section.sec, i, 1))
+    print(f"Time elapsed: {time.time() - start}")
+
+
+medium()
+fast()

docs/source/_static/cython_demo/setup.py

@@ -0,0 +1,20 @@
+from setuptools import setup, Extension
+
+import petsc4py
+import petsctools
+
+
+extension = Extension(
+    name="fast",
+    language="c",
+    sources=["fast.pyx"],
+    include_dirs=[
+        petsc4py.get_include(),
+        *petsctools.get_petsc_dirs(subdir="include"),
+    ],
+    library_dirs=petsctools.get_petsc_dirs(subdir="lib"),
+    runtime_library_dirs=petsctools.get_petsc_dirs(subdir="lib"),
+    libraries=["petsc", "mpi"],
+)
+
+setup(ext_modules=[extension])

docs/source/_static/cython_demo/slow.py

@@ -0,0 +1,18 @@
+import time
+
+from petsc4py import PETSc
+
+
+def slow():
+    N = int(1e8)
+    section = PETSc.Section().create()
+    section.setChart(0, N)
+
+    start = time.time()
+    for i in range(N):
+        if i % 2 == 0:
+            section.setDof(i, 1)
+    print(f"Time elapsed: {time.time() - start}")
+
+
+slow()

docs/source/cython.rst

@@ -0,0 +1,99 @@
+PETSc C bindings
+----------------
+
+In some circumstances it is desirable to write PETSc C code instead of using
+the petsc4py Python bindings. For example:
+
+* The overhead involved in calling Python bindings may be unacceptable. This
+  may be the case in very tight loops.
+* The desired API functionality is not available in petsc4py.
+
+To support this, petsctools makes some of PETSc's C API available to use
+through `Cython <https://cython.org/>`_. It does this in a similar way
+to petsc4py but emphasises readability and faithfulness to the API.
+
+Demo
+~~~~
+
+To demonstrate this, consider the following simple piece of PETSc code:
+
+.. literalinclude:: _static/cython_demo/slow.py
+    :language: python3
+
+This code is written in Python and consists of petsc4py API calls. The
+Python overhead is therefore maximised. Run on the author's machine this
+code takes 4.0s to run to completion.
+
+Compare this to the examples given in this Cython file:
+
+.. literalinclude:: _static/cython_demo/fast.pyx
+    :language: cython
+
+Here we have two cases. The first (``medium``) is very similar to
+``slow`` but is able to partially compile itself because it is
+written in Cython. The second (``fast``) avoids petsc4py entirely
+and instead uses the PETSc C API directly as exposed through
+petsctools.
+
+Switching to Cython and using the C API directly both lead to
+performance improvements. ``medium`` takes 1.8s and ``fast``
+takes 0.78s.
+
+Compiling Cython code
+~~~~~~~~~~~~~~~~~~~~~
+
+To compile the Cython code it must be registered as a compiled
+extension inside a ``setup.py`` file. A working example for the
+extension provided above looks like:
+
+.. literalinclude:: _static/cython_demo/setup.py
+    :language: python3
+
+This file should be added to your project along with a suitable
+``pyproject.toml`` that indicates ``setuptools`` as the
+``build-backend``. ``setuptools``, ``cython``, ``petsc4py`` and
+``petsctools`` are all necessary build-time dependencies.
+
+.. note::
+   If you encounter errors like::
+
+      /.../petsc/include/petscsys.h:124:12: fatal error: mpi.h: No such file or directory
+        124 |   #include <mpi.h>
+            |            ^~~~~~~
+      compilation terminated.
+      error: command '/usr/bin/x86_64-linux-gnu-gcc' failed with exit code 1
+
+   then this usually means that setuptools cannot find your MPI distribution.
+   To fix this simply set the environment variable ``CC=mpicc`` and
+   try again.
+
+Conventions used
+~~~~~~~~~~~~~~~~
+
+* All objects are available inside the ``cpetsc`` namespace after
+  ``cimport``-ing it (e.g. ``cpetsc.PetscSectionSetDof``).
+
+* petsc4py PETSc objects are renamed to their C API equivalent with
+  a ``_py`` suffix. For example ``cpetsc.PetscSection`` represents the
+  C type and ``cpetsc.PetscSection_py`` the Cython type.
+
+* The C handle of the petsc4py objects are available through a specific
+  attribute that depends on the type. Examples include:
+
+   * ``cpetsc.Mat_py.mat``  ⟷  ``cpetsc.Mat``
+   * ``cpetsc.Vec_py.vec``  ⟷  ``cpetsc.Vec``
+   * ``cpetsc.IS_py.iset``  ⟷  ``cpetsc.IS``
+   * ``cpetsc.PetscSection_py.sec``  ⟷  ``cpetsc.PetscSection``
+
+  For more information you will have to refer to the `petsc4py
+  source code <https://gitlab.com/petsc/petsc/-/blob/main/src/binding/petsc4py/src/petsc4py/PETSc.pxd>`__.
+
+Adding more functions
+~~~~~~~~~~~~~~~~~~~~~
+
+Adding additional PETSc functions to petsctools is straightforward. You simply
+have to add the bindings to the
+`definitions file <https://github.com/firedrakeproject/petsctools/blob/main/petsctools/cpetsc.pxd>`__
+in a declarative way, just as you would write any other C file. For more
+information please refer to the `Cython documentation
+<https://cython.readthedocs.io/en/latest/src/tutorial/pxd_files.html>`__.

docs/source/index.rst

@@ -8,4 +8,5 @@ petsctools provides Pythonic extensions for petsc4py and slepc4py.
    :maxdepth: 2
 
    examples
+   cython
    generated/modules

petsctools/cpetsc.pxd

@@ -0,0 +1,60 @@
+"""This file exposes the PETSc API as a module for use in Cython."""
+
+# IMPORTANT: This file cannot be accessed if petsctools is installed in editable mode.
+
+from petsc4py cimport PETSc as _PETSc
+
+# clearer aliases from petsc4py, so the names here match the C API
+ctypedef _PETSc.PetscMat Mat
+ctypedef _PETSc.Mat Mat_py
+ctypedef _PETSc.PetscSF PetscSF
+ctypedef _PETSc.SF PetscSF_py
+ctypedef _PETSc.PetscSection PetscSection
+ctypedef _PETSc.Section PetscSection_py
+ctypedef _PETSc.PetscIS IS
+ctypedef _PETSc.IS IS_py
+
+# other PETSc imports
+from petsc4py.PETSc cimport (
+    CHKERR,
+    PetscErrorCode,
+)
+
+
+cdef extern from "petsc.h":
+    # fundamental types
+    ctypedef long PetscInt
+    ctypedef double PetscReal
+    ctypedef double PetscScalar
+    ctypedef enum PetscBool:
+        PETSC_TRUE
+        PETSC_FALSE
+    ctypedef enum InsertMode:
+        INSERT_VALUES
+        ADD_VALUES
+    ctypedef enum PetscCopyMode:
+        PETSC_COPY_VALUES
+        PETSC_OWN_POINTER
+        PETSC_USE_POINTER
+
+    # memory management
+    PetscErrorCode PetscCalloc1(size_t,void*)
+    PetscErrorCode PetscMalloc1(size_t,void*)
+    PetscErrorCode PetscFree(void*)
+
+    # Mat
+    PetscErrorCode MatSetValue(Mat,PetscInt,PetscInt,const PetscScalar,InsertMode)
+    PetscErrorCode MatSetValuesBlockedLocal(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode)
+
+    # PetscSF
+    ctypedef struct PetscSFNode:
+        PetscInt rank
+        PetscInt index
+
+    PetscErrorCode PetscSFGetGraph(PetscSF,PetscInt*,PetscInt*,const PetscInt**,const PetscSFNode**)
+    PetscErrorCode PetscSFSetGraph(PetscSF,PetscInt,PetscInt,PetscInt*,PetscCopyMode,PetscSFNode*,PetscCopyMode)
+
+    # PetscSection
+    PetscErrorCode PetscSectionGetDof(PetscSection,PetscInt,PetscInt*)
+    PetscErrorCode PetscSectionSetDof(PetscSection,PetscInt,PetscInt)
+    PetscErrorCode PetscSectionGetOffset(PetscSection,PetscInt,PetscInt*)

pyproject.toml

@@ -39,8 +39,13 @@ ci = [
   {include-group = "docs"},
   {include-group = "lint"},
   {include-group = "test"},
+  "cython",
+  "setuptools",
 ]
 
+[tool.setuptools.package-data]
+petsctools = ["__init__.pxd", "cpetsc.pxd"]
+
 [tool.ruff]
 line-length = 79