Fixed bug for z-pencils with z-implicit diffusion on CPUs. (CaNS-World#148)

Author: p-costa

src/solver.f90

@@ -7,7 +7,8 @@
 module mod_solver
   use, intrinsic :: iso_c_binding, only: C_PTR
   use decomp_2d
-  use mod_fft   , only: fft
+  use mod_common_mpi, only: ipencil_axis
+  use mod_fft       , only: fft
   use mod_types
   implicit none
   private
@@ -26,30 +27,32 @@ subroutine solver(n,ng,arrplan,normfft,lambdaxy,a,b,c,bc,c_or_f,p)
     character(len=1), dimension(0:1,3), intent(in) :: bc
     character(len=1), intent(in), dimension(3) :: c_or_f
     real(rp), intent(inout), dimension(0:,0:,0:) :: p
-    real(rp), dimension(xsize(1),xsize(2),xsize(3)) :: px
-    real(rp), dimension(ysize(1),ysize(2),ysize(3)) :: py
-    real(rp), dimension(zsize(1),zsize(2),zsize(3)) :: pz
+    real(rp), allocatable, dimension(:,:,:) :: px,py,pz
     integer :: q
     integer, dimension(3) :: n_z
     !
     n_z(:) = zsize(:)
-#if !defined(_DECOMP_Y) && !defined(_DECOMP_Z)
-    !$OMP PARALLEL WORKSHARE
-    px(:,:,:) = p(1:n(1),1:n(2),1:n(3))
-    !$OMP END PARALLEL WORKSHARE
-#elif defined(_DECOMP_Y)
-    !$OMP PARALLEL WORKSHARE
-    py(:,:,:) = p(1:n(1),1:n(2),1:n(3))
-    !$OMP END PARALLEL WORKSHARE
-    call transpose_y_to_x(py,px)
-#elif defined(_DECOMP_Z)
-    !$OMP PARALLEL WORKSHARE
-    pz(:,:,:) = p(1:n(1),1:n(2),1:n(3))
-    !$OMP END PARALLEL WORKSHARE
-    !call transpose_z_to_x(pz,px)
-    call transpose_z_to_y(pz,py)
-    call transpose_y_to_x(py,px)
-#endif
+    allocate(px(xsize(1),xsize(2),xsize(3)))
+    allocate(py(ysize(1),ysize(2),ysize(3)))
+    allocate(pz(zsize(1),zsize(2),zsize(3)))
+    select case(ipencil_axis)
+    case(1)
+      !$OMP PARALLEL WORKSHARE
+      px(:,:,:) = p(1:n(1),1:n(2),1:n(3))
+      !$OMP END PARALLEL WORKSHARE
+    case(2)
+      !$OMP PARALLEL WORKSHARE
+      py(:,:,:) = p(1:n(1),1:n(2),1:n(3))
+      !$OMP END PARALLEL WORKSHARE
+      call transpose_y_to_x(py,px)
+    case(3)
+      !$OMP PARALLEL WORKSHARE
+      pz(:,:,:) = p(1:n(1),1:n(2),1:n(3))
+      !$OMP END PARALLEL WORKSHARE
+      !call transpose_z_to_x(pz,px)
+      call transpose_z_to_y(pz,py)
+      call transpose_y_to_x(py,px)
+    end select
     call fft(arrplan(1,1),px) ! fwd transform in x
     !
     call transpose_x_to_y(px,py)
@@ -210,30 +213,31 @@ subroutine solver_gaussel_z(n,a,b,c,bcz,c_or_f,p)
     character(len=1), dimension(0:1), intent(in) :: bcz
     character(len=1), intent(in), dimension(3) :: c_or_f
     real(rp), intent(inout), dimension(0:,0:,0:) :: p
-    real(rp), allocatable, dimension(:,:,:), save :: px,py,pz
+    real(rp), allocatable, dimension(:,:,:) :: px,py,pz
     integer :: q
     integer, dimension(3) :: n_z
     logical :: is_no_decomp_z
     !
     n_z(:) = zsize(:)
-    is_no_decomp_z = xsize(3) == n_z(3) ! not decomposed along z: xsize(3) == ysize(3) == ng(3) when dims(2) = 1
+    is_no_decomp_z = xsize(3) == n_z(3).or.ipencil_axis == 3 ! not decomposed along z: xsize(3) == ysize(3) == ng(3) when dims(2) = 1
     if(.not.is_no_decomp_z) then
-      if(.not.allocated(px)) allocate(px(xsize(1),xsize(2),xsize(3)))
-      if(.not.allocated(py)) allocate(py(ysize(1),ysize(2),ysize(3)))
-      if(.not.allocated(pz)) allocate(pz(zsize(1),zsize(2),zsize(3)))
-#if !defined(_DECOMP_Y) && !defined(_DECOMP_Z)
-      !$OMP PARALLEL WORKSHARE
-      px(:,:,:) = p(1:n(1),1:n(2),1:n(3))
-      !$OMP END PARALLEL WORKSHARE
-      !call transpose_x_to_z(px,pz)
-      call transpose_x_to_y(px,py)
-      call transpose_y_to_z(py,pz)
-#elif defined(_DECOMP_Y)
-      !$OMP PARALLEL WORKSHARE
-      py(:,:,:) = p(1:n(1),1:n(2),1:n(3))
-      !$OMP END PARALLEL WORKSHARE
-      call transpose_y_to_z(py,pz)
-#endif
+      allocate(px(xsize(1),xsize(2),xsize(3)))
+      allocate(py(ysize(1),ysize(2),ysize(3)))
+      allocate(pz(zsize(1),zsize(2),zsize(3)))
+      select case(ipencil_axis)
+      case(1)
+        !$OMP PARALLEL WORKSHARE
+        px(:,:,:) = p(1:n(1),1:n(2),1:n(3))
+        !$OMP END PARALLEL WORKSHARE
+        !call transpose_x_to_z(px,pz)
+        call transpose_x_to_y(px,py)
+        call transpose_y_to_z(py,pz)
+      case(2)
+        !$OMP PARALLEL WORKSHARE
+        py(:,:,:) = p(1:n(1),1:n(2),1:n(3))
+        !$OMP END PARALLEL WORKSHARE
+        call transpose_y_to_z(py,pz)
+      end select
     end if
     q = 0
     if(c_or_f(3) == 'f'.and.bcz(1) == 'D') q = 1
@@ -244,27 +248,29 @@ subroutine solver_gaussel_z(n,a,b,c,bcz,c_or_f,p)
         call gaussel(         n_z(1),n_z(2),n_z(3)-q,0,a,b,c,pz)
       end if
     else
+      n_z(:) = n(:)
       if(bcz(0)//bcz(1) == 'PP') then
-        call gaussel_periodic(n(1),n(2),n(3)-q,1,a,b,c,p)
+        call gaussel_periodic(n_z(1),n_z(2),n_z(3)-q,1,a,b,c,p )
       else
-        call gaussel(         n(1),n(2),n(3)-q,1,a,b,c,p)
+        call gaussel(         n_z(1),n_z(2),n_z(3)-q,1,a,b,c,p )
       end if
     end if
     !
     if(.not.is_no_decomp_z) then
-#if !defined(_DECOMP_Y) && !defined(_DECOMP_Z)
-      !call transpose_z_to_x(pz,px)
-      call transpose_z_to_y(pz,py)
-      call transpose_y_to_x(py,px)
-      !$OMP PARALLEL WORKSHARE
-      p(1:n(1),1:n(2),1:n(3)) = px(:,:,:)
-      !$OMP END PARALLEL WORKSHARE
-#elif defined(_DECOMP_Y)
-      call transpose_z_to_y(pz,py)
-      !$OMP PARALLEL WORKSHARE
-      p(1:n(1),1:n(2),1:n(3)) = py(:,:,:)
-      !$OMP END PARALLEL WORKSHARE
-#endif
+      select case(ipencil_axis)
+      case(1)
+        !call transpose_z_to_x(pz,px)
+        call transpose_z_to_y(pz,py)
+        call transpose_y_to_x(py,px)
+        !$OMP PARALLEL WORKSHARE
+        p(1:n(1),1:n(2),1:n(3)) = px(:,:,:)
+        !$OMP END PARALLEL WORKSHARE
+      case(2)
+        call transpose_z_to_y(pz,py)
+        !$OMP PARALLEL WORKSHARE
+        p(1:n(1),1:n(2),1:n(3)) = py(:,:,:)
+        !$OMP END PARALLEL WORKSHARE
+      end select
     end if
   end subroutine solver_gaussel_z
   !

src/solver_gpu.f90

@@ -398,7 +398,7 @@ subroutine solver_gaussel_z_gpu(n,a,b,c,bcz,c_or_f,p)
     n_x(:) = ap_x%shape(:)
     n_y(:) = ap_y%shape(:)
     n_z(:) = ap_z%shape(:)
-    is_no_decomp_z = n_x(3) == n_z(3) ! not decomposed along z: xsize(3) == ysize(3) == ng(3) when dims(2) = 1
+    is_no_decomp_z = n_x(3) == n_z(3).or.ipencil_axis == 3 ! not decomposed along z: xsize(3) == ysize(3) == ng(3) when dims(2) = 1
     if(.not.is_no_decomp_z) then
       px(1:n_x(1),1:n_x(2),1:n_x(3)) => solver_buf_0(1:product(n_x(:)))
       if(cudecomp_is_t_in_place) then
@@ -444,21 +444,20 @@ subroutine solver_gaussel_z_gpu(n,a,b,c,bcz,c_or_f,p)
       end select
     end if
     !
-    if(ipencil_axis /= 3 .and. .not.is_no_decomp_z) then
-      q = 0
-      if(c_or_f(3) == 'f'.and.bcz(1) == 'D') q = 1
+    q = 0
+    if(c_or_f(3) == 'f'.and.bcz(1) == 'D') q = 1
+    if(.not.is_no_decomp_z) then
       if(bcz(0)//bcz(1) == 'PP') then
         call gaussel_periodic_gpu(n_z_0(1),n_z_0(2),n_z_0(3)-q,0,a,b,c,pz,work,pz_aux_1,pz_aux_2)
       else
         call gaussel_gpu(         n_z_0(1),n_z_0(2),n_z_0(3)-q,0,a,b,c,pz,work)
       end if
     else
-      q = 0
-      if(c_or_f(3) == 'f'.and.bcz(1) == 'D') q = 1
+      n_z_0(:) = n(:)
       if(bcz(0)//bcz(1) == 'PP') then
-        call gaussel_periodic_gpu(n(1),n(2),n(3)-q,1,a,b,c,p,work,pz_aux_1,pz_aux_2)
+        call gaussel_periodic_gpu(n_z_0(1),n_z_0(2),n_z_0(3)-q,1,a,b,c,p ,work,pz_aux_1,pz_aux_2)
       else
-        call gaussel_gpu(         n(1),n(2),n(3)-q,1,a,b,c,p,work)
+        call gaussel_gpu(         n_z_0(1),n_z_0(2),n_z_0(3)-q,1,a,b,c,p ,work)
       end if
     end if
     !