FFT for OMP backend (via 2decomp&fft)

.github/workflows/build_and_test.yml

@@ -9,7 +9,7 @@ on:
 
 jobs:
   unit-tests:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     env:
       OMPI_VERSION: 4.1.5
     steps:

CMakeLists.txt

@@ -2,5 +2,18 @@ cmake_minimum_required(VERSION 3.10)
 project(x3d2 LANGUAGES Fortran)
 enable_testing()
 
+#
+# Set the Poisson solver choice
+#
+set(POISSON_SOLVER FFT CACHE STRING
+  "Select the Poisson solver: FFT or ITER")
+
+if(${CMAKE_Fortran_COMPILER_ID} STREQUAL "PGI" OR
+   ${CMAKE_Fortran_COMPILER_ID} STREQUAL "NVHPC")
+   set(BACKEND CUDA)
+else()
+   set(BACKEND OMP)
+endif()
+
 add_subdirectory(src)
 add_subdirectory(tests)

cmake/Finddecomp2d.cmake

@@ -0,0 +1,27 @@
+# - Find the 2decomp-fft library
+find_package(decomp2d
+             PATHS ${CMAKE_SOURCE_DIR}/decomp2d/build)
+if (decomp2d_FOUND)
+  message(STATUS "2decomp-fft FOUND")
+else(decomp2d_FOUND)
+  message(STATUS "2decomp-fft PATH not available we'll try to download and install")
+  configure_file(${CMAKE_SOURCE_DIR}/cmake/decomp2d/downloadBuild2decomp.cmake.in decomp2d-build/CMakeLists.txt)
+  #message("Second CMAKE_GENERATOR ${CMAKE_GENERATOR}") 
+  execute_process(COMMAND ${CMAKE_COMMAND} -G "${CMAKE_GENERATOR}" .
+          RESULT_VARIABLE result
+          WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/decomp2d-build )
+  if(result)
+      message(FATAL_ERROR "CMake step for 2decomp-fft failed: ${result}")
+  else()
+      message("CMake step for 2decomp-fft completed (${result}).")
+  endif()
+  execute_process(COMMAND ${CMAKE_COMMAND} --build .
+         RESULT_VARIABLE result
+          WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/decomp2d-build )
+  if(result)
+      message(FATAL_ERROR "Build step for 2decomp-fft failed: ${result}")
+  endif()
+  set(D2D_ROOT ${CMAKE_CURRENT_BINARY_DIR}/decomp2d-build/downloadBuild2decomp-prefix/src/downloadBuild2decomp-build)
+  find_package(decomp2d REQUIRED
+	  PATHS ${D2D_ROOT})
+endif(decomp2d_FOUND)

cmake/decomp2d/downloadBuild2decomp.cmake.in

@@ -0,0 +1,26 @@
+# downloadBuild2decomp.cmake.in
+#
+cmake_minimum_required(VERSION 3.0.2)
+
+project(downloadBuild2decomp NONE)
+
+include(ExternalProject)
+
+ExternalProject_Add(downloadBuild2decomp
+    GIT_REPOSITORY    "https://github.com/xcompact3d/2decomp-fft"
+    GIT_TAG           "v2.0.3"
+    SOURCE_DIR        "${CMAKE_CURRENT_BINARY_DIR}/decomp2d-src"
+    INSTALL_DIR       "${CMAKE_CURRENT_BINARY_DIR}/decomp2d-opt"
+    TEST_COMMAND      ""
+)
+
+#ExternalProject_Add(downloadBuild2decomp
+#    GIT_REPOSITORY    "https://github.com/xcompact3d/2decomp-fft"
+#    GIT_TAG           "main"
+#    CONFIGURE_COMMAND "cmake -S ${CMAKE_CURRENT_BINARY_DIR}/decomp2d-src "
+#    BUILD_COMMAND     ""
+#    INSTALL_COMMAND   ""
+#    TEST_COMMAND      ""
+#    SOURCE_DIR        "${CMAKE_CURRENT_BINARY_DIR}/decomp2d-src"
+#    BINARY_DIR        ""
+#    INSTALL_DIR       ""

src/2decompfft/decomp.f90

