Computes the redistribute mapping from the g_lo data decomposition to the lz_lo decomposition
Type | Intent | Optional | Attributes | Name | ||
---|---|---|---|---|---|---|
type(g_layout_type), | intent(in) | :: | g_lo | |||
type(lz_layout_type), | intent(in) | :: | lz_lo | |||
type(redist_type), | intent(inout) | :: | lambda_map |
subroutine setup_lambda_redistribute_local(g_lo, lz_lo, lambda_map)
use mp, only: nproc
use layouts_type, only: g_layout_type, lz_layout_type
use gs2_layouts, only: idx_local, proc_id
use gs2_layouts, only: ik_idx, it_idx, ie_idx, is_idx, il_idx, idx, ig_idx,isign_idx
use redistribute, only: index_list_type, init_redist, delete_list
use sorting, only: quicksort
implicit none
type(g_layout_type), intent(in) :: g_lo
type(lz_layout_type), intent(in) :: lz_lo
type(redist_type), intent(in out) :: lambda_map
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(2) :: to_high
integer :: to_low
integer :: ig, isign, iglo, il, ilz, il0
integer :: je, ilz_bak
integer :: n, ip
!Initialise data counters
nn_to = 0
nn_from = 0
!First count the data to be send | g_lo-->lz_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 lz_lo idx for ig=-ntgrid
ilz=idx(lz_lo,-g_lo%ntgrid,ik_idx(g_lo,iglo), &
it_idx(g_lo,iglo),ie_idx(g_lo,iglo),is_idx(g_lo,iglo))
!Loop over other local dimensions, noting that ig->ig+1 ==> ilz->ilz+1
!Note that ilz is independent ofi sign so we just add two pieces of data per point instead
do ig=-g_lo%ntgrid,g_lo%ntgrid
!Increment the data sent counter for the processor with ilz
nn_from(proc_id(lz_lo,ilz))=nn_from(proc_id(lz_lo,ilz))+2
!Increment ilz
ilz=ilz+1
enddo
enddo
endif
!Now count how much data to receive | lz_lo<--g_lo
!Protect against procs with no data
if(lz_lo%ulim_proc>=lz_lo%llim_proc)then
do ilz = lz_lo%llim_proc, lz_lo%ulim_alloc
do il=1,g_lo%nlambda !lz_lo%?
!Get iglo
iglo=idx(g_lo,ik_idx(lz_lo,ilz),it_idx(lz_lo,ilz),&
il,ie_idx(lz_lo,ilz),is_idx(lz_lo,ilz))
!Increment the data to receive count for this proc
!Note we increment by two due to independence of isign
nn_to(proc_id(g_lo,iglo))=nn_to(proc_id(g_lo,iglo))+2
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)))
!For sorting messsages later
allocate (sort_list(ip)%first(nn_to(ip)))
end if
end do
!Reinitialise counters
nn_to = 0
nn_from = 0
!First fill in the 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
!Convert to ilz for ig=-ntgrid
ilz=idx(lz_lo,-g_lo%ntgrid,ik_idx(g_lo,iglo), &
it_idx(g_lo,iglo),ie_idx(g_lo,iglo),is_idx(g_lo,iglo))
!Store backup of ilz value
ilz_bak=ilz
!Loop over other local dimensions
do isign=1,2
do ig=-g_lo%ntgrid,g_lo%ntgrid
!Get proc id
ip=proc_id(lz_lo,ilz)
!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
!Increment ilz
ilz=ilz+1
enddo
!Restore ilz
ilz=ilz_bak
enddo
enddo
endif
!Now fill in the receiving indices, these must match the message order (through sorting later)
!NOTE: Not all procs have data in lz_lo so protect against this
if(lz_lo%ulim_proc>=lz_lo%llim_proc)then
do ilz=lz_lo%llim_proc,lz_lo%ulim_alloc
ig=ig_idx(lz_lo,ilz)
!Whilst lz_lo is independent of sign we actually have a lambda dimension double that of other layouts, which is how the sign dependent information is stored, so loop over sign here.
do isign=1,2
!Should the upper limit actually be max(nlambda,ng2+1)?
do il0=1,g_lo%nlambda !lz_lo%?
je=jend(ig)
!Pick the correct value of il
il=il0
if (je==0) then
if (isign==2) il=2*g_lo%nlambda+1-il !lz_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 !lz_lo%?
endif
else
if(isign==2) il=2*je-il
endif
endif
!Get iglo value. Note we use il0 and not il
iglo=idx(g_lo,ik_idx(lz_lo,ilz),it_idx(lz_lo,ilz),&
il0,ie_idx(lz_lo,ilz),is_idx(lz_lo,ilz))
!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)=ilz
!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
endif
!Now sort receive indices into message order
do ip=0,nproc-1
!Only worry about the cases where we're receiving data
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)
endif
enddo
!Now setup array range values
from_low (1) = -g_lo%ntgrid
from_low (2) = 1
from_low (3) = g_lo%llim_proc
to_low = lz_lo%llim_proc
to_high(1) = max(2*nlambda, 2*ng2+1) !lz_lo%?
to_high(2) = lz_lo%ulim_alloc
from_high(1) = g_lo%ntgrid
from_high(2) = 2
from_high(3) = g_lo%ulim_alloc
!Create lambda map redistribute objects
call init_redist (lambda_map, 'c', to_low, to_high, to_list, &
from_low, from_high, from_list)
!Deallocate the list objects
call delete_list (to_list)
call delete_list (from_list)
call delete_list (sort_list)
end subroutine setup_lambda_redistribute_local