Skip to content

Feature/general indices gatscat #1832

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
abrooks1085 wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

Feature/general indices gatscat #1832

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1236dd2
Generalized pelist_gather_2D/3D for general dim order. Added wrappers…
abrooks1085 Feb 3, 2026
81758a7
Generalized pelist_gather_2D/3D to handle general dim order. Added wr…
abrooks1085 Feb 3, 2026
a1084e4
Added new pelist_gather/scatter routines to generic interfaces mpp_ga…
abrooks1085 Feb 3, 2026
8948f99
Require contiguous assumed-shape arrays for GEN_2D gather/scatter
abrooks1085 Feb 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions mpp/include/mpp_comm.inc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -380,9 +380,13 @@
#define MPP_GATHER_1D_ mpp_gather_logical_1d
#define MPP_GATHER_1DV_ mpp_gather_logical_1dv
#undef MPP_GATHER_PELIST_2D_
#undef MPP_GATHER_PELIST_GEN_2D_
#undef MPP_GATHER_PELIST_3D_
#undef MPP_GATHER_PELIST_GEN_3D_
#define MPP_GATHER_PELIST_2D_ mpp_gather_pelist_logical_2d
#define MPP_GATHER_PELIST_GEN_2D_ mpp_gather_pelist_logical_gen_2d
#define MPP_GATHER_PELIST_3D_ mpp_gather_pelist_logical_3d
#define MPP_GATHER_PELIST_GEN_3D_ mpp_gather_pelist_logical_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_LOGICAL
#include <mpp_gather.fh>
Expand All @@ -394,9 +398,13 @@
#define MPP_GATHER_1D_ mpp_gather_int4_1d
#define MPP_GATHER_1DV_ mpp_gather_int4_1dv
#undef MPP_GATHER_PELIST_2D_
#undef MPP_GATHER_PELIST_GEN_2D_
#undef MPP_GATHER_PELIST_3D_
#undef MPP_GATHER_PELIST_GEN_3D_
#define MPP_GATHER_PELIST_2D_ mpp_gather_pelist_int4_2d
#define MPP_GATHER_PELIST_GEN_2D_ mpp_gather_pelist_int4_gen_2d
#define MPP_GATHER_PELIST_3D_ mpp_gather_pelist_int4_3d
#define MPP_GATHER_PELIST_GEN_3D_ mpp_gather_pelist_int4_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_INTEGER4
#include <mpp_gather.fh>
Expand All @@ -409,9 +417,13 @@
#define MPP_GATHER_1D_ mpp_gather_int8_1d
#define MPP_GATHER_1DV_ mpp_gather_int8_1dv
#undef MPP_GATHER_PELIST_2D_
#undef MPP_GATHER_PELIST_GEN_2D_
#undef MPP_GATHER_PELIST_3D_
#undef MPP_GATHER_PELIST_GEN_3D_
#define MPP_GATHER_PELIST_2D_ mpp_gather_pelist_int8_2d
#define MPP_GATHER_PELIST_GEN_2D_ mpp_gather_pelist_int8_gen_2d
#define MPP_GATHER_PELIST_3D_ mpp_gather_pelist_int8_3d
#define MPP_GATHER_PELIST_GEN_3D_ mpp_gather_pelist_int8_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_INTEGER8
#include <mpp_gather.fh>
Expand All @@ -424,9 +436,13 @@
#define MPP_GATHER_1D_ mpp_gather_real4_1d
#define MPP_GATHER_1DV_ mpp_gather_real4_1dv
#undef MPP_GATHER_PELIST_2D_
#undef MPP_GATHER_PELIST_GEN_2D_
#undef MPP_GATHER_PELIST_3D_
#undef MPP_GATHER_PELIST_GEN_3D_
#define MPP_GATHER_PELIST_2D_ mpp_gather_pelist_real4_2d
#define MPP_GATHER_PELIST_GEN_2D_ mpp_gather_pelist_real4_gen_2d
#define MPP_GATHER_PELIST_3D_ mpp_gather_pelist_real4_3d
#define MPP_GATHER_PELIST_GEN_3D_ mpp_gather_pelist_real4_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_REAL4
#include <mpp_gather.fh>
Expand All @@ -438,50 +454,70 @@
#define MPP_GATHER_1D_ mpp_gather_real8_1d
#define MPP_GATHER_1DV_ mpp_gather_real8_1dv
#undef MPP_GATHER_PELIST_2D_
#undef MPP_GATHER_PELIST_GEN_2D_
#undef MPP_GATHER_PELIST_3D_
#undef MPP_GATHER_PELIST_GEN_3D_
#define MPP_GATHER_PELIST_2D_ mpp_gather_pelist_real8_2d
#define MPP_GATHER_PELIST_GEN_2D_ mpp_gather_pelist_real8_gen_2d
#define MPP_GATHER_PELIST_3D_ mpp_gather_pelist_real8_3d
#define MPP_GATHER_PELIST_GEN_3D_ mpp_gather_pelist_real8_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_REAL8
#include <mpp_gather.fh>

!#################################################
#undef MPP_SCATTER_PELIST_2D_
#undef MPP_SCATTER_PELIST_GEN_2D_
#undef MPP_SCATTER_PELIST_3D_
#undef MPP_SCATTER_PELIST_GEN_3D_
#undef MPP_TYPE_
#define MPP_TYPE_ integer(i4_kind)
#define MPP_SCATTER_PELIST_2D_ mpp_scatter_pelist_int4_2d
#define MPP_SCATTER_PELIST_GEN_2D_ mpp_scatter_pelist_int4_gen_2d
#define MPP_SCATTER_PELIST_3D_ mpp_scatter_pelist_int4_3d
#define MPP_SCATTER_PELIST_GEN_3D_ mpp_scatter_pelist_int4_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_INTEGER4
#include <mpp_scatter.fh>

#undef MPP_SCATTER_PELIST_2D_
#undef MPP_SCATTER_PELIST_GEN_2D_
#undef MPP_SCATTER_PELIST_3D_
#undef MPP_SCATTER_PELIST_GEN_3D_
#undef MPP_TYPE_
#define MPP_TYPE_ integer(i8_kind)
#define MPP_SCATTER_PELIST_2D_ mpp_scatter_pelist_int8_2d
#define MPP_SCATTER_PELIST_GEN_2D_ mpp_scatter_pelist_int8_gen_2d
#define MPP_SCATTER_PELIST_3D_ mpp_scatter_pelist_int8_3d
#define MPP_SCATTER_PELIST_GEN_3D_ mpp_scatter_pelist_int8_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_INTEGER8
#include <mpp_scatter.fh>

#undef MPP_SCATTER_PELIST_2D_
#undef MPP_SCATTER_PELIST_GEN_2D_
#undef MPP_SCATTER_PELIST_3D_
#undef MPP_SCATTER_PELIST_GEN_3D_
#undef MPP_TYPE_
#define MPP_TYPE_ real(r4_kind)
#define MPP_SCATTER_PELIST_2D_ mpp_scatter_pelist_real4_2d
#define MPP_SCATTER_PELIST_GEN_2D_ mpp_scatter_pelist_real4_gen_2d
#define MPP_SCATTER_PELIST_3D_ mpp_scatter_pelist_real4_3d
#define MPP_SCATTER_PELIST_GEN_3D_ mpp_scatter_pelist_real4_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_REAL4
#include <mpp_scatter.fh>

#undef MPP_SCATTER_PELIST_2D_
#undef MPP_SCATTER_PELIST_GEN_2D_
#undef MPP_SCATTER_PELIST_3D_
#undef MPP_SCATTER_PELIST_GEN_3D_
#undef MPP_TYPE_
#define MPP_TYPE_ real(r8_kind)
#define MPP_SCATTER_PELIST_2D_ mpp_scatter_pelist_real8_2d
#define MPP_SCATTER_PELIST_GEN_2D_ mpp_scatter_pelist_real8_gen_2d
#define MPP_SCATTER_PELIST_3D_ mpp_scatter_pelist_real8_3d
#define MPP_SCATTER_PELIST_GEN_3D_ mpp_scatter_pelist_real8_gen_3d
#undef MPI_TYPE_
#define MPI_TYPE_ MPI_REAL8
#include <mpp_scatter.fh>
Expand Down
89 changes: 75 additions & 14 deletions mpp/include/mpp_gather.fh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,6 @@ subroutine MPP_GATHER_1DV_(sbuf, ssize, rbuf, rsize, pelist)
deallocate(displs)
end subroutine MPP_GATHER_1DV_


