phys_grid.f90

CCSM Research Tools: Community Atmosphere Model (CAM)

#include <misc.h>
module phys_grid
!----------------------------------------------------------------------- 
! 
! Purpose: Definition of physics computational horizontal grid.
!
! Method: Variables are private; interface routines used to extract
!         information for use in user code.
! 
! Entry points:
!      phy_grid_init       initialize chunk'ed data structure
!
!      get_chunk_indices_p get local chunk index range
!      get_ncols_p         get number of columns for a given chunk
!      get_xxx_all_p       get global indices or coordinates for a given
!                          chunk
!      get_xxx_vec_p       get global indices or coordinates for a subset
!                          of the columns in a chunk
!      get_xxx_p           get global indices or coordinates for a single
!                          column
!      where xxx is
!       lat                for global latitude index
!       lon                for global longitude index
!       rlat               for latitude coordinate (in radians)
!       rlon               for longitude coordinate (in radians)
!
!      get_chunk_coord_p   get local chunk and column indices
!                          for given (lon,lat) coordinates
!
!      scatter_field_to_chunk
!                          distribute longitude/latitude field
!                          to decomposed chunk data structure
!      gather_chunk_to_field
!                          reconstruct longitude/latitude field
!                          from decomposed chunk data structure
!
!      read_chunk_from_field
!                          read and distribute longitude/latitude field
!                          to decomposed chunk data structure
!      write_field_from_chunk
!                          write longitude/latitude field
!                          from decomposed chunk data structure
!
!      block_to_chunk_send_pters
!                          return pointers into send buffer where data
!                          from decomposed longitude/latitude fields should
!                          be copied to
!      block_to_chunk_recv_pters
!                          return pointers into receive buffer where data
!                          for decomposed chunk data structures should
!                          be copied from
!      transpose_block_to_chunk
!                          transpose buffer containing decomposed 
!                          longitude/latitude fields to buffer
!                          containing decomposed chunk data structures
!
!      chunk_to_block_send_pters
!                          return pointers into send buffer where data
!                          from decomposed chunk data structures should
!                          be copied to
!      chunk_to_block_recv_pters
!                          return pointers into receive buffer where data
!                          for decomposed longitude/latitude fields should
!                          be copied from
!      transpose_chunk_to_block
!                          transpose buffer containing decomposed
!                          chunk data structures to buffer
!                          containing decomposed longitude/latitude fields
!
!      chunk_index         identify whether index is for a latitude or
!                          a chunk
!
! Author: John Drake and Patrick Worley
! 
!-----------------------------------------------------------------------
   use precision
   use ppgrid, only: pcols, pver, begchunk, endchunk
   use pmgrid, only: plon, plat, beglat, endlat
#if ( defined SPMD )
   use spmd_dyn, only: proc, npes
   use mpishorthand
#endif

   save

#if ( ! defined SPMD )
   integer :: npes = 1
#endif
   integer :: nlthreads                ! number of local OpenMP threads
   integer, dimension(:), allocatable, private :: npthreads
                                       ! number of OpenMP threads per process
   integer :: ngthreads                ! total number of threads

! chunk data structures
   type chunk
     integer  :: ncols                 ! number of vertical columns
     integer  :: lon(pcols)            ! global longitude indices
     integer  :: lat(pcols)            ! global latitude indices
     integer  :: owner                 ! id of process where chunk assigned
     integer  :: lchunk                ! local chunk index
   end type chunk

   integer :: nchunks                  ! global chunk count
   type (chunk), dimension(:), allocatable, private :: chunks  
                                       ! global computational grid

   integer, private :: nlchunks        ! local chunk count
   integer, dimension(:), allocatable, private :: lchunks 
                                       ! local chunks

   type knuhc
     integer  :: chunkid               ! chunk id
     integer  :: col                   ! column index in chunk
   end type knuhc

   type (knuhc), dimension(:,:), allocatable, private :: knuhcs
                                       ! map from global (lon,lat) coordinates
                                       ! to chunk'ed grid

! column mapping data structures
   type column_map
     integer  :: chunk                 ! global chunk index
     integer  :: ccol                  ! column ordering in chunk
   end type column_map

   integer :: ngcols                   ! global column count
   integer :: nlcols                   ! local column count
   type (column_map), dimension(:), allocatable, private :: pgcols
                                       ! ordered list of columns (for use in gather/scatter)
                                       ! NOTE: consistent with local ordering

