Merge pull request #1034 from FlorianDeconinck/feature/fdeconinck/MKIAU_fortran_python_bridge

[MKIAU] Initial prototype of a CFFI-based fortran/python bridge

Author: sdrabenh

GEOSmkiau_GridComp/CMakeLists.txt

@@ -1,5 +1,7 @@
 esma_set_this()
 
+option(BUILD_PYMKIAU_INTERFACE "Build pyMKIAU interface" OFF)
+
 set (srcs
   IAU_GridCompMod.F90
   GEOS_mkiauGridComp.F90
@@ -8,6 +10,90 @@ set (srcs
   DynVec_GridComp.F90
   )
 
-set(dependencies MAPL_cfio_r4 NCEP_sp_r4i4 GEOS_Shared GMAO_mpeu MAPL FVdycoreCubed_GridComp ESMF::ESMF NetCDF::NetCDF_Fortran)
-esma_add_library (${this} SRCS ${srcs} DEPENDENCIES ${dependencies})
+if (BUILD_PYMKIAU_INTERFACE)
+  list (APPEND srcs
+    pyMKIAU/interface/interface.f90
+    pyMKIAU/interface/interface.c)
+
+  message(STATUS "Building pyMKIAU interface")
+
+  add_definitions(-DPYMKIAU_INTEGRATION)
+
+  # The Python library creation requires mpiexec/mpirun to run on a
+  # compute node. Probably a weird SLURM thing?
+  find_package(Python3 COMPONENTS Interpreter REQUIRED)
+
+  # Set up some variables in case names change
+  set(PYMKIAU_INTERFACE_LIBRARY     ${CMAKE_CURRENT_BINARY_DIR}/libpyMKIAU_interface_py.so)
+  set(PYMKIAU_INTERFACE_HEADER_FILE ${CMAKE_CURRENT_BINARY_DIR}/pyMKIAU_interface_py.h)
+  set(PYMKIAU_INTERFACE_FLAG_HEADER_FILE ${CMAKE_CURRENT_SOURCE_DIR}/pyMKIAU/interface/interface.h)
+  set(PYMKIAU_INTERFACE_SRCS        ${CMAKE_CURRENT_SOURCE_DIR}/pyMKIAU/interface/interface.py)
+
+  # This command creates the shared object library from Python
+  add_custom_command(
+    OUTPUT ${PYMKIAU_INTERFACE_LIBRARY}
+    # Note below is essentially:
+    #  mpirun -np 1 python file
+    # but we use the CMake options as much as we can for flexibility
+    COMMAND ${CMAKE_COMMAND} -E copy_if_different ${PYMKIAU_INTERFACE_FLAG_HEADER_FILE} ${CMAKE_CURRENT_BINARY_DIR}
+    COMMAND ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 1 ${Python3_EXECUTABLE} ${PYMKIAU_INTERFACE_SRCS}
+    BYPRODUCTS ${PYMKIAU_INTERFACE_HEADER_FILE}
+    WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
+    MAIN_DEPENDENCY ${PYMKIAU_INTERFACE_SRCS}
+    COMMENT "Building pyMKIAU interface library with Python"
+    VERBATIM
+  )
+
+  # This creates a target we can use for dependencies and post build
+  add_custom_target(generate_pyMKIAU_interface_library DEPENDS ${PYMKIAU_INTERFACE_LIBRARY})
 
+  # Because of the weird hacking of INTERFACE libraries below, we cannot
+  # use the "usual" CMake calls to install() the .so. I think it's because
+  # INTERFACE libraries don't actually produce any artifacts as far as
+  # CMake is concerned. So we add a POST_BUILD custom command to "install"
+  # the library into install/lib
+  add_custom_command(TARGET generate_pyMKIAU_interface_library
+    POST_BUILD
+    # We first need to make a lib dir if it doesn't exist. If not, then
+    # the next command can copy the script into a *file* called lib because
+    # of a race condition (if install/lib/ isn't mkdir'd first)
+    COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_INSTALL_PREFIX}/lib
+    # Now we copy the file (if different...though not sure if this is useful)
+    COMMAND ${CMAKE_COMMAND} -E copy_if_different "${PYMKIAU_INTERFACE_LIBRARY}" ${CMAKE_INSTALL_PREFIX}/lib
+    )
+
+  # We use INTERFACE libraries to create a sort of "fake" target library we can use
+  # to make libFVdycoreCubed_GridComp.a depend on. It seems to work!
+  add_library(pyMKIAU_interface_py INTERFACE)
+
+  # The target_include_directories bits were essentially stolen from the esma_add_library
+  # code...
+  target_include_directories(pyMKIAU_interface_py INTERFACE
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}> # stubs
+    # modules and copied *.h, *.inc
+    $<BUILD_INTERFACE:${esma_include}/${this}>
+    $<INSTALL_INTERFACE:include/${this}>
+    )
+  target_link_libraries(pyMKIAU_interface_py INTERFACE ${PYMKIAU_INTERFACE_LIBRARY})
+
+  # This makes sure the library is built first
+  add_dependencies(pyMKIAU_interface_py generate_pyMKIAU_interface_library)
+
+  # This bit is to resolve an issue and Google told me to do this. I'm not
+  # sure that the LIBRARY DESTINATION bit actually does anything since
+  # this is using INTERFACE
+  install(TARGETS pyMKIAU_interface_py
+    EXPORT ${PROJECT_NAME}-targets
+    LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib
+    )
+
+endif ()
+
+if (BUILD_PYMKIAU_INTERFACE)
+  set(dependencies pyMKIAU_interface_py MAPL_cfio_r4 NCEP_sp_r4i4 GEOS_Shared GMAO_mpeu MAPL FVdycoreCubed_GridComp ESMF::ESMF NetCDF::NetCDF_Fortran)
+else ()
+  set(dependencies MAPL_cfio_r4 NCEP_sp_r4i4 GEOS_Shared GMAO_mpeu MAPL FVdycoreCubed_GridComp ESMF::ESMF NetCDF::NetCDF_Fortran)
+endif ()
+
+esma_add_library (${this} SRCS ${srcs} DEPENDENCIES ${dependencies})

