Skip to content
Snippets Groups Projects
Commit f5613bcc authored by Juan Escobar's avatar Juan Escobar
Browse files

Juan 16/11/2021:ZSOLVER, add routine zt1d_discretisation_allocate3d , for fast...

Juan 16/11/2021:ZSOLVER, add routine zt1d_discretisation_allocate3d , for fast GPU memory allocation in MutiGrid
parent 88f14ec3
No related branches found
No related tags found
No related merge requests found
......@@ -135,6 +135,29 @@ private
! Stoarge for vertical coefficients
type(vertical_coefficients) :: vert_coeff
type Temp_jacobi
real(kind=rl), pointer , contiguous :: r(:)
real(kind=rl), pointer , contiguous :: c(:), utmp(:)
real(kind=rl), pointer , contiguous :: u0(:,:,:)
real(kind=rl), pointer , contiguous :: ut0(:,:,:)
type(scalar3d) , pointer :: Sr,Sc,Sut0,Sutmp
end type Temp_jacobi
integer , parameter :: max_lev = 128
type (Temp_jacobi) , save , dimension(max_lev) :: Tjacobi
real(kind=rl), pointer , contiguous , dimension(:) :: zt1d_discretisation
INTEGER :: zt1d_discretisation_size = 10000000
INTEGER , parameter :: zt1d_discretisation_size_factor = 2 * 4
INTEGER :: nt1d_discretisation_top = 0 , nt1d_discretisation_max = 0
INTEGER , ALLOCATABLE, DIMENSION (:) :: nt1d_discretisation_pointer , nt1d_discretisation_pointer_ksize
INTEGER , parameter :: zt1d_discretisation_pointer_size = 128
INTEGER :: nt1d_discretisation_pointer_top = 0 , nt1d_discretisation_pointer_max = 0
logical :: gfirst_call_zt1d_discretisation = .true.
public::discretisation_initialise_mnh
public::discretisation_initialise
public::discretisation_finalise
......@@ -158,12 +181,91 @@ public::DIRECTION_FORWARD
public::DIRECTION_BACKWARD
public::ORDERING_LEX
public::ORDERING_RB
public::zt1d_discretisation_init
contains
!==================================================================
! Initialise module
!==================================================================
subroutine zt1d_discretisation_init(kiu,kju,kku)
implicit none
integer , optional , intent(in) :: kiu,kju,kku ! size of 1 3D MNH array
if (gfirst_call_zt1d_discretisation) then
!
gfirst_call_zt1d_discretisation = .false.
!
if (present(kiu)) then
zt1d_discretisation_size = kiu*kju*kku * zt1d_discretisation_size_factor
end if
allocate(zt1d_discretisation(zt1d_discretisation_size))
zt1d_discretisation = 0.0
!$acc enter data copyin(zt1d_discretisation)
!
allocate(nt1d_discretisation_pointer(zt1d_discretisation_pointer_size))
nt1d_discretisation_pointer = 0.0
allocate(nt1d_discretisation_pointer_ksize(zt1d_discretisation_pointer_size))
nt1d_discretisation_pointer_ksize = 0.0
!
end if
end subroutine zt1d_discretisation_init
function zt1d_discretisation_allocate3d(ptab3d,kib,kie,kjb,kje,kkb,kke) result(kindex)
implicit none
real(kind=rl), pointer , contiguous , dimension(:,:,:) :: ptab3d
!real(kind=rl), pointer , dimension(:,:,:) :: ptab3d
integer , intent(in) :: kib,kie,kjb,kje,kkb,kke
integer :: kindex
! local var
integer :: ksize, kbeg, kend
if ( nt1d_discretisation_pointer_top == zt1d_discretisation_pointer_size ) then
print*,"ERROR zt1d_discretisation_init:: zt1d_discretisation_pointer_size to small=",zt1d_discretisation_pointer_size
print*,"ERROR zt1d_discretisation_init:: Augment it "
stop
end if
ksize = (kie-kib+1)*(kje-kjb+1)*(kke-kkb+1)
kbeg = nt1d_discretisation_top + 1
kend = nt1d_discretisation_top + ksize
if ( kend > zt1d_discretisation_size ) then
print*,"ERROR zt1d_discretisation_allocate3d:: zt1d_discretisation_size to small=",zt1d_discretisation_size
print*,"ERROR zt1d_discretisation_allocate3d:: Augment it "
stop
end if
ptab3d(kib:kie,kjb:kje,kkb:kke) => zt1d_discretisation(kbeg:)
nt1d_discretisation_pointer_top = nt1d_discretisation_pointer_top + 1
nt1d_discretisation_top = kend
nt1d_discretisation_pointer(nt1d_discretisation_pointer_top) = kend
nt1d_discretisation_pointer_ksize(nt1d_discretisation_pointer_top) = ksize
if ( nt1d_discretisation_pointer_top > nt1d_discretisation_pointer_max ) then
nt1d_discretisation_pointer_max = nt1d_discretisation_pointer_top
print*,"zt1d_discretisation_allocate3d:: nt1d_discretisation_pointer_max=",nt1d_discretisation_pointer_max
call flush(6)
endif
if ( nt1d_discretisation_top > nt1d_discretisation_max ) then
nt1d_discretisation_max = nt1d_discretisation_top
print*,"zt1d_discretisation_allocate3d:: nt1d_discretisation_max=",nt1d_discretisation_max
call flush(6)
endif
end function zt1d_discretisation_allocate3d
subroutine discretisation_initialise_mnh(grid_param_in, &
model_param_in, &
smoother_param_in, &
......@@ -1326,9 +1428,8 @@ end subroutine construct_vertical_coeff
integer :: iib,iie,ijb,ije
real , dimension(:,:,:) , pointer , contiguous :: zSut0_st , zu_st
real , dimension(:,:,:) , pointer , contiguous :: zSut0_st , zu_st, zSr_st, zSutmp_st
integer , parameter :: max_lev = 128
logical , save :: Ljacob_first_call = .true.
logical , save , dimension(max_lev) :: Ljacob_first_call_level = .true.
......@@ -1338,15 +1439,9 @@ end subroutine construct_vertical_coeff
real(kind=rl), pointer , contiguous :: ut0(:,:,:)
type(scalar3d) , pointer :: Sr,Sc,Sut0,Sutmp
type Temp_jacobi
real(kind=rl), pointer , contiguous :: r(:)
real(kind=rl), pointer , contiguous :: c(:), utmp(:)
real(kind=rl), pointer , contiguous :: u0(:,:,:)
real(kind=rl), pointer , contiguous :: ut0(:,:,:)
type(scalar3d) , pointer :: Sr,Sc,Sut0,Sutmp
end type Temp_jacobi
type (Temp_jacobi) , save , dimension(max_lev) :: Tjacobi
integer :: nib,nie,njb,nje,nkb,nke
integer :: ii,ij,ik
integer :: iindex
!
! init size , param
......@@ -1354,7 +1449,14 @@ end subroutine construct_vertical_coeff
nz = u%grid_param%nz
nlocal = u%ix_max -u%ix_min + 1
halo_size = u%halo_size
nib = 1-halo_size
nie = nlocal+halo_size
njb = 1-halo_size
nje = nlocal+halo_size
nkb = 0
nke = nz + 1
! Set optimal smoothing parameter on each level
!rho = 2.0_rl/(2.0_rl+4.0_rl*model_param%omega2*u%grid_param%n**2/(1.0_rl+4.0_rl*model_param%omega2*u%grid_param%n**2))
rho = smoother_param%rho
......@@ -1369,42 +1471,62 @@ end subroutine construct_vertical_coeff
!$acc enter data create(Tjacobi(level)%u0)
end if
if (LUseT) then
allocate(Tjacobi(level)%ut0(1-halo_size:nlocal+halo_size, &
1-halo_size:nlocal+halo_size, &
0:u%grid_param%nz+1) )
!$acc enter data create(Tjacobi(level)%ut0)
print*,"allocate(ut0),nib,nie,njb,nje,nkb,nke,level",nib,nie,njb,nje,nkb,nke, level
!!$ allocate(Tjacobi(level)%ut0(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
!!$ ut0 => Tjacobi(level)%ut0
!!$ allocate(ut0(nib:nie,njb:nje,nkb:nke) )
!!$ !$acc enter data create(ut0)
iindex = zt1d_discretisation_allocate3d(ut0,nib,nie,njb,nje,nkb,nke)
Tjacobi(level)%ut0 => ut0
end if
! Create residual vector
allocate(Tjacobi(level)%r(nz))
!$acc enter data create(Tjacobi(level)%r)
r => Tjacobi(level)%r
!$acc enter data create(r)
! Allocate memory for auxiliary vectors for Thomas algorithm
allocate(Tjacobi(level)%c(nz))
!$acc enter data create(Tjacobi(level)%c)
c => Tjacobi(level)%c
!$acc enter data create(c)
allocate(Tjacobi(level)%utmp(nz))
!$acc enter data create(Tjacobi(level)%utmp)
utmp => Tjacobi(level)%utmp
!$acc enter data create(utmp)
if (LUseT) then
allocate(Tjacobi(level)%Sr)
allocate(Tjacobi(level)%Sr%st(1-halo_size:nlocal+halo_size, &
1-halo_size:nlocal+halo_size, &
0:u%grid_param%nz+1) )
!$acc enter data create(Tjacobi(level)%Sr%st)
!!$ allocate(Tjacobi(level)%Sr%st(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
!!$ zSr_st => Tjacobi(level)%Sr%st
!!$ !$acc enter data create(zSr_st)
iindex = zt1d_discretisation_allocate3d(zSr_st,nib,nie,njb,nje,nkb,nke)
Tjacobi(level)%Sr%st => zSr_st
allocate(Tjacobi(level)%Sut0)
allocate(Tjacobi(level)%Sut0%st(1-halo_size:nlocal+halo_size, &
1-halo_size:nlocal+halo_size, &
0:u%grid_param%nz+1) )
!$acc enter data create(Tjacobi(level)%Sut0%st)
!!$ allocate(Tjacobi(level)%Sut0%st(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
!!$ zSut0_st => Tjacobi(level)%Sut0%st
!!$ !$acc enter data create(zSut0_st)
iindex = zt1d_discretisation_allocate3d(zSut0_st,nib,nie,njb,nje,nkb,nke)
Tjacobi(level)%Sut0%st => zSut0_st
allocate(Tjacobi(level)%Sutmp)
allocate(Tjacobi(level)%Sutmp%st(1-halo_size:nlocal+halo_size, &
1-halo_size:nlocal+halo_size, &
0:u%grid_param%nz+1) )
!$acc enter data create(Tjacobi(level)%Sutmp%st)
!!$ allocate(Tjacobi(level)%Sutmp%st(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
!!$ zSutmp_st => Tjacobi(level)%Sutmp%st
!!$ !$acc enter data create(zSutmp_st)
iindex = zt1d_discretisation_allocate3d(zSutmp_st,nib,nie,njb,nje,nkb,nke)
Tjacobi(level)%Sutmp%st => zSutmp_st
end if
end if
call boundary_mnh(u)
#ifdef OVERLAPCOMMS
overlap_comms = (nlocal > 2)
#else
......@@ -1448,40 +1570,27 @@ end subroutine construct_vertical_coeff
! Temporary vector
if (LUseO) then
!!$ allocate(u0(0:u%grid_param%nz+1, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size) )
u0 => Tjacobi(level)%u0
end if
if (LUseT) then
!!$ allocate(ut0(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
ut0 => Tjacobi(level)%ut0
end if
if (LUseO) u0(:,:,:) = u%s(:,:,:)
if (LUseT) then
zu_st => u%st
!$acc kernels
ut0(:,:,:) = zu_st(:,:,:)
!$acc kernels loop independent collapse(3)
do concurrent (ii=nib:nie,ij=njb:nje,ik=nkb:nke)
ut0(ii,ij,ik) = zu_st(ii,ij,ik)
end do
!$acc end kernels
end if
! Create residual vector
!!$ allocate(r(nz))
r => Tjacobi(level)%r
! Allocate memory for auxiliary vectors for Thomas algorithm
!!$ allocate(c(nz))
c => Tjacobi(level)%c
!!$ allocate(utmp(nz))
utmp => Tjacobi(level)%utmp
if (LUseT) then
!!$ allocate(Sr%st(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
Sr => Tjacobi(level)%Sr
!!$ allocate(Sut0%st(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
Sut0 => Tjacobi(level)%Sut0
zSut0_st => Sut0%st
......@@ -1491,9 +1600,6 @@ end subroutine construct_vertical_coeff
zSut0_st(:,:,:) = zu_st(:,:,:)
!$acc end kernels
!!$ allocate(Sutmp%st(1-halo_size:nlocal+halo_size, &
!!$ 1-halo_size:nlocal+halo_size, &
!!$ 0:u%grid_param%nz+1) )
Sutmp => Tjacobi(level)%Sutmp
endif
......
......@@ -73,6 +73,7 @@ contains
real(kind=rl) , optional , intent (in) :: PL,PH
real(kind=rl) , optional , intent (in) :: PA_K(:),PB_K(:),PC_K(:),PD_K(:)
call mg_init_mnh(KN,KNZ,PL,PH,PA_K,PB_K,PC_K,PD_K)
! Initialise ghosts in initial solution, as mg_solve assumes that they
......@@ -113,6 +114,7 @@ implicit none
integer , optional , intent(in) :: KIB,KIE,KIU,KJB,KJE,KJU,KKU
real(kind=rl) , optional , intent(in) :: PY(:,:,:)
call zt1d_discretisation_init(KIU,KJU,KKU)
call initialise_rhs_mnh(grid_param,model_param,xb_fine,KIB,KIE,KIU,KJB,KJE,KJU,KKU,PY)
call haloswap_mnh(mg_param%n_lev,pproc,xb_fine)
......@@ -179,7 +181,7 @@ end subroutine mg_main_get_u_mnh
call mg_main_initialise_u_mnh()
if (i_am_master_mpi) then
write(STDOUT,*),'************************ IT=',it,' ***************************'
write(STDOUT,*) '************************ IT=',it,' ***************************'
call flush(STDOUT)
end if
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment