Merge pull request xcompact3d#19 from Nanoseb/omp-tridiag-kernel

Add tridiagonal solver to omp backend

Author: semi-h

src/CMakeLists.txt

@@ -8,6 +8,8 @@ set(SRC
   omp/backend.f90
   omp/common.f90
   omp/kernels_dist.f90
+  omp/sendrecv.f90
+  omp/exec_dist.f90
 )
 set(CUDASRC
   cuda/backend.f90

src/omp/exec_dist.f90

@@ -0,0 +1,65 @@
+module m_omp_exec_dist
+   use mpi
+
+   use m_common, only: dp
+   use m_omp_common, only: SZ
+   use m_omp_kernels_dist, only: der_univ_dist, der_univ_subs
+   use m_tdsops, only: tdsops_t
+   use m_omp_sendrecv, only: sendrecv_fields
+
+   implicit none
+
+contains
+
+   subroutine exec_dist_tds_compact( &
+      du, u, u_recv_s, u_recv_e, du_send_s, du_send_e, du_recv_s, du_recv_e, &
+      tdsops, nproc, pprev, pnext, n_block &
+      )
+      implicit none
+
+      ! du = d(u)
+      real(dp), dimension(:, :, :), intent(out) :: du
+      real(dp), dimension(:, :, :), intent(in) :: u, u_recv_s, u_recv_e
+
+      ! The ones below are intent(out) just so that we can write data in them,
+      ! not because we actually need the data they store later where this
+      ! subroutine is called. We absolutely don't care about the data they pass back
+      real(dp), dimension(:, :, :), intent(out) :: &
+         du_send_s, du_send_e, du_recv_s, du_recv_e
+
+      type(tdsops_t), intent(in) :: tdsops
+      integer, intent(in) :: nproc, pprev, pnext
+      integer, intent(in) :: n_block
+
+      integer :: n_data
+      integer :: k
+
+      n_data = SZ*n_block
+
+      !$omp parallel do
+      do k = 1, n_block
+         call der_univ_dist( &
+            du(:, :, k), du_send_s(:, :, k), du_send_e(:, :, k), u(:, :, k), &
+            u_recv_s(:, :, k), u_recv_e(:, :, k), &
+            tdsops%coeffs_s, tdsops%coeffs_e, tdsops%coeffs, tdsops%n, &
+            tdsops%dist_fw, tdsops%dist_bw, tdsops%dist_af &
+            )
+      end do
+      !$omp end parallel do
+
+      ! halo exchange for 2x2 systems
+      call sendrecv_fields(du_recv_s, du_recv_e, du_send_s, du_send_e, &
+                           n_data, nproc, pprev, pnext)
+
+      !$omp parallel do
+      do k = 1, n_block
+         call der_univ_subs(du(:, :, k), &
+                            du_recv_s(:, :, k), du_recv_e(:, :, k), &
+                            tdsops%n, tdsops%dist_sa, tdsops%dist_sc)
+      end do
+      !$omp end parallel do
+
+   end subroutine exec_dist_tds_compact
+
+end module m_omp_exec_dist
+

src/omp/kernels_dist.f90

@@ -8,7 +8,7 @@ module m_omp_kernels_dist
 
 contains
 
-   subroutine der_univ_dist_omp( &
+   subroutine der_univ_dist( &
       du, send_u_s, send_u_e, u, u_s, u_e, coeffs_s, coeffs_e, coeffs, n, &
       ffr, fbc, faf &
       )
@@ -134,9 +134,9 @@ subroutine der_univ_dist_omp( &
       end do
       !$omp end simd
 
-   end subroutine der_univ_dist_omp
+   end subroutine der_univ_dist
 
-   subroutine der_univ_subs_omp(du, recv_u_s, recv_u_e, n, dist_sa, dist_sc)
+   subroutine der_univ_subs(du, recv_u_s, recv_u_e, n, dist_sa, dist_sc)
       implicit none
 
       ! Arguments
@@ -193,6 +193,6 @@ subroutine der_univ_subs_omp(du, recv_u_s, recv_u_e, n, dist_sa, dist_sc)
       end do
       !$omp end simd
 
-   end subroutine der_univ_subs_omp
+   end subroutine der_univ_subs
 
 end module m_omp_kernels_dist

src/omp/sendrecv.f90

@@ -0,0 +1,38 @@
+module m_omp_sendrecv
+   use mpi
+
+   use m_common, only: dp
+
+   implicit none
+
+contains
+
+   subroutine sendrecv_fields(f_recv_s, f_recv_e, f_send_s, f_send_e, &
+                              n_data, nproc, prev, next)
+      implicit none
+
+      real(dp), dimension(:, :, :), intent(out) :: f_recv_s, f_recv_e
+      real(dp), dimension(:, :, :), intent(in) :: f_send_s, f_send_e
+      integer, intent(in) :: n_data, nproc, prev, next
+
+      integer :: req(4), err(4), ierr, tag = 1234
+
+      if (nproc == 1) then
+         f_recv_s = f_send_e
+         f_recv_e = f_send_s
+      else
+         call MPI_Isend(f_send_s, n_data, MPI_DOUBLE_PRECISION, &
+                        prev, tag, MPI_COMM_WORLD, req(1), err(1))
+         call MPI_Irecv(f_recv_e, n_data, MPI_DOUBLE_PRECISION, &
+                        next, tag, MPI_COMM_WORLD, req(2), err(2))
+         call MPI_Isend(f_send_e, n_data, MPI_DOUBLE_PRECISION, &
+                        next, tag, MPI_COMM_WORLD, req(3), err(3))
+         call MPI_Irecv(f_recv_s, n_data, MPI_DOUBLE_PRECISION, &
+                        prev, tag, MPI_COMM_WORLD, req(4), err(4))
+
+         call MPI_Waitall(4, req, MPI_STATUSES_IGNORE, ierr)
+      end if
+
+   end subroutine sendrecv_fields
+
+end module m_omp_sendrecv