GEOSmkiau_GridComp/GEOS_mkiauGridComp.F90

@@ -19,6 +19,10 @@ module GEOS_mkiauGridCompMod
   use GEOS_UtilsMod
 ! use GEOS_RemapMod, only: myremap => remap
   use m_set_eta, only: set_eta
+#ifdef PYMKIAU_INTEGRATION
+  use pyMKIAU_interface_mod
+  use ieee_exceptions, only: ieee_get_halting_mode, ieee_set_halting_mode, ieee_all
+#endif
   implicit none
   private
 
@@ -91,8 +95,15 @@ subroutine SetServices ( GC, RC )
     type (ESMF_Config)                      :: CF
 
     logical                                 :: BLEND_AT_PBL
-
-!=============================================================================
+#ifdef PYMKIAU_INTEGRATION
+    ! IEEE trapping see below
+    logical                                 :: halting_mode(5)
+    ! BOGUS DATA TO SHOW USAGE
+    type(a_pod_struct_type) :: options
+    real, allocatable, dimension(:,:,:) :: in_buffer
+    real, allocatable, dimension(:,:,:) :: out_buffer
+#endif
+  !=============================================================================
 
 ! Begin...
 
@@ -459,6 +470,25 @@ subroutine SetServices ( GC, RC )
     call MAPL_GenericSetServices    ( gc, RC=STATUS)
     VERIFY_(STATUS)
 
+#ifdef PYMKIAU_INTEGRATION
+    ! Spin the interface - we have to deactivate the ieee error
+    ! to be able to load numpy, scipy and other numpy packages
+    ! that generate NaN as an init mechanism for numerical solving
+    call ieee_get_halting_mode(ieee_all, halting_mode)
+    call ieee_set_halting_mode(ieee_all, .false.)
+    call pyMKIAU_interface_f_setservice()
+    call ieee_set_halting_mode(ieee_all, halting_mode)
+
+    ! BOGUS CODE TO SHOW USAGE
+    options%npx = 10
+    options%npy = 11
+    options%npz = 12
+    allocate (in_buffer(10,11,12), source = 42.42 )
+    allocate (out_buffer(10,11,12), source = 0.0 )
+    call pyMKIAU_interface_f_run(options, in_buffer, out_buffer)
+    write(*,*) "[pyMKIAU] From fortran OUT[5,5,5] is ", out_buffer(5,5,5)
+#endif
+
     RETURN_(ESMF_SUCCESS)
 
   end subroutine SetServices

