@@ -45,7 +45,7 @@ subroutine MPP_GATHER_1D_(sbuf, rbuf, pelist)
4545 op_root = pelist2(1)
4646
4747 cnt = size(sbuf(:))
48- if(size(rbuf(:)) < cnt*nproc) call mpp_error(FATAL, &
48+ if((mpp_pe().eq.op_root).AND.( size(rbuf(:)) < cnt*nproc) ) call mpp_error(FATAL, &
4949 "MPP_GATHER_1D_: size(rbuf) must be at least npes*size(sbuf) ")
5050
5151 call mpp_gather( sbuf, rbuf, size(sbuf), op_root, pelist2, ierr )
@@ -132,13 +132,13 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d
132132 logical, intent(in) :: is_root_pe
133133 integer, optional, intent(in) :: ishift, jshift
134134
135- integer :: i, msgsize, root_pe, root_pe_test
135+ integer :: i, j, k
136+ integer :: root_pe, root_pe_test
136137 integer :: i1, i2, j1, j2, ioff, joff
137- integer :: my_ind(4), gind(4,size(pelist))
138- type array3D
139- MPP_TYPE_, dimension(:,:,:), allocatable :: data3D_type
140- endtype array3D
141- type(array3d), dimension(:), allocatable :: temp
138+ integer :: base_idx, send_count, msg_start
139+ integer :: blocksize_i, blocksize_j, blocksize
140+ integer, dimension(:), allocatable :: gind, counts
141+ MPP_TYPE_, dimension(:), allocatable :: rbuf
142142
143143 if (.not.ANY(mpp_pe().eq.pelist(:))) return
144144
@@ -166,76 +166,51 @@ subroutine MPP_GATHER_PELIST_3D_(is, ie, js, je, nk, pelist, array_seg, gather_d
166166 if (present(ishift)) ioff=ishift
167167 if (present(jshift)) joff=jshift
168168
169- my_ind(1) = is
170- my_ind(2) = ie
171- my_ind(3) = js
172- my_ind(4) = je
169+ ! gather indices into global index on root_pe
170+ if (is_root_pe) allocate(gind(4*size(pelist)))
171+ call mpp_gather((/is, ie, js, je/), gind, pelist)
173172
174- ! gather indices into global index on root_pe
173+ ! Compute and allocate counts and 1d recv buffer (rbuf)
175174 if (is_root_pe) then
176- allocate(temp(1:size(pelist)))
177- do i = 1, size(pelist)
178- ! root_pe data copy - no send to self
179- if (pelist(i).eq.root_pe) then
180- gind(:,i) = my_ind(:)
181- else
182- call mpp_recv(gind(:,i:i), 4, pelist(i), .FALSE., COMM_TAG_1)
183- endif
184- enddo
185- call mpp_sync_self(check=EVENT_RECV)
186- gind(1,:)=gind(1,:)+ioff
187- gind(2,:)=gind(2,:)+ioff
188- gind(3,:)=gind(3,:)+joff
189- gind(4,:)=gind(4,:)+joff
190- ! check indices to make sure they are within the range of "data"
191- if ((minval(gind).lt.1) .OR. (maxval(gind(1:2,:)).gt.size(gather_data,1)) .OR. &
192- (maxval(gind(3:4,:)).gt.size(gather_data,2))) &
193- call mpp_error(FATAL,"fms_io(mpp_gather_pelist): specified indices (with shift) are outside of the &
194- &range of the receiving array")
195- else
196- ! non root_pe's send indices to root_pe
197- call mpp_send(my_ind(:), 4, root_pe, COMM_TAG_1)
198- call mpp_sync_self(check=EVENT_SEND)
175+ allocate(counts(size(pelist)))
176+
177+ do k = 1, size(pelist)
178+ base_idx = 4*(k-1)
179+ i1 = gind( base_idx + 1 ) ;; i2 = gind( base_idx + 2 )
180+ j1 = gind( base_idx + 3 ) ;; j2 = gind( base_idx + 4 )
181+ counts(k) = (i2 - i1 + 1) * (j2 - j1 + 1) * nk
182+ enddo
183+
184+ allocate(rbuf(sum(counts)))
199185 endif
200186
201- ! gather segments into data based on indices
187+ send_count = (ie-is+1)*(je-js+1)*nk
188+
189+ ! gather data into 1d recv buffer
190+ call mpp_gather(reshape(array_seg(is:ie,js:je,1:nk),[send_count]), send_count, rbuf, counts, pelist)
191+
192+ ! Unpack recv buffer into return array (gather_data)
202193 if (is_root_pe) then
203- do i = 1, size(pelist)
204- if (pelist(i).ne.root_pe) then ! no send to self
205- i1 = gind(1,i)
206- i2 = gind(2,i)
207- j1 = gind(3,i)
208- j2 = gind(4,i)
209- msgsize = (i2-i1+1)*(j2-j1+1)*nk
210- allocate(temp(i)%data3D_type(i1:i2,j1:j2,1:nk))
211- call mpp_recv(temp(i)%data3D_type(i1:i2,j1:j2,1:nk), msgsize, pelist(i), .FALSE., COMM_TAG_2)
212- endif
213- enddo
214- call mpp_sync_self(check=EVENT_RECV)
215- ! unbuffer/copy the data into the return array
216- do i = 1, size(pelist)
217- if (pelist(i).eq.root_pe) then
218- ! data copy - no send to self
219- gather_data(is+ioff:ie+ioff,js+joff:je+joff,1:nk) = array_seg(is:ie,js:je,1:nk)
220- else
221- i1 = gind(1,i)
222- i2 = gind(2,i)
223- j1 = gind(3,i)
224- j2 = gind(4,i)
225- gather_data(i1:i2,j1:j2,1:nk)=temp(i)%data3D_type(i1:i2,j1:j2,1:nk)
226- deallocate(temp(i)%data3D_type)
227- endif
228- enddo
229- deallocate(temp)
230- else
231- ! non root_pe's send data to root_pe
232- msgsize = (my_ind(2)-my_ind(1)+1) * (my_ind(4)-my_ind(3)+1) * nk
233- call mpp_send(array_seg, msgsize, root_pe, COMM_TAG_2)
234- call mpp_sync_self(check=EVENT_SEND)
194+ msg_start = 1
195+ do k = 1, size(pelist)
196+ base_idx = 4*(k-1)
197+ i1 = gind( base_idx + 1 ) + ioff ;; i2 = gind( base_idx + 2 ) + ioff
198+ j1 = gind( base_idx + 3 ) + joff ;; j2 = gind( base_idx + 4 ) + joff
199+
200+ blocksize_i = i2 - i1 + 1
201+ blocksize_j = j2 - j1 + 1
202+ blocksize = blocksize_i * blocksize_j * nk
203+
204+ gather_data(i1:i2, j1:j2, 1:nk) = reshape(rbuf(msg_start:msg_start+blocksize-1), &
205+ [blocksize_i, blocksize_j, nk])
206+
207+ msg_start = msg_start + blocksize
208+ enddo
209+
210+ deallocate(gind, rbuf, counts)
235211 endif
236212
237213 call mpp_sync_self()
238- return
239214
240215end subroutine MPP_GATHER_PELIST_3D_
241216!> @}
0 commit comments