Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Routines to reshape / redistribute tensors
10 : !> \author Patrick Seewald
11 : ! **************************************************************************************************
12 : MODULE dbt_reshape_ops
13 : #:include "dbt_macros.fypp"
14 : #:set maxdim = maxrank
15 : #:set ndims = range(2,maxdim+1)
16 :
17 : USE OMP_LIB, ONLY: omp_get_num_threads, omp_get_thread_num
18 : USE dbt_block, ONLY: &
19 : dbt_iterator_type, dbt_iterator_next_block, dbt_iterator_blocks_left, dbt_iterator_start, &
20 : dbt_iterator_stop, dbt_get_block_data, dbt_put_block_data, dbt_reserve_blocks
21 : USE dbt_types, ONLY: dbt_blk_sizes, &
22 : dbt_type, &
23 : dbt_get_info, &
24 : ndims_tensor, &
25 : dbt_get_stored_coordinates, &
26 : dbt_clear
27 : USE kinds, ONLY: dp
28 : USE matrix_communication, ONLY: exchange_matrix_buffers, matrix_buffer_type
29 : USE message_passing, ONLY: mp_comm_type
30 : USE util, ONLY: get_limit
31 :
32 : #include "../base/base_uses.f90"
33 :
34 : IMPLICIT NONE
35 : PRIVATE
36 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_reshape_ops'
37 :
38 : PUBLIC :: dbt_reshape, dbt_copy_indexed, dbt_reshape_release
39 :
40 : ! Non-owning references allow scalar copies and image gathers to share one transport.
41 : TYPE tensor_source_type
42 : TYPE(dbt_type), POINTER :: tensor => NULL()
43 : INTEGER :: stack_offset = 0
44 : END TYPE tensor_source_type
45 :
46 : INTERFACE dbt_copy_indexed
47 : MODULE PROCEDURE dbt_reshape, dbt_gather
48 : END INTERFACE dbt_copy_indexed
49 :
50 : ! Communication buffers persist across calls such that pages are faulted in once rather than per
51 : ! call; per block, the index holds the data offset (relative to the peer's data) and nd-index.
52 : TYPE(matrix_buffer_type), SAVE :: buffer_send, buffer_recv
53 :
54 : CONTAINS
55 :
56 : ! **************************************************************************************************
57 : !> \brief copy data (involves reshape)
58 : !> tensor_out = tensor_out + tensor_in move_data memory optimization:
59 : !> transfer data from tensor_in to tensor_out s.t. tensor_in is empty on return
60 : !> \param bounds crop tensor data while copying (as dbt_crop): start and end index per dimension
61 : !> \param tensor_in source tensor on the parent communicator
62 : !> \param tensor_out destination tensor; smaller communicators denote equal, contiguous subgroups
63 : !> \param summation add to destination instead of replacing it
64 : !> \param move_data clear the source after copying
65 : !> \param index_map optional source-to-destination block indices; zero skips a block
66 : !> \param group_mask optional subgroup selection by source block index
67 : !> \param group_dim source dimension indexing group_mask
68 : !> \author Ole Schuett
69 : ! **************************************************************************************************
70 319676 : SUBROUTINE dbt_reshape(tensor_in, tensor_out, summation, move_data, bounds, index_map, group_mask, group_dim)
71 :
72 : TYPE(dbt_type), TARGET, INTENT(INOUT) :: tensor_in
73 : TYPE(dbt_type), INTENT(INOUT) :: tensor_out
74 : LOGICAL, INTENT(IN), OPTIONAL :: summation, move_data
75 : INTEGER, INTENT(IN), OPTIONAL :: bounds(:, :), index_map(:, :), group_dim
76 : LOGICAL, INTENT(IN), OPTIONAL :: group_mask(:, :)
77 :
78 639352 : TYPE(tensor_source_type) :: sources(1)
79 :
80 319676 : sources(1)%tensor => tensor_in
81 : CALL dbt_reshape_sources(sources, tensor_out, tensor_in%pgrid%mp_comm_2d, summation, move_data, bounds, &
82 1266152 : index_map=index_map, group_mask=group_mask, group_dim=group_dim)
83 319676 : END SUBROUTINE dbt_reshape
84 :
85 : ! **************************************************************************************************
86 : !> \brief Gather whole blocks from selected tensors, optionally relabelling each block index.
87 : !> Zero images or mapped indices are skipped. Blocks retain their sizes. Each image occupies
88 : !> one source tensor extent along stack_dim, including skipped images. Retained index maps
89 : !> must be injective; the stack-dimension map stays within one source extent.
90 : !> \param tensor_in source tensors on one communicator, with identical shapes
91 : !> \param tensor_out destination tensor
92 : !> \param images source tensor indices; zero skips an image
93 : !> \param index_map source-to-destination block indices, padded to the largest source dimension
94 : !> \param stack_dim dimension along which images are stacked
95 : ! **************************************************************************************************
96 11214 : SUBROUTINE dbt_gather(tensor_in, tensor_out, images, index_map, stack_dim)
97 : TYPE(dbt_type), TARGET, INTENT(INOUT) :: tensor_in(:)
98 : TYPE(dbt_type), INTENT(INOUT) :: tensor_out
99 : INTEGER, INTENT(IN) :: images(:), index_map(:, :), stack_dim
100 :
101 22428 : INTEGER :: i, nblks(SIZE(index_map, 2))
102 178131 : TYPE(tensor_source_type) :: sources(SIZE(images))
103 :
104 11214 : CPASSERT(SIZE(tensor_in) > 0)
105 11214 : CPASSERT(SIZE(index_map, 2) == ndims_tensor(tensor_out))
106 11214 : CPASSERT(stack_dim >= 1 .AND. stack_dim <= SIZE(nblks))
107 11214 : CALL dbt_get_info(tensor_in(1), nblks_total=nblks)
108 44856 : CPASSERT(SIZE(index_map, 1) >= MAXVAL(nblks))
109 39966 : CPASSERT(ALL(index_map(:nblks(stack_dim), stack_dim) <= nblks(stack_dim)))
110 178131 : DO i = 1, SIZE(images)
111 166917 : CPASSERT(images(i) >= 0 .AND. images(i) <= SIZE(tensor_in))
112 166917 : IF (images(i) == 0) CYCLE
113 95904 : sources(i)%tensor => tensor_in(images(i))
114 178131 : sources(i)%stack_offset = (i - 1)*nblks(stack_dim)
115 : END DO
116 : CALL dbt_reshape_sources(sources, tensor_out, tensor_in(1)%pgrid%mp_comm_2d, &
117 11214 : index_map=index_map, stack_dim=stack_dim)
118 11214 : END SUBROUTINE dbt_gather
119 :
120 : ! **************************************************************************************************
121 : !> \brief Shared block transport for copies, indexed gathers and subgroup replication.
122 : !> Subgroups have identical layouts and contiguous rank ranges in comm. group_mask selects
123 : !> replicas by a source block index; with no mask, every subgroup receives the block.
124 : ! **************************************************************************************************
125 661780 : SUBROUTINE dbt_reshape_sources(sources, tensor_out, comm, summation, move_data, bounds, &
126 330890 : index_map, stack_dim, group_mask, group_dim)
127 :
128 : TYPE(tensor_source_type), INTENT(IN) :: sources(:)
129 : TYPE(dbt_type), INTENT(INOUT) :: tensor_out
130 : CLASS(mp_comm_type), INTENT(IN) :: comm
131 : LOGICAL, INTENT(IN), OPTIONAL :: summation
132 : LOGICAL, INTENT(IN), OPTIONAL :: move_data
133 : INTEGER, INTENT(IN), OPTIONAL :: bounds(:, :), index_map(:, :), stack_dim, group_dim
134 : LOGICAL, INTENT(IN), OPTIONAL :: group_mask(:, :)
135 :
136 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_reshape'
137 :
138 : INTEGER :: iproc, numnodes, nindex, &
139 : handle, handle2, iblk, offset, ndata, &
140 : isource, d, igroup, ngroups, group_size, peer, first_offset
141 : INTEGER, DIMENSION(2) :: block_range
142 : TYPE(dbt_iterator_type) :: iter
143 : CHARACTER(LEN=32) :: name_pack, name_unpack
144 330890 : INTEGER, DIMENSION(ndims_tensor(tensor_out)) :: blk_size, blk_offset, ind_nd, dest, dest_size
145 661780 : INTEGER, DIMENSION(2, ndims_tensor(tensor_out)) :: bounds_prv
146 : LOGICAL :: found, summation_prv, move_prv, do_bounds, overlap, inside
147 : TYPE(dbt_type), POINTER :: tensor_in
148 661780 : INTEGER :: source_shape(ndims_tensor(tensor_out))
149 :
150 330890 : INTEGER, ALLOCATABLE :: counts_send(:, :), counts_recv(:, :)
151 330890 : INTEGER, ALLOCATABLE :: nblks_send_mythread(:), ndata_send_mythread(:)
152 :
153 330890 : CALL timeset(routineN, handle)
154 :
155 330890 : IF (PRESENT(summation)) THEN
156 115937 : summation_prv = summation
157 : ELSE
158 : summation_prv = .FALSE.
159 : END IF
160 :
161 330890 : IF (PRESENT(move_data)) THEN
162 299890 : move_prv = move_data
163 : ELSE
164 : move_prv = .FALSE.
165 : END IF
166 :
167 330890 : do_bounds = PRESENT(bounds)
168 330890 : IF (do_bounds) THEN
169 9380 : bounds_prv(:, :) = bounds(:, :)
170 : ELSE
171 3029238 : bounds_prv(:, :) = 0
172 : END IF
173 :
174 330890 : CPASSERT(tensor_out%valid)
175 330890 : group_size = tensor_out%pgrid%mp_comm_2d%num_pe
176 330890 : CPASSERT(MOD(comm%num_pe, group_size) == 0)
177 330890 : CPASSERT(MOD(comm%mepos, group_size) == tensor_out%pgrid%mp_comm_2d%mepos)
178 330890 : ngroups = comm%num_pe/group_size
179 330890 : CPASSERT(PRESENT(group_mask) .EQV. PRESENT(group_dim))
180 330890 : IF (PRESENT(group_mask)) THEN
181 11614 : CPASSERT(SIZE(group_mask, 1) == ngroups)
182 11614 : CPASSERT(group_dim >= 1 .AND. group_dim <= SIZE(ind_nd))
183 : END IF
184 330890 : IF (PRESENT(index_map)) THEN
185 11214 : CPASSERT(SIZE(index_map, 2) == SIZE(ind_nd))
186 5736024 : CPASSERT(ALL(index_map >= 0))
187 : END IF
188 :
189 330890 : IF (.NOT. summation_prv) THEN
190 215013 : CALL timeset(routineN//"_clear", handle2)
191 215013 : CALL dbt_clear(tensor_out)
192 215013 : CALL timestop(handle2)
193 : END IF
194 :
195 330890 : numnodes = comm%num_pe
196 330890 : nindex = ndims_tensor(tensor_out) + 1
197 : ! separate timers for blocks that need a reshape (index order differs from matrix layout)
198 330890 : name_pack = routineN//"_pack"
199 817483 : DO isource = 1, SIZE(sources)
200 486593 : IF (.NOT. ASSOCIATED(sources(isource)%tensor)) CYCLE
201 415580 : CALL dbt_get_info(sources(isource)%tensor, nblks_total=source_shape)
202 415580 : IF (PRESENT(index_map)) THEN
203 383616 : CPASSERT(SIZE(index_map, 1) >= MAXVAL(source_shape))
204 : END IF
205 415580 : IF (PRESENT(group_mask)) THEN
206 11614 : CPASSERT(SIZE(group_mask, 2) == source_shape(group_dim))
207 : END IF
208 746470 : IF (permuted(sources(isource)%tensor)) name_pack = routineN//"_pack_perm"
209 : END DO
210 330890 : name_unpack = routineN//"_unpack"
211 330890 : IF (permuted(tensor_out)) name_unpack = routineN//"_unpack_perm"
212 1323560 : ALLOCATE (counts_send(2, 0:numnodes - 1), counts_recv(2, 0:numnodes - 1), source=0)
213 :
214 : !$OMP PARALLEL DEFAULT(NONE) &
215 : !$OMP SHARED(sources,tensor_out,summation,comm,numnodes,nindex,name_pack,name_unpack) &
216 : !$OMP SHARED(index_map,stack_dim,group_mask,group_dim,group_size,ngroups) &
217 : !$OMP SHARED(do_bounds,bounds_prv,buffer_send,buffer_recv) &
218 : !$OMP SHARED(counts_send,counts_recv) &
219 : !$OMP PRIVATE(nblks_send_mythread,ndata_send_mythread) &
220 : !$OMP PRIVATE(iter,ind_nd,blk_size,blk_offset,found,overlap,inside,iproc,handle2) &
221 330890 : !$OMP PRIVATE(block_range,offset,ndata,iblk,tensor_in,isource,d,dest,dest_size,igroup,peer,first_offset)
222 : CALL timeset(routineN//"_count", handle2)
223 : ALLOCATE (nblks_send_mythread(0:numnodes - 1), ndata_send_mythread(0:numnodes - 1), source=0)
224 :
225 : DO isource = 1, SIZE(sources)
226 : tensor_in => sources(isource)%tensor
227 : IF (.NOT. ASSOCIATED(tensor_in)) CYCLE
228 : CALL dbt_iterator_start(iter, tensor_in)
229 : DO WHILE (dbt_iterator_blocks_left(iter))
230 : CALL dbt_iterator_next_block(iter, ind_nd, blk_size=blk_size, blk_offset=blk_offset)
231 : IF (do_bounds) THEN
232 : CALL block_bounds(bounds_prv, blk_offset, blk_size, overlap, inside)
233 : IF (.NOT. overlap) CYCLE
234 : END IF
235 : dest = ind_nd
236 : IF (PRESENT(index_map)) THEN
237 : DO d = 1, SIZE(dest)
238 : dest(d) = index_map(ind_nd(d), d)
239 : END DO
240 : IF (ANY(dest == 0)) CYCLE
241 : END IF
242 : IF (PRESENT(stack_dim)) dest(stack_dim) = dest(stack_dim) + sources(isource)%stack_offset
243 : IF (PRESENT(index_map)) THEN
244 : CALL dbt_blk_sizes(tensor_out, dest, dest_size)
245 : CPASSERT(ALL(dest_size == blk_size))
246 : END IF
247 : CALL dbt_get_stored_coordinates(tensor_out, dest, iproc)
248 : DO igroup = 1, ngroups
249 : IF (PRESENT(group_mask)) THEN
250 : IF (.NOT. group_mask(igroup, ind_nd(group_dim))) CYCLE
251 : END IF
252 : peer = (igroup - 1)*group_size + iproc
253 : nblks_send_mythread(peer) = nblks_send_mythread(peer) + 1
254 : ndata_send_mythread(peer) = ndata_send_mythread(peer) + PRODUCT(blk_size)
255 : END DO
256 : END DO
257 : CALL dbt_iterator_stop(iter)
258 : END DO
259 : !$OMP CRITICAL(omp_dbt_reshape)
260 : counts_send(1, :) = counts_send(1, :) + ndata_send_mythread(:)
261 : counts_send(2, :) = counts_send(2, :) + nblks_send_mythread(:)
262 : ndata_send_mythread(:) = counts_send(1, :) ! current totals indicate slot for this thread
263 : nblks_send_mythread(:) = counts_send(2, :)
264 : !$OMP END CRITICAL(omp_dbt_reshape)
265 : !$OMP BARRIER
266 :
267 : !$OMP MASTER
268 : CALL comm%alltoall(counts_send, counts_recv, 2)
269 : !$OMP END MASTER
270 : !$OMP BARRIER
271 : CALL timestop(handle2)
272 :
273 : CALL timeset(TRIM(name_pack), handle2)
274 : !$OMP MASTER
275 : CALL grow_buffer(buffer_send, counts_send(1, :), counts_send(2, :), nindex)
276 : CALL grow_buffer(buffer_recv, counts_recv(1, :), counts_recv(2, :), nindex)
277 : !$OMP END MASTER
278 : !$OMP BARRIER
279 :
280 : DO isource = 1, SIZE(sources)
281 : tensor_in => sources(isource)%tensor
282 : IF (.NOT. ASSOCIATED(tensor_in)) CYCLE
283 : CALL dbt_iterator_start(iter, tensor_in)
284 : DO WHILE (dbt_iterator_blocks_left(iter))
285 : CALL dbt_iterator_next_block(iter, ind_nd, blk_size=blk_size, blk_offset=blk_offset)
286 : inside = .TRUE.
287 : IF (do_bounds) THEN
288 : CALL block_bounds(bounds_prv, blk_offset, blk_size, overlap, inside)
289 : IF (.NOT. overlap) CYCLE
290 : END IF
291 : dest = ind_nd
292 : IF (PRESENT(index_map)) THEN
293 : DO d = 1, SIZE(dest)
294 : dest(d) = index_map(ind_nd(d), d)
295 : END DO
296 : IF (ANY(dest == 0)) CYCLE
297 : END IF
298 : IF (PRESENT(stack_dim)) dest(stack_dim) = dest(stack_dim) + sources(isource)%stack_offset
299 : CALL dbt_get_stored_coordinates(tensor_out, dest, iproc)
300 : first_offset = -1
301 : DO igroup = 1, ngroups
302 : IF (PRESENT(group_mask)) THEN
303 : IF (.NOT. group_mask(igroup, ind_nd(group_dim))) CYCLE
304 : END IF
305 : peer = (igroup - 1)*group_size + iproc
306 : ! Insert block data, reading the source only once for all replicas.
307 : ndata = PRODUCT(blk_size)
308 : ndata_send_mythread(peer) = ndata_send_mythread(peer) - ndata
309 : offset = buffer_send%msg_offset(peer) + ndata_send_mythread(peer)
310 : IF (first_offset < 0) THEN
311 : CALL dbt_get_block_data(tensor_in, ind_nd, blk_size, buffer_send%msg(offset + 1:offset + ndata), found)
312 : CPASSERT(found)
313 : IF (.NOT. inside) THEN
314 : CALL block_crop(bounds_prv, blk_offset, blk_size, buffer_send%msg(offset + 1:offset + ndata))
315 : END IF
316 : first_offset = offset
317 : ELSE
318 : buffer_send%msg(offset + 1:offset + ndata) = buffer_send%msg(first_offset + 1:first_offset + ndata)
319 : END IF
320 : ! Payload offsets in the metadata are relative to the peer's segment.
321 : nblks_send_mythread(peer) = nblks_send_mythread(peer) - 1
322 : iblk = buffer_send%indx_offset(peer) + nblks_send_mythread(peer) + 1
323 : buffer_send%indx(2:, iblk) = dest(:)
324 : buffer_send%indx(1, iblk) = ndata_send_mythread(peer)
325 : END DO
326 : END DO
327 : CALL dbt_iterator_stop(iter)
328 : END DO
329 : DEALLOCATE (nblks_send_mythread, ndata_send_mythread)
330 : !$OMP BARRIER
331 : CALL timestop(handle2)
332 :
333 : !$OMP MASTER
334 : IF (numnodes > 1) THEN
335 : CALL exchange_matrix_buffers(comm, buffer_recv, buffer_send)
336 : ELSE
337 : ! A single rank receives what it sends: the buffers trade places.
338 : CALL swap_buffers(buffer_send, buffer_recv)
339 : END IF
340 : !$OMP END MASTER
341 : !$OMP BARRIER
342 :
343 : CALL timeset(routineN//"_reserve", handle2)
344 : block_range = get_limit(buffer_recv%indx_offset(numnodes), omp_get_num_threads(), omp_get_thread_num())
345 : CALL dbt_reserve_blocks(tensor_out, TRANSPOSE(buffer_recv%indx(2:, block_range(1):block_range(2))))
346 : !$OMP BARRIER
347 : CALL timestop(handle2)
348 :
349 : CALL timeset(TRIM(name_unpack), handle2)
350 : DO iproc = 0, numnodes - 1
351 : !$OMP DO
352 : DO iblk = buffer_recv%indx_offset(iproc) + 1, buffer_recv%indx_offset(iproc + 1)
353 : ind_nd(:) = buffer_recv%indx(2:, iblk)
354 : CALL dbt_blk_sizes(tensor_out, ind_nd, blk_size)
355 : offset = buffer_recv%msg_offset(iproc) + buffer_recv%indx(1, iblk)
356 : ndata = PRODUCT(blk_size)
357 : CALL dbt_put_block_data(tensor_out, ind_nd, blk_size, buffer_recv%msg(offset + 1:offset + ndata), &
358 : summation=summation)
359 : END DO
360 : !$OMP END DO
361 : END DO
362 : CALL timestop(handle2)
363 : !$OMP END PARALLEL
364 :
365 330890 : DEALLOCATE (counts_send, counts_recv)
366 :
367 330890 : IF (move_prv) THEN
368 210866 : CALL timeset(routineN//"_clear", handle2)
369 421732 : DO isource = 1, SIZE(sources)
370 421732 : IF (ASSOCIATED(sources(isource)%tensor)) CALL dbt_clear(sources(isource)%tensor)
371 : END DO
372 210866 : CALL timestop(handle2)
373 : END IF
374 :
375 330890 : CALL timestop(handle)
376 661780 : END SUBROUTINE dbt_reshape_sources
377 :
378 : ! **************************************************************************************************
379 : !> \brief Release the communication buffers that persist across calls of dbt_reshape.
380 : !> \author Hans Pabst
381 : ! **************************************************************************************************
382 2910 : SUBROUTINE dbt_reshape_release()
383 2910 : CALL release_buffer(buffer_send)
384 2910 : CALL release_buffer(buffer_recv)
385 2910 : END SUBROUTINE dbt_reshape_release
386 :
387 : ! **************************************************************************************************
388 : !> \brief Whether the tensor's index order differs from its matrix layout (blocks need a reshape).
389 : !> \author Hans Pabst
390 : ! **************************************************************************************************
391 746470 : FUNCTION permuted(tensor)
392 : TYPE(dbt_type), INTENT(IN) :: tensor
393 : LOGICAL :: permuted
394 :
395 : INTEGER :: i
396 :
397 : ASSOCIATE (map1_2d => tensor%nd_index_blk%map1_2d, map2_2d => tensor%nd_index_blk%map2_2d)
398 8698222 : permuted = .NOT. ALL([map1_2d, map2_2d] == [(i, i=1, SIZE(map1_2d) + SIZE(map2_2d))])
399 : END ASSOCIATE
400 746470 : END FUNCTION permuted
401 :
402 : ! **************************************************************************************************
403 : !> \brief Whether a block (element offset and size per dimension) overlaps the bounds, and whether
404 : !> it lies entirely inside them.
405 : !> \author Hans Pabst
406 : ! **************************************************************************************************
407 173988 : PURE SUBROUTINE block_bounds(bounds, blk_offset, blk_size, overlap, inside)
408 : INTEGER, DIMENSION(:, :), INTENT(IN) :: bounds
409 : INTEGER, DIMENSION(:), INTENT(IN) :: blk_offset, blk_size
410 : LOGICAL, INTENT(OUT) :: overlap, inside
411 :
412 430562 : overlap = ALL(blk_offset <= bounds(2, :) .AND. bounds(1, :) <= blk_offset + blk_size - 1)
413 430542 : inside = ALL(bounds(1, :) <= blk_offset .AND. blk_offset + blk_size - 1 <= bounds(2, :))
414 173988 : END SUBROUTINE block_bounds
415 :
416 : ! **************************************************************************************************
417 : !> \brief Zero the elements of a block (column-major data) outside the bounds (as dbt_crop).
418 : !> \author Hans Pabst
419 : ! **************************************************************************************************
420 5 : PURE SUBROUTINE block_crop(bounds, blk_offset, blk_size, data)
421 : INTEGER, DIMENSION(:, :), INTENT(IN) :: bounds
422 : INTEGER, DIMENSION(:), INTENT(IN) :: blk_offset, blk_size
423 : REAL(dp), DIMENSION(:), INTENT(INOUT) :: data
424 :
425 : INTEGER :: d, i
426 10 : INTEGER, DIMENSION(SIZE(blk_size)) :: hi, idx, lo
427 :
428 20 : lo(:) = MAX(bounds(1, :) - blk_offset + 1, 1)
429 20 : hi(:) = MIN(bounds(2, :) - blk_offset + 1, blk_size)
430 20 : idx(:) = 1
431 5801 : DO i = 1, SIZE(data)
432 13608 : IF (ANY(idx < lo .OR. hi < idx)) data(i) = 0.0_dp
433 6574 : DO d = 1, SIZE(idx) ! next column-major index
434 6564 : idx(d) = idx(d) + 1
435 6564 : IF (idx(d) <= blk_size(d)) EXIT
436 6569 : idx(d) = 1
437 : END DO
438 : END DO
439 5 : END SUBROUTINE block_crop
440 :
441 : ! **************************************************************************************************
442 : !> \brief Lay out a persistent buffer for the given per-peer counts (as allocate_matrix_buffers),
443 : !> growing its arrays only when they are too small.
444 : !> \author Hans Pabst
445 : ! **************************************************************************************************
446 661780 : SUBROUTINE grow_buffer(buffer, ndata, nblks, nindex)
447 : TYPE(matrix_buffer_type), INTENT(INOUT) :: buffer
448 : INTEGER, DIMENSION(0:), INTENT(IN) :: ndata, nblks
449 : INTEGER, INTENT(IN) :: nindex
450 :
451 : INTEGER :: iproc, nprocs
452 :
453 661780 : nprocs = SIZE(ndata)
454 661780 : IF (ALLOCATED(buffer%msg_offset)) THEN
455 660680 : IF (SIZE(buffer%msg_offset) /= nprocs + 1) DEALLOCATE (buffer%msg_offset, buffer%indx_offset)
456 : END IF
457 670084 : IF (.NOT. ALLOCATED(buffer%msg_offset)) ALLOCATE (buffer%msg_offset(0:nprocs), buffer%indx_offset(0:nprocs))
458 661780 : buffer%msg_offset(0) = 0
459 661780 : buffer%indx_offset(0) = 0
460 1824508 : DO iproc = 0, nprocs - 1
461 1162728 : buffer%msg_offset(iproc + 1) = buffer%msg_offset(iproc) + ndata(iproc)
462 1824508 : buffer%indx_offset(iproc + 1) = buffer%indx_offset(iproc) + nblks(iproc)
463 : END DO
464 :
465 661780 : IF (ALLOCATED(buffer%msg)) THEN
466 660680 : IF (SIZE(buffer%msg) < buffer%msg_offset(nprocs)) DEALLOCATE (buffer%msg)
467 : END IF
468 661780 : IF (.NOT. ALLOCATED(buffer%msg)) THEN
469 6321 : ALLOCATE (buffer%msg(MAX(buffer%msg_offset(nprocs) + buffer%msg_offset(nprocs)/4, 1)))
470 : END IF
471 : ! records are exchanged as whole columns: the leading dimension must match
472 661780 : IF (ALLOCATED(buffer%indx)) THEN
473 660680 : IF (SIZE(buffer%indx, 1) /= nindex .OR. SIZE(buffer%indx, 2) < buffer%indx_offset(nprocs)) THEN
474 54243 : DEALLOCATE (buffer%indx)
475 : END IF
476 : END IF
477 661780 : IF (.NOT. ALLOCATED(buffer%indx)) THEN
478 221372 : ALLOCATE (buffer%indx(nindex, MAX(buffer%indx_offset(nprocs) + buffer%indx_offset(nprocs)/4, 1)))
479 : END IF
480 661780 : END SUBROUTINE grow_buffer
481 :
482 : ! **************************************************************************************************
483 : !> \brief Exchange the contents of two buffers without copying.
484 : !> \author Hans Pabst
485 : ! **************************************************************************************************
486 80416 : SUBROUTINE swap_buffers(a, b)
487 : TYPE(matrix_buffer_type), INTENT(INOUT) :: a, b
488 :
489 80416 : TYPE(matrix_buffer_type) :: t
490 :
491 80416 : CALL MOVE_ALLOC(a%msg, t%msg)
492 80416 : CALL MOVE_ALLOC(b%msg, a%msg)
493 80416 : CALL MOVE_ALLOC(t%msg, b%msg)
494 80416 : CALL MOVE_ALLOC(a%indx, t%indx)
495 80416 : CALL MOVE_ALLOC(b%indx, a%indx)
496 80416 : CALL MOVE_ALLOC(t%indx, b%indx)
497 80416 : CALL MOVE_ALLOC(a%msg_offset, t%msg_offset)
498 80416 : CALL MOVE_ALLOC(b%msg_offset, a%msg_offset)
499 80416 : CALL MOVE_ALLOC(t%msg_offset, b%msg_offset)
500 80416 : CALL MOVE_ALLOC(a%indx_offset, t%indx_offset)
501 80416 : CALL MOVE_ALLOC(b%indx_offset, a%indx_offset)
502 80416 : CALL MOVE_ALLOC(t%indx_offset, b%indx_offset)
503 80416 : END SUBROUTINE swap_buffers
504 :
505 : ! **************************************************************************************************
506 : !> \brief Deallocate the arrays of a buffer.
507 : !> \author Hans Pabst
508 : ! **************************************************************************************************
509 5820 : SUBROUTINE release_buffer(buffer)
510 : TYPE(matrix_buffer_type), INTENT(INOUT) :: buffer
511 :
512 5820 : IF (ALLOCATED(buffer%msg)) DEALLOCATE (buffer%msg)
513 5820 : IF (ALLOCATED(buffer%indx)) DEALLOCATE (buffer%indx)
514 5820 : IF (ALLOCATED(buffer%msg_offset)) DEALLOCATE (buffer%msg_offset)
515 5820 : IF (ALLOCATED(buffer%indx_offset)) DEALLOCATE (buffer%indx_offset)
516 5820 : END SUBROUTINE release_buffer
517 :
518 0 : END MODULE dbt_reshape_ops
|
