Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief DBT tensor framework for block-sparse tensor contraction.
10 : !> Representation of n-rank tensors as DBT tall-and-skinny matrices.
11 : !> Support for arbitrary redistribution between different representations.
12 : !> Support for arbitrary tensor contractions
13 : !> \todo implement checks and error messages
14 : !> \author Patrick Seewald
15 : ! **************************************************************************************************
16 : MODULE dbt_methods
17 : #:include "dbt_macros.fypp"
18 : #:set maxdim = maxrank
19 : #:set ndims = range(2,maxdim+1)
20 :
21 : USE cp_dbcsr_api, ONLY: &
22 : dbcsr_type, dbcsr_release, &
23 : dbcsr_iterator_type, dbcsr_iterator_start, dbcsr_iterator_blocks_left, dbcsr_iterator_next_block, &
24 : dbcsr_has_symmetry, dbcsr_desymmetrize, dbcsr_put_block, dbcsr_clear, dbcsr_iterator_stop
25 : USE dbt_allocate_wrap, ONLY: &
26 : allocate_any
27 : USE dbt_array_list_methods, ONLY: &
28 : get_arrays, reorder_arrays, get_ith_array, array_list, array_sublist, check_equal, array_eq_i, &
29 : create_array_list, destroy_array_list, sizes_of_arrays
30 : USE dbm_api, ONLY: &
31 : dbm_clear
32 : USE dbt_tas_types, ONLY: &
33 : dbt_tas_split_info
34 : USE dbt_tas_base, ONLY: &
35 : dbt_tas_copy, dbt_tas_finalize, dbt_tas_get_info, dbt_tas_info
36 : USE dbt_tas_mm, ONLY: &
37 : dbt_tas_multiply, dbt_tas_batched_mm_init, dbt_tas_batched_mm_finalize, &
38 : dbt_tas_batched_mm_complete, dbt_tas_set_batched_state
39 : USE dbt_block, ONLY: &
40 : dbt_iterator_type, dbt_get_block, dbt_put_block, dbt_iterator_start, &
41 : dbt_iterator_blocks_left, dbt_iterator_stop, dbt_iterator_next_block, &
42 : ndims_iterator, dbt_reserve_blocks, block_nd, destroy_block, checker_tr
43 : USE dbt_index, ONLY: &
44 : dbt_get_mapping_info, nd_to_2d_mapping, dbt_inverse_order, permute_index, get_nd_indices_tensor, &
45 : ndims_mapping_row, ndims_mapping_column, ndims_mapping
46 : USE dbt_types, ONLY: &
47 : dbt_create, dbt_type, ndims_tensor, dims_tensor, &
48 : dbt_distribution_type, dbt_distribution, dbt_nd_mp_comm, dbt_destroy, &
49 : dbt_distribution_destroy, dbt_distribution_new_expert, dbt_get_stored_coordinates, &
50 : blk_dims_tensor, dbt_hold, dbt_pgrid_type, mp_environ_pgrid, dbt_filter, &
51 : dbt_clear, dbt_finalize, dbt_get_num_blocks, dbt_scale, &
52 : dbt_get_num_blocks_total, dbt_get_info, ndims_matrix_row, ndims_matrix_column, &
53 : dbt_max_nblks_local, dbt_default_distvec, dbt_contraction_storage, dbt_nblks_total, &
54 : dbt_distribution_new, dbt_copy_contraction_storage, dbt_pgrid_destroy, dbt_get_nze
55 : USE kinds, ONLY: &
56 : dp, default_string_length, int_8, dp
57 : USE message_passing, ONLY: &
58 : mp_cart_type
59 : USE util, ONLY: &
60 : sort
61 : USE dbt_reshape_ops, ONLY: &
62 : dbt_reshape, dbt_reshape_release
63 : USE dbt_tas_split, ONLY: &
64 : dbt_tas_mp_comm, rowsplit, colsplit, dbt_tas_info_hold, dbt_tas_release_info, default_nsplit_accept_ratio, &
65 : default_pdims_accept_ratio, dbt_tas_create_split
66 : USE dbt_split, ONLY: &
67 : dbt_split_copyback, dbt_make_compatible_blocks, dbt_crop
68 : USE dbt_io, ONLY: &
69 : dbt_write_tensor_info, dbt_write_tensor_dist, prep_output_unit, dbt_write_split_info
70 : USE message_passing, ONLY: mp_comm_type
71 :
72 : #include "../base/base_uses.f90"
73 :
74 : IMPLICIT NONE
75 : PRIVATE
76 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_methods'
77 :
78 : PUBLIC :: &
79 : dbt_contract, &
80 : dbt_copy, &
81 : dbt_get_block, &
82 : dbt_get_stored_coordinates, &
83 : dbt_inverse_order, &
84 : dbt_iterator_blocks_left, &
85 : dbt_iterator_next_block, &
86 : dbt_iterator_start, &
87 : dbt_iterator_stop, &
88 : dbt_iterator_type, &
89 : dbt_put_block, &
90 : dbt_reserve_blocks, &
91 : dbt_copy_matrix_to_tensor, &
92 : dbt_copy_tensor_to_matrix, &
93 : dbt_batched_contract_init, &
94 : dbt_batched_contract_finalize, &
95 : dbt_library_init, &
96 : dbt_library_finalize, &
97 : dbt_library_print_stats, &
98 : dbt_clear_mempools
99 :
100 : ! Per contraction (tensor names): remapped and cropped tensors, i.e., data copied
101 : ! before the multiplication (local to the rank)
102 : TYPE contract_stats_type
103 : CHARACTER(LEN=3*default_string_length + 6) :: names = ""
104 : INTEGER(KIND=int_8) :: ncalls = 0, nempty = 0, flop = 0
105 : INTEGER(KIND=int_8), DIMENSION(3) :: nremap = 0, nzremap = 0
106 : INTEGER(KIND=int_8), DIMENSION(2) :: ncrop = 0, nzcrop = 0, ncrop_remap = 0
107 : END TYPE contract_stats_type
108 :
109 : TYPE(contract_stats_type), DIMENSION(:), ALLOCATABLE, SAVE :: contract_stats
110 : INTEGER, SAVE :: ncontract_stats = 0
111 :
112 : CONTAINS
113 :
114 : ! **************************************************************************************************
115 : !> \brief Initialize the DBT library (after dbm_library_init).
116 : !> \author Hans Pabst
117 : ! **************************************************************************************************
118 1474 : SUBROUTINE dbt_library_init()
119 1474 : END SUBROUTINE dbt_library_init
120 :
121 : ! **************************************************************************************************
122 : !> \brief Finalize the DBT library (before dbm_library_finalize), which releases its buffers.
123 : !> \author Hans Pabst
124 : ! **************************************************************************************************
125 1474 : SUBROUTINE dbt_library_finalize()
126 1474 : CALL dbt_clear_mempools()
127 1474 : IF (ALLOCATED(contract_stats)) DEALLOCATE (contract_stats)
128 1474 : ncontract_stats = 0
129 1474 : END SUBROUTINE dbt_library_finalize
130 :
131 : ! **************************************************************************************************
132 : !> \brief Release the buffers that persist across calls (as dbcsr_clear_mempools), e.g., once a
133 : !> phase with large tensor operations ends. They are allocated again when needed.
134 : !> \author Hans Pabst
135 : ! **************************************************************************************************
136 2910 : SUBROUTINE dbt_clear_mempools()
137 2910 : CALL dbt_reshape_release()
138 2910 : END SUBROUTINE dbt_clear_mempools
139 :
140 : ! **************************************************************************************************
141 : !> \brief Print per-contraction statistics (rank-local table, totals over all ranks).
142 : !> \param mpi_comm communicator of all ranks calling this routine
143 : !> \param output_unit unit to write to (rank-local table is written if positive)
144 : !> \author Hans Pabst
145 : ! **************************************************************************************************
146 11306 : SUBROUTINE dbt_library_print_stats(mpi_comm, output_unit)
147 : TYPE(mp_comm_type), INTENT(IN) :: mpi_comm
148 : INTEGER, INTENT(IN) :: output_unit
149 :
150 : INTEGER :: i, j, k
151 : INTEGER(KIND=int_8) :: ncalls, nzcopied
152 11306 : INTEGER, DIMENSION(:), ALLOCATABLE :: order
153 11306 : REAL(KIND=dp), DIMENSION(:), ALLOCATABLE :: copied
154 :
155 11306 : ncalls = 0
156 11306 : nzcopied = 0
157 14324 : DO i = 1, ncontract_stats
158 3018 : ncalls = ncalls + contract_stats(i)%ncalls
159 29414 : nzcopied = nzcopied + SUM(contract_stats(i)%nzremap) + SUM(contract_stats(i)%nzcrop)
160 : END DO
161 11306 : CALL mpi_comm%sum(ncalls)
162 11306 : CALL mpi_comm%sum(nzcopied)
163 :
164 11306 : IF (output_unit > 0 .AND. 0 < ncalls) THEN
165 1285 : ALLOCATE (order(ncontract_stats), copied(ncontract_stats))
166 1766 : DO i = 1, ncontract_stats
167 1509 : order(i) = i
168 9311 : copied(i) = -REAL(SUM(contract_stats(i)%nzremap) + SUM(contract_stats(i)%nzcrop), dp)
169 : END DO
170 257 : CALL sort(copied, ncontract_stats, order) ! ascending: most copied data first
171 257 : WRITE (output_unit, '(/,T2,A)') REPEAT("-", 79)
172 257 : WRITE (output_unit, '(T2,A,T70,A)') "DBT CONTRACTION STATISTICS (rank-local)", "copied [MB]"
173 : WRITE (output_unit, '(T2,A,T47,A,T55,A,T62,A,T68,A,T75,A)') &
174 257 : "A x B = C", "calls", "remap", "crop", "both", "total"
175 1766 : DO k = 1, ncontract_stats
176 1509 : i = order(k)
177 257 : ASSOCIATE (s => contract_stats(i))
178 : WRITE (output_unit, '(T2,A,T44,I8,T52,3I3,T61,2I3,T67,2I3,T73,F8.1)') &
179 1509 : s%names(1:MIN(LEN_TRIM(s%names), 41)), s%ncalls, &
180 12072 : (MIN(s%nremap(j), 99_int_8), j=1, 3), (MIN(s%ncrop(j), 99_int_8), j=1, 2), &
181 4527 : (MIN(s%ncrop_remap(j), 99_int_8), j=1, 2), &
182 25653 : 8.0_dp*REAL(SUM(s%nzremap) + SUM(s%nzcrop), dp)/1.0E6_dp
183 : END ASSOCIATE
184 : END DO
185 257 : WRITE (output_unit, '(T2,A,T44,I8,T73,F8.1)') "all ranks", ncalls, 8.0_dp*REAL(nzcopied, dp)/1.0E6_dp
186 257 : WRITE (output_unit, '(T2,A)') REPEAT("-", 79)
187 : END IF
188 11306 : END SUBROUTINE dbt_library_print_stats
189 :
190 : ! **************************************************************************************************
191 : !> \brief Index of the statistics entry for the given tensor names (created if missing).
192 : !> \author Hans Pabst
193 : ! **************************************************************************************************
194 257420 : FUNCTION contract_stats_index(name_1, name_2, name_3) RESULT(idx)
195 : CHARACTER(LEN=*), INTENT(IN) :: name_1, name_2, name_3
196 : INTEGER :: idx
197 :
198 : CHARACTER(LEN=3*default_string_length + 6) :: names
199 : INTEGER :: i
200 257420 : TYPE(contract_stats_type), DIMENSION(:), ALLOCATABLE :: grown
201 :
202 257420 : names = TRIM(name_1)//" x "//TRIM(name_2)//" = "//TRIM(name_3)
203 257420 : idx = 0
204 1361425 : DO i = 1, ncontract_stats
205 1360923 : IF (contract_stats(i)%names == names) idx = i
206 1361425 : IF (0 < idx) EXIT
207 : END DO
208 257420 : IF (0 == idx) THEN
209 3022 : IF (.NOT. ALLOCATED(contract_stats)) ALLOCATE (contract_stats(16))
210 502 : IF (SIZE(contract_stats) <= ncontract_stats) THEN
211 0 : ALLOCATE (grown(2*SIZE(contract_stats)))
212 0 : grown(1:ncontract_stats) = contract_stats(1:ncontract_stats)
213 0 : CALL MOVE_ALLOC(grown, contract_stats)
214 : END IF
215 502 : ncontract_stats = ncontract_stats + 1
216 502 : idx = ncontract_stats
217 502 : contract_stats(idx)%names = names
218 : END IF
219 257420 : END FUNCTION contract_stats_index
220 :
221 : ! **************************************************************************************************
222 : !> \brief Copy tensor data.
223 : !> Redistributes tensor data according to distributions of target and source tensor.
224 : !> Permutes tensor index according to `order` argument (if present).
225 : !> Source and target tensor formats are arbitrary as long as the following requirements are met:
226 : !> * source and target tensors have the same rank and the same sizes in each dimension in terms
227 : !> of tensor elements (block sizes don't need to be the same).
228 : !> If `order` argument is present, sizes must match after index permutation.
229 : !> OR
230 : !> * target tensor is not yet created, in this case an exact copy of source tensor is returned.
231 : !> \param tensor_in Source
232 : !> \param tensor_out Target
233 : !> \param order Permutation of target tensor index.
234 : !> Exact same convention as order argument of RESHAPE intrinsic.
235 : !> \param bounds crop tensor data: start and end index for each tensor dimension
236 : !> \author Patrick Seewald
237 : ! **************************************************************************************************
238 1165878 : SUBROUTINE dbt_copy(tensor_in, tensor_out, order, summation, bounds, move_data, unit_nr)
239 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_in, tensor_out
240 : INTEGER, DIMENSION(ndims_tensor(tensor_in)), &
241 : INTENT(IN), OPTIONAL :: order
242 : LOGICAL, INTENT(IN), OPTIONAL :: summation, move_data
243 : INTEGER, DIMENSION(2, ndims_tensor(tensor_in)), &
244 : INTENT(IN), OPTIONAL :: bounds
245 : INTEGER, INTENT(IN), OPTIONAL :: unit_nr
246 : INTEGER :: handle
247 :
248 582939 : CALL tensor_in%pgrid%mp_comm_2d%sync()
249 582939 : CALL timeset("dbt_total", handle)
250 :
251 : ! make sure that it is safe to use dbt_copy during a batched contraction
252 582939 : CALL dbt_tas_batched_mm_complete(tensor_in%matrix_rep, warn=.TRUE.)
253 582939 : CALL dbt_tas_batched_mm_complete(tensor_out%matrix_rep, warn=.TRUE.)
254 :
255 582939 : CALL dbt_copy_expert(tensor_in, tensor_out, order, summation, bounds, move_data, unit_nr)
256 582939 : CALL tensor_in%pgrid%mp_comm_2d%sync()
257 582939 : CALL timestop(handle)
258 582939 : END SUBROUTINE dbt_copy
259 :
260 : ! **************************************************************************************************
261 : !> \brief expert routine for copying a tensor. For internal use only.
262 : !> \author Patrick Seewald
263 : ! **************************************************************************************************
264 5772 : SUBROUTINE dbt_copy_expert(tensor_in, tensor_out, order, summation, bounds, move_data, unit_nr)
265 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_in, tensor_out
266 : INTEGER, DIMENSION(ndims_tensor(tensor_in)), &
267 : INTENT(IN), OPTIONAL :: order
268 : LOGICAL, INTENT(IN), OPTIONAL :: summation, move_data
269 : INTEGER, DIMENSION(2, ndims_tensor(tensor_in)), &
270 : INTENT(IN), OPTIONAL :: bounds
271 : INTEGER, INTENT(IN), OPTIONAL :: unit_nr
272 :
273 : TYPE(dbt_type), POINTER :: in_tmp_1, in_tmp_2, &
274 : in_tmp_3, out_tmp_1
275 : INTEGER :: handle, unit_nr_prv
276 641540 : INTEGER, DIMENSION(:), ALLOCATABLE :: map1_in_1, map1_in_2, map2_in_1, map2_in_2
277 :
278 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_copy'
279 : LOGICAL :: dist_compatible_tas, dist_compatible_tensor, &
280 : summation_prv, new_in_1, new_in_2, &
281 : new_in_3, new_out_1, block_compatible, &
282 : move_prv, crop_reshape
283 641540 : TYPE(array_list) :: blk_sizes_in
284 :
285 641540 : CALL timeset(routineN, handle)
286 :
287 641540 : CPASSERT(tensor_out%valid)
288 :
289 641540 : unit_nr_prv = prep_output_unit(unit_nr)
290 :
291 641540 : IF (PRESENT(move_data)) THEN
292 458637 : move_prv = move_data
293 : ELSE
294 182903 : move_prv = .FALSE.
295 : END IF
296 :
297 641540 : dist_compatible_tas = .FALSE.
298 641540 : dist_compatible_tensor = .FALSE.
299 641540 : block_compatible = .FALSE.
300 641540 : new_in_1 = .FALSE.
301 641540 : new_in_2 = .FALSE.
302 641540 : new_in_3 = .FALSE.
303 641540 : new_out_1 = .FALSE.
304 :
305 641540 : IF (PRESENT(summation)) THEN
306 187145 : summation_prv = summation
307 : ELSE
308 : summation_prv = .FALSE.
309 : END IF
310 :
311 : ! crop while reshaping (single copy) if that is the copy anyway
312 641540 : crop_reshape = .FALSE.
313 641540 : IF (PRESENT(bounds) .AND. .NOT. PRESENT(order)) crop_reshape = copy_reshapes(tensor_in, tensor_out)
314 641540 : IF (PRESENT(bounds) .AND. .NOT. crop_reshape) THEN
315 33838 : ALLOCATE (in_tmp_1)
316 4834 : CALL dbt_crop(tensor_in, in_tmp_1, bounds=bounds, move_data=move_prv)
317 4834 : new_in_1 = .TRUE.
318 4834 : move_prv = .TRUE.
319 : ELSE
320 : in_tmp_1 => tensor_in
321 : END IF
322 :
323 641540 : IF (PRESENT(order)) THEN
324 175725 : CALL reorder_arrays(in_tmp_1%blk_sizes, blk_sizes_in, order=order)
325 175725 : block_compatible = check_equal(blk_sizes_in, tensor_out%blk_sizes)
326 : ELSE
327 465815 : block_compatible = check_equal(in_tmp_1%blk_sizes, tensor_out%blk_sizes)
328 : END IF
329 :
330 641540 : IF (.NOT. block_compatible) THEN
331 1075503 : ALLOCATE (in_tmp_2, out_tmp_1)
332 : CALL dbt_make_compatible_blocks(in_tmp_1, tensor_out, in_tmp_2, out_tmp_1, order=order, &
333 82731 : nodata2=.NOT. summation_prv, move_data=move_prv)
334 82731 : new_in_2 = .TRUE.; new_out_1 = .TRUE.
335 82731 : move_prv = .TRUE.
336 : ELSE
337 : in_tmp_2 => in_tmp_1
338 : out_tmp_1 => tensor_out
339 : END IF
340 :
341 641540 : IF (PRESENT(order)) THEN
342 1230075 : ALLOCATE (in_tmp_3)
343 175725 : CALL dbt_permute_index(in_tmp_2, in_tmp_3, order)
344 175725 : new_in_3 = .TRUE.
345 : ELSE
346 : in_tmp_3 => in_tmp_2
347 : END IF
348 :
349 1924620 : ALLOCATE (map1_in_1(ndims_matrix_row(in_tmp_3)))
350 1924620 : ALLOCATE (map1_in_2(ndims_matrix_column(in_tmp_3)))
351 641540 : CALL dbt_get_mapping_info(in_tmp_3%nd_index, map1_2d=map1_in_1, map2_2d=map1_in_2)
352 :
353 1924620 : ALLOCATE (map2_in_1(ndims_matrix_row(out_tmp_1)))
354 1924620 : ALLOCATE (map2_in_2(ndims_matrix_column(out_tmp_1)))
355 641540 : CALL dbt_get_mapping_info(out_tmp_1%nd_index, map1_2d=map2_in_1, map2_2d=map2_in_2)
356 :
357 641540 : IF (.NOT. PRESENT(order)) THEN
358 465815 : IF (array_eq_i(map1_in_1, map2_in_1) .AND. array_eq_i(map1_in_2, map2_in_2)) THEN
359 385512 : dist_compatible_tas = check_equal(in_tmp_3%nd_dist, out_tmp_1%nd_dist)
360 1043763 : ELSE IF (array_eq_i([map1_in_1, map1_in_2], [map2_in_1, map2_in_2])) THEN
361 33900 : dist_compatible_tensor = check_equal(in_tmp_3%nd_dist, out_tmp_1%nd_dist)
362 : END IF
363 : END IF
364 :
365 385512 : IF (dist_compatible_tas) THEN
366 317440 : CALL dbt_tas_copy(out_tmp_1%matrix_rep, in_tmp_3%matrix_rep, summation)
367 317440 : IF (move_prv) CALL dbt_clear(in_tmp_3)
368 324100 : ELSE IF (dist_compatible_tensor) THEN
369 24210 : CALL dbt_copy_nocomm(in_tmp_3, out_tmp_1, summation)
370 24210 : IF (move_prv) CALL dbt_clear(in_tmp_3)
371 299890 : ELSE IF (crop_reshape) THEN
372 938 : CALL dbt_reshape(in_tmp_3, out_tmp_1, summation, move_data=move_prv, bounds=bounds)
373 : ELSE
374 298952 : CALL dbt_reshape(in_tmp_3, out_tmp_1, summation, move_data=move_prv)
375 : END IF
376 :
377 641540 : IF (new_in_1) THEN
378 4834 : CALL dbt_destroy(in_tmp_1)
379 4834 : DEALLOCATE (in_tmp_1)
380 : END IF
381 :
382 641540 : IF (new_in_2) THEN
383 82731 : CALL dbt_destroy(in_tmp_2)
384 82731 : DEALLOCATE (in_tmp_2)
385 : END IF
386 :
387 641540 : IF (new_in_3) THEN
388 175725 : CALL dbt_destroy(in_tmp_3)
389 175725 : DEALLOCATE (in_tmp_3)
390 : END IF
391 :
392 641540 : IF (new_out_1) THEN
393 82731 : IF (unit_nr_prv /= 0) THEN
394 0 : CALL dbt_write_tensor_dist(out_tmp_1, unit_nr)
395 : END IF
396 82731 : CALL dbt_split_copyback(out_tmp_1, tensor_out, summation)
397 82731 : CALL dbt_destroy(out_tmp_1)
398 82731 : DEALLOCATE (out_tmp_1)
399 : END IF
400 :
401 641540 : CALL timestop(handle)
402 :
403 1288852 : END SUBROUTINE dbt_copy_expert
404 :
405 : ! **************************************************************************************************
406 : !> \brief Whether dbt_copy_expert (without order) copies by dbt_reshape: compatible blocks, but
407 : !> neither the matrix layout nor the distribution allow a copy without communication.
408 : !> \author Hans Pabst
409 : ! **************************************************************************************************
410 4622 : FUNCTION copy_reshapes(tensor_in, tensor_out)
411 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in, tensor_out
412 : LOGICAL :: copy_reshapes
413 :
414 9244 : INTEGER, DIMENSION(ndims_matrix_row(tensor_in)) :: map1_in_1
415 9244 : INTEGER, DIMENSION(ndims_matrix_column(tensor_in)) :: map1_in_2
416 9244 : INTEGER, DIMENSION(ndims_matrix_row(tensor_out)) :: map2_in_1
417 9244 : INTEGER, DIMENSION(ndims_matrix_column(tensor_out)) :: map2_in_2
418 :
419 4622 : copy_reshapes = check_equal(tensor_in%blk_sizes, tensor_out%blk_sizes)
420 4622 : IF (copy_reshapes) THEN
421 4622 : CALL dbt_get_mapping_info(tensor_in%nd_index, map1_2d=map1_in_1, map2_2d=map1_in_2)
422 4622 : CALL dbt_get_mapping_info(tensor_out%nd_index, map1_2d=map2_in_1, map2_2d=map2_in_2)
423 60054 : IF (array_eq_i([map1_in_1, map1_in_2], [map2_in_1, map2_in_2])) THEN
424 4618 : copy_reshapes = .NOT. check_equal(tensor_in%nd_dist, tensor_out%nd_dist)
425 : END IF
426 : END IF
427 4622 : END FUNCTION copy_reshapes
428 :
429 : ! **************************************************************************************************
430 : !> \brief copy without communication, requires that both tensors have same process grid and distribution
431 : !> \param summation Whether to sum matrices b = a + b
432 : !> \author Patrick Seewald
433 : ! **************************************************************************************************
434 24210 : SUBROUTINE dbt_copy_nocomm(tensor_in, tensor_out, summation)
435 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in
436 : TYPE(dbt_type), INTENT(INOUT) :: tensor_out
437 : LOGICAL, INTENT(IN), OPTIONAL :: summation
438 : TYPE(dbt_iterator_type) :: iter
439 24210 : INTEGER, DIMENSION(ndims_tensor(tensor_in)) :: ind_nd
440 24210 : TYPE(block_nd) :: blk_data
441 : LOGICAL :: found
442 :
443 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_copy_nocomm'
444 : INTEGER :: handle
445 :
446 24210 : CALL timeset(routineN, handle)
447 24210 : CPASSERT(tensor_out%valid)
448 :
449 24210 : IF (PRESENT(summation)) THEN
450 11376 : IF (.NOT. summation) CALL dbt_clear(tensor_out)
451 : ELSE
452 12834 : CALL dbt_clear(tensor_out)
453 : END IF
454 :
455 24210 : CALL dbt_reserve_blocks(tensor_in, tensor_out)
456 :
457 : !$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,tensor_out,summation) &
458 24210 : !$OMP PRIVATE(iter,ind_nd,blk_data,found)
459 : CALL dbt_iterator_start(iter, tensor_in)
460 : DO WHILE (dbt_iterator_blocks_left(iter))
461 : CALL dbt_iterator_next_block(iter, ind_nd)
462 : CALL dbt_get_block(tensor_in, ind_nd, blk_data, found)
463 : CPASSERT(found)
464 : CALL dbt_put_block(tensor_out, ind_nd, blk_data, summation=summation)
465 : CALL destroy_block(blk_data)
466 : END DO
467 : CALL dbt_iterator_stop(iter)
468 : !$OMP END PARALLEL
469 :
470 24210 : CALL timestop(handle)
471 48420 : END SUBROUTINE dbt_copy_nocomm
472 :
473 : ! **************************************************************************************************
474 : !> \brief copy matrix to tensor.
475 : !> \param summation tensor_out = tensor_out + matrix_in
476 : !> \author Patrick Seewald
477 : ! **************************************************************************************************
478 97109 : SUBROUTINE dbt_copy_matrix_to_tensor(matrix_in, tensor_out, summation)
479 : TYPE(dbcsr_type), TARGET, INTENT(IN) :: matrix_in
480 : TYPE(dbt_type), INTENT(INOUT) :: tensor_out
481 : LOGICAL, INTENT(IN), OPTIONAL :: summation
482 : TYPE(dbcsr_type), POINTER :: matrix_in_desym
483 :
484 : INTEGER, DIMENSION(2) :: ind_2d
485 97109 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: block_arr
486 97109 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: block
487 : TYPE(dbcsr_iterator_type) :: iter
488 :
489 : INTEGER :: handle
490 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_copy_matrix_to_tensor'
491 :
492 97109 : CALL timeset(routineN, handle)
493 97109 : CPASSERT(tensor_out%valid)
494 :
495 97109 : NULLIFY (block)
496 :
497 97109 : IF (dbcsr_has_symmetry(matrix_in)) THEN
498 4950 : ALLOCATE (matrix_in_desym)
499 4950 : CALL dbcsr_desymmetrize(matrix_in, matrix_in_desym)
500 : ELSE
501 : matrix_in_desym => matrix_in
502 : END IF
503 :
504 97109 : IF (PRESENT(summation)) THEN
505 0 : IF (.NOT. summation) CALL dbt_clear(tensor_out)
506 : ELSE
507 97109 : CALL dbt_clear(tensor_out)
508 : END IF
509 :
510 97109 : CALL dbt_reserve_blocks(matrix_in_desym, tensor_out)
511 :
512 : !$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in_desym,tensor_out,summation) &
513 97109 : !$OMP PRIVATE(iter,ind_2d,block,block_arr)
514 : CALL dbcsr_iterator_start(iter, matrix_in_desym)
515 : DO WHILE (dbcsr_iterator_blocks_left(iter))
516 : CALL dbcsr_iterator_next_block(iter, ind_2d(1), ind_2d(2), block)
517 : CALL allocate_any(block_arr, source=block)
518 : CALL dbt_put_block(tensor_out, ind_2d, SHAPE(block_arr), block_arr, summation=summation)
519 : DEALLOCATE (block_arr)
520 : END DO
521 : CALL dbcsr_iterator_stop(iter)
522 : !$OMP END PARALLEL
523 :
524 97109 : IF (dbcsr_has_symmetry(matrix_in)) THEN
525 4950 : CALL dbcsr_release(matrix_in_desym)
526 4950 : DEALLOCATE (matrix_in_desym)
527 : END IF
528 :
529 97109 : CALL timestop(handle)
530 :
531 194218 : END SUBROUTINE dbt_copy_matrix_to_tensor
532 :
533 : ! **************************************************************************************************
534 : !> \brief copy tensor to matrix
535 : !> \param summation matrix_out = matrix_out + tensor_in
536 : !> \author Patrick Seewald
537 : ! **************************************************************************************************
538 65292 : SUBROUTINE dbt_copy_tensor_to_matrix(tensor_in, matrix_out, summation)
539 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in
540 : TYPE(dbcsr_type), INTENT(INOUT) :: matrix_out
541 : LOGICAL, INTENT(IN), OPTIONAL :: summation
542 : TYPE(dbt_iterator_type) :: iter
543 : INTEGER :: handle
544 : INTEGER, DIMENSION(2) :: ind_2d
545 65292 : REAL(KIND=dp), DIMENSION(:, :), ALLOCATABLE :: block
546 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_copy_tensor_to_matrix'
547 : LOGICAL :: found
548 :
549 65292 : CALL timeset(routineN, handle)
550 :
551 65292 : IF (PRESENT(summation)) THEN
552 5910 : IF (.NOT. summation) CALL dbcsr_clear(matrix_out)
553 : ELSE
554 59382 : CALL dbcsr_clear(matrix_out)
555 : END IF
556 :
557 65292 : CALL dbt_reserve_blocks(tensor_in, matrix_out)
558 :
559 : !$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,matrix_out,summation) &
560 65292 : !$OMP PRIVATE(iter,ind_2d,block,found)
561 : CALL dbt_iterator_start(iter, tensor_in)
562 : DO WHILE (dbt_iterator_blocks_left(iter))
563 : CALL dbt_iterator_next_block(iter, ind_2d)
564 : IF (dbcsr_has_symmetry(matrix_out) .AND. checker_tr(ind_2d(1), ind_2d(2))) CYCLE
565 :
566 : CALL dbt_get_block(tensor_in, ind_2d, block, found)
567 : CPASSERT(found)
568 :
569 : IF (dbcsr_has_symmetry(matrix_out) .AND. ind_2d(1) > ind_2d(2)) THEN
570 : CALL dbcsr_put_block(matrix_out, ind_2d(2), ind_2d(1), TRANSPOSE(block), summation=summation)
571 : ELSE
572 : CALL dbcsr_put_block(matrix_out, ind_2d(1), ind_2d(2), block, summation=summation)
573 : END IF
574 : DEALLOCATE (block)
575 : END DO
576 : CALL dbt_iterator_stop(iter)
577 : !$OMP END PARALLEL
578 :
579 65292 : CALL timestop(handle)
580 :
581 130584 : END SUBROUTINE dbt_copy_tensor_to_matrix
582 :
583 : ! **************************************************************************************************
584 : !> \brief Contract tensors by multiplying matrix representations.
585 : !> tensor_3(map_1, map_2) := alpha * tensor_1(notcontract_1, contract_1)
586 : !> * tensor_2(contract_2, notcontract_2)
587 : !> + beta * tensor_3(map_1, map_2)
588 : !>
589 : !> \note
590 : !> note 1: block sizes of the corresponding indices need to be the same in all tensors.
591 : !>
592 : !> note 2: for best performance the tensors should have been created in matrix layouts
593 : !> compatible with the contraction, e.g. tensor_1 should have been created with either
594 : !> map1_2d == contract_1 and map2_2d == notcontract_1 or map1_2d == notcontract_1 and
595 : !> map2_2d == contract_1 (the same with tensor_2 and contract_2 / notcontract_2 and with
596 : !> tensor_3 and map_1 / map_2).
597 : !> Furthermore the two largest tensors involved in the contraction should map both to either
598 : !> tall or short matrices: the largest matrix dimension should be "on the same side"
599 : !> and should have identical distribution (which is always the case if the distributions were
600 : !> obtained with dbt_default_distvec).
601 : !>
602 : !> note 3: if the same tensor occurs in multiple contractions, a different tensor object should
603 : !> be created for each contraction and the data should be copied between the tensors by use of
604 : !> dbt_copy. If the same tensor object is used in multiple contractions,
605 : !> matrix layouts are not compatible for all contractions (see note 2).
606 : !>
607 : !> note 4: automatic optimizations are enabled by using the feature of batched contraction, see
608 : !> dbt_batched_contract_init, dbt_batched_contract_finalize.
609 : !> The arguments bounds_1, bounds_2, bounds_3 give the index ranges of the batches.
610 : !>
611 : !> \param tensor_1 first tensor (in)
612 : !> \param tensor_2 second tensor (in)
613 : !> \param contract_1 indices of tensor_1 to contract
614 : !> \param contract_2 indices of tensor_2 to contract (1:1 with contract_1)
615 : !> \param map_1 which indices of tensor_3 map to non-contracted indices of tensor_1 (1:1 with notcontract_1)
616 : !> \param map_2 which indices of tensor_3 map to non-contracted indices of tensor_2 (1:1 with notcontract_2)
617 : !> \param notcontract_1 indices of tensor_1 not to contract
618 : !> \param notcontract_2 indices of tensor_2 not to contract
619 : !> \param tensor_3 contracted tensor (out)
620 : !> \param bounds_1 bounds corresponding to contract_1 AKA contract_2:
621 : !> start and end index of an index range over which to contract.
622 : !> For use in batched contraction.
623 : !> \param bounds_2 bounds corresponding to notcontract_1: start and end index of an index range.
624 : !> For use in batched contraction.
625 : !> \param bounds_3 bounds corresponding to notcontract_2: start and end index of an index range.
626 : !> For use in batched contraction.
627 : !> \param optimize_dist Whether distribution should be optimized internally. In the current
628 : !> implementation this guarantees optimal parameters only for dense matrices.
629 : !> \param pgrid_opt_1 Optionally return optimal process grid for tensor_1.
630 : !> This can be used to choose optimal process grids for subsequent tensor
631 : !> contractions with tensors of similar shape and sparsity. Under some conditions,
632 : !> pgrid_opt_1 can not be returned, in this case the pointer is not associated.
633 : !> \param pgrid_opt_2 Optionally return optimal process grid for tensor_2.
634 : !> \param pgrid_opt_3 Optionally return optimal process grid for tensor_3.
635 : !> \param filter_eps As in DBM mm
636 : !> \param flop As in DBM mm
637 : !> \param move_data memory optimization: transfer data such that tensor_1 and tensor_2 are empty on return
638 : !> \param retain_sparsity enforce the sparsity pattern of the existing tensor_3; default is no
639 : !> \param unit_nr output unit for logging
640 : !> set it to -1 on ranks that should not write (and any valid unit number on
641 : !> ranks that should write output) if 0 on ALL ranks, no output is written
642 : !> \param log_verbose verbose logging (for testing only)
643 : !> \author Patrick Seewald
644 : ! **************************************************************************************************
645 514840 : SUBROUTINE dbt_contract(alpha, tensor_1, tensor_2, beta, tensor_3, &
646 257420 : contract_1, notcontract_1, &
647 257420 : contract_2, notcontract_2, &
648 257420 : map_1, map_2, &
649 257420 : bounds_1, bounds_2, bounds_3, &
650 : optimize_dist, pgrid_opt_1, pgrid_opt_2, pgrid_opt_3, &
651 : filter_eps, flop, move_data, retain_sparsity, unit_nr, log_verbose)
652 : REAL(dp), INTENT(IN) :: alpha
653 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_1
654 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_2
655 : REAL(dp), INTENT(IN) :: beta
656 : INTEGER, DIMENSION(:), INTENT(IN) :: contract_1
657 : INTEGER, DIMENSION(:), INTENT(IN) :: contract_2
658 : INTEGER, DIMENSION(:), INTENT(IN) :: map_1
659 : INTEGER, DIMENSION(:), INTENT(IN) :: map_2
660 : INTEGER, DIMENSION(:), INTENT(IN) :: notcontract_1
661 : INTEGER, DIMENSION(:), INTENT(IN) :: notcontract_2
662 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_3
663 : INTEGER, DIMENSION(2, SIZE(contract_1)), &
664 : INTENT(IN), OPTIONAL :: bounds_1
665 : INTEGER, DIMENSION(2, SIZE(notcontract_1)), &
666 : INTENT(IN), OPTIONAL :: bounds_2
667 : INTEGER, DIMENSION(2, SIZE(notcontract_2)), &
668 : INTENT(IN), OPTIONAL :: bounds_3
669 : LOGICAL, INTENT(IN), OPTIONAL :: optimize_dist
670 : TYPE(dbt_pgrid_type), INTENT(OUT), &
671 : POINTER, OPTIONAL :: pgrid_opt_1
672 : TYPE(dbt_pgrid_type), INTENT(OUT), &
673 : POINTER, OPTIONAL :: pgrid_opt_2
674 : TYPE(dbt_pgrid_type), INTENT(OUT), &
675 : POINTER, OPTIONAL :: pgrid_opt_3
676 : REAL(KIND=dp), INTENT(IN), OPTIONAL :: filter_eps
677 : INTEGER(KIND=int_8), INTENT(OUT), OPTIONAL :: flop
678 : LOGICAL, INTENT(IN), OPTIONAL :: move_data
679 : LOGICAL, INTENT(IN), OPTIONAL :: retain_sparsity
680 : INTEGER, OPTIONAL, INTENT(IN) :: unit_nr
681 : LOGICAL, INTENT(IN), OPTIONAL :: log_verbose
682 :
683 : INTEGER :: handle
684 :
685 257420 : CALL tensor_1%pgrid%mp_comm_2d%sync()
686 257420 : CALL timeset("dbt_total", handle)
687 : CALL dbt_contract_expert(alpha, tensor_1, tensor_2, beta, tensor_3, &
688 : contract_1, notcontract_1, &
689 : contract_2, notcontract_2, &
690 : map_1, map_2, &
691 : bounds_1=bounds_1, &
692 : bounds_2=bounds_2, &
693 : bounds_3=bounds_3, &
694 : optimize_dist=optimize_dist, &
695 : pgrid_opt_1=pgrid_opt_1, &
696 : pgrid_opt_2=pgrid_opt_2, &
697 : pgrid_opt_3=pgrid_opt_3, &
698 : filter_eps=filter_eps, &
699 : flop=flop, &
700 : move_data=move_data, &
701 : retain_sparsity=retain_sparsity, &
702 : unit_nr=unit_nr, &
703 257420 : log_verbose=log_verbose)
704 257420 : CALL tensor_1%pgrid%mp_comm_2d%sync()
705 257420 : CALL timestop(handle)
706 :
707 257420 : END SUBROUTINE dbt_contract
708 :
709 : ! **************************************************************************************************
710 : !> \brief expert routine for tensor contraction. For internal use only.
711 : !> \param nblks_local number of local blocks on this MPI rank
712 : !> \author Patrick Seewald
713 : ! **************************************************************************************************
714 257420 : SUBROUTINE dbt_contract_expert(alpha, tensor_1, tensor_2, beta, tensor_3, &
715 257420 : contract_1, notcontract_1, &
716 257420 : contract_2, notcontract_2, &
717 257420 : map_1, map_2, &
718 257420 : bounds_1, bounds_2, bounds_3, &
719 : optimize_dist, pgrid_opt_1, pgrid_opt_2, pgrid_opt_3, &
720 : filter_eps, flop, move_data, retain_sparsity, &
721 : nblks_local, unit_nr, log_verbose)
722 : REAL(dp), INTENT(IN) :: alpha
723 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_1
724 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_2
725 : REAL(dp), INTENT(IN) :: beta
726 : INTEGER, DIMENSION(:), INTENT(IN) :: contract_1
727 : INTEGER, DIMENSION(:), INTENT(IN) :: contract_2
728 : INTEGER, DIMENSION(:), INTENT(IN) :: map_1
729 : INTEGER, DIMENSION(:), INTENT(IN) :: map_2
730 : INTEGER, DIMENSION(:), INTENT(IN) :: notcontract_1
731 : INTEGER, DIMENSION(:), INTENT(IN) :: notcontract_2
732 : TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_3
733 : INTEGER, DIMENSION(2, SIZE(contract_1)), &
734 : INTENT(IN), OPTIONAL :: bounds_1
735 : INTEGER, DIMENSION(2, SIZE(notcontract_1)), &
736 : INTENT(IN), OPTIONAL :: bounds_2
737 : INTEGER, DIMENSION(2, SIZE(notcontract_2)), &
738 : INTENT(IN), OPTIONAL :: bounds_3
739 : LOGICAL, INTENT(IN), OPTIONAL :: optimize_dist
740 : TYPE(dbt_pgrid_type), INTENT(OUT), &
741 : POINTER, OPTIONAL :: pgrid_opt_1
742 : TYPE(dbt_pgrid_type), INTENT(OUT), &
743 : POINTER, OPTIONAL :: pgrid_opt_2
744 : TYPE(dbt_pgrid_type), INTENT(OUT), &
745 : POINTER, OPTIONAL :: pgrid_opt_3
746 : REAL(KIND=dp), INTENT(IN), OPTIONAL :: filter_eps
747 : INTEGER(KIND=int_8), INTENT(OUT), OPTIONAL :: flop
748 : LOGICAL, INTENT(IN), OPTIONAL :: move_data
749 : LOGICAL, INTENT(IN), OPTIONAL :: retain_sparsity
750 : INTEGER, INTENT(OUT), OPTIONAL :: nblks_local
751 : INTEGER, OPTIONAL, INTENT(IN) :: unit_nr
752 : LOGICAL, INTENT(IN), OPTIONAL :: log_verbose
753 :
754 : TYPE(dbt_type), POINTER :: tensor_contr_1, tensor_contr_2, tensor_contr_3
755 4890980 : TYPE(dbt_type), TARGET :: tensor_algn_1, tensor_algn_2, tensor_algn_3
756 : TYPE(dbt_type), POINTER :: tensor_crop_1, tensor_crop_2
757 : TYPE(dbt_type), POINTER :: tensor_small, tensor_large
758 :
759 : LOGICAL :: assert_stmt, tensors_remapped
760 : INTEGER :: max_mm_dim, max_tensor, &
761 : unit_nr_prv, ref_tensor, handle
762 257420 : TYPE(mp_cart_type) :: mp_comm_opt
763 514840 : INTEGER, DIMENSION(SIZE(contract_1)) :: contract_1_mod
764 514840 : INTEGER, DIMENSION(SIZE(notcontract_1)) :: notcontract_1_mod
765 514840 : INTEGER, DIMENSION(SIZE(contract_2)) :: contract_2_mod
766 514840 : INTEGER, DIMENSION(SIZE(notcontract_2)) :: notcontract_2_mod
767 514840 : INTEGER, DIMENSION(SIZE(map_1)) :: map_1_mod
768 514840 : INTEGER, DIMENSION(SIZE(map_2)) :: map_2_mod
769 : LOGICAL :: trans_1, trans_2, trans_3
770 : LOGICAL :: new_1, new_2, new_3, move_data_1, move_data_2
771 : INTEGER :: ndims1, ndims2, ndims3
772 : INTEGER, DIMENSION(2) :: occ_1, occ_2 ! blocks in bounds, sticking out
773 257420 : INTEGER, DIMENSION(:), ALLOCATABLE :: dims1, dims2, dims3
774 :
775 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_contract'
776 257420 : CHARACTER(LEN=1), DIMENSION(:), ALLOCATABLE :: indchar1, indchar2, indchar3, indchar1_mod, &
777 257420 : indchar2_mod, indchar3_mod
778 : CHARACTER(LEN=1), DIMENSION(15), SAVE :: alph = &
779 : ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o']
780 514840 : INTEGER, DIMENSION(2, ndims_tensor(tensor_1)) :: bounds_t1
781 514840 : INTEGER, DIMENSION(2, ndims_tensor(tensor_2)) :: bounds_t2
782 : LOGICAL :: do_crop_1, do_crop_2, do_write_3, nodata_3, do_batched, pgrid_changed, &
783 : pgrid_changed_any, do_change_pgrid(2)
784 3346460 : TYPE(dbt_tas_split_info) :: split_opt, split, split_opt_avg
785 : INTEGER, DIMENSION(2) :: pdims_2d_opt, pdims_sub, pdims_sub_opt
786 : REAL(dp) :: pdim_ratio, pdim_ratio_opt
787 : INTEGER :: istat
788 : INTEGER(KIND=int_8) :: flop_prv
789 514840 : INTEGER, DIMENSION(ndims_tensor(tensor_1)) :: order_1
790 514840 : INTEGER, DIMENSION(ndims_tensor(tensor_2)) :: order_2
791 257420 : INTEGER, DIMENSION(:, :), ALLOCATABLE :: bounds_a1, bounds_a2 ! unallocated means absent
792 : LOGICAL :: cropped_1, cropped_2
793 :
794 257420 : NULLIFY (tensor_contr_1, tensor_contr_2, tensor_contr_3, tensor_crop_1, tensor_crop_2, &
795 257420 : tensor_small)
796 :
797 257420 : CALL timeset(routineN, handle)
798 :
799 257420 : CPASSERT(tensor_1%valid)
800 257420 : CPASSERT(tensor_2%valid)
801 257420 : CPASSERT(tensor_3%valid)
802 :
803 257420 : assert_stmt = SIZE(contract_1) == SIZE(contract_2)
804 257420 : CPASSERT(assert_stmt)
805 :
806 257420 : assert_stmt = SIZE(map_1) == SIZE(notcontract_1)
807 257420 : CPASSERT(assert_stmt)
808 :
809 257420 : assert_stmt = SIZE(map_2) == SIZE(notcontract_2)
810 257420 : CPASSERT(assert_stmt)
811 :
812 257420 : assert_stmt = SIZE(notcontract_1) + SIZE(contract_1) == ndims_tensor(tensor_1)
813 257420 : CPASSERT(assert_stmt)
814 :
815 257420 : assert_stmt = SIZE(notcontract_2) + SIZE(contract_2) == ndims_tensor(tensor_2)
816 257420 : CPASSERT(assert_stmt)
817 :
818 257420 : assert_stmt = SIZE(map_1) + SIZE(map_2) == ndims_tensor(tensor_3)
819 257420 : CPASSERT(assert_stmt)
820 :
821 257420 : unit_nr_prv = prep_output_unit(unit_nr)
822 :
823 257420 : IF (PRESENT(flop)) flop = 0
824 257420 : IF (PRESENT(nblks_local)) nblks_local = 0
825 :
826 257420 : IF (PRESENT(move_data)) THEN
827 61916 : move_data_1 = move_data
828 61916 : move_data_2 = move_data
829 : ELSE
830 195504 : move_data_1 = .FALSE.
831 195504 : move_data_2 = .FALSE.
832 : END IF
833 :
834 257420 : nodata_3 = .TRUE.
835 257420 : IF (PRESENT(retain_sparsity)) THEN
836 4930 : IF (retain_sparsity) nodata_3 = .FALSE.
837 : END IF
838 :
839 : CALL dbt_map_bounds_to_tensors(tensor_1, tensor_2, &
840 : contract_1, notcontract_1, &
841 : contract_2, notcontract_2, &
842 : bounds_t1, bounds_t2, &
843 : bounds_1=bounds_1, bounds_2=bounds_2, bounds_3=bounds_3, &
844 257420 : do_crop_1=do_crop_1, do_crop_2=do_crop_2)
845 :
846 : ! cropping is deferred: a tensor that is remapped is cropped while copied (single copy)
847 257420 : tensor_crop_1 => tensor_1
848 257420 : tensor_crop_2 => tensor_2
849 257420 : cropped_1 = .FALSE.
850 257420 : cropped_2 = .FALSE.
851 :
852 257420 : istat = contract_stats_index(tensor_1%name, tensor_2%name, tensor_3%name)
853 : ASSOCIATE (s => contract_stats(istat))
854 257420 : s%ncalls = s%ncalls + 1
855 : END ASSOCIATE
856 :
857 : ! shortcut for empty tensors
858 : ! this is needed to avoid unnecessary work in case user contracts different portions of a
859 : ! tensor consecutively to save memory
860 : ASSOCIATE (mp_comm => tensor_crop_1%pgrid%mp_comm_2d)
861 257420 : IF (do_crop_1) THEN
862 86941 : occ_1 = num_blocks_in_bounds(tensor_1, bounds_t1)
863 : ELSE
864 511437 : occ_1 = [dbt_get_num_blocks(tensor_1), 0]
865 : END IF
866 257420 : CALL mp_comm%max(occ_1)
867 257420 : IF (do_crop_2) THEN
868 82861 : occ_2 = num_blocks_in_bounds(tensor_2, bounds_t2)
869 : ELSE
870 523677 : occ_2 = [dbt_get_num_blocks(tensor_2), 0]
871 : END IF
872 514840 : CALL mp_comm%max(occ_2)
873 : END ASSOCIATE
874 :
875 257420 : IF (occ_1(1) == 0 .OR. occ_2(1) == 0) THEN
876 34916 : contract_stats(istat)%nempty = contract_stats(istat)%nempty + 1
877 34916 : CALL dbt_scale(tensor_3, beta)
878 : ! as if cropped with move_data (tensor data was transferred into the crop)
879 34916 : IF (do_crop_1 .AND. move_data_1) CALL dbt_clear(tensor_1)
880 34916 : IF (do_crop_2 .AND. move_data_2) CALL dbt_clear(tensor_2)
881 :
882 34916 : CALL timestop(handle)
883 34916 : RETURN
884 : END IF
885 :
886 : ! no block sticks out on any rank: a crop would copy the tensor as is
887 222504 : do_crop_1 = do_crop_1 .AND. occ_1(2) > 0
888 222504 : do_crop_2 = do_crop_2 .AND. occ_2(2) > 0
889 : ASSOCIATE (s => contract_stats(istat))
890 34900 : IF (do_crop_1) s%ncrop(1) = s%ncrop(1) + 1
891 257404 : IF (do_crop_2) s%ncrop(2) = s%ncrop(2) + 1
892 : END ASSOCIATE
893 :
894 222504 : IF (unit_nr_prv /= 0) THEN
895 55120 : IF (unit_nr_prv > 0) THEN
896 10 : WRITE (unit_nr_prv, '(A)') repeat("-", 80)
897 10 : WRITE (unit_nr_prv, '(A,1X,A,1X,A,1X,A,1X,A,1X,A)') "DBT TENSOR CONTRACTION:", &
898 20 : TRIM(tensor_crop_1%name), 'x', TRIM(tensor_crop_2%name), '=', TRIM(tensor_3%name)
899 10 : WRITE (unit_nr_prv, '(A)') repeat("-", 80)
900 : END IF
901 55120 : CALL dbt_write_tensor_info(tensor_crop_1, unit_nr_prv, full_info=log_verbose)
902 55120 : CALL dbt_write_tensor_dist(tensor_crop_1, unit_nr_prv)
903 55120 : CALL dbt_write_tensor_info(tensor_crop_2, unit_nr_prv, full_info=log_verbose)
904 55120 : CALL dbt_write_tensor_dist(tensor_crop_2, unit_nr_prv)
905 : END IF
906 :
907 : ! align tensor index with data, tensor data is not modified
908 222504 : ndims1 = ndims_tensor(tensor_crop_1)
909 222504 : ndims2 = ndims_tensor(tensor_crop_2)
910 222504 : ndims3 = ndims_tensor(tensor_3)
911 890016 : ALLOCATE (indchar1(ndims1), indchar1_mod(ndims1))
912 890016 : ALLOCATE (indchar2(ndims2), indchar2_mod(ndims2))
913 890016 : ALLOCATE (indchar3(ndims3), indchar3_mod(ndims3))
914 :
915 : ! labeling tensor index with letters
916 :
917 1923816 : indchar1([notcontract_1, contract_1]) = alph(1:ndims1) ! arb. choice
918 535710 : indchar2(notcontract_2) = alph(ndims1 + 1:ndims1 + SIZE(notcontract_2)) ! arb. choice
919 519825 : indchar2(contract_2) = indchar1(contract_1)
920 492287 : indchar3(map_1) = indchar1(notcontract_1)
921 535710 : indchar3(map_2) = indchar2(notcontract_2)
922 :
923 222504 : IF (unit_nr_prv /= 0) CALL dbt_print_contraction_index(tensor_crop_1, indchar1, &
924 : tensor_crop_2, indchar2, &
925 55120 : tensor_3, indchar3, unit_nr_prv)
926 222504 : IF (unit_nr_prv > 0) THEN
927 10 : WRITE (unit_nr_prv, '(T2,A)') "aligning tensor index with data"
928 : END IF
929 :
930 : CALL align_tensor(tensor_crop_1, contract_1, notcontract_1, &
931 222504 : tensor_algn_1, contract_1_mod, notcontract_1_mod, indchar1, indchar1_mod, order_1)
932 :
933 : CALL align_tensor(tensor_crop_2, contract_2, notcontract_2, &
934 222504 : tensor_algn_2, contract_2_mod, notcontract_2_mod, indchar2, indchar2_mod, order_2)
935 :
936 : ! bounds in the aligned index order (index i of tensor_X is index order_X(i) of tensor_algn_X)
937 222504 : IF (do_crop_1) THEN
938 104700 : ALLOCATE (bounds_a1(2, ndims_tensor(tensor_1)))
939 248392 : bounds_a1(:, order_1) = bounds_t1(:, :)
940 : END IF
941 222504 : IF (do_crop_2) THEN
942 120780 : ALLOCATE (bounds_a2(2, ndims_tensor(tensor_2)))
943 366798 : bounds_a2(:, order_2) = bounds_t2(:, :)
944 : END IF
945 :
946 : CALL align_tensor(tensor_3, map_1, map_2, &
947 222504 : tensor_algn_3, map_1_mod, map_2_mod, indchar3, indchar3_mod)
948 :
949 222504 : IF (unit_nr_prv /= 0) CALL dbt_print_contraction_index(tensor_algn_1, indchar1_mod, &
950 : tensor_algn_2, indchar2_mod, &
951 55120 : tensor_algn_3, indchar3_mod, unit_nr_prv)
952 :
953 667512 : ALLOCATE (dims1(ndims1))
954 667512 : ALLOCATE (dims2(ndims2))
955 667512 : ALLOCATE (dims3(ndims3))
956 :
957 : ! ideally we should consider block sizes and occupancy to measure tensor sizes but current solution should work for most
958 : ! cases and is more elegant. Note that we can not easily consider occupancy since it is unknown for result tensor
959 222504 : CALL blk_dims_tensor(tensor_crop_1, dims1)
960 222504 : CALL blk_dims_tensor(tensor_crop_2, dims2)
961 222504 : CALL blk_dims_tensor(tensor_3, dims3)
962 :
963 : max_mm_dim = MAXLOC([PRODUCT(INT(dims1(notcontract_1), int_8)), &
964 : PRODUCT(INT(dims1(contract_1), int_8)), &
965 1770326 : PRODUCT(INT(dims2(notcontract_2), int_8))], DIM=1)
966 2650636 : max_tensor = MAXLOC([PRODUCT(INT(dims1, int_8)), PRODUCT(INT(dims2, int_8)), PRODUCT(INT(dims3, int_8))], DIM=1)
967 57065 : SELECT CASE (max_mm_dim)
968 : CASE (1)
969 57065 : IF (unit_nr_prv > 0) THEN
970 3 : WRITE (unit_nr_prv, '(T2,A)') "large tensors: 1, 3; small tensor: 2"
971 3 : WRITE (unit_nr_prv, '(T2,A)') "sorting contraction indices"
972 : END IF
973 57065 : CALL index_linked_sort(contract_1_mod, contract_2_mod)
974 57065 : CALL index_linked_sort(map_2_mod, notcontract_2_mod)
975 56577 : SELECT CASE (max_tensor)
976 : CASE (1)
977 56577 : CALL index_linked_sort(notcontract_1_mod, map_1_mod)
978 : CASE (3)
979 488 : CALL index_linked_sort(map_1_mod, notcontract_1_mod)
980 : CASE DEFAULT
981 57065 : CPABORT("should not happen")
982 : END SELECT
983 :
984 : CALL reshape_mm_compatible(tensor_algn_1, tensor_algn_3, tensor_contr_1, tensor_contr_3, &
985 : contract_1_mod, notcontract_1_mod, map_2_mod, map_1_mod, &
986 : trans_1, trans_3, new_1, new_3, ref_tensor, nodata2=nodata_3, optimize_dist=optimize_dist, &
987 : move_data_1=move_data_1, unit_nr=unit_nr_prv, &
988 57065 : bounds1=bounds_a1, crop1=cropped_1)
989 :
990 : CALL reshape_mm_small(tensor_algn_2, contract_2_mod, notcontract_2_mod, tensor_contr_2, trans_2, &
991 : new_2, move_data=move_data_2, unit_nr=unit_nr_prv, &
992 57065 : bounds=bounds_a2, crop=cropped_2)
993 :
994 56577 : SELECT CASE (ref_tensor)
995 : CASE (1)
996 56577 : tensor_large => tensor_contr_1
997 : CASE (2)
998 57065 : tensor_large => tensor_contr_3
999 : END SELECT
1000 57065 : tensor_small => tensor_contr_2
1001 :
1002 : CASE (2)
1003 75589 : IF (unit_nr_prv > 0) THEN
1004 5 : WRITE (unit_nr_prv, '(T2,A)') "large tensors: 1, 2; small tensor: 3"
1005 5 : WRITE (unit_nr_prv, '(T2,A)') "sorting contraction indices"
1006 : END IF
1007 :
1008 75589 : CALL index_linked_sort(notcontract_1_mod, map_1_mod)
1009 75589 : CALL index_linked_sort(notcontract_2_mod, map_2_mod)
1010 74403 : SELECT CASE (max_tensor)
1011 : CASE (1)
1012 74403 : CALL index_linked_sort(contract_1_mod, contract_2_mod)
1013 : CASE (2)
1014 1186 : CALL index_linked_sort(contract_2_mod, contract_1_mod)
1015 : CASE DEFAULT
1016 75589 : CPABORT("should not happen")
1017 : END SELECT
1018 :
1019 : CALL reshape_mm_compatible(tensor_algn_1, tensor_algn_2, tensor_contr_1, tensor_contr_2, &
1020 : notcontract_1_mod, contract_1_mod, notcontract_2_mod, contract_2_mod, &
1021 : trans_1, trans_2, new_1, new_2, ref_tensor, optimize_dist=optimize_dist, &
1022 : move_data_1=move_data_1, move_data_2=move_data_2, unit_nr=unit_nr_prv, &
1023 75589 : bounds1=bounds_a1, bounds2=bounds_a2, crop1=cropped_1, crop2=cropped_2)
1024 75589 : trans_1 = .NOT. trans_1
1025 :
1026 : CALL reshape_mm_small(tensor_algn_3, map_1_mod, map_2_mod, tensor_contr_3, trans_3, &
1027 75589 : new_3, nodata=nodata_3, unit_nr=unit_nr_prv)
1028 :
1029 74403 : SELECT CASE (ref_tensor)
1030 : CASE (1)
1031 74403 : tensor_large => tensor_contr_1
1032 : CASE (2)
1033 75589 : tensor_large => tensor_contr_2
1034 : END SELECT
1035 75589 : tensor_small => tensor_contr_3
1036 :
1037 : CASE (3)
1038 89850 : IF (unit_nr_prv > 0) THEN
1039 2 : WRITE (unit_nr_prv, '(T2,A)') "large tensors: 2, 3; small tensor: 1"
1040 2 : WRITE (unit_nr_prv, '(T2,A)') "sorting contraction indices"
1041 : END IF
1042 89850 : CALL index_linked_sort(map_1_mod, notcontract_1_mod)
1043 89850 : CALL index_linked_sort(contract_2_mod, contract_1_mod)
1044 89390 : SELECT CASE (max_tensor)
1045 : CASE (2)
1046 89390 : CALL index_linked_sort(notcontract_2_mod, map_2_mod)
1047 : CASE (3)
1048 460 : CALL index_linked_sort(map_2_mod, notcontract_2_mod)
1049 : CASE DEFAULT
1050 89850 : CPABORT("should not happen")
1051 : END SELECT
1052 :
1053 : CALL reshape_mm_compatible(tensor_algn_2, tensor_algn_3, tensor_contr_2, tensor_contr_3, &
1054 : contract_2_mod, notcontract_2_mod, map_1_mod, map_2_mod, &
1055 : trans_2, trans_3, new_2, new_3, ref_tensor, nodata2=nodata_3, optimize_dist=optimize_dist, &
1056 : move_data_1=move_data_2, unit_nr=unit_nr_prv, &
1057 89850 : bounds1=bounds_a2, crop1=cropped_2)
1058 :
1059 89850 : trans_2 = .NOT. trans_2
1060 89850 : trans_3 = .NOT. trans_3
1061 :
1062 : CALL reshape_mm_small(tensor_algn_1, notcontract_1_mod, contract_1_mod, tensor_contr_1, &
1063 : trans_1, new_1, move_data=move_data_1, unit_nr=unit_nr_prv, &
1064 89850 : bounds=bounds_a1, crop=cropped_1)
1065 :
1066 89390 : SELECT CASE (ref_tensor)
1067 : CASE (1)
1068 89390 : tensor_large => tensor_contr_2
1069 : CASE (2)
1070 89850 : tensor_large => tensor_contr_3
1071 : END SELECT
1072 312354 : tensor_small => tensor_contr_1
1073 :
1074 : END SELECT
1075 :
1076 222504 : IF (unit_nr_prv /= 0) CALL dbt_print_contraction_index(tensor_contr_1, indchar1_mod, &
1077 : tensor_contr_2, indchar2_mod, &
1078 55120 : tensor_contr_3, indchar3_mod, unit_nr_prv)
1079 222504 : IF (unit_nr_prv /= 0) THEN
1080 55120 : IF (new_1) CALL dbt_write_tensor_info(tensor_contr_1, unit_nr_prv, full_info=log_verbose)
1081 55120 : IF (new_1) CALL dbt_write_tensor_dist(tensor_contr_1, unit_nr_prv)
1082 55120 : IF (new_2) CALL dbt_write_tensor_info(tensor_contr_2, unit_nr_prv, full_info=log_verbose)
1083 55120 : IF (new_2) CALL dbt_write_tensor_dist(tensor_contr_2, unit_nr_prv)
1084 : END IF
1085 :
1086 : ASSOCIATE (s => contract_stats(istat))
1087 34900 : IF (do_crop_1) s%nzcrop(1) = s%nzcrop(1) + dbt_get_nze(tensor_contr_1)
1088 222504 : IF (do_crop_2) s%nzcrop(2) = s%nzcrop(2) + dbt_get_nze(tensor_contr_2)
1089 222504 : IF (new_1) THEN
1090 23228 : s%nremap(1) = s%nremap(1) + 1
1091 23228 : s%nzremap(1) = s%nzremap(1) + dbt_get_nze(tensor_contr_1)
1092 23228 : IF (do_crop_1) s%ncrop_remap(1) = s%ncrop_remap(1) + 1
1093 : END IF
1094 445008 : IF (new_2) THEN
1095 11157 : s%nremap(2) = s%nremap(2) + 1
1096 11157 : s%nzremap(2) = s%nzremap(2) + dbt_get_nze(tensor_contr_2)
1097 11157 : IF (do_crop_2) s%ncrop_remap(2) = s%ncrop_remap(2) + 1
1098 : END IF
1099 : END ASSOCIATE
1100 :
1101 : CALL dbt_tas_multiply(trans_1, trans_2, trans_3, alpha, &
1102 : tensor_contr_1%matrix_rep, tensor_contr_2%matrix_rep, &
1103 : beta, &
1104 : tensor_contr_3%matrix_rep, filter_eps=filter_eps, flop=flop_prv, &
1105 : unit_nr=unit_nr_prv, log_verbose=log_verbose, &
1106 : split_opt=split_opt, &
1107 222504 : move_data_a=move_data_1, move_data_b=move_data_2, retain_sparsity=retain_sparsity)
1108 222504 : IF (PRESENT(flop)) flop = flop_prv
1109 : ASSOCIATE (s => contract_stats(istat))
1110 222504 : s%flop = s%flop + flop_prv
1111 222504 : IF (new_3) THEN ! result is copied back into tensor_3
1112 23276 : s%nremap(3) = s%nremap(3) + 1
1113 23276 : s%nzremap(3) = s%nzremap(3) + dbt_get_nze(tensor_contr_3)
1114 : END IF
1115 : END ASSOCIATE
1116 :
1117 222504 : IF (PRESENT(pgrid_opt_1)) THEN
1118 0 : IF (.NOT. new_1) THEN
1119 0 : ALLOCATE (pgrid_opt_1)
1120 0 : pgrid_opt_1 = opt_pgrid(tensor_1, split_opt)
1121 : END IF
1122 : END IF
1123 :
1124 222504 : IF (PRESENT(pgrid_opt_2)) THEN
1125 0 : IF (.NOT. new_2) THEN
1126 0 : ALLOCATE (pgrid_opt_2)
1127 0 : pgrid_opt_2 = opt_pgrid(tensor_2, split_opt)
1128 : END IF
1129 : END IF
1130 :
1131 222504 : IF (PRESENT(pgrid_opt_3)) THEN
1132 0 : IF (.NOT. new_3) THEN
1133 0 : ALLOCATE (pgrid_opt_3)
1134 0 : pgrid_opt_3 = opt_pgrid(tensor_3, split_opt)
1135 : END IF
1136 : END IF
1137 :
1138 222504 : do_batched = tensor_small%matrix_rep%do_batched > 0
1139 :
1140 222504 : tensors_remapped = .FALSE.
1141 222504 : IF (new_1 .OR. new_2 .OR. new_3) tensors_remapped = .TRUE.
1142 :
1143 45627 : IF (tensors_remapped .AND. do_batched) THEN
1144 : CALL cp_warn(__LOCATION__, &
1145 0 : "Internal process grid optimization disabled because tensors are not in contraction-compatible format")
1146 : END IF
1147 :
1148 : ! optimize process grid during batched contraction
1149 222504 : do_change_pgrid(:) = .FALSE.
1150 222504 : IF ((.NOT. tensors_remapped) .AND. do_batched) THEN
1151 : ASSOCIATE (storage => tensor_small%contraction_storage)
1152 0 : CPASSERT(storage%static)
1153 93532 : split = dbt_tas_info(tensor_large%matrix_rep)
1154 : do_change_pgrid(:) = &
1155 93532 : update_contraction_storage(storage, split_opt, split)
1156 :
1157 372780 : IF (ANY(do_change_pgrid)) THEN
1158 674 : mp_comm_opt = dbt_tas_mp_comm(tensor_small%pgrid%mp_comm_2d, split_opt%split_rowcol, NINT(storage%nsplit_avg))
1159 : CALL dbt_tas_create_split(split_opt_avg, mp_comm_opt, split_opt%split_rowcol, &
1160 674 : NINT(storage%nsplit_avg), own_comm=.TRUE.)
1161 2022 : pdims_2d_opt = split_opt_avg%mp_comm%num_pe_cart
1162 : END IF
1163 :
1164 : END ASSOCIATE
1165 :
1166 93532 : IF (do_change_pgrid(1) .AND. .NOT. do_change_pgrid(2)) THEN
1167 : ! check if new grid has better subgrid, if not there is no need to change process grid
1168 2022 : pdims_sub_opt = split_opt_avg%mp_comm_group%num_pe_cart
1169 2022 : pdims_sub = split%mp_comm_group%num_pe_cart
1170 :
1171 2696 : pdim_ratio = MAXVAL(REAL(pdims_sub, dp))/MINVAL(pdims_sub)
1172 2696 : pdim_ratio_opt = MAXVAL(REAL(pdims_sub_opt, dp))/MINVAL(pdims_sub_opt)
1173 674 : IF (pdim_ratio/pdim_ratio_opt <= default_pdims_accept_ratio**2) THEN
1174 0 : do_change_pgrid(1) = .FALSE.
1175 0 : CALL dbt_tas_release_info(split_opt_avg)
1176 : END IF
1177 : END IF
1178 : END IF
1179 :
1180 222504 : IF (unit_nr_prv /= 0) THEN
1181 55120 : do_write_3 = .TRUE.
1182 55120 : IF (tensor_contr_3%matrix_rep%do_batched > 0) THEN
1183 23146 : IF (tensor_contr_3%matrix_rep%mm_storage%batched_out) do_write_3 = .FALSE.
1184 : END IF
1185 : IF (do_write_3) THEN
1186 32034 : CALL dbt_write_tensor_info(tensor_contr_3, unit_nr_prv, full_info=log_verbose)
1187 32034 : CALL dbt_write_tensor_dist(tensor_contr_3, unit_nr_prv)
1188 : END IF
1189 : END IF
1190 :
1191 222504 : IF (new_3) THEN
1192 : ! need redistribute if we created new tensor for tensor 3
1193 23276 : CALL dbt_scale(tensor_algn_3, beta)
1194 23276 : CALL dbt_copy_expert(tensor_contr_3, tensor_algn_3, summation=.TRUE., move_data=.TRUE.)
1195 23276 : IF (PRESENT(filter_eps)) CALL dbt_filter(tensor_algn_3, filter_eps)
1196 : ! tensor_3 automatically has correct data because tensor_algn_3 contains a matrix
1197 : ! pointer to data of tensor_3
1198 : END IF
1199 :
1200 : ! transfer contraction storage
1201 222504 : CALL dbt_copy_contraction_storage(tensor_contr_1, tensor_1)
1202 222504 : CALL dbt_copy_contraction_storage(tensor_contr_2, tensor_2)
1203 222504 : CALL dbt_copy_contraction_storage(tensor_contr_3, tensor_3)
1204 :
1205 222504 : IF (unit_nr_prv /= 0) THEN
1206 55120 : IF (new_3 .AND. do_write_3) CALL dbt_write_tensor_info(tensor_3, unit_nr_prv, full_info=log_verbose)
1207 55120 : IF (new_3 .AND. do_write_3) CALL dbt_write_tensor_dist(tensor_3, unit_nr_prv)
1208 : END IF
1209 :
1210 222504 : CALL dbt_destroy(tensor_algn_1)
1211 222504 : CALL dbt_destroy(tensor_algn_2)
1212 222504 : CALL dbt_destroy(tensor_algn_3)
1213 :
1214 222504 : IF (new_1 .OR. cropped_1) THEN
1215 58126 : CALL dbt_destroy(tensor_contr_1)
1216 58126 : DEALLOCATE (tensor_contr_1)
1217 : END IF
1218 222504 : IF (new_2 .OR. cropped_2) THEN
1219 51415 : CALL dbt_destroy(tensor_contr_2)
1220 51415 : DEALLOCATE (tensor_contr_2)
1221 : END IF
1222 222504 : IF (new_3) THEN
1223 23276 : CALL dbt_destroy(tensor_contr_3)
1224 23276 : DEALLOCATE (tensor_contr_3)
1225 : END IF
1226 :
1227 222504 : IF (PRESENT(move_data)) THEN
1228 56059 : IF (move_data) THEN
1229 39640 : CALL dbt_clear(tensor_1)
1230 39640 : CALL dbt_clear(tensor_2)
1231 : END IF
1232 : END IF
1233 :
1234 222504 : IF (unit_nr_prv > 0) THEN
1235 10 : WRITE (unit_nr_prv, '(A)') repeat("-", 80)
1236 10 : WRITE (unit_nr_prv, '(A)') "TENSOR CONTRACTION DONE"
1237 10 : WRITE (unit_nr_prv, '(A)') repeat("-", 80)
1238 : END IF
1239 :
1240 666164 : IF (ANY(do_change_pgrid)) THEN
1241 674 : pgrid_changed_any = .FALSE.
1242 0 : SELECT CASE (max_mm_dim)
1243 : CASE (1)
1244 0 : IF (ALLOCATED(tensor_1%contraction_storage) .AND. ALLOCATED(tensor_3%contraction_storage)) THEN
1245 : CALL dbt_change_pgrid_2d(tensor_1, tensor_1%pgrid%mp_comm_2d, pdims=pdims_2d_opt, &
1246 : nsplit=split_opt_avg%ngroup, dimsplit=split_opt_avg%split_rowcol, &
1247 : pgrid_changed=pgrid_changed, &
1248 0 : unit_nr=unit_nr_prv)
1249 0 : IF (pgrid_changed) pgrid_changed_any = .TRUE.
1250 : CALL dbt_change_pgrid_2d(tensor_3, tensor_3%pgrid%mp_comm_2d, pdims=pdims_2d_opt, &
1251 : nsplit=split_opt_avg%ngroup, dimsplit=split_opt_avg%split_rowcol, &
1252 : pgrid_changed=pgrid_changed, &
1253 0 : unit_nr=unit_nr_prv)
1254 0 : IF (pgrid_changed) pgrid_changed_any = .TRUE.
1255 : END IF
1256 0 : IF (pgrid_changed_any) THEN
1257 0 : IF (tensor_2%matrix_rep%do_batched == 3) THEN
1258 : ! set flag that process grid has been optimized to make sure that no grid optimizations are done
1259 : ! in TAS multiply algorithm
1260 0 : CALL dbt_tas_batched_mm_complete(tensor_2%matrix_rep)
1261 : END IF
1262 : END IF
1263 : CASE (2)
1264 184 : IF (ALLOCATED(tensor_1%contraction_storage) .AND. ALLOCATED(tensor_2%contraction_storage)) THEN
1265 : CALL dbt_change_pgrid_2d(tensor_1, tensor_1%pgrid%mp_comm_2d, pdims=pdims_2d_opt, &
1266 : nsplit=split_opt_avg%ngroup, dimsplit=split_opt_avg%split_rowcol, &
1267 : pgrid_changed=pgrid_changed, &
1268 184 : unit_nr=unit_nr_prv)
1269 184 : IF (pgrid_changed) pgrid_changed_any = .TRUE.
1270 : CALL dbt_change_pgrid_2d(tensor_2, tensor_2%pgrid%mp_comm_2d, pdims=pdims_2d_opt, &
1271 : nsplit=split_opt_avg%ngroup, dimsplit=split_opt_avg%split_rowcol, &
1272 : pgrid_changed=pgrid_changed, &
1273 184 : unit_nr=unit_nr_prv)
1274 184 : IF (pgrid_changed) pgrid_changed_any = .TRUE.
1275 : END IF
1276 8 : IF (pgrid_changed_any) THEN
1277 184 : IF (tensor_3%matrix_rep%do_batched == 3) THEN
1278 172 : CALL dbt_tas_batched_mm_complete(tensor_3%matrix_rep)
1279 : END IF
1280 : END IF
1281 : CASE (3)
1282 490 : IF (ALLOCATED(tensor_2%contraction_storage) .AND. ALLOCATED(tensor_3%contraction_storage)) THEN
1283 : CALL dbt_change_pgrid_2d(tensor_2, tensor_2%pgrid%mp_comm_2d, pdims=pdims_2d_opt, &
1284 : nsplit=split_opt_avg%ngroup, dimsplit=split_opt_avg%split_rowcol, &
1285 : pgrid_changed=pgrid_changed, &
1286 222 : unit_nr=unit_nr_prv)
1287 222 : IF (pgrid_changed) pgrid_changed_any = .TRUE.
1288 : CALL dbt_change_pgrid_2d(tensor_3, tensor_3%pgrid%mp_comm_2d, pdims=pdims_2d_opt, &
1289 : nsplit=split_opt_avg%ngroup, dimsplit=split_opt_avg%split_rowcol, &
1290 : pgrid_changed=pgrid_changed, &
1291 222 : unit_nr=unit_nr_prv)
1292 222 : IF (pgrid_changed) pgrid_changed_any = .TRUE.
1293 : END IF
1294 674 : IF (pgrid_changed_any) THEN
1295 222 : IF (tensor_1%matrix_rep%do_batched == 3) THEN
1296 222 : CALL dbt_tas_batched_mm_complete(tensor_1%matrix_rep)
1297 : END IF
1298 : END IF
1299 : END SELECT
1300 674 : CALL dbt_tas_release_info(split_opt_avg)
1301 : END IF
1302 :
1303 222504 : IF ((.NOT. tensors_remapped) .AND. do_batched) THEN
1304 : ! freeze TAS process grids if tensor grids were optimized
1305 93532 : CALL dbt_tas_set_batched_state(tensor_1%matrix_rep, opt_grid=.TRUE.)
1306 93532 : CALL dbt_tas_set_batched_state(tensor_2%matrix_rep, opt_grid=.TRUE.)
1307 93532 : CALL dbt_tas_set_batched_state(tensor_3%matrix_rep, opt_grid=.TRUE.)
1308 : END IF
1309 :
1310 222504 : CALL dbt_tas_release_info(split_opt)
1311 :
1312 222504 : CALL timestop(handle)
1313 :
1314 737344 : END SUBROUTINE dbt_contract_expert
1315 :
1316 : ! **************************************************************************************************
1317 : !> \brief align tensor index with data
1318 : !> \author Patrick Seewald
1319 : ! **************************************************************************************************
1320 6007608 : SUBROUTINE align_tensor(tensor_in, contract_in, notcontract_in, &
1321 667512 : tensor_out, contract_out, notcontract_out, indp_in, indp_out, order)
1322 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in
1323 : INTEGER, DIMENSION(:), INTENT(IN) :: contract_in, notcontract_in
1324 : TYPE(dbt_type), INTENT(OUT) :: tensor_out
1325 : INTEGER, DIMENSION(SIZE(contract_in)), &
1326 : INTENT(OUT) :: contract_out
1327 : INTEGER, DIMENSION(SIZE(notcontract_in)), &
1328 : INTENT(OUT) :: notcontract_out
1329 : CHARACTER(LEN=1), DIMENSION(ndims_tensor(tensor_in)), INTENT(IN) :: indp_in
1330 : CHARACTER(LEN=1), DIMENSION(ndims_tensor(tensor_in)), INTENT(OUT) :: indp_out
1331 : INTEGER, DIMENSION(ndims_tensor(tensor_in)), INTENT(OUT), OPTIONAL :: order
1332 667512 : INTEGER, DIMENSION(ndims_tensor(tensor_in)) :: align
1333 :
1334 667512 : CALL dbt_align_index(tensor_in, tensor_out, order=align)
1335 1845143 : IF (PRESENT(order)) order = align
1336 1531937 : contract_out = align(contract_in)
1337 1563707 : notcontract_out = align(notcontract_in)
1338 2428132 : indp_out(align) = indp_in
1339 :
1340 667512 : END SUBROUTINE align_tensor
1341 :
1342 : ! **************************************************************************************************
1343 : !> \brief Prepare tensor for contraction: redistribute to a 2d format which can be contracted by
1344 : !> matrix multiplication. This routine reshapes the two largest of the three tensors.
1345 : !> Redistribution is avoided if tensors already in a consistent layout.
1346 : !> \param ind1_free indices of tensor 1 that are "free" (not linked to any index of tensor 2)
1347 : !> \param ind1_linked indices of tensor 1 that are linked to indices of tensor 2
1348 : !> 1:1 correspondence with ind1_linked
1349 : !> \param trans1 transpose flag of matrix rep. tensor 1
1350 : !> \param trans2 transpose flag of matrix rep. tensor 2
1351 : !> \param new1 whether a new tensor 1 was created
1352 : !> \param new2 whether a new tensor 2 was created
1353 : !> \param nodata1 don't copy data of tensor 1
1354 : !> \param nodata2 don't copy data of tensor 2
1355 : !> \param move_data_1 memory optimization: transfer data s.t. tensor1 may be empty on return
1356 : !> \param move_data_2 memory optimization: transfer data s.t. tensor2 may be empty on return
1357 : !> \param optimize_dist experimental: optimize distribution
1358 : !> \param unit_nr output unit
1359 : !> \author Patrick Seewald
1360 : ! **************************************************************************************************
1361 222504 : SUBROUTINE reshape_mm_compatible(tensor1, tensor2, tensor1_out, tensor2_out, ind1_free, ind1_linked, &
1362 222504 : ind2_free, ind2_linked, trans1, trans2, new1, new2, ref_tensor, &
1363 : nodata1, nodata2, move_data_1, &
1364 : move_data_2, optimize_dist, unit_nr, bounds1, bounds2, crop1, crop2)
1365 : TYPE(dbt_type), TARGET, INTENT(INOUT) :: tensor1
1366 : TYPE(dbt_type), TARGET, INTENT(INOUT) :: tensor2
1367 : TYPE(dbt_type), POINTER, INTENT(OUT) :: tensor1_out, tensor2_out
1368 : INTEGER, DIMENSION(:), INTENT(IN) :: ind1_free, ind2_free
1369 : INTEGER, DIMENSION(:), INTENT(IN) :: ind1_linked, ind2_linked
1370 : LOGICAL, INTENT(OUT) :: trans1, trans2
1371 : LOGICAL, INTENT(OUT) :: new1, new2
1372 : INTEGER, INTENT(OUT) :: ref_tensor
1373 : LOGICAL, INTENT(IN), OPTIONAL :: nodata1, nodata2
1374 : LOGICAL, INTENT(INOUT), OPTIONAL :: move_data_1, move_data_2
1375 : LOGICAL, INTENT(IN), OPTIONAL :: optimize_dist
1376 : INTEGER, INTENT(IN), OPTIONAL :: unit_nr
1377 : INTEGER, DIMENSION(2, ndims_tensor(tensor1)), INTENT(IN), OPTIONAL :: bounds1
1378 : INTEGER, DIMENSION(2, ndims_tensor(tensor2)), INTENT(IN), OPTIONAL :: bounds2
1379 : LOGICAL, INTENT(OUT), OPTIONAL :: crop1, crop2
1380 : LOGICAL :: crop1_prv, crop2_prv
1381 : INTEGER :: compat1, compat1_old, compat2, compat2_old, &
1382 : unit_nr_prv
1383 222504 : TYPE(mp_cart_type) :: comm_2d
1384 222504 : TYPE(array_list) :: dist_list
1385 222504 : INTEGER, DIMENSION(:), ALLOCATABLE :: mp_dims
1386 1557528 : TYPE(dbt_distribution_type) :: dist_in
1387 : INTEGER(KIND=int_8) :: nblkrows, nblkcols
1388 : LOGICAL :: optimize_dist_prv
1389 445008 : INTEGER, DIMENSION(ndims_tensor(tensor1)) :: dims1
1390 222504 : INTEGER, DIMENSION(ndims_tensor(tensor2)) :: dims2
1391 :
1392 222504 : NULLIFY (tensor1_out, tensor2_out)
1393 :
1394 222504 : unit_nr_prv = prep_output_unit(unit_nr)
1395 :
1396 222504 : CALL blk_dims_tensor(tensor1, dims1)
1397 222504 : CALL blk_dims_tensor(tensor2, dims2)
1398 :
1399 1535832 : IF (PRODUCT(int(dims1, int_8)) >= PRODUCT(int(dims2, int_8))) THEN
1400 220370 : ref_tensor = 1
1401 : ELSE
1402 2134 : ref_tensor = 2
1403 : END IF
1404 :
1405 222504 : IF (PRESENT(optimize_dist)) THEN
1406 476 : optimize_dist_prv = optimize_dist
1407 : ELSE
1408 : optimize_dist_prv = .FALSE.
1409 : END IF
1410 :
1411 222504 : compat1 = compat_map(tensor1%nd_index, ind1_linked)
1412 222504 : compat2 = compat_map(tensor2%nd_index, ind2_linked)
1413 222504 : compat1_old = compat1
1414 222504 : compat2_old = compat2
1415 :
1416 222504 : IF (unit_nr_prv > 0) THEN
1417 10 : WRITE (unit_nr_prv, '(T2,A,1X,A,A,1X)', advance='no') "compatibility of", TRIM(tensor1%name), ":"
1418 6 : SELECT CASE (compat1)
1419 : CASE (0)
1420 6 : WRITE (unit_nr_prv, '(A)') "Not compatible"
1421 : CASE (1)
1422 3 : WRITE (unit_nr_prv, '(A)') "Normal"
1423 : CASE (2)
1424 10 : WRITE (unit_nr_prv, '(A)') "Transposed"
1425 : END SELECT
1426 10 : WRITE (unit_nr_prv, '(T2,A,1X,A,A,1X)', advance='no') "compatibility of", TRIM(tensor2%name), ":"
1427 5 : SELECT CASE (compat2)
1428 : CASE (0)
1429 5 : WRITE (unit_nr_prv, '(A)') "Not compatible"
1430 : CASE (1)
1431 4 : WRITE (unit_nr_prv, '(A)') "Normal"
1432 : CASE (2)
1433 10 : WRITE (unit_nr_prv, '(A)') "Transposed"
1434 : END SELECT
1435 : END IF
1436 :
1437 222504 : new1 = .FALSE.
1438 222504 : new2 = .FALSE.
1439 222504 : crop1_prv = .FALSE.
1440 222504 : crop2_prv = .FALSE.
1441 :
1442 222504 : IF (compat1 == 0 .OR. optimize_dist_prv) THEN
1443 34103 : new1 = .TRUE.
1444 : END IF
1445 :
1446 222504 : IF (compat2 == 0 .OR. optimize_dist_prv) THEN
1447 23456 : new2 = .TRUE.
1448 : END IF
1449 :
1450 222504 : IF (ref_tensor == 1) THEN ! tensor 1 is reference and tensor 2 is reshaped compatible with tensor 1
1451 220370 : IF (compat1 == 0 .OR. optimize_dist_prv) THEN ! tensor 1 is not contraction compatible --> reshape
1452 33833 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "Redistribution of", TRIM(tensor1%name)
1453 101499 : nblkrows = PRODUCT(INT(dims1(ind1_linked), KIND=int_8))
1454 67668 : nblkcols = PRODUCT(INT(dims1(ind1_free), KIND=int_8))
1455 33833 : comm_2d = dbt_tas_mp_comm(tensor1%pgrid%mp_comm_2d, nblkrows, nblkcols)
1456 236831 : ALLOCATE (tensor1_out)
1457 : CALL dbt_remap(tensor1, ind1_linked, ind1_free, tensor1_out, comm_2d=comm_2d, &
1458 : nodata=nodata1, move_data=move_data_1, &
1459 33833 : bounds=bounds1)
1460 33833 : CALL comm_2d%free()
1461 33833 : compat1 = 1
1462 : ELSE
1463 186537 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "No redistribution of", TRIM(tensor1%name)
1464 186537 : CALL crop_or_point(tensor1, tensor1_out, crop1_prv, bounds1, move_data_1)
1465 : END IF
1466 220370 : IF (compat2 == 0 .OR. optimize_dist_prv) THEN ! tensor 2 is not contraction compatible --> reshape
1467 23126 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A,1X,A,1X,A)') "Redistribution of", &
1468 8 : TRIM(tensor2%name), "compatible with", TRIM(tensor1%name)
1469 23122 : dist_in = dbt_distribution(tensor1_out)
1470 23122 : dist_list = array_sublist(dist_in%nd_dist, ind1_linked)
1471 23122 : IF (compat1 == 1) THEN ! linked index is first 2d dimension
1472 : ! get distribution of linked index, tensor 2 must adopt this distribution
1473 : ! get grid dimensions of linked index
1474 35220 : ALLOCATE (mp_dims(ndims_mapping_row(dist_in%pgrid%nd_index_grid)))
1475 11740 : CALL dbt_get_mapping_info(dist_in%pgrid%nd_index_grid, dims1_2d=mp_dims)
1476 82180 : ALLOCATE (tensor2_out)
1477 : CALL dbt_remap(tensor2, ind2_linked, ind2_free, tensor2_out, comm_2d=dist_in%pgrid%mp_comm_2d, &
1478 : dist1=dist_list, mp_dims_1=mp_dims, nodata=nodata2, move_data=move_data_2, &
1479 11740 : bounds=bounds2)
1480 11382 : ELSE IF (compat1 == 2) THEN ! linked index is second 2d dimension
1481 : ! get distribution of linked index, tensor 2 must adopt this distribution
1482 : ! get grid dimensions of linked index
1483 34146 : ALLOCATE (mp_dims(ndims_mapping_column(dist_in%pgrid%nd_index_grid)))
1484 11382 : CALL dbt_get_mapping_info(dist_in%pgrid%nd_index_grid, dims2_2d=mp_dims)
1485 79674 : ALLOCATE (tensor2_out)
1486 : CALL dbt_remap(tensor2, ind2_free, ind2_linked, tensor2_out, comm_2d=dist_in%pgrid%mp_comm_2d, &
1487 : dist2=dist_list, mp_dims_2=mp_dims, nodata=nodata2, move_data=move_data_2, &
1488 11382 : bounds=bounds2)
1489 : ELSE
1490 0 : CPABORT("should not happen")
1491 : END IF
1492 23122 : compat2 = compat1
1493 : ELSE
1494 197248 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "No redistribution of", TRIM(tensor2%name)
1495 197248 : CALL crop_or_point(tensor2, tensor2_out, crop2_prv, bounds2, move_data_2)
1496 : END IF
1497 : ELSE ! tensor 2 is reference and tensor 1 is reshaped compatible with tensor 2
1498 2134 : IF (compat2 == 0 .OR. optimize_dist_prv) THEN ! tensor 2 is not contraction compatible --> reshape
1499 334 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "Redistribution of", TRIM(tensor2%name)
1500 754 : nblkrows = PRODUCT(INT(dims2(ind2_linked), KIND=int_8))
1501 730 : nblkcols = PRODUCT(INT(dims2(ind2_free), KIND=int_8))
1502 334 : comm_2d = dbt_tas_mp_comm(tensor2%pgrid%mp_comm_2d, nblkrows, nblkcols)
1503 2338 : ALLOCATE (tensor2_out)
1504 : CALL dbt_remap(tensor2, ind2_linked, ind2_free, tensor2_out, nodata=nodata2, move_data=move_data_2, &
1505 334 : bounds=bounds2)
1506 334 : CALL comm_2d%free()
1507 334 : compat2 = 1
1508 : ELSE
1509 1800 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "No redistribution of", TRIM(tensor2%name)
1510 1800 : CALL crop_or_point(tensor2, tensor2_out, crop2_prv, bounds2, move_data_2)
1511 : END IF
1512 2134 : IF (compat1 == 0 .OR. optimize_dist_prv) THEN ! tensor 1 is not contraction compatible --> reshape
1513 273 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A,1X,A,1X,A)') "Redistribution of", TRIM(tensor1%name), &
1514 6 : "compatible with", TRIM(tensor2%name)
1515 270 : dist_in = dbt_distribution(tensor2_out)
1516 270 : dist_list = array_sublist(dist_in%nd_dist, ind2_linked)
1517 270 : IF (compat2 == 1) THEN
1518 804 : ALLOCATE (mp_dims(ndims_mapping_row(dist_in%pgrid%nd_index_grid)))
1519 268 : CALL dbt_get_mapping_info(dist_in%pgrid%nd_index_grid, dims1_2d=mp_dims)
1520 1876 : ALLOCATE (tensor1_out)
1521 : CALL dbt_remap(tensor1, ind1_linked, ind1_free, tensor1_out, comm_2d=dist_in%pgrid%mp_comm_2d, &
1522 : dist1=dist_list, mp_dims_1=mp_dims, nodata=nodata1, move_data=move_data_1, &
1523 268 : bounds=bounds1)
1524 2 : ELSE IF (compat2 == 2) THEN
1525 6 : ALLOCATE (mp_dims(ndims_mapping_column(dist_in%pgrid%nd_index_grid)))
1526 2 : CALL dbt_get_mapping_info(dist_in%pgrid%nd_index_grid, dims2_2d=mp_dims)
1527 14 : ALLOCATE (tensor1_out)
1528 : CALL dbt_remap(tensor1, ind1_free, ind1_linked, tensor1_out, comm_2d=dist_in%pgrid%mp_comm_2d, &
1529 : dist2=dist_list, mp_dims_2=mp_dims, nodata=nodata1, move_data=move_data_1, &
1530 2 : bounds=bounds1)
1531 : ELSE
1532 0 : CPABORT("should not happen")
1533 : END IF
1534 270 : compat1 = compat2
1535 : ELSE
1536 1864 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "No redistribution of", TRIM(tensor1%name)
1537 1864 : CALL crop_or_point(tensor1, tensor1_out, crop1_prv, bounds1, move_data_1)
1538 : END IF
1539 : END IF
1540 :
1541 156549 : SELECT CASE (compat1)
1542 : CASE (1)
1543 156549 : trans1 = .FALSE.
1544 : CASE (2)
1545 65955 : trans1 = .TRUE.
1546 : CASE DEFAULT
1547 222504 : CPABORT("should not happen")
1548 : END SELECT
1549 :
1550 149701 : SELECT CASE (compat2)
1551 : CASE (1)
1552 149701 : trans2 = .FALSE.
1553 : CASE (2)
1554 72803 : trans2 = .TRUE.
1555 : CASE DEFAULT
1556 222504 : CPABORT("should not happen")
1557 : END SELECT
1558 :
1559 222504 : IF (unit_nr_prv > 0) THEN
1560 10 : IF (compat1 /= compat1_old) THEN
1561 6 : WRITE (unit_nr_prv, '(T2,A,1X,A,A,1X)', advance='no') "compatibility of", TRIM(tensor1_out%name), ":"
1562 0 : SELECT CASE (compat1)
1563 : CASE (0)
1564 0 : WRITE (unit_nr_prv, '(A)') "Not compatible"
1565 : CASE (1)
1566 5 : WRITE (unit_nr_prv, '(A)') "Normal"
1567 : CASE (2)
1568 6 : WRITE (unit_nr_prv, '(A)') "Transposed"
1569 : END SELECT
1570 : END IF
1571 10 : IF (compat2 /= compat2_old) THEN
1572 5 : WRITE (unit_nr_prv, '(T2,A,1X,A,A,1X)', advance='no') "compatibility of", TRIM(tensor2_out%name), ":"
1573 0 : SELECT CASE (compat2)
1574 : CASE (0)
1575 0 : WRITE (unit_nr_prv, '(A)') "Not compatible"
1576 : CASE (1)
1577 4 : WRITE (unit_nr_prv, '(A)') "Normal"
1578 : CASE (2)
1579 5 : WRITE (unit_nr_prv, '(A)') "Transposed"
1580 : END SELECT
1581 : END IF
1582 : END IF
1583 :
1584 222504 : IF (new1 .AND. PRESENT(move_data_1)) move_data_1 = .TRUE.
1585 222504 : IF (new2 .AND. PRESENT(move_data_2)) move_data_2 = .TRUE.
1586 :
1587 222504 : IF (PRESENT(crop1)) crop1 = crop1_prv
1588 222504 : IF (PRESENT(crop2)) crop2 = crop2_prv
1589 :
1590 222504 : END SUBROUTINE reshape_mm_compatible
1591 :
1592 : ! **************************************************************************************************
1593 : !> \brief Prepare tensor for contraction: redistribute to a 2d format which can be contracted by
1594 : !> matrix multiplication. This routine reshapes the smallest of the three tensors.
1595 : !> \param ind1 index that should be mapped to first matrix dimension
1596 : !> \param ind2 index that should be mapped to second matrix dimension
1597 : !> \param trans transpose flag of matrix rep.
1598 : !> \param new whether a new tensor was created for tensor_out
1599 : !> \param nodata don't copy tensor data
1600 : !> \param move_data memory optimization: transfer data s.t. tensor_in may be empty on return
1601 : !> \param unit_nr output unit
1602 : !> \author Patrick Seewald
1603 : ! **************************************************************************************************
1604 222504 : SUBROUTINE reshape_mm_small(tensor_in, ind1, ind2, tensor_out, trans, new, nodata, move_data, unit_nr, &
1605 : bounds, crop)
1606 : TYPE(dbt_type), TARGET, INTENT(INOUT) :: tensor_in
1607 : INTEGER, DIMENSION(:), INTENT(IN) :: ind1, ind2
1608 : TYPE(dbt_type), POINTER, INTENT(OUT) :: tensor_out
1609 : LOGICAL, INTENT(OUT) :: trans
1610 : LOGICAL, INTENT(OUT) :: new
1611 : LOGICAL, INTENT(IN), OPTIONAL :: nodata
1612 : LOGICAL, INTENT(INOUT), OPTIONAL :: move_data
1613 : INTEGER, INTENT(IN), OPTIONAL :: unit_nr
1614 : INTEGER, DIMENSION(2, ndims_tensor(tensor_in)), INTENT(IN), OPTIONAL :: bounds
1615 : LOGICAL, INTENT(OUT), OPTIONAL :: crop
1616 : LOGICAL :: crop_prv
1617 : INTEGER :: compat1, compat2, compat1_old, compat2_old, unit_nr_prv
1618 : LOGICAL :: nodata_prv
1619 :
1620 222504 : NULLIFY (tensor_out)
1621 : IF (PRESENT(nodata)) THEN
1622 222504 : nodata_prv = nodata
1623 : ELSE
1624 222504 : nodata_prv = .FALSE.
1625 : END IF
1626 :
1627 222504 : unit_nr_prv = prep_output_unit(unit_nr)
1628 :
1629 222504 : new = .FALSE.
1630 222504 : crop_prv = .FALSE.
1631 222504 : compat1 = compat_map(tensor_in%nd_index, ind1)
1632 222504 : compat2 = compat_map(tensor_in%nd_index, ind2)
1633 222504 : compat1_old = compat1; compat2_old = compat2
1634 222504 : IF (unit_nr_prv > 0) THEN
1635 10 : WRITE (unit_nr_prv, '(T2,A,1X,A,A,1X)', advance='no') "compatibility of", TRIM(tensor_in%name), ":"
1636 10 : IF (compat1 == 1 .AND. compat2 == 2) THEN
1637 4 : WRITE (unit_nr_prv, '(A)') "Normal"
1638 6 : ELSE IF (compat1 == 2 .AND. compat2 == 1) THEN
1639 2 : WRITE (unit_nr_prv, '(A)') "Transposed"
1640 : ELSE
1641 4 : WRITE (unit_nr_prv, '(A)') "Not compatible"
1642 : END IF
1643 : END IF
1644 222504 : IF (compat1 == 0 .or. compat2 == 0) THEN ! index mapping not compatible with contract index
1645 :
1646 102 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "Redistribution of", TRIM(tensor_in%name)
1647 :
1648 714 : ALLOCATE (tensor_out)
1649 102 : CALL dbt_remap(tensor_in, ind1, ind2, tensor_out, nodata=nodata, move_data=move_data, bounds=bounds)
1650 102 : CALL dbt_copy_contraction_storage(tensor_in, tensor_out)
1651 102 : compat1 = 1
1652 102 : compat2 = 2
1653 102 : new = .TRUE.
1654 : ELSE
1655 222402 : IF (unit_nr_prv > 0) WRITE (unit_nr_prv, '(T2,A,1X,A)') "No redistribution of", TRIM(tensor_in%name)
1656 222402 : CALL crop_or_point(tensor_in, tensor_out, crop_prv, bounds, move_data)
1657 : END IF
1658 :
1659 222504 : IF (compat1 == 1 .AND. compat2 == 2) THEN
1660 151378 : trans = .FALSE.
1661 71126 : ELSE IF (compat1 == 2 .AND. compat2 == 1) THEN
1662 71126 : trans = .TRUE.
1663 : ELSE
1664 0 : CPABORT("this should not happen")
1665 : END IF
1666 :
1667 222504 : IF (unit_nr_prv > 0) THEN
1668 10 : IF (compat1_old /= compat1 .OR. compat2_old /= compat2) THEN
1669 4 : WRITE (unit_nr_prv, '(T2,A,1X,A,A,1X)', advance='no') "compatibility of", TRIM(tensor_out%name), ":"
1670 4 : IF (compat1 == 1 .AND. compat2 == 2) THEN
1671 4 : WRITE (unit_nr_prv, '(A)') "Normal"
1672 0 : ELSE IF (compat1 == 2 .AND. compat2 == 1) THEN
1673 0 : WRITE (unit_nr_prv, '(A)') "Transposed"
1674 : ELSE
1675 0 : WRITE (unit_nr_prv, '(A)') "Not compatible"
1676 : END IF
1677 : END IF
1678 : END IF
1679 :
1680 222504 : IF (PRESENT(crop)) crop = crop_prv
1681 :
1682 222504 : END SUBROUTINE reshape_mm_small
1683 :
1684 : ! **************************************************************************************************
1685 : !> \brief Crop tensor_in into a new tensor_out if bounds are present, otherwise point to tensor_in.
1686 : !> \param cropped whether tensor_out is a new tensor
1687 : !> \param move_data memory optimization: transfer data s.t. tensor_in may be empty on return
1688 : !> \author Hans Pabst
1689 : ! **************************************************************************************************
1690 609851 : SUBROUTINE crop_or_point(tensor_in, tensor_out, cropped, bounds, move_data)
1691 : TYPE(dbt_type), TARGET, INTENT(INOUT) :: tensor_in
1692 : TYPE(dbt_type), POINTER, INTENT(OUT) :: tensor_out
1693 : LOGICAL, INTENT(OUT) :: cropped
1694 : INTEGER, DIMENSION(2, ndims_tensor(tensor_in)), INTENT(IN), OPTIONAL :: bounds
1695 : LOGICAL, INTENT(INOUT), OPTIONAL :: move_data
1696 :
1697 609851 : cropped = PRESENT(bounds)
1698 609851 : IF (cropped) THEN
1699 751560 : ALLOCATE (tensor_out)
1700 75156 : CALL dbt_crop(tensor_in, tensor_out, bounds, move_data=move_data)
1701 75156 : IF (PRESENT(move_data)) move_data = .TRUE. ! tensor_out is temporary
1702 : ELSE
1703 534695 : tensor_out => tensor_in
1704 : END IF
1705 609851 : END SUBROUTINE crop_or_point
1706 :
1707 : ! **************************************************************************************************
1708 : !> \brief Number of local blocks overlapping the bounds (blocks dbt_crop would keep), and number
1709 : !> of local blocks not entirely within the bounds (none: dbt_crop would copy as is).
1710 : !> \author Hans Pabst
1711 : ! **************************************************************************************************
1712 339604 : FUNCTION num_blocks_in_bounds(tensor, bounds) RESULT(nblocks)
1713 : TYPE(dbt_type), INTENT(INOUT) :: tensor
1714 : INTEGER, DIMENSION(2, ndims_tensor(tensor)), INTENT(IN) :: bounds
1715 : INTEGER, DIMENSION(2) :: nblocks
1716 :
1717 : TYPE(dbt_iterator_type) :: iter
1718 169802 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: blk_index, blk_offset, blk_size
1719 : INTEGER :: nin, nout
1720 :
1721 169802 : nin = 0
1722 169802 : nout = 0
1723 : !$OMP PARALLEL DEFAULT(NONE) SHARED(tensor,bounds) PRIVATE(iter,blk_index,blk_size,blk_offset) &
1724 169802 : !$OMP REDUCTION(+:nin,nout)
1725 : CALL dbt_iterator_start(iter, tensor)
1726 : DO WHILE (dbt_iterator_blocks_left(iter))
1727 : CALL dbt_iterator_next_block(iter, blk_index, blk_size=blk_size, blk_offset=blk_offset)
1728 : IF (ALL(blk_offset <= bounds(2, :) .AND. bounds(1, :) <= blk_offset + blk_size - 1)) THEN
1729 : nin = nin + 1
1730 : END IF
1731 : IF (ANY(blk_offset < bounds(1, :) .OR. bounds(2, :) < blk_offset + blk_size - 1)) THEN
1732 : nout = nout + 1
1733 : END IF
1734 : END DO
1735 : CALL dbt_iterator_stop(iter)
1736 : !$OMP END PARALLEL
1737 509406 : nblocks = [nin, nout]
1738 169802 : END FUNCTION num_blocks_in_bounds
1739 :
1740 : ! **************************************************************************************************
1741 : !> \brief update contraction storage that keeps track of process grids during a batched contraction
1742 : !> and decide if tensor process grid needs to be optimized
1743 : !> \param split_opt optimized TAS process grid
1744 : !> \param split current TAS process grid
1745 : !> \author Patrick Seewald
1746 : ! **************************************************************************************************
1747 93532 : FUNCTION update_contraction_storage(storage, split_opt, split) RESULT(do_change_pgrid)
1748 : TYPE(dbt_contraction_storage), INTENT(INOUT) :: storage
1749 : TYPE(dbt_tas_split_info), INTENT(IN) :: split_opt
1750 : TYPE(dbt_tas_split_info), INTENT(IN) :: split
1751 : INTEGER, DIMENSION(2) :: pdims, pdims_sub
1752 : LOGICAL, DIMENSION(2) :: do_change_pgrid
1753 : REAL(kind=dp) :: change_criterion, pdims_ratio
1754 : INTEGER :: nsplit_opt, nsplit
1755 :
1756 93532 : CPASSERT(ALLOCATED(split_opt%ngroup_opt))
1757 93532 : nsplit_opt = split_opt%ngroup_opt
1758 93532 : nsplit = split%ngroup
1759 :
1760 280596 : pdims = split%mp_comm%num_pe_cart
1761 :
1762 93532 : storage%ibatch = storage%ibatch + 1
1763 :
1764 : storage%nsplit_avg = (storage%nsplit_avg*REAL(storage%ibatch - 1, dp) + REAL(nsplit_opt, dp)) &
1765 93532 : /REAL(storage%ibatch, dp)
1766 :
1767 93532 : SELECT CASE (split_opt%split_rowcol)
1768 : CASE (rowsplit)
1769 93532 : pdims_ratio = REAL(pdims(1), dp)/pdims(2)
1770 : CASE (colsplit)
1771 93532 : pdims_ratio = REAL(pdims(2), dp)/pdims(1)
1772 : END SELECT
1773 :
1774 280596 : do_change_pgrid(:) = .FALSE.
1775 :
1776 : ! check for process grid dimensions
1777 280596 : pdims_sub = split%mp_comm_group%num_pe_cart
1778 467660 : change_criterion = MAXVAL(REAL(pdims_sub, dp))/MINVAL(pdims_sub)
1779 93532 : IF (change_criterion > default_pdims_accept_ratio**2) do_change_pgrid(1) = .TRUE.
1780 :
1781 : ! check for split factor
1782 93532 : change_criterion = MAX(REAL(nsplit, dp)/storage%nsplit_avg, REAL(storage%nsplit_avg, dp)/nsplit)
1783 93532 : IF (change_criterion > default_nsplit_accept_ratio) do_change_pgrid(2) = .TRUE.
1784 :
1785 93532 : END FUNCTION update_contraction_storage
1786 :
1787 : ! **************************************************************************************************
1788 : !> \brief Check if 2d index is compatible with tensor index
1789 : !> \author Patrick Seewald
1790 : ! **************************************************************************************************
1791 890016 : FUNCTION compat_map(nd_index, compat_ind)
1792 : TYPE(nd_to_2d_mapping), INTENT(IN) :: nd_index
1793 : INTEGER, DIMENSION(:), INTENT(IN) :: compat_ind
1794 1780032 : INTEGER, DIMENSION(ndims_mapping_row(nd_index)) :: map1
1795 1780032 : INTEGER, DIMENSION(ndims_mapping_column(nd_index)) :: map2
1796 : INTEGER :: compat_map
1797 :
1798 890016 : CALL dbt_get_mapping_info(nd_index, map1_2d=map1, map2_2d=map2)
1799 :
1800 890016 : compat_map = 0
1801 890016 : IF (array_eq_i(map1, compat_ind)) THEN
1802 : compat_map = 1
1803 407349 : ELSE IF (array_eq_i(map2, compat_ind)) THEN
1804 350548 : compat_map = 2
1805 : END IF
1806 :
1807 890016 : END FUNCTION compat_map
1808 :
1809 : ! **************************************************************************************************
1810 : !> \brief
1811 : !> \author Patrick Seewald
1812 : ! **************************************************************************************************
1813 667512 : SUBROUTINE index_linked_sort(ind_ref, ind_dep)
1814 : INTEGER, DIMENSION(:), INTENT(INOUT) :: ind_ref, ind_dep
1815 1335024 : INTEGER, DIMENSION(SIZE(ind_ref)) :: sort_indices
1816 :
1817 667512 : CALL sort(ind_ref, SIZE(ind_ref), sort_indices)
1818 3095644 : ind_dep(:) = ind_dep(sort_indices)
1819 :
1820 667512 : END SUBROUTINE index_linked_sort
1821 :
1822 : ! **************************************************************************************************
1823 : !> \brief
1824 : !> \author Patrick Seewald
1825 : ! **************************************************************************************************
1826 0 : FUNCTION opt_pgrid(tensor, tas_split_info)
1827 : TYPE(dbt_type), INTENT(IN) :: tensor
1828 : TYPE(dbt_tas_split_info), INTENT(IN) :: tas_split_info
1829 0 : INTEGER, DIMENSION(ndims_matrix_row(tensor)) :: map1
1830 0 : INTEGER, DIMENSION(ndims_matrix_column(tensor)) :: map2
1831 : TYPE(dbt_pgrid_type) :: opt_pgrid
1832 0 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: dims
1833 :
1834 0 : CALL dbt_get_mapping_info(tensor%pgrid%nd_index_grid, map1_2d=map1, map2_2d=map2)
1835 0 : CALL blk_dims_tensor(tensor, dims)
1836 0 : opt_pgrid = dbt_nd_mp_comm(tas_split_info%mp_comm, map1, map2, tdims=dims)
1837 :
1838 0 : ALLOCATE (opt_pgrid%tas_split_info, SOURCE=tas_split_info)
1839 0 : CALL dbt_tas_info_hold(opt_pgrid%tas_split_info)
1840 0 : END FUNCTION opt_pgrid
1841 :
1842 : ! **************************************************************************************************
1843 : !> \brief Copy tensor to tensor with modified index mapping
1844 : !> \param map1_2d new index mapping
1845 : !> \param map2_2d new index mapping
1846 : !> \author Patrick Seewald
1847 : ! **************************************************************************************************
1848 518949 : SUBROUTINE dbt_remap(tensor_in, map1_2d, map2_2d, tensor_out, comm_2d, dist1, dist2, &
1849 57661 : mp_dims_1, mp_dims_2, name, nodata, move_data, bounds)
1850 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in
1851 : INTEGER, DIMENSION(2, ndims_tensor(tensor_in)), INTENT(IN), OPTIONAL :: bounds
1852 : INTEGER, DIMENSION(:), INTENT(IN) :: map1_2d, map2_2d
1853 : TYPE(dbt_type), INTENT(OUT) :: tensor_out
1854 : CHARACTER(len=*), INTENT(IN), OPTIONAL :: name
1855 : LOGICAL, INTENT(IN), OPTIONAL :: nodata, move_data
1856 : CLASS(mp_comm_type), INTENT(IN), OPTIONAL :: comm_2d
1857 : TYPE(array_list), INTENT(IN), OPTIONAL :: dist1, dist2
1858 : INTEGER, DIMENSION(SIZE(map1_2d)), OPTIONAL :: mp_dims_1
1859 : INTEGER, DIMENSION(SIZE(map2_2d)), OPTIONAL :: mp_dims_2
1860 : CHARACTER(len=default_string_length) :: name_tmp
1861 57661 : INTEGER, DIMENSION(:), ALLOCATABLE :: ${varlist("blk_sizes")}$, &
1862 57661 : ${varlist("nd_dist")}$
1863 403627 : TYPE(dbt_distribution_type) :: dist
1864 57661 : TYPE(mp_cart_type) :: comm_2d_prv
1865 : INTEGER :: handle, i
1866 57661 : INTEGER, DIMENSION(ndims_tensor(tensor_in)) :: pdims, myploc
1867 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_remap'
1868 : LOGICAL :: nodata_prv
1869 172983 : TYPE(dbt_pgrid_type) :: comm_nd
1870 :
1871 57661 : CALL timeset(routineN, handle)
1872 :
1873 57661 : IF (PRESENT(name)) THEN
1874 0 : name_tmp = name
1875 : ELSE
1876 57661 : name_tmp = tensor_in%name
1877 : END IF
1878 57661 : IF (PRESENT(dist1)) THEN
1879 12008 : CPASSERT(PRESENT(mp_dims_1))
1880 : END IF
1881 :
1882 57661 : IF (PRESENT(dist2)) THEN
1883 11384 : CPASSERT(PRESENT(mp_dims_2))
1884 : END IF
1885 :
1886 57661 : IF (PRESENT(comm_2d)) THEN
1887 57225 : comm_2d_prv = comm_2d
1888 : ELSE
1889 436 : comm_2d_prv = tensor_in%pgrid%mp_comm_2d
1890 : END IF
1891 :
1892 57661 : comm_nd = dbt_nd_mp_comm(comm_2d_prv, map1_2d, map2_2d, dims1_nd=mp_dims_1, dims2_nd=mp_dims_2)
1893 57661 : CALL mp_environ_pgrid(comm_nd, pdims, myploc)
1894 :
1895 : #:for ndim in ndims
1896 115318 : IF (ndims_tensor(tensor_in) == ${ndim}$) THEN
1897 57661 : CALL get_arrays(tensor_in%blk_sizes, ${varlist("blk_sizes", nmax=ndim)}$)
1898 : END IF
1899 : #:endfor
1900 :
1901 : #:for ndim in ndims
1902 115318 : IF (ndims_tensor(tensor_in) == ${ndim}$) THEN
1903 : #:for idim in range(1, ndim+1)
1904 172799 : IF (PRESENT(dist1)) THEN
1905 71860 : IF (ANY(map1_2d == ${idim}$)) THEN
1906 71296 : i = MINLOC(map1_2d, dim=1, mask=map1_2d == ${idim}$) ! i is location of idim in map1_2d
1907 23828 : CALL get_ith_array(dist1, i, nd_dist_${idim}$)
1908 : END IF
1909 : END IF
1910 :
1911 172799 : IF (PRESENT(dist2)) THEN
1912 68276 : IF (ANY(map2_2d == ${idim}$)) THEN
1913 68192 : i = MINLOC(map2_2d, dim=1, mask=map2_2d == ${idim}$) ! i is location of idim in map2_2d
1914 22740 : CALL get_ith_array(dist2, i, nd_dist_${idim}$)
1915 : END IF
1916 : END IF
1917 :
1918 172799 : IF (.NOT. ALLOCATED(nd_dist_${idim}$)) THEN
1919 378693 : ALLOCATE (nd_dist_${idim}$ (SIZE(blk_sizes_${idim}$)))
1920 126231 : CALL dbt_default_distvec(SIZE(blk_sizes_${idim}$), pdims(${idim}$), blk_sizes_${idim}$, nd_dist_${idim}$)
1921 : END IF
1922 : #:endfor
1923 : CALL dbt_distribution_new_expert(dist, comm_nd, map1_2d, map2_2d, &
1924 57661 : ${varlist("nd_dist", nmax=ndim)}$, own_comm=.TRUE.)
1925 : END IF
1926 : #:endfor
1927 :
1928 : #:for ndim in ndims
1929 115318 : IF (ndims_tensor(tensor_in) == ${ndim}$) THEN
1930 : CALL dbt_create(tensor_out, name_tmp, dist, map1_2d, map2_2d, &
1931 57661 : ${varlist("blk_sizes", nmax=ndim)}$)
1932 : END IF
1933 : #:endfor
1934 :
1935 57661 : IF (PRESENT(nodata)) THEN
1936 23276 : nodata_prv = nodata
1937 : ELSE
1938 : nodata_prv = .FALSE.
1939 : END IF
1940 :
1941 57661 : IF (.NOT. nodata_prv) CALL dbt_copy_expert(tensor_in, tensor_out, move_data=move_data, bounds=bounds)
1942 57661 : CALL dbt_distribution_destroy(dist)
1943 :
1944 57661 : CALL timestop(handle)
1945 172983 : END SUBROUTINE dbt_remap
1946 :
1947 : ! **************************************************************************************************
1948 : !> \brief Align index with data
1949 : !> \param order permutation resulting from alignment
1950 : !> \author Patrick Seewald
1951 : ! **************************************************************************************************
1952 5340096 : SUBROUTINE dbt_align_index(tensor_in, tensor_out, order)
1953 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in
1954 : TYPE(dbt_type), INTENT(OUT) :: tensor_out
1955 1335024 : INTEGER, DIMENSION(ndims_matrix_row(tensor_in)) :: map1_2d
1956 1335024 : INTEGER, DIMENSION(ndims_matrix_column(tensor_in)) :: map2_2d
1957 : INTEGER, DIMENSION(ndims_tensor(tensor_in)), &
1958 : INTENT(OUT), OPTIONAL :: order
1959 667512 : INTEGER, DIMENSION(ndims_tensor(tensor_in)) :: order_prv
1960 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_align_index'
1961 : INTEGER :: handle
1962 :
1963 667512 : CALL timeset(routineN, handle)
1964 :
1965 667512 : CALL dbt_get_mapping_info(tensor_in%nd_index_blk, map1_2d=map1_2d, map2_2d=map2_2d)
1966 4188752 : order_prv = dbt_inverse_order([map1_2d, map2_2d])
1967 667512 : CALL dbt_permute_index(tensor_in, tensor_out, order=order_prv)
1968 :
1969 2428132 : IF (PRESENT(order)) order = order_prv
1970 :
1971 667512 : CALL timestop(handle)
1972 667512 : END SUBROUTINE dbt_align_index
1973 :
1974 : ! **************************************************************************************************
1975 : !> \brief Create new tensor by reordering index, data is copied exactly (shallow copy)
1976 : !> \author Patrick Seewald
1977 : ! **************************************************************************************************
1978 7589133 : SUBROUTINE dbt_permute_index(tensor_in, tensor_out, order)
1979 : TYPE(dbt_type), INTENT(INOUT) :: tensor_in
1980 : TYPE(dbt_type), INTENT(OUT) :: tensor_out
1981 : INTEGER, DIMENSION(ndims_tensor(tensor_in)), &
1982 : INTENT(IN) :: order
1983 :
1984 4216185 : TYPE(nd_to_2d_mapping) :: nd_index_blk_rs, nd_index_rs
1985 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_permute_index'
1986 : INTEGER :: handle
1987 : INTEGER :: ndims
1988 :
1989 843237 : CALL timeset(routineN, handle)
1990 :
1991 843237 : ndims = ndims_tensor(tensor_in)
1992 :
1993 843237 : CALL permute_index(tensor_in%nd_index, nd_index_rs, order)
1994 843237 : CALL permute_index(tensor_in%nd_index_blk, nd_index_blk_rs, order)
1995 843237 : CALL permute_index(tensor_in%pgrid%nd_index_grid, tensor_out%pgrid%nd_index_grid, order)
1996 :
1997 843237 : tensor_out%matrix_rep => tensor_in%matrix_rep
1998 843237 : tensor_out%owns_matrix = .FALSE.
1999 :
2000 843237 : tensor_out%nd_index = nd_index_rs
2001 843237 : tensor_out%nd_index_blk = nd_index_blk_rs
2002 843237 : tensor_out%pgrid%mp_comm_2d = tensor_in%pgrid%mp_comm_2d
2003 843237 : IF (ALLOCATED(tensor_in%pgrid%tas_split_info)) THEN
2004 843237 : ALLOCATE (tensor_out%pgrid%tas_split_info, SOURCE=tensor_in%pgrid%tas_split_info)
2005 : END IF
2006 843237 : tensor_out%refcount => tensor_in%refcount
2007 843237 : CALL dbt_hold(tensor_out)
2008 :
2009 843237 : CALL reorder_arrays(tensor_in%blk_sizes, tensor_out%blk_sizes, order)
2010 843237 : CALL reorder_arrays(tensor_in%blk_offsets, tensor_out%blk_offsets, order)
2011 843237 : CALL reorder_arrays(tensor_in%nd_dist, tensor_out%nd_dist, order)
2012 843237 : CALL reorder_arrays(tensor_in%blks_local, tensor_out%blks_local, order)
2013 2529711 : ALLOCATE (tensor_out%nblks_local(ndims))
2014 1686474 : ALLOCATE (tensor_out%nfull_local(ndims))
2015 3097470 : tensor_out%nblks_local(order) = tensor_in%nblks_local(:)
2016 3097470 : tensor_out%nfull_local(order) = tensor_in%nfull_local(:)
2017 843237 : tensor_out%name = tensor_in%name
2018 843237 : tensor_out%valid = .TRUE.
2019 :
2020 843237 : IF (ALLOCATED(tensor_in%contraction_storage)) THEN
2021 324663 : ALLOCATE (tensor_out%contraction_storage, SOURCE=tensor_in%contraction_storage)
2022 324663 : CALL destroy_array_list(tensor_out%contraction_storage%batch_ranges)
2023 324663 : CALL reorder_arrays(tensor_in%contraction_storage%batch_ranges, tensor_out%contraction_storage%batch_ranges, order)
2024 : END IF
2025 :
2026 843237 : CALL timestop(handle)
2027 1686474 : END SUBROUTINE dbt_permute_index
2028 :
2029 : ! **************************************************************************************************
2030 : !> \brief Map contraction bounds to bounds referring to tensor indices
2031 : !> see dbt_contract for docu of dummy arguments
2032 : !> \param bounds_t1 bounds mapped to tensor_1
2033 : !> \param bounds_t2 bounds mapped to tensor_2
2034 : !> \param do_crop_1 whether tensor 1 should be cropped
2035 : !> \param do_crop_2 whether tensor 2 should be cropped
2036 : !> \author Patrick Seewald
2037 : ! **************************************************************************************************
2038 257420 : SUBROUTINE dbt_map_bounds_to_tensors(tensor_1, tensor_2, &
2039 257420 : contract_1, notcontract_1, &
2040 514840 : contract_2, notcontract_2, &
2041 257420 : bounds_t1, bounds_t2, &
2042 162026 : bounds_1, bounds_2, bounds_3, &
2043 : do_crop_1, do_crop_2)
2044 :
2045 : TYPE(dbt_type), INTENT(IN) :: tensor_1, tensor_2
2046 : INTEGER, DIMENSION(:), INTENT(IN) :: contract_1, contract_2, &
2047 : notcontract_1, notcontract_2
2048 : INTEGER, DIMENSION(2, ndims_tensor(tensor_1)), &
2049 : INTENT(OUT) :: bounds_t1
2050 : INTEGER, DIMENSION(2, ndims_tensor(tensor_2)), &
2051 : INTENT(OUT) :: bounds_t2
2052 : INTEGER, DIMENSION(2, SIZE(contract_1)), &
2053 : INTENT(IN), OPTIONAL :: bounds_1
2054 : INTEGER, DIMENSION(2, SIZE(notcontract_1)), &
2055 : INTENT(IN), OPTIONAL :: bounds_2
2056 : INTEGER, DIMENSION(2, SIZE(notcontract_2)), &
2057 : INTENT(IN), OPTIONAL :: bounds_3
2058 : LOGICAL, INTENT(OUT), OPTIONAL :: do_crop_1, do_crop_2
2059 : LOGICAL, DIMENSION(2) :: do_crop
2060 :
2061 257420 : do_crop = .FALSE.
2062 :
2063 913027 : bounds_t1(1, :) = 1
2064 913027 : CALL dbt_get_info(tensor_1, nfull_total=bounds_t1(2, :))
2065 :
2066 963356 : bounds_t2(1, :) = 1
2067 963356 : CALL dbt_get_info(tensor_2, nfull_total=bounds_t2(2, :))
2068 :
2069 257420 : IF (PRESENT(bounds_1)) THEN
2070 239588 : bounds_t1(:, contract_1) = bounds_1
2071 105770 : do_crop(1) = .TRUE.
2072 239588 : bounds_t2(:, contract_2) = bounds_1
2073 257420 : do_crop(2) = .TRUE.
2074 : END IF
2075 :
2076 257420 : IF (PRESENT(bounds_2)) THEN
2077 314774 : bounds_t1(:, notcontract_1) = bounds_2
2078 257420 : do_crop(1) = .TRUE.
2079 : END IF
2080 :
2081 257420 : IF (PRESENT(bounds_3)) THEN
2082 329482 : bounds_t2(:, notcontract_2) = bounds_3
2083 257420 : do_crop(2) = .TRUE.
2084 : END IF
2085 :
2086 257420 : IF (PRESENT(do_crop_1)) do_crop_1 = do_crop(1)
2087 257420 : IF (PRESENT(do_crop_2)) do_crop_2 = do_crop(2)
2088 :
2089 578355 : END SUBROUTINE dbt_map_bounds_to_tensors
2090 :
2091 : ! **************************************************************************************************
2092 : !> \brief print tensor contraction indices in a human readable way
2093 : !> \param indchar1 characters printed for index of tensor 1
2094 : !> \param indchar2 characters printed for index of tensor 2
2095 : !> \param indchar3 characters printed for index of tensor 3
2096 : !> \param unit_nr output unit
2097 : !> \author Patrick Seewald
2098 : ! **************************************************************************************************
2099 165360 : SUBROUTINE dbt_print_contraction_index(tensor_1, indchar1, tensor_2, indchar2, tensor_3, indchar3, unit_nr)
2100 : TYPE(dbt_type), INTENT(IN) :: tensor_1, tensor_2, tensor_3
2101 : CHARACTER(LEN=1), DIMENSION(ndims_tensor(tensor_1)), INTENT(IN) :: indchar1
2102 : CHARACTER(LEN=1), DIMENSION(ndims_tensor(tensor_2)), INTENT(IN) :: indchar2
2103 : CHARACTER(LEN=1), DIMENSION(ndims_tensor(tensor_3)), INTENT(IN) :: indchar3
2104 : INTEGER, INTENT(IN) :: unit_nr
2105 330720 : INTEGER, DIMENSION(ndims_matrix_row(tensor_1)) :: map11
2106 330720 : INTEGER, DIMENSION(ndims_matrix_column(tensor_1)) :: map12
2107 330720 : INTEGER, DIMENSION(ndims_matrix_row(tensor_2)) :: map21
2108 330720 : INTEGER, DIMENSION(ndims_matrix_column(tensor_2)) :: map22
2109 330720 : INTEGER, DIMENSION(ndims_matrix_row(tensor_3)) :: map31
2110 330720 : INTEGER, DIMENSION(ndims_matrix_column(tensor_3)) :: map32
2111 : INTEGER :: ichar1, ichar2, ichar3, unit_nr_prv
2112 :
2113 165360 : unit_nr_prv = prep_output_unit(unit_nr)
2114 :
2115 165360 : IF (unit_nr_prv /= 0) THEN
2116 165360 : CALL dbt_get_mapping_info(tensor_1%nd_index_blk, map1_2d=map11, map2_2d=map12)
2117 165360 : CALL dbt_get_mapping_info(tensor_2%nd_index_blk, map1_2d=map21, map2_2d=map22)
2118 165360 : CALL dbt_get_mapping_info(tensor_3%nd_index_blk, map1_2d=map31, map2_2d=map32)
2119 : END IF
2120 :
2121 165360 : IF (unit_nr_prv > 0) THEN
2122 30 : WRITE (unit_nr_prv, '(T2,A)') "INDEX INFO"
2123 30 : WRITE (unit_nr_prv, '(T15,A)', advance='no') "tensor index: ("
2124 123 : DO ichar1 = 1, SIZE(indchar1)
2125 123 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar1(ichar1)
2126 : END DO
2127 30 : WRITE (unit_nr_prv, '(A)', advance='no') ") x ("
2128 120 : DO ichar2 = 1, SIZE(indchar2)
2129 120 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar2(ichar2)
2130 : END DO
2131 30 : WRITE (unit_nr_prv, '(A)', advance='no') ") = ("
2132 123 : DO ichar3 = 1, SIZE(indchar3)
2133 123 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar3(ichar3)
2134 : END DO
2135 30 : WRITE (unit_nr_prv, '(A)') ")"
2136 :
2137 30 : WRITE (unit_nr_prv, '(T15,A)', advance='no') "matrix index: ("
2138 82 : DO ichar1 = 1, SIZE(map11)
2139 82 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar1(map11(ichar1))
2140 : END DO
2141 30 : WRITE (unit_nr_prv, '(A1)', advance='no') "|"
2142 71 : DO ichar1 = 1, SIZE(map12)
2143 71 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar1(map12(ichar1))
2144 : END DO
2145 30 : WRITE (unit_nr_prv, '(A)', advance='no') ") x ("
2146 76 : DO ichar2 = 1, SIZE(map21)
2147 76 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar2(map21(ichar2))
2148 : END DO
2149 30 : WRITE (unit_nr_prv, '(A1)', advance='no') "|"
2150 74 : DO ichar2 = 1, SIZE(map22)
2151 74 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar2(map22(ichar2))
2152 : END DO
2153 30 : WRITE (unit_nr_prv, '(A)', advance='no') ") = ("
2154 79 : DO ichar3 = 1, SIZE(map31)
2155 79 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar3(map31(ichar3))
2156 : END DO
2157 30 : WRITE (unit_nr_prv, '(A1)', advance='no') "|"
2158 74 : DO ichar3 = 1, SIZE(map32)
2159 74 : WRITE (unit_nr_prv, '(A1)', advance='no') indchar3(map32(ichar3))
2160 : END DO
2161 30 : WRITE (unit_nr_prv, '(A)') ")"
2162 : END IF
2163 :
2164 165360 : END SUBROUTINE dbt_print_contraction_index
2165 :
2166 : ! **************************************************************************************************
2167 : !> \brief Initialize batched contraction for this tensor.
2168 : !>
2169 : !> Explanation: A batched contraction is a contraction performed in several consecutive steps
2170 : !> by specification of bounds in dbt_contract. This can be used to reduce memory by
2171 : !> a large factor. The routines dbt_batched_contract_init and
2172 : !> dbt_batched_contract_finalize should be called to define the scope of a batched
2173 : !> contraction as this enables important optimizations (adapting communication scheme to
2174 : !> batches and adapting process grid to multiplication algorithm). The routines
2175 : !> dbt_batched_contract_init and dbt_batched_contract_finalize must be
2176 : !> called before the first and after the last contraction step on all 3 tensors.
2177 : !>
2178 : !> Requirements:
2179 : !> - the tensors are in a compatible matrix layout (see documentation of
2180 : !> `dbt_contract`, note 2 & 3). If they are not, process grid optimizations are
2181 : !> disabled and a warning is issued.
2182 : !> - within the scope of a batched contraction, it is not allowed to access or change tensor
2183 : !> data except by calling the routines dbt_contract & dbt_copy.
2184 : !> - the bounds affecting indices of the smallest tensor must not change in the course of a
2185 : !> batched contraction (todo: get rid of this requirement).
2186 : !>
2187 : !> Side effects:
2188 : !> - the parallel layout (process grid and distribution) of all tensors may change. In order
2189 : !> to disable the process grid optimization including this side effect, call this routine
2190 : !> only on the smallest of the 3 tensors.
2191 : !>
2192 : !> \note
2193 : !> Note 1: for an example of batched contraction see `examples/dbt_example.F`.
2194 : !> (todo: the example is outdated and should be updated).
2195 : !>
2196 : !> Note 2: it is meaningful to use this feature if the contraction consists of one batch only
2197 : !> but if multiple contractions involving the same 3 tensors are performed
2198 : !> (batched_contract_init and batched_contract_finalize must then be called before/after each
2199 : !> contraction call). The process grid is then optimized after the first contraction
2200 : !> and future contraction may profit from this optimization.
2201 : !>
2202 : !> \param batch_range_i refers to the ith tensor dimension and contains all block indices starting
2203 : !> a new range. The size should be the number of ranges plus one, the last
2204 : !> element being the block index plus one of the last block in the last range.
2205 : !> For internal load balancing optimizations, optionally specify the index
2206 : !> ranges of batched contraction.
2207 : !> \author Patrick Seewald
2208 : ! **************************************************************************************************
2209 127879 : SUBROUTINE dbt_batched_contract_init(tensor, ${varlist("batch_range")}$)
2210 : TYPE(dbt_type), INTENT(INOUT) :: tensor
2211 : INTEGER, DIMENSION(:), OPTIONAL, INTENT(IN) :: ${varlist("batch_range")}$
2212 255758 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: tdims
2213 127879 : INTEGER, DIMENSION(:), ALLOCATABLE :: ${varlist("batch_range_prv")}$
2214 : LOGICAL :: static_range
2215 :
2216 127879 : CALL dbt_get_info(tensor, nblks_total=tdims)
2217 :
2218 127879 : static_range = .TRUE.
2219 : #:for idim in range(1, maxdim+1)
2220 127879 : IF (ndims_tensor(tensor) >= ${idim}$) THEN
2221 303330 : IF (PRESENT(batch_range_${idim}$)) THEN
2222 460972 : ALLOCATE (batch_range_prv_${idim}$, source=batch_range_${idim}$)
2223 303330 : static_range = .FALSE.
2224 : ELSE
2225 232704 : ALLOCATE (batch_range_prv_${idim}$ (2))
2226 232704 : batch_range_prv_${idim}$ (1) = 1
2227 232704 : batch_range_prv_${idim}$ (2) = tdims(${idim}$) + 1
2228 : END IF
2229 : END IF
2230 : #:endfor
2231 :
2232 127879 : ALLOCATE (tensor%contraction_storage)
2233 127879 : tensor%contraction_storage%static = static_range
2234 127879 : IF (static_range) THEN
2235 86249 : CALL dbt_tas_batched_mm_init(tensor%matrix_rep)
2236 : END IF
2237 127879 : tensor%contraction_storage%nsplit_avg = 0.0_dp
2238 127879 : tensor%contraction_storage%ibatch = 0
2239 :
2240 : #:for ndim in range(1, maxdim+1)
2241 255758 : IF (ndims_tensor(tensor) == ${ndim}$) THEN
2242 : CALL create_array_list(tensor%contraction_storage%batch_ranges, ${ndim}$, &
2243 127879 : ${varlist("batch_range_prv", nmax=ndim)}$)
2244 : END IF
2245 : #:endfor
2246 :
2247 127879 : END SUBROUTINE dbt_batched_contract_init
2248 :
2249 : ! **************************************************************************************************
2250 : !> \brief finalize batched contraction. This performs all communication that has been postponed in
2251 : !> the contraction calls.
2252 : !> \author Patrick Seewald
2253 : ! **************************************************************************************************
2254 255758 : SUBROUTINE dbt_batched_contract_finalize(tensor, unit_nr)
2255 : TYPE(dbt_type), INTENT(INOUT) :: tensor
2256 : INTEGER, INTENT(IN), OPTIONAL :: unit_nr
2257 : LOGICAL :: do_write
2258 : INTEGER :: unit_nr_prv, handle
2259 :
2260 127879 : CALL tensor%pgrid%mp_comm_2d%sync()
2261 127879 : CALL timeset("dbt_total", handle)
2262 127879 : unit_nr_prv = prep_output_unit(unit_nr)
2263 :
2264 127879 : do_write = .FALSE.
2265 :
2266 127879 : IF (tensor%contraction_storage%static) THEN
2267 86249 : IF (tensor%matrix_rep%do_batched > 0) THEN
2268 86249 : IF (tensor%matrix_rep%mm_storage%batched_out) do_write = .TRUE.
2269 : END IF
2270 86249 : CALL dbt_tas_batched_mm_finalize(tensor%matrix_rep)
2271 : END IF
2272 :
2273 127879 : IF (do_write .AND. unit_nr_prv /= 0) THEN
2274 16666 : IF (unit_nr_prv > 0) THEN
2275 : WRITE (unit_nr_prv, "(T2,A)") &
2276 0 : "FINALIZING BATCHED PROCESSING OF MATMUL"
2277 : END IF
2278 16666 : CALL dbt_write_tensor_info(tensor, unit_nr_prv)
2279 16666 : CALL dbt_write_tensor_dist(tensor, unit_nr_prv)
2280 : END IF
2281 :
2282 127879 : CALL destroy_array_list(tensor%contraction_storage%batch_ranges)
2283 127879 : DEALLOCATE (tensor%contraction_storage)
2284 127879 : CALL tensor%pgrid%mp_comm_2d%sync()
2285 127879 : CALL timestop(handle)
2286 :
2287 127879 : END SUBROUTINE dbt_batched_contract_finalize
2288 :
2289 : ! **************************************************************************************************
2290 : !> \brief change the process grid of a tensor
2291 : !> \param nodata optionally don't copy the tensor data (then tensor is empty on returned)
2292 : !> \param batch_range_i refers to the ith tensor dimension and contains all block indices starting
2293 : !> a new range. The size should be the number of ranges plus one, the last
2294 : !> element being the block index plus one of the last block in the last range.
2295 : !> For internal load balancing optimizations, optionally specify the index
2296 : !> ranges of batched contraction.
2297 : !> \author Patrick Seewald
2298 : ! **************************************************************************************************
2299 812 : SUBROUTINE dbt_change_pgrid(tensor, pgrid, ${varlist("batch_range")}$, &
2300 : nodata, pgrid_changed, unit_nr)
2301 : TYPE(dbt_type), INTENT(INOUT) :: tensor
2302 : TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid
2303 : INTEGER, DIMENSION(:), OPTIONAL, INTENT(IN) :: ${varlist("batch_range")}$
2304 : !!
2305 : LOGICAL, INTENT(IN), OPTIONAL :: nodata
2306 : LOGICAL, INTENT(OUT), OPTIONAL :: pgrid_changed
2307 : INTEGER, INTENT(IN), OPTIONAL :: unit_nr
2308 : CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_change_pgrid'
2309 : CHARACTER(default_string_length) :: name
2310 : INTEGER :: handle
2311 812 : INTEGER, ALLOCATABLE, DIMENSION(:) :: ${varlist("bs")}$, &
2312 812 : ${varlist("dist")}$
2313 1624 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: pcoord, pcoord_ref, pdims, pdims_ref, &
2314 1624 : tdims
2315 5684 : TYPE(dbt_type) :: t_tmp
2316 5684 : TYPE(dbt_distribution_type) :: dist
2317 1624 : INTEGER, DIMENSION(ndims_matrix_row(tensor)) :: map1
2318 : INTEGER, &
2319 1624 : DIMENSION(ndims_matrix_column(tensor)) :: map2
2320 1624 : LOGICAL, DIMENSION(ndims_tensor(tensor)) :: mem_aware
2321 812 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: nbatch
2322 : INTEGER :: ind1, ind2, batch_size, ibatch
2323 :
2324 812 : IF (PRESENT(pgrid_changed)) pgrid_changed = .FALSE.
2325 812 : CALL mp_environ_pgrid(pgrid, pdims, pcoord)
2326 812 : CALL mp_environ_pgrid(tensor%pgrid, pdims_ref, pcoord_ref)
2327 :
2328 836 : IF (ALL(pdims == pdims_ref)) THEN
2329 8 : IF (ALLOCATED(pgrid%tas_split_info) .AND. ALLOCATED(tensor%pgrid%tas_split_info)) THEN
2330 8 : IF (pgrid%tas_split_info%ngroup == tensor%pgrid%tas_split_info%ngroup) THEN
2331 : RETURN
2332 : END IF
2333 : END IF
2334 : END IF
2335 :
2336 804 : CALL timeset(routineN, handle)
2337 :
2338 : #:for idim in range(1, maxdim+1)
2339 3216 : IF (ndims_tensor(tensor) >= ${idim}$) THEN
2340 2412 : mem_aware(${idim}$) = PRESENT(batch_range_${idim}$)
2341 2412 : IF (mem_aware(${idim}$)) nbatch(${idim}$) = SIZE(batch_range_${idim}$) - 1
2342 : END IF
2343 : #:endfor
2344 :
2345 804 : CALL dbt_get_info(tensor, nblks_total=tdims, name=name)
2346 :
2347 : #:for idim in range(1, maxdim+1)
2348 3216 : IF (ndims_tensor(tensor) >= ${idim}$) THEN
2349 7236 : ALLOCATE (bs_${idim}$ (dbt_nblks_total(tensor, ${idim}$)))
2350 2412 : CALL get_ith_array(tensor%blk_sizes, ${idim}$, bs_${idim}$)
2351 7236 : ALLOCATE (dist_${idim}$ (tdims(${idim}$)))
2352 2412 : dist_${idim}$ = 0
2353 2412 : IF (mem_aware(${idim}$)) THEN
2354 6556 : DO ibatch = 1, nbatch(${idim}$)
2355 4144 : ind1 = batch_range_${idim}$ (ibatch)
2356 4144 : ind2 = batch_range_${idim}$ (ibatch + 1) - 1
2357 4144 : batch_size = ind2 - ind1 + 1
2358 : CALL dbt_default_distvec(batch_size, pdims(${idim}$), &
2359 6556 : bs_${idim}$ (ind1:ind2), dist_${idim}$ (ind1:ind2))
2360 : END DO
2361 : ELSE
2362 0 : CALL dbt_default_distvec(tdims(${idim}$), pdims(${idim}$), bs_${idim}$, dist_${idim}$)
2363 : END IF
2364 : END IF
2365 : #:endfor
2366 :
2367 804 : CALL dbt_get_mapping_info(tensor%nd_index_blk, map1_2d=map1, map2_2d=map2)
2368 : #:for ndim in ndims
2369 1608 : IF (ndims_tensor(tensor) == ${ndim}$) THEN
2370 804 : CALL dbt_distribution_new(dist, pgrid, ${varlist("dist", nmax=ndim)}$)
2371 804 : CALL dbt_create(t_tmp, name, dist, map1, map2, ${varlist("bs", nmax=ndim)}$)
2372 : END IF
2373 : #:endfor
2374 804 : CALL dbt_distribution_destroy(dist)
2375 :
2376 804 : IF (PRESENT(nodata)) THEN
2377 0 : IF (.NOT. nodata) CALL dbt_copy_expert(tensor, t_tmp, move_data=.TRUE.)
2378 : ELSE
2379 804 : CALL dbt_copy_expert(tensor, t_tmp, move_data=.TRUE.)
2380 : END IF
2381 :
2382 804 : CALL dbt_copy_contraction_storage(tensor, t_tmp)
2383 :
2384 804 : CALL dbt_destroy(tensor)
2385 804 : tensor = t_tmp
2386 :
2387 804 : IF (PRESENT(unit_nr)) THEN
2388 804 : IF (unit_nr > 0) THEN
2389 0 : WRITE (unit_nr, "(T2,A,1X,A)") "OPTIMIZED PGRID INFO FOR", TRIM(tensor%name)
2390 0 : WRITE (unit_nr, "(T4,A,1X,3I6)") "process grid dimensions:", pdims
2391 0 : CALL dbt_write_split_info(pgrid, unit_nr)
2392 : END IF
2393 : END IF
2394 :
2395 804 : IF (PRESENT(pgrid_changed)) pgrid_changed = .TRUE.
2396 :
2397 804 : CALL timestop(handle)
2398 812 : END SUBROUTINE dbt_change_pgrid
2399 :
2400 : ! **************************************************************************************************
2401 : !> \brief map tensor to a new 2d process grid for the matrix representation.
2402 : !> \author Patrick Seewald
2403 : ! **************************************************************************************************
2404 812 : SUBROUTINE dbt_change_pgrid_2d(tensor, mp_comm, pdims, nodata, nsplit, dimsplit, pgrid_changed, unit_nr)
2405 : TYPE(dbt_type), INTENT(INOUT) :: tensor
2406 : TYPE(mp_cart_type), INTENT(IN) :: mp_comm
2407 : INTEGER, DIMENSION(2), INTENT(IN), OPTIONAL :: pdims
2408 : LOGICAL, INTENT(IN), OPTIONAL :: nodata
2409 : INTEGER, INTENT(IN), OPTIONAL :: nsplit, dimsplit
2410 : LOGICAL, INTENT(OUT), OPTIONAL :: pgrid_changed
2411 : INTEGER, INTENT(IN), OPTIONAL :: unit_nr
2412 1624 : INTEGER, DIMENSION(ndims_matrix_row(tensor)) :: map1
2413 1624 : INTEGER, DIMENSION(ndims_matrix_column(tensor)) :: map2
2414 1624 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: dims, nbatches
2415 2436 : TYPE(dbt_pgrid_type) :: pgrid
2416 812 : INTEGER, DIMENSION(:), ALLOCATABLE :: ${varlist("batch_range")}$
2417 812 : INTEGER, DIMENSION(:), ALLOCATABLE :: array
2418 : INTEGER :: idim
2419 :
2420 812 : CALL dbt_get_mapping_info(tensor%pgrid%nd_index_grid, map1_2d=map1, map2_2d=map2)
2421 812 : CALL blk_dims_tensor(tensor, dims)
2422 :
2423 812 : IF (ALLOCATED(tensor%contraction_storage)) THEN
2424 : ASSOCIATE (batch_ranges => tensor%contraction_storage%batch_ranges)
2425 3248 : nbatches = sizes_of_arrays(tensor%contraction_storage%batch_ranges) - 1
2426 : ! for good load balancing the process grid dimensions should be chosen adapted to the
2427 : ! tensor dimenions. For batched contraction the tensor dimensions should be divided by
2428 : ! the number of batches (number of index ranges).
2429 4060 : DO idim = 1, ndims_tensor(tensor)
2430 2436 : CALL get_ith_array(tensor%contraction_storage%batch_ranges, idim, array)
2431 2436 : dims(idim) = array(nbatches(idim) + 1) - array(1)
2432 2436 : DEALLOCATE (array)
2433 2436 : dims(idim) = dims(idim)/nbatches(idim)
2434 3248 : IF (dims(idim) <= 0) dims(idim) = 1
2435 : END DO
2436 : END ASSOCIATE
2437 : END IF
2438 :
2439 812 : pgrid = dbt_nd_mp_comm(mp_comm, map1, map2, pdims_2d=pdims, tdims=dims, nsplit=nsplit, dimsplit=dimsplit)
2440 812 : IF (ALLOCATED(tensor%contraction_storage)) THEN
2441 : #:for ndim in range(1, maxdim+1)
2442 1624 : IF (ndims_tensor(tensor) == ${ndim}$) THEN
2443 812 : CALL get_arrays(tensor%contraction_storage%batch_ranges, ${varlist("batch_range", nmax=ndim)}$)
2444 : CALL dbt_change_pgrid(tensor, pgrid, ${varlist("batch_range", nmax=ndim)}$, &
2445 812 : nodata=nodata, pgrid_changed=pgrid_changed, unit_nr=unit_nr)
2446 : END IF
2447 : #:endfor
2448 : ELSE
2449 0 : CALL dbt_change_pgrid(tensor, pgrid, nodata=nodata, pgrid_changed=pgrid_changed, unit_nr=unit_nr)
2450 : END IF
2451 812 : CALL dbt_pgrid_destroy(pgrid)
2452 :
2453 812 : END SUBROUTINE dbt_change_pgrid_2d
2454 :
2455 81865 : END MODULE dbt_methods
|
