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 : block_nd, create_block, destroy_block, dbt_iterator_type, dbt_iterator_next_block, &
20 : dbt_iterator_blocks_left, dbt_iterator_start, dbt_iterator_stop, dbt_get_block, &
21 : dbt_reserve_blocks, dbt_put_block
22 : USE dbt_types, ONLY: dbt_blk_sizes, &
23 : dbt_type, &
24 : ndims_tensor, &
25 : dbt_get_stored_coordinates, &
26 : dbt_clear
27 : USE matrix_communication, ONLY: allocate_matrix_buffers, exchange_matrix_buffers, matrix_buffer_type
28 : USE message_passing, ONLY: mp_comm_type
29 : USE util, ONLY: get_limit
30 :
31 : #include "../base/base_uses.f90"
32 :
33 : IMPLICIT NONE
34 : PRIVATE
35 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_reshape_ops'
36 :
37 : PUBLIC :: dbt_reshape
38 :
39 : CONTAINS
40 :
41 : ! **************************************************************************************************
42 : !> \brief copy data (involves reshape)
43 : !> tensor_out = tensor_out + tensor_in move_data memory optimization:
44 : !> transfer data from tensor_in to tensor_out s.t. tensor_in is empty on return
45 : !> \author Ole Schuett
46 : ! **************************************************************************************************
47 292087 : SUBROUTINE dbt_reshape(tensor_in, tensor_out, summation, move_data)
48 :
49 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in, tensor_out
50 : LOGICAL, INTENT(IN), OPTIONAL :: summation
51 : LOGICAL, INTENT(IN), OPTIONAL :: move_data
52 :
53 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_reshape'
54 :
55 : INTEGER :: iproc, numnodes, &
56 : handle, iblk, offset, ndata
57 : INTEGER, DIMENSION(2) :: block_range
58 : TYPE(dbt_iterator_type) :: iter
59 292087 : TYPE(block_nd) :: blk_data
60 292087 : TYPE(matrix_buffer_type) :: buffer_recv, buffer_send
61 292087 : INTEGER, DIMENSION(ndims_tensor(tensor_in)) :: blk_size, ind_nd
62 : LOGICAL :: found, summation_prv, move_prv
63 :
64 292087 : INTEGER, ALLOCATABLE, DIMENSION(:) :: nblks_send_total, ndata_send_total, &
65 292087 : nblks_recv_total, ndata_recv_total, &
66 292087 : nblks_send_mythread, ndata_send_mythread
67 : TYPE(mp_comm_type) :: mp_comm
68 :
69 292087 : CALL timeset(routineN, handle)
70 :
71 292087 : IF (PRESENT(summation)) THEN
72 103487 : summation_prv = summation
73 : ELSE
74 : summation_prv = .FALSE.
75 : END IF
76 :
77 292087 : IF (PRESENT(move_data)) THEN
78 292087 : move_prv = move_data
79 : ELSE
80 : move_prv = .FALSE.
81 : END IF
82 :
83 292087 : CPASSERT(tensor_out%valid)
84 :
85 292087 : IF (.NOT. summation_prv) CALL dbt_clear(tensor_out)
86 :
87 292087 : mp_comm = tensor_in%pgrid%mp_comm_2d
88 292087 : numnodes = mp_comm%num_pe
89 1168348 : ALLOCATE (nblks_send_total(0:numnodes - 1), ndata_send_total(0:numnodes - 1), source=0)
90 876261 : ALLOCATE (nblks_recv_total(0:numnodes - 1), ndata_recv_total(0:numnodes - 1), source=0)
91 :
92 : !$OMP PARALLEL DEFAULT(OMP_DEFAULT_NONE_WITH_OOP) &
93 : !$OMP SHARED(tensor_in,tensor_out,summation) &
94 : !$OMP SHARED(buffer_send,buffer_recv,mp_comm,numnodes) &
95 : !$OMP SHARED(nblks_send_total,ndata_send_total,nblks_recv_total,ndata_recv_total) &
96 : !$OMP PRIVATE(nblks_send_mythread,ndata_send_mythread) &
97 : !$OMP PRIVATE(iter,ind_nd,blk_size,blk_data,found,iproc) &
98 292087 : !$OMP PRIVATE(block_range,offset,ndata,iblk)
99 : ALLOCATE (nblks_send_mythread(0:numnodes - 1), ndata_send_mythread(0:numnodes - 1), source=0)
100 :
101 : CALL dbt_iterator_start(iter, tensor_in)
102 : DO WHILE (dbt_iterator_blocks_left(iter))
103 : CALL dbt_iterator_next_block(iter, ind_nd, blk_size=blk_size)
104 : CALL dbt_get_stored_coordinates(tensor_out, ind_nd, iproc)
105 : nblks_send_mythread(iproc) = nblks_send_mythread(iproc) + 1
106 : ndata_send_mythread(iproc) = ndata_send_mythread(iproc) + PRODUCT(blk_size)
107 : END DO
108 : CALL dbt_iterator_stop(iter)
109 : !$OMP CRITICAL(omp_dbt_reshape)
110 : nblks_send_total(:) = nblks_send_total(:) + nblks_send_mythread(:)
111 : ndata_send_total(:) = ndata_send_total(:) + ndata_send_mythread(:)
112 : nblks_send_mythread(:) = nblks_send_total(:) ! current totals indicate slot for this thread
113 : ndata_send_mythread(:) = ndata_send_total(:)
114 : !$OMP END CRITICAL(omp_dbt_reshape)
115 : !$OMP BARRIER
116 :
117 : !$OMP MASTER
118 : CALL mp_comm%alltoall(nblks_send_total, nblks_recv_total, 1)
119 : CALL mp_comm%alltoall(ndata_send_total, ndata_recv_total, 1)
120 : CALL allocate_matrix_buffers(ndata_send_total, ndata_recv_total, ndims_tensor(tensor_in) + 1, &
121 : buffer_send, buffer_recv, nblks_send_total, nblks_recv_total)
122 : !$OMP END MASTER
123 : !$OMP BARRIER
124 :
125 : CALL dbt_iterator_start(iter, tensor_in)
126 : DO WHILE (dbt_iterator_blocks_left(iter))
127 : CALL dbt_iterator_next_block(iter, ind_nd, blk_size=blk_size)
128 : CALL dbt_get_stored_coordinates(tensor_out, ind_nd, iproc)
129 : CALL dbt_get_block(tensor_in, ind_nd, blk_data, found)
130 : CPASSERT(found)
131 : ! insert block data
132 : ndata = PRODUCT(blk_size)
133 : ndata_send_mythread(iproc) = ndata_send_mythread(iproc) - ndata
134 : offset = buffer_send%msg_offset(iproc) + ndata_send_mythread(iproc)
135 : buffer_send%msg(offset + 1:offset + ndata) = blk_data%blk(:)
136 : ! insert block index
137 : nblks_send_mythread(iproc) = nblks_send_mythread(iproc) - 1
138 : iblk = buffer_send%indx_offset(iproc) + nblks_send_mythread(iproc) + 1
139 : buffer_send%indx(2:, iblk) = ind_nd(:)
140 : ! Payload offsets in the metadata are relative to the peer's segment.
141 : buffer_send%indx(1, iblk) = ndata_send_mythread(iproc)
142 : CALL destroy_block(blk_data)
143 : END DO
144 : CALL dbt_iterator_stop(iter)
145 : DEALLOCATE (nblks_send_mythread, ndata_send_mythread)
146 : !$OMP BARRIER
147 :
148 : !$OMP MASTER
149 : IF (numnodes > 1) THEN
150 : CALL exchange_matrix_buffers(mp_comm, buffer_recv, buffer_send)
151 : DEALLOCATE (buffer_send%msg, buffer_send%indx)
152 : ELSE
153 : ! The single-rank receive buffer takes ownership of the packed arrays.
154 : CALL MOVE_ALLOC(buffer_send%msg, buffer_recv%msg)
155 : CALL MOVE_ALLOC(buffer_send%indx, buffer_recv%indx)
156 : END IF
157 : DEALLOCATE (buffer_send%msg_offset, buffer_send%indx_offset)
158 : !$OMP END MASTER
159 : !$OMP BARRIER
160 :
161 : block_range = get_limit(SIZE(buffer_recv%indx, 2), omp_get_num_threads(), omp_get_thread_num())
162 : CALL dbt_reserve_blocks(tensor_out, TRANSPOSE(buffer_recv%indx(2:, block_range(1):block_range(2))))
163 :
164 : DO iproc = 0, numnodes - 1
165 : !$OMP DO
166 : DO iblk = buffer_recv%indx_offset(iproc) + 1, buffer_recv%indx_offset(iproc + 1)
167 : ind_nd(:) = buffer_recv%indx(2:, iblk)
168 : CALL dbt_blk_sizes(tensor_out, ind_nd, blk_size)
169 : offset = buffer_recv%msg_offset(iproc) + buffer_recv%indx(1, iblk)
170 : ndata = PRODUCT(blk_size)
171 : CALL create_block(blk_data, blk_size, &
172 : array=buffer_recv%msg(offset + 1:offset + ndata))
173 : CALL dbt_put_block(tensor_out, ind_nd, blk_data, summation=summation)
174 : CALL destroy_block(blk_data)
175 : END DO
176 : !$OMP END DO
177 : END DO
178 :
179 : !$OMP MASTER
180 : DEALLOCATE (buffer_recv%msg, buffer_recv%indx, buffer_recv%msg_offset, buffer_recv%indx_offset)
181 : !$OMP END MASTER
182 : !$OMP END PARALLEL
183 :
184 292087 : DEALLOCATE (nblks_recv_total, ndata_recv_total)
185 292087 : DEALLOCATE (nblks_send_total, ndata_send_total)
186 :
187 292087 : IF (move_prv) CALL dbt_clear(tensor_in)
188 :
189 292087 : CALL timestop(handle)
190 584174 : END SUBROUTINE dbt_reshape
191 :
192 : END MODULE dbt_reshape_ops
|