@@ -0,0 +1,113 @@
+submodule(m_decomp) m_decomp_2decompfft
+!! Parallel decomposition provided by 2decomp&FFT
+
+  use mpi
+  implicit none
+
+  type, extends(decomp_t) :: decomp_2decompfft_t
+  contains
+    procedure :: decomposition => decomposition_2decompfft
+  end type
+
+  contains
+
+  module subroutine init_decomp(decomp)
+    class(decomp_t), allocatable, intent(out):: decomp
+
+    allocate(decomp_2decompfft_t :: decomp)
+  end subroutine
+
+  module subroutine decomposition_2decompfft(self, grid, par)
+    !! Performs 2D mesh decomposition using 2decomp&fft
+    use m_mesh_content, only: par_t, grid_t
+    use decomp_2d, only: decomp_2d_init, DECOMP_2D_COMM_CART_X, xsize, xstart
+    use decomp_2d_mpi, only: nrank, nproc
+
+    class(decomp_2decompfft_t) :: self
+    class(grid_t), intent(inout) :: grid
+    class(par_t), intent(inout) :: par 
+    integer :: p_col, p_row
+    integer, allocatable, dimension(:, :, :) :: global_ranks
+    integer, allocatable, dimension(:) :: global_ranks_lin
+    integer :: nproc
+    integer, dimension(3) :: subd_pos, subd_pos_prev, subd_pos_next
+    logical, dimension(3) :: periodic_bc
+    integer :: dir
+    logical :: is_last_domain
+    integer :: nx, ny, nz
+    integer :: ierr
+    integer :: cart_rank
+    integer, dimension(2) :: coords
+
+    if (par%is_root()) then
+      print*, "Domain decomposition by 2decomp&fft"
+    end if
+    nrank = par%nrank
+    nproc = par%nproc
+
+    nx = grid%global_cell_dims(1)
+    ny = grid%global_cell_dims(2)
+    nz = grid%global_cell_dims(3)
+
+    p_row = par%nproc_dir(2)
+    p_col = par%nproc_dir(3)
+    if (p_row*p_col /= par%nproc) then
+      error stop "Decomposition in X not supported by 2decomp&fft backend"
+    end if
+    periodic_bc(:) = grid%periodic_BC(:)
+    call decomp_2d_init(nx, ny, nz, p_row, p_col, periodic_bc)
+
+    ! Get global_ranks
+    allocate(global_ranks(1, p_row, p_col))
+    allocate(global_ranks_lin(p_row*p_col))
+    global_ranks_lin(:) = 0
+
+    call MPI_Comm_rank(DECOMP_2D_COMM_CART_X, cart_rank, ierr)
+    call MPI_Cart_coords(DECOMP_2D_COMM_CART_X, cart_rank, 2, coords, ierr)
+
+    global_ranks_lin(coords(1)+1 + p_row*(coords(2))) = par%nrank
+
+    call MPI_Allreduce(MPI_IN_PLACE, global_ranks_lin, p_row*p_col, MPI_INTEGER, MPI_SUM, MPI_COMM_WORLD, ierr)
+
+    global_ranks = reshape(global_ranks_lin, shape=[1, p_row, p_col])
+
+    ! subdomain position in the global domain
+    subd_pos = findloc(global_ranks, par%nrank)
+
+    ! local/directional position of the subdomain
+    par%nrank_dir(:) = subd_pos(:) - 1
+
+    ! Get local domain size and offset from 2decomp
+    grid%cell_dims(:) = xsize(:)
+    par%n_offset(:) = xstart(:)
+
+    ! compute vert_dims from cell_dims
+    do dir = 1, 3
+      is_last_domain = (par%nrank_dir(dir) + 1 == par%nproc_dir(dir))
+      if (is_last_domain .and. (.not. grid%periodic_BC(dir))) then
+        grid%vert_dims(dir) = grid%cell_dims(dir) +1
+      else
+        grid%vert_dims(dir) = grid%cell_dims(dir)
+      end if
+    end do
+
+    ! Get neighbour ranks
+    do dir = 1, 3
+      nproc = par%nproc_dir(dir)
+      subd_pos_prev(:) = subd_pos(:)
+      subd_pos_prev(dir) = modulo(subd_pos(dir) - 2, nproc) + 1
+      par%pprev(dir) = global_ranks(subd_pos_prev(1), &
+                                         subd_pos_prev(2), &
+                                         subd_pos_prev(3))
+
+      subd_pos_next(:) = subd_pos(:)
+      subd_pos_next(dir) = modulo(subd_pos(dir) - nproc, nproc) + 1
+      par%pnext(dir) = global_ranks(subd_pos_next(1), &
+                                         subd_pos_next(2), &
+                                         subd_pos_next(3))
+    end do
+
+  end subroutine decomposition_2decompfft
+
+
+end submodule

