Homec4science
Diffusion wireles-gpu (master) Browse

wireles-gpu/src/press_stag_gpu.cuf
105cd9b56f1cmaster

Last Change
Raw File

press_stag_gpu.cuf
View Options

module poisson
use precision
use dimen
! Pressure solver
integer npsx,npsy,ni,nj,lsize
parameter(npsx=1, npsy=nz/(nz2-2))
parameter(ni=(nx/2+1)/npsx, nj=ny/npsy)
parameter(lsize=ni*nj*nzb)
contains
subroutine press_stag_gpu(p,rhs_p,dpdx,dpdy,me,nall,t,plan)
!---------------------------------------------------------------------------
! declaration
!---------------------------------------------------------------------------
use cudafor
use cufft_m
implicit none
integer,intent(in) :: t,me,nall
real(fp),dimension(nx,ny,nz2),intent(in),device:: rhs_p
real(fp),dimension(nx,ny,nz2),intent(inout),device:: p,dpdx,dpdy
real(fp),dimension(nx,ny,nz2),device:: p_test
real(fp),dimension(nx,ny,nz2) :: p_host
complex(fp), dimension(:,:,:),allocatable,device:: h_x,h_2,h_3
complex(fp), dimension(nx/2+1,ny,nzb) :: h_x_host
complex(fp), dimension(:,:,:),allocatable,device:: dpdx_hat,dpdy_hat
complex(fp), dimension(ni,nj,nz) :: h_2_host,h_3_host
integer i_global,j_global
integer i,j,k,ii,jj,size,istat
type(c_ptr) plan(2)
! Cuda variable
type(dim3) :: grid, tBlock, grid1,tBlock1, grid_3d,tBlock_3d
grid= dim3((nx/2)/16,nj/16,1)
tBlock= dim3(16,16,1)
tBlock1 = dim3(16,16,1)
grid1 = dim3((nx/2)/16,ny/16,1) ! Notice the difference between nj and ny
tBlock_3d = dim3(16,16,2)
grid_3d = dim3((nx/2)/16,ny/16,nzb/2) ! Notice the difference between nzb and nz2
save h_x,h_2,h_3,dpdx_hat,dpdy_hat
!---------------------------------------------------------------------------
! main code
!---------------------------------------------------------------------------
!---------------------------------------------------------------------------
! init
if(t.eq.1)then
! Allocate internal variable
allocate(h_x(nx/2+1,ny,nzb))
allocate(dpdx_hat(nx/2+1,ny,nzb))
allocate(dpdy_hat(nx/2+1,ny,nzb))
allocate(h_2(ni,nj,nz))
allocate(h_3(ni,nj,nz))
endif
! ! !---------------------------------------------------------------------------
! ! ! compute forward fft interior
! ! print *,rhs_p(:,3,3)
call cufftExec(plan(1),rhs_p(:,:,2:nzb+1),h_x)
! !---------------------------------------------------------------------------
! ! redistribute mpi
call slab_transpose_z2y<<<grid1, tBlock1>>>(h_x,h_3)
! ! call mpi_alltoall(h_3(1,1,1),lsize,mpi_double_complex(fp), &
! ! h_2(1,1,1),lsize,mpi_double_complex(fp),nall,ierr)
! h_2 = h_3
!---------------------------------------------------------------------------
! solve tri-diag matrix
call solve_tridiag<<<grid, tBlock>>>(h_3,me)
!-------------------------------------------------------------------------
! redistribute mpi
! h_3 = h_2
call slab_transpose_y2z<<<grid, tBlock>>>(h_3,h_x)
! h_3_host = h_3
! !-------------------------------------------------------------------------
! ! compute p_hat (i.e h_z), dpdx_hat, dpdy_hat
! ! x
call press_grad<<<grid_3d, tBlock_3d>>>(h_x,dpdx_hat,dpdy_hat)
call cufftExec(plan(2),dpdx_hat,dpdx(:,:,2:nzb+1))
call cufftExec(plan(2),dpdy_hat,dpdy(:,:,2:nzb+1))
call cufftExec(plan(2),h_x,p(:,:,2:nzb+1))
end subroutine press_stag_gpu
attributes(global) subroutine slab_transpose_z2y(h_x,h_3)
! Transpose z slab to y slab
implicit none
complex(fp), dimension(nx/2+1,ny,nzb),intent(in) :: h_x
complex(fp), dimension(ni,nj,nz),intent(inout) :: h_3
integer :: i,j,k,i_block,j_block
i = (blockIdx%x - 1) * blockDim%x + threadIdx%x;
j = (blockIdx%y - 1) * blockDim%y + threadIdx%y;
do j_block=1,npsy
if (i==nx/2) h_3(nx/2+1,j,(j_block-1)*nzb+1:j_block*nzb)= h_x(nx/2+1,j+(j_block-1)*nj,:)*inxny
h_3(i,j,(j_block-1)*nzb+1:j_block*nzb)= h_x(i,j+(j_block-1)*nj,:)*inxny
end do
end subroutine
attributes(global) subroutine slab_transpose_y2z(h_3,h_x)
! Transpose y slab to z slab
implicit none
complex(fp), dimension(nx/2+1,ny,nzb),intent(inout) :: h_x
complex(fp), dimension(ni,nj,nz),intent(in) :: h_3
integer :: i,j,k,i_block,j_block
i = (blockIdx%x - 1) * blockDim%x + threadIdx%x;
j = (blockIdx%y - 1) * blockDim%y + threadIdx%y;
do j_block=1,npsy
if (i==nx/2) h_x(nx/2+1,j+(j_block-1)*nj,:) = h_3(nx/2+1,j,(j_block-1)*nzb+1:j_block*nzb)
h_x(i,j+(j_block-1)*nj,:) = h_3(i,j,(j_block-1)*nzb+1:j_block*nzb)
end do
end subroutine
attributes(global) subroutine solve_tridiag(h_2,device_id)
! Solve tridiag system
implicit none
complex(fp), dimension(ni,nj,nz),intent(inout) :: h_2
real(fp),dimension(nz+1) :: rhs_col,a1,b1,c1,p_colr,p_coli
integer :: i,j,k,i_global,j_global,ii,jj
integer,value,intent(in) :: device_id
i = (blockIdx%x - 1) * blockDim%x + threadIdx%x;
j = (blockIdx%y - 1) * blockDim%y + threadIdx%y;
! TODO: JUST FOR ONE GPU!!!!!!
j_global = j + device_id*nj
jj = j_global - 1
i_global=i! + (mei-1)*ni
ii=i_global-1
if(jj.gt.(ny/2)) jj=jj-ny
jj=jj*l_r
! Assumble the tridiag matrix
call mat_assemble(a1,b1,c1,ii,jj)
! RHS: real(fp) part
rhs_col(1)=0._fp
rhs_col(2:nz)=real(h_2(i,j,1:nz-1))
rhs_col(nz+1)=0._fp
! Solve the tridiag
call tridag_gpu(a1,b1,c1,rhs_col,p_colr,nz+1)
! RHS: aimag part
rhs_col(1)=0._fp
rhs_col(2:nz)=aimag(h_2(i,j,1:nz-1))
rhs_col(nz+1)=0._fp
! Solve the tridiag
call tridag_gpu(a1,b1,c1,rhs_col,p_coli,nz+1)
h_2(i,j,1:nz)=dcmplx(p_colr(2:nz+1),p_coli(2:nz+1))
! Handling the leftover edge (nx/2+1)
! I don't know if there is any elegant way to do it
if (i == 1) then
call mat_assemble(a1,b1,c1,nx/2,jj)
rhs_col(1)=0._fp!1.0*dz*real(bottomw(i_global,j_global))
rhs_col(2:nz)=real(h_2(nx/2+1,j,1:nz-1))
rhs_col(nz+1)=0._fp!1.0*real(topw(i_global,j_global))*dz
! Solve the tridiag
call tridag_gpu (a1,b1,c1,rhs_col,p_colr,nz+1)
! RHS: aimag part
rhs_col(1)=0._fp!1.0*dz*aimag(bottomw(i_global,j_global))
rhs_col(2:nz)=aimag(h_2(nx/2+1,j,1:nz-1))
rhs_col(nz+1)=0._fp!1.0*aimag(topw(i_global,j_global))*dz
! Solve the tridiag
call tridag_gpu (a1,b1,c1,rhs_col,p_coli,nz+1)
h_2(nx/2+1,j,1:nz)=dcmplx(p_colr(2:nz+1),p_coli(2:nz+1))
endif
end subroutine
attributes(device) SUBROUTINE mat_assemble(a1,b1,c1,ii,jj)
! Assemble the tridiag matrix along the z pencil
implicit none
INTEGER ii,jj,k
real(fp) a1(nz+1),b1(nz+1),c1(nz+1)
! bot
a1(1)=0._fp
b1(1)=-1.0
c1(1)=1.0
if ((ii.eq.0).and.(jj.eq.0)) then
a1(1)=0._fp
b1(1)=1.0
c1(1)=0._fp
end if
! interior
do k=2,nz
a1(k)=1.0/(dz**2)
b1(k)=(-ii*ii-jj*jj-2.0/(dz**2))
c1(k)=1.0/(dz**2)
end do
! top
a1(nz+1)= -1.0
b1(nz+1)= 1.0
c1(nz+1)= 0._fp
return
END SUBROUTINE mat_assemble
attributes(device) SUBROUTINE tridag_gpu(a,b,c,r,u,n)
INTEGER*4 n,NMAX
real(fp) a(n),b(n),c(n),r(n),u(n)
PARAMETER (NMAX=500)
INTEGER*4 j
real(fp) bet,gam(NMAX)
! if(b(1).eq.0.)pause !'tridag: rewrite equations'
bet=b(1)
u(1)=r(1)/bet
do 11 j=2,n
gam(j)=c(j-1)/bet
bet=b(j)-a(j)*gam(j)
! if(bet.eq.0.) then
! print *, 'tridag failed at k=',j
! print *,'a, b, c, gam, and bet=',a(j),b(j),c(j),gam(j),bet
! pause
! end if
u(j)=(r(j)-a(j)*u(j-1))/bet
11 continue
do 12 j=n-1,1,-1
u(j)=u(j)-gam(j+1)*u(j+1)
12 continue
return
END SUBROUTINE tridag_gpu
attributes(global) subroutine press_grad(f_hat,dfdx_hat,dfdy_hat)
implicit none
real(fp),dimension(nz2) :: a,b,c
complex(fp), dimension(nx/2+1,ny,nzb),intent(inout) :: f_hat
complex(fp), dimension(nx/2+1,ny,nzb),intent(inout) :: dfdx_hat,dfdy_hat
integer :: i,j,k,ii,jj
i = (blockIdx%x - 1) * blockDim%x + threadIdx%x;
j = (blockIdx%y - 1) * blockDim%y + threadIdx%y;
k = (blockIdx%z - 1) * blockDim%z + threadIdx%z;
! ! Shift wave number
ii = i - 1
jj = j - 1
if(jj.gt.nint(ny/2.)) jj=jj-ny
jj=jj*l_r
! ! dfdx_hat
dfdx_hat(i,j,k)=dcmplx(aimag(f_hat(i,j,k))*(-1.0), &
real(f_hat(i,j,k)))*ii
! ! dfdy_hat
dfdy_hat(i,j,k)=dcmplx(aimag(f_hat(i,j,k))*(-1.0), &
real(f_hat(i,j,k)))*jj
if (j == ny/2+1) then
dfdy_hat(i,ny/2+1,:)=dcmplx(0._fp)
end if
if (i == nx/2) then
dfdx_hat(nx/2+1,j,k)=dcmplx(0._fp)
dfdy_hat(nx/2+1,j,k)=dcmplx(aimag(f_hat(nx/2+1,j,k))*(-1.0), &
real(f_hat(nx/2+1,j,k)))*jj
dfdy_hat(nx/2+1,ny/2+1,:)=dcmplx(0._fp)
end if
end subroutine press_grad
attributes(global) subroutine velo_corr(u,dpdi)
implicit none
real(fp), dimension(nx,ny,nz2),intent(inout) :: u
real(fp), dimension(nx,ny,nz2),intent(in) :: dpdi
integer :: i,j,k
i = (blockIdx%x - 1) * blockDim%x + threadIdx%x;
j = (blockIdx%y - 1) * blockDim%y + threadIdx%y;
k = (blockIdx%z - 1) * blockDim%z + threadIdx%z;
if (k>1 .and. k<nz2) u(i,j,k)= u(i,j,k) - dpdi(i,j,k)*dt
end subroutine velo_corr
end module poisson
c4science · Help