fix: Allow dist. algorithm to solve TDSs with non-uniform off diagonals.

Author: semi-h

src/cuda/kernels_dist.f90

@@ -9,70 +9,69 @@ module m_cuda_kernels_dist
 
    attributes(global) subroutine der_univ_dist( &
       du, send_u_b, send_u_e, u, u_b, u_e, coeffs_b, coeffs_e, coeffs, n, &
-      alfa, ffr, fbc &
+      ffr, fbc, faf &
       )
       implicit none
 
       ! Arguments
       real(dp), device, intent(out), dimension(:, :, :) :: du, send_u_b, &
                                                            send_u_e
       real(dp), device, intent(in), dimension(:, :, :) :: u, u_b, u_e
-      real(dp), device, intent(in), dimension(:) :: ffr, fbc
       real(dp), device, intent(in), dimension(:, :) :: coeffs_b, coeffs_e
       real(dp), device, intent(in), dimension(:) :: coeffs
-      real(dp), value, intent(in) :: alfa
       integer, value, intent(in) :: n
+      real(dp), device, intent(in), dimension(:) :: ffr, fbc, faf
 
       ! Local variables
       integer :: i, j, b, k, lj
       integer :: jm2, jm1, jp1, jp2
-      integer :: n_s, n_m, n_b, n_e !stencil, middle, begin, end
 
-      real(dp) :: temp_du, c_m4, c_m3, c_m2, c_m1, c_j, c_p1, c_p2, c_p3, c_p4
+      real(dp) :: c_m4, c_m3, c_m2, c_m1, c_j, c_p1, c_p2, c_p3, c_p4, &
+                  temp_du, alpha, last_r
 
       i = threadIdx%x
       b = blockIdx%x
 
-      n_s = (size(coeffs)-1)/2
-      n_m = size(coeffs)
-      n_b = size(coeffs_b, dim=2)
-      n_e = size(coeffs_e, dim=2)
-
       ! store bulk coeffs in the registers
       c_m4 = coeffs(1); c_m3 = coeffs(2); c_m2 = coeffs(3); c_m1 = coeffs(4)
       c_j = coeffs(5)
       c_p1 = coeffs(6); c_p2 = coeffs(7); c_p3 = coeffs(8); c_p4 = coeffs(9)
+      last_r = ffr(1)
 
       du(i, 1, b) = coeffs(1)*u_b(i, 1, b) + coeffs(2)*u_b(i, 2, b) &
                     + coeffs(3)*u_b(i, 3, b) + coeffs(4)*u_b(i, 4, b) &
                     + coeffs(5)*u(i, 1, b) &
                     + coeffs(6)*u(i, 2, b) + coeffs(7)*u(i, 3, b) &
                     + coeffs(8)*u(i, 4, b) + coeffs(9)*u(i, 5, b)
+      du(i, 1, b) = du(i, 1, b)*faf(1)
       du(i, 2, b) = coeffs(1)*u_b(i, 2, b) + coeffs(2)*u_b(i, 3, b) &
                     + coeffs(3)*u_b(i, 4, b) + coeffs(4)*u(i, 1, b) &
                     + coeffs(5)*u(i, 2, b) &
                     + coeffs(6)*u(i, 3, b) + coeffs(7)*u(i, 4, b) &
                     + coeffs(8)*u(i, 5, b) + coeffs(9)*u(i, 6, b)
+      du(i, 2, b) = du(i, 2, b)*faf(2)
       du(i, 3, b) = coeffs(1)*u_b(i, 3, b) + coeffs(2)*u_b(i, 4, b) &
                     + coeffs(3)*u(i, 1, b) + coeffs(4)*u(i, 2, b) &
                     + coeffs(5)*u(i, 3, b) &
                     + coeffs(6)*u(i, 4, b) + coeffs(7)*u(i, 5, b) &
                     + coeffs(8)*u(i, 6, b) + coeffs(9)*u(i, 7, b)
-      du(i, 3, b) = ffr(3)*(du(i, 3, b) - alfa*du(i, 2, b))
+      du(i, 3, b) = ffr(3)*(du(i, 3, b) - faf(3)*du(i, 2, b))
       du(i, 4, b) = coeffs(1)*u_b(i, 4, b) + coeffs(2)*u(i, 1, b) &
                     + coeffs(3)*u(i, 2, b) + coeffs(4)*u(i, 3, b) &
                     + coeffs(5)*u(i, 4, b) &
                     + coeffs(6)*u(i, 5, b) + coeffs(7)*u(i, 6, b) &
                     + coeffs(8)*u(i, 7, b) + coeffs(9)*u(i, 8, b)