src/omp/poisson_fft.f90 → src/2decompfft/omp/poisson_fft.f90

@@ -1,9 +1,13 @@
 module m_omp_poisson_fft
+
+  use decomp_2d_constants, only: PHYSICAL_IN_X
+  use decomp_2d_fft, only: decomp_2d_fft_init, decomp_2d_fft_3d
   use m_allocator, only: field_t
   use m_common, only: dp
   use m_poisson_fft, only: poisson_fft_t
   use m_tdsops, only: dirps_t
   use m_mesh, only: mesh_t
+  use m_omp_spectral, only: process_spectral_div_u
 
   implicit none
 
@@ -35,26 +39,44 @@ function init(mesh, xdirps, ydirps, zdirps) result(poisson_fft)
 
     call poisson_fft%base_init(mesh, xdirps, ydirps, zdirps)
 
+    if (mesh%par%is_root()) then
+      print*, "Initialising 2decomp&fft"
+    end if
+
+    call decomp_2d_fft_init(PHYSICAL_IN_X)
+    allocate(poisson_fft%c_x(poisson_fft%nx_spec, poisson_fft%ny_spec, poisson_fft%nz_spec))
+
   end function init
 
   subroutine fft_forward_omp(self, f_in)
     implicit none
 
     class(omp_poisson_fft_t) :: self
     class(field_t), intent(in) :: f_in
+
+    call decomp_2d_fft_3d(f_in%data, self%c_x)
+
   end subroutine fft_forward_omp
 
   subroutine fft_backward_omp(self, f_out)
     implicit none
 
     class(omp_poisson_fft_t) :: self
     class(field_t), intent(inout) :: f_out
+
+    call decomp_2d_fft_3d(self%c_x, f_out%data)
+
   end subroutine fft_backward_omp
 
   subroutine fft_postprocess_omp(self)
     implicit none
 
     class(omp_poisson_fft_t) :: self
+
+    call process_spectral_div_u(self%c_x, self%waves, self%nx_spec, self%ny_spec, self%nz_spec, &
+      self%y_sp_st, self%nx_glob, self%ny_glob, self%nz_glob, &
+      self%ax, self%bx, self%ay, self%by, self%az, self%bz)
+
   end subroutine fft_postprocess_omp
 
 end module m_omp_poisson_fft

src/CMakeLists.txt

@@ -9,11 +9,16 @@ set(SRC
   solver.f90
   tdsops.f90
   time_integrator.f90
+  ordering.f90
+  mesh.f90
+  mesh_content.f90
+  decomp.f90
+  field.f90
   vector_calculus.f90
   omp/backend.f90
   omp/common.f90
   omp/kernels/distributed.f90
-  omp/poisson_fft.f90
+  omp/kernels/spectral_processing.f90
   omp/sendrecv.f90
   omp/exec_dist.f90
 )
@@ -31,21 +36,31 @@ set(CUDASRC
   cuda/sendrecv.f90
   cuda/tdsops.f90
 )
+set(2DECOMPFFTSRC
+  2decompfft/omp/poisson_fft.f90
+  2decompfft/decomp.f90
+)
+set(GENERICDECOMPSRC
+  decomp_generic.f90
+)
 
-if(${CMAKE_Fortran_COMPILER_ID} STREQUAL "PGI" OR
-   ${CMAKE_Fortran_COMPILER_ID} STREQUAL "NVHPC")
+if(${BACKEND} STREQUAL "CUDA")
   list(APPEND SRC ${CUDASRC})
 endif()
 