GEOSmkiau_GridComp/pyMKIAU/.gitignore

@@ -0,0 +1,12 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.pytest_cache
+*.egg-info/
+test_data/
+.gt_cache_*
+.translate-*/
+.vscode
+test_data/
+sandbox/
+*.mod

GEOSmkiau_GridComp/pyMKIAU/README.md

@@ -0,0 +1,40 @@
+# Fortran - Python bridge prototype
+
+Nomenclatura: we call the brige "fpy" and "c", "f" and "py" denotes functions in their respective language.
+
+Building: you have to pass `-DBUILD_PYMKIAU_INTERFACE=ON` to your `cmake` command to turn on the interface build and execution.
+
+## Pipeline
+
+Here's a quick rundown of how a buffer travels through the interface and back.
+  
+- From Fortran in `GEOS_MKIAUGridComp:488` we call `pyMKIAU_interface_f_run` with the buffer passed as argument
+- This pings the interface, located at `pyMKIAU/interface/interface.f90`. This interface uses the `iso_c_binding` to marshall the parameters downward (careful about the user type, look at the code)
+- Fortran then call into C at `pyMKIAU/interface/interface.c`. Those functions now expect that a few `extern` hooks have been made available on the python side, they are define in `pyMKIAU/interface/interface.h`
+- At runtime, the hooks are found and code carries to the python thanks to cffi. The .so that exposes the hooks is in `pyMKIAU/interface/interface.py`. Within this code, we: expose extern functions via `ffi.extern`, build a shared library to link for runtime and pass the code down to the `pyMKIAU` python package which lives at `pyMKIAU/pyMKIAU`
+- In the package, the `serservices` or `run` function is called.
+
+## Fortran <--> C: iso_c_binding
+
+We leverage Fortan `iso_c_binding` extension to do conform Fortran and C calling structure. Which comes with a bunch of easy type casting and some pretty steep potholes.
+The two big ones are:
+
+- strings need to be send/received as a buffer plus a length,
+- pointers/buffers are _not_ able to be pushed into a user type.
+
+## C <->Python: CFFI based glue
+
+The interface is based on CFFI which is reponsible for the heavy lifting of
+
+- spinning a python interpreter
+- passing memory between C and Python without a copy
+
+## Running python
+
+The last trick is to make sure your package is callable by the `interface.py`. Basically your code has to be accessible by the interpreter, be via virtual env, conda env or PYTHONPATH.
+The easy way to know is that you need to be able to get into your environment and run in a python terminal:
+
+```python
+from pyMKIAU.core import pyMKIAU_init
+pyMKIAU_init()
+```

GEOSmkiau_GridComp/pyMKIAU/interface/interface.c

@@ -0,0 +1,31 @@
+#include <stdio.h>
+#include <time.h>
+#include "interface.h"
+
+extern int pyMKIAU_interface_c_setservice()
+{
+    // Check magic number
+    int return_code = pyMKIAU_interface_py_setservices();
+
+    if (return_code < 0)
+    {
+        exit(return_code);
+    }
+}
+
+extern int pyMKIAU_interface_c_run(a_pod_struct_t *options, const float *in_buffer, float *out_buffer)
+{
+    // Check magic number
+    if (options->mn_123456789 != 123456789)
+    {
+        printf("Magic number failed, pyMKIAU interface is broken on the C side\n");
+        exit(-1);
+    }
+
+    int return_code = pyMKIAU_interface_py_run(options, in_buffer, out_buffer);
+
+    if (return_code < 0)
+    {
+        exit(return_code);
+    }
+}

GEOSmkiau_GridComp/pyMKIAU/interface/interface.f90

