(git:e7e05ae)
dbt_types.F
Go to the documentation of this file.
1 !--------------------------------------------------------------------------------------------------!
2 ! CP2K: A general program to perform molecular dynamics simulations !
3 ! Copyright 2000-2024 CP2K developers group <https://cp2k.org> !
4 ! !
5 ! SPDX-License-Identifier: GPL-2.0-or-later !
6 !--------------------------------------------------------------------------------------------------!
7 
8 ! **************************************************************************************************
9 !> \brief DBT tensor framework for block-sparse tensor contraction: Types and create/destroy routines.
10 !> \author Patrick Seewald
11 ! **************************************************************************************************
12 MODULE dbt_types
13 
14 
15  USE dbcsr_api, ONLY: dbcsr_type, dbcsr_get_info, dbcsr_distribution_type, dbcsr_distribution_get
16  USE dbt_array_list_methods, ONLY: &
17  array_list, array_offsets, create_array_list, destroy_array_list, get_array_elements, &
18  sizes_of_arrays, sum_of_arrays, array_sublist, get_arrays, get_ith_array, array_eq_i
19  USE dbm_api, ONLY: &
20  dbm_distribution_obj, dbm_type
21  USE kinds, ONLY: dp, dp, default_string_length
22  USE dbt_tas_base, ONLY: &
23  dbt_tas_create, dbt_tas_distribution_new, &
24  dbt_tas_distribution_destroy, dbt_tas_finalize, dbt_tas_get_info, &
25  dbt_tas_destroy, dbt_tas_get_stored_coordinates, dbt_tas_filter, &
26  dbt_tas_get_num_blocks, dbt_tas_get_num_blocks_total, dbt_tas_get_nze, &
27  dbt_tas_get_nze_total, dbt_tas_clear
28  USE dbt_tas_types, ONLY: &
29  dbt_tas_type, dbt_tas_distribution_type, dbt_tas_split_info, dbt_tas_mm_storage
30  USE dbt_tas_mm, ONLY: dbt_tas_set_batched_state
31  USE dbt_index, ONLY: &
32  get_2d_indices_tensor, get_nd_indices_pgrid, create_nd_to_2d_mapping, destroy_nd_to_2d_mapping, &
33  dbt_get_mapping_info, nd_to_2d_mapping, split_tensor_index, combine_tensor_index, combine_pgrid_index, &
34  split_pgrid_index, ndims_mapping, ndims_mapping_row, ndims_mapping_column
35  USE dbt_tas_split, ONLY: &
36  dbt_tas_create_split_rows_or_cols, dbt_tas_release_info, dbt_tas_info_hold, &
37  dbt_tas_create_split, dbt_tas_get_split_info, dbt_tas_set_strict_split
38  USE kinds, ONLY: default_string_length, int_8, dp
39  USE message_passing, ONLY: &
40  mp_cart_type, mp_dims_create, mp_comm_type
41  USE dbt_tas_global, ONLY: dbt_tas_distribution, dbt_tas_rowcol_data, dbt_tas_default_distvec
42  USE dbt_allocate_wrap, ONLY: allocate_any
43  USE dbm_api, ONLY: dbm_scale
44  USE util, ONLY: sort
45 #include "../base/base_uses.f90"
46 
47  IMPLICIT NONE
48  PRIVATE
49  CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_types'
50 
51  PUBLIC :: &
52  blk_dims_tensor, &
53  dbt_blk_offsets, &
54  dbt_blk_sizes, &
55  dbt_clear, &
56  dbt_create, &
57  dbt_destroy, &
58  dbt_distribution, &
59  dbt_distribution_destroy, &
60  dbt_distribution_new, &
61  dbt_distribution_new_expert, &
62  dbt_distribution_type, &
63  dbt_filter, &
64  dbt_finalize, &
65  dbt_get_info, &
66  dbt_get_num_blocks, &
67  dbt_get_num_blocks_total, &
68  dbt_get_nze, &
69  dbt_get_nze_total, &
70  dbt_get_stored_coordinates, &
71  dbt_hold, &
72  dbt_mp_dims_create, &
73  dbt_nd_mp_comm, &
74  dbt_nd_mp_free, &
75  dbt_pgrid_change_dims, &
76  dbt_pgrid_create, &
77  dbt_pgrid_create_expert, &
78  dbt_pgrid_destroy, &
79  dbt_pgrid_type, &
80  dbt_pgrid_set_strict_split, &
81  dbt_scale, &
82  dbt_type, &
83  dims_tensor, &
84  mp_environ_pgrid, &
85  ndims_tensor, &
86  ndims_matrix_row, &
87  ndims_matrix_column, &
88  dbt_nblks_local, &
89  dbt_nblks_total, &
90  dbt_blk_size, &
91  dbt_max_nblks_local, &
92  dbt_default_distvec, &
93  dbt_contraction_storage, &
94  dbt_copy_contraction_storage
95 
96  TYPE dbt_pgrid_type
97  TYPE(nd_to_2d_mapping) :: nd_index_grid
98  TYPE(mp_cart_type) :: mp_comm_2d
99  TYPE(dbt_tas_split_info), ALLOCATABLE :: tas_split_info
100  INTEGER :: nproc = -1
101  END TYPE
102 
103  TYPE dbt_contraction_storage
104  REAL(dp) :: nsplit_avg = 0.0_dp
105  INTEGER :: ibatch = -1
106  TYPE(array_list) :: batch_ranges
107  LOGICAL :: static = .false.
108  END TYPE
109 
110  TYPE dbt_type
111  TYPE(dbt_tas_type), POINTER :: matrix_rep => null()
112  TYPE(nd_to_2d_mapping) :: nd_index_blk
113  TYPE(nd_to_2d_mapping) :: nd_index
114  TYPE(array_list) :: blk_sizes
115  TYPE(array_list) :: blk_offsets
116  TYPE(array_list) :: nd_dist
117  TYPE(dbt_pgrid_type) :: pgrid
118  TYPE(array_list) :: blks_local
119  INTEGER, DIMENSION(:), ALLOCATABLE :: nblks_local
120  INTEGER, DIMENSION(:), ALLOCATABLE :: nfull_local
121  LOGICAL :: valid = .false.
122  LOGICAL :: owns_matrix = .false.
123  CHARACTER(LEN=default_string_length) :: name = ""
124  ! lightweight reference counting for communicators:
125  INTEGER, POINTER :: refcount => null()
126  TYPE(dbt_contraction_storage), ALLOCATABLE :: contraction_storage
127  END TYPE dbt_type
128 
129  TYPE dbt_distribution_type
130  TYPE(dbt_tas_distribution_type) :: dist
131  TYPE(dbt_pgrid_type) :: pgrid
132  TYPE(array_list) :: nd_dist
133  ! lightweight reference counting for communicators:
134  INTEGER, POINTER :: refcount => null()
135  END TYPE
136 
137 ! **************************************************************************************************
138 !> \brief tas matrix distribution function object for one matrix index
139 !> \var dims tensor dimensions only for this matrix dimension
140 !> \var dims_grid grid dimensions only for this matrix dimension
141 !> \var nd_dist dist only for tensor dimensions belonging to this matrix dimension
142 !> \var tas_dist_t map matrix index to process grid
143 !> \var tas_rowcols_t map process grid to matrix index
144 ! **************************************************************************************************
145  TYPE, EXTENDS(dbt_tas_distribution) :: dbt_tas_dist_t
146  INTEGER, DIMENSION(:), ALLOCATABLE :: dims
147  INTEGER, DIMENSION(:), ALLOCATABLE :: dims_grid
148  TYPE(array_list) :: nd_dist
149  CONTAINS
150  PROCEDURE :: dist => tas_dist_t
151  PROCEDURE :: rowcols => tas_rowcols_t
152  END TYPE
153 
154 ! **************************************************************************************************
155 !> \brief block size object for one matrix index
156 !> \var dims tensor dimensions only for this matrix dimension
157 !> \var blk_size block size only for this matrix dimension
158 ! **************************************************************************************************
159  TYPE, EXTENDS(dbt_tas_rowcol_data) :: dbt_tas_blk_size_t
160  INTEGER, DIMENSION(:), ALLOCATABLE :: dims
161  TYPE(array_list) :: blk_size
162  CONTAINS
163  PROCEDURE :: data => tas_blk_size_t
164  END TYPE
165 
166  INTERFACE dbt_create
167  MODULE PROCEDURE dbt_create_new
168  MODULE PROCEDURE dbt_create_template
169  MODULE PROCEDURE dbt_create_matrix
170  END INTERFACE
171 
172  INTERFACE dbt_tas_dist_t
173  MODULE PROCEDURE new_dbt_tas_dist_t
174  END INTERFACE
175 
176  INTERFACE dbt_tas_blk_size_t
177  MODULE PROCEDURE new_dbt_tas_blk_size_t
178  END INTERFACE
179 
180 CONTAINS
181 
182 ! **************************************************************************************************
183 !> \brief Create distribution object for one matrix dimension
184 !> \param nd_dist arrays for distribution vectors along all dimensions
185 !> \param map_blks tensor to matrix mapping object for blocks
186 !> \param map_grid tensor to matrix mapping object for process grid
187 !> \param which_dim for which dimension (1 or 2) distribution should be created
188 !> \return distribution object
189 !> \author Patrick Seewald
190 ! **************************************************************************************************
191  FUNCTION new_dbt_tas_dist_t(nd_dist, map_blks, map_grid, which_dim)
192  TYPE(array_list), INTENT(IN) :: nd_dist
193  TYPE(nd_to_2d_mapping), INTENT(IN) :: map_blks, map_grid
194  INTEGER, INTENT(IN) :: which_dim
195 
196  TYPE(dbt_tas_dist_t) :: new_dbt_tas_dist_t
197  INTEGER, DIMENSION(2) :: grid_dims
198  INTEGER(KIND=int_8), DIMENSION(2) :: matrix_dims
199  INTEGER, DIMENSION(:), ALLOCATABLE :: index_map
200 
201  IF (which_dim == 1) THEN
202  ALLOCATE (new_dbt_tas_dist_t%dims(ndims_mapping_row(map_blks)))
203  ALLOCATE (index_map(ndims_mapping_row(map_blks)))
204  CALL dbt_get_mapping_info(map_blks, &
205  dims_2d_i8=matrix_dims, &
206  map1_2d=index_map, &
207  dims1_2d=new_dbt_tas_dist_t%dims)
208  ALLOCATE (new_dbt_tas_dist_t%dims_grid(ndims_mapping_row(map_grid)))
209  CALL dbt_get_mapping_info(map_grid, &
210  dims_2d=grid_dims, &
211  dims1_2d=new_dbt_tas_dist_t%dims_grid)
212  ELSEIF (which_dim == 2) THEN
213  ALLOCATE (new_dbt_tas_dist_t%dims(ndims_mapping_column(map_blks)))
214  ALLOCATE (index_map(ndims_mapping_column(map_blks)))
215  CALL dbt_get_mapping_info(map_blks, &
216  dims_2d_i8=matrix_dims, &
217  map2_2d=index_map, &
218  dims2_2d=new_dbt_tas_dist_t%dims)
219  ALLOCATE (new_dbt_tas_dist_t%dims_grid(ndims_mapping_column(map_grid)))
220  CALL dbt_get_mapping_info(map_grid, &
221  dims_2d=grid_dims, &
222  dims2_2d=new_dbt_tas_dist_t%dims_grid)
223  ELSE
224  cpabort("Unknown value for which_dim")
225  END IF
226 
227  new_dbt_tas_dist_t%nd_dist = array_sublist(nd_dist, index_map)
228  new_dbt_tas_dist_t%nprowcol = grid_dims(which_dim)
229  new_dbt_tas_dist_t%nmrowcol = matrix_dims(which_dim)
230  END FUNCTION
231 
232 ! **************************************************************************************************
233 !> \author Patrick Seewald
234 ! **************************************************************************************************
235  FUNCTION tas_dist_t(t, rowcol)
236  CLASS(dbt_tas_dist_t), INTENT(IN) :: t
237  INTEGER(KIND=int_8), INTENT(IN) :: rowcol
238  INTEGER, DIMENSION(4) :: ind_blk
239  INTEGER, DIMENSION(4) :: dist_blk
240  INTEGER :: tas_dist_t
241 
242  ind_blk(:SIZE(t%dims)) = split_tensor_index(rowcol, t%dims)
243  dist_blk(:SIZE(t%dims)) = get_array_elements(t%nd_dist, ind_blk(:SIZE(t%dims)))
244  tas_dist_t = combine_pgrid_index(dist_blk(:SIZE(t%dims)), t%dims_grid)
245  END FUNCTION
246 
247 ! **************************************************************************************************
248 !> \author Patrick Seewald
249 ! **************************************************************************************************
250  FUNCTION tas_rowcols_t(t, dist)
251  CLASS(dbt_tas_dist_t), INTENT(IN) :: t
252  INTEGER, INTENT(IN) :: dist
253  INTEGER(KIND=int_8), DIMENSION(:), ALLOCATABLE :: tas_rowcols_t
254  INTEGER, DIMENSION(4) :: dist_blk
255  INTEGER, DIMENSION(:), ALLOCATABLE :: dist_1, dist_2, dist_3, dist_4, blks_1, blks_2, blks_3, blks_4, blks_tmp, nd_ind
256  INTEGER :: i_1, i_2, i_3, i_4, i, iblk, iblk_count, nblks
257  INTEGER(KIND=int_8) :: nrowcols
258  TYPE(array_list) :: blks
259 
260  dist_blk(:SIZE(t%dims)) = split_pgrid_index(dist, t%dims_grid)
261 
262  IF (SIZE(t%dims) == 1) THEN
263  CALL get_arrays(t%nd_dist, dist_1)
264  END IF
265  IF (SIZE(t%dims) == 2) THEN
266  CALL get_arrays(t%nd_dist, dist_1, dist_2)
267  END IF
268  IF (SIZE(t%dims) == 3) THEN
269  CALL get_arrays(t%nd_dist, dist_1, dist_2, dist_3)
270  END IF
271  IF (SIZE(t%dims) == 4) THEN
272  CALL get_arrays(t%nd_dist, dist_1, dist_2, dist_3, dist_4)
273  END IF
274 
275  IF (SIZE(t%dims) .GE. 1) THEN
276  nblks = SIZE(dist_1)
277  ALLOCATE (blks_tmp(nblks))
278  iblk_count = 0
279  DO iblk = 1, nblks
280  IF (dist_1(iblk) == dist_blk(1)) THEN
281  iblk_count = iblk_count + 1
282  blks_tmp(iblk_count) = iblk
283  END IF
284  END DO
285  ALLOCATE (blks_1(iblk_count))
286  blks_1(:) = blks_tmp(:iblk_count)
287  DEALLOCATE (blks_tmp)
288  END IF
289  IF (SIZE(t%dims) .GE. 2) THEN
290  nblks = SIZE(dist_2)
291  ALLOCATE (blks_tmp(nblks))
292  iblk_count = 0
293  DO iblk = 1, nblks
294  IF (dist_2(iblk) == dist_blk(2)) THEN
295  iblk_count = iblk_count + 1
296  blks_tmp(iblk_count) = iblk
297  END IF
298  END DO
299  ALLOCATE (blks_2(iblk_count))
300  blks_2(:) = blks_tmp(:iblk_count)
301  DEALLOCATE (blks_tmp)
302  END IF
303  IF (SIZE(t%dims) .GE. 3) THEN
304  nblks = SIZE(dist_3)
305  ALLOCATE (blks_tmp(nblks))
306  iblk_count = 0
307  DO iblk = 1, nblks
308  IF (dist_3(iblk) == dist_blk(3)) THEN
309  iblk_count = iblk_count + 1
310  blks_tmp(iblk_count) = iblk
311  END IF
312  END DO
313  ALLOCATE (blks_3(iblk_count))
314  blks_3(:) = blks_tmp(:iblk_count)
315  DEALLOCATE (blks_tmp)
316  END IF
317  IF (SIZE(t%dims) .GE. 4) THEN
318  nblks = SIZE(dist_4)
319  ALLOCATE (blks_tmp(nblks))
320  iblk_count = 0
321  DO iblk = 1, nblks
322  IF (dist_4(iblk) == dist_blk(4)) THEN
323  iblk_count = iblk_count + 1
324  blks_tmp(iblk_count) = iblk
325  END IF
326  END DO
327  ALLOCATE (blks_4(iblk_count))
328  blks_4(:) = blks_tmp(:iblk_count)
329  DEALLOCATE (blks_tmp)
330  END IF
331 
332  IF (SIZE(t%dims) == 1) THEN
333  CALL create_array_list(blks, 1, blks_1)
334  END IF
335  IF (SIZE(t%dims) == 2) THEN
336  CALL create_array_list(blks, 2, blks_1, blks_2)
337  END IF
338  IF (SIZE(t%dims) == 3) THEN
339  CALL create_array_list(blks, 3, blks_1, blks_2, blks_3)
340  END IF
341  IF (SIZE(t%dims) == 4) THEN
342  CALL create_array_list(blks, 4, blks_1, blks_2, blks_3, blks_4)
343  END IF
344 
345  nrowcols = product(int(sizes_of_arrays(blks), int_8))
346  ALLOCATE (tas_rowcols_t(nrowcols))
347 
348  IF (SIZE(t%dims) == 1) THEN
349  ALLOCATE (nd_ind(1))
350  i = 0
351  DO i_1 = 1, SIZE(blks_1)
352  i = i + 1
353 
354  nd_ind(:) = get_array_elements(blks, [i_1])
355  tas_rowcols_t(i) = combine_tensor_index(nd_ind, t%dims)
356  END DO
357  END IF
358  IF (SIZE(t%dims) == 2) THEN
359  ALLOCATE (nd_ind(2))
360  i = 0
361  DO i_1 = 1, SIZE(blks_1)
362  DO i_2 = 1, SIZE(blks_2)
363  i = i + 1
364 
365  nd_ind(:) = get_array_elements(blks, [i_1, i_2])
366  tas_rowcols_t(i) = combine_tensor_index(nd_ind, t%dims)
367  END DO
368  END DO
369  END IF
370  IF (SIZE(t%dims) == 3) THEN
371  ALLOCATE (nd_ind(3))
372  i = 0
373  DO i_1 = 1, SIZE(blks_1)
374  DO i_2 = 1, SIZE(blks_2)
375  DO i_3 = 1, SIZE(blks_3)
376  i = i + 1
377 
378  nd_ind(:) = get_array_elements(blks, [i_1, i_2, i_3])
379  tas_rowcols_t(i) = combine_tensor_index(nd_ind, t%dims)
380  END DO
381  END DO
382  END DO
383  END IF
384  IF (SIZE(t%dims) == 4) THEN
385  ALLOCATE (nd_ind(4))
386  i = 0
387  DO i_1 = 1, SIZE(blks_1)
388  DO i_2 = 1, SIZE(blks_2)
389  DO i_3 = 1, SIZE(blks_3)
390  DO i_4 = 1, SIZE(blks_4)
391  i = i + 1
392 
393  nd_ind(:) = get_array_elements(blks, [i_1, i_2, i_3, i_4])
394  tas_rowcols_t(i) = combine_tensor_index(nd_ind, t%dims)
395  END DO
396  END DO
397  END DO
398  END DO
399  END IF
400 
401  END FUNCTION
402 
403 ! **************************************************************************************************
404 !> \brief Create block size object for one matrix dimension
405 !> \param blk_size arrays for block sizes along all dimensions
406 !> \param map_blks tensor to matrix mapping object for blocks
407 !> \param which_dim for which dimension (1 or 2) distribution should be created
408 !> \return block size object
409 !> \author Patrick Seewald
410 ! **************************************************************************************************
411  FUNCTION new_dbt_tas_blk_size_t(blk_size, map_blks, which_dim)
412  TYPE(array_list), INTENT(IN) :: blk_size
413  TYPE(nd_to_2d_mapping), INTENT(IN) :: map_blks
414  INTEGER, INTENT(IN) :: which_dim
415  INTEGER(KIND=int_8), DIMENSION(2) :: matrix_dims
416  INTEGER, DIMENSION(:), ALLOCATABLE :: index_map
417  TYPE(dbt_tas_blk_size_t) :: new_dbt_tas_blk_size_t
418 
419  IF (which_dim == 1) THEN
420  ALLOCATE (index_map(ndims_mapping_row(map_blks)))
421  ALLOCATE (new_dbt_tas_blk_size_t%dims(ndims_mapping_row(map_blks)))
422  CALL dbt_get_mapping_info(map_blks, &
423  dims_2d_i8=matrix_dims, &
424  map1_2d=index_map, &
425  dims1_2d=new_dbt_tas_blk_size_t%dims)
426  ELSEIF (which_dim == 2) THEN
427  ALLOCATE (index_map(ndims_mapping_column(map_blks)))
428  ALLOCATE (new_dbt_tas_blk_size_t%dims(ndims_mapping_column(map_blks)))
429  CALL dbt_get_mapping_info(map_blks, &
430  dims_2d_i8=matrix_dims, &
431  map2_2d=index_map, &
432  dims2_2d=new_dbt_tas_blk_size_t%dims)
433  ELSE
434  cpabort("Unknown value for which_dim")
435  END IF
436 
437  new_dbt_tas_blk_size_t%blk_size = array_sublist(blk_size, index_map)
438  new_dbt_tas_blk_size_t%nmrowcol = matrix_dims(which_dim)
439 
440  new_dbt_tas_blk_size_t%nfullrowcol = product(int(sum_of_arrays(new_dbt_tas_blk_size_t%blk_size), &
441  kind=int_8))
442  END FUNCTION
443 
444 ! **************************************************************************************************
445 !> \author Patrick Seewald
446 ! **************************************************************************************************
447  FUNCTION tas_blk_size_t(t, rowcol)
448  CLASS(dbt_tas_blk_size_t), INTENT(IN) :: t
449  INTEGER(KIND=int_8), INTENT(IN) :: rowcol
450  INTEGER :: tas_blk_size_t
451  INTEGER, DIMENSION(SIZE(t%dims)) :: ind_blk
452  INTEGER, DIMENSION(SIZE(t%dims)) :: blk_size
453 
454  ind_blk(:) = split_tensor_index(rowcol, t%dims)
455  blk_size(:) = get_array_elements(t%blk_size, ind_blk)
456  tas_blk_size_t = product(blk_size)
457 
458  END FUNCTION
459 
460 ! **************************************************************************************************
461 !> \brief load balancing criterion whether to accept process grid dimension based on total number of
462 !> cores and tensor dimension
463 !> \param pdims_avail available process grid dimensions (total number of cores)
464 !> \param pdim process grid dimension to test
465 !> \param tdim tensor dimension corresponding to pdim
466 !> \param lb_ratio load imbalance acceptance factor
467 !> \author Patrick Seewald
468 ! **************************************************************************************************
469  PURE FUNCTION accept_pdims_loadbalancing(pdims_avail, pdim, tdim, lb_ratio)
470  INTEGER, INTENT(IN) :: pdims_avail
471  INTEGER, INTENT(IN) :: pdim
472  INTEGER, INTENT(IN) :: tdim
473  REAL(dp), INTENT(IN) :: lb_ratio
474  LOGICAL :: accept_pdims_loadbalancing
475 
476  accept_pdims_loadbalancing = .false.
477  IF (mod(pdims_avail, pdim) == 0) THEN
478  IF (real(tdim, dp)*lb_ratio < real(pdim, dp)) THEN
479  IF (mod(tdim, pdim) == 0) accept_pdims_loadbalancing = .true.
480  ELSE
481  accept_pdims_loadbalancing = .true.
482  END IF
483  END IF
484 
485  END FUNCTION
486 
487 ! **************************************************************************************************
488 !> \brief Create process grid dimensions corresponding to one dimension of the matrix representation
489 !> of a tensor, imposing that no process grid dimension is greater than the corresponding
490 !> tensor dimension.
491 !> \param nodes Total number of nodes available for this matrix dimension
492 !> \param dims process grid dimension corresponding to tensor_dims
493 !> \param tensor_dims tensor dimensions
494 !> \param lb_ratio load imbalance acceptance factor
495 !> \author Patrick Seewald
496 ! **************************************************************************************************
497  RECURSIVE SUBROUTINE dbt_mp_dims_create(nodes, dims, tensor_dims, lb_ratio)
498  INTEGER, INTENT(IN) :: nodes
499  INTEGER, DIMENSION(:), INTENT(INOUT) :: dims
500  INTEGER, DIMENSION(:), INTENT(IN) :: tensor_dims
501  REAL(dp), INTENT(IN), OPTIONAL :: lb_ratio
502 
503  INTEGER, DIMENSION(:), ALLOCATABLE :: tensor_dims_sorted, sort_indices, dims_store
504  REAL(dp), DIMENSION(:), ALLOCATABLE :: sort_key
505  INTEGER :: pdims_rem, idim, pdim
506  REAL(dp) :: lb_ratio_prv
507 
508  IF (PRESENT(lb_ratio)) THEN
509  lb_ratio_prv = lb_ratio
510  ELSE
511  lb_ratio_prv = 0.1_dp
512  END IF
513 
514  ALLOCATE (dims_store, source=dims)
515 
516  ! get default process grid dimensions
517  IF (any(dims == 0)) THEN
518  CALL mp_dims_create(nodes, dims)
519  END IF
520 
521  ! sort dimensions such that problematic grid dimensions (those who should be corrected) come first
522  ALLOCATE (sort_key(SIZE(tensor_dims)))
523  sort_key(:) = real(tensor_dims, dp)/dims
524 
525  ALLOCATE (tensor_dims_sorted, source=tensor_dims)
526  ALLOCATE (sort_indices(SIZE(sort_key)))
527  CALL sort(sort_key, SIZE(sort_key), sort_indices)
528  tensor_dims_sorted(:) = tensor_dims_sorted(sort_indices)
529  dims(:) = dims(sort_indices)
530 
531  ! remaining number of nodes
532  pdims_rem = nodes
533 
534  DO idim = 1, SIZE(tensor_dims_sorted)
535  IF (.NOT. accept_pdims_loadbalancing(pdims_rem, dims(idim), tensor_dims_sorted(idim), lb_ratio_prv)) THEN
536  pdim = tensor_dims_sorted(idim)
537  DO WHILE (.NOT. accept_pdims_loadbalancing(pdims_rem, pdim, tensor_dims_sorted(idim), lb_ratio_prv))
538  pdim = pdim - 1
539  END DO
540  dims(idim) = pdim
541  pdims_rem = pdims_rem/dims(idim)
542 
543  IF (idim .NE. SIZE(tensor_dims_sorted)) THEN
544  dims(idim + 1:) = 0
545  CALL mp_dims_create(pdims_rem, dims(idim + 1:))
546  ELSEIF (lb_ratio_prv < 0.5_dp) THEN
547  ! resort to a less strict load imbalance factor
548  dims(:) = dims_store
549  CALL dbt_mp_dims_create(nodes, dims, tensor_dims, 0.5_dp)
550  RETURN
551  ELSE
552  ! resort to default process grid dimensions
553  dims(:) = dims_store
554  CALL mp_dims_create(nodes, dims)
555  RETURN
556  END IF
557 
558  ELSE
559  pdims_rem = pdims_rem/dims(idim)
560  END IF
561  END DO
562 
563  dims(sort_indices) = dims
564 
565  END SUBROUTINE
566 
567 ! **************************************************************************************************
568 !> \brief Create an n-dimensional process grid.
569 !> We can not use a n-dimensional MPI cartesian grid for tensors since the mapping between
570 !> n-dim. and 2-dim. index allows for an arbitrary reordering of tensor index. Therefore we
571 !> can not use n-dim. MPI Cartesian grid because it may not be consistent with the respective
572 !> 2d grid. The 2d Cartesian MPI grid is the reference grid (since tensor data is stored as
573 !> DBM matrix) and this routine creates an object that is a n-dim. interface to this grid.
574 !> map1_2d and map2_2d don't need to be specified (correctly), grid may be redefined in
575 !> dbt_distribution_new. Note that pgrid is equivalent to a MPI cartesian grid only
576 !> if map1_2d and map2_2d don't reorder indices (which is the case if
577 !> [map1_2d, map2_2d] == [1, 2, ..., ndims]). Otherwise the mapping of grid coordinates to
578 !> processes depends on the ordering of the indices and is not equivalent to a MPI cartesian
579 !> grid.
580 !> \param mp_comm simple MPI Communicator
581 !> \param dims grid dimensions - if entries are 0, dimensions are chosen automatically.
582 !> \param pgrid n-dimensional grid object
583 !> \param map1_2d which nd-indices map to first matrix index and in which order
584 !> \param map2_2d which nd-indices map to first matrix index and in which order
585 !> \param tensor_dims tensor block dimensions. If present, process grid dimensions are created such
586 !> that good load balancing is ensured even if some of the tensor dimensions are
587 !> small (i.e. on the same order or smaller than nproc**(1/ndim) where ndim is
588 !> the tensor rank)
589 !> \param nsplit impose a constant split factor
590 !> \param dimsplit which matrix dimension to split
591 !> \author Patrick Seewald
592 ! **************************************************************************************************
593  SUBROUTINE dbt_pgrid_create_expert(mp_comm, dims, pgrid, map1_2d, map2_2d, tensor_dims, nsplit, dimsplit)
594  CLASS(mp_comm_type), INTENT(IN) :: mp_comm
595  INTEGER, DIMENSION(:), INTENT(INOUT) :: dims
596  TYPE(dbt_pgrid_type), INTENT(OUT) :: pgrid
597  INTEGER, DIMENSION(:), INTENT(IN) :: map1_2d, map2_2d
598  INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: tensor_dims
599  INTEGER, INTENT(IN), OPTIONAL :: nsplit, dimsplit
600  INTEGER, DIMENSION(2) :: pdims_2d
601  INTEGER :: nproc, ndims, handle
602  TYPE(dbt_tas_split_info) :: info
603 
604  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_pgrid_create_expert'
605 
606  CALL timeset(routinen, handle)
607 
608  ndims = SIZE(dims)
609 
610  nproc = mp_comm%num_pe
611  IF (any(dims == 0)) THEN
612  IF (.NOT. PRESENT(tensor_dims)) THEN
613  CALL mp_dims_create(nproc, dims)
614  ELSE
615  CALL dbt_mp_dims_create(nproc, dims, tensor_dims)
616  END IF
617  END IF
618  CALL create_nd_to_2d_mapping(pgrid%nd_index_grid, dims, map1_2d, map2_2d, base=0, col_major=.false.)
619  CALL dbt_get_mapping_info(pgrid%nd_index_grid, dims_2d=pdims_2d)
620  CALL pgrid%mp_comm_2d%create(mp_comm, 2, pdims_2d)
621 
622  IF (PRESENT(nsplit)) THEN
623  cpassert(PRESENT(dimsplit))
624  CALL dbt_tas_create_split(info, pgrid%mp_comm_2d, dimsplit, nsplit, opt_nsplit=.false.)
625  ALLOCATE (pgrid%tas_split_info, source=info)
626  END IF
627 
628  ! store number of MPI ranks because we need it for PURE function dbt_max_nblks_local
629  pgrid%nproc = nproc
630 
631  CALL timestop(handle)
632  END SUBROUTINE
633 
634 ! **************************************************************************************************
635 !> \brief Create a default nd process topology that is consistent with a given 2d topology.
636 !> Purpose: a nd tensor defined on the returned process grid can be represented as a DBM
637 !> matrix with the given 2d topology.
638 !> This is needed to enable contraction of 2 tensors (must have the same 2d process grid).
639 !> \param comm_2d communicator with 2-dimensional topology
640 !> \param map1_2d which nd-indices map to first matrix index and in which order
641 !> \param map2_2d which nd-indices map to second matrix index and in which order
642 !> \param dims_nd nd dimensions
643 !> \param pdims_2d if comm_2d does not have a cartesian topology associated, can input dimensions
644 !> with pdims_2d
645 !> \param tdims tensor block dimensions. If present, process grid dimensions are created such that
646 !> good load balancing is ensured even if some of the tensor dimensions are small
647 !> (i.e. on the same order or smaller than nproc**(1/ndim) where ndim is the tensor rank)
648 !> \return with nd cartesian grid
649 !> \author Patrick Seewald
650 ! **************************************************************************************************
651  FUNCTION dbt_nd_mp_comm(comm_2d, map1_2d, map2_2d, dims_nd, dims1_nd, dims2_nd, pdims_2d, tdims, &
652  nsplit, dimsplit)
653  CLASS(mp_comm_type), INTENT(IN) :: comm_2d
654  INTEGER, DIMENSION(:), INTENT(IN) :: map1_2d, map2_2d
655  INTEGER, DIMENSION(SIZE(map1_2d) + SIZE(map2_2d)), &
656  INTENT(IN), OPTIONAL :: dims_nd
657  INTEGER, DIMENSION(SIZE(map1_2d)), INTENT(IN), OPTIONAL :: dims1_nd
658  INTEGER, DIMENSION(SIZE(map2_2d)), INTENT(IN), OPTIONAL :: dims2_nd
659  INTEGER, DIMENSION(2), INTENT(IN), OPTIONAL :: pdims_2d
660  INTEGER, DIMENSION(SIZE(map1_2d) + SIZE(map2_2d)), &
661  INTENT(IN), OPTIONAL :: tdims
662  INTEGER, INTENT(IN), OPTIONAL :: nsplit, dimsplit
663  INTEGER :: ndim1, ndim2
664  INTEGER, DIMENSION(2) :: dims_2d
665 
666  INTEGER, DIMENSION(SIZE(map1_2d)) :: dims1_nd_prv
667  INTEGER, DIMENSION(SIZE(map2_2d)) :: dims2_nd_prv
668  INTEGER, DIMENSION(SIZE(map1_2d) + SIZE(map2_2d)) :: dims_nd_prv
669  INTEGER :: handle
670  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_nd_mp_comm'
671  TYPE(dbt_pgrid_type) :: dbt_nd_mp_comm
672 
673  CALL timeset(routinen, handle)
674 
675  ndim1 = SIZE(map1_2d); ndim2 = SIZE(map2_2d)
676 
677  IF (PRESENT(pdims_2d)) THEN
678  dims_2d(:) = pdims_2d
679  ELSE
680 ! This branch allows us to call this routine with a plain mp_comm_type without actually requiring an mp_cart_type
681 ! In a few cases in CP2K, this prevents erroneous calls to mpi_cart_get with a non-cartesian communicator
682  SELECT TYPE (comm_2d)
683  CLASS IS (mp_cart_type)
684  dims_2d = comm_2d%num_pe_cart
685  CLASS DEFAULT
686  CALL cp_abort(__location__, "If the argument pdims_2d is not given, the "// &
687  "communicator comm_2d must be of class mp_cart_type.")
688  END SELECT
689  END IF
690 
691  IF (.NOT. PRESENT(dims_nd)) THEN
692  dims1_nd_prv = 0; dims2_nd_prv = 0
693  IF (PRESENT(dims1_nd)) THEN
694  dims1_nd_prv(:) = dims1_nd
695  ELSE
696 
697  IF (PRESENT(tdims)) THEN
698  CALL dbt_mp_dims_create(dims_2d(1), dims1_nd_prv, tdims(map1_2d))
699  ELSE
700  CALL mp_dims_create(dims_2d(1), dims1_nd_prv)
701  END IF
702  END IF
703 
704  IF (PRESENT(dims2_nd)) THEN
705  dims2_nd_prv(:) = dims2_nd
706  ELSE
707  IF (PRESENT(tdims)) THEN
708  CALL dbt_mp_dims_create(dims_2d(2), dims2_nd_prv, tdims(map2_2d))
709  ELSE
710  CALL mp_dims_create(dims_2d(2), dims2_nd_prv)
711  END IF
712  END IF
713  dims_nd_prv(map1_2d) = dims1_nd_prv
714  dims_nd_prv(map2_2d) = dims2_nd_prv
715  ELSE
716  cpassert(product(dims_nd(map1_2d)) == dims_2d(1))
717  cpassert(product(dims_nd(map2_2d)) == dims_2d(2))
718  dims_nd_prv = dims_nd
719  END IF
720 
721  CALL dbt_pgrid_create_expert(comm_2d, dims_nd_prv, dbt_nd_mp_comm, &
722  tensor_dims=tdims, map1_2d=map1_2d, map2_2d=map2_2d, nsplit=nsplit, dimsplit=dimsplit)
723 
724  CALL timestop(handle)
725 
726  END FUNCTION
727 
728 ! **************************************************************************************************
729 !> \brief Release the MPI communicator.
730 !> \author Patrick Seewald
731 ! **************************************************************************************************
732  SUBROUTINE dbt_nd_mp_free(mp_comm)
733  TYPE(mp_comm_type), INTENT(INOUT) :: mp_comm
734 
735  CALL mp_comm%free()
736  END SUBROUTINE dbt_nd_mp_free
737 
738 ! **************************************************************************************************
739 !> \brief remap a process grid (needed when mapping between tensor and matrix index is changed)
740 !> \param map1_2d new mapping
741 !> \param map2_2d new mapping
742 !> \author Patrick Seewald
743 ! **************************************************************************************************
744  SUBROUTINE dbt_pgrid_remap(pgrid_in, map1_2d, map2_2d, pgrid_out)
745  TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid_in
746  INTEGER, DIMENSION(:), INTENT(IN) :: map1_2d, map2_2d
747  TYPE(dbt_pgrid_type), INTENT(OUT) :: pgrid_out
748  INTEGER, DIMENSION(:), ALLOCATABLE :: dims
749  INTEGER, DIMENSION(ndims_mapping_row(pgrid_in%nd_index_grid)) :: map1_2d_old
750  INTEGER, DIMENSION(ndims_mapping_column(pgrid_in%nd_index_grid)) :: map2_2d_old
751 
752  ALLOCATE (dims(SIZE(map1_2d) + SIZE(map2_2d)))
753  CALL dbt_get_mapping_info(pgrid_in%nd_index_grid, dims_nd=dims, map1_2d=map1_2d_old, map2_2d=map2_2d_old)
754  CALL dbt_pgrid_create_expert(pgrid_in%mp_comm_2d, dims, pgrid_out, map1_2d=map1_2d, map2_2d=map2_2d)
755  IF (array_eq_i(map1_2d_old, map1_2d) .AND. array_eq_i(map2_2d_old, map2_2d)) THEN
756  IF (ALLOCATED(pgrid_in%tas_split_info)) THEN
757  ALLOCATE (pgrid_out%tas_split_info, source=pgrid_in%tas_split_info)
758  CALL dbt_tas_info_hold(pgrid_out%tas_split_info)
759  END IF
760  END IF
761  END SUBROUTINE
762 
763 ! **************************************************************************************************
764 !> \brief as mp_environ but for special pgrid type
765 !> \author Patrick Seewald
766 ! **************************************************************************************************
767  SUBROUTINE mp_environ_pgrid(pgrid, dims, task_coor)
768  TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid
769  INTEGER, DIMENSION(ndims_mapping(pgrid%nd_index_grid)), INTENT(OUT) :: dims
770  INTEGER, DIMENSION(ndims_mapping(pgrid%nd_index_grid)), INTENT(OUT) :: task_coor
771  INTEGER, DIMENSION(2) :: task_coor_2d
772 
773  task_coor_2d = pgrid%mp_comm_2d%mepos_cart
774  CALL dbt_get_mapping_info(pgrid%nd_index_grid, dims_nd=dims)
775  task_coor = get_nd_indices_pgrid(pgrid%nd_index_grid, task_coor_2d)
776  END SUBROUTINE
777 
778 ! **************************************************************************************************
779 !> \brief Create a tensor distribution.
780 !> \param pgrid process grid
781 !> \param map1_2d which nd-indices map to first matrix index and in which order
782 !> \param map2_2d which nd-indices map to second matrix index and in which order
783 !> \param own_comm whether distribution should own communicator
784 !> \author Patrick Seewald
785 ! **************************************************************************************************
786  SUBROUTINE dbt_distribution_new_expert(dist, pgrid, map1_2d, map2_2d, nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4, own_comm)
787  TYPE(dbt_distribution_type), INTENT(OUT) :: dist
788  TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid
789  INTEGER, DIMENSION(:), INTENT(IN) :: map1_2d
790  INTEGER, DIMENSION(:), INTENT(IN) :: map2_2d
791  INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4
792  LOGICAL, INTENT(IN), OPTIONAL :: own_comm
793  INTEGER :: ndims
794  TYPE(mp_cart_type) :: comm_2d
795  INTEGER, DIMENSION(2) :: pdims_2d_check, &
796  pdims_2d
797  INTEGER, DIMENSION(SIZE(map1_2d) + SIZE(map2_2d)) :: dims, nblks_nd, task_coor
798  TYPE(array_list) :: nd_dist
799  TYPE(nd_to_2d_mapping) :: map_blks, map_grid
800  INTEGER :: handle
801  TYPE(dbt_tas_dist_t) :: row_dist_obj, col_dist_obj
802  TYPE(dbt_pgrid_type) :: pgrid_prv
803  LOGICAL :: need_pgrid_remap
804  INTEGER, DIMENSION(ndims_mapping_row(pgrid%nd_index_grid)) :: map1_2d_check
805  INTEGER, DIMENSION(ndims_mapping_column(pgrid%nd_index_grid)) :: map2_2d_check
806  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_distribution_new_expert'
807 
808  CALL timeset(routinen, handle)
809  ndims = SIZE(map1_2d) + SIZE(map2_2d)
810  cpassert(ndims .GE. 2 .AND. ndims .LE. 4)
811 
812  CALL create_array_list(nd_dist, ndims, nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4)
813 
814  nblks_nd(:) = sizes_of_arrays(nd_dist)
815 
816  need_pgrid_remap = .true.
817  IF (PRESENT(own_comm)) THEN
818  CALL dbt_get_mapping_info(pgrid%nd_index_grid, map1_2d=map1_2d_check, map2_2d=map2_2d_check)
819  IF (own_comm) THEN
820  IF (.NOT. array_eq_i(map1_2d_check, map1_2d) .OR. .NOT. array_eq_i(map2_2d_check, map2_2d)) THEN
821  cpabort("map1_2d / map2_2d are not consistent with pgrid")
822  END IF
823  pgrid_prv = pgrid
824  need_pgrid_remap = .false.
825  END IF
826  END IF
827 
828  IF (need_pgrid_remap) CALL dbt_pgrid_remap(pgrid, map1_2d, map2_2d, pgrid_prv)
829 
830  ! check that 2d process topology is consistent with nd topology.
831  CALL mp_environ_pgrid(pgrid_prv, dims, task_coor)
832 
833  ! process grid index mapping
834  CALL create_nd_to_2d_mapping(map_grid, dims, map1_2d, map2_2d, base=0, col_major=.false.)
835 
836  ! blk index mapping
837  CALL create_nd_to_2d_mapping(map_blks, nblks_nd, map1_2d, map2_2d)
838 
839  row_dist_obj = dbt_tas_dist_t(nd_dist, map_blks, map_grid, 1)
840  col_dist_obj = dbt_tas_dist_t(nd_dist, map_blks, map_grid, 2)
841 
842  CALL dbt_get_mapping_info(map_grid, dims_2d=pdims_2d)
843 
844  comm_2d = pgrid_prv%mp_comm_2d
845 
846  pdims_2d_check = comm_2d%num_pe_cart
847  IF (any(pdims_2d_check .NE. pdims_2d)) THEN
848  cpabort("inconsistent process grid dimensions")
849  END IF
850 
851  IF (ALLOCATED(pgrid_prv%tas_split_info)) THEN
852  CALL dbt_tas_distribution_new(dist%dist, comm_2d, row_dist_obj, col_dist_obj, split_info=pgrid_prv%tas_split_info)
853  ELSE
854  CALL dbt_tas_distribution_new(dist%dist, comm_2d, row_dist_obj, col_dist_obj)
855  ALLOCATE (pgrid_prv%tas_split_info, source=dist%dist%info)
856  CALL dbt_tas_info_hold(pgrid_prv%tas_split_info)
857  END IF
858 
859  dist%nd_dist = nd_dist
860  dist%pgrid = pgrid_prv
861 
862  ALLOCATE (dist%refcount)
863  dist%refcount = 1
864  CALL timestop(handle)
865 
866  CONTAINS
867  PURE FUNCTION array_eq_i(arr1, arr2)
868  INTEGER, INTENT(IN), DIMENSION(:) :: arr1
869  INTEGER, INTENT(IN), DIMENSION(:) :: arr2
870  LOGICAL :: array_eq_i
871 
872  array_eq_i = .false.
873  IF (SIZE(arr1) .EQ. SIZE(arr2)) array_eq_i = all(arr1 == arr2)
874 
875  END FUNCTION
876 
877  END SUBROUTINE
878 
879 ! **************************************************************************************************
880 !> \brief Create a tensor distribution.
881 !> \param pgrid process grid
882 !> \param nd_dist_i distribution vectors for all tensor dimensions
883 !> \author Patrick Seewald
884 ! **************************************************************************************************
885  SUBROUTINE dbt_distribution_new(dist, pgrid, nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4)
886  TYPE(dbt_distribution_type), INTENT(OUT) :: dist
887  TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid
888  INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4
889  INTEGER, DIMENSION(ndims_mapping_row(pgrid%nd_index_grid)) :: map1_2d
890  INTEGER, DIMENSION(ndims_mapping_column(pgrid%nd_index_grid)) :: map2_2d
891  INTEGER :: ndims
892 
893  CALL dbt_get_mapping_info(pgrid%nd_index_grid, map1_2d=map1_2d, map2_2d=map2_2d, ndim_nd=ndims)
894 
895  CALL dbt_distribution_new_expert(dist, pgrid, map1_2d, map2_2d, nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4)
896 
897  END SUBROUTINE
898 
899 ! **************************************************************************************************
900 !> \brief destroy process grid
901 !> \param keep_comm if .TRUE. communicator is not freed
902 !> \author Patrick Seewald
903 ! **************************************************************************************************
904  SUBROUTINE dbt_pgrid_destroy(pgrid, keep_comm)
905  TYPE(dbt_pgrid_type), INTENT(INOUT) :: pgrid
906  LOGICAL, INTENT(IN), OPTIONAL :: keep_comm
907  LOGICAL :: keep_comm_prv
908  IF (PRESENT(keep_comm)) THEN
909  keep_comm_prv = keep_comm
910  ELSE
911  keep_comm_prv = .false.
912  END IF
913  IF (.NOT. keep_comm_prv) CALL pgrid%mp_comm_2d%free()
914  CALL destroy_nd_to_2d_mapping(pgrid%nd_index_grid)
915  IF (ALLOCATED(pgrid%tas_split_info) .AND. .NOT. keep_comm_prv) THEN
916  CALL dbt_tas_release_info(pgrid%tas_split_info)
917  DEALLOCATE (pgrid%tas_split_info)
918  END IF
919  END SUBROUTINE
920 
921 ! **************************************************************************************************
922 !> \brief Destroy tensor distribution
923 !> \author Patrick Seewald
924 ! **************************************************************************************************
925  SUBROUTINE dbt_distribution_destroy(dist)
926  TYPE(dbt_distribution_type), INTENT(INOUT) :: dist
927  INTEGER :: handle
928  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_distribution_destroy'
929  LOGICAL :: abort
930 
931  CALL timeset(routinen, handle)
932  CALL dbt_tas_distribution_destroy(dist%dist)
933  CALL destroy_array_list(dist%nd_dist)
934 
935  abort = .false.
936  IF (.NOT. ASSOCIATED(dist%refcount)) THEN
937  abort = .true.
938  ELSEIF (dist%refcount < 1) THEN
939  abort = .true.
940  END IF
941 
942  IF (abort) THEN
943  cpabort("can not destroy non-existing tensor distribution")
944  END IF
945 
946  dist%refcount = dist%refcount - 1
947 
948  IF (dist%refcount == 0) THEN
949  CALL dbt_pgrid_destroy(dist%pgrid)
950  DEALLOCATE (dist%refcount)
951  ELSE
952  CALL dbt_pgrid_destroy(dist%pgrid, keep_comm=.true.)
953  END IF
954 
955  CALL timestop(handle)
956  END SUBROUTINE
957 
958 ! **************************************************************************************************
959 !> \brief reference counting for distribution
960 !> (only needed for communicator handle that must be freed when no longer needed)
961 !> \author Patrick Seewald
962 ! **************************************************************************************************
963  SUBROUTINE dbt_distribution_hold(dist)
964  TYPE(dbt_distribution_type), INTENT(IN) :: dist
965  INTEGER, POINTER :: ref
966 
967  IF (dist%refcount < 1) THEN
968  cpabort("can not hold non-existing tensor distribution")
969  END IF
970  ref => dist%refcount
971  ref = ref + 1
972  END SUBROUTINE
973 
974 ! **************************************************************************************************
975 !> \brief get distribution from tensor
976 !> \return distribution
977 !> \author Patrick Seewald
978 ! **************************************************************************************************
979  FUNCTION dbt_distribution(tensor)
980  TYPE(dbt_type), INTENT(IN) :: tensor
981  TYPE(dbt_distribution_type) :: dbt_distribution
982 
983  CALL dbt_tas_get_info(tensor%matrix_rep, distribution=dbt_distribution%dist)
984  dbt_distribution%pgrid = tensor%pgrid
985  dbt_distribution%nd_dist = tensor%nd_dist
986  dbt_distribution%refcount => dbt_distribution%refcount
987  END FUNCTION
988 
989 ! **************************************************************************************************
990 !> \author Patrick Seewald
991 ! **************************************************************************************************
992  SUBROUTINE dbt_distribution_remap(dist_in, map1_2d, map2_2d, dist_out)
993  TYPE(dbt_distribution_type), INTENT(IN) :: dist_in
994  INTEGER, DIMENSION(:), INTENT(IN) :: map1_2d, map2_2d
995  TYPE(dbt_distribution_type), INTENT(OUT) :: dist_out
996  INTEGER, DIMENSION(:), ALLOCATABLE :: dist_1, dist_2, dist_3, dist_4
997  INTEGER :: ndims
998  ndims = SIZE(map1_2d) + SIZE(map2_2d)
999  IF (ndims == 1) THEN
1000  CALL get_arrays(dist_in%nd_dist, dist_1)
1001  CALL dbt_distribution_new_expert(dist_out, dist_in%pgrid, map1_2d, map2_2d, dist_1)
1002  END IF
1003  IF (ndims == 2) THEN
1004  CALL get_arrays(dist_in%nd_dist, dist_1, dist_2)
1005  CALL dbt_distribution_new_expert(dist_out, dist_in%pgrid, map1_2d, map2_2d, dist_1, dist_2)
1006  END IF
1007  IF (ndims == 3) THEN
1008  CALL get_arrays(dist_in%nd_dist, dist_1, dist_2, dist_3)
1009  CALL dbt_distribution_new_expert(dist_out, dist_in%pgrid, map1_2d, map2_2d, dist_1, dist_2, dist_3)
1010  END IF
1011  IF (ndims == 4) THEN
1012  CALL get_arrays(dist_in%nd_dist, dist_1, dist_2, dist_3, dist_4)
1013  CALL dbt_distribution_new_expert(dist_out, dist_in%pgrid, map1_2d, map2_2d, dist_1, dist_2, dist_3, dist_4)
1014  END IF
1015  END SUBROUTINE
1016 
1017 ! **************************************************************************************************
1018 !> \brief create a tensor.
1019 !> For performance, the arguments map1_2d and map2_2d (controlling matrix representation of
1020 !> tensor) should be consistent with the the contraction to be performed (see documentation
1021 !> of dbt_contract).
1022 !> \param map1_2d which nd-indices to map to first 2d index and in which order
1023 !> \param map2_2d which nd-indices to map to first 2d index and in which order
1024 !> \param blk_size_i blk sizes in each dimension
1025 !> \author Patrick Seewald
1026 ! **************************************************************************************************
1027  SUBROUTINE dbt_create_new(tensor, name, dist, map1_2d, map2_2d, &
1028  blk_size_1, blk_size_2, blk_size_3, blk_size_4)
1029  TYPE(dbt_type), INTENT(OUT) :: tensor
1030  CHARACTER(len=*), INTENT(IN) :: name
1031  TYPE(dbt_distribution_type), INTENT(INOUT) :: dist
1032  INTEGER, DIMENSION(:), INTENT(IN) :: map1_2d
1033  INTEGER, DIMENSION(:), INTENT(IN) :: map2_2d
1034  INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: blk_size_1, blk_size_2, blk_size_3, blk_size_4
1035  INTEGER :: ndims
1036  INTEGER(KIND=int_8), DIMENSION(2) :: dims_2d
1037  INTEGER, DIMENSION(SIZE(map1_2d) + SIZE(map2_2d)) :: dims, pdims, task_coor
1038  TYPE(dbt_tas_blk_size_t) :: col_blk_size_obj, row_blk_size_obj
1039  TYPE(dbt_distribution_type) :: dist_new
1040  TYPE(array_list) :: blk_size, blks_local
1041  TYPE(nd_to_2d_mapping) :: map
1042  INTEGER :: handle
1043  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_create_new'
1044  INTEGER, DIMENSION(:), ALLOCATABLE :: blks_local_1, blks_local_2, blks_local_3, blks_local_4
1045  INTEGER, DIMENSION(:), ALLOCATABLE :: dist_1, dist_2, dist_3, dist_4
1046  INTEGER :: iblk_count, iblk
1047  INTEGER, DIMENSION(:), ALLOCATABLE :: nblks_local, nfull_local
1048 
1049  CALL timeset(routinen, handle)
1050  ndims = SIZE(map1_2d) + SIZE(map2_2d)
1051  CALL create_array_list(blk_size, ndims, blk_size_1, blk_size_2, blk_size_3, blk_size_4)
1052  dims = sizes_of_arrays(blk_size)
1053 
1054  CALL create_nd_to_2d_mapping(map, dims, map1_2d, map2_2d)
1055  CALL dbt_get_mapping_info(map, dims_2d_i8=dims_2d)
1056 
1057  row_blk_size_obj = dbt_tas_blk_size_t(blk_size, map, 1)
1058  col_blk_size_obj = dbt_tas_blk_size_t(blk_size, map, 2)
1059 
1060  CALL dbt_distribution_remap(dist, map1_2d, map2_2d, dist_new)
1061 
1062  ALLOCATE (tensor%matrix_rep)
1063  CALL dbt_tas_create(matrix=tensor%matrix_rep, &
1064  name=trim(name)//" matrix", &
1065  dist=dist_new%dist, &
1066  row_blk_size=row_blk_size_obj, &
1067  col_blk_size=col_blk_size_obj)
1068 
1069  tensor%owns_matrix = .true.
1070 
1071  tensor%nd_index_blk = map
1072  tensor%name = name
1073 
1074  CALL dbt_tas_finalize(tensor%matrix_rep)
1075  CALL destroy_nd_to_2d_mapping(map)
1076 
1077  ! map element-wise tensor index
1078  CALL create_nd_to_2d_mapping(map, sum_of_arrays(blk_size), map1_2d, map2_2d)
1079  tensor%nd_index = map
1080  tensor%blk_sizes = blk_size
1081 
1082  CALL mp_environ_pgrid(dist_new%pgrid, pdims, task_coor)
1083 
1084  IF (ndims == 1) THEN
1085  CALL get_arrays(dist_new%nd_dist, dist_1)
1086  END IF
1087  IF (ndims == 2) THEN
1088  CALL get_arrays(dist_new%nd_dist, dist_1, dist_2)
1089  END IF
1090  IF (ndims == 3) THEN
1091  CALL get_arrays(dist_new%nd_dist, dist_1, dist_2, dist_3)
1092  END IF
1093  IF (ndims == 4) THEN
1094  CALL get_arrays(dist_new%nd_dist, dist_1, dist_2, dist_3, dist_4)
1095  END IF
1096 
1097  ALLOCATE (nblks_local(ndims))
1098  ALLOCATE (nfull_local(ndims))
1099  nfull_local(:) = 0
1100  IF (ndims .GE. 1) THEN
1101  nblks_local(1) = count(dist_1 == task_coor(1))
1102  ALLOCATE (blks_local_1(nblks_local(1)))
1103  iblk_count = 0
1104  DO iblk = 1, SIZE(dist_1)
1105  IF (dist_1(iblk) == task_coor(1)) THEN
1106  iblk_count = iblk_count + 1
1107  blks_local_1(iblk_count) = iblk
1108  nfull_local(1) = nfull_local(1) + blk_size_1(iblk)
1109  END IF
1110  END DO
1111  END IF
1112  IF (ndims .GE. 2) THEN
1113  nblks_local(2) = count(dist_2 == task_coor(2))
1114  ALLOCATE (blks_local_2(nblks_local(2)))
1115  iblk_count = 0
1116  DO iblk = 1, SIZE(dist_2)
1117  IF (dist_2(iblk) == task_coor(2)) THEN
1118  iblk_count = iblk_count + 1
1119  blks_local_2(iblk_count) = iblk
1120  nfull_local(2) = nfull_local(2) + blk_size_2(iblk)
1121  END IF
1122  END DO
1123  END IF
1124  IF (ndims .GE. 3) THEN
1125  nblks_local(3) = count(dist_3 == task_coor(3))
1126  ALLOCATE (blks_local_3(nblks_local(3)))
1127  iblk_count = 0
1128  DO iblk = 1, SIZE(dist_3)
1129  IF (dist_3(iblk) == task_coor(3)) THEN
1130  iblk_count = iblk_count + 1
1131  blks_local_3(iblk_count) = iblk
1132  nfull_local(3) = nfull_local(3) + blk_size_3(iblk)
1133  END IF
1134  END DO
1135  END IF
1136  IF (ndims .GE. 4) THEN
1137  nblks_local(4) = count(dist_4 == task_coor(4))
1138  ALLOCATE (blks_local_4(nblks_local(4)))
1139  iblk_count = 0
1140  DO iblk = 1, SIZE(dist_4)
1141  IF (dist_4(iblk) == task_coor(4)) THEN
1142  iblk_count = iblk_count + 1
1143  blks_local_4(iblk_count) = iblk
1144  nfull_local(4) = nfull_local(4) + blk_size_4(iblk)
1145  END IF
1146  END DO
1147  END IF
1148 
1149  IF (ndims == 1) THEN
1150  CALL create_array_list(blks_local, 1, blks_local_1)
1151  END IF
1152  IF (ndims == 2) THEN
1153  CALL create_array_list(blks_local, 2, blks_local_1, blks_local_2)
1154  END IF
1155  IF (ndims == 3) THEN
1156  CALL create_array_list(blks_local, 3, blks_local_1, blks_local_2, blks_local_3)
1157  END IF
1158  IF (ndims == 4) THEN
1159  CALL create_array_list(blks_local, 4, blks_local_1, blks_local_2, blks_local_3, blks_local_4)
1160  END IF
1161 
1162  ALLOCATE (tensor%nblks_local(ndims))
1163  ALLOCATE (tensor%nfull_local(ndims))
1164  tensor%nblks_local(:) = nblks_local
1165  tensor%nfull_local(:) = nfull_local
1166 
1167  tensor%blks_local = blks_local
1168 
1169  tensor%nd_dist = dist_new%nd_dist
1170  tensor%pgrid = dist_new%pgrid
1171 
1172  CALL dbt_distribution_hold(dist_new)
1173  tensor%refcount => dist_new%refcount
1174  CALL dbt_distribution_destroy(dist_new)
1175 
1176  CALL array_offsets(tensor%blk_sizes, tensor%blk_offsets)
1177 
1178  tensor%valid = .true.
1179  CALL timestop(handle)
1180  END SUBROUTINE
1181 
1182 ! **************************************************************************************************
1183 !> \brief reference counting for tensors
1184 !> (only needed for communicator handle that must be freed when no longer needed)
1185 !> \author Patrick Seewald
1186 ! **************************************************************************************************
1187  SUBROUTINE dbt_hold(tensor)
1188  TYPE(dbt_type), INTENT(IN) :: tensor
1189  INTEGER, POINTER :: ref
1190 
1191  IF (tensor%refcount < 1) THEN
1192  cpabort("can not hold non-existing tensor")
1193  END IF
1194  ref => tensor%refcount
1195  ref = ref + 1
1196 
1197  END SUBROUTINE
1198 
1199 ! **************************************************************************************************
1200 !> \brief how many tensor dimensions are mapped to matrix row
1201 !> \author Patrick Seewald
1202 ! **************************************************************************************************
1203  PURE FUNCTION ndims_matrix_row(tensor)
1204  TYPE(dbt_type), INTENT(IN) :: tensor
1205  INTEGER(int_8) :: ndims_matrix_row
1206 
1207  ndims_matrix_row = ndims_mapping_row(tensor%nd_index_blk)
1208 
1209  END FUNCTION
1210 
1211 ! **************************************************************************************************
1212 !> \brief how many tensor dimensions are mapped to matrix column
1213 !> \author Patrick Seewald
1214 ! **************************************************************************************************
1215  PURE FUNCTION ndims_matrix_column(tensor)
1216  TYPE(dbt_type), INTENT(IN) :: tensor
1217  INTEGER(int_8) :: ndims_matrix_column
1218 
1219  ndims_matrix_column = ndims_mapping_column(tensor%nd_index_blk)
1220  END FUNCTION
1221 
1222 ! **************************************************************************************************
1223 !> \brief tensor rank
1224 !> \author Patrick Seewald
1225 ! **************************************************************************************************
1226  PURE FUNCTION ndims_tensor(tensor)
1227  TYPE(dbt_type), INTENT(IN) :: tensor
1228  INTEGER :: ndims_tensor
1229 
1230  ndims_tensor = tensor%nd_index%ndim_nd
1231  END FUNCTION
1232 
1233 ! **************************************************************************************************
1234 !> \brief tensor dimensions
1235 !> \author Patrick Seewald
1236 ! **************************************************************************************************
1237  SUBROUTINE dims_tensor(tensor, dims)
1238  TYPE(dbt_type), INTENT(IN) :: tensor
1239  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1240  INTENT(OUT) :: dims
1241 
1242  cpassert(tensor%valid)
1243  dims = tensor%nd_index%dims_nd
1244  END SUBROUTINE
1245 
1246 ! **************************************************************************************************
1247 !> \brief create a tensor from template
1248 !> \author Patrick Seewald
1249 ! **************************************************************************************************
1250  SUBROUTINE dbt_create_template(tensor_in, tensor, name, dist, map1_2d, map2_2d)
1251  TYPE(dbt_type), INTENT(INOUT) :: tensor_in
1252  TYPE(dbt_type), INTENT(OUT) :: tensor
1253  CHARACTER(len=*), INTENT(IN), OPTIONAL :: name
1254  TYPE(dbt_distribution_type), &
1255  INTENT(INOUT), OPTIONAL :: dist
1256  INTEGER, DIMENSION(:), INTENT(IN), &
1257  OPTIONAL :: map1_2d, map2_2d
1258  INTEGER :: handle
1259  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_create_template'
1260  INTEGER, DIMENSION(:), ALLOCATABLE :: bsize_1, bsize_2, bsize_3, bsize_4
1261  INTEGER, DIMENSION(:), ALLOCATABLE :: map1_2d_prv, map2_2d_prv
1262  CHARACTER(len=default_string_length) :: name_prv
1263  TYPE(dbt_distribution_type) :: dist_prv
1264 
1265  CALL timeset(routinen, handle)
1266 
1267  IF (PRESENT(dist) .OR. PRESENT(map1_2d) .OR. PRESENT(map2_2d)) THEN
1268  ! need to create matrix representation from scratch
1269  IF (PRESENT(dist)) THEN
1270  dist_prv = dist
1271  ELSE
1272  dist_prv = dbt_distribution(tensor_in)
1273  END IF
1274  IF (PRESENT(map1_2d) .AND. PRESENT(map2_2d)) THEN
1275  ALLOCATE (map1_2d_prv, source=map1_2d)
1276  ALLOCATE (map2_2d_prv, source=map2_2d)
1277  ELSE
1278  ALLOCATE (map1_2d_prv(ndims_matrix_row(tensor_in)))
1279  ALLOCATE (map2_2d_prv(ndims_matrix_column(tensor_in)))
1280  CALL dbt_get_mapping_info(tensor_in%nd_index_blk, map1_2d=map1_2d_prv, map2_2d=map2_2d_prv)
1281  END IF
1282  IF (PRESENT(name)) THEN
1283  name_prv = name
1284  ELSE
1285  name_prv = tensor_in%name
1286  END IF
1287 
1288  IF (ndims_tensor(tensor_in) == 1) THEN
1289  CALL get_arrays(tensor_in%blk_sizes, bsize_1)
1290  CALL dbt_create(tensor, name_prv, dist_prv, map1_2d_prv, map2_2d_prv, &
1291  bsize_1)
1292  END IF
1293  IF (ndims_tensor(tensor_in) == 2) THEN
1294  CALL get_arrays(tensor_in%blk_sizes, bsize_1, bsize_2)
1295  CALL dbt_create(tensor, name_prv, dist_prv, map1_2d_prv, map2_2d_prv, &
1296  bsize_1, bsize_2)
1297  END IF
1298  IF (ndims_tensor(tensor_in) == 3) THEN
1299  CALL get_arrays(tensor_in%blk_sizes, bsize_1, bsize_2, bsize_3)
1300  CALL dbt_create(tensor, name_prv, dist_prv, map1_2d_prv, map2_2d_prv, &
1301  bsize_1, bsize_2, bsize_3)
1302  END IF
1303  IF (ndims_tensor(tensor_in) == 4) THEN
1304  CALL get_arrays(tensor_in%blk_sizes, bsize_1, bsize_2, bsize_3, bsize_4)
1305  CALL dbt_create(tensor, name_prv, dist_prv, map1_2d_prv, map2_2d_prv, &
1306  bsize_1, bsize_2, bsize_3, bsize_4)
1307  END IF
1308  ELSE
1309  ! create matrix representation from template
1310  ALLOCATE (tensor%matrix_rep)
1311  IF (.NOT. PRESENT(name)) THEN
1312  CALL dbt_tas_create(tensor_in%matrix_rep, tensor%matrix_rep, &
1313  name=trim(tensor_in%name)//" matrix")
1314  ELSE
1315  CALL dbt_tas_create(tensor_in%matrix_rep, tensor%matrix_rep, name=trim(name)//" matrix")
1316  END IF
1317  tensor%owns_matrix = .true.
1318  CALL dbt_tas_finalize(tensor%matrix_rep)
1319 
1320  tensor%nd_index_blk = tensor_in%nd_index_blk
1321  tensor%nd_index = tensor_in%nd_index
1322  tensor%blk_sizes = tensor_in%blk_sizes
1323  tensor%blk_offsets = tensor_in%blk_offsets
1324  tensor%nd_dist = tensor_in%nd_dist
1325  tensor%blks_local = tensor_in%blks_local
1326  ALLOCATE (tensor%nblks_local(ndims_tensor(tensor_in)))
1327  tensor%nblks_local(:) = tensor_in%nblks_local
1328  ALLOCATE (tensor%nfull_local(ndims_tensor(tensor_in)))
1329  tensor%nfull_local(:) = tensor_in%nfull_local
1330  tensor%pgrid = tensor_in%pgrid
1331 
1332  tensor%refcount => tensor_in%refcount
1333  CALL dbt_hold(tensor)
1334 
1335  tensor%valid = .true.
1336  IF (PRESENT(name)) THEN
1337  tensor%name = name
1338  ELSE
1339  tensor%name = tensor_in%name
1340  END IF
1341  END IF
1342  CALL timestop(handle)
1343  END SUBROUTINE
1344 
1345 ! **************************************************************************************************
1346 !> \brief Create 2-rank tensor from matrix.
1347 !> \author Patrick Seewald
1348 ! **************************************************************************************************
1349  SUBROUTINE dbt_create_matrix(matrix_in, tensor, order, name)
1350  TYPE(dbcsr_type), INTENT(IN) :: matrix_in
1351  TYPE(dbt_type), INTENT(OUT) :: tensor
1352  INTEGER, DIMENSION(2), INTENT(IN), OPTIONAL :: order
1353  CHARACTER(len=*), INTENT(IN), OPTIONAL :: name
1354 
1355  CHARACTER(len=default_string_length) :: name_in
1356  INTEGER, DIMENSION(2) :: order_in
1357  TYPE(mp_comm_type) :: comm_2d
1358  TYPE(dbcsr_distribution_type) :: matrix_dist
1359  TYPE(dbt_distribution_type) :: dist
1360  INTEGER, DIMENSION(:), POINTER :: row_blk_size, col_blk_size
1361  INTEGER, DIMENSION(:), POINTER :: col_dist, row_dist
1362  INTEGER :: handle, comm_2d_handle
1363  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_create_matrix'
1364  TYPE(dbt_pgrid_type) :: comm_nd
1365  INTEGER, DIMENSION(2) :: pdims_2d
1366 
1367  CALL timeset(routinen, handle)
1368 
1369  NULLIFY (row_blk_size, col_blk_size, col_dist, row_dist)
1370  IF (PRESENT(name)) THEN
1371  name_in = name
1372  ELSE
1373  CALL dbcsr_get_info(matrix_in, name=name_in)
1374  END IF
1375 
1376  IF (PRESENT(order)) THEN
1377  order_in = order
1378  ELSE
1379  order_in = [1, 2]
1380  END IF
1381 
1382  CALL dbcsr_get_info(matrix_in, distribution=matrix_dist)
1383  CALL dbcsr_distribution_get(matrix_dist, group=comm_2d_handle, row_dist=row_dist, col_dist=col_dist, &
1384  nprows=pdims_2d(1), npcols=pdims_2d(2))
1385  CALL comm_2d%set_handle(comm_2d_handle)
1386  comm_nd = dbt_nd_mp_comm(comm_2d, [order_in(1)], [order_in(2)], pdims_2d=pdims_2d)
1387 
1388  CALL dbt_distribution_new_expert( &
1389  dist, &
1390  comm_nd, &
1391  [order_in(1)], [order_in(2)], &
1392  row_dist, col_dist, own_comm=.true.)
1393 
1394  CALL dbcsr_get_info(matrix_in, row_blk_size=row_blk_size, col_blk_size=col_blk_size)
1395 
1396  CALL dbt_create_new(tensor, name_in, dist, &
1397  [order_in(1)], [order_in(2)], &
1398  row_blk_size, &
1399  col_blk_size)
1400 
1401  CALL dbt_distribution_destroy(dist)
1402  CALL timestop(handle)
1403  END SUBROUTINE
1404 
1405 ! **************************************************************************************************
1406 !> \brief Destroy a tensor
1407 !> \author Patrick Seewald
1408 ! **************************************************************************************************
1409  SUBROUTINE dbt_destroy(tensor)
1410  TYPE(dbt_type), INTENT(INOUT) :: tensor
1411  INTEGER :: handle
1412  CHARACTER(LEN=*), PARAMETER :: routinen = 'dbt_destroy'
1413  LOGICAL :: abort
1414 
1415  CALL timeset(routinen, handle)
1416  IF (tensor%owns_matrix) THEN
1417  CALL dbt_tas_destroy(tensor%matrix_rep)
1418  DEALLOCATE (tensor%matrix_rep)
1419  ELSE
1420  NULLIFY (tensor%matrix_rep)
1421  END IF
1422  tensor%owns_matrix = .false.
1423 
1424  CALL destroy_nd_to_2d_mapping(tensor%nd_index_blk)
1425  CALL destroy_nd_to_2d_mapping(tensor%nd_index)
1426  !CALL destroy_nd_to_2d_mapping(tensor%nd_index_grid)
1427  CALL destroy_array_list(tensor%blk_sizes)
1428  CALL destroy_array_list(tensor%blk_offsets)
1429  CALL destroy_array_list(tensor%nd_dist)
1430  CALL destroy_array_list(tensor%blks_local)
1431 
1432  DEALLOCATE (tensor%nblks_local, tensor%nfull_local)
1433 
1434  abort = .false.
1435  IF (.NOT. ASSOCIATED(tensor%refcount)) THEN
1436  abort = .true.
1437  ELSEIF (tensor%refcount < 1) THEN
1438  abort = .true.
1439  END IF
1440 
1441  IF (abort) THEN
1442  cpabort("can not destroy non-existing tensor")
1443  END IF
1444 
1445  tensor%refcount = tensor%refcount - 1
1446 
1447  IF (tensor%refcount == 0) THEN
1448  CALL dbt_pgrid_destroy(tensor%pgrid)
1449  !CALL tensor%comm_2d%free()
1450  !CALL tensor%comm_nd%free()
1451  DEALLOCATE (tensor%refcount)
1452  ELSE
1453  CALL dbt_pgrid_destroy(tensor%pgrid, keep_comm=.true.)
1454  END IF
1455 
1456  tensor%valid = .false.
1457  tensor%name = ""
1458  CALL timestop(handle)
1459  END SUBROUTINE
1460 
1461 ! **************************************************************************************************
1462 !> \brief tensor block dimensions
1463 !> \author Patrick Seewald
1464 ! **************************************************************************************************
1465  SUBROUTINE blk_dims_tensor(tensor, dims)
1466  TYPE(dbt_type), INTENT(IN) :: tensor
1467  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1468  INTENT(OUT) :: dims
1469 
1470  cpassert(tensor%valid)
1471  dims = tensor%nd_index_blk%dims_nd
1472  END SUBROUTINE
1473 
1474 ! **************************************************************************************************
1475 !> \brief Size of tensor block
1476 !> \author Patrick Seewald
1477 ! **************************************************************************************************
1478  SUBROUTINE dbt_blk_sizes(tensor, ind, blk_size)
1479  TYPE(dbt_type), INTENT(IN) :: tensor
1480  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1481  INTENT(IN) :: ind
1482  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1483  INTENT(OUT) :: blk_size
1484 
1485  blk_size(:) = get_array_elements(tensor%blk_sizes, ind)
1486  END SUBROUTINE
1487 
1488 ! **************************************************************************************************
1489 !> \brief offset of tensor block
1490 !> \param ind block index
1491 !> \param blk_offset block offset
1492 !> \author Patrick Seewald
1493 ! **************************************************************************************************
1494  SUBROUTINE dbt_blk_offsets(tensor, ind, blk_offset)
1495  TYPE(dbt_type), INTENT(IN) :: tensor
1496  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1497  INTENT(IN) :: ind
1498  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1499  INTENT(OUT) :: blk_offset
1500 
1501  cpassert(tensor%valid)
1502  blk_offset(:) = get_array_elements(tensor%blk_offsets, ind)
1503  END SUBROUTINE
1504 
1505 ! **************************************************************************************************
1506 !> \brief Generalization of block_get_stored_coordinates for tensors.
1507 !> \author Patrick Seewald
1508 ! **************************************************************************************************
1509  SUBROUTINE dbt_get_stored_coordinates(tensor, ind_nd, processor)
1510  TYPE(dbt_type), INTENT(IN) :: tensor
1511  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1512  INTENT(IN) :: ind_nd
1513  INTEGER, INTENT(OUT) :: processor
1514 
1515  INTEGER(KIND=int_8), DIMENSION(2) :: ind_2d
1516 
1517  ind_2d(:) = get_2d_indices_tensor(tensor%nd_index_blk, ind_nd)
1518  CALL dbt_tas_get_stored_coordinates(tensor%matrix_rep, ind_2d(1), ind_2d(2), processor)
1519  END SUBROUTINE
1520 
1521 ! **************************************************************************************************
1522 !> \author Patrick Seewald
1523 ! **************************************************************************************************
1524  SUBROUTINE dbt_pgrid_create(mp_comm, dims, pgrid, tensor_dims)
1525  CLASS(mp_comm_type), INTENT(IN) :: mp_comm
1526  INTEGER, DIMENSION(:), INTENT(INOUT) :: dims
1527  TYPE(dbt_pgrid_type), INTENT(OUT) :: pgrid
1528  INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: tensor_dims
1529  INTEGER, DIMENSION(:), ALLOCATABLE :: map1_2d, map2_2d
1530  INTEGER :: i, ndims
1531 
1532  ndims = SIZE(dims)
1533 
1534  ALLOCATE (map1_2d(ndims/2))
1535  ALLOCATE (map2_2d(ndims - ndims/2))
1536  map1_2d(:) = (/(i, i=1, SIZE(map1_2d))/)
1537  map2_2d(:) = (/(i, i=SIZE(map1_2d) + 1, SIZE(map1_2d) + SIZE(map2_2d))/)
1538 
1539  CALL dbt_pgrid_create_expert(mp_comm, dims, pgrid, map1_2d, map2_2d, tensor_dims)
1540 
1541  END SUBROUTINE
1542 
1543 ! **************************************************************************************************
1544 !> \brief freeze current split factor such that it is never changed during contraction
1545 !> \author Patrick Seewald
1546 ! **************************************************************************************************
1547  SUBROUTINE dbt_pgrid_set_strict_split(pgrid)
1548  TYPE(dbt_pgrid_type), INTENT(INOUT) :: pgrid
1549  IF (ALLOCATED(pgrid%tas_split_info)) CALL dbt_tas_set_strict_split(pgrid%tas_split_info)
1550  END SUBROUTINE
1551 
1552 ! **************************************************************************************************
1553 !> \brief change dimensions of an existing process grid.
1554 !> \param pgrid process grid to be changed
1555 !> \param pdims new process grid dimensions, should all be set > 0
1556 !> \author Patrick Seewald
1557 ! **************************************************************************************************
1558  SUBROUTINE dbt_pgrid_change_dims(pgrid, pdims)
1559  TYPE(dbt_pgrid_type), INTENT(INOUT) :: pgrid
1560  INTEGER, DIMENSION(:), INTENT(INOUT) :: pdims
1561  TYPE(dbt_pgrid_type) :: pgrid_tmp
1562  INTEGER :: nsplit, dimsplit
1563  INTEGER, DIMENSION(ndims_mapping_row(pgrid%nd_index_grid)) :: map1_2d
1564  INTEGER, DIMENSION(ndims_mapping_column(pgrid%nd_index_grid)) :: map2_2d
1565  TYPe(nd_to_2d_mapping) :: nd_index_grid
1566  INTEGER, DIMENSION(2) :: pdims_2d
1567 
1568  cpassert(all(pdims > 0))
1569  CALL dbt_tas_get_split_info(pgrid%tas_split_info, nsplit=nsplit, split_rowcol=dimsplit)
1570  CALL dbt_get_mapping_info(pgrid%nd_index_grid, map1_2d=map1_2d, map2_2d=map2_2d)
1571  CALL create_nd_to_2d_mapping(nd_index_grid, pdims, map1_2d, map2_2d, base=0, col_major=.false.)
1572  CALL dbt_get_mapping_info(nd_index_grid, dims_2d=pdims_2d)
1573  IF (mod(pdims_2d(dimsplit), nsplit) == 0) THEN
1574  CALL dbt_pgrid_create_expert(pgrid%mp_comm_2d, pdims, pgrid_tmp, map1_2d=map1_2d, map2_2d=map2_2d, &
1575  nsplit=nsplit, dimsplit=dimsplit)
1576  ELSE
1577  CALL dbt_pgrid_create_expert(pgrid%mp_comm_2d, pdims, pgrid_tmp, map1_2d=map1_2d, map2_2d=map2_2d)
1578  END IF
1579  CALL dbt_pgrid_destroy(pgrid)
1580  pgrid = pgrid_tmp
1581  END SUBROUTINE
1582 
1583 ! **************************************************************************************************
1584 !> \brief As block_filter
1585 !> \author Patrick Seewald
1586 ! **************************************************************************************************
1587  SUBROUTINE dbt_filter(tensor, eps)
1588  TYPE(dbt_type), INTENT(INOUT) :: tensor
1589  REAL(dp), INTENT(IN) :: eps
1590 
1591  CALL dbt_tas_filter(tensor%matrix_rep, eps)
1592 
1593  END SUBROUTINE
1594 
1595 ! **************************************************************************************************
1596 !> \brief local number of blocks along dimension idim
1597 !> \author Patrick Seewald
1598 ! **************************************************************************************************
1599  PURE FUNCTION dbt_nblks_local(tensor, idim)
1600  TYPE(dbt_type), INTENT(IN) :: tensor
1601  INTEGER, INTENT(IN) :: idim
1602  INTEGER :: dbt_nblks_local
1603 
1604  IF (idim > ndims_tensor(tensor)) THEN
1605  dbt_nblks_local = 0
1606  ELSE
1607  dbt_nblks_local = tensor%nblks_local(idim)
1608  END IF
1609 
1610  END FUNCTION
1611 
1612 ! **************************************************************************************************
1613 !> \brief total numbers of blocks along dimension idim
1614 !> \author Patrick Seewald
1615 ! **************************************************************************************************
1616  PURE FUNCTION dbt_nblks_total(tensor, idim)
1617  TYPE(dbt_type), INTENT(IN) :: tensor
1618  INTEGER, INTENT(IN) :: idim
1619  INTEGER :: dbt_nblks_total
1620 
1621  IF (idim > ndims_tensor(tensor)) THEN
1622  dbt_nblks_total = 0
1623  ELSE
1624  dbt_nblks_total = tensor%nd_index_blk%dims_nd(idim)
1625  END IF
1626  END FUNCTION
1627 
1628 ! **************************************************************************************************
1629 !> \brief As block_get_info but for tensors
1630 !> \param nblks_total number of blocks along each dimension
1631 !> \param nfull_total number of elements along each dimension
1632 !> \param nblks_local local number of blocks along each dimension
1633 !> \param nfull_local local number of elements along each dimension
1634 !> \param my_ploc process coordinates in process grid
1635 !> \param pdims process grid dimensions
1636 !> \param blks_local_4 local blocks along dimension 4
1637 !> \param proc_dist_4 distribution along dimension 4
1638 !> \param blk_size_4 block sizes along dimension 4
1639 !> \param blk_offset_4 block offsets along dimension 4
1640 !> \param distribution distribution object
1641 !> \param name name of tensor
1642 !> \author Patrick Seewald
1643 ! **************************************************************************************************
1644  SUBROUTINE dbt_get_info(tensor, nblks_total, &
1645  nfull_total, &
1646  nblks_local, &
1647  nfull_local, &
1648  pdims, &
1649  my_ploc, &
1650  blks_local_1, blks_local_2, blks_local_3, blks_local_4, &
1651  proc_dist_1, proc_dist_2, proc_dist_3, proc_dist_4, &
1652  blk_size_1, blk_size_2, blk_size_3, blk_size_4, &
1653  blk_offset_1, blk_offset_2, blk_offset_3, blk_offset_4, &
1654  distribution, &
1655  name)
1656  TYPE(dbt_type), INTENT(IN) :: tensor
1657  INTEGER, INTENT(OUT), OPTIONAL, DIMENSION(ndims_tensor(tensor)) :: nblks_total
1658  INTEGER, INTENT(OUT), OPTIONAL, DIMENSION(ndims_tensor(tensor)) :: nfull_total
1659  INTEGER, INTENT(OUT), OPTIONAL, DIMENSION(ndims_tensor(tensor)) :: nblks_local
1660  INTEGER, INTENT(OUT), OPTIONAL, DIMENSION(ndims_tensor(tensor)) :: nfull_local
1661  INTEGER, INTENT(OUT), OPTIONAL, DIMENSION(ndims_tensor(tensor)) :: my_ploc
1662  INTEGER, INTENT(OUT), OPTIONAL, DIMENSION(ndims_tensor(tensor)) :: pdims
1663  INTEGER, DIMENSION(dbt_nblks_local(tensor, 1)), INTENT(OUT), OPTIONAL :: blks_local_1
1664  INTEGER, DIMENSION(dbt_nblks_total(tensor, 1)), INTENT(OUT), OPTIONAL :: proc_dist_1
1665  INTEGER, DIMENSION(dbt_nblks_total(tensor, 1)), INTENT(OUT), OPTIONAL :: blk_size_1
1666  INTEGER, DIMENSION(dbt_nblks_total(tensor, 1)), INTENT(OUT), OPTIONAL :: blk_offset_1
1667  INTEGER, DIMENSION(dbt_nblks_local(tensor, 2)), INTENT(OUT), OPTIONAL :: blks_local_2
1668  INTEGER, DIMENSION(dbt_nblks_total(tensor, 2)), INTENT(OUT), OPTIONAL :: proc_dist_2
1669  INTEGER, DIMENSION(dbt_nblks_total(tensor, 2)), INTENT(OUT), OPTIONAL :: blk_size_2
1670  INTEGER, DIMENSION(dbt_nblks_total(tensor, 2)), INTENT(OUT), OPTIONAL :: blk_offset_2
1671  INTEGER, DIMENSION(dbt_nblks_local(tensor, 3)), INTENT(OUT), OPTIONAL :: blks_local_3
1672  INTEGER, DIMENSION(dbt_nblks_total(tensor, 3)), INTENT(OUT), OPTIONAL :: proc_dist_3
1673  INTEGER, DIMENSION(dbt_nblks_total(tensor, 3)), INTENT(OUT), OPTIONAL :: blk_size_3
1674  INTEGER, DIMENSION(dbt_nblks_total(tensor, 3)), INTENT(OUT), OPTIONAL :: blk_offset_3
1675  INTEGER, DIMENSION(dbt_nblks_local(tensor, 4)), INTENT(OUT), OPTIONAL :: blks_local_4
1676  INTEGER, DIMENSION(dbt_nblks_total(tensor, 4)), INTENT(OUT), OPTIONAL :: proc_dist_4
1677  INTEGER, DIMENSION(dbt_nblks_total(tensor, 4)), INTENT(OUT), OPTIONAL :: blk_size_4
1678  INTEGER, DIMENSION(dbt_nblks_total(tensor, 4)), INTENT(OUT), OPTIONAL :: blk_offset_4
1679  TYPE(dbt_distribution_type), INTENT(OUT), OPTIONAL :: distribution
1680  CHARACTER(len=*), INTENT(OUT), OPTIONAL :: name
1681  INTEGER, DIMENSION(ndims_tensor(tensor)) :: pdims_tmp, my_ploc_tmp
1682 
1683  IF (PRESENT(nblks_total)) CALL dbt_get_mapping_info(tensor%nd_index_blk, dims_nd=nblks_total)
1684  IF (PRESENT(nfull_total)) CALL dbt_get_mapping_info(tensor%nd_index, dims_nd=nfull_total)
1685  IF (PRESENT(nblks_local)) nblks_local(:) = tensor%nblks_local
1686  IF (PRESENT(nfull_local)) nfull_local(:) = tensor%nfull_local
1687 
1688  IF (PRESENT(my_ploc) .OR. PRESENT(pdims)) CALL mp_environ_pgrid(tensor%pgrid, pdims_tmp, my_ploc_tmp)
1689  IF (PRESENT(my_ploc)) my_ploc = my_ploc_tmp
1690  IF (PRESENT(pdims)) pdims = pdims_tmp
1691 
1692  IF (1 <= ndims_tensor(tensor)) THEN
1693  IF (PRESENT(blks_local_1)) CALL get_ith_array(tensor%blks_local, 1, &
1694  dbt_nblks_local(tensor, 1), &
1695  blks_local_1)
1696  IF (PRESENT(proc_dist_1)) CALL get_ith_array(tensor%nd_dist, 1, &
1697  dbt_nblks_total(tensor, 1), &
1698  proc_dist_1)
1699  IF (PRESENT(blk_size_1)) CALL get_ith_array(tensor%blk_sizes, 1, &
1700  dbt_nblks_total(tensor, 1), &
1701  blk_size_1)
1702  IF (PRESENT(blk_offset_1)) CALL get_ith_array(tensor%blk_offsets, 1, &
1703  dbt_nblks_total(tensor, 1), &
1704  blk_offset_1)
1705  END IF
1706  IF (2 <= ndims_tensor(tensor)) THEN
1707  IF (PRESENT(blks_local_2)) CALL get_ith_array(tensor%blks_local, 2, &
1708  dbt_nblks_local(tensor, 2), &
1709  blks_local_2)
1710  IF (PRESENT(proc_dist_2)) CALL get_ith_array(tensor%nd_dist, 2, &
1711  dbt_nblks_total(tensor, 2), &
1712  proc_dist_2)
1713  IF (PRESENT(blk_size_2)) CALL get_ith_array(tensor%blk_sizes, 2, &
1714  dbt_nblks_total(tensor, 2), &
1715  blk_size_2)
1716  IF (PRESENT(blk_offset_2)) CALL get_ith_array(tensor%blk_offsets, 2, &
1717  dbt_nblks_total(tensor, 2), &
1718  blk_offset_2)
1719  END IF
1720  IF (3 <= ndims_tensor(tensor)) THEN
1721  IF (PRESENT(blks_local_3)) CALL get_ith_array(tensor%blks_local, 3, &
1722  dbt_nblks_local(tensor, 3), &
1723  blks_local_3)
1724  IF (PRESENT(proc_dist_3)) CALL get_ith_array(tensor%nd_dist, 3, &
1725  dbt_nblks_total(tensor, 3), &
1726  proc_dist_3)
1727  IF (PRESENT(blk_size_3)) CALL get_ith_array(tensor%blk_sizes, 3, &
1728  dbt_nblks_total(tensor, 3), &
1729  blk_size_3)
1730  IF (PRESENT(blk_offset_3)) CALL get_ith_array(tensor%blk_offsets, 3, &
1731  dbt_nblks_total(tensor, 3), &
1732  blk_offset_3)
1733  END IF
1734  IF (4 <= ndims_tensor(tensor)) THEN
1735  IF (PRESENT(blks_local_4)) CALL get_ith_array(tensor%blks_local, 4, &
1736  dbt_nblks_local(tensor, 4), &
1737  blks_local_4)
1738  IF (PRESENT(proc_dist_4)) CALL get_ith_array(tensor%nd_dist, 4, &
1739  dbt_nblks_total(tensor, 4), &
1740  proc_dist_4)
1741  IF (PRESENT(blk_size_4)) CALL get_ith_array(tensor%blk_sizes, 4, &
1742  dbt_nblks_total(tensor, 4), &
1743  blk_size_4)
1744  IF (PRESENT(blk_offset_4)) CALL get_ith_array(tensor%blk_offsets, 4, &
1745  dbt_nblks_total(tensor, 4), &
1746  blk_offset_4)
1747  END IF
1748 
1749  IF (PRESENT(distribution)) distribution = dbt_distribution(tensor)
1750  IF (PRESENT(name)) name = tensor%name
1751 
1752  END SUBROUTINE
1753 
1754 ! **************************************************************************************************
1755 !> \brief As block_get_num_blocks: get number of local blocks
1756 !> \author Patrick Seewald
1757 ! **************************************************************************************************
1758  PURE FUNCTION dbt_get_num_blocks(tensor) RESULT(num_blocks)
1759  TYPE(dbt_type), INTENT(IN) :: tensor
1760  INTEGER :: num_blocks
1761  num_blocks = dbt_tas_get_num_blocks(tensor%matrix_rep)
1762  END FUNCTION
1763 
1764 ! **************************************************************************************************
1765 !> \brief Get total number of blocks
1766 !> \author Patrick Seewald
1767 ! **************************************************************************************************
1768  FUNCTION dbt_get_num_blocks_total(tensor) RESULT(num_blocks)
1769  TYPE(dbt_type), INTENT(IN) :: tensor
1770  INTEGER(KIND=int_8) :: num_blocks
1771  num_blocks = dbt_tas_get_num_blocks_total(tensor%matrix_rep)
1772  END FUNCTION
1773 
1774 ! **************************************************************************************************
1775 !> \brief Clear tensor (s.t. it does not contain any blocks)
1776 !> \author Patrick Seewald
1777 ! **************************************************************************************************
1778  SUBROUTINE dbt_clear(tensor)
1779  TYPE(dbt_type), INTENT(INOUT) :: tensor
1780 
1781  CALL dbt_tas_clear(tensor%matrix_rep)
1782  END SUBROUTINE
1783 
1784 ! **************************************************************************************************
1785 !> \brief Finalize tensor, as block_finalize. This should be taken care of internally in DBT
1786 !> tensors, there should not be any need to call this routine outside of DBT tensors.
1787 !> \author Patrick Seewald
1788 ! **************************************************************************************************
1789  SUBROUTINE dbt_finalize(tensor)
1790  TYPE(dbt_type), INTENT(INOUT) :: tensor
1791  CALL dbt_tas_finalize(tensor%matrix_rep)
1792  END SUBROUTINE
1793 
1794 ! **************************************************************************************************
1795 !> \brief as block_scale
1796 !> \author Patrick Seewald
1797 ! **************************************************************************************************
1798  SUBROUTINE dbt_scale(tensor, alpha)
1799  TYPE(dbt_type), INTENT(INOUT) :: tensor
1800  REAL(dp), INTENT(IN) :: alpha
1801  CALL dbm_scale(tensor%matrix_rep%matrix, alpha)
1802  END SUBROUTINE
1803 
1804 ! **************************************************************************************************
1805 !> \author Patrick Seewald
1806 ! **************************************************************************************************
1807  PURE FUNCTION dbt_get_nze(tensor)
1808  TYPE(dbt_type), INTENT(IN) :: tensor
1809  INTEGER :: dbt_get_nze
1810  dbt_get_nze = dbt_tas_get_nze(tensor%matrix_rep)
1811  END FUNCTION
1812 
1813 ! **************************************************************************************************
1814 !> \author Patrick Seewald
1815 ! **************************************************************************************************
1816  FUNCTION dbt_get_nze_total(tensor)
1817  TYPE(dbt_type), INTENT(IN) :: tensor
1818  INTEGER(KIND=int_8) :: dbt_get_nze_total
1819  dbt_get_nze_total = dbt_tas_get_nze_total(tensor%matrix_rep)
1820  END FUNCTION
1821 
1822 ! **************************************************************************************************
1823 !> \brief block size of block with index ind along dimension idim
1824 !> \author Patrick Seewald
1825 ! **************************************************************************************************
1826  PURE FUNCTION dbt_blk_size(tensor, ind, idim)
1827  TYPE(dbt_type), INTENT(IN) :: tensor
1828  INTEGER, DIMENSION(ndims_tensor(tensor)), &
1829  INTENT(IN) :: ind
1830  INTEGER, INTENT(IN) :: idim
1831  INTEGER, DIMENSION(ndims_tensor(tensor)) :: blk_size
1832  INTEGER :: dbt_blk_size
1833 
1834  IF (idim > ndims_tensor(tensor)) THEN
1835  dbt_blk_size = 0
1836  ELSE
1837  blk_size(:) = get_array_elements(tensor%blk_sizes, ind)
1838  dbt_blk_size = blk_size(idim)
1839  END IF
1840  END FUNCTION
1841 
1842 ! **************************************************************************************************
1843 !> \brief returns an estimate of maximum number of local blocks in tensor
1844 !> (irrespective of the actual number of currently present blocks)
1845 !> this estimate is based on the following assumption: tensor data is dense and
1846 !> load balancing is within a factor of 2
1847 !> \author Patrick Seewald
1848 ! **************************************************************************************************
1849  PURE FUNCTION dbt_max_nblks_local(tensor) RESULT(blk_count)
1850  TYPE(dbt_type), INTENT(IN) :: tensor
1851  INTEGER :: blk_count, nproc
1852  INTEGER, DIMENSION(ndims_tensor(tensor)) :: bdims
1853  INTEGER(int_8) :: blk_count_total
1854  INTEGER, PARAMETER :: max_load_imbalance = 2
1855 
1856  CALL dbt_get_mapping_info(tensor%nd_index_blk, dims_nd=bdims)
1857 
1858  blk_count_total = product(int(bdims, int_8))
1859 
1860  ! can not call an MPI routine due to PURE
1861  nproc = tensor%pgrid%nproc
1862 
1863  blk_count = int(blk_count_total/nproc*max_load_imbalance)
1864 
1865  END FUNCTION
1866 
1867 ! **************************************************************************************************
1868 !> \brief get a load-balanced and randomized distribution along one tensor dimension
1869 !> \param nblk number of blocks (along one tensor dimension)
1870 !> \param nproc number of processes (along one process grid dimension)
1871 !> \param blk_size block sizes
1872 !> \param dist distribution
1873 !> \author Patrick Seewald
1874 ! **************************************************************************************************
1875  SUBROUTINE dbt_default_distvec(nblk, nproc, blk_size, dist)
1876  INTEGER, INTENT(IN) :: nblk
1877  INTEGER, INTENT(IN) :: nproc
1878  INTEGER, DIMENSION(nblk), INTENT(IN) :: blk_size
1879  INTEGER, DIMENSION(nblk), INTENT(OUT) :: dist
1880 
1881  CALL dbt_tas_default_distvec(nblk, nproc, blk_size, dist)
1882  END SUBROUTINE
1883 
1884 ! **************************************************************************************************
1885 !> \author Patrick Seewald
1886 ! **************************************************************************************************
1887  SUBROUTINE dbt_copy_contraction_storage(tensor_in, tensor_out)
1888  TYPE(dbt_type), INTENT(IN) :: tensor_in
1889  TYPE(dbt_type), INTENT(INOUT) :: tensor_out
1890  TYPE(dbt_contraction_storage), ALLOCATABLE :: tensor_storage_tmp
1891  TYPE(dbt_tas_mm_storage), ALLOCATABLE :: tas_storage_tmp
1892 
1893  IF (tensor_in%matrix_rep%do_batched > 0) THEN
1894  ALLOCATE (tas_storage_tmp, source=tensor_in%matrix_rep%mm_storage)
1895  ! transfer data for batched contraction
1896  IF (ALLOCATED(tensor_out%matrix_rep%mm_storage)) DEALLOCATE (tensor_out%matrix_rep%mm_storage)
1897  CALL move_alloc(tas_storage_tmp, tensor_out%matrix_rep%mm_storage)
1898  END IF
1899  CALL dbt_tas_set_batched_state(tensor_out%matrix_rep, state=tensor_in%matrix_rep%do_batched, &
1900  opt_grid=tensor_in%matrix_rep%has_opt_pgrid)
1901  IF (ALLOCATED(tensor_in%contraction_storage)) THEN
1902  ALLOCATE (tensor_storage_tmp, source=tensor_in%contraction_storage)
1903  END IF
1904  IF (ALLOCATED(tensor_out%contraction_storage)) DEALLOCATE (tensor_out%contraction_storage)
1905  IF (ALLOCATED(tensor_storage_tmp)) CALL move_alloc(tensor_storage_tmp, tensor_out%contraction_storage)
1906 
1907  END SUBROUTINE
1908 
1909  END MODULE
tensor
struct tensor_ tensor
dbm_api
Definition: dbm_api.F:8
dbm_api::dbm_scale
subroutine, public dbm_scale(matrix, alpha)
Multiplies all entries in the given matrix by the given factor alpha.
Definition: dbm_api.F:631
dbt_allocate_wrap
Wrapper for allocating, copying and reshaping arrays.
Definition: dbt_allocate_wrap.F:15
dbt_array_list_methods
Representation of arbitrary number of 1d integer arrays with arbitrary sizes. This is needed for gene...
Definition: dbt_array_list_methods.F:14
dbt_array_list_methods::array_eq_i
pure logical function, public array_eq_i(arr1, arr2)
check whether two arrays are equal
Definition: dbt_array_list_methods.F:369
dbt_array_list_methods::sum_of_arrays
integer function, dimension(:), allocatable, public sum_of_arrays(list)
sum of all elements for each array stored in list
Definition: dbt_array_list_methods.F:289
dbt_array_list_methods::get_arrays
subroutine, public get_arrays(list, data_1, data_2, data_3, data_4, i_selected)
Get all arrays contained in list.
Definition: dbt_array_list_methods.F:197
dbt_array_list_methods::create_array_list
subroutine, public create_array_list(list, ndata, data_1, data_2, data_3, data_4)
collects any number of arrays of different sizes into a single array (listcol_data),...
Definition: dbt_array_list_methods.F:94
dbt_array_list_methods::destroy_array_list
subroutine, public destroy_array_list(list)
destroy array list.
Definition: dbt_array_list_methods.F:185
dbt_array_list_methods::sizes_of_arrays
integer function, dimension(:), allocatable, public sizes_of_arrays(list)
sizes of arrays stored in list
Definition: dbt_array_list_methods.F:272
dbt_array_list_methods::array_offsets
subroutine, public array_offsets(list_in, list_out)
partial sums of array elements.
Definition: dbt_array_list_methods.F:307
dbt_array_list_methods::get_array_elements
pure integer function, dimension(number_of_arrays(list)), public get_array_elements(list, indices)
Get an element for each array.
Definition: dbt_array_list_methods.F:72
dbt_array_list_methods::get_ith_array
subroutine, public get_ith_array(list, i, array_size, array)
get ith array
Definition: dbt_array_list_methods.F:237
dbt_array_list_methods::array_sublist
type(array_list) function, public array_sublist(list, i_selected)
extract a subset of arrays
Definition: dbt_array_list_methods.F:154
dbt_index
tensor index and mapping to DBM index
Definition: dbt_index.F:12
dbt_index::create_nd_to_2d_mapping
subroutine, public create_nd_to_2d_mapping(map, dims, map1_2d, map2_2d, base, col_major)
Create all data needed to quickly map between nd index and 2d index.
Definition: dbt_index.F:73
dbt_index::get_nd_indices_pgrid
pure integer function, dimension(map%ndim_nd), public get_nd_indices_pgrid(map, ind_in)
transform 2d index to nd index, using info from index mapping.
Definition: dbt_index.F:396
dbt_index::ndims_mapping_row
pure integer function, public ndims_mapping_row(map)
how many tensor dimensions are mapped to matrix row
Definition: dbt_index.F:141
dbt_index::get_2d_indices_tensor
pure integer(kind=int_8) function, dimension(2), public get_2d_indices_tensor(map, ind_in)
transform nd index to 2d index, using info from index mapping.
Definition: dbt_index.F:318
dbt_index::ndims_mapping
pure integer function, public ndims_mapping(map)
Definition: dbt_index.F:130
dbt_index::combine_tensor_index
pure integer(kind=int_8) function, public combine_tensor_index(ind_in, dims)
transform nd index to flat index
Definition: dbt_index.F:235
dbt_index::dbt_get_mapping_info
pure subroutine, public dbt_get_mapping_info(map, ndim_nd, ndim1_2d, ndim2_2d, dims_2d_i8, dims_2d, dims_nd, dims1_2d, dims2_2d, map1_2d, map2_2d, map_nd, base, col_major)
get mapping info
Definition: dbt_index.F:176
dbt_index::combine_pgrid_index
pure integer function, public combine_pgrid_index(ind_in, dims)
transform nd index to flat index
Definition: dbt_index.F:254
dbt_index::ndims_mapping_column
pure integer function, public ndims_mapping_column(map)
how many tensor dimensions are mapped to matrix column
Definition: dbt_index.F:151
dbt_index::destroy_nd_to_2d_mapping
subroutine, public destroy_nd_to_2d_mapping(map)
Definition: dbt_index.F:115
dbt_index::split_tensor_index
pure integer function, dimension(size(dims)), public split_tensor_index(ind_in, dims)
transform flat index to nd index
Definition: dbt_index.F:273
dbt_index::split_pgrid_index
pure integer function, dimension(size(dims)), public split_pgrid_index(ind_in, dims)
transform flat index to nd index
Definition: dbt_index.F:296
dbt_tas_base
Tall-and-skinny matrices: base routines similar to DBM API, mostly wrappers around existing DBM routi...
Definition: dbt_tas_base.F:13
dbt_tas_base::dbt_tas_get_nze_total
integer(kind=int_8) function, public dbt_tas_get_nze_total(matrix)
Get total number of non-zero elements.
Definition: dbt_tas_base.F:958
dbt_tas_base::dbt_tas_distribution_destroy
subroutine, public dbt_tas_distribution_destroy(dist)
...
Definition: dbt_tas_base.F:433
dbt_tas_base::dbt_tas_get_stored_coordinates
subroutine, public dbt_tas_get_stored_coordinates(matrix, row, column, processor)
As dbt_get_stored_coordinates.
Definition: dbt_tas_base.F:462
dbt_tas_base::dbt_tas_get_nze
pure integer function, public dbt_tas_get_nze(matrix)
As dbt_get_nze: get number of local non-zero elements.
Definition: dbt_tas_base.F:944
dbt_tas_base::dbt_tas_get_info
subroutine, public dbt_tas_get_info(matrix, nblkrows_total, nblkcols_total, local_rows, local_cols, proc_row_dist, proc_col_dist, row_blk_size, col_blk_size, distribution, name)
...
Definition: dbt_tas_base.F:999
dbt_tas_base::dbt_tas_finalize
subroutine, public dbt_tas_finalize(matrix)
...
Definition: dbt_tas_base.F:327
dbt_tas_base::dbt_tas_distribution_new
subroutine, public dbt_tas_distribution_new(dist, mp_comm, row_dist, col_dist, split_info, nosplit)
create new distribution. Exactly like dbm_distribution_new but with custom types for row_dist and col...
Definition: dbt_tas_base.F:346
dbt_tas_base::dbt_tas_filter
subroutine, public dbt_tas_filter(matrix, eps)
As dbm_filter.
Definition: dbt_tas_base.F:1128
dbt_tas_base::dbt_tas_clear
subroutine, public dbt_tas_clear(matrix)
Clear matrix (erase all data)
Definition: dbt_tas_base.F:974
dbt_tas_base::dbt_tas_destroy
subroutine, public dbt_tas_destroy(matrix)
...
Definition: dbt_tas_base.F:233
dbt_tas_base::dbt_tas_get_num_blocks_total
integer(kind=int_8) function, public dbt_tas_get_num_blocks_total(matrix)
get total number of blocks
Definition: dbt_tas_base.F:926
dbt_tas_base::dbt_tas_get_num_blocks
pure integer function, public dbt_tas_get_num_blocks(matrix)
As dbt_get_num_blocks: get number of local blocks.
Definition: dbt_tas_base.F:913
dbt_tas_global
Global data (distribution and block sizes) for tall-and-skinny matrices For very sparse matrices with...
Definition: dbt_tas_global.F:17
dbt_tas_global::dbt_tas_default_distvec
subroutine, public dbt_tas_default_distvec(nblk, nproc, blk_size, dist)
get a load-balanced and randomized distribution along one tensor dimension
Definition: dbt_tas_global.F:485
dbt_tas_mm
Matrix multiplication for tall-and-skinny matrices. This uses the k-split (non-recursive) CARMA algor...
Definition: dbt_tas_mm.F:18
dbt_tas_mm::dbt_tas_set_batched_state
subroutine, public dbt_tas_set_batched_state(matrix, state, opt_grid)
set state flags during batched multiplication
Definition: dbt_tas_mm.F:1699
dbt_tas_split
methods to split tall-and-skinny matrices along longest dimension. Basically, we are splitting proces...
Definition: dbt_tas_split.F:13
dbt_tas_split::dbt_tas_release_info
subroutine, public dbt_tas_release_info(split_info)
...
Definition: dbt_tas_split.F:457
dbt_tas_split::dbt_tas_get_split_info
subroutine, public dbt_tas_get_split_info(info, mp_comm, nsplit, igroup, mp_comm_group, split_rowcol, pgrid_offset)
Get info on split.
Definition: dbt_tas_split.F:427
dbt_tas_split::dbt_tas_create_split
subroutine, public dbt_tas_create_split(split_info, mp_comm, split_rowcol, nsplit, own_comm, opt_nsplit)
Split Cartesian process grid using a default split heuristic.
Definition: dbt_tas_split.F:333
dbt_tas_split::dbt_tas_create_split_rows_or_cols
subroutine, public dbt_tas_create_split_rows_or_cols(split_info, mp_comm, ngroup, igroup, split_rowcol, own_comm)
split mpi grid by rows or columns
Definition: dbt_tas_split.F:72
dbt_tas_split::dbt_tas_info_hold
subroutine, public dbt_tas_info_hold(split_info)
...
Definition: dbt_tas_split.F:492
dbt_tas_split::dbt_tas_set_strict_split
subroutine, public dbt_tas_set_strict_split(info)
freeze current split factor such that it is never changed during multiplication
Definition: dbt_tas_split.F:713
dbt_tas_types
DBT tall-and-skinny base types. Mostly wrappers around existing DBM routines.
Definition: dbt_tas_types.F:13
dbt_types
DBT tensor framework for block-sparse tensor contraction: Types and create/destroy routines.
Definition: dbt_types.F:12
dbt_types::dbt_pgrid_destroy
subroutine, public dbt_pgrid_destroy(pgrid, keep_comm)
destroy process grid
Definition: dbt_types.F:905
dbt_types::dbt_distribution_new
subroutine, public dbt_distribution_new(dist, pgrid, nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4)
Create a tensor distribution.
Definition: dbt_types.F:886
dbt_types::blk_dims_tensor
subroutine, public blk_dims_tensor(tensor, dims)
tensor block dimensions
Definition: dbt_types.F:1466
dbt_types::dims_tensor
subroutine, public dims_tensor(tensor, dims)
tensor dimensions
Definition: dbt_types.F:1238
dbt_types::dbt_copy_contraction_storage
subroutine, public dbt_copy_contraction_storage(tensor_in, tensor_out)
Definition: dbt_types.F:1888
dbt_types::dbt_nd_mp_comm
type(dbt_pgrid_type) function, public dbt_nd_mp_comm(comm_2d, map1_2d, map2_2d, dims_nd, dims1_nd, dims2_nd, pdims_2d, tdims, nsplit, dimsplit)
Create a default nd process topology that is consistent with a given 2d topology. Purpose: a nd tenso...
Definition: dbt_types.F:653
dbt_types::dbt_blk_sizes
subroutine, public dbt_blk_sizes(tensor, ind, blk_size)
Size of tensor block.
Definition: dbt_types.F:1479
dbt_types::dbt_destroy
subroutine, public dbt_destroy(tensor)
Destroy a tensor.
Definition: dbt_types.F:1410
dbt_types::dbt_max_nblks_local
pure integer function, public dbt_max_nblks_local(tensor)
returns an estimate of maximum number of local blocks in tensor (irrespective of the actual number of...
Definition: dbt_types.F:1850
dbt_types::dbt_mp_dims_create
recursive subroutine, public dbt_mp_dims_create(nodes, dims, tensor_dims, lb_ratio)
Create process grid dimensions corresponding to one dimension of the matrix representation of a tenso...
Definition: dbt_types.F:498
dbt_types::dbt_get_info
subroutine, public dbt_get_info(tensor, nblks_total, nfull_total, nblks_local, nfull_local, pdims, my_ploc, blks_local_1, blks_local_2, blks_local_3, blks_local_4, proc_dist_1, proc_dist_2, proc_dist_3, proc_dist_4, blk_size_1, blk_size_2, blk_size_3, blk_size_4, blk_offset_1, blk_offset_2, blk_offset_3, blk_offset_4, distribution, name)
As block_get_info but for tensors.
Definition: dbt_types.F:1656
dbt_types::dbt_distribution_new_expert
subroutine, public dbt_distribution_new_expert(dist, pgrid, map1_2d, map2_2d, nd_dist_1, nd_dist_2, nd_dist_3, nd_dist_4, own_comm)
Create a tensor distribution.
Definition: dbt_types.F:787
dbt_types::dbt_distribution
type(dbt_distribution_type) function, public dbt_distribution(tensor)
get distribution from tensor
Definition: dbt_types.F:980
dbt_types::ndims_tensor
pure integer function, public ndims_tensor(tensor)
tensor rank
Definition: dbt_types.F:1227
dbt_types::dbt_pgrid_set_strict_split
subroutine, public dbt_pgrid_set_strict_split(pgrid)
freeze current split factor such that it is never changed during contraction
Definition: dbt_types.F:1548
dbt_types::dbt_nblks_total
pure integer function, public dbt_nblks_total(tensor, idim)
total numbers of blocks along dimension idim
Definition: dbt_types.F:1617
dbt_types::dbt_blk_size
pure integer function, public dbt_blk_size(tensor, ind, idim)
block size of block with index ind along dimension idim
Definition: dbt_types.F:1827
dbt_types::dbt_get_num_blocks
pure integer function, public dbt_get_num_blocks(tensor)
As block_get_num_blocks: get number of local blocks.
Definition: dbt_types.F:1759
dbt_types::dbt_default_distvec
subroutine, public dbt_default_distvec(nblk, nproc, blk_size, dist)
get a load-balanced and randomized distribution along one tensor dimension
Definition: dbt_types.F:1876
dbt_types::dbt_hold
subroutine, public dbt_hold(tensor)
reference counting for tensors (only needed for communicator handle that must be freed when no longer...
Definition: dbt_types.F:1188
dbt_types::dbt_pgrid_create_expert
subroutine, public dbt_pgrid_create_expert(mp_comm, dims, pgrid, map1_2d, map2_2d, tensor_dims, nsplit, dimsplit)
Create an n-dimensional process grid. We can not use a n-dimensional MPI cartesian grid for tensors s...
Definition: dbt_types.F:594
dbt_types::dbt_clear
subroutine, public dbt_clear(tensor)
Clear tensor (s.t. it does not contain any blocks)
Definition: dbt_types.F:1779
dbt_types::dbt_finalize
subroutine, public dbt_finalize(tensor)
Finalize tensor, as block_finalize. This should be taken care of internally in DBT tensors,...
Definition: dbt_types.F:1790
dbt_types::dbt_get_nze_total
integer(kind=int_8) function, public dbt_get_nze_total(tensor)
Definition: dbt_types.F:1817
dbt_types::dbt_nblks_local
pure integer function, public dbt_nblks_local(tensor, idim)
local number of blocks along dimension idim
Definition: dbt_types.F:1600
dbt_types::mp_environ_pgrid
subroutine, public mp_environ_pgrid(pgrid, dims, task_coor)
as mp_environ but for special pgrid type
Definition: dbt_types.F:768
dbt_types::dbt_get_stored_coordinates
subroutine, public dbt_get_stored_coordinates(tensor, ind_nd, processor)
Generalization of block_get_stored_coordinates for tensors.
Definition: dbt_types.F:1510
dbt_types::dbt_get_nze
pure integer function, public dbt_get_nze(tensor)
Definition: dbt_types.F:1808
dbt_types::dbt_pgrid_create
subroutine, public dbt_pgrid_create(mp_comm, dims, pgrid, tensor_dims)
Definition: dbt_types.F:1525
dbt_types::dbt_get_num_blocks_total
integer(kind=int_8) function, public dbt_get_num_blocks_total(tensor)
Get total number of blocks.
Definition: dbt_types.F:1769
dbt_types::ndims_matrix_row
pure integer(int_8) function, public ndims_matrix_row(tensor)
how many tensor dimensions are mapped to matrix row
Definition: dbt_types.F:1204
dbt_types::dbt_pgrid_change_dims
subroutine, public dbt_pgrid_change_dims(pgrid, pdims)
change dimensions of an existing process grid.
Definition: dbt_types.F:1559
dbt_types::ndims_matrix_column
pure integer(int_8) function, public ndims_matrix_column(tensor)
how many tensor dimensions are mapped to matrix column
Definition: dbt_types.F:1216
dbt_types::dbt_nd_mp_free
subroutine, public dbt_nd_mp_free(mp_comm)
Release the MPI communicator.
Definition: dbt_types.F:733
dbt_types::dbt_blk_offsets
subroutine, public dbt_blk_offsets(tensor, ind, blk_offset)
offset of tensor block
Definition: dbt_types.F:1495
dbt_types::dbt_filter
subroutine, public dbt_filter(tensor, eps)
As block_filter.
Definition: dbt_types.F:1588
dbt_types::dbt_distribution_destroy
subroutine, public dbt_distribution_destroy(dist)
Destroy tensor distribution.
Definition: dbt_types.F:926
dbt_types::dbt_scale
subroutine, public dbt_scale(tensor, alpha)
as block_scale
Definition: dbt_types.F:1799
kinds
Defines the basic variable types.
Definition: kinds.F:23
kinds::int_8
integer, parameter, public int_8
Definition: kinds.F:54
kinds::dp
integer, parameter, public dp
Definition: kinds.F:34
kinds::default_string_length
integer, parameter, public default_string_length
Definition: kinds.F:57
message_passing
Interface to the message passing library MPI.
Definition: message_passing.F:23
message_passing::mp_dims_create
subroutine, public mp_dims_create(nodes, dims)
wrapper to MPI_Dims_create
Definition: message_passing.F:2132
util
All kind of helpful little routines.
Definition: util.F:14