! column remap data structures
   integer, dimension(:), allocatable, private :: gs_col_num
                                       ! number of columns scattered to each process in
                                       ! field_to_chunk scatter
   integer, dimension(:), allocatable, private :: gs_col_offset
                                       ! offset of columns (-1) in pgcols scattered to
                                       ! each process in field_to_chunk scatter

   integer, dimension(:), allocatable, private :: btofc_blk_num
                                       ! number of grid points scattered to each process in
                                       ! block_to_chunk alltoallv, and gathered from each
                                       ! process in chunk_to_block alltoallv

   integer, dimension(:), allocatable, private :: btofc_chk_num
                                       ! number of grid points gathered from each process in
                                       ! block_to_chunk alltoallv, and scattered to each
                                       ! process in chunk_to_block alltoallv

   type btofc_pters
     integer :: ncols                  ! number of columns in block
     integer :: nlvls                  ! number of levels in columns
     integer, dimension(:,:), pointer :: pter 
   end type btofc_pters
   type (btofc_pters), dimension(:), allocatable, private :: btofc_blk_offset
                                       ! offset in btoc send array (-1) where 
                                       ! (blockid, bcid, k) column should be packed in
                                       ! block_to_chunk alltoallv, AND
                                       ! offset in ctob receive array (-1) from which
                                       ! (blockid, bcid, k) column should be unpacked in
                                       ! chunk_to_block alltoallv

   type (btofc_pters), dimension(:), allocatable, private :: btofc_chk_offset
                                       ! offset in btoc receive array (-1) from which
                                       ! (lchnk, i, k) data should be unpacked in
                                       ! block_to_chunk alltoallv, AND
                                       ! offset in ctob send array (-1) where
                                       ! (lchnk, i, k) data should be packed in
                                       ! chunk_to_block alltoallv

   integer :: block_buf_nrecs          ! number of local grid points (lon,lat,lev)
                                       ! in dynamics decomposition (including level 0)
   integer :: chunk_buf_nrecs          ! number of local grid points (lon,lat,lev)
                                       ! in physics decomposition (including level 0)

! miscellaneous phys_grid data
   real(r8) :: clat_p(plat)            ! physics grid latitudes (radians)
   integer  :: nlon_p(plat)            ! num longitudes per latitude
   real(r8) :: clon_p(plon,plat)       ! physics grid longitudes (radians)
   logical :: physgrid_set = .false.   ! flag indicates physics grid has been set
   logical :: local_dp_map = .false.   ! flag indicates that mapping between dynamics 
                                       ! and physics decompositions does not require 
                                       ! interprocessor communication

contains
!========================================================================

   subroutine phys_grid_init(opt, chunks_per_thread)
!----------------------------------------------------------------------- 
! 
! Purpose: Physics mapping initialization routine:  
! 
! Method: 
! 
! Author: John Drake and Patrick Worley
! 
!-----------------------------------------------------------------------
   use precision
   use pmgrid, only: iam, plev, plond, platd
   use pspect, only: pmmax, pnmax
   use rgrid, only: nlon
   use commap, only: clat, clon
   use dyn_grid, only: get_block_coord_cnt_d, get_block_coord_d, &
                       get_block_col_cnt_d, get_block_lvl_cnt_d, &
                       get_lon_d, get_lat_d, get_block_bounds_d, &
                       get_block_owner_d, get_block_levels_d

   implicit none
!
!------------------------------Arguments--------------------------------
!
   integer, intent(in)  :: chunks_per_thread ! target number of chunks
                                         !  per thread
   integer, intent(in)  :: opt           ! grid optimization option
                                         ! -1: each chunk is a latitude line
                                         !  0: chunks do not cross latitude boundaries
!
!---------------------------Local workspace-----------------------------
!
   integer :: i, j, jb, k, lchnk, p      ! loop indices
   integer :: tchunks                    ! target number of chunks per thread
   integer :: cbeg                       ! beginning longitude index for 
                                         !  current chunk
   integer :: cid                        ! chunk id
   integer :: pchunkid                   ! chunk global ordering
   integer :: begpchunk, endpchunk       ! segment of chunk global ordering on 
                                         !  a given process
   integer :: plchunks                   ! number of chunks for a given process
   integer :: curgcol                    ! current global column index
   integer :: firstblock, lastblock      ! global block indices
   integer :: blksiz                     ! current block size
   integer :: glbcnt, curcnt             ! running grid point counts
   integer :: curp                       ! current process id
   integer :: block_cnt                  ! number of blocks containing data
                                         ! for a given vertical column
   integer :: numlvl                     ! number of vertical levels in block 
                                         ! column
   integer :: levels(plev+1)             ! vertical level indices
   integer :: owner_d                    ! processor owning given block column
   integer :: owner_p                    ! processor owning given chunk column
   integer :: ncol                       ! number of columns in current chunk
   integer :: blockids(plev+1)           ! block indices
   integer :: bcids(plev+1)              ! block column indices
   integer :: glon, glat                 ! global (lon,lat) indices
   integer :: ntmp1, ntmp2               ! work variables
