File: C:\NOAA\NEMS_11731\src\chem\gocart\src\GMAO_Shared\MAPL_Base\arrayscatter.H

1     ! $Id: arrayscatter.H,v 1.8 2009/01/26 13:23:24 trayanov Exp $
2     
3     #ifdef NAME_
4     #undef NAME_
5     #endif
6     
7     #define NAME_ ArrayScatter_
8     
9     #include "overload.macro"
10     
11       subroutine SUB_(local_array, global_array, grid, mask, depe, hw, rc)
12     
13     ! Mask is really a permutation on the first dimension
14     
15         TYPE_(kind=EKIND_),         intent(  OUT) :: local_array  DIMENSIONS_
16         TYPE_(kind=EKIND_), target, intent(IN   ) :: global_array DIMENSIONS_
17         type (ESMF_Grid)                          :: grid
18         integer, optional,          intent(IN   ) :: mask(:)
19         integer, optional,          intent(IN   ) :: depe
20         integer, optional,          intent(IN   ) :: hw
21         integer, optional,          intent(  OUT) :: rc
22         
23     ! Local variables
24     
25         integer                               :: status
26         character(len=ESMF_MAXSTR)            :: IAm='ArrayScatter'
27     
28         TYPE_(kind=EKIND_),    pointer        :: myglob DIMENSIONS_ => null()
29         TYPE_(kind=EKIND_),    pointer        :: VAR(:)
30         type (ESMF_DistGrid)                  :: distGrid
31         type(ESMF_DELayout)                   :: LAYOUT
32         integer,               allocatable    :: AL(:,:)
33         integer,               allocatable    :: AU(:,:)
34         integer, dimension(:), allocatable    :: SENDCOUNTS, DISPLS, KK
35         integer                               :: nDEs
36         integer                               :: recvcount
37         integer                               :: I, J, K, II,  JJ, de, deId
38         integer                               :: I1, IN, J1, JN
39         integer                               :: gridRank
40         integer                               :: LX, LY
41         integer                               :: srcPE
42         integer                               :: MYHW, ISZ, JSZ, ISZL
43         integer                               :: deList(1)
44         logical                               :: alloc_var
45     
46     ! Works only on 1D and 2D arrays
47     ! Note: for tile variables the gridRank is 1 
48     ! and the case RANK_=2 needs additional attention 
49     
50         ASSERT_(RANK_ <= 2)
51     
52     ! Optional change of source PE. Default=MAPL_Root
53     
54         if(present(depe)) then
55            srcPE = depe 
56         else
57            srcPE = MAPL_Root
58         end if
59     
60     ! Optional single halo width
61     
62         if(present(hw)) then
63            myhw = hw
64         else
65            myhw = 0
66         end if
67     
68     ! Some halo limitations
69     
70         if(myhw > 0) then
71            ASSERT_(RANK_ == 2        ) ! No halo allowed on 1D
72            ASSERT_(.not.present(MASK)) ! No halo allowed if 1st dim is permutted
73         end if
74      
75     ! Initialize
76         alloc_var=.false.
77     
78     ! Get grid and layout information
79     
80         call ESMF_GridGet    (GRID,   dimCount=gridRank, rc=STATUS);VERIFY_(STATUS)
81         call ESMF_GridGet    (GRID,   distGrid=distGrid, rc=STATUS);VERIFY_(STATUS)
82         call ESMF_DistGridGet(distGRID, delayout=layout, rc=STATUS);VERIFY_(STATUS)
83         call ESMF_DELayoutGet(layout, deCount =nDEs, localDeList=deList, rc=status)
84         VERIFY_(STATUS)
85         deId = deList(1)
86     
87         allocate (AL(gridRank,0:nDEs-1),  stat=status)
88         VERIFY_(STATUS)
89         allocate (AU(gridRank,0:nDEs-1),  stat=status)
90         VERIFY_(STATUS)
91         allocate (sendcounts(0:nDEs-1), stat=status)
92         VERIFY_(STATUS)
93         call ESMF_DistGridGet(distgrid, &
94              minIndexPDimPDe=AL, maxIndexPDimPDe=AU, rc=status)
95         VERIFY_(STATUS)
96     
97         ISZ = size(GLOBAL_ARRAY,1)
98     
99     #if (RANK_ == 2)
100            JSZ = size(GLOBAL_ARRAY,2)
101     #else
102            JSZ = 1
103     #endif
104     
105     ! Compute count to be sent to each PE
106     
107         if(present(mask)) then
108            sendcounts = 0
109            do II = 1,ISZ
110               sendcounts(mask(ii)) = sendcounts(mask(ii)) + 1
111            enddo
112            sendcounts = sendcounts*JSZ
113     
114         else
115            do I = 0,nDEs-1
116               LX = AU(1,I) - AL(1,I) + 1 + 2*MYHW
117     #if (RANK_ == 1)
118               sendcounts(I) = LX
119     #else
120               LY = AU(2,I) - AL(2,I) + 1 + 2*MYHW
121               sendcounts(I) = LX*LY
122     #endif
123            end do
124         end if
125     
126     ! Count I will recieve
127     
128         recvcount = sendcounts(deId)
129     
130     ! Put VAR together at the srcPE 
131     
132         if (deId == srcPE) then
133     
134            allocate(DISPLS(0:nDEs          ), stat=status)
135            VERIFY_(STATUS)
136     
137     ! Compute displacements into the VAR vector
138     
139            displs(0) = 0
140            do I = 1,nDEs
141               displs(I) = displs(I-1) + sendcounts(I-1)
142            end do
143     !ALT       ASSERT_(displs(nDEs) == (ISZ+2*myhw)*(JSZ+2*myhw))
144     
145     ! If there is a halo, make a haloed copy of the global array.
146     !   otherwise just copy the pointer.
147     
148            myglob => global_array
149            
150     #if (RANK_ == 2)
151            if (myhw > 0) then
152               allocate(myglob(1-myhw:isz+myhw,1-myhw:jsz+myhw), stat=status)
153               VERIFY_(STATUS)
154               myglob(1:isz,1:jsz) = GLOBAL_ARRAY
155     
156     ! Fill the halo (I is cyclic)
157            
158               do j=1,myhw
159                  myglob(1  -j,:) = myglob(isz-j+1,:)
160                  myglob(isz+j,:) = myglob(    j  ,:)
161                  myglob(:,1  -j) = MAPL_Undef
162                  myglob(:,jsz+j) = MAPL_Undef
163               enddo
164            endif
165     #endif
166     
167     ! Fill the VAR vector
168            
169            alloc_var = .true.
170            if (present(mask)) then
171               allocate(VAR(0:displs(nDEs)-1), stat=status)
172               VERIFY_(STATUS)
173               allocate(KK (0:nDEs-1        ), stat=status)
174               VERIFY_(STATUS)
175               KK = DISPLS(0:nDEs-1)
176     
177               do I=1,ISZ
178                  K = MASK(I)
179                  II = KK(K)
180     #if (RANK_ == 1)
181                  VAR(II) = MYGLOB(I)
182     #else
183     	    LX = AU(1,K) - AL(1,K) + 1 
184                  do J=1,JSZ
185                     VAR(II+LX*(J-1)) = MYGLOB(I,J)
186                  end do
187     #endif
188                  KK(MASK(I)) = KK(MASK(I)) + 1 
189               end do
190     
191               deallocate(KK, stat=status)
192               VERIFY_(STATUS)
193     
194            else
195     
196     #if (RANK_ == 1)
197               var => myglob
198               alloc_var = .false.
199     #else
200               allocate(VAR(0:displs(nDEs)-1), stat=status)
201               VERIFY_(STATUS)
202     
203               if (gridRank == 1) then
204                  J1 = lbound(local_array,2)
205                  JN = ubound(local_array,2)
206               endif
207               do I = 0,nDEs-1
208                  I1 = AL(1,I) - myhw
209                  IN = AU(1,I) + myhw
210                  if (gridRank > 1) then
211                     J1 = AL(2,I) - myhw
212                     JN = AU(2,I) + myhw
213                  end if
214     
215                  K = displs(I)
216                  do JJ=J1,JN
217                     do II=I1,IN
218                        var(K) = MYglob(II,JJ)
219                        K = K+1
220                     end do
221                  end do
222               end do
223     #endif
224     
225            endif !  present(mask)
226     
227            if (myhw > 0) then
228               deallocate(myglob, stat=status)
229               VERIFY_(STATUS)
230            end if
231         
232          else
233             allocate(var(0:1), stat=status)
234             VERIFY_(STATUS)
235             allocate(DISPLS(0:nDEs), stat=status)
236             VERIFY_(STATUS)
237          end if !  I am srcPEa
238     
239     
240     ! Do the communications
241         
242         call MAPL_CommsScatterV(layout, var, sendcounts, displs, &
243                                 local_array, recvcount, srcPE, status)
244         VERIFY_(STATUS)
245     
246     ! Clean-up
247     
248         deallocate(displs, stat=status)
249         VERIFY_(STATUS)
250         if(alloc_var) then
251            deallocate(VAR, stat=status)
252            VERIFY_(STATUS)
253         end if
254     
255         deallocate(sendcounts, stat=status)
256         VERIFY_(STATUS)
257         deallocate(AU,         stat=status)
258         VERIFY_(STATUS)
259         deallocate(AL,         stat=status)
260         VERIFY_(STATUS)
261     
262     ! All done
263     
264         RETURN_(ESMF_SUCCESS)
265       end subroutine SUB_
266     
267     #undef NAME_
268     #undef DIMENSIONS_
269     #undef RANK_
270     #undef VARTYPE_
271