-      du(i, 4, b) = ffr(4)*(du(i, 4, b) - alfa*du(i, 3, b))
+      du(i, 4, b) = ffr(4)*(du(i, 4, b) - faf(3)*du(i, 3, b))
+
+      alpha = faf(5)
 
-      do j = n_s+1, n-n_s
+      do j = 5, n-4
          temp_du = c_m4*u(i, j-4, b) + c_m3*u(i, j-3, b) &
                  + c_m2*u(i, j-2, b) + c_m1*u(i, j-1, b) &
                  + c_j*u(i, j, b) &
                  + c_p1*u(i, j+1, b) + c_p2*u(i, j+2, b) &
                  + c_p3*u(i, j+3, b) + c_p4*u(i, j+4, b)
-         du(i, j, b) = ffr(j)*(temp_du - alfa*du(i, j-1, b))
+         du(i, j, b) = ffr(j)*(temp_du - alpha*du(i, j-1, b))
       end do
 
       j = n-3
@@ -81,49 +80,48 @@ attributes(global) subroutine der_univ_dist( &
                   + coeffs(5)*u(i, j, b) &
                   + coeffs(6)*u(i, j+1, b) + coeffs(7)*u(i, j+2, b) &
                   + coeffs(8)*u(i, j+3, b) + coeffs(9)*u_e(i, 1, b)
-      du(i, j, b) = ffr(j)*(du(i, j, b) - alfa*du(i, j-1, b))
+      du(i, j, b) = ffr(j)*(du(i, j, b) - faf(j)*du(i, j-1, b))
       j = n-2
       du(i, j, b) = coeffs(1)*u(i, j-4, b) + coeffs(2)*u(i, j-3, b) &
                     + coeffs(3)*u(i, j-2, b) + coeffs(4)*u(i, j-1, b) &
                     + coeffs(5)*u(i, j, b) &
                     + coeffs(6)*u(i, j+1, b) + coeffs(7)*u(i, j+2, b) &
                     + coeffs(8)*u_e(i, 1, b) + coeffs(9)*u_e(i, 2, b)
-      du(i, j, b) = ffr(j)*(du(i, j, b) - alfa*du(i, j-1, b))
+      du(i, j, b) = ffr(j)*(du(i, j, b) - faf(j)*du(i, j-1, b))
       j = n-1
       du(i, j, b) = coeffs(1)*u(i, j-4, b) + coeffs(2)*u(i, j-3, b) &
                     + coeffs(3)*u(i, j-2, b) + coeffs(4)*u(i, j-1, b) &
                     + coeffs(5)*u(i, j, b) &
                     + coeffs(6)*u(i, j+1, b) + coeffs(7)*u_e(i, 1, b) &
                     + coeffs(8)*u_e(i, 2, b) + coeffs(9)*u_e(i, 3, b)
-      du(i, j, b) = ffr(j)*(du(i, j, b) - alfa*du(i, j-1, b))
+      du(i, j, b) = ffr(j)*(du(i, j, b) - faf(j)*du(i, j-1, b))
       j = n
       du(i, j, b) = coeffs(1)*u(i, j-4, b) + coeffs(2)*u(i, j-3, b) &
                     + coeffs(3)*u(i, j-2, b) + coeffs(4)*u(i, j-1, b) &
                     + coeffs(5)*u(i, j, b) &
                     + coeffs(6)*u_e(i, 1, b) + coeffs(7)*u_e(i, 2, b) &
                     + coeffs(8)*u_e(i, 3, b) + coeffs(9)*u_e(i, 4, b)
-      du(i, j, b) = ffr(j)*(du(i, j, b) - alfa*du(i, j-1, b))
+      du(i, j, b) = ffr(j)*(du(i, j, b) - faf(j)*du(i, j-1, b))
 
       send_u_e(i, 1, b) = du(i, n, b)
 
       ! Backward pass of the hybrid algorithm
       do j = n - 2, 2, -1
          du(i, j, b) = du(i, j, b) - fbc(j)*du(i, j + 1, b)
       end do
-      du(i, 1, b) = ffr(1)*(du(i, 1, b) - fbc(1)*du(i, 2, b))
+      du(i, 1, b) = last_r*(du(i, 1, b) - fbc(1)*du(i, 2, b))
       send_u_b(i, 1, b) = du(i, 1, b)
 
    end subroutine der_univ_dist
 
    attributes(global) subroutine der_univ_subs(du, recv_u_b, recv_u_e, &
-                                               n, alfa, dist_sa, dist_sc)
+                                               n, dist_sa, dist_sc)
       implicit none
 
       ! Arguments
       real(dp), device, intent(out), dimension(:, :, :) :: du
       real(dp), device, intent(in), dimension(:, :, :) :: recv_u_b, recv_u_e
       real(dp), device, intent(in), dimension(:) :: dist_sa, dist_sc