!-----------------------------------------------------------------------
!
! Initialize physics grid, using dynamics grid
!
   do j=1,plat
      clat_p(j) = clat(j)
      nlon_p(j) = nlon(j)
      do i=1,nlon(j)
         clon_p(i,j) = clon(i,j)
      enddo
   enddo
!
! Determine total number of columns and block index bounds
!
   ngcols = 0
   do j=1,plat
      ngcols = ngcols + nlon_p(j)
   enddo
   call get_block_bounds_d(firstblock,lastblock)
!
! Option -1: each latitude line is a single chunk, same as 1D dynamics decompositions.
!            
   if (opt == -1) then
!
! Check that pcols == plon
!
      if (pcols /= plon) then
         write(6,*) "PHYS_GRID_INIT error: opt -1 specified, but PCOLS /= PLON"
         call endrun()
      endif
!
! Determine total number of chunks
!
      nchunks = plat
!
! Allocate and initialize chunks and knuhcs data structures
!
      allocate ( chunks(1:nchunks) )
      allocate ( knuhcs(1:plond, 1:platd) )

      cid = 0
      do j=1,plat
         chunks(j)%ncols = nlon_p(j)
         do i=1,chunks(j)%ncols
            chunks(j)%lon(i) = i
            chunks(j)%lat(i) = j
            knuhcs(i,j)%chunkid = j
            knuhcs(i,j)%col = i
         enddo
      enddo
!
! Determine parallel decomposition (assuming 1D latitude decomposition in dynamics)
!
      do j=1,plat
#if (defined SPMD)
         chunks(j)%owner = proc(j)
#else
         chunks(j)%owner = 0
#endif
      enddo
!
! (including allocating and initializing data structures for gather/scatter)
!  
      allocate ( pgcols(1:ngcols) )
      allocate ( gs_col_num(0:npes-1) )
      allocate ( gs_col_offset(0:npes) )

      pchunkid = 0
      endpchunk = 0
      curgcol = 0
      do p=0,npes-1
         gs_col_offset(p) = curgcol + 1
         begpchunk = endpchunk + 1
         plchunks = 0
         gs_col_num(p) = 0
         do cid=1,nchunks
            if (chunks(cid)%owner == p) then
               pchunkid = pchunkid + 1
               plchunks = plchunks + 1

               do i=1,chunks(cid)%ncols
                  curgcol = curgcol + 1
                  pgcols(curgcol)%chunk = cid
                  pgcols(curgcol)%ccol = i
                  gs_col_num(p) = gs_col_num(p) + 1
               enddo

            endif
         enddo
         endpchunk = begpchunk + plchunks - 1
      enddo
      gs_col_offset(npes) = curgcol + 1

      do j=1,plat
         chunks(j)%lchunk = j
      enddo
      nlchunks = endlat-beglat+1
      nlcols = gs_col_num(iam)
!
! Local chunk indices are identical to global latitudes {beglat,...,endlat}
!
      begchunk = beglat
      endchunk = endlat
      allocate ( lchunks(begchunk:endchunk) )
      do j=begchunk,endchunk
         lchunks(j) = j
      enddo
!
! Set flag indicating columns in physics and dynamics 
! decompositions reside on the same processors
!
      local_dp_map = .true. 
!
   else
!
! Option == 0: split local longitude/latitude blocks into chunks,
!               while attempting to create load-balanced chunks
! Option == 1: load balance chunks and assignment, attempting to
!                also minimize communication costs
! Option == 2: split local longitude/latitude blocks into chunks,
!               assigning columns using block ordering
! Option == 3: split indiviudal longitude/latitude blocks into chunks,
!               assigning columns using block ordering (default)
!
! Allocate and initialize chunks and knuhcs data structures.
!
      call create_chunks(opt, chunks_per_thread)
!
! Assign chunks to processes.
!
      call assign_chunks(opt)
!
! Determine whether dynamics and physics decompositions
! are colocated, not requiring any interprocessor communication
! in the coupling.
      local_dp_map = .true.   
      do cid=1,nchunks
         do i=1,chunks(cid)%ncols
            glon = chunks(cid)%lon(i)