subroutine MPP_GATHER_PELIST_2D_(is, ie, js, je, pelist, array_seg, gather_data, is_root_pe, &
ishift, jshift)
integer, intent(in) :: is, ie, js, je
Expand All @@ -106,6 +105,26 @@ subroutine MPP_GATHER_PELIST_2D_(is, ie, js, je, pelist, array_seg, gather_data,
logical, intent(in) :: is_root_pe
integer, optional, intent(in) :: ishift, jshift

integer, dimension(3) :: dim_order

dim_order = (/1,2,3/)

call mpp_gather(is, ie, js, je, pelist, array_seg, gather_data, dim_order, is_root_pe, &
ishift, jshift)
return

end subroutine MPP_GATHER_PELIST_2D_

subroutine MPP_GATHER_PELIST_GEN_2D_(is, ie, js, je, pelist, array_seg, gather_data, dim_order, is_root_pe, &
ishift, jshift)
integer, intent(in) :: is, ie, js, je
integer, dimension(:), intent(in) :: pelist
MPP_TYPE_, dimension(:,:), contiguous, target, intent(in) :: array_seg
MPP_TYPE_, dimension(:,:), contiguous, target, intent(inout) :: gather_data
integer, dimension(3), intent(in) :: dim_order
logical, intent(in) :: is_root_pe
integer, optional, intent(in) :: ishift, jshift

MPP_TYPE_, pointer :: arr3D(:,:,:)
MPP_TYPE_, pointer :: data3D(:,:,:)

Expand All @@ -116,11 +135,11 @@ subroutine MPP_GATHER_PELIST_2D_(is, ie, js, je, pelist, array_seg, gather_data,
data3D => null()
endif

call mpp_gather(is, ie, js, je, 1, pelist, arr3D, data3D, is_root_pe, &
call mpp_gather(is, ie, js, je, 1, pelist, arr3D, data3D, dim_order, is_root_pe, &
ishift, jshift)
return

end subroutine MPP_GATHER_PELIST_2D_
end subroutine MPP_GATHER_PELIST_GEN_2D_


subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_data, is_root_pe, &
Expand All @@ -132,16 +151,43 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d
logical, intent(in) :: is_root_pe
integer, optional, intent(in) :: ishift, jshift

integer :: i, j, k
integer, dimension(3) :: dim_order

dim_order = (/1, 2, 3/)

call mpp_gather(is, ie, js, je, nk, pelist, array_seg, gather_data, dim_order, is_root_pe, &
ishift, jshift)
return

end subroutine MPP_GATHER_PELIST_3D_

subroutine MPP_GATHER_PELIST_GEN_3D_(is, ie, js, je, nk, pelist, array_seg, gather_data, dim_order, is_root_pe, &
ishift, jshift)
integer, intent(in) :: is, ie, js, je, nk
integer, dimension(:), intent(in) :: pelist
MPP_TYPE_, dimension(:,:,:), intent(in) :: array_seg
MPP_TYPE_, dimension(:,:,:), intent(inout) :: gather_data
integer, dimension(3), intent(in) :: dim_order
logical, intent(in) :: is_root_pe
integer, optional, intent(in) :: ishift, jshift

integer :: root_pe, root_pe_test
integer :: k, us, ue, vs, ve, ws, we
integer :: i1, i2, j1, j2, ioff, joff
integer :: base_idx, send_count, msg_start
integer :: blocksize_i, blocksize_j, blocksize
integer :: blocksize_u, blocksize_v, blocksize_w, blocksize
integer, dimension(3) :: start_idx, stop_idx
integer, dimension(:), allocatable :: gind, counts
MPP_TYPE_, dimension(:), allocatable :: rbuf

if (.not.ANY(mpp_pe().eq.pelist(:))) return

! Check dim_order is a permutation of 1..3
if ( any(dim_order < 1) .or. any(dim_order > 3) ) call mpp_error(FATAL, &
"fms_io(mpp_gather_pelist): dim_order entries must be in {1,2,3}")
if ( dim_order(1) == dim_order(2) .or. dim_order(1) == dim_order(3) .or. dim_order(2) == dim_order(3) ) &
call mpp_error(FATAL, "fms_io(mpp_gather_pelist): dim_order must be a permutation of 1,2,3")

if (is_root_pe) then
root_pe = mpp_pe()
root_pe_test = 999
Expand All @@ -160,7 +206,6 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d
if ((is_root_pe) .and. (mpp_pe().ne.root_pe)) call mpp_error(FATAL, &
"fms_io(mpp_gather_pelist): too many root_pes specified")


ioff=0
joff=0
if (present(ishift)) ioff=ishift
Expand All @@ -170,7 +215,7 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d
if (is_root_pe) allocate(gind(4*size(pelist)))
call mpp_gather((/is, ie, js, je/), gind, pelist)

! Compute and allocate counts and 1d recv buffer (rbuf)
! Compute recv counts and allocate 1d recv buffer (rbuf)
if (is_root_pe) then
allocate(counts(size(pelist)))

Expand All @@ -186,8 +231,14 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d

send_count = (ie-is+1)*(je-js+1)*nk

! Get generalized stop indicies for array_seg
stop_idx = (/ie-is+1, je-js+1, nk/)
ue = stop_idx(dim_order(1))
ve = stop_idx(dim_order(2))
we = stop_idx(dim_order(3))

! gather data into 1d recv buffer
call mpp_gather(reshape(array_seg(is:ie,js:je,1:nk),[send_count]), send_count, rbuf, counts, pelist)
call mpp_gather(reshape(array_seg(1:ue,1:ve,1:we),[send_count]), send_count, rbuf, counts, pelist)

! Unpack recv buffer into return array (gather_data)
if (is_root_pe) then
Expand All @@ -197,12 +248,22 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d
i1 = gind( base_idx + 1 ) + ioff ;; i2 = gind( base_idx + 2 ) + ioff
j1 = gind( base_idx + 3 ) + joff ;; j2 = gind( base_idx + 4 ) + joff

blocksize_i = i2 - i1 + 1
blocksize_j = j2 - j1 + 1
blocksize = blocksize_i * blocksize_j * nk
! Get generalized start/stop indicies
start_idx = (/i1,j1,1/)
stop_idx = (/i2,j2,nk/)

us = start_idx(dim_order(1)) ;; ue = stop_idx(dim_order(1))
vs = start_idx(dim_order(2)) ;; ve = stop_idx(dim_order(2))
ws = start_idx(dim_order(3)) ;; we = stop_idx(dim_order(3))

gather_data(i1:i2, j1:j2, 1:nk) = reshape(rbuf(msg_start:msg_start+blocksize-1), &
[blocksize_i, blocksize_j, nk])
! Compute block sizes
blocksize_u = ue - us + 1
blocksize_v = ve - vs + 1
blocksize_w = we - ws + 1
blocksize = blocksize_u * blocksize_v * blocksize_w

gather_data(us:ue, vs:ve, ws:we) = reshape(rbuf(msg_start:msg_start+blocksize-1), &
[blocksize_u, blocksize_v, blocksize_w])

msg_start = msg_start + blocksize
enddo
Expand All @@ -212,5 +273,5 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d

call mpp_sync_self()

end subroutine MPP_GATHER_PELIST_3D_
end subroutine MPP_GATHER_PELIST_GEN_3D_
!> @}
66 changes: 53 additions & 13 deletions mpp/include/mpp_scatter.fh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,25 @@ subroutine MPP_SCATTER_PELIST_2D_(is, ie, js, je, pelist, array_seg, input_data,
MPP_TYPE_, dimension(:,:), contiguous, target, intent(in) :: input_data !< 2D array of input data
logical, intent(in) :: is_root_pe !< operational root pe

integer, dimension(3) :: dim_order

dim_order = (/1,2,3/)

call mpp_scatter(is, ie, js, je, pelist, array_seg, input_data, dim_order, is_root_pe)

return

end subroutine MPP_SCATTER_PELIST_2D_

subroutine MPP_SCATTER_PELIST_GEN_2D_(is, ie, js, je, pelist, array_seg, input_data, dim_order, is_root_pe)
integer, intent(in) :: is, ie, js, je !< indices of segment array
integer, dimension(:), intent(in) :: pelist !<PE list of target pes,
!! must be in monotonic increasing order
MPP_TYPE_, dimension(:,:), contiguous, target, intent(inout) :: array_seg !< 2D array of output data
MPP_TYPE_, dimension(:,:), contiguous, target, intent(in) :: input_data !< 2D array of input data
integer, dimension(3), intent(in) :: dim_order
logical, intent(in) :: is_root_pe !< operational root pe

MPP_TYPE_, pointer :: arr3D(:,:,:)
MPP_TYPE_, pointer :: data3D(:,:,:)

Expand All @@ -41,24 +60,44 @@ subroutine MPP_SCATTER_PELIST_2D_(is, ie, js, je, pelist, array_seg, input_data,
data3D => null()
endif

call mpp_scatter(is, ie, js, je, 1, pelist, arr3D, data3D, is_root_pe)
call mpp_scatter(is, ie, js, je, 1, pelist, arr3D, data3D, dim_order, is_root_pe)

return

end subroutine MPP_SCATTER_PELIST_2D_
end subroutine MPP_SCATTER_PELIST_GEN_2D_

subroutine MPP_SCATTER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, input_data, is_root_pe)
integer, intent(in) :: is, ie, js, je, nk
integer, dimension(:), intent(in) :: pelist
MPP_TYPE_, dimension(is:ie,js:je,1:nk), intent(inout) :: array_seg
MPP_TYPE_, dimension(:,:,:), intent(inout) :: array_seg
MPP_TYPE_, dimension(:,:,:), intent(in) :: input_data
logical, intent(in) :: is_root_pe

integer, dimension(3) :: dim_order

dim_order = (/1,2,3/)

call mpp_scatter(is, ie, js, je, nk, pelist, array_seg, input_data, dim_order, is_root_pe)

return

end subroutine MPP_SCATTER_PELIST_3D_

subroutine MPP_SCATTER_PELIST_GEN_3D_(is, ie, js, je, nk, pelist, array_seg, input_data, dim_order, is_root_pe)
integer, intent(in) :: is, ie, js, je, nk
integer, dimension(:), intent(in) :: pelist
MPP_TYPE_, dimension(:,:,:), intent(inout) :: array_seg
MPP_TYPE_, dimension(:,:,:), intent(in) :: input_data
integer, dimension(3), intent(in) :: dim_order
logical, intent(in) :: is_root_pe

integer :: i, j, k, n, m, ierr, base_idx
integer :: i1, i2, j1, j2
integer :: us, ue, vs, ve, ws, we
integer :: root_pe, root_pe_test, recv_count
integer, dimension(size(pelist)) :: counts, displs
integer, dimension(4*size(pelist)) :: gind
integer, dimension(3) :: start_idx, end_idx
MPP_TYPE_, dimension(:), allocatable :: temp

if (.not.ANY(mpp_pe().eq.pelist(:))) return
Expand Down Expand Up @@ -104,14 +143,16 @@ subroutine MPP_SCATTER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, input_d
base_idx = 4*(n-1)
i1 = gind( base_idx + 1 ) ;; i2 = gind( base_idx + 2 )
j1 = gind( base_idx + 3 ) ;; j2 = gind( base_idx + 4 )
do k = 1, nk
do j = j1, j2
do i = i1, i2
temp(m) = input_data(i,j,k)
m = m + 1
enddo
enddo
enddo

start_idx = (/i1, j1, 1/)
end_idx = (/i2, j2, nk/)

us = start_idx(dim_order(1)) ;; ue = end_idx(dim_order(1))
vs = start_idx(dim_order(2)) ;; ve = end_idx(dim_order(2))
ws = start_idx(dim_order(3)) ;; we = end_idx(dim_order(3))

temp(m:m+counts(n)-1) = reshape( input_data(us:ue, vs:ve, ws:we), [counts(n)] )
m = m + counts(n)
enddo
else
allocate(temp(1))
Expand All @@ -128,6 +169,5 @@ subroutine MPP_SCATTER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, input_d

return

end subroutine MPP_SCATTER_PELIST_3D_

end subroutine MPP_SCATTER_PELIST_GEN_3D_
!> @}
Loading
Loading