@@ -0,0 +1,43 @@
+module pyMKIAU_interface_mod
+
+   use iso_c_binding, only: c_int, c_float, c_double, c_bool, c_ptr
+
+   implicit none
+
+   private
+   public :: pyMKIAU_interface_f_setservice, pyMKIAU_interface_f_run
+   public :: a_pod_struct_type
+
+   !-----------------------------------------------------------------------
+   ! See `interface.h` for explanation of the POD-strict struct
+   !-----------------------------------------------------------------------
+   type, bind(c) :: a_pod_struct_type
+      integer(kind=c_int) :: npx
+      integer(kind=c_int) :: npy
+      integer(kind=c_int) :: npz
+      ! Magic number
+      integer(kind=c_int) :: make_flags_C_interop = 123456789
+   end type
+
+
+   interface
+
+      subroutine pyMKIAU_interface_f_setservice() bind(c, name='pyMKIAU_interface_c_setservice')
+      end subroutine pyMKIAU_interface_f_setservice
+
+      subroutine pyMKIAU_interface_f_run(options, in_buffer, out_buffer) bind(c, name='pyMKIAU_interface_c_run')
+
+         import c_float, a_pod_struct_type
+
+         implicit none
+         ! This is an interface to a C function, the intent ARE NOT enforced
+         ! by the compiler. Consider them developer hints
+         type(a_pod_struct_type), intent(in) :: options
+         real(kind=c_float), dimension(*), intent(in) :: in_buffer
+         real(kind=c_float), dimension(*), intent(out) :: out_buffer
+
+      end subroutine pyMKIAU_interface_f_run
+   
+   end interface
+
+end module pyMKIAU_interface_mod

GEOSmkiau_GridComp/pyMKIAU/interface/interface.h

@@ -0,0 +1,40 @@
+#pragma once
+
+/***
+ * C Header for the interface to python.
+ * Define here any POD-strict structures and external functions
+ * that will get exported by cffi from python (see interface.py)
+ ***/
+
+#include <stdbool.h>
+#include <stdlib.h>
+
+// POD-strict structure to pack options and flags efficiently
+// Struct CANNOT hold pointers. The iso_c_binding does not allow for foolproof
+// pointer memory packing.
+// We use the low-embedded trick of the magic number to attempt to catch
+// any type mismatch betweeen Fortran and C. This is not a foolproof method
+// but it bring a modicum of check at the cost of a single integer.
+typedef struct
+{
+	int npx;
+	int npy;
+	int npz;
+	// Magic number needs to be last item
+	int mn_123456789;
+} a_pod_struct_t;
+
+// For complex type that can be exported with different
+// types (like the MPI communication object), you can rely on C `union`
+typedef union
+{
+	int comm_int;
+	void *comm_ptr;
+} MPI_Comm_t;
+
+// Python hook functions: defined as external so that the .so can link out ot them
+// Though we define `in_buffer` as a `const float*` it is _not_ enforced
+// by the interface. Treat as a developer hint only.
+
+extern int pyMKIAU_interface_py_run(a_pod_struct_t *options, const float *in_buffer, float *out_buffer);
+extern int pyMKIAU_interface_py_setservices();

GEOSmkiau_GridComp/pyMKIAU/interface/interface.py

@@ -0,0 +1,46 @@
+import cffi  # type: ignore
+
+TMPFILEBASE = "pyMKIAU_interface_py"
+
+ffi = cffi.FFI()
+
+source = """
+from {} import ffi
+from datetime import datetime
+from pyMKIAU.core import pyMKIAU_init, pyMKIAU_run #< User code starts here
+import traceback
+
+@ffi.def_extern()
+def pyMKIAU_interface_py_setservices() -> int:
+
+    try:
+        # Calling out off the bridge into the python
+        pyMKIAU_init()
+    except Exception as err:
+        print("Error in Python:")
+        print(traceback.format_exc())
+        return -1
+    return 0
+
+@ffi.def_extern()
+def pyMKIAU_interface_py_run(options, in_buffer, out_buffer) -> int:
+
+    try:
+        # Calling out off the bridge into the python
+        pyMKIAU_run(options, in_buffer, out_buffer)
+    except Exception as err:
+        print("Error in Python:")
+        print(traceback.format_exc())
+        return -1
+    return 0
+""".format(TMPFILEBASE)
+
+with open("interface.h") as f:
+    data = "".join([line for line in f if not line.startswith("#")])
+    data = data.replace("CFFI_DLLEXPORT", "")
+    ffi.embedding_api(data)
+
+ffi.set_source(TMPFILEBASE, '#include "interface.h"')
+
+ffi.embedding_init_code(source)
+ffi.compile(target="lib" + TMPFILEBASE + ".so", verbose=True)