Construct a redistribute for g_lo -> le_lo
Type | Intent | Optional | Attributes | Name | ||
---|---|---|---|---|---|---|
type(g_layout_type), | intent(in) | :: | g_lo | |||
type(le_layout_type), | intent(in) | :: | le_lo | |||
type(redist_type), | intent(inout) | :: | g2le |
subroutine setup_g2le_redistribute_local(g_lo, le_lo, g2le)
use layouts_type, only: g_layout_type, le_layout_type
use mp, only: nproc
use gs2_layouts, only: idx_local, proc_id
use gs2_layouts, only: ig_idx, isign_idx
use gs2_layouts, only: ik_idx, it_idx, ie_idx, is_idx, il_idx, idx
use sorting, only: quicksort
use redistribute, only: index_list_type, init_redist, delete_list
implicit none
type(g_layout_type), intent(in) :: g_lo
type(le_layout_type), intent(in) :: le_lo
type(redist_type), intent(in out) :: g2le
type (index_list_type), dimension(0:nproc-1) :: to_list, from_list, sort_list
integer, dimension (0:nproc-1) :: nn_to, nn_from
integer, dimension (3) :: from_low, from_high
integer, dimension (3) :: to_high
integer :: to_low
integer :: ig, isign, iglo, il, ile
integer :: ie
integer :: n, ip, je
integer :: ile_bak, il0
!> The following is for debug/testing to force all grid
!> points to be communicated if we wish.
logical, parameter :: skip_forbidden_region = .true.
!Initialise the data counters
nn_to = 0
nn_from = 0
!First count the data to be sent | g_lo-->le_lo
!Protect against procs with no data
if(g_lo%ulim_proc>=g_lo%llim_proc)then
do iglo = g_lo%llim_proc,g_lo%ulim_alloc
!Get le_lo idx for ig=-ntgrid
ile=idx(le_lo,-g_lo%ntgrid,ik_idx(g_lo,iglo),&
it_idx(g_lo,iglo),is_idx(g_lo,iglo))
il = il_idx(g_lo, iglo)
!Loop over remaining dimensions, note ile is independent of isign
!so add two to count
do ig=-g_lo%ntgrid, g_lo%ntgrid
if ((.not. forbid(ig, il)) .or. .not. skip_forbidden_region) then
!Increment the data sent counter for this proc
nn_from(proc_id(le_lo,ile))=nn_from(proc_id(le_lo,ile))+2
end if
!Increment ile
ile=ile+1
enddo
enddo
endif
!Now count how much data to receive | le_lo<--g_lo
!Protect against procs with no data
if(le_lo%ulim_proc>=le_lo%llim_proc)then
do ile=le_lo%llim_proc,le_lo%ulim_alloc
ig = ig_idx(le_lo, ile)
!Loop over local dimensions, adding 2 to account for each sign
do ie=1,g_lo%negrid !le_lo%?
do il=1,g_lo%nlambda !le_lo%?
if (forbid(ig, il) .and. skip_forbidden_region) cycle
!Get index
iglo=idx(g_lo,ik_idx(le_lo,ile),it_idx(le_lo,ile),&
il,ie,is_idx(le_lo,ile))
!Increment the data to receive counter
nn_to(proc_id(g_lo,iglo))=nn_to(proc_id(g_lo,iglo))+2
enddo
enddo
enddo
endif
!Now allocate storage for index arrays
do ip = 0, nproc-1
if (nn_from(ip) > 0) then
allocate (from_list(ip)%first (nn_from(ip)))
allocate (from_list(ip)%second(nn_from(ip)))
allocate (from_list(ip)%third (nn_from(ip)))
end if
if (nn_to(ip) > 0) then
allocate (to_list(ip)%first (nn_to(ip)))
allocate (to_list(ip)%second(nn_to(ip)))
allocate (to_list(ip)%third (nn_to(ip)))
!For sorting message order later
allocate (sort_list(ip)%first(nn_to(ip)))
end if
end do
!Reinitialise counters
nn_to = 0
nn_from = 0
!First fill in sending indices, these define the message order
!Protect against procs with no data
if(g_lo%ulim_proc>=g_lo%llim_proc)then
do iglo=g_lo%llim_proc,g_lo%ulim_alloc
!Get ile for ig=-ntgrid
ile=idx(le_lo,-g_lo%ntgrid,ik_idx(g_lo,iglo),&
it_idx(g_lo,iglo),is_idx(g_lo,iglo))
ile_bak=ile
il = il_idx(g_lo, iglo)
!Loop over sign
do isign=1,2
do ig=-g_lo%ntgrid, g_lo%ntgrid
if ((.not. forbid(ig, il)) .or. .not. skip_forbidden_region ) then
!Get proc id
ip=proc_id(le_lo,ile)
!Increment procs message counter
n=nn_from(ip)+1
nn_from(ip)=n
!Store indices
from_list(ip)%first(n)=ig
from_list(ip)%second(n)=isign
from_list(ip)%third(n)=iglo
end if
!Increment ile
ile=ile+1
enddo
!Restore ile
ile=ile_bak
enddo
enddo
endif
!Now fill in the receiving indices, these must match message data order
!Protect against procs with no data
if(le_lo%ulim_proc>=le_lo%llim_proc)then
do ile=le_lo%llim_proc,le_lo%ulim_alloc
!Get ig index
ig=ig_idx(le_lo,ile)
!Loop over local dimensions,Whilst ile is independent of sign this information
!is in lambda so loop over sign included here
do ie=1,g_lo%negrid !le_lo%?
do isign=1,2
do il0=1,g_lo%nlambda !le_lo%?
if (forbid(ig, il0) .and. skip_forbidden_region) cycle
!Pick correct extended lambda value
je=jend(ig)
il=il0
if (je==0) then
if (isign==2) il=2*g_lo%nlambda+1-il !le_lo%?
else
if(il==je) then
if(isign==1) il=2*je
else if(il>je) then
if(isign==1) then
il=il+je
else
il=2*g_lo%nlambda+1-il+je !le_lo%?
endif
else
if(isign==2) il=2*je-il
endif
endif
!Get iglo value
iglo=idx(g_lo,ik_idx(le_lo,ile),it_idx(le_lo,ile),&
il0,ie,is_idx(le_lo,ile))
!Get proc_id
ip=proc_id(g_lo,iglo)
!Increment counter
n=nn_to(ip)+1
nn_to(ip)=n
!Store indices
to_list(ip)%first(n)=il
to_list(ip)%second(n)=ie
to_list(ip)%third(n)=ile
!Store sorting index
sort_list(ip)%first(n)=ig+g_lo%ntgrid-1+g_lo%ntgridtotal*(isign-1+2*(iglo-g_lo%llim_world))
enddo
enddo
enddo
enddo
endif
!Now sort receive indices into message order
do ip=0,nproc-1
if(nn_to(ip)>0) then
!Apply quicksort
call quicksort(nn_to(ip),sort_list(ip)%first,to_list(ip)%first,to_list(ip)%second,to_list(ip)%third)
endif
enddo
!Now setup array range values
from_low (1) = -g_lo%ntgrid
from_low (2) = 1
from_low (3) = g_lo%llim_proc
from_high(1) = g_lo%ntgrid
from_high(2) = 2
from_high(3) = g_lo%ulim_alloc
to_low = le_lo%llim_proc
! Note: Need to revisit next two so not dependent on module level nlambda/negrid
to_high(1) = max(2*nlambda, 2*ng2+1)
to_high(2) = negrid + 1 ! TT: just followed convention with +1.
! TT: It may be good to avoid bank conflict.
to_high(3) = le_lo%ulim_alloc
!Create g2le redist object
call init_redist (g2le, 'c', to_low, to_high, to_list, from_low, from_high, from_list)
!Deallocate lists
call delete_list (to_list)
call delete_list (from_list)
call delete_list (sort_list)
end subroutine setup_g2le_redistribute_local