No more Cython DEF and IF + prepare MPI

fluiddyn · Feb 2, 2024 · a768d26 · a768d26
1 parent 910b562
commit a768d26
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 27 deletions.
diff --git a/Makefile b/Makefile
@@ -1,4 +1,4 @@
-.PHONY: clean cleanall cleanmako cleancython develop build_ext_inplace list-sessions tests
+.PHONY: clean cleanall develop build_ext_inplace list-sessions tests
 
 develop: sync
 	pdm run pip install -e . --no-deps --no-build-isolation -v
@@ -17,13 +17,7 @@ cleanso:
 cleanpythran:
 	find src -name __pythran__ -type d -exec rm -rf "{}" +
 
-cleancython:
-	find src -name "*_cy.cpp" -delete
-
-cleanmako:
-	python -c "from fluidfft_builder.src_cy.make_files_with_mako import clean_files as c; c()"
-
-cleanall: clean cleanso cleanmako cleancython cleanpythran
+cleanall: clean cleanso cleanpythran
 
 black:
 	pdm run black

diff --git a/plugins/fluidfft-builder/fluidfft_builder/include_cy/base.pyx b/plugins/fluidfft-builder/fluidfft_builder/include_cy/base.pyx
@@ -1,8 +1,6 @@
 # cython: embedsignature=True
 # cython: language_level=3
 
-DEF MPI4PY = 0
-
 cimport cython
 
 from cython cimport view
@@ -21,16 +19,6 @@ else:
     nb_proc = comm.size
     rank = comm.Get_rank()
 
-IF MPI4PY:
-    from mpi4py cimport MPI
-
-    #from mpi4py.mpi_c cimport *
-    # for mpi4py > 2.0
-    from mpi4py.libmpi cimport *
-
-    # fix a bug arising when using a recent version of mpi4py
-    cdef extern from 'mpi-compat.h': pass
-
 # we define python and c types for physical and Fourier spaces
 DTYPEb = np.uint8
 ctypedef np.uint8_t DTYPEb_t

diff --git a/plugins/fluidfft-builder/fluidfft_builder/include_cy/base_mpi.pyx b/plugins/fluidfft-builder/fluidfft_builder/include_cy/base_mpi.pyx
@@ -0,0 +1,5 @@
+from mpi4py cimport MPI
+from mpi4py.libmpi cimport *
+
+# fix a bug arising when using a recent version of mpi4py
+cdef extern from 'mpi-compat.h': pass
diff --git a/plugins/fluidfft-builder/fluidfft_builder/make_cy_files.py b/plugins/fluidfft-builder/fluidfft_builder/make_cy_files.py
@@ -47,15 +47,23 @@ def make_file(path_output, class_name, numpy_api="numpy"):
     name_output = path_output.name
     module_name, extension = name_output.split(".")
 
-    template_name = f"template{dimension}d.{extension}"
+    if name_output.startswith("mpi"):
+        cdef_public_mpi_comm = "cdef public MPI.Comm comm"
+        include_base_mpi_pyx = "include 'base_mpi.pyx'"
+    else:
+        cdef_public_mpi_comm = ""
+        include_base_mpi_pyx = ""
 
+    template_name = f"template{dimension}d.{extension}"
     template = load_template(template_name)
 
     content = template.substitute(
         {
             "module_name": module_name,
             "class_name": class_name,
             "numpy_api": numpy_api,
+            "cdef_public_mpi_comm": cdef_public_mpi_comm,
+            "include_base_mpi_pyx": include_base_mpi_pyx,
         }
     )
 

diff --git a/plugins/fluidfft-builder/fluidfft_builder/templates/template2d.pyx b/plugins/fluidfft-builder/fluidfft_builder/templates/template2d.pyx
@@ -2,7 +2,7 @@
 # cython: language_level=3
 
 include 'base.pyx'
-
+${include_base_mpi_pyx}
 
 from fft2d_${module_name} cimport (
     ${class_name} as mycppclass,
@@ -31,8 +31,7 @@ cdef class ${class_name}:
     cdef tuple _shapeK_loc, _shapeX_loc, _shapeK_seq, _shapeX_seq
     cdef int _is_mpi_lib
 
-    IF MPI4PY:
-        cdef public MPI.Comm comm
+    ${cdef_public_mpi_comm}
     cdef public int nb_proc, rank
 
     def __cinit__(self, int n0=2, int n1=2):

diff --git a/plugins/fluidfft-builder/fluidfft_builder/templates/template3d.pyx b/plugins/fluidfft-builder/fluidfft_builder/templates/template3d.pyx
@@ -2,7 +2,7 @@
 # cython: language_level=3
 
 include 'base.pyx'
-
+${include_base_mpi_pyx}
 
 from fft3d_${module_name} cimport (
     ${class_name} as mycppclass,
@@ -60,8 +60,8 @@ cdef class ${class_name}:
     cdef int dim_first_fft
     cdef int _is_mpi_lib
 
-    IF MPI4PY:
-        cdef public MPI.Comm comm
+    ${cdef_public_mpi_comm}
+
     cdef public int nb_proc, rank
 
     def __cinit__(self, int n0=2, int n1=2, int n2=4):

diff --git a/plugins/fluidfft-fftw/Makefile b/plugins/fluidfft-fftw/Makefile
@@ -1,3 +1,6 @@
 
 develop:
 	pip install -e . -vv --no-build-isolation --no-deps
+
+clean:
+	rm -rf build