+if (${POISSON_SOLVER} STREQUAL "FFT" AND ${BACKEND} STREQUAL "OMP")
+  list(APPEND SRC ${2DECOMPFFTSRC})
+else()
+  list(APPEND SRC ${GENERICDECOMPSRC})
+endif()
+
 add_library(x3d2 STATIC ${SRC})
 target_include_directories(x3d2 INTERFACE ${CMAKE_CURRENT_BINARY_DIR})
 
 add_executable(xcompact xcompact.f90)
 target_link_libraries(xcompact PRIVATE x3d2)
 
-if(${CMAKE_Fortran_COMPILER_ID} STREQUAL "PGI" OR
-   ${CMAKE_Fortran_COMPILER_ID} STREQUAL "NVHPC")
-
+if(${BACKEND} STREQUAL "CUDA")
   set(CMAKE_Fortran_FLAGS "-cpp -cuda")
   set(CMAKE_Fortran_FLAGS_DEBUG "-g -O0 -traceback -Mbounds -Mchkptr -Ktrap=fp")
   set(CMAKE_Fortran_FLAGS_RELEASE "-O3 -fast")
@@ -59,11 +74,19 @@ elseif(${CMAKE_Fortran_COMPILER_ID} STREQUAL "GNU")
   set(CMAKE_Fortran_FLAGS_RELEASE "-O3 -ffast-math")
 endif()
 
+if (${POISSON_SOLVER} STREQUAL "FFT" AND ${BACKEND} STREQUAL "OMP")
+  message(STATUS "Using the FFT poisson solver with 2decomp&fft")
+  set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_SOURCE_DIR}/cmake")
+  find_package(decomp2d REQUIRED)
+  include_directories(${decomp2d_INCLUDE_DIRS})
+  target_link_libraries(decomp2d)
+  target_link_libraries(x3d2 PRIVATE decomp2d)
+endif()
+
 find_package(OpenMP REQUIRED)
 target_link_libraries(x3d2 PRIVATE OpenMP::OpenMP_Fortran)
 target_link_libraries(xcompact PRIVATE OpenMP::OpenMP_Fortran)
 
 find_package(MPI REQUIRED)
 target_link_libraries(x3d2 PRIVATE MPI::MPI_Fortran)
 target_link_libraries(xcompact PRIVATE MPI::MPI_Fortran)
-

src/allocator.f90

@@ -64,7 +64,7 @@ module m_allocator
 contains
 
   function allocator_init(mesh, sz) result(allocator)
-    type(mesh_t), target, intent(inout) :: mesh
+    class(mesh_t), target, intent(inout) :: mesh
     integer, intent(in) :: sz
     type(allocator_t) :: allocator
 

src/decomp.f90

@@ -0,0 +1,54 @@
+module m_decomp
+  use m_mesh_content, only: par_t, grid_t
+  implicit none
+
+  ! Object implementing the actual decomposition
+  type, abstract :: decomp_t 
+  contains
+    procedure(decomposition), public, deferred :: decomposition
+  end type decomp_t
+
+  ! Abstraction layer to allow different version of the decomposition computation
+  type, public :: decomp_mod_t
+    class(decomp_t), allocatable :: decomp
+    contains
+    procedure, public :: decomp_grid
+  end type
+
+  interface decomp_mod_t
+    module procedure init_decomp_mod
+  end interface
+
+  interface
+    module subroutine init_decomp(decomp)
+      class(decomp_t), allocatable, intent(out) :: decomp
+    end subroutine
+
+    module subroutine decomposition(self, grid, par)
+      use m_mesh_content, only: par_t, grid_t
+      class(decomp_t) :: self
+      class(grid_t), intent(inout) :: grid
+      class(par_t), intent(inout) :: par
+    end subroutine
+
+   end interface
+
+  contains
+
+  function init_decomp_mod() result(decomp)
+    type(decomp_mod_t) :: decomp
+
+    call init_decomp(decomp%decomp)
+
+  end function
+
+  module subroutine decomp_grid(self, grid, par)
+      class(decomp_mod_t) :: self
+      class(grid_t), intent(inout) :: grid
+      class(par_t), intent(inout) :: par
+
+      call self%decomp%decomposition(grid, par)
+  end subroutine
+
+
+end module m_decomp