-      real(dp), value, intent(in) :: alfa
       integer, value, intent(in) :: n
 
       ! Local variables

src/derparams.f90

@@ -10,21 +10,20 @@ subroutine der_1_vv()
    end subroutine der_1_vv
 
    subroutine der_2_vv(coeffs, coeffs_b, coeffs_e, &
-                       dist_fr, dist_bc, dist_sa, dist_sc, &
-                       n_halo, alfa, dx2, n, bcond)
+                       dist_fr, dist_bc, dist_af, dist_sa, dist_sc, &
+                       n_halo, dx2, n, bcond)
       implicit none
 
       real(dp), allocatable, dimension(:), intent(out) :: coeffs, &
-         dist_fr, dist_bc, dist_sa, dist_sc
+         dist_fr, dist_bc, dist_af, dist_sa, dist_sc
       real(dp), allocatable, dimension(:,:), intent(out) :: coeffs_b, coeffs_e
       integer, intent(out) :: n_halo
-      real(dp), intent(out) :: alfa
       real(dp), intent(in) :: dx2
       integer, intent(in) :: n
       character(len=*), intent(in) :: bcond
 
       real(dp), allocatable :: dist_b(:)
-      real(dp) :: asi, bsi, csi, dsi
+      real(dp) :: alfa, asi, bsi, csi, dsi
       integer :: i, n_stencil
 
       allocate(dist_sa(n), dist_sc(n), dist_b(n))
@@ -55,14 +54,16 @@ subroutine der_2_vv(coeffs, coeffs_b, coeffs_e, &
          print*, 'Boundary condition is not recognized :', bcond
       end select
 
-      call process_dist(dist_fr, dist_bc, dist_sa, dist_sc, dist_b, n)
+      call process_dist(dist_fr, dist_bc, dist_af, dist_sa, dist_sc, dist_b, n)
 
    end subroutine der_2_vv
 
-   subroutine process_dist(dist_fr, dist_bc, dist_sa, dist_sc, dist_b, n)
+   subroutine process_dist(dist_fr, dist_bc, dist_af, &
+                           dist_sa, dist_sc, dist_b, n)
       implicit none
 
-      real(dp), allocatable, dimension(:), intent(out) :: dist_fr, dist_bc
+      real(dp), allocatable, dimension(:), intent(out) :: dist_fr, dist_bc, &
+                                                          dist_af
       real(dp), dimension(:), intent(inout) :: dist_sa, dist_sc, dist_b
       integer, intent(in) :: n
 
@@ -71,7 +72,8 @@ subroutine process_dist(dist_fr, dist_bc, dist_sa, dist_sc, dist_b, n)
       m = n
       nrank = 0; nproc = 1
 
-      allocate(dist_fr(n), dist_bc(n))
+      ! forward factors, backward factors, and auxiliary factor
+      allocate(dist_fr(n), dist_bc(n), dist_af(n))
 
       do nrank = 0, nproc-1
 
@@ -80,9 +82,11 @@ subroutine process_dist(dist_fr, dist_bc, dist_sa, dist_sc, dist_b, n)
             dist_sa(i) = dist_sa(i)/dist_b(i)
             dist_sc(i) = dist_sc(i)/dist_b(i)
             dist_bc(i) = dist_sc(i)
+            dist_af(i) = 1._dp/dist_b(i)
          end do
          do i = 3+m*nrank, m+m*nrank
             dist_fr(i) = 1.d0/(dist_b(i)-dist_sa(i)*dist_sc(i-1))
+            dist_af(i) = dist_sa(i)
             dist_sa(i) = -dist_fr(i)*dist_sa(i)*dist_sa(i-1)
             dist_sc(i) = dist_fr(i)*dist_sc(i)
             !dist_bc(i) = dist_sc(i)
@@ -92,7 +96,7 @@ subroutine process_dist(dist_fr, dist_bc, dist_sa, dist_sc, dist_b, n)
             dist_bc(i) = dist_sc(i)
             dist_sc(i) = -dist_sc(i)*dist_sc(i+1)
          end do
-         ! this is not good
+         ! dist_fr(1) is never used, so store this extra factor instead.
          dist_fr(1+m*nrank) = 1.d0/(1.d0-dist_sc(1+m*nrank)*dist_sa(2+m*nrank))
 
          dist_sa(1+m*nrank) = dist_fr(1+m*nrank)*dist_sa(1+m*nrank)