1!--------------------------------------------------------------------------------------------------!
2!   CP2K: A general program to perform molecular dynamics simulations                              !
3!   Copyright (C) 2000 - 2019  CP2K developers group                                               !
4!--------------------------------------------------------------------------------------------------!
5
6MODULE qs_fb_com_tasks_types
7
8   USE cp_para_types,                   ONLY: cp_para_env_type
9   USE dbcsr_api,                       ONLY: dbcsr_get_block_p,&
10                                              dbcsr_get_info,&
11                                              dbcsr_put_block,&
12                                              dbcsr_type
13   USE kinds,                           ONLY: dp,&
14                                              int_4,&
15                                              int_8
16   USE memory_utilities,                ONLY: reallocate
17   USE message_passing,                 ONLY: mp_alltoall
18   USE qs_fb_matrix_data_types,         ONLY: fb_matrix_data_add,&
19                                              fb_matrix_data_get,&
20                                              fb_matrix_data_has_data,&
21                                              fb_matrix_data_obj
22   USE util,                            ONLY: sort
23#include "./base/base_uses.f90"
24
25   IMPLICIT NONE
26
27   PRIVATE
28
29! public parameters:
30   PUBLIC :: TASK_N_RECORDS, &
31             TASK_DEST, &
32             TASK_SRC, &
33             TASK_PAIR, &
34             TASK_COST
35
36! public types
37   PUBLIC :: fb_com_tasks_obj, &
38             fb_com_atom_pairs_obj
39
40! public methods
41!API
42   PUBLIC :: fb_com_tasks_release, &
43             fb_com_tasks_nullify, &
44             fb_com_tasks_create, &
45             fb_com_tasks_get, &
46             fb_com_tasks_set, &
47             fb_com_tasks_transpose_dest_src, &
48             fb_com_tasks_build_atom_pairs, &
49             fb_com_tasks_encode_pair, &
50             fb_com_tasks_decode_pair, &
51             fb_com_atom_pairs_release, &
52             fb_com_atom_pairs_nullify, &
53             fb_com_atom_pairs_has_data, &
54             fb_com_atom_pairs_create, &
55             fb_com_atom_pairs_init, &
56             fb_com_atom_pairs_get, &
57             fb_com_atom_pairs_decode, &
58             fb_com_atom_pairs_calc_buffer_sizes, &
59             fb_com_atom_pairs_gather_blks, &
60             fb_com_atom_pairs_distribute_blks
61
62   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'qs_fb_com_tasks_types'
63   INTEGER, PRIVATE, SAVE :: last_fb_com_tasks_id = 0
64   INTEGER, PRIVATE, SAVE :: last_fb_com_atom_pairs_id = 0
65
66! **********************************************************************
67! explaination on format of task lists (same for tasks_recv and tasks_send):
68! tasks_recv has dimension (4, ntasks_recv), and stores information on
69! the block to be copied or transfered
70! - tasks_recv(TASK_DEST,itask) = destination MPI rank of itask-th task
71! - tasks_recv(TASK_SRC,itask) = source MPI rank of itask-th task
72! - tasks_recv(TASK_PAIR,itask) = compressed pair indices of the block of itask-th task
73! - tasks_recv(TASK_COST,itask) = the cost of itask-th task
74!
75! number of record slots in each task in the task lists
76   INTEGER, PARAMETER :: TASK_N_RECORDS = 4
77! the indices for the records (1:TASK_DIM) in a task
78   INTEGER, PARAMETER :: TASK_DEST = 1, &
79                         TASK_SRC = 2, &
80                         TASK_PAIR = 3, &
81                         TASK_COST = 4
82! **********************************************************************
83
84! **********************************************************************
85!> \brief data content for communication tasks used for send and receive
86!>        matrix blocks
87!> \param id_nr     : unique id for the object
88!> \param ref_count : reference count on the object
89!> \param tasks     : the list of communication tasks, which is
90!>                    represented by a 2D array, first dim stores
91!>                    info for the communication: src and desc procs
92!>                    and the atomic pair indexing the matrix block
93!>                    to be communicated, etc.
94!> \param task_dim  : the size of the first dimension of tasks
95!> \param ntasks    : total number of local tasks
96!> \param nencode   : the total number of atoms used for encoding
97!>                    the block coordinates (iatom, jatom)
98!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
99! **********************************************************************
100   TYPE fb_com_tasks_data
101      INTEGER :: id_nr, ref_count
102      ! use pure integer arrays to facilitate easier MPI coms
103      INTEGER(KIND=int_8), DIMENSION(:, :), POINTER :: tasks
104      INTEGER :: task_dim
105      INTEGER :: ntasks
106      INTEGER :: nencode
107   END TYPE fb_com_tasks_data
108
109!**********************************************************************
110!> \brief defines a fb_com_tasks object
111!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
112!**********************************************************************
113   TYPE fb_com_tasks_obj
114      TYPE(fb_com_tasks_data), POINTER, PRIVATE :: obj
115   END TYPE fb_com_tasks_obj
116
117! **********************************************************************
118!> \brief data content for the list of block coordinates with the
119!>        associated src/dest proc id for communication. These will be
120!>        generated from the fb_com_tasks object
121!> \param id_nr         : unique id for the object
122!> \param ref_count     : reference count on the object
123!> \param pairs         : the list of communication tasks, which is
124!>                        represented by a 2D array, first dim stores
125!>                        info for the communication: src and desc procs
126!>                        and the atomic pair indexing the matrix block
127!>                        to be communicated, etc.
128!> \param npairs        : number of blks to be communicated in the atom
129!>                        pair list
130!> \param natoms_encode : the total number of atoms used for encoding
131!>                        the proc + block coordinates (pe, iatom, jatom)
132!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
133! **********************************************************************
134   TYPE fb_com_atom_pairs_data
135      INTEGER :: id_nr, ref_count
136      INTEGER(KIND=int_8), DIMENSION(:), POINTER :: pairs
137      INTEGER :: npairs
138      INTEGER :: natoms_encode
139   END TYPE fb_com_atom_pairs_data
140
141! **********************************************************************
142!> \brief defines a fb_com_atom_pairs object
143!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
144! **********************************************************************
145   TYPE fb_com_atom_pairs_obj
146      TYPE(fb_com_atom_pairs_data), POINTER, PRIVATE :: obj
147   END TYPE fb_com_atom_pairs_obj
148
149CONTAINS
150
151! **********************************************************************
152!> \brief Retains an fb_com_tasks object
153!> \param com_tasks the fb_com_tasks object, its content must not be
154!>                   NULL or UNDEFINED
155!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
156! **************************************************************************************************
157   SUBROUTINE fb_com_tasks_retain(com_tasks)
158      TYPE(fb_com_tasks_obj), INTENT(IN)                 :: com_tasks
159
160      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_retain', &
161         routineP = moduleN//':'//routineN
162
163      CPASSERT(ASSOCIATED(com_tasks%obj))
164      CPASSERT(com_tasks%obj%ref_count > 0)
165      com_tasks%obj%ref_count = com_tasks%obj%ref_count + 1
166   END SUBROUTINE fb_com_tasks_retain
167
168! **********************************************************************
169!> \brief Retains an fb_com_atom_pairs object
170!> \param atom_pairs the fb_com_atom_pairs object, its content must not be
171!>                   NULL or UNDEFINED
172!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
173! **************************************************************************************************
174   SUBROUTINE fb_com_atom_pairs_retain(atom_pairs)
175      TYPE(fb_com_atom_pairs_obj), INTENT(IN)            :: atom_pairs
176
177      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_retain', &
178         routineP = moduleN//':'//routineN
179
180      CPASSERT(ASSOCIATED(atom_pairs%obj))
181      CPASSERT(atom_pairs%obj%ref_count > 0)
182      atom_pairs%obj%ref_count = atom_pairs%obj%ref_count + 1
183   END SUBROUTINE fb_com_atom_pairs_retain
184
185! **********************************************************************
186!> \brief Releases an fb_com_tasks object
187!> \param com_tasks the fb_com_tasks object, its content must not be
188!>                   UNDEFINED, and the subroutine does nothing if the
189!>                   content points to NULL
190!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
191! **************************************************************************************************
192   SUBROUTINE fb_com_tasks_release(com_tasks)
193      TYPE(fb_com_tasks_obj), INTENT(INOUT)              :: com_tasks
194
195      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_release', &
196         routineP = moduleN//':'//routineN
197
198      IF (ASSOCIATED(com_tasks%obj)) THEN
199         CPASSERT(com_tasks%obj%ref_count > 0)
200         com_tasks%obj%ref_count = com_tasks%obj%ref_count - 1
201         IF (com_tasks%obj%ref_count == 0) THEN
202            com_tasks%obj%ref_count = 1
203            IF (ASSOCIATED(com_tasks%obj%tasks)) THEN
204               DEALLOCATE (com_tasks%obj%tasks)
205            END IF
206            com_tasks%obj%ref_count = 0
207            DEALLOCATE (com_tasks%obj)
208         END IF
209      ELSE
210         NULLIFY (com_tasks%obj)
211      END IF
212   END SUBROUTINE fb_com_tasks_release
213
214! **********************************************************************
215!> \brief Releases an fb_com_atom_pairs object
216!> \param atom_pairs the fb_com_atom_pairs object, its content must not
217!>                    be UNDEFINED, and the subroutine does nothing if
218!>                    the content points to NULL
219!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
220! **************************************************************************************************
221   SUBROUTINE fb_com_atom_pairs_release(atom_pairs)
222      TYPE(fb_com_atom_pairs_obj), INTENT(INOUT)         :: atom_pairs
223
224      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_release', &
225         routineP = moduleN//':'//routineN
226
227      IF (ASSOCIATED(atom_pairs%obj)) THEN
228         CPASSERT(atom_pairs%obj%ref_count > 0)
229         atom_pairs%obj%ref_count = atom_pairs%obj%ref_count - 1
230         IF (atom_pairs%obj%ref_count == 0) THEN
231            atom_pairs%obj%ref_count = 1
232            IF (ASSOCIATED(atom_pairs%obj%pairs)) THEN
233               DEALLOCATE (atom_pairs%obj%pairs)
234            END IF
235            atom_pairs%obj%ref_count = 0
236            DEALLOCATE (atom_pairs%obj)
237         END IF
238      ELSE
239         NULLIFY (atom_pairs%obj)
240      END IF
241   END SUBROUTINE fb_com_atom_pairs_release
242
243! **********************************************************************
244!> \brief Nullifies a fb_com_tasks object, note that it does not release
245!>        the original object. This procedure is used to nullify the
246!>        pointer contained in the object which is used to associate to
247!>        the actual object content
248!> \param com_tasks the com_tasks object
249!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
250! **************************************************************************************************
251   SUBROUTINE fb_com_tasks_nullify(com_tasks)
252      TYPE(fb_com_tasks_obj), INTENT(INOUT)              :: com_tasks
253
254      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_nullify', &
255         routineP = moduleN//':'//routineN
256
257      NULLIFY (com_tasks%obj)
258   END SUBROUTINE fb_com_tasks_nullify
259
260! **********************************************************************
261!> \brief Nullifies a fb_com_atom_pairs object, note that it does not
262!>        release the original object. This procedure is used to nullify
263!>        the pointer contained in the object which is used to associate
264!>        to the actual object content
265!> \param atom_pairs the fb_com_atom_pairs object
266!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
267! **************************************************************************************************
268   SUBROUTINE fb_com_atom_pairs_nullify(atom_pairs)
269      TYPE(fb_com_atom_pairs_obj), INTENT(INOUT)         :: atom_pairs
270
271      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_nullify', &
272         routineP = moduleN//':'//routineN
273
274      NULLIFY (atom_pairs%obj)
275   END SUBROUTINE fb_com_atom_pairs_nullify
276
277! **********************************************************************
278!> \brief Associates one fb_com_tasks object to another
279!> \param a the fb_com_tasks object to be associated
280!> \param b the fb_com_tasks object that a is to be associated to
281!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
282! **************************************************************************************************
283   SUBROUTINE fb_com_tasks_associate(a, b)
284      TYPE(fb_com_tasks_obj), INTENT(OUT)                :: a
285      TYPE(fb_com_tasks_obj), INTENT(IN)                 :: b
286
287      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_associate', &
288         routineP = moduleN//':'//routineN
289
290      a%obj => b%obj
291   END SUBROUTINE fb_com_tasks_associate
292
293! **********************************************************************
294!> \brief Associates one fb_com_atom_pairs object to another
295!> \param a the fb_com_atom_pairs object to be associated
296!> \param b the fb_com_atom_pairs object that a is to be associated to
297!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
298! **************************************************************************************************
299   SUBROUTINE fb_com_atom_pairs_associate(a, b)
300      TYPE(fb_com_atom_pairs_obj), INTENT(OUT)           :: a
301      TYPE(fb_com_atom_pairs_obj), INTENT(IN)            :: b
302
303      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_associate', &
304         routineP = moduleN//':'//routineN
305
306      a%obj => b%obj
307   END SUBROUTINE fb_com_atom_pairs_associate
308
309! **********************************************************************
310!> \brief Checks if a fb_com_tasks object is associated with an actual
311!>        data content or not
312!> \param com_tasks the fb_com_tasks object
313!> \return ...
314!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
315! **************************************************************************************************
316   FUNCTION fb_com_tasks_has_data(com_tasks) RESULT(res)
317      TYPE(fb_com_tasks_obj), INTENT(IN)                 :: com_tasks
318      LOGICAL                                            :: res
319
320      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_has_data', &
321         routineP = moduleN//':'//routineN
322
323      res = ASSOCIATED(com_tasks%obj)
324   END FUNCTION fb_com_tasks_has_data
325
326! **********************************************************************
327!> \brief Checks if a fb_com_atom_pairs object is associated with an actual
328!>        data content or not
329!> \param atom_pairs the fb_com_atom_pairs object
330!> \return ...
331!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
332! **************************************************************************************************
333   FUNCTION fb_com_atom_pairs_has_data(atom_pairs) RESULT(res)
334      TYPE(fb_com_atom_pairs_obj), INTENT(IN)            :: atom_pairs
335      LOGICAL                                            :: res
336
337      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_has_data', &
338         routineP = moduleN//':'//routineN
339
340      res = ASSOCIATED(atom_pairs%obj)
341   END FUNCTION fb_com_atom_pairs_has_data
342
343! **********************************************************************
344!> \brief Creates and initialises an empty fb_com_tasks object
345!> \param com_tasks the fb_com_tasks object, its content must be NULL
346!>                   and cannot be UNDEFINED
347!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
348! **************************************************************************************************
349   SUBROUTINE fb_com_tasks_create(com_tasks)
350      TYPE(fb_com_tasks_obj), INTENT(INOUT)              :: com_tasks
351
352      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_create', &
353         routineP = moduleN//':'//routineN
354
355      CPASSERT(.NOT. ASSOCIATED(com_tasks%obj))
356      ALLOCATE (com_tasks%obj)
357      com_tasks%obj%task_dim = TASK_N_RECORDS
358      com_tasks%obj%ntasks = 0
359      com_tasks%obj%nencode = 0
360      NULLIFY (com_tasks%obj%tasks)
361      com_tasks%obj%ref_count = 1
362      com_tasks%obj%id_nr = last_fb_com_tasks_id + 1
363      last_fb_com_tasks_id = com_tasks%obj%id_nr
364   END SUBROUTINE fb_com_tasks_create
365
366! **********************************************************************
367!> \brief Creates and initialises an empty fb_com_atom_pairs object
368!> \param atom_pairs the fb_com_atom_pairs object, its content must be
369!>                    NULL and cannot be UNDEFINED
370!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
371! **************************************************************************************************
372   SUBROUTINE fb_com_atom_pairs_create(atom_pairs)
373      TYPE(fb_com_atom_pairs_obj), INTENT(INOUT)         :: atom_pairs
374
375      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_create', &
376         routineP = moduleN//':'//routineN
377
378      CPASSERT(.NOT. ASSOCIATED(atom_pairs%obj))
379      ALLOCATE (atom_pairs%obj)
380      atom_pairs%obj%npairs = 0
381      atom_pairs%obj%natoms_encode = 0
382      NULLIFY (atom_pairs%obj%pairs)
383      atom_pairs%obj%ref_count = 1
384      atom_pairs%obj%id_nr = last_fb_com_atom_pairs_id + 1
385      last_fb_com_atom_pairs_id = atom_pairs%obj%id_nr
386   END SUBROUTINE fb_com_atom_pairs_create
387
388! **********************************************************************
389!> \brief Initialises an fb_com_tasks object, and makes it empty
390!> \param com_tasks the fb_com_tasks object, its content must not be
391!>                   NULL or UNDEFINED
392!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
393! **************************************************************************************************
394   SUBROUTINE fb_com_tasks_init(com_tasks)
395      TYPE(fb_com_tasks_obj), INTENT(INOUT)              :: com_tasks
396
397      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_init', &
398         routineP = moduleN//':'//routineN
399
400      CPASSERT(ASSOCIATED(com_tasks%obj))
401      IF (ASSOCIATED(com_tasks%obj%tasks)) THEN
402         DEALLOCATE (com_tasks%obj%tasks)
403      END IF
404      com_tasks%obj%task_dim = TASK_N_RECORDS
405      com_tasks%obj%ntasks = 0
406      com_tasks%obj%nencode = 0
407   END SUBROUTINE fb_com_tasks_init
408
409! **********************************************************************
410!> \brief Initialises an fb_com_atom_pairs object, and makes it empty
411!> \param atom_pairs the fb_com_atom_pairs object, its content must not
412!>                    be NULL or UNDEFINED
413!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
414! **************************************************************************************************
415   SUBROUTINE fb_com_atom_pairs_init(atom_pairs)
416      TYPE(fb_com_atom_pairs_obj), INTENT(INOUT)         :: atom_pairs
417
418      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_init', &
419         routineP = moduleN//':'//routineN
420
421      CPASSERT(ASSOCIATED(atom_pairs%obj))
422      IF (ASSOCIATED(atom_pairs%obj%pairs)) THEN
423         DEALLOCATE (atom_pairs%obj%pairs)
424      END IF
425      atom_pairs%obj%npairs = 0
426      atom_pairs%obj%natoms_encode = 0
427   END SUBROUTINE fb_com_atom_pairs_init
428
429! **********************************************************************
430!> \brief Gets attributes from a fb_com_tasks object, one should only
431!>        access the data content in a fb_com_tasks object outside this
432!>        module via this procedure.
433!> \param com_tasks the fb_com_tasks object, its content must not be
434!>                   NULL or UNDEFINED
435!> \param task_dim [OPTIONAL]: if present, outputs com_tasks%obj%task_dim
436!> \param ntasks [OPTIONAL]: if present, outputs com_tasks%obj%ntasks
437!> \param nencode [OPTIONAL]: if present, outputs com_tasks%obj%nencode
438!> \param tasks [OPTIONAL]: if present, outputs pointer com_tasks%obj%tasks
439!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
440! **************************************************************************************************
441   SUBROUTINE fb_com_tasks_get(com_tasks, &
442                               task_dim, &
443                               ntasks, &
444                               nencode, &
445                               tasks)
446      TYPE(fb_com_tasks_obj), INTENT(IN)                 :: com_tasks
447      INTEGER, INTENT(OUT), OPTIONAL                     :: task_dim, ntasks, nencode
448      INTEGER(KIND=int_8), DIMENSION(:, :), OPTIONAL, &
449         POINTER                                         :: tasks
450
451      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_get', &
452         routineP = moduleN//':'//routineN
453
454      CPASSERT(ASSOCIATED(com_tasks%obj))
455      IF (PRESENT(task_dim)) task_dim = com_tasks%obj%task_dim
456      IF (PRESENT(ntasks)) ntasks = com_tasks%obj%ntasks
457      IF (PRESENT(nencode)) nencode = com_tasks%obj%nencode
458      IF (PRESENT(tasks)) tasks => com_tasks%obj%tasks
459   END SUBROUTINE fb_com_tasks_get
460
461! **********************************************************************
462!> \brief Gets attributes from a fb_com_atom_pairs object, one should
463!>        only access the data content in a fb_com_atom_pairs object
464!>        outside this module via this procedure.
465!> \param atom_pairs the fb_com_atom_pairs object, its content must not
466!>                    be NULL or UNDEFINED
467!> \param npairs [OPTIONAL]: if present, outputs atom_pairs%obj%npairs
468!> \param natoms_encode [OPTIONAL]: if present, outputs atom_pairs%obj%natoms_encode
469!> \param pairs [OPTIONAL]: if present, outputs pointer atom_pairs%obj%pairs
470!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
471! **************************************************************************************************
472   SUBROUTINE fb_com_atom_pairs_get(atom_pairs, &
473                                    npairs, &
474                                    natoms_encode, &
475                                    pairs)
476      TYPE(fb_com_atom_pairs_obj), INTENT(IN)            :: atom_pairs
477      INTEGER, INTENT(OUT), OPTIONAL                     :: npairs, natoms_encode
478      INTEGER(KIND=int_8), DIMENSION(:), OPTIONAL, &
479         POINTER                                         :: pairs
480
481      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_get', &
482         routineP = moduleN//':'//routineN
483
484      CPASSERT(ASSOCIATED(atom_pairs%obj))
485      IF (PRESENT(npairs)) npairs = atom_pairs%obj%npairs
486      IF (PRESENT(natoms_encode)) natoms_encode = atom_pairs%obj%natoms_encode
487      IF (PRESENT(pairs)) pairs => atom_pairs%obj%pairs
488   END SUBROUTINE fb_com_atom_pairs_get
489
490! **********************************************************************
491!> \brief Sets attributes in a fb_com_tasks object, one should only
492!>        access the data content in a fb_com_tasks object outside this
493!>        module via this procedure.
494!> \param com_tasks the fb_com_tasks object, its content must not be
495!>                   NULL or UNDEFINED
496!> \param task_dim [OPTIONAL]: if present, sets com_tasks%obj%task_dim
497!> \param ntasks [OPTIONAL]: if present, sets com_tasks%obj%ntasks
498!> \param nencode [OPTIONAL]: if present, sets com_tasks%obj%nencode
499!> \param tasks [OPTIONAL]: if present, associates pointer com_tasks%obj%tasks
500!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
501! **************************************************************************************************
502   SUBROUTINE fb_com_tasks_set(com_tasks, &
503                               task_dim, &
504                               ntasks, &
505                               nencode, &
506                               tasks)
507      TYPE(fb_com_tasks_obj), INTENT(INOUT)              :: com_tasks
508      INTEGER, INTENT(IN), OPTIONAL                      :: task_dim, ntasks, nencode
509      INTEGER(KIND=int_8), DIMENSION(:, :), OPTIONAL, &
510         POINTER                                         :: tasks
511
512      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_tasks_set', &
513         routineP = moduleN//':'//routineN
514
515      CPASSERT(ASSOCIATED(com_tasks%obj))
516      IF (PRESENT(task_dim)) com_tasks%obj%task_dim = task_dim
517      IF (PRESENT(ntasks)) com_tasks%obj%ntasks = ntasks
518      IF (PRESENT(nencode)) com_tasks%obj%nencode = nencode
519      IF (PRESENT(tasks)) THEN
520         IF (ASSOCIATED(com_tasks%obj%tasks)) THEN
521            DEALLOCATE (com_tasks%obj%tasks)
522         END IF
523         com_tasks%obj%tasks => tasks
524      END IF
525   END SUBROUTINE fb_com_tasks_set
526
527! **********************************************************************
528!> \brief Sets attributes in a fb_com_atom_pairs object, one should only
529!>        access the data content in a fb_com_atom_pairs object outside
530!>        this module via this procedure.
531!> \param atom_pairs the fb_com_atom_pairs object, its content must not
532!>                    be NULL or UNDEFINED
533!> \param npairs [OPTIONAL]: if present, sets atom_pairs%obj%npairs
534!> \param natoms_encode [OPTIONAL]: if present, sets atom_pairs%obj%natoms_encode
535!> \param pairs [OPTIONAL]: if present, associates pointer atom_pairs%obj%pairs
536!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
537! **************************************************************************************************
538   SUBROUTINE fb_com_atom_pairs_set(atom_pairs, &
539                                    npairs, &
540                                    natoms_encode, &
541                                    pairs)
542      TYPE(fb_com_atom_pairs_obj), INTENT(INOUT)         :: atom_pairs
543      INTEGER, INTENT(IN), OPTIONAL                      :: npairs, natoms_encode
544      INTEGER(KIND=int_8), DIMENSION(:), OPTIONAL, &
545         POINTER                                         :: pairs
546
547      CHARACTER(len=*), PARAMETER :: routineN = 'fb_com_atom_pairs_set', &
548         routineP = moduleN//':'//routineN
549
550      CPASSERT(ASSOCIATED(atom_pairs%obj))
551      IF (PRESENT(npairs)) atom_pairs%obj%npairs = npairs
552      IF (PRESENT(natoms_encode)) atom_pairs%obj%natoms_encode = natoms_encode
553      IF (PRESENT(pairs)) THEN
554         IF (ASSOCIATED(atom_pairs%obj%pairs)) THEN
555            DEALLOCATE (atom_pairs%obj%pairs)
556         END IF
557         atom_pairs%obj%pairs => pairs
558      END IF
559   END SUBROUTINE fb_com_atom_pairs_set
560
561! **********************************************************************
562!> \brief Start from a local set of tasks that has desc/src process equal
563!>        to the local MPI rank, communicate with other processes so
564!>        that a new local set of tasks is constructed with src/desc
565!>        process equal to the local MPI rank
566!> \param tasks_dest_is_me the local com_task object with all tasks
567!>                          having the desc process id equal to my_id
568!> \param direction direction of operation:
569!>                   ">" means from tasks_dest_is_me construct tasks_src_is_me
570!>                   "<" means from tasks_src_is_me construct tasks_dest_is_me
571!> \param tasks_src_is_me the local com_task object with all tasks
572!>                          having the src process id equal to my_id
573!> \param para_env CP2K parallel environment object that stores MPI related
574!>                  information of the current run
575!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
576! **************************************************************************************************
577   SUBROUTINE fb_com_tasks_transpose_dest_src(tasks_dest_is_me, &
578                                              direction, &
579                                              tasks_src_is_me, &
580                                              para_env)
581      TYPE(fb_com_tasks_obj), INTENT(INOUT)              :: tasks_dest_is_me
582      CHARACTER, INTENT(IN)                              :: direction
583      TYPE(fb_com_tasks_obj), INTENT(INOUT)              :: tasks_src_is_me
584      TYPE(cp_para_env_type), POINTER                    :: para_env
585
586      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_tasks_transpose_dest_src', &
587         routineP = moduleN//':'//routineN
588
589      INTEGER                                            :: handle, ii, ind, ipe, itask, jj, &
590                                                            nencode, ntasks_in, ntasks_out, rank, &
591                                                            rank_pos, task_dim
592      INTEGER(KIND=int_8), DIMENSION(:, :), POINTER      :: tasks_in, tasks_out
593      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: recv_buf, recv_disps, recv_sizes, &
594                                                            send_buf, send_disps, send_sizes
595
596      CALL timeset(routineN, handle)
597
598      NULLIFY (tasks_in, tasks_out)
599
600      IF (direction == "<") THEN
601         CALL fb_com_tasks_get(com_tasks=tasks_src_is_me, &
602                               task_dim=task_dim, &
603                               ntasks=ntasks_in, &
604                               tasks=tasks_in, &
605                               nencode=nencode)
606         rank_pos = TASK_DEST
607      ELSE
608         CALL fb_com_tasks_get(com_tasks=tasks_dest_is_me, &
609                               task_dim=task_dim, &
610                               ntasks=ntasks_in, &
611                               tasks=tasks_in, &
612                               nencode=nencode)
613         rank_pos = TASK_SRC
614      END IF
615
616      ! allocate local arrays
617      ALLOCATE (send_sizes(para_env%num_pe))
618      ALLOCATE (send_disps(para_env%num_pe))
619      ALLOCATE (send_buf(para_env%num_pe))
620
621      ALLOCATE (recv_sizes(para_env%num_pe))
622      ALLOCATE (recv_disps(para_env%num_pe))
623      ALLOCATE (recv_buf(para_env%num_pe))
624
625      ! first count how many local recv/send tasks need to be sent to
626      ! other processes, and share this information with the other
627      ! processes.  using send_buf as a temporary array for counting
628      send_buf = 0
629      ! looping over local task list
630      DO itask = 1, ntasks_in
631         rank = INT(tasks_in(rank_pos, itask)) + 1
632         send_buf(rank) = send_buf(rank) + 1
633      END DO
634
635      CALL mp_alltoall(send_buf, recv_buf, 1, para_env%group)
636
637      ! now that we know how many recv/send tasks to send, pack the
638      ! tasks, and send them around, so that the recv/send tasks are
639      ! sent to the correct src/dest processes, and these then are
640      ! collected into the send/recv tasks list on each of the src/dest
641      ! processes
642
643      send_sizes = 0
644      send_disps = 0
645      recv_sizes = 0
646      recv_disps = 0
647
648      ! work out the sizes of send and recv buffers and allocate them
649      send_sizes(1) = send_buf(1)*task_dim
650      recv_sizes(1) = recv_buf(1)*task_dim
651      DO ipe = 2, para_env%num_pe
652         send_sizes(ipe) = send_buf(ipe)*task_dim
653         send_disps(ipe) = send_disps(ipe - 1) + send_sizes(ipe - 1)
654         recv_sizes(ipe) = recv_buf(ipe)*task_dim
655         recv_disps(ipe) = recv_disps(ipe - 1) + recv_sizes(ipe - 1)
656      END DO
657
658      ! reallocate send and recv buffers to the correct sizes for
659      ! transferring the actual tasks
660      DEALLOCATE (send_buf)
661      DEALLOCATE (recv_buf)
662      ALLOCATE (send_buf(SUM(send_sizes)))
663      ALLOCATE (recv_buf(SUM(recv_sizes)))
664
665      ! now that the send buffer is of correct size, do packing
666      ! send_buf and recv_buf may be zero sized
667      IF (SIZE(send_buf) > 0) send_buf = 0
668      IF (SIZE(recv_buf) > 0) recv_buf = 0
669      send_sizes = 0
670      DO itask = 1, ntasks_in
671         rank = INT(tasks_in(rank_pos, itask)) + 1
672         DO ii = 1, task_dim
673            ind = send_disps(rank) + send_sizes(rank) + ii
674            send_buf(ind) = INT(tasks_in(ii, itask))
675         END DO
676         send_sizes(rank) = send_sizes(rank) + task_dim
677      END DO
678      ! do communication
679      CALL mp_alltoall(send_buf, send_sizes, send_disps, &
680                       recv_buf, recv_sizes, recv_disps, &
681                       para_env%group)
682
683      ! deallocate send buffers
684      DEALLOCATE (send_buf)
685      DEALLOCATE (send_sizes)
686      DEALLOCATE (send_disps)
687
688      ! allocate the output task list
689      ntasks_out = SUM(recv_sizes)/task_dim
690      ! this will not be deallocated in this subroutine
691      ALLOCATE (tasks_out(task_dim, ntasks_out))
692
693      ! do unpacking
694      itask = 0
695      DO ipe = 1, para_env%num_pe
696         DO ii = 0, recv_sizes(ipe)/task_dim - 1
697            itask = itask + 1
698            DO jj = 1, task_dim
699               ind = recv_disps(ipe) + ii*task_dim + jj
700               tasks_out(jj, itask) = recv_buf(ind)
701            END DO
702         END DO
703      END DO
704
705      ! set output tasks
706      IF (direction == "<") THEN
707         CALL fb_com_tasks_set(com_tasks=tasks_dest_is_me, &
708                               task_dim=task_dim, &
709                               ntasks=ntasks_out, &
710                               tasks=tasks_out, &
711                               nencode=nencode)
712      ELSE
713         CALL fb_com_tasks_set(com_tasks=tasks_src_is_me, &
714                               task_dim=task_dim, &
715                               ntasks=ntasks_out, &
716                               tasks=tasks_out, &
717                               nencode=nencode)
718      END IF
719
720      ! deallocate recv buffers
721      DEALLOCATE (recv_buf)
722      DEALLOCATE (recv_sizes)
723      DEALLOCATE (recv_disps)
724
725      CALL timestop(handle)
726
727   END SUBROUTINE fb_com_tasks_transpose_dest_src
728
729! **********************************************************************
730!> \brief Generate send or receive atom_pair lists from a com_tasks
731!>        object. atom_pair list is used as a condensed index for the
732!>        local/remote matrix blocks to be sent/received.
733!> \param com_tasks the com_tasks object
734!> \param atom_pairs fb_com_atom_pairs_obj containing  list of encoded
735!>                    atomic pair indices and the dest/src proc id for
736!>                    the matrix block to be sent/received.
737!> \param natoms_encode the total number of atoms the atomic pair indices
738!>                       corresponds to, and it is used for encode the
739!>                       atom_pairs values
740!> \param send_or_recv whether the atom_pair to be generated is for
741!>                      the local matrix blocks to be sent or the
742!>                      remote matrix blocks to be received for this MPI
743!>                      process
744!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
745! **************************************************************************************************
746   SUBROUTINE fb_com_tasks_build_atom_pairs(com_tasks, &
747                                            atom_pairs, &
748                                            natoms_encode, &
749                                            send_or_recv)
750      TYPE(fb_com_tasks_obj), INTENT(IN)                 :: com_tasks
751      TYPE(fb_com_atom_pairs_obj), INTENT(INOUT)         :: atom_pairs
752      INTEGER, INTENT(IN)                                :: natoms_encode
753      CHARACTER(len=*), INTENT(IN)                       :: send_or_recv
754
755      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_tasks_build_atom_pairs', &
756         routineP = moduleN//':'//routineN
757
758      INTEGER                                            :: handle, iatom, ii, itask, jatom, npairs, &
759                                                            ntasks, rank, rank_pos
760      INTEGER(KIND=int_8)                                :: pair
761      INTEGER(KIND=int_8), DIMENSION(:), POINTER         :: pairs
762      INTEGER(KIND=int_8), DIMENSION(:, :), POINTER      :: tasks
763      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: tmp_index
764      LOGICAL                                            :: check_ok
765
766      CALL timeset(routineN, handle)
767
768      NULLIFY (pairs, tasks)
769
770      check_ok = fb_com_atom_pairs_has_data(atom_pairs)
771      CPASSERT(check_ok)
772
773      ! initialise atom_pairs
774      CALL fb_com_atom_pairs_init(atom_pairs)
775
776      IF (TRIM(send_or_recv) == "send") THEN
777         rank_pos = TASK_DEST
778      ELSE
779         rank_pos = TASK_SRC
780      END IF
781
782      CALL fb_com_tasks_get(com_tasks=com_tasks, &
783                            ntasks=ntasks, &
784                            tasks=tasks)
785
786      ALLOCATE (pairs(ntasks))
787      ! we can have cases where ntasks == 0
788      IF (SIZE(pairs) > 0) pairs = 0
789      npairs = ntasks
790
791      DO itask = 1, ntasks
792         pair = tasks(TASK_PAIR, itask)
793         CALL fb_com_tasks_decode_pair(pair, iatom, jatom, natoms_encode)
794         rank = INT(tasks(rank_pos, itask))
795         CALL fb_com_atom_pairs_encode(pairs(itask), &
796                                       rank, iatom, jatom, natoms_encode)
797      END DO
798
799      ! sort atom_pairs so that the pairs are ordered process blocks and
800      ! that possible duplicates may be found (we don't want to send or
801      ! receive same information to the same destination or source more
802      ! than once)
803      IF (npairs > 0) THEN
804         ALLOCATE (tmp_index(npairs))
805         ! only sort the actual pairs recorded in the send list
806         CALL sort(pairs, npairs, tmp_index)
807         DEALLOCATE (tmp_index)
808      END IF
809
810      ! remove duplicates
811      IF (npairs > 1) THEN
812         npairs = 1
813         ! first atom pair must be allowed
814         DO ii = 2, ntasks
815            IF (pairs(ii) > pairs(ii - 1)) THEN
816               npairs = npairs + 1
817               pairs(npairs) = pairs(ii)
818            END IF
819         END DO
820         ! reallocate the pairs list
821         CALL reallocate(pairs, 1, npairs)
822      END IF
823
824      CALL fb_com_atom_pairs_set(atom_pairs=atom_pairs, &
825                                 pairs=pairs, &
826                                 npairs=npairs, &
827                                 natoms_encode=natoms_encode)
828
829      CALL timestop(handle)
830
831   END SUBROUTINE fb_com_tasks_build_atom_pairs
832
833! **********************************************************************
834!> \brief Encodes (iatom, jatom) pair index of a block into a single
835!>        integer
836!> \param ind encoded integer
837!> \param iatom the first index of the (iatom, jatom) block index
838!> \param jatom the second index of the (iatom, jatom) block index
839!> \param natoms the total number of atoms iatom and jatom indexes
840!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
841! **************************************************************************************************
842   SUBROUTINE fb_com_tasks_encode_pair(ind, iatom, jatom, natoms)
843      INTEGER(KIND=int_8), INTENT(OUT)                   :: ind
844      INTEGER, INTENT(IN)                                :: iatom, jatom, natoms
845
846      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_tasks_encode_pair', &
847         routineP = moduleN//':'//routineN
848
849      INTEGER(KIND=int_8)                                :: iatom8, jatom8, natoms8
850
851      natoms8 = INT(natoms, int_8)
852      iatom8 = INT(iatom, int_8)
853      jatom8 = INT(jatom, int_8)
854
855      ind = (iatom8 - 1_int_8)*natoms8 + (jatom8 - 1_int_8)
856   END SUBROUTINE fb_com_tasks_encode_pair
857
858! **********************************************************************
859!> \brief Dncodes a single integer into (iatom, jatom) pair index of
860!>        a block into a single
861!> \param ind encoded integer
862!> \param iatom the first index of the (iatom, jatom) block index
863!> \param jatom the second index of the (iatom, jatom) block index
864!> \param natoms the total number of atoms iatom and jatom indexes
865!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
866! **************************************************************************************************
867   SUBROUTINE fb_com_tasks_decode_pair(ind, iatom, jatom, natoms)
868      INTEGER(KIND=int_8), INTENT(IN)                    :: ind
869      INTEGER, INTENT(OUT)                               :: iatom, jatom
870      INTEGER, INTENT(IN)                                :: natoms
871
872      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_tasks_decode_pair', &
873         routineP = moduleN//':'//routineN
874
875      INTEGER(KIND=int_8)                                :: iatom8, jatom8, natoms8
876
877      natoms8 = INT(natoms, int_8)
878      iatom8 = ind/natoms8 + 1_int_8
879      jatom8 = MOD(ind, natoms8) + 1_int_8
880      iatom = INT(iatom8, int_4)
881      jatom = INT(jatom8, int_4)
882   END SUBROUTINE fb_com_tasks_decode_pair
883
884! **********************************************************************
885!> \brief Encodes (rank, iatom, jatom) index of a communication task---to
886!>         send/receive a block to/from a process---into a single integer
887!> \param ind encoded integer
888!> \param pe the rank of the process the block to be send to or receive
889!>            from
890!> \param iatom the first index of the (iatom, jatom) block index
891!> \param jatom the second index of the (iatom, jatom) block index
892!> \param natoms the total number of atoms iatom and jatom indexes
893!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
894! **************************************************************************************************
895   SUBROUTINE fb_com_atom_pairs_encode(ind, pe, iatom, jatom, natoms)
896      INTEGER(KIND=int_8), INTENT(OUT)                   :: ind
897      INTEGER, INTENT(IN)                                :: pe, iatom, jatom, natoms
898
899      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_atom_pairs_encode', &
900         routineP = moduleN//':'//routineN
901
902      INTEGER(KIND=int_8)                                :: natoms8, pair
903
904! pe must start count from 0 (i.e same as MPI convension)
905
906      natoms8 = INT(natoms, int_8)
907      CALL fb_com_tasks_encode_pair(pair, iatom, jatom, natoms)
908      ind = INT(pe, int_8)*natoms8*natoms8 + pair
909   END SUBROUTINE fb_com_atom_pairs_encode
910
911! **********************************************************************
912!> \brief Decodes a single integer into the (rank, iatom, jatom) index
913!>        of a communication task to send/receive a block to/from a
914!>        process
915!> \param ind    : encoded integer
916!> \param pe     : the rank of the process the block to be send to or receive
917!>            from
918!> \param iatom  : the first index of the (iatom, jatom) block index
919!> \param jatom  : the second index of the (iatom, jatom) block index
920!> \param natoms : the total number of atoms iatom and jatom indexes
921!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
922! **************************************************************************************************
923   SUBROUTINE fb_com_atom_pairs_decode(ind, pe, iatom, jatom, natoms)
924      INTEGER(KIND=int_8), INTENT(IN)                    :: ind
925      INTEGER, INTENT(OUT)                               :: pe, iatom, jatom
926      INTEGER, INTENT(IN)                                :: natoms
927
928      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_atom_pairs_decode', &
929         routineP = moduleN//':'//routineN
930
931      INTEGER(KIND=int_8)                                :: natoms8, pair
932
933! pe start count from 0 (i.e same as MPI convension)
934
935      natoms8 = INT(natoms, int_8)
936      pe = INT(ind/(natoms8*natoms8), int_4)
937      pair = MOD(ind, natoms8*natoms8)
938      CALL fb_com_tasks_decode_pair(pair, iatom, jatom, natoms)
939   END SUBROUTINE fb_com_atom_pairs_decode
940
941! **********************************************************************
942!> \brief Calculate the MPI send or recv buffer sizes according to the
943!>        communication pairs (atom_pairs) and DBCSR matrix data.
944!>        Each atom_pair corresponds to one DBCSR matrix block that
945!>        needs to be sent or recerived.
946!> \param atom_pairs : the communication pair object for either sending
947!>                     or receiving
948!> \param nprocs : total number of MPI processes in communicator
949!> \param row_blk_sizes : row_blk_sizes(iblkrow) = number of element rows
950!>                        in each block in the iblkrow-th block row of
951!>                        the DBCSR matrix
952!> \param col_blk_sizes : col_blk_sizes(iblkcol) = number of element cols
953!>                        in each block in the iblkcol-th block col of
954!>                        the DBCSR matrix
955!> \param sendrecv_sizes : size required for the send of recv buffer
956!>                         for each dest/src process
957!> \param sendrecv_disps : sendrecv_disps(ipe) + 1 = starting location
958!>                         in send/recv buffer for data destined for
959!>                         process ipe
960!> \param sendrecv_pair_counts : sendrecv_pair_counts(ipe) = number of
961!>                               pairs (blocks) to be sent to or recv
962!>                               from process ipe
963!> \param sendrecv_pair_disps send_recv_pair_disps(ipe) + 1 = start
964!>                               location in atom_pairs array for
965!>                               all the pairs to be sent to or recv
966!>                               from process ipe
967!> \param row_map : optional blk row map for the DBCSR blocks
968!> \param col_map : optional blk col map for the DBCSR blocks
969!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
970! **************************************************************************************************
971   SUBROUTINE fb_com_atom_pairs_calc_buffer_sizes(atom_pairs, &
972                                                  nprocs, &
973                                                  row_blk_sizes, &
974                                                  col_blk_sizes, &
975                                                  sendrecv_sizes, &
976                                                  sendrecv_disps, &
977                                                  sendrecv_pair_counts, &
978                                                  sendrecv_pair_disps, &
979                                                  row_map, &
980                                                  col_map)
981      TYPE(fb_com_atom_pairs_obj), INTENT(IN)            :: atom_pairs
982      INTEGER, INTENT(IN)                                :: nprocs
983      INTEGER, DIMENSION(:), INTENT(IN)                  :: row_blk_sizes, col_blk_sizes
984      INTEGER, DIMENSION(:), INTENT(OUT)                 :: sendrecv_sizes, sendrecv_disps, &
985                                                            sendrecv_pair_counts, &
986                                                            sendrecv_pair_disps
987      INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL        :: row_map, col_map
988
989      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_atom_pairs_calc_buffer_sizes', &
990         routineP = moduleN//':'//routineN
991
992      INTEGER                                            :: iatom, ipair, ipe, jatom, natoms_encode, &
993                                                            ncols_blk, npairs, nrows_blk, pe
994      INTEGER(KIND=int_8), DIMENSION(:), POINTER         :: pairs
995      LOGICAL                                            :: check_ok
996
997      NULLIFY (pairs)
998
999      check_ok = SIZE(sendrecv_sizes) == nprocs .AND. &
1000                 SIZE(sendrecv_disps) == nprocs .AND. &
1001                 SIZE(sendrecv_pair_counts) == nprocs .AND. &
1002                 SIZE(sendrecv_pair_disps) == nprocs
1003      CPASSERT(check_ok)
1004
1005      check_ok = fb_com_atom_pairs_has_data(atom_pairs)
1006      CPASSERT(check_ok)
1007
1008      CALL fb_com_atom_pairs_get(atom_pairs=atom_pairs, &
1009                                 pairs=pairs, &
1010                                 npairs=npairs, &
1011                                 natoms_encode=natoms_encode)
1012
1013      sendrecv_sizes = 0
1014      sendrecv_pair_counts = 0
1015      DO ipair = 1, npairs
1016         ! decode processor and (iatom, jatom) information
1017         CALL fb_com_atom_pairs_decode(pairs(ipair), &
1018                                       pe, iatom, jatom, natoms_encode)
1019         pe = pe + 1 ! we need proc to count from 1
1020         IF (PRESENT(row_map)) iatom = row_map(iatom)
1021         IF (PRESENT(col_map)) jatom = row_map(jatom)
1022         nrows_blk = row_blk_sizes(iatom)
1023         ncols_blk = col_blk_sizes(jatom)
1024         sendrecv_sizes(pe) = sendrecv_sizes(pe) + nrows_blk*ncols_blk
1025         sendrecv_pair_counts(pe) = sendrecv_pair_counts(pe) + 1
1026      END DO
1027      ! calculate displacements of the data of each destibation pe in
1028      ! send buffer and in the list of pairs to be sent
1029      sendrecv_disps = 0
1030      sendrecv_pair_disps = 0
1031      DO ipe = 2, nprocs
1032         sendrecv_disps(ipe) = sendrecv_disps(ipe - 1) + sendrecv_sizes(ipe - 1)
1033         sendrecv_pair_disps(ipe) = sendrecv_pair_disps(ipe - 1) + sendrecv_pair_counts(ipe - 1)
1034      END DO
1035
1036   END SUBROUTINE fb_com_atom_pairs_calc_buffer_sizes
1037
1038! ****************************************************************************
1039!> \brief Given send and recv fb_com_atom_pair object, gather all the
1040!>        relevant DBCSR matrix blocks together, and add them to
1041!>        a fb_matrix_data object for storage
1042!> \param dbcsr_mat : the DBCSR matrix where the matrix blocks will be
1043!>                    obtained from
1044!> \param atom_pairs_send : prescription on exactly which DBCSR blocks
1045!>                          are to be sent to where
1046!> \param atom_pairs_recv : prescription on exactly which DBCSR blocks
1047!>                          are to be received from where
1048!> \param para_env        : CP2K parallel environment
1049!> \param matrix_storage  : the fb_matrix_data object to store the
1050!>                          received DBCSR matrix blocks
1051!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
1052! **************************************************************************************************
1053   SUBROUTINE fb_com_atom_pairs_gather_blks(dbcsr_mat, &
1054                                            atom_pairs_send, &
1055                                            atom_pairs_recv, &
1056                                            para_env, &
1057                                            matrix_storage)
1058      TYPE(dbcsr_type), POINTER                          :: dbcsr_mat
1059      TYPE(fb_com_atom_pairs_obj), INTENT(IN)            :: atom_pairs_send, atom_pairs_recv
1060      TYPE(cp_para_env_type), POINTER                    :: para_env
1061      TYPE(fb_matrix_data_obj), INTENT(INOUT)            :: matrix_storage
1062
1063      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_atom_pairs_gather_blks', &
1064         routineP = moduleN//':'//routineN
1065
1066      INTEGER :: handle, iatom, ii, ind, ipair, ipe, jatom, jj, ncols_blk, ncols_blk_max, &
1067         npairs_recv, npairs_send, nrows_blk, nrows_blk_max, numprocs, pe, recv_encode, send_encode
1068      INTEGER(KIND=int_8), DIMENSION(:), POINTER         :: pairs_recv, pairs_send
1069      INTEGER, ALLOCATABLE, DIMENSION(:) :: recv_disps, recv_pair_count, recv_pair_disps, &
1070         recv_sizes, send_disps, send_pair_count, send_pair_disps, send_sizes
1071      INTEGER, DIMENSION(:), POINTER                     :: col_block_size_data, row_block_size_data
1072      LOGICAL                                            :: check_ok, found
1073      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: recv_buf, send_buf
1074      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: mat_block
1075
1076      CALL timeset(routineN, handle)
1077
1078      NULLIFY (pairs_send, pairs_recv, mat_block, &
1079               row_block_size_data, col_block_size_data)
1080
1081      check_ok = fb_com_atom_pairs_has_data(atom_pairs_send)
1082      CPASSERT(check_ok)
1083      check_ok = fb_com_atom_pairs_has_data(atom_pairs_send)
1084      CPASSERT(check_ok)
1085      check_ok = fb_matrix_data_has_data(matrix_storage)
1086      CPASSERT(check_ok)
1087
1088      ! get com pair informations
1089      CALL fb_com_atom_pairs_get(atom_pairs=atom_pairs_send, &
1090                                 pairs=pairs_send, &
1091                                 npairs=npairs_send, &
1092                                 natoms_encode=send_encode)
1093      CALL fb_com_atom_pairs_get(atom_pairs=atom_pairs_recv, &
1094                                 pairs=pairs_recv, &
1095                                 npairs=npairs_recv, &
1096                                 natoms_encode=recv_encode)
1097      ! get para_env info
1098      numprocs = para_env%num_pe
1099
1100      ! get dbcsr row and col block sizes
1101      CALL dbcsr_get_info(dbcsr_mat, row_blk_size=row_block_size_data, col_blk_size=col_block_size_data)
1102
1103      ! allocate temporary arrays for send
1104      ALLOCATE (send_sizes(numprocs))
1105      ALLOCATE (send_disps(numprocs))
1106      ALLOCATE (send_pair_count(numprocs))
1107      ALLOCATE (send_pair_disps(numprocs))
1108
1109      ! setup send buffer sizes
1110      CALL fb_com_atom_pairs_calc_buffer_sizes(atom_pairs_send, &
1111                                               numprocs, &
1112                                               row_block_size_data, &
1113                                               col_block_size_data, &
1114                                               send_sizes, &
1115                                               send_disps, &
1116                                               send_pair_count, &
1117                                               send_pair_disps)
1118
1119      ! allocate send buffer
1120      ALLOCATE (send_buf(SUM(send_sizes)))
1121
1122      ! allocate temporary arrays for recv
1123      ALLOCATE (recv_sizes(numprocs))
1124      ALLOCATE (recv_disps(numprocs))
1125      ALLOCATE (recv_pair_count(numprocs))
1126      ALLOCATE (recv_pair_disps(numprocs))
1127
1128      ! setup recv buffer sizes
1129      CALL fb_com_atom_pairs_calc_buffer_sizes(atom_pairs_recv, &
1130                                               numprocs, &
1131                                               row_block_size_data, &
1132                                               col_block_size_data, &
1133                                               recv_sizes, &
1134                                               recv_disps, &
1135                                               recv_pair_count, &
1136                                               recv_pair_disps)
1137
1138      ! allocate recv buffer
1139      ALLOCATE (recv_buf(SUM(recv_sizes)))
1140
1141      ! do packing
1142      DO ipe = 1, numprocs
1143         ! need to reuse send_sizes as an accumulative displacement, so recalculate
1144         send_sizes(ipe) = 0
1145         DO ipair = 1, send_pair_count(ipe)
1146            CALL fb_com_atom_pairs_decode(pairs_send(send_pair_disps(ipe) + ipair), &
1147                                          pe, iatom, jatom, send_encode)
1148            nrows_blk = row_block_size_data(iatom)
1149            ncols_blk = col_block_size_data(jatom)
1150            CALL dbcsr_get_block_p(matrix=dbcsr_mat, &
1151                                   row=iatom, col=jatom, block=mat_block, &
1152                                   found=found)
1153            IF (.NOT. found) THEN
1154               CPABORT("Matrix block not found")
1155            ELSE
1156               ! we have found the matrix block
1157               DO jj = 1, ncols_blk
1158                  DO ii = 1, nrows_blk
1159                     ! column major format in blocks
1160                     ind = send_disps(ipe) + send_sizes(ipe) + ii + (jj - 1)*nrows_blk
1161                     send_buf(ind) = mat_block(ii, jj)
1162                  END DO ! ii
1163               END DO ! jj
1164               send_sizes(ipe) = send_sizes(ipe) + nrows_blk*ncols_blk
1165            END IF
1166         END DO ! ipair
1167      END DO ! ipe
1168
1169      ! do communication
1170      CALL mp_alltoall(send_buf, send_sizes, send_disps, &
1171                       recv_buf, recv_sizes, recv_disps, &
1172                       para_env%group)
1173
1174      ! cleanup temporary arrays no longer needed
1175      DEALLOCATE (send_buf)
1176      DEALLOCATE (send_sizes)
1177      DEALLOCATE (send_disps)
1178      DEALLOCATE (send_pair_count)
1179      DEALLOCATE (send_pair_disps)
1180
1181      ! unpack into matrix_data object
1182      NULLIFY (mat_block)
1183      nrows_blk_max = MAXVAL(row_block_size_data)
1184      ncols_blk_max = MAXVAL(col_block_size_data)
1185      ALLOCATE (mat_block(nrows_blk_max, ncols_blk_max))
1186      DO ipe = 1, numprocs
1187         recv_sizes(ipe) = 0
1188         DO ipair = 1, recv_pair_count(ipe)
1189            CALL fb_com_atom_pairs_decode(pairs_recv(recv_pair_disps(ipe) + ipair), &
1190                                          pe, iatom, jatom, recv_encode)
1191            nrows_blk = row_block_size_data(iatom)
1192            ncols_blk = col_block_size_data(jatom)
1193            ! ALLOCATE(mat_block(nrows_blk,ncols_blk), STAT=stat)
1194            ! CPPostcondition(stat==0, cp_failure_level, routineP,failure)
1195            mat_block(:, :) = 0.0_dp
1196            DO jj = 1, ncols_blk
1197               DO ii = 1, nrows_blk
1198                  ! column major format in blocks
1199                  ind = recv_disps(ipe) + recv_sizes(ipe) + ii + (jj - 1)*nrows_blk
1200                  mat_block(ii, jj) = recv_buf(ind)
1201               END DO ! ii
1202            END DO ! jj
1203            CALL fb_matrix_data_add(matrix_storage, &
1204                                    iatom, jatom, &
1205                                    mat_block(1:nrows_blk, 1:ncols_blk))
1206            recv_sizes(ipe) = recv_sizes(ipe) + nrows_blk*ncols_blk
1207            ! DEALLOCATE(mat_block, STAT=stat)
1208            ! CPPostcondition(stat==0, cp_failure_level, routineP,failure)
1209         END DO ! ipair
1210      END DO ! ipe
1211      DEALLOCATE (mat_block)
1212
1213      ! cleanup rest of the temporary arrays
1214      DEALLOCATE (recv_buf)
1215      DEALLOCATE (recv_sizes)
1216      DEALLOCATE (recv_disps)
1217      DEALLOCATE (recv_pair_count)
1218      DEALLOCATE (recv_pair_disps)
1219
1220      CALL timestop(handle)
1221
1222   END SUBROUTINE fb_com_atom_pairs_gather_blks
1223
1224! ****************************************************************************
1225!> \brief Given send and recv fb_com_atom_pair object, distribute the matrix
1226!>        blocks stored in a fb_matrix_data object to a compatable DBCSR
1227!>        matrix. It is assumed in this subroutine that the sizes of each
1228!>        block stored in fb_matrix_data object is consistent with the
1229!>        pre-defined block sizes in the DBCSR matrix.
1230!> \param matrix_storage  : the fb_matrix_data object
1231!> \param atom_pairs_send : prescription on exactly which DBCSR blocks
1232!>                          are to be sent to where
1233!> \param atom_pairs_recv : prescription on exactly which DBCSR blocks
1234!>                          are to be received from where
1235!> \param para_env        : CP2K parallel environment
1236!> \param dbcsr_mat : the DBCSR matrix where the matrix blocks will be
1237!>                    distributed to
1238!> \author Lianheng Tong (LT) lianheng.tong@kcl.ac.uk
1239! **************************************************************************************************
1240   SUBROUTINE fb_com_atom_pairs_distribute_blks(matrix_storage, &
1241                                                atom_pairs_send, &
1242                                                atom_pairs_recv, &
1243                                                para_env, &
1244                                                dbcsr_mat)
1245      TYPE(fb_matrix_data_obj), INTENT(IN)               :: matrix_storage
1246      TYPE(fb_com_atom_pairs_obj), INTENT(IN)            :: atom_pairs_send, atom_pairs_recv
1247      TYPE(cp_para_env_type), POINTER                    :: para_env
1248      TYPE(dbcsr_type), POINTER                          :: dbcsr_mat
1249
1250      CHARACTER(LEN=*), PARAMETER :: routineN = 'fb_com_atom_pairs_distribute_blks', &
1251         routineP = moduleN//':'//routineN
1252
1253      INTEGER :: handle, iatom, ii, ind, ipair, ipe, jatom, jj, ncols_blk, npairs_recv, &
1254         npairs_send, nrows_blk, numprocs, pe, recv_encode, send_encode
1255      INTEGER(KIND=int_8), DIMENSION(:), POINTER         :: pairs_recv, pairs_send
1256      INTEGER, ALLOCATABLE, DIMENSION(:) :: recv_disps, recv_pair_count, recv_pair_disps, &
1257         recv_sizes, send_disps, send_pair_count, send_pair_disps, send_sizes
1258      INTEGER, DIMENSION(:), POINTER                     :: col_block_size_data, row_block_size_data
1259      LOGICAL                                            :: check_ok, found
1260      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: recv_buf, send_buf
1261      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: mat_block
1262
1263      CALL timeset(routineN, handle)
1264
1265      NULLIFY (pairs_send, pairs_recv, mat_block, &
1266               row_block_size_data, col_block_size_data)
1267
1268      check_ok = fb_com_atom_pairs_has_data(atom_pairs_send)
1269      CPASSERT(check_ok)
1270      check_ok = fb_com_atom_pairs_has_data(atom_pairs_send)
1271      CPASSERT(check_ok)
1272      check_ok = fb_matrix_data_has_data(matrix_storage)
1273      CPASSERT(check_ok)
1274
1275      ! get com pair informations
1276      CALL fb_com_atom_pairs_get(atom_pairs=atom_pairs_send, &
1277                                 pairs=pairs_send, &
1278                                 npairs=npairs_send, &
1279                                 natoms_encode=send_encode)
1280      CALL fb_com_atom_pairs_get(atom_pairs=atom_pairs_recv, &
1281                                 pairs=pairs_recv, &
1282                                 npairs=npairs_recv, &
1283                                 natoms_encode=recv_encode)
1284      ! get para_env info
1285      numprocs = para_env%num_pe
1286
1287      ! get dbcsr row and col block sizes
1288      CALL dbcsr_get_info(dbcsr_mat, row_blk_size=row_block_size_data, col_blk_size=col_block_size_data)
1289
1290      ! allocate temporary arrays for send
1291      ALLOCATE (send_sizes(numprocs))
1292      ALLOCATE (send_disps(numprocs))
1293      ALLOCATE (send_pair_count(numprocs))
1294      ALLOCATE (send_pair_disps(numprocs))
1295
1296      ! setup send buffer sizes
1297      CALL fb_com_atom_pairs_calc_buffer_sizes(atom_pairs_send, &
1298                                               numprocs, &
1299                                               row_block_size_data, &
1300                                               col_block_size_data, &
1301                                               send_sizes, &
1302                                               send_disps, &
1303                                               send_pair_count, &
1304                                               send_pair_disps)
1305
1306      ! allocate send buffer
1307      ALLOCATE (send_buf(SUM(send_sizes)))
1308
1309      ! allocate temporary arrays for recv
1310      ALLOCATE (recv_sizes(numprocs))
1311      ALLOCATE (recv_disps(numprocs))
1312      ALLOCATE (recv_pair_count(numprocs))
1313      ALLOCATE (recv_pair_disps(numprocs))
1314
1315      ! setup recv buffer sizes
1316      CALL fb_com_atom_pairs_calc_buffer_sizes(atom_pairs_recv, &
1317                                               numprocs, &
1318                                               row_block_size_data, &
1319                                               col_block_size_data, &
1320                                               recv_sizes, &
1321                                               recv_disps, &
1322                                               recv_pair_count, &
1323                                               recv_pair_disps)
1324
1325      ! allocate recv buffer
1326      ALLOCATE (recv_buf(SUM(recv_sizes)))
1327
1328      ! do packing
1329      DO ipe = 1, numprocs
1330         ! need to reuse send_sizes as an accumulative displacement, so recalculate
1331         send_sizes(ipe) = 0
1332         DO ipair = 1, send_pair_count(ipe)
1333            CALL fb_com_atom_pairs_decode(pairs_send(send_pair_disps(ipe) + ipair), &
1334                                          pe, iatom, jatom, send_encode)
1335            CALL fb_matrix_data_get(matrix_storage, &
1336                                    iatom, jatom, &
1337                                    mat_block, found)
1338            IF (.NOT. found) THEN
1339               CPABORT("Matrix block not found")
1340            ELSE
1341               nrows_blk = row_block_size_data(iatom)
1342               ncols_blk = col_block_size_data(jatom)
1343               DO jj = 1, ncols_blk
1344                  DO ii = 1, nrows_blk
1345                     ! column major format in blocks
1346                     ind = send_disps(ipe) + send_sizes(ipe) + ii + (jj - 1)*nrows_blk
1347                     send_buf(ind) = mat_block(ii, jj)
1348                  END DO ! ii
1349               END DO ! jj
1350               send_sizes(ipe) = send_sizes(ipe) + nrows_blk*ncols_blk
1351            END IF
1352         END DO ! ipair
1353      END DO ! ipe
1354
1355      ! do communication
1356      CALL mp_alltoall(send_buf, send_sizes, send_disps, &
1357                       recv_buf, recv_sizes, recv_disps, &
1358                       para_env%group)
1359
1360      ! cleanup temporary arrays no longer needed
1361      DEALLOCATE (send_buf)
1362      DEALLOCATE (send_sizes)
1363      DEALLOCATE (send_disps)
1364      DEALLOCATE (send_pair_count)
1365      DEALLOCATE (send_pair_disps)
1366
1367      ! unpack into DBCSR matrix
1368      DO ipe = 1, numprocs
1369         recv_sizes(ipe) = 0
1370         DO ipair = 1, recv_pair_count(ipe)
1371            CALL fb_com_atom_pairs_decode(pairs_recv(recv_pair_disps(ipe) + ipair), &
1372                                          pe, iatom, jatom, recv_encode)
1373            nrows_blk = row_block_size_data(iatom)
1374            ncols_blk = col_block_size_data(jatom)
1375            ind = recv_disps(ipe) + recv_sizes(ipe)
1376            CALL dbcsr_put_block(dbcsr_mat, &
1377                                 iatom, jatom, &
1378                                 recv_buf((ind + 1):(ind + nrows_blk*ncols_blk)))
1379            recv_sizes(ipe) = recv_sizes(ipe) + nrows_blk*ncols_blk
1380         END DO ! ipair
1381      END DO ! ipe
1382
1383      ! cleanup rest of the temporary arrays
1384      DEALLOCATE (recv_buf)
1385      DEALLOCATE (recv_sizes)
1386      DEALLOCATE (recv_disps)
1387      DEALLOCATE (recv_pair_count)
1388      DEALLOCATE (recv_pair_disps)
1389
1390      ! dbcsr matrix is not finalised in this subroutine
1391
1392      CALL timestop(handle)
1393
1394   END SUBROUTINE fb_com_atom_pairs_distribute_blks
1395
1396END MODULE qs_fb_com_tasks_types
1397