diff --git a/src/dbm/dbm_api.F b/src/dbm/dbm_api.F index 38ac5f31c7..fff4dfea97 100644 --- a/src/dbm/dbm_api.F +++ b/src/dbm/dbm_api.F @@ -83,6 +83,7 @@ MODULE dbm_api PUBLIC :: dbm_iterator PUBLIC :: dbm_iterator_start PUBLIC :: dbm_iterator_stop + PUBLIC :: dbm_iterator_num_blocks PUBLIC :: dbm_iterator_blocks_left PUBLIC :: dbm_iterator_next_block @@ -850,6 +851,28 @@ CONTAINS CALL validate(matrix) END SUBROUTINE dbm_iterator_start +! ************************************************************************************************** +!> \brief Returns number of blocks the iterator will provide to calling thread. +!> \param iterator ... +!> \return ... +!> \author Ole Schuett +! ************************************************************************************************** + FUNCTION dbm_iterator_num_blocks(iterator) RESULT(num_blocks) + TYPE(dbm_iterator), INTENT(IN) :: iterator + INTEGER :: num_blocks + + INTERFACE + FUNCTION dbm_iterator_num_blocks_c(iterator) & + BIND(C, name="dbm_iterator_num_blocks") + IMPORT :: C_PTR, C_INT + TYPE(C_PTR), VALUE :: iterator + INTEGER(kind=C_INT) :: dbm_iterator_num_blocks_c + END FUNCTION dbm_iterator_num_blocks_c + END INTERFACE + + num_blocks = dbm_iterator_num_blocks_c(iterator%c_ptr) + END FUNCTION dbm_iterator_num_blocks + ! ************************************************************************************************** !> \brief Tests whether the given iterator has any block left. !> \param iterator ... diff --git a/src/dbm/dbm_matrix.c b/src/dbm/dbm_matrix.c index cc75c0fe4a..9565013092 100644 --- a/src/dbm/dbm_matrix.c +++ b/src/dbm/dbm_matrix.c @@ -368,27 +368,32 @@ void dbm_filter(dbm_matrix_t *matrix, const double eps) { /******************************************************************************* * \brief Adds list of blocks efficiently. The blocks will be filled with zeros. + * This routine must always be called within an OpenMP parallel region. * \author Ole Schuett ******************************************************************************/ void dbm_reserve_blocks(dbm_matrix_t *matrix, const int nblocks, const int rows[], const int cols[]) { - assert(omp_get_num_threads() == 1); + assert(omp_get_num_threads() == omp_get_max_threads() && + "Please call dbm_reserve_blocks within an OpenMP parallel region."); const int my_rank = matrix->dist->my_rank; + for (int i = 0; i < nblocks; i++) { + const int ishard = rows[i] % matrix->nshards; + dbm_shard_t *shard = &matrix->shards[ishard]; + omp_set_lock(&shard->lock); + assert(0 <= rows[i] && rows[i] < matrix->nrows); + assert(0 <= cols[i] && cols[i] < matrix->ncols); + assert(dbm_get_stored_coordinates(matrix, rows[i], cols[i]) == my_rank); + const int row_size = matrix->row_sizes[rows[i]]; + const int col_size = matrix->col_sizes[cols[i]]; + const int block_size = row_size * col_size; + dbm_shard_get_or_promise_block(shard, rows[i], cols[i], block_size); + omp_unset_lock(&shard->lock); + } +#pragma omp barrier -#pragma omp parallel for schedule(dynamic) +#pragma omp for schedule(dynamic) for (int ishard = 0; ishard < matrix->nshards; ishard++) { dbm_shard_t *shard = &matrix->shards[ishard]; - for (int i = 0; i < nblocks; i++) { - if (rows[i] % matrix->nshards == ishard) { - assert(0 <= rows[i] && rows[i] < matrix->nrows); - assert(0 <= cols[i] && cols[i] < matrix->ncols); - assert(dbm_get_stored_coordinates(matrix, rows[i], cols[i]) == my_rank); - const int row_size = matrix->row_sizes[rows[i]]; - const int col_size = matrix->col_sizes[cols[i]]; - const int block_size = row_size * col_size; - dbm_shard_get_or_promise_block(shard, rows[i], cols[i], block_size); - } - } dbm_shard_allocate_promised_blocks(shard); } } @@ -474,30 +479,42 @@ void dbm_add(dbm_matrix_t *matrix_a, const dbm_matrix_t *matrix_b) { /******************************************************************************* * \brief Creates an iterator for the blocks of the given matrix. * The iteration order is not stable. + * This routine must always be called within an OpenMP parallel region. * \author Ole Schuett ******************************************************************************/ void dbm_iterator_start(dbm_iterator_t **iter_out, const dbm_matrix_t *matrix) { - assert(omp_get_num_threads() == 1); + assert(omp_get_num_threads() == omp_get_max_threads() && + "Please call dbm_iterator_start within an OpenMP parallel region."); dbm_iterator_t *iter = malloc(sizeof(dbm_iterator_t)); iter->matrix = matrix; iter->next_block = 0; - iter->next_shard = 0; - while (iter->next_shard < matrix->nshards) { - if (matrix->shards[iter->next_shard].nblocks > 0) { - break; - } - iter->next_shard++; + iter->next_shard = omp_get_thread_num(); + while (iter->next_shard < matrix->nshards && + matrix->shards[iter->next_shard].nblocks == 0) { + iter->next_shard += omp_get_num_threads(); } assert(*iter_out == NULL); *iter_out = iter; } +/******************************************************************************* + * \brief Returns number of blocks the iterator will provide to calling thread. + * \author Ole Schuett + ******************************************************************************/ +int dbm_iterator_num_blocks(const dbm_iterator_t *iter) { + int num_blocks = 0; + for (int ishard = omp_get_thread_num(); ishard < iter->matrix->nshards; + ishard += omp_get_num_threads()) { + num_blocks += iter->matrix->shards[ishard].nblocks; + } + return num_blocks; +} + /******************************************************************************* * \brief Tests whether the given iterator has any block left. * \author Ole Schuett ******************************************************************************/ bool dbm_iterator_blocks_left(const dbm_iterator_t *iter) { - assert(omp_get_num_threads() == 1); return iter->next_shard < iter->matrix->nshards; } @@ -507,7 +524,6 @@ bool dbm_iterator_blocks_left(const dbm_iterator_t *iter) { ******************************************************************************/ void dbm_iterator_next_block(dbm_iterator_t *iter, int *row, int *col, double **block, int *row_size, int *col_size) { - assert(omp_get_num_threads() == 1); const dbm_matrix_t *matrix = iter->matrix; assert(iter->next_shard < matrix->nshards); const dbm_shard_t *shard = &matrix->shards[iter->next_shard]; @@ -524,10 +540,10 @@ void dbm_iterator_next_block(dbm_iterator_t *iter, int *row, int *col, iter->next_block++; if (iter->next_block >= shard->nblocks) { // Advance to the next non-empty shard... - iter->next_shard++; + iter->next_shard += omp_get_num_threads(); while (iter->next_shard < matrix->nshards && matrix->shards[iter->next_shard].nblocks == 0) { - iter->next_shard++; + iter->next_shard += omp_get_num_threads(); } iter->next_block = 0; // ...and continue with its first block. } @@ -537,10 +553,7 @@ void dbm_iterator_next_block(dbm_iterator_t *iter, int *row, int *col, * \brief Releases the given iterator. * \author Ole Schuett ******************************************************************************/ -void dbm_iterator_stop(dbm_iterator_t *iter) { - assert(omp_get_num_threads() == 1); - free(iter); -} +void dbm_iterator_stop(dbm_iterator_t *iter) { free(iter); } /******************************************************************************* * \brief Computes a checksum of the given matrix. diff --git a/src/dbm/dbm_matrix.h b/src/dbm/dbm_matrix.h index b167f46266..7459083f7a 100644 --- a/src/dbm/dbm_matrix.h +++ b/src/dbm/dbm_matrix.h @@ -104,6 +104,7 @@ void dbm_filter(dbm_matrix_t *matrix, const double eps); /******************************************************************************* * \brief Adds list of blocks efficiently. The blocks will be filled with zeros. + * This routine must always be called within an OpenMP parallel region. * \author Ole Schuett ******************************************************************************/ void dbm_reserve_blocks(dbm_matrix_t *matrix, const int nblocks, @@ -130,10 +131,17 @@ void dbm_add(dbm_matrix_t *matrix_a, const dbm_matrix_t *matrix_b); /******************************************************************************* * \brief Creates an iterator for the blocks of the given matrix. * The iteration order is not stable. + * This routine must always be called within an OpenMP parallel region. * \author Ole Schuett ******************************************************************************/ void dbm_iterator_start(dbm_iterator_t **iter_out, const dbm_matrix_t *matrix); +/******************************************************************************* + * \brief Returns number of blocks the iterator will provide to calling thread. + * \author Ole Schuett + ******************************************************************************/ +int dbm_iterator_num_blocks(const dbm_iterator_t *iter); + /******************************************************************************* * \brief Tests whether the given iterator has any block left. * \author Ole Schuett diff --git a/src/dbm/dbm_miniapp.c b/src/dbm/dbm_miniapp.c index 6a1a27f229..e9df0c2d55 100644 --- a/src/dbm/dbm_miniapp.c +++ b/src/dbm/dbm_miniapp.c @@ -66,21 +66,25 @@ static void reserve_all_blocks(dbm_matrix_t *matrix) { dbm_get_row_sizes(matrix, &nrows, &row_sizes); dbm_get_col_sizes(matrix, &ncols, &col_sizes); - const int nblocks = nrows * ncols; - int reserve_row[nblocks]; - int reserve_col[nblocks]; - - int nblocks_reserve = 0; - for (int i = 0; i < nblocks; i++) { - const int row = i % nrows; - const int col = i % ncols; - if (dbm_get_stored_coordinates(matrix, row, col) == matrix->dist->my_rank) { - reserve_row[nblocks_reserve] = row; - reserve_col[nblocks_reserve] = col; - nblocks_reserve++; +#pragma omp parallel + { + const int nblocks = nrows * ncols; + int reserve_row[nblocks / omp_get_num_threads() + 1]; + int reserve_col[nblocks / omp_get_num_threads() + 1]; + int nblocks_reserve = 0; +#pragma omp for + for (int i = 0; i < nblocks; i++) { + const int row = i % nrows; + const int col = i % ncols; + if (dbm_get_stored_coordinates(matrix, row, col) == + matrix->dist->my_rank) { + reserve_row[nblocks_reserve] = row; + reserve_col[nblocks_reserve] = col; + nblocks_reserve++; + } } + dbm_reserve_blocks(matrix, nblocks_reserve, reserve_row, reserve_col); } - dbm_reserve_blocks(matrix, nblocks_reserve, reserve_row, reserve_col); } /******************************************************************************* @@ -88,18 +92,21 @@ static void reserve_all_blocks(dbm_matrix_t *matrix) { * \author Ole Schuett ******************************************************************************/ static void set_all_blocks(dbm_matrix_t *matrix) { - dbm_iterator_t *iter = NULL; - dbm_iterator_start(&iter, matrix); - while (dbm_iterator_blocks_left(iter)) { - int row, col, row_size, col_size; - double *block; - dbm_iterator_next_block(iter, &row, &col, &block, &row_size, &col_size); - const int block_size = row_size * col_size; - for (int i = 0; i < block_size; i++) { - block[i] = 1.0; +#pragma omp parallel + { + dbm_iterator_t *iter = NULL; + dbm_iterator_start(&iter, matrix); + while (dbm_iterator_blocks_left(iter)) { + int row, col, row_size, col_size; + double *block; + dbm_iterator_next_block(iter, &row, &col, &block, &row_size, &col_size); + const int block_size = row_size * col_size; + for (int i = 0; i < block_size; i++) { + block[i] = 1.0; + } } + dbm_iterator_stop(iter); } - dbm_iterator_stop(iter); } /******************************************************************************* diff --git a/src/dbt/dbt_api.F b/src/dbt/dbt_api.F index cf857707c9..2f0de4d0a7 100644 --- a/src/dbt/dbt_api.F +++ b/src/dbt/dbt_api.F @@ -18,19 +18,19 @@ MODULE dbt_api USE dbt_block, ONLY: dbt_iterator_blocks_left,& dbt_iterator_next_block,& + dbt_iterator_num_blocks,& dbt_iterator_start,& dbt_iterator_stop,& - dbt_iterator_type,& - dbt_reserved_block_indices + dbt_iterator_type USE dbt_index, ONLY: dbt_get_mapping_info USE dbt_io, ONLY: dbt_write_blocks,& dbt_write_split_info,& dbt_write_tensor_dist,& dbt_write_tensor_info USE dbt_methods, ONLY: & - dbt_batched_contract_finalize, dbt_batched_contract_init, dbt_contract, & - dbt_contract_index, dbt_copy, dbt_copy_matrix_to_tensor, dbt_copy_tensor_to_matrix, & - dbt_get_block, dbt_get_stored_coordinates, dbt_put_block, dbt_reserve_blocks + dbt_batched_contract_finalize, dbt_batched_contract_init, dbt_contract, dbt_copy, & + dbt_copy_matrix_to_tensor, dbt_copy_tensor_to_matrix, dbt_get_block, & + dbt_get_stored_coordinates, dbt_put_block, dbt_reserve_blocks USE dbt_split, ONLY: dbt_split_blocks USE dbt_test, ONLY: dbt_checksum,& dbt_contract_test @@ -69,6 +69,7 @@ MODULE dbt_api PUBLIC :: dbt_iterator_stop PUBLIC :: dbt_iterator_start PUBLIC :: dbt_iterator_type + PUBLIC :: dbt_iterator_num_blocks PUBLIC :: dbt_split_blocks PUBLIC :: dbt_pgrid_type PUBLIC :: dbt_pgrid_create @@ -98,8 +99,6 @@ MODULE dbt_api PUBLIC :: dbt_ndims PUBLIC :: dbt_dims PUBLIC :: dbt_pgrid_change_dims - PUBLIC :: dbt_reserved_block_indices - PUBLIC :: dbt_contract_index PUBLIC :: dbt_ndims_matrix_row PUBLIC :: dbt_ndims_matrix_column PUBLIC :: dbt_nblks_local diff --git a/src/dbt/dbt_block.F b/src/dbt/dbt_block.F index b23fb09057..2d18ca2b72 100644 --- a/src/dbt/dbt_block.F +++ b/src/dbt/dbt_block.F @@ -26,7 +26,8 @@ MODULE dbt_block dbt_tas_iterator USE dbt_tas_base, ONLY: & dbt_tas_iterator_next_block, dbt_tas_iterator_blocks_left, dbt_tas_iterator_start, & - dbt_tas_iterator_stop, dbt_tas_get_block_p, dbt_tas_put_block, dbt_tas_reserve_blocks + dbt_tas_iterator_stop, dbt_tas_get_block_p, dbt_tas_put_block, dbt_tas_reserve_blocks, & + dbt_tas_iterator_num_blocks USE kinds, ONLY: dp, int_8, dp USE dbt_index, ONLY: & nd_to_2d_mapping, ndims_mapping, get_nd_indices_tensor, destroy_nd_to_2d_mapping, get_2d_indices_tensor, & @@ -48,6 +49,7 @@ MODULE dbt_block block_nd, & create_block, & dbt_get_block, & + dbt_iterator_num_blocks, & dbt_iterator_blocks_left, & dbt_iterator_next_block, & dbt_iterator_start, & @@ -55,16 +57,13 @@ MODULE dbt_block dbt_iterator_type, & dbt_put_block, & dbt_reserve_blocks, & - dbt_reserved_block_indices, & destroy_block, & checker_tr, & ndims_iterator TYPE dbt_iterator_type TYPE(dbt_tas_iterator) :: iter - TYPE(nd_to_2d_mapping) :: nd_index_blk - TYPE(nd_to_2d_mapping) :: nd_index - TYPE(array_list) :: blk_sizes, blk_offsets + TYPE(dbt_type), POINTER :: tensor => NULL() END TYPE dbt_iterator_type TYPE block_nd @@ -119,16 +118,11 @@ CONTAINS ! ************************************************************************************************** SUBROUTINE dbt_iterator_start(iterator, tensor) TYPE(dbt_iterator_type), INTENT(OUT) :: iterator - TYPE(dbt_type), INTENT(IN) :: tensor + TYPE(dbt_type), INTENT(IN), TARGET :: tensor CPASSERT(tensor%valid) - CALL dbt_tas_iterator_start(iterator%iter, tensor%matrix_rep) - iterator%nd_index_blk = tensor%nd_index_blk - iterator%nd_index = tensor%nd_index - iterator%blk_sizes = tensor%blk_sizes - iterator%blk_offsets = tensor%blk_offsets - + iterator%tensor => tensor END SUBROUTINE ! ************************************************************************************************** @@ -139,11 +133,6 @@ CONTAINS TYPE(dbt_iterator_type), INTENT(INOUT) :: iterator CALL dbt_tas_iterator_stop(iterator%iter) - CALL destroy_nd_to_2d_mapping(iterator%nd_index) - CALL destroy_nd_to_2d_mapping(iterator%nd_index_blk) - CALL destroy_array_list(iterator%blk_sizes) - CALL destroy_array_list(iterator%blk_offsets) - END SUBROUTINE ! ************************************************************************************************** @@ -156,7 +145,7 @@ CONTAINS TYPE(dbt_iterator_type), INTENT(IN) :: iterator INTEGER :: ndims_iterator - ndims_iterator = iterator%nd_index%ndim_nd + ndims_iterator = iterator%tensor%nd_index%ndim_nd END FUNCTION ! ************************************************************************************************** @@ -179,14 +168,26 @@ CONTAINS CALL dbt_tas_iterator_next_block(iterator%iter, ind_2d(1), ind_2d(2)) - ind_nd(:) = get_nd_indices_tensor(iterator%nd_index_blk, ind_2d) - IF (PRESENT(blk_size)) blk_size(:) = get_array_elements(iterator%blk_sizes, ind_nd) + ind_nd(:) = get_nd_indices_tensor(iterator%tensor%nd_index_blk, ind_2d) + IF (PRESENT(blk_size)) blk_size(:) = get_array_elements(iterator%tensor%blk_sizes, ind_nd) ! note: blk_offset needs to be determined by tensor metadata, can not be derived from 2d row/col ! offset since block index mapping is not consistent with element index mapping - IF (PRESENT(blk_offset)) blk_offset(:) = get_array_elements(iterator%blk_offsets, ind_nd) + IF (PRESENT(blk_offset)) blk_offset(:) = get_array_elements(iterator%tensor%blk_offsets, ind_nd) END SUBROUTINE +! ************************************************************************************************** +!> \brief Generalization of block_iterator_num_blocks for tensors. +!> \author Ole Schuett +! ************************************************************************************************** + FUNCTION dbt_iterator_num_blocks(iterator) + TYPE(dbt_iterator_type), INTENT(IN) :: iterator + INTEGER :: dbt_iterator_num_blocks + + dbt_iterator_num_blocks = dbt_tas_iterator_num_blocks(iterator%iter) + + END FUNCTION + ! ************************************************************************************************** !> \brief Generalization of block_iterator_blocks_left for tensors. !> \author Patrick Seewald @@ -261,17 +262,30 @@ CONTAINS !> \author Patrick Seewald ! ************************************************************************************************** SUBROUTINE dbt_reserve_blocks_template(tensor_in, tensor_out) - TYPE(dbt_type), INTENT(IN) :: tensor_in - TYPE(dbt_type), INTENT(INOUT) :: tensor_out - INTEGER :: handle + TYPE(dbt_type), INTENT(IN) :: tensor_in + TYPE(dbt_type), INTENT(INOUT) :: tensor_out - INTEGER, DIMENSION(dbt_get_num_blocks(tensor_in), ndims_tensor(tensor_in)) :: blk_ind - CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_reserve_blocks_template' + CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_reserve_blocks_template' + + TYPE(dbt_iterator_type) :: iter + INTEGER :: handle, nblk, iblk + INTEGER, DIMENSION(:, :), ALLOCATABLE :: blk_ind CALL timeset(routineN, handle) - CALL dbt_reserved_block_indices(tensor_in, blk_ind) +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,tensor_out) & +!$OMP PRIVATE(iter,nblk,iblk,blk_ind) + CALL dbt_iterator_start(iter, tensor_in) + nblk = dbt_iterator_num_blocks(iter) + ALLOCATE (blk_ind(nblk, ndims_tensor(tensor_in))) + DO iblk = 1, nblk + CALL dbt_iterator_next_block(iter, ind_nd=blk_ind(iblk, :)) + END DO + CPASSERT(.NOT. dbt_iterator_blocks_left(iter)) + CALL dbt_iterator_stop(iter) + CALL dbt_reserve_blocks(tensor_out, blk_ind) +!$OMP END PARALLEL CALL timestop(handle) END SUBROUTINE @@ -310,7 +324,10 @@ CONTAINS END DO CALL dbcsr_iterator_stop(iter) +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_out,blk_ind_1,blk_ind_2) CALL dbt_reserve_blocks(tensor_out, blk_ind_1, blk_ind_2) +!$OMP END PARALLEL IF (dbcsr_has_symmetry(matrix_in)) THEN CALL dbcsr_release(matrix_in_desym) @@ -338,20 +355,25 @@ CONTAINS nblk = dbt_get_num_blocks(tensor_in) ALLOCATE (blk_ind_1(nblk), blk_ind_2(nblk)) - CALL dbt_iterator_start(iter, tensor_in) - iblk = 0 + +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,matrix_out,iblk,blk_ind_1,blk_ind_2) & +!$OMP PRIVATE(iter,ind_2d) + CALL dbt_iterator_start(iter, tensor_in) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind_2d) IF (dbcsr_has_symmetry(matrix_out)) THEN IF (checker_tr(ind_2d(1), ind_2d(2))) CYCLE IF (ind_2d(1) > ind_2d(2)) CALL swap(ind_2d(1), ind_2d(2)) END IF - +!$OMP CRITICAL iblk = iblk + 1 - blk_ind_1(iblk) = ind_2d(1); blk_ind_2(iblk) = ind_2d(2) + blk_ind_1(iblk) = ind_2d(1) + blk_ind_2(iblk) = ind_2d(2) +!$OMP END CRITICAL END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL dbcsr_reserve_blocks(matrix_out, blk_ind_1(:iblk), blk_ind_2(:iblk)) CALL dbcsr_finalize(matrix_out) @@ -372,30 +394,6 @@ CONTAINS b = tmp END SUBROUTINE swap -! ************************************************************************************************** -!> \brief indices of non-zero blocks -!> \author Patrick Seewald -! ************************************************************************************************** - SUBROUTINE dbt_reserved_block_indices(tensor, blk_ind) - TYPE(dbt_type), INTENT(IN) :: tensor - INTEGER :: iblk, nblk - TYPE(dbt_iterator_type) :: iterator - INTEGER, DIMENSION(ndims_tensor(tensor)) :: ind_nd - INTEGER, DIMENSION(dbt_get_num_blocks(tensor), ndims_tensor(tensor)), INTENT(OUT) :: blk_ind - - CPASSERT(tensor%valid) - - nblk = dbt_get_num_blocks(tensor) - - CALL dbt_iterator_start(iterator, tensor) - DO iblk = 1, nblk - CALL dbt_iterator_next_block(iterator, ind_nd) - blk_ind(iblk, :) = ind_nd(:) - END DO - CALL dbt_iterator_stop(iterator) - - END SUBROUTINE - ! ************************************************************************************************** !> \brief Create block from array, array can be n-dimensional. !> \author Patrick Seewald diff --git a/src/dbt/dbt_methods.F b/src/dbt/dbt_methods.F index de62534c96..5a472f7ded 100644 --- a/src/dbt/dbt_methods.F +++ b/src/dbt/dbt_methods.F @@ -34,7 +34,7 @@ MODULE dbt_methods USE dbt_tas_base, ONLY: & dbt_tas_copy, dbt_tas_finalize, dbt_tas_get_info, dbt_tas_info USE dbt_tas_mm, ONLY: & - dbt_tas_multiply, dbt_tas_batched_mm_init, dbt_tas_batched_mm_finalize, dbt_tas_result_index, & + dbt_tas_multiply, dbt_tas_batched_mm_init, dbt_tas_batched_mm_finalize, & dbt_tas_batched_mm_complete, dbt_tas_set_batched_state USE dbt_block, ONLY: & dbt_iterator_type, dbt_get_block, dbt_put_block, dbt_iterator_start, & @@ -89,7 +89,6 @@ MODULE dbt_methods dbt_reserve_blocks, & dbt_copy_matrix_to_tensor, & dbt_copy_tensor_to_matrix, & - dbt_contract_index, & dbt_batched_contract_init, & dbt_batched_contract_finalize @@ -302,6 +301,8 @@ CONTAINS CALL dbt_reserve_blocks(tensor_in, tensor_out) +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,tensor_out,summation) & +!$OMP PRIVATE(iter,ind_nd,blk_data,found) CALL dbt_iterator_start(iter, tensor_in) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind_nd) @@ -311,6 +312,7 @@ CONTAINS CALL destroy_block(blk_data) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) END SUBROUTINE @@ -355,6 +357,8 @@ CONTAINS CALL dbt_reserve_blocks(matrix_in_desym, tensor_out) +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in_desym,tensor_out,summation) & +!$OMP PRIVATE(iter,ind_2d,block,tr,block_arr) CALL dbcsr_iterator_start(iter, matrix_in_desym) DO WHILE (dbcsr_iterator_blocks_left(iter)) CALL dbcsr_iterator_next_block(iter, ind_2d(1), ind_2d(2), block, tr) @@ -363,6 +367,7 @@ CONTAINS DEALLOCATE (block_arr) END DO CALL dbcsr_iterator_stop(iter) +!$OMP END PARALLEL IF (dbcsr_has_symmetry(matrix_in)) THEN CALL dbcsr_release(matrix_in_desym) @@ -399,6 +404,8 @@ CONTAINS CALL dbt_reserve_blocks(tensor_in, matrix_out) +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,matrix_out,summation) & +!$OMP PRIVATE(iter,ind_2d,block,found) CALL dbt_iterator_start(iter, tensor_in) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind_2d) @@ -415,6 +422,7 @@ CONTAINS DEALLOCATE (block) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) @@ -549,8 +557,6 @@ CONTAINS ! ************************************************************************************************** !> \brief expert routine for tensor contraction. For internal use only. !> \param nblks_local number of local blocks on this MPI rank -!> \param result_index get indices of non-zero tensor blocks for tensor_3 without actually -!> performing contraction this is an estimate based on block norm multiplication !> \author Patrick Seewald ! ************************************************************************************************** SUBROUTINE dbt_contract_expert(alpha, tensor_1, tensor_2, beta, tensor_3, & @@ -560,7 +566,7 @@ CONTAINS bounds_1, bounds_2, bounds_3, & optimize_dist, pgrid_opt_1, pgrid_opt_2, pgrid_opt_3, & filter_eps, flop, move_data, retain_sparsity, & - nblks_local, result_index, unit_nr, log_verbose) + nblks_local, unit_nr, log_verbose) REAL(dp), INTENT(IN) :: alpha TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_1 TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_2 @@ -590,8 +596,6 @@ CONTAINS LOGICAL, INTENT(IN), OPTIONAL :: move_data LOGICAL, INTENT(IN), OPTIONAL :: retain_sparsity INTEGER, INTENT(OUT), OPTIONAL :: nblks_local - INTEGER, DIMENSION(dbt_max_nblks_local(tensor_3), ndims_tensor(tensor_3)), & - OPTIONAL, INTENT(OUT) :: result_index INTEGER, OPTIONAL, INTENT(IN) :: unit_nr LOGICAL, INTENT(IN), OPTIONAL :: log_verbose @@ -600,11 +604,9 @@ CONTAINS TYPE(dbt_type), POINTER :: tensor_crop_1, tensor_crop_2 TYPE(dbt_type), POINTER :: tensor_small, tensor_large - INTEGER(int_8), DIMENSION(:, :), ALLOCATABLE :: result_index_2d LOGICAL :: assert_stmt, tensors_remapped INTEGER :: max_mm_dim, max_tensor, mp_comm, & - iblk, nblk, unit_nr_prv, ref_tensor, mp_comm_opt, & - handle + unit_nr_prv, ref_tensor, mp_comm_opt, handle INTEGER, DIMENSION(SIZE(contract_1)) :: contract_1_mod INTEGER, DIMENSION(SIZE(notcontract_1)) :: notcontract_1_mod INTEGER, DIMENSION(SIZE(contract_2)) :: contract_2_mod @@ -661,7 +663,6 @@ CONTAINS unit_nr_prv = prep_output_unit(unit_nr) IF (PRESENT(flop)) flop = 0 - IF (PRESENT(result_index)) result_index = 0 IF (PRESENT(nblks_local)) nblks_local = 0 IF (PRESENT(move_data)) THEN @@ -901,58 +902,13 @@ CONTAINS IF (new_2) CALL dbt_write_tensor_dist(tensor_contr_2, unit_nr_prv) END IF - IF (.NOT. PRESENT(result_index)) THEN - CALL dbt_tas_multiply(trans_1, trans_2, trans_3, alpha, & - tensor_contr_1%matrix_rep, tensor_contr_2%matrix_rep, & - beta, & - tensor_contr_3%matrix_rep, filter_eps=filter_eps, flop=flop, & - unit_nr=unit_nr_prv, log_verbose=log_verbose, & - split_opt=split_opt, & - move_data_a=move_data_1, move_data_b=move_data_2, retain_sparsity=retain_sparsity) - ELSE - - CALL dbt_tas_result_index(trans_1, trans_2, trans_3, tensor_contr_1%matrix_rep, tensor_contr_2%matrix_rep, & - tensor_contr_3%matrix_rep, filter_eps=filter_eps, blk_ind=result_index_2d) - - nblk = SIZE(result_index_2d, 1) - IF (PRESENT(nblks_local)) nblks_local = nblk - IF (SIZE(result_index, 1) < nblk) THEN - CALL cp_abort(__LOCATION__, & - "allocated size of `result_index` is too small. This error occurs due to a high load imbalance of distributed tensor data.") - END IF - - DO iblk = 1, nblk - result_index(iblk, :) = get_nd_indices_tensor(tensor_contr_3%nd_index_blk, result_index_2d(iblk, :)) - END DO - - IF (new_1) THEN - CALL dbt_destroy(tensor_contr_1) - DEALLOCATE (tensor_contr_1) - END IF - IF (new_2) THEN - CALL dbt_destroy(tensor_contr_2) - DEALLOCATE (tensor_contr_2) - END IF - IF (new_3) THEN - CALL dbt_destroy(tensor_contr_3) - DEALLOCATE (tensor_contr_3) - END IF - IF (do_crop_1) THEN - CALL dbt_destroy(tensor_crop_1) - DEALLOCATE (tensor_crop_1) - END IF - IF (do_crop_2) THEN - CALL dbt_destroy(tensor_crop_2) - DEALLOCATE (tensor_crop_2) - END IF - - CALL dbt_destroy(tensor_algn_1) - CALL dbt_destroy(tensor_algn_2) - CALL dbt_destroy(tensor_algn_3) - - CALL timestop(handle) - RETURN - END IF + CALL dbt_tas_multiply(trans_1, trans_2, trans_3, alpha, & + tensor_contr_1%matrix_rep, tensor_contr_2%matrix_rep, & + beta, & + tensor_contr_3%matrix_rep, filter_eps=filter_eps, flop=flop, & + unit_nr=unit_nr_prv, log_verbose=log_verbose, & + split_opt=split_opt, & + move_data_a=move_data_1, move_data_b=move_data_2, retain_sparsity=retain_sparsity) IF (PRESENT(pgrid_opt_1)) THEN IF (.NOT. new_1) THEN @@ -1796,57 +1752,6 @@ CONTAINS CALL timestop(handle) END SUBROUTINE -! ************************************************************************************************** -!> \brief get indices of non-zero tensor blocks for contraction result without actually -!> performing contraction. -!> this is an estimate based on block norm multiplication. -!> See documentation of dbt_contract. -!> \param nblks_local number of local blocks on this MPI rank -!> \param result_index indices of local non-zero tensor blocks for tensor_3 -!> only the elements result_index(:nblks_local, :) are relevant (all others are set to 0) -!> \author Patrick Seewald -! ************************************************************************************************** - SUBROUTINE dbt_contract_index(alpha, tensor_1, tensor_2, beta, tensor_3, & - contract_1, notcontract_1, & - contract_2, notcontract_2, & - map_1, map_2, & - bounds_1, bounds_2, bounds_3, & - filter_eps, & - nblks_local, result_index) - REAL(dp), INTENT(IN) :: alpha - TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_1 - TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_2 - REAL(dp), INTENT(IN) :: beta - INTEGER, DIMENSION(:), INTENT(IN) :: contract_1 - INTEGER, DIMENSION(:), INTENT(IN) :: contract_2 - INTEGER, DIMENSION(:), INTENT(IN) :: map_1 - INTEGER, DIMENSION(:), INTENT(IN) :: map_2 - INTEGER, DIMENSION(:), INTENT(IN) :: notcontract_1 - INTEGER, DIMENSION(:), INTENT(IN) :: notcontract_2 - TYPE(dbt_type), INTENT(INOUT), TARGET :: tensor_3 - INTEGER, DIMENSION(2, SIZE(contract_1)), & - INTENT(IN), OPTIONAL :: bounds_1 - INTEGER, DIMENSION(2, SIZE(notcontract_1)), & - INTENT(IN), OPTIONAL :: bounds_2 - INTEGER, DIMENSION(2, SIZE(notcontract_2)), & - INTENT(IN), OPTIONAL :: bounds_3 - REAL(KIND=dp), INTENT(IN), OPTIONAL :: filter_eps - INTEGER, INTENT(OUT) :: nblks_local - INTEGER, DIMENSION(dbt_max_nblks_local(tensor_3), ndims_tensor(tensor_3)), & - INTENT(OUT) :: result_index - - CALL dbt_contract_expert(alpha, tensor_1, tensor_2, beta, tensor_3, & - contract_1, notcontract_1, & - contract_2, notcontract_2, & - map_1, map_2, & - bounds_1=bounds_1, & - bounds_2=bounds_2, & - bounds_3=bounds_3, & - filter_eps=filter_eps, & - nblks_local=nblks_local, & - result_index=result_index) - END SUBROUTINE - ! ************************************************************************************************** !> \brief Map contraction bounds to bounds referring to tensor indices !> see dbt_contract for docu of dummy arguments diff --git a/src/dbt/dbt_reshape_ops.F b/src/dbt/dbt_reshape_ops.F index 173bbf4465..8d6e8e4a6f 100644 --- a/src/dbt/dbt_reshape_ops.F +++ b/src/dbt/dbt_reshape_ops.F @@ -104,14 +104,21 @@ CONTAINS ALLOCATE (num_send(0:2*numnodes - 1)) num_send(:) = 0 ALLOCATE (req_array(1:numnodes, 4)) + +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,tensor_out,num_send) & +!$OMP PRIVATE(iter,ind_nd,blk_size,iproc) CALL dbt_iterator_start(iter, tensor_in) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind_nd, blk_size=blk_size) CALL dbt_get_stored_coordinates(tensor_out, ind_nd, iproc) +!$OMP ATOMIC num_send(2*iproc) = num_send(2*iproc) + PRODUCT(blk_size) +!$OMP ATOMIC num_send(2*iproc + 1) = num_send(2*iproc + 1) + 1 END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL + CALL mp_alltoall(num_send, num_rec, 2, mp_comm) DO iproc = 0, numnodes - 1 num_entries_recv(iproc) = num_rec(2*iproc) @@ -124,16 +131,22 @@ CONTAINS CALL block_buffer_create(buffer_recv(iproc), num_blocks_recv(iproc), num_entries_recv(iproc), & ndims_tensor(tensor_in)) END DO + +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,tensor_out,buffer_send) & +!$OMP PRIVATE(iter,ind_nd,blk_size,blk_data,found,iproc) CALL dbt_iterator_start(iter, tensor_in) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind_nd, blk_size=blk_size) CALL dbt_get_block(tensor_in, ind_nd, blk_data, found) CPASSERT(found) CALL dbt_get_stored_coordinates(tensor_out, ind_nd, iproc) +!$OMP CRITICAL CALL block_buffer_add_anyd_block(buffer_send(iproc), ind_nd, blk_data) +!$OMP END CRITICAL CALL destroy_block(blk_data) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL IF (move_prv) CALL dbt_clear(tensor_in) @@ -153,7 +166,10 @@ CONTAINS DEALLOCATE (index_recv) END DO +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_out,blks_to_allocate) CALL dbt_reserve_blocks(tensor_out, blks_to_allocate) +!$OMP END PARALLEL DEALLOCATE (blks_to_allocate) DO iproc = 0, numnodes - 1 diff --git a/src/dbt/dbt_split.F b/src/dbt/dbt_split.F index f95f845bda..fc93b826de 100644 --- a/src/dbt/dbt_split.F +++ b/src/dbt/dbt_split.F @@ -21,11 +21,11 @@ MODULE dbt_split dbt_get_block, & dbt_put_block, & dbt_iterator_start, & + dbt_iterator_num_blocks, & dbt_iterator_blocks_left, & dbt_iterator_stop, & dbt_iterator_next_block, & - dbt_reserve_blocks, & - dbt_reserved_block_indices + dbt_reserve_blocks USE dbt_index, ONLY: dbt_get_mapping_info, & dbt_inverse_order USE dbt_types, ONLY: dbt_create, & @@ -172,6 +172,15 @@ CONTAINS END IF END IF +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(tensor_in,tensor_out) & +!$OMP SHARED(${varlist("blk_nsplit", nmax=ndim)}$) & +!$OMP SHARED(${varlist("inblock_offset", nmax=ndim)}$) & +!$OMP SHARED(${varlist("blk_size", nmax=ndim)}$) & +!$OMP SHARED(${varlist("index_split_offset", nmax=ndim)}$) & +!$OMP PRIVATE(iter,found,bcount,blks_to_allocate,bi_split,inblock_offset) & +!$OMP PRIVATE(blk_index,blk_size,blk_offset,blk_shape) & +!$OMP PRIVATE(block_2d,block_3d,block_4d) CALL dbt_iterator_start(iter, tensor_in) bcount = 0 @@ -254,6 +263,7 @@ CONTAINS #:endfor END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL dbt_finalize(tensor_out) @@ -341,7 +351,7 @@ CONTAINS INTEGER, DIMENSION(:), ALLOCATABLE :: ${varlist("inblock_offset")}$, ${varlist("blk_size_split")}$ INTEGER, DIMENSION(:, :), ALLOCATABLE :: blks_to_allocate INTEGER :: idim, iblk, bcount - INTEGER :: ${varlist("iblk")}$, isplit_sum, splitsum, nblk + INTEGER :: ${varlist("iblk")}$, isplit_sum, splitsum TYPE(dbt_iterator_type) :: iter INTEGER, DIMENSION(ndims_tensor(tensor_out)) :: blk_index, blk_size, blk_offset, blk_shape, blk_index_n LOGICAL :: found @@ -406,9 +416,17 @@ CONTAINS DEALLOCATE (blk_size_split_d, blk_size_d) END DO - nblk = dbt_get_num_blocks(tensor_split_in) - ALLOCATE (blks_to_allocate(nblk, ndims_tensor(tensor_split_in))) +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(tensor_split_in,tensor_out,summation) & +!$OMP SHARED(${varlist("split", nmax=ndim)}$) & +!$OMP SHARED(${varlist("first_split", nmax=ndim)}$) & +!$OMP SHARED(${varlist("last_split", nmax=ndim)}$) & +!$OMP SHARED(${varlist("inblock_offset", nmax=ndim)}$) & +!$OMP PRIVATE(iter,blks_to_allocate,bcount,blk_index_n) & +!$OMP PRIVATE(blk_index,blk_size,blk_shape,blk_offset,inblock_offset,found) & +!$OMP PRIVATE(block_2d,block_3d,block_4d,block_split_2d,block_split_3d,block_split_4d) CALL dbt_iterator_start(iter, tensor_split_in) + ALLOCATE (blks_to_allocate(dbt_iterator_num_blocks(iter), ndims_tensor(tensor_split_in))) bcount = 0 DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, blk_index, blk_size=blk_size) @@ -458,6 +476,7 @@ CONTAINS #:endfor END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) @@ -599,16 +618,21 @@ CONTAINS TYPE(dbt_type), INTENT(OUT) :: tensor_out INTEGER, DIMENSION(2, ndims_tensor(tensor_in)), INTENT(IN) :: bounds LOGICAL, INTENT(IN), OPTIONAL :: move_data + + CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_crop' + INTEGER, DIMENSION(2, ndims_tensor(tensor_in)) :: blk_bounds TYPE(dbt_iterator_type) :: iter INTEGER, DIMENSION(ndims_tensor(tensor_in)) :: blk_index, blk_size, blk_offset LOGICAL :: found, move_data_prv - INTEGER :: idim, iblk, iblk_all, nblk - INTEGER, DIMENSION(:, :), ALLOCATABLE :: blk_ind, blk_ind_tmp + INTEGER :: handle, idim, iblk_out + INTEGER, DIMENSION(:, :), ALLOCATABLE :: blk_ind_out #:for ndim in ndims REAL(dp), DIMENSION(${shape_colon(n=ndim)}$), ALLOCATABLE :: block_${ndim}$d, block_put_${ndim}$d #:endfor + CALL timeset(routineN, handle) + IF (PRESENT(move_data)) THEN move_data_prv = move_data ELSE @@ -617,29 +641,28 @@ CONTAINS CALL dbt_create(tensor_in, tensor_out) +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor_in,tensor_out,bounds) & +!$OMP PRIVATE(iter,blk_ind_out,iblk_out,blk_index,blk_size,blk_offset,found,blk_bounds) & +!$OMP PRIVATE(block_2d,block_put_2d,block_3d,block_put_3d,block_4d,block_put_4d) + ! reserve blocks inside bounds - ALLOCATE (blk_ind(dbt_get_num_blocks(tensor_in), ndims_tensor(tensor_in))) - CALL dbt_reserved_block_indices(tensor_in, blk_ind) - nblk = dbt_get_num_blocks(tensor_in) - ALLOCATE (blk_ind_tmp(dbt_get_num_blocks(tensor_in), ndims_tensor(tensor_in))) - blk_ind_tmp(:, :) = 0 - iblk = 0 - blk_loop: DO iblk_all = 1, nblk - CALL dbt_blk_offsets(tensor_in, blk_ind(iblk_all, :), blk_offset) - CALL dbt_blk_sizes(tensor_in, blk_ind(iblk_all, :), blk_size) + CALL dbt_iterator_start(iter, tensor_in) + ALLOCATE (blk_ind_out(dbt_iterator_num_blocks(iter), ndims_tensor(tensor_in))) + blk_ind_out(:, :) = 0 + iblk_out = 0 + blk_loop: DO WHILE (dbt_iterator_blocks_left(iter)) + CALL dbt_iterator_next_block(iter, blk_index, blk_size=blk_size, blk_offset=blk_offset) DO idim = 1, ndims_tensor(tensor_in) IF (bounds(1, idim) > blk_offset(idim) - 1 + blk_size(idim)) CYCLE blk_loop IF (bounds(2, idim) < blk_offset(idim)) CYCLE blk_loop END DO - iblk = iblk + 1 - blk_ind_tmp(iblk, :) = blk_ind(iblk_all, :) + iblk_out = iblk_out + 1 + blk_ind_out(iblk_out, :) = blk_index END DO blk_loop + CALL dbt_iterator_stop(iter) - DEALLOCATE (blk_ind) - ALLOCATE (blk_ind(iblk, ndims_tensor(tensor_in))) - blk_ind(:, :) = blk_ind_tmp(:iblk, :) - - CALL dbt_reserve_blocks(tensor_out, blk_ind) + CALL dbt_reserve_blocks(tensor_out, blk_ind_out(1:iblk_out, :)) + DEALLOCATE (blk_ind_out) ! copy blocks CALL dbt_iterator_start(iter, tensor_out) @@ -666,6 +689,7 @@ CONTAINS #:endfor END DO iter_loop CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL dbt_finalize(tensor_out) IF (move_data_prv) CALL dbt_clear(tensor_in) @@ -673,6 +697,7 @@ CONTAINS ! transfer data for batched contraction CALL dbt_copy_contraction_storage(tensor_in, tensor_out) + CALL timestop(handle) END SUBROUTINE END MODULE diff --git a/src/dbt/dbt_test.F b/src/dbt/dbt_test.F index f59aa5edac..865085621b 100644 --- a/src/dbt/dbt_test.F +++ b/src/dbt/dbt_test.F @@ -88,6 +88,8 @@ CONTAINS dbt_equal = .TRUE. +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor1,tensor2_tmp,dbt_equal) & +!$OMP PRIVATE(iter,ind_nd, blk_size,blk_data1,blk_data2,found) CALL dbt_iterator_start(iter, tensor1) DO WHILE (dbt_iterator_blocks_left(iter)) @@ -98,11 +100,14 @@ CONTAINS CPASSERT(found) IF (.NOT. blocks_equal(blk_data1, blk_data2)) THEN +!$OMP CRITICAL dbt_equal = .FALSE. +!$OMP END CRITICAL END IF END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL dbt_destroy(tensor2_tmp) END FUNCTION @@ -396,7 +401,16 @@ CONTAINS END IF ALLOCATE (ind_nd(nblks_alloc, ndims_tensor(tensor))) + +!$OMP PARALLEL DEFAULT(NONE) SHARED(ind_nd,nblks_alloc,tensor,mynode,enumerate,randmat_counter) & +!$OMP SHARED(${varlist("blk_ind")}$) & +!$OMP PRIVATE(my_nblks_alloc,ib,proc,i,iterator,blk_offset,blk_size,blk_index) & +!$OMP PRIVATE(blk_index_2d,nblks_2d,iseed,nze,tensor_dims) & +!$OMP PRIVATE(${varlist("blk_values", nmin=2)}$) & +!$OMP PRIVATE(${varlist("my_blk_ind")}$) + my_nblks_alloc = 0 +!$OMP DO DO ib = 1, nblks_alloc #:for ndim in ndims IF (ndims_tensor(tensor) == ${ndim}$) THEN @@ -408,6 +422,7 @@ CONTAINS my_nblks_alloc = my_nblks_alloc + 1 END IF END DO +!$OMP END DO #:for dim in range(1, maxdim+1) IF (ndims_tensor(tensor) >= ${dim}$) THEN @@ -416,6 +431,7 @@ CONTAINS #:endfor i = 0 +!$OMP DO DO ib = 1, nblks_alloc CALL dbt_get_stored_coordinates(tensor, ind_nd(ib, :), proc) IF (proc == mynode) THEN @@ -427,6 +443,7 @@ CONTAINS #:endfor END IF END DO +!$OMP END DO #:for ndim in ndims IF (ndims_tensor(tensor) == ${ndim}$) THEN @@ -460,6 +477,7 @@ CONTAINS #:endfor END DO CALL dbt_iterator_stop(iterator) +!$OMP END PARALLEL END SUBROUTINE @@ -523,6 +541,8 @@ CONTAINS CALL allocate_any(array, shape_spec=dims_nd) array(${shape_colon(ndim)}$) = 0.0_dp +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor,array) & +!$OMP PRIVATE(iterator,ind_nd,blk_size,blk_offset,block,found,idim,blk_start,blk_end) CALL dbt_iterator_start(iterator, tensor) DO WHILE (dbt_iterator_blocks_left(iterator)) CALL dbt_iterator_next_block(iterator, ind_nd, blk_size=blk_size, blk_offset=blk_offset) @@ -539,6 +559,7 @@ CONTAINS DEALLOCATE (block) END DO CALL dbt_iterator_stop(iterator) +!$OMP END PARALLEL CALL mp_sum(array, tensor%pgrid%mp_comm_2d) END SUBROUTINE diff --git a/src/dbt/tas/dbt_tas_base.F b/src/dbt/tas/dbt_tas_base.F index e9d9a5902b..954041aa59 100644 --- a/src/dbt/tas/dbt_tas_base.F +++ b/src/dbt/tas/dbt_tas_base.F @@ -18,8 +18,8 @@ MODULE dbt_tas_base dbm_get_block_p, dbm_get_col_block_sizes, dbm_get_distribution, dbm_get_local_cols, & dbm_get_local_rows, dbm_get_name, dbm_get_num_blocks, dbm_get_nze, & dbm_get_row_block_sizes, dbm_iterator, dbm_iterator_blocks_left, dbm_iterator_next_block, & - dbm_iterator_start, dbm_iterator_stop, dbm_put_block, dbm_release, dbm_reserve_blocks, & - dbm_type + dbm_iterator_num_blocks, dbm_iterator_start, dbm_iterator_stop, dbm_put_block, & + dbm_release, dbm_reserve_blocks, dbm_type USE dbt_tas_global, ONLY: dbt_tas_blk_size_arb,& dbt_tas_dist_arb,& dbt_tas_distribution,& @@ -36,7 +36,6 @@ MODULE dbt_tas_base dbt_tas_iterator,& dbt_tas_split_info,& dbt_tas_type - USE dbt_tas_util, ONLY: index_unique USE kinds, ONLY: default_string_length,& dp,& int_8 @@ -69,6 +68,7 @@ MODULE dbt_tas_base dbt_tas_get_num_blocks_total, & dbt_tas_get_stored_coordinates, & dbt_tas_info, & + dbt_tas_iterator_num_blocks, & dbt_tas_iterator_blocks_left, & dbt_tas_iterator_next_block, & dbt_tas_iterator_start, & @@ -266,12 +266,16 @@ CONTAINS CALL dbt_tas_reserve_blocks(matrix_a, matrix_b) +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_a,matrix_b,summation) & +!$OMP PRIVATE(iter,row,column,block) CALL dbt_tas_iterator_start(iter, matrix_a) DO WHILE (dbt_tas_iterator_blocks_left(iter)) CALL dbt_tas_iterator_next_block(iter, row, column, block) CALL dbt_tas_put_block(matrix_b, row, column, block, summation=summation) END DO CALL dbt_tas_iterator_stop(iter) +!$OMP END PARALLEL + CALL timestop(handle) END SUBROUTINE @@ -286,23 +290,29 @@ CONTAINS TYPE(dbt_tas_type), INTENT(IN) :: matrix_in TYPE(dbt_tas_type), INTENT(INOUT) :: matrix_out - INTEGER :: iblk, nblk - INTEGER(KIND=int_8) :: column, row - INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: col_res, row_res + CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_tas_reserve_blocks_template' + + INTEGER :: handle, iblk, nblk + INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: columns, rows TYPE(dbt_tas_iterator) :: iter - nblk = dbt_tas_get_num_blocks(matrix_in) - ALLOCATE (row_res(nblk), col_res(nblk)) + CALL timeset(routineN, handle) + +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out) & +!$OMP PRIVATE(iter,nblk,rows,columns) CALL dbt_tas_iterator_start(iter, matrix_in) + nblk = dbt_tas_iterator_num_blocks(iter) + ALLOCATE (rows(nblk), columns(nblk)) DO iblk = 1, nblk - CALL dbt_tas_iterator_next_block(iter, row, column) - row_res(iblk) = row - col_res(iblk) = column + CALL dbt_tas_iterator_next_block(iter, row=rows(iblk), column=columns(iblk)) END DO CPASSERT(.NOT. dbt_tas_iterator_blocks_left(iter)) CALL dbt_tas_iterator_stop(iter) - CALL dbt_tas_reserve_blocks(matrix_out, row_res, col_res) + CALL dbt_tas_reserve_blocks_index(matrix_out, rows=rows, columns=columns) +!$OMP END PARALLEL + + CALL timestop(handle) END SUBROUTINE ! ************************************************************************************************** @@ -554,8 +564,11 @@ CONTAINS CALL dbm_distribution_release(dist) DEALLOCATE (row_size_vec, col_size_vec) - nblks_local = dbt_tas_get_num_blocks(matrix_rect) + +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_rect,matrix_dbm) & +!$OMP PRIVATE(iter,nblks_local,nz_rows,nz_cols,rb_count,row,col,block) CALL dbt_tas_iterator_start(iter, matrix_rect) + nblks_local = dbt_tas_iterator_num_blocks(iter) ALLOCATE (nz_rows(nblks_local), nz_cols(nblks_local)) rb_count = 0 DO WHILE (dbt_tas_iterator_blocks_left(iter)) @@ -564,17 +577,17 @@ CONTAINS nz_rows(rb_count) = INT(row) nz_cols(rb_count) = INT(col) END DO + CALL dbt_tas_iterator_stop(iter) CALL dbm_reserve_blocks(matrix_dbm, nz_rows, nz_cols) - CALL dbt_tas_iterator_stop(iter) CALL dbt_tas_iterator_start(iter, matrix_rect) DO WHILE (dbt_tas_iterator_blocks_left(iter)) CALL dbt_tas_iterator_next_block(iter, row, col, block) CALL dbm_put_block(matrix_dbm, INT(row), INT(col), block) END DO - CALL dbt_tas_iterator_stop(iter) +!$OMP END PARALLEL CALL dbm_finalize(matrix_dbm) @@ -630,12 +643,14 @@ CONTAINS CALL dbt_tas_create(matrix_rect, TRIM(name)//"_compressed", & dist, row_blk_size_obj, col_blk_size_obj) +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_dbm,matrix_rect) PRIVATE(iter,row,col,block) CALL dbm_iterator_start(iter, matrix_dbm) DO WHILE (dbm_iterator_blocks_left(iter)) CALL dbm_iterator_next_block(iter, row, col, block) CALL dbt_tas_put_block(matrix_rect, INT(row, KIND=int_8), INT(col, KIND=int_8), block) END DO CALL dbm_iterator_stop(iter) +!$OMP END PARALLEL CALL dbt_tas_finalize(matrix_rect) @@ -650,13 +665,26 @@ CONTAINS ! ************************************************************************************************** SUBROUTINE dbt_tas_iterator_start(iter, matrix_in) TYPE(dbt_tas_iterator), INTENT(INOUT) :: iter - TYPE(dbt_tas_type), INTENT(IN) :: matrix_in + TYPE(dbt_tas_type), INTENT(IN), TARGET :: matrix_in CALL dbm_iterator_start(iter%iter, matrix_in%matrix) - iter%dist = matrix_in%dist + iter%dist => matrix_in%dist END SUBROUTINE +! ************************************************************************************************** +!> \brief As dbm_iterator_num_blocks +!> \param iter ... +!> \return ... +!> \author Ole Schuett +! ************************************************************************************************** + FUNCTION dbt_tas_iterator_num_blocks(iter) + TYPE(dbt_tas_iterator), INTENT(IN) :: iter + INTEGER :: dbt_tas_iterator_num_blocks + + dbt_tas_iterator_num_blocks = dbm_iterator_num_blocks(iter%iter) + END FUNCTION + ! ************************************************************************************************** !> \brief As dbm_iterator_blocks_left !> \param iter ... @@ -718,28 +746,19 @@ CONTAINS CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_tas_reserve_blocks_index' - INTEGER :: handle, icol, irow - INTEGER, ALLOCATABLE, DIMENSION(:, :) :: rowcols_group_u - INTEGER, DIMENSION(SIZE(rows), 2) :: rowcols_group + INTEGER :: handle, i + INTEGER, DIMENSION(SIZE(rows)) :: columns_group, rows_group CALL timeset(routineN, handle) - DO irow = 1, SIZE(rows) - CALL dbt_index_global_to_local(dbt_tas_info(matrix), matrix%dist, row=rows(irow), & - row_group=rowcols_group(irow, 1)) + CPASSERT(SIZE(rows) == SIZE(columns)) + DO i = 1, SIZE(rows) + CALL dbt_index_global_to_local(dbt_tas_info(matrix), matrix%dist, & + row=rows(i), row_group=rows_group(i), & + column=columns(i), column_group=columns_group(i)) END DO - DO icol = 1, SIZE(columns) - CALL dbt_index_global_to_local(dbt_tas_info(matrix), matrix%dist, column=columns(icol), & - column_group=rowcols_group(icol, 2)) - END DO - - ! TODO: While this was needed for DBCSR, it is no longer necessary for DBM. - CALL index_unique(rowcols_group, rowcols_group_u) ! make sure that index is unique, not sure - ! if this is really needed or whether DBM - ! takes care of duplicate indices - - CALL dbm_reserve_blocks(matrix%matrix, rowcols_group_u(:, 1), rowcols_group_u(:, 2)) + CALL dbm_reserve_blocks(matrix%matrix, rows_group, columns_group) CALL timestop(handle) END SUBROUTINE @@ -786,16 +805,8 @@ CONTAINS INTEGER(KIND=int_8), INTENT(IN) :: row, col INTEGER, INTENT(OUT) :: row_size, col_size - CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_tas_blk_sizes' - - INTEGER :: handle - - CALL timeset(routineN, handle) - row_size = matrix%row_blk_size%data(row) col_size = matrix%col_blk_size%data(col) - - CALL timestop(handle) END SUBROUTINE ! ************************************************************************************************** diff --git a/src/dbt/tas/dbt_tas_mm.F b/src/dbt/tas/dbt_tas_mm.F index d0e447f238..f8cbff7b98 100644 --- a/src/dbt/tas/dbt_tas_mm.F +++ b/src/dbt/tas/dbt_tas_mm.F @@ -23,10 +23,10 @@ MODULE dbt_tas_mm dbm_redistribute, dbm_release, dbm_scale, dbm_type, dbm_zero USE dbt_tas_base, ONLY: & dbt_tas_clear, dbt_tas_copy, dbt_tas_create, dbt_tas_destroy, dbt_tas_distribution_new, & - dbt_tas_filter, dbt_tas_get_info, dbt_tas_get_num_blocks, dbt_tas_get_nze_total, & - dbt_tas_info, dbt_tas_iterator_blocks_left, dbt_tas_iterator_next_block, & - dbt_tas_iterator_start, dbt_tas_iterator_stop, dbt_tas_nblkcols_total, & - dbt_tas_nblkrows_total, dbt_tas_put_block, dbt_tas_reserve_blocks + dbt_tas_filter, dbt_tas_get_info, dbt_tas_get_nze_total, dbt_tas_info, & + dbt_tas_iterator_blocks_left, dbt_tas_iterator_next_block, dbt_tas_iterator_start, & + dbt_tas_iterator_stop, dbt_tas_nblkcols_total, dbt_tas_nblkrows_total, dbt_tas_put_block, & + dbt_tas_reserve_blocks USE dbt_tas_global, ONLY: dbt_tas_blk_size_one,& dbt_tas_default_distvec,& dbt_tas_dist_arb,& @@ -71,7 +71,6 @@ MODULE dbt_tas_mm dbt_tas_multiply, & dbt_tas_batched_mm_init, & dbt_tas_batched_mm_finalize, & - dbt_tas_result_index, & dbt_tas_set_batched_state, & dbt_tas_batched_mm_complete @@ -101,14 +100,13 @@ CONTAINS !> (for internal use only) !> \param retain_sparsity ... !> \param simple_split ... -!> \param result_index ... !> \param unit_nr unit number for logging output !> \param log_verbose only for testing: verbose output !> \author Patrick Seewald ! ************************************************************************************************** RECURSIVE SUBROUTINE dbt_tas_multiply(transa, transb, transc, alpha, matrix_a, matrix_b, beta, matrix_c, & optimize_dist, split_opt, filter_eps, flop, move_data_a, & - move_data_b, retain_sparsity, simple_split, result_index, unit_nr, log_verbose) + move_data_b, retain_sparsity, simple_split, unit_nr, log_verbose) LOGICAL, INTENT(IN) :: transa, transb, transc REAL(dp), INTENT(IN) :: alpha @@ -121,8 +119,6 @@ CONTAINS INTEGER(KIND=int_8), INTENT(OUT), OPTIONAL :: flop LOGICAL, INTENT(IN), OPTIONAL :: move_data_a, move_data_b, & retain_sparsity, simple_split - INTEGER(int_8), ALLOCATABLE, DIMENSION(:, :), & - INTENT(OUT), OPTIONAL :: result_index INTEGER, INTENT(IN), OPTIONAL :: unit_nr LOGICAL, INTENT(IN), OPTIONAL :: log_verbose @@ -266,15 +262,9 @@ CONTAINS nze_b = dbt_tas_get_nze_total(matrix_b) IF (.NOT. simple_split_prv) THEN - CALL dbt_tas_result_index(transa, transb, transc, matrix_a, matrix_b, matrix_c, filter_eps, & - blk_ind=result_index, nze=nze_c, retain_sparsity=retain_sparsity) - - IF (PRESENT(result_index)) THEN - CALL mp_sync(matrix_a%dist%info%mp_comm) - CALL timestop(handle2) - CALL timestop(handle) - RETURN - END IF + CALL dbt_tas_estimate_result_nze(transa, transb, transc, matrix_a, matrix_b, matrix_c, & + estimated_nze=nze_c, filter_eps=filter_eps, & + retain_sparsity=retain_sparsity) max_mm_dim = MAXLOC(dims, 1) nsplit = split_factor_estimate(max_mm_dim, nze_a, nze_b, nze_c, numproc) @@ -1402,28 +1392,24 @@ CONTAINS !> \param matrix_a ... !> \param matrix_b ... !> \param matrix_c ... +!> \param estimated_nze ... !> \param filter_eps ... !> \param unit_nr ... -!> \param blk_ind ... -!> \param nze ... !> \param retain_sparsity ... !> \author Patrick Seewald ! ************************************************************************************************** - SUBROUTINE dbt_tas_result_index(transa, transb, transc, matrix_a, matrix_b, matrix_c, filter_eps, & - unit_nr, blk_ind, nze, retain_sparsity) + SUBROUTINE dbt_tas_estimate_result_nze(transa, transb, transc, matrix_a, matrix_b, matrix_c, & + estimated_nze, filter_eps, unit_nr, retain_sparsity) LOGICAL, INTENT(IN) :: transa, transb, transc TYPE(dbt_tas_type), INTENT(INOUT), TARGET :: matrix_a, matrix_b, matrix_c + INTEGER(int_8), INTENT(OUT) :: estimated_nze REAL(KIND=dp), INTENT(IN), OPTIONAL :: filter_eps INTEGER, INTENT(IN), OPTIONAL :: unit_nr - INTEGER(int_8), ALLOCATABLE, DIMENSION(:, :), & - INTENT(OUT), OPTIONAL :: blk_ind - INTEGER(int_8), INTENT(OUT), OPTIONAL :: nze LOGICAL, INTENT(IN), OPTIONAL :: retain_sparsity - CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_tas_result_index' + CHARACTER(LEN=*), PARAMETER :: routineN = 'dbt_tas_estimate_result_nze' - INTEGER :: col_size, handle, iblk, mp_comm, nblk, & - row_size + INTEGER :: col_size, handle, mp_comm, row_size INTEGER(int_8) :: col, row LOGICAL :: retain_sparsity_prv TYPE(dbt_tas_iterator) :: iter @@ -1456,24 +1442,21 @@ CONTAINS matrix_c_bnorm => matrix_c END IF - nblk = dbt_tas_get_num_blocks(matrix_c_bnorm) - IF (PRESENT(blk_ind)) ALLOCATE (blk_ind(nblk, 2)) - + estimated_nze = 0 +!$OMP PARALLEL DEFAULT(NONE) REDUCTION(+:estimated_nze) SHARED(matrix_c_bnorm,matrix_c) & +!$OMP PRIVATE(iter,row,col,row_size,col_size) CALL dbt_tas_iterator_start(iter, matrix_c_bnorm) - IF (PRESENT(nze)) nze = 0 - DO iblk = 1, nblk + DO WHILE (dbt_tas_iterator_blocks_left(iter)) CALL dbt_tas_iterator_next_block(iter, row, col) row_size = matrix_c%row_blk_size%data(row) col_size = matrix_c%col_blk_size%data(col) - IF (PRESENT(nze)) nze = nze + row_size*col_size - IF (PRESENT(blk_ind)) blk_ind(iblk, :) = [row, col] + estimated_nze = estimated_nze + row_size*col_size END DO CALL dbt_tas_iterator_stop(iter) +!$OMP END PARALLEL - IF (PRESENT(nze)) THEN - CALL dbt_tas_get_split_info(dbt_tas_info(matrix_a), mp_comm=mp_comm) - CALL mp_sum(nze, mp_comm) - END IF + CALL dbt_tas_get_split_info(dbt_tas_info(matrix_a), mp_comm=mp_comm) + CALL mp_sum(estimated_nze, mp_comm) IF (.NOT. retain_sparsity_prv) THEN CALL dbt_tas_destroy(matrix_c_bnorm) @@ -1541,8 +1524,8 @@ CONTAINS CHARACTER(len=default_string_length) :: name INTEGER(KIND=int_8) :: column, nblkcols, nblkrows, row LOGICAL :: nodata_prv - REAL(dp), DIMENSION(1, 1) :: dbt_put - REAL(dp), DIMENSION(:, :), POINTER :: dbt_get + REAL(dp), DIMENSION(1, 1) :: blk_put + REAL(dp), DIMENSION(:, :), POINTER :: blk_get TYPE(dbt_tas_blk_size_one) :: col_blk_size, row_blk_size TYPE(dbt_tas_iterator) :: iter @@ -1565,17 +1548,16 @@ CONTAINS IF (.NOT. nodata_prv) THEN CALL dbt_tas_reserve_blocks(matrix_in, matrix_out) - +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out) & +!$OMP PRIVATE(iter,row,column,blk_get,blk_put) CALL dbt_tas_iterator_start(iter, matrix_in) - DO WHILE (dbt_tas_iterator_blocks_left(iter)) - CALL dbt_tas_iterator_next_block(iter, row, column, dbt_get) - !CALL dbt_tas_get_block_p(matrix_out, row, column, dbt_put) - !CPASSERT(ASSOCIATED(dbt_put)) - dbt_put(1, 1) = NORM2(dbt_get) - CALL dbt_tas_put_block(matrix_out, row, column, dbt_put) + CALL dbt_tas_iterator_next_block(iter, row, column, blk_get) + blk_put(1, 1) = NORM2(blk_get) + CALL dbt_tas_put_block(matrix_out, row, column, blk_put) END DO CALL dbt_tas_iterator_stop(iter) +!$OMP END PARALLEL END IF END SUBROUTINE diff --git a/src/dbt/tas/dbt_tas_reshape_ops.F b/src/dbt/tas/dbt_tas_reshape_ops.F index fe8808f322..8d29c09b15 100644 --- a/src/dbt/tas/dbt_tas_reshape_ops.F +++ b/src/dbt/tas/dbt_tas_reshape_ops.F @@ -34,8 +34,7 @@ MODULE dbt_tas_reshape_ops dbt_tas_iterator,& dbt_tas_split_info,& dbt_tas_type - USE dbt_tas_util, ONLY: index_unique,& - swap + USE dbt_tas_util, ONLY: swap USE kinds, ONLY: dp,& int_8 USE message_passing, ONLY: mp_alltoall,& @@ -43,6 +42,9 @@ MODULE dbt_tas_reshape_ops mp_irecv,& mp_isend,& mp_waitall + +!$ USE OMP_LIB, ONLY: omp_lock_kind, omp_init_lock + #include "../../base/base_uses.f90" IMPLICIT NONE @@ -83,6 +85,8 @@ CONTAINS mynode, nblk, ndata, numnodes INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: blks_to_allocate, index_recv INTEGER(KIND=int_8), DIMENSION(2) :: blk_index + INTEGER(kind=omp_lock_kind), ALLOCATABLE, & + DIMENSION(:) :: locks INTEGER, ALLOCATABLE, DIMENSION(:) :: num_blocks_recv, num_blocks_send, & num_entries_recv, num_entries_send, & num_rec, num_send @@ -131,23 +135,30 @@ CONTAINS ALLOCATE (num_send(0:2*numnodes - 1)) num_send(:) = 0 ALLOCATE (req_array(1:numnodes, 4)) - CALL dbt_tas_iterator_start(iter, matrix_in) + ALLOCATE (locks(0:numnodes - 1)) + DO iproc = 0, numnodes - 1 + CALL omp_init_lock(locks(iproc)) + END DO CALL timeset(routineN//"_get_coord", handle2) +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out,tr_in,num_send) & +!$OMP PRIVATE(iter,blk_index,blk_size,iproc) + CALL dbt_tas_iterator_start(iter, matrix_in) DO WHILE (dbt_tas_iterator_blocks_left(iter)) CALL dbt_tas_iterator_next_block(iter, blk_index(1), blk_index(2), & row_size=blk_size(1), col_size=blk_size(2)) - IF (tr_in) THEN CALL dbt_tas_get_stored_coordinates(matrix_out, blk_index(2), blk_index(1), iproc) ELSE CALL dbt_tas_get_stored_coordinates(matrix_out, blk_index(1), blk_index(2), iproc) END IF - +!$OMP ATOMIC num_send(2*iproc) = num_send(2*iproc) + PRODUCT(blk_size) +!$OMP ATOMIC num_send(2*iproc + 1) = num_send(2*iproc + 1) + 1 END DO CALL dbt_tas_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle2) CALL timeset(routineN//"_alltoall", handle2) @@ -167,6 +178,8 @@ CONTAINS END DO +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out,tr_in,buffer_send,locks) & +!$OMP PRIVATE(iter,blk_index,blk_size,block,iproc) CALL dbt_tas_iterator_start(iter, matrix_in) DO WHILE (dbt_tas_iterator_blocks_left(iter)) CALL dbt_tas_iterator_next_block(iter, blk_index(1), blk_index(2), block, & @@ -176,9 +189,12 @@ CONTAINS ELSE CALL dbt_tas_get_stored_coordinates(matrix_out, blk_index(1), blk_index(2), iproc) END IF + CALL omp_set_lock(locks(iproc)) CALL dbt_buffer_add_block(buffer_send(iproc), blk_index, block, transposed=tr_in) + CALL omp_unset_lock(locks(iproc)) END DO CALL dbt_tas_iterator_stop(iter) +!$OMP END PARALLEL IF (move_prv) CALL dbt_tas_clear(matrix_in) @@ -195,6 +211,7 @@ CONTAINS CALL timeset(routineN//"_buffer_obtain", handle2) + ! TODO Add OpenMP to the buffer unpacking. nblk = SUM(num_blocks_recv) ALLOCATE (blks_to_allocate(nblk, 2)) @@ -206,7 +223,10 @@ CONTAINS DEALLOCATE (index_recv) END DO +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_out,blks_to_allocate) CALL dbt_tas_reserve_blocks(matrix_out, blks_to_allocate(:, 1), blks_to_allocate(:, 2)) +!$OMP END PARALLEL DEALLOCATE (blks_to_allocate) DO iproc = 0, numnodes - 1 @@ -257,7 +277,8 @@ CONTAINS TYPE(dbt_tas_blk_size_repl), TARGET :: repl_blksize TYPE(dbt_tas_blk_size_arb), TARGET :: dir_blksize TYPE(dbt_tas_distribution_type) :: dist - INTEGER :: mp_comm, numnodes, mynode + INTEGER :: mp_comm, numnodes, mynode, ngroup + INTEGER(kind=omp_lock_kind), ALLOCATABLE, DIMENSION(:) :: locks TYPE(dbt_buffer_type), ALLOCATABLE, DIMENSION(:) :: buffer_recv, buffer_send INTEGER, ALLOCATABLE, DIMENSION(:) :: num_blocks_recv, num_blocks_send, & num_entries_recv, num_entries_send, & @@ -304,6 +325,7 @@ CONTAINS col_dist => dbm_distribution_col_dist(dbm_dist) mp_comm = info%mp_comm + ngroup = info%ngroup CALL mp_environ(numnodes, mynode, mp_comm) CALL mp_environ(numnodes, pdims, pcoord, mp_comm) @@ -349,19 +371,29 @@ CONTAINS ALLOCATE (num_send(0:2*numnodes - 1)) num_send(:) = 0 ALLOCATE (req_array(1:numnodes, 4)) + ALLOCATE (locks(0:numnodes - 1)) + DO iproc = 0, numnodes - 1 + CALL omp_init_lock(locks(iproc)) + END DO - ALLOCATE (iprocs(info%ngroup)) +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out,num_send,ngroup) & +!$OMP PRIVATE(iter,blk_index,blk_size,iprocs) + ALLOCATE (iprocs(ngroup)) CALL dbm_iterator_start(iter, matrix_in) DO WHILE (dbm_iterator_blocks_left(iter)) CALL dbm_iterator_next_block(iter, blk_index(1), blk_index(2), & row_size=blk_size(1), col_size=blk_size(2)) CALL dbt_repl_get_stored_coordinates(matrix_out, blk_index(1), blk_index(2), iprocs) DO i = 1, SIZE(iprocs) +!$OMP ATOMIC num_send(2*iprocs(i)) = num_send(2*iprocs(i)) + PRODUCT(blk_size) +!$OMP ATOMIC num_send(2*iprocs(i) + 1) = num_send(2*iprocs(i) + 1) + 1 END DO END DO CALL dbm_iterator_stop(iter) + DEALLOCATE (iprocs) +!$OMP END PARALLEL CALL timeset(routineN//"_alltoall", handle2) CALL mp_alltoall(num_send, num_rec, 2, mp_comm) @@ -379,16 +411,23 @@ CONTAINS END DO +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out,buffer_send,locks,ngroup) & +!$OMP PRIVATE(iter,blk_index,blk_size,block,iprocs) + ALLOCATE (iprocs(ngroup)) CALL dbm_iterator_start(iter, matrix_in) DO WHILE (dbm_iterator_blocks_left(iter)) CALL dbm_iterator_next_block(iter, blk_index(1), blk_index(2), block, & row_size=blk_size(1), col_size=blk_size(2)) CALL dbt_repl_get_stored_coordinates(matrix_out, blk_index(1), blk_index(2), iprocs) DO i = 1, SIZE(iprocs) + CALL omp_set_lock(locks(iprocs(i))) CALL dbt_buffer_add_block(buffer_send(iprocs(i)), INT(blk_index, KIND=int_8), block) + CALL omp_unset_lock(locks(iprocs(i))) END DO END DO CALL dbm_iterator_stop(iter) + DEALLOCATE (iprocs) +!$OMP END PARALLEL IF (move_prv) CALL dbm_clear(matrix_in) @@ -401,6 +440,7 @@ CONTAINS CALL timestop(handle2) + ! TODO Add OpenMP to the buffer unpacking. nblk = SUM(num_blocks_recv) ALLOCATE (blks_to_allocate(nblk, 2)) @@ -412,7 +452,10 @@ CONTAINS DEALLOCATE (index_recv) END DO +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_out,blks_to_allocate) CALL dbm_reserve_blocks(matrix_out%matrix, blks_to_allocate(:, 1), blks_to_allocate(:, 2)) +!$OMP END PARALLEL DEALLOCATE (blks_to_allocate) DO iproc = 0, numnodes - 1 @@ -454,11 +497,12 @@ CONTAINS mynode, nblk, ndata, numnodes INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: index_recv INTEGER(KIND=int_8), DIMENSION(2) :: blk_index_i8 - INTEGER, ALLOCATABLE, DIMENSION(:) :: iprocs, num_blocks_recv, & - num_blocks_send, num_entries_recv, & - num_entries_send, num_rec, num_send - INTEGER, ALLOCATABLE, DIMENSION(:, :) :: blks_to_allocate, blks_to_allocate_u, & - req_array + INTEGER(kind=omp_lock_kind), ALLOCATABLE, & + DIMENSION(:) :: locks + INTEGER, ALLOCATABLE, DIMENSION(:) :: num_blocks_recv, num_blocks_send, & + num_entries_recv, num_entries_send, & + num_rec, num_send + INTEGER, ALLOCATABLE, DIMENSION(:, :) :: blks_to_allocate, req_array INTEGER, DIMENSION(2) :: blk_index, blk_size INTEGER, DIMENSION(:), POINTER :: col_block_sizes, row_block_sizes LOGICAL :: move_prv @@ -497,18 +541,25 @@ CONTAINS ALLOCATE (num_send(0:2*numnodes - 1)) num_send(:) = 0 ALLOCATE (req_array(1:numnodes, 4)) + ALLOCATE (locks(0:numnodes - 1)) + DO iproc = 0, numnodes - 1 + CALL omp_init_lock(locks(iproc)) + END DO - ALLOCATE (iprocs(info%ngroup)) - +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out,num_send) & +!$OMP PRIVATE(iter,blk_index,blk_size,iproc) CALL dbm_iterator_start(iter, matrix_in%matrix) DO WHILE (dbm_iterator_blocks_left(iter)) CALL dbm_iterator_next_block(iter, blk_index(1), blk_index(2), & row_size=blk_size(1), col_size=blk_size(2)) CALL dbm_get_stored_coordinates(matrix_out, blk_index(1), blk_index(2), iproc) +!$OMP ATOMIC num_send(2*iproc) = num_send(2*iproc) + PRODUCT(blk_size) +!$OMP ATOMIC num_send(2*iproc + 1) = num_send(2*iproc + 1) + 1 END DO CALL dbm_iterator_stop(iter) +!$OMP END PARALLEL CALL timeset(routineN//"_alltoall", handle2) CALL mp_alltoall(num_send, num_rec, 2, mp_comm) @@ -526,15 +577,19 @@ CONTAINS END DO +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_in,matrix_out,buffer_send,locks) & +!$OMP PRIVATE(iter,blk_index,blk_size,block,iproc) CALL dbm_iterator_start(iter, matrix_in%matrix) DO WHILE (dbm_iterator_blocks_left(iter)) CALL dbm_iterator_next_block(iter, blk_index(1), blk_index(2), block, & row_size=blk_size(1), col_size=blk_size(2)) CALL dbm_get_stored_coordinates(matrix_out, blk_index(1), blk_index(2), iproc) + CALL omp_set_lock(locks(iproc)) CALL dbt_buffer_add_block(buffer_send(iproc), INT(blk_index, KIND=int_8), block) + CALL omp_unset_lock(locks(iproc)) END DO - CALL dbm_iterator_stop(iter) +!$OMP END PARALLEL IF (move_prv) CALL dbt_tas_clear(matrix_in) @@ -547,6 +602,7 @@ CONTAINS CALL timestop(handle2) + ! TODO Add OpenMP to the buffer unpacking. nblk = SUM(num_blocks_recv) ALLOCATE (blks_to_allocate(nblk, 2)) @@ -558,10 +614,11 @@ CONTAINS DEALLOCATE (index_recv) END DO - CALL index_unique(blks_to_allocate, blks_to_allocate_u) - - CALL dbm_reserve_blocks(matrix_out, blks_to_allocate_u(:, 1), blks_to_allocate_u(:, 2)) - DEALLOCATE (blks_to_allocate, blks_to_allocate_u) +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(matrix_out,blks_to_allocate) + CALL dbm_reserve_blocks(matrix_out, blks_to_allocate(:, 1), blks_to_allocate(:, 2)) +!$OMP END PARALLEL + DEALLOCATE (blks_to_allocate) DO iproc = 0, numnodes - 1 ! First, we need to get the index to create block diff --git a/src/dbt/tas/dbt_tas_types.F b/src/dbt/tas/dbt_tas_types.F index 563c3f11d4..6cd3e832f1 100644 --- a/src/dbt/tas/dbt_tas_types.F +++ b/src/dbt/tas/dbt_tas_types.F @@ -87,8 +87,7 @@ MODULE dbt_tas_types END TYPE TYPE dbt_tas_iterator - TYPE(dbt_tas_split_info) :: info - TYPE(dbt_tas_distribution_type) :: dist + TYPE(dbt_tas_distribution_type), POINTER :: dist => NULL() TYPE(dbm_iterator) :: iter END TYPE dbt_tas_iterator diff --git a/src/dbt/tas/dbt_tas_util.F b/src/dbt/tas/dbt_tas_util.F index aa20674183..54948145e3 100644 --- a/src/dbt/tas/dbt_tas_util.F +++ b/src/dbt/tas/dbt_tas_util.F @@ -29,7 +29,6 @@ MODULE dbt_tas_util PUBLIC :: & array_eq, & - index_unique, & swap, & generate_larnv_seed @@ -75,115 +74,6 @@ CONTAINS arr(2) = tmp END SUBROUTINE -! ************************************************************************************************** -!> \brief Get all unique elements in index_in -!> \param index_in ... -!> \param index_out ... -!> \author Patrick Seewald -! ************************************************************************************************** - SUBROUTINE index_unique(index_in, index_out) - INTEGER, DIMENSION(:, :), INTENT(IN) :: index_in - INTEGER, ALLOCATABLE, DIMENSION(:, :), INTENT(OUT) :: index_out - - INTEGER :: blk, count, orig_size - INTEGER, ALLOCATABLE, DIMENSION(:, :) :: index_tmp - INTEGER, DIMENSION(2) :: prev_index - INTEGER, DIMENSION(1:SIZE(index_in, 1), 1:SIZE(& - index_in, 2)) :: index_sorted - - orig_size = SIZE(index_in, 1) - ALLOCATE (index_tmp(orig_size, 2)) - index_sorted(:, :) = index_in(:, :) - CALL sort_indices(orig_size, index_sorted(:, 1), index_sorted(:, 2)) - count = 0 - prev_index(:) = [0, 0] - DO blk = 1, orig_size - IF (ANY(index_sorted(blk, :) .NE. prev_index(:))) THEN - count = count + 1 - index_tmp(count, :) = index_sorted(blk, :) - prev_index(:) = index_sorted(blk, :) - END IF - END DO - - ALLOCATE (index_out(count, 2)) - index_out(:, :) = index_tmp(1:count, :) - END SUBROUTINE - -! ************************************************************************************************** -!> \brief Sorts the rows & columns of a work matrix -!> Sorts the row and column indices so that the rows monotonically -!> increase and the columns monotonically increase within each row. -!> Passing the blk_p array rearranges the block pointers accordingly. -!> This must be done if they are pointing to valid data, otherwise -!> they become invalid. -!> TODO: Simplify using array_sort.fypp and merge with index_unique above. -!> \param n number of blocks (elements) to sort -!> \param row_i row indices -!> \param col_i column indices -!> \param blk_p block pointers -!> \param blk_d data storage -!> \author Patrick Seewald -! ************************************************************************************************** - SUBROUTINE sort_indices(n, row_i, col_i, blk_p, blk_d) - INTEGER, INTENT(IN) :: n - INTEGER, DIMENSION(1:), INTENT(INOUT) :: row_i, col_i - INTEGER, DIMENSION(1:), INTENT(INOUT), OPTIONAL :: blk_p, blk_d - - CHARACTER(len=*), PARAMETER :: routineN = 'sort_indices' - INTEGER(KIND=int_8), PARAMETER :: lmask8 = 4294967295_int_8 - - INTEGER :: error_handle, i - INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: sort_keys - INTEGER, ALLOCATABLE, DIMENSION(:) :: buf, buf_d - -! --------------------------------------------------------------------------- - - IF (SIZE(row_i) .EQ. 0) RETURN - - CALL timeset(routineN, error_handle) - - IF (SIZE(row_i) < n) CPABORT('row_i too small') - IF (SIZE(col_i) < n) CPABORT('col_i too small') - IF (PRESENT(blk_p)) THEN - IF (SIZE(blk_p) < n) CPABORT('blk_p too small') - ALLOCATE (buf(n)) - buf(1:n) = blk_p(1:n) - END IF - IF (PRESENT(blk_d)) THEN - ALLOCATE (buf_d(n)) - buf_d(1:n) = blk_d(1:n) - END IF - ! Create an ordering for both rows and columns. If the blk_p must - ! be rearranged, then the col_i array will be used as a - ! permutation vector. - ALLOCATE (sort_keys(n)) - sort_keys(:) = IOR(ISHFT(INT(row_i(1:n), int_8), 32), INT(col_i(1:n), int_8)) - IF (PRESENT(blk_p)) col_i(1:n) = (/(i, i=1, n)/) - ! Now do a nice quicksort. - CALL sort(sort_keys, n, col_i) - ! Since blk_d is usually not present we can have two loops that - ! are essentially the same. - IF (PRESENT(blk_p)) THEN - DO i = 1, n - blk_p(i) = buf(col_i(i)) - END DO - DEALLOCATE (buf) - END IF - IF (PRESENT(blk_d)) THEN - DO i = 1, n - blk_d(i) = buf_d(col_i(i)) - END DO - DEALLOCATE (buf_d) - END IF - DO i = 1, n - col_i(i) = INT(IAND(sort_keys(i), lmask8), int_4) - row_i(i) = INT(ISHFT(sort_keys(i), -32), int_4) - END DO - DEALLOCATE (sort_keys) - CALL timestop(error_handle) - - END SUBROUTINE sort_indices - ! ************************************************************************************************** !> \brief ... !> \param arr1 ... diff --git a/src/hfx_ri.F b/src/hfx_ri.F index dfb7dd7ebd..ee827c0d7b 100644 --- a/src/hfx_ri.F +++ b/src/hfx_ri.F @@ -3133,6 +3133,10 @@ CONTAINS deriv_dim_prv = 1 IF (PRESENT(deriv_dim)) deriv_dim_prv = deriv_dim +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(t_3c_der,t_3c_contr,work_virial,force,use_virial,do_mp2_prv,deriv_dim_prv) & +!$OMP SHARED(pref,idx_to_at,atom_of_kind,kind_of,particle_set,cell) & +!$OMP PRIVATE(i_xyz,j_xyz,iter,ind,der_blk,contr_blk,found,new_force,iat,iat_of_kind,ikind,scoord) DO i_xyz = 1, 3 CALL dbt_iterator_start(iter, t_3c_der(i_xyz)) DO WHILE (dbt_iterator_blocks_left(iter)) @@ -3153,9 +3157,11 @@ CONTAINS ikind = kind_of(iat) IF (.NOT. do_mp2_prv) THEN +!$OMP ATOMIC force(ikind)%fock_4c(i_xyz, iat_of_kind) = force(ikind)%fock_4c(i_xyz, iat_of_kind) & + new_force ELSE +!$OMP ATOMIC force(ikind)%mp2_non_sep(i_xyz, iat_of_kind) = force(ikind)%mp2_non_sep(i_xyz, iat_of_kind) & + new_force END IF @@ -3164,6 +3170,7 @@ CONTAINS CALL real_to_scaled(scoord, particle_set(iat)%r, cell) DO j_xyz = 1, 3 +!$OMP ATOMIC work_virial(i_xyz, j_xyz) = work_virial(i_xyz, j_xyz) + new_force*scoord(j_xyz) END DO END IF @@ -3173,6 +3180,7 @@ CONTAINS END DO !iter CALL dbt_iterator_stop(iter) END DO +!$OMP END PARALLEL CALL timestop(handle) END SUBROUTINE get_force_from_3c_trace @@ -3220,6 +3228,10 @@ CONTAINS use_virial = .FALSE. IF (PRESENT(work_virial) .AND. PRESENT(cell) .AND. PRESENT(particle_set)) use_virial = .TRUE. +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(t_2c_MO_AO,t_mo_coeff,pref,use_virial,particle_set,cell,force,work_virial) & +!$OMP SHARED(idx_to_at,atom_of_kind,kind_of,i_xyz) & +!$OMP PRIVATE(iter,ind,mo_ao_blk,mo_coeff_blk,found,new_force,iat,iat_of_kind,ikind,scoord,j_xyz) CALL dbt_iterator_start(iter, t_2c_MO_AO) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind) @@ -3236,6 +3248,7 @@ CONTAINS iat_of_kind = atom_of_kind(iat) ikind = kind_of(iat) +!$OMP ATOMIC force(ikind)%fock_4c(i_xyz, iat_of_kind) = force(ikind)%fock_4c(i_xyz, iat_of_kind) & + new_force @@ -3244,6 +3257,7 @@ CONTAINS CALL real_to_scaled(scoord, particle_set(iat)%r, cell) DO j_xyz = 1, 3 +!$OMP ATOMIC work_virial(i_xyz, j_xyz) = work_virial(i_xyz, j_xyz) + new_force*scoord(j_xyz) END DO END IF @@ -3254,6 +3268,7 @@ CONTAINS DEALLOCATE (mo_ao_blk) END DO !iter CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) @@ -3315,6 +3330,11 @@ CONTAINS do_ovlp_prv = .FALSE. IF (PRESENT(do_ovlp)) do_ovlp_prv = do_ovlp +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(t_2c_der,t_2c_contr,work_virial,force,use_virial,do_mp2_prv,do_ovlp_prv) & +!$OMP SHARED(pref,idx_to_at,atom_of_kind,kind_of,particle_set,cell) & +!$OMP PRIVATE(i_xyz,j_xyz,iter,ind,der_blk,contr_blk,found,new_force) & +!$OMP PRIVATE(iat,jat,iat_of_kind,jat_of_kind,ikind,jkind,scoord) DO i_xyz = 1, 3 CALL dbt_iterator_start(iter, t_2c_der(i_xyz)) DO WHILE (dbt_iterator_blocks_left(iter)) @@ -3337,12 +3357,15 @@ CONTAINS ikind = kind_of(iat) IF (do_mp2_prv) THEN +!$OMP ATOMIC force(ikind)%mp2_non_sep(i_xyz, iat_of_kind) = force(ikind)%mp2_non_sep(i_xyz, iat_of_kind) & + new_force ELSE IF (do_ovlp_prv) THEN +!$OMP ATOMIC force(ikind)%overlap(i_xyz, iat_of_kind) = force(ikind)%overlap(i_xyz, iat_of_kind) & + new_force ELSE +!$OMP ATOMIC force(ikind)%fock_4c(i_xyz, iat_of_kind) = force(ikind)%fock_4c(i_xyz, iat_of_kind) & + new_force END IF @@ -3352,6 +3375,7 @@ CONTAINS CALL real_to_scaled(scoord, particle_set(iat)%r, cell) DO j_xyz = 1, 3 +!$OMP ATOMIC work_virial(i_xyz, j_xyz) = work_virial(i_xyz, j_xyz) + new_force*scoord(j_xyz) END DO END IF @@ -3361,13 +3385,15 @@ CONTAINS jkind = kind_of(jat) IF (do_mp2_prv) THEN +!$OMP ATOMIC force(jkind)%mp2_non_sep(i_xyz, jat_of_kind) = force(jkind)%mp2_non_sep(i_xyz, jat_of_kind) & - new_force ELSE IF (do_ovlp_prv) THEN +!$OMP ATOMIC force(jkind)%overlap(i_xyz, jat_of_kind) = force(jkind)%overlap(i_xyz, jat_of_kind) & - new_force ELSE - +!$OMP ATOMIC force(jkind)%fock_4c(i_xyz, jat_of_kind) = force(jkind)%fock_4c(i_xyz, jat_of_kind) & - new_force END IF @@ -3377,6 +3403,7 @@ CONTAINS CALL real_to_scaled(scoord, particle_set(jat)%r, cell) DO j_xyz = 1, 3 +!$OMP ATOMIC work_virial(i_xyz, j_xyz) = work_virial(i_xyz, j_xyz) - new_force*scoord(j_xyz) END DO END IF @@ -3389,7 +3416,7 @@ CONTAINS CALL dbt_iterator_stop(iter) END DO !i_xyz - +!$OMP END PARALLEL CALL timestop(handle) END SUBROUTINE get_2c_der_force diff --git a/src/mp2_integrals.F b/src/mp2_integrals.F index 18f1de809c..0e30a17593 100644 --- a/src/mp2_integrals.F +++ b/src/mp2_integrals.F @@ -1001,7 +1001,9 @@ CONTAINS END IF ! reserve (R|ai) block +!$OMP PARALLEL DEFAULT(NONE) SHARED(tB_in,loc_P,loc_a) CALL dbt_reserve_blocks(tB_in, [loc_P], [loc_a], [1]) +!$OMP END PARALLEL ! reserve (R|P) blocks ! in my_Lrows, R index is replicated. For (R|P), we distribute quadratic blocks cyclically over @@ -1019,7 +1021,11 @@ CONTAINS block_ind_L_P(ii) = loc_P END IF END DO + +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(tL,block_ind_L_R,block_ind_L_P,ii) CALL dbt_reserve_blocks(tL, block_ind_L_R(1:ii), block_ind_L_P(1:ii)) +!$OMP END PARALLEL ! insert (R|ai) block CALL dbt_put_block(tB_in, [loc_P, loc_a, 1], SHAPE(BIb_C), BIb_C) diff --git a/src/qs_tensors.F b/src/qs_tensors.F index 1817b0ca94..dde565d461 100644 --- a/src/qs_tensors.F +++ b/src/qs_tensors.F @@ -35,7 +35,8 @@ MODULE qs_tensors USE dbt_api, ONLY: & dbt_blk_sizes, dbt_clear, dbt_copy, dbt_create, dbt_destroy, dbt_filter, dbt_get_block, & dbt_get_info, dbt_get_num_blocks, dbt_get_nze_total, dbt_get_stored_coordinates, & - dbt_ndims, dbt_put_block, dbt_reserve_blocks, dbt_reserved_block_indices, dbt_type + dbt_iterator_next_block, dbt_iterator_num_blocks, dbt_iterator_start, dbt_iterator_stop, & + dbt_iterator_type, dbt_ndims, dbt_put_block, dbt_reserve_blocks, dbt_type USE distribution_1d_types, ONLY: distribution_1d_type USE distribution_2d_types, ONLY: distribution_2d_type USE gamma, ONLY: init_md_ftable @@ -623,10 +624,10 @@ CONTAINS CHARACTER(LEN=*), PARAMETER :: routineN = 'alloc_block_3c' - INTEGER :: handle, i, i_img, iatom, ikind, iproc, & - j_img, jatom, jcell, jkind, katom, & - kcell, kkind, natom, nimg, op_ij, & - op_jk, op_pos_prv + INTEGER :: handle, i_img, iatom, ikind, j_img, & + jatom, jcell, jkind, katom, kcell, & + kkind, natom, nimg, op_ij, op_jk, & + op_pos_prv INTEGER(int_8), ALLOCATABLE, DIMENSION(:, :) :: nblk INTEGER, DIMENSION(3) :: cell_j, cell_k, kp_index_lbounds, & kp_index_ubounds @@ -806,17 +807,11 @@ CONTAINS END DO CALL neighbor_list_3c_iterator_destroy(nl_3c_iter) +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(t3c,nimg,alloc_i,alloc_j,alloc_k) PRIVATE(i_img,j_img) DO j_img = 1, nimg DO i_img = 1, nimg IF (ALLOCATED(alloc_i(i_img, j_img)%array)) THEN - DO i = 1, SIZE(alloc_i(i_img, j_img)%array) - CALL dbt_get_stored_coordinates(t3c(i_img, j_img), & - [alloc_i(i_img, j_img)%array(i), alloc_j(i_img, j_img)%array(i), & - alloc_k(i_img, j_img)%array(i)], & - iproc) - CPASSERT(iproc .EQ. para_env%mepos) - END DO - CALL dbt_reserve_blocks(t3c(i_img, j_img), & alloc_i(i_img, j_img)%array, & alloc_j(i_img, j_img)%array, & @@ -824,6 +819,7 @@ CONTAINS END IF END DO END DO +!$OMP END PARALLEL CALL timestop(handle) @@ -2761,7 +2757,8 @@ CONTAINS REAL(dp), INTENT(INOUT) :: memory INTEGER :: buffer_left, buffer_size, buffer_start, & - i, memory_usage, nbits, nints + i, iblk, memory_usage, nbits, nblk, & + nints, offset, shared_offset INTEGER(int_8) :: estimate_to_store_int, & storage_counter_integrals INTEGER, DIMENSION(3) :: ind @@ -2769,6 +2766,7 @@ CONTAINS REAL(dp) :: spherical_estimate REAL(dp), ALLOCATABLE, DIMENSION(:, :, :), TARGET :: blk_data REAL(dp), DIMENSION(:), POINTER :: blk_data_1d + TYPE(dbt_iterator_type) :: iter TYPE(hfx_cache_type), DIMENSION(:), POINTER :: integral_caches TYPE(hfx_cache_type), POINTER :: maxval_cache TYPE(hfx_container_type), DIMENSION(:), POINTER :: integral_containers @@ -2790,7 +2788,21 @@ CONTAINS IF (ALLOCATED(blk_indices)) DEALLOCATE (blk_indices) ALLOCATE (blk_indices(dbt_get_num_blocks(tensor), 3)) - CALL dbt_reserved_block_indices(tensor, blk_indices) + shared_offset = 0 +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor,blk_indices,shared_offset) & +!$OMP PRIVATE(iter,ind,offset,nblk,iblk) + CALL dbt_iterator_start(iter, tensor) + nblk = dbt_iterator_num_blocks(iter) +!$OMP CRITICAL + offset = shared_offset + shared_offset = shared_offset + nblk +!$OMP END CRITICAL + DO iblk = 1, nblk + CALL dbt_iterator_next_block(iter, ind) + blk_indices(offset + iblk, :) = ind(:) + END DO + CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL ! Can not use the tensor iterator here because the order of the blocks is not guaranteed. DO i = 1, SIZE(blk_indices, 1) @@ -2897,7 +2909,10 @@ CONTAINS memory_usage, .FALSE.) END DO +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(tensor,blk_indices) CALL dbt_reserve_blocks(tensor, blk_indices) +!$OMP END PARALLEL ! Can not use the tensor iterator here because the order of the blocks is not guaranteed. DO i = 1, SIZE(blk_indices, 1) diff --git a/src/rpa_gw.F b/src/rpa_gw.F index 9a7c2b2e41..23ee3355ba 100644 --- a/src/rpa_gw.F +++ b/src/rpa_gw.F @@ -5508,6 +5508,10 @@ CONTAINS vec_Sigma_prv = 0.0_dp +!$OMP PARALLEL DEFAULT(NONE) REDUCTION(+:vec_Sigma_prv) & +!$OMP SHARED(t3c_1_redist,t3c_2,mo_bounds) & +!$OMP PRIVATE(iter,ind,bsize,boff,block_1,block_2,found) & +!$OMP PRIVATE(n_start_block,n_start,n_end_block,n_end,trace_shape) CALL dbt_iterator_start(iter, t3c_1_redist) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind, blk_size=bsize, blk_offset=boff) @@ -5541,6 +5545,7 @@ CONTAINS DEALLOCATE (block_1, block_2) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL dbt_destroy(t3c_1_redist) diff --git a/src/rpa_gw_ic.F b/src/rpa_gw_ic.F index 2f5fb6e97c..a1c5dde004 100644 --- a/src/rpa_gw_ic.F +++ b/src/rpa_gw_ic.F @@ -230,10 +230,19 @@ CONTAINS INTEGER, DIMENSION(3) :: boff, bsize, ind LOGICAL :: found REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: block_1, block_2 + REAL(KIND=dp), & + DIMENSION(mo_bounds(2)-mo_bounds(1)+1) :: Delta_Sigma_Neaton_prv TYPE(dbt_iterator_type) :: iter CALL timeset(routineN, handle) + CPASSERT(SIZE(Delta_Sigma_Neaton_prv) == SIZE(Delta_Sigma_Neaton)) + Delta_Sigma_Neaton_prv = 0.0_dp + +!$OMP PARALLEL DEFAULT(NONE) REDUCTION(+:Delta_Sigma_Neaton_prv) & +!$OMP SHARED(t3c_1,t3c_2,mo_bounds) & +!$OMP PRIVATE(iter,ind,bsize,boff,block_1,block_2,found) & +!$OMP PRIVATE(n_start_block,n_start,n_end_block,n_end) CALL dbt_iterator_start(iter, t3c_1) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind, blk_size=bsize, blk_offset=boff) @@ -259,15 +268,17 @@ CONTAINS n_end = boff(3) + bsize(3) - mo_bounds(1) END IF - Delta_Sigma_Neaton(n_start:n_end) = & - Delta_Sigma_Neaton(n_start:n_end) + & + Delta_Sigma_Neaton_prv(n_start:n_end) = & + Delta_Sigma_Neaton_prv(n_start:n_end) + & (/(DOT_PRODUCT(block_1(:, n, n), & block_2(:, n, n)), & n=n_start_block, n_end_block)/) DEALLOCATE (block_1, block_2) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL + Delta_Sigma_Neaton = Delta_Sigma_Neaton + Delta_Sigma_Neaton_prv CALL mp_sum(Delta_Sigma_Neaton, para_env%group) CALL timestop(handle) diff --git a/src/xas_tdp_atom.F b/src/xas_tdp_atom.F index 0f175d45ad..7144491dfd 100644 --- a/src/xas_tdp_atom.F +++ b/src/xas_tdp_atom.F @@ -997,6 +997,11 @@ CONTAINS !Do the actual contraction: coeff_y = sum_pq sum_x P_pq (phi_p phi_q xi_x) S_xy^-1 +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(pqX,rho_ao,ri_sinv,xas_atom_env,blk_size_ri,idx_to_nb,nspins,nnb,neighbors,iex,factor) & +!$OMP PRIVATE(iter,ind,t_block,tensor_found,iatom,jatom,katom,inb,prefac,ispin) & +!$OMP PRIVATE(pmat_block,pmat_found,pmat_blockt,pmat_foundt,work1,work2,jnb,ri_at) & +!$OMP PRIVATE(sinv_block,sinv_found,sinv_blockt,sinv_foundt) CALL dbt_iterator_start(iter, pqX) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind) @@ -1055,9 +1060,11 @@ CONTAINS work2(:) = work2(:) + factor*work1(i, 1)*sinv_blockt(:, i) END DO END IF - - xas_atom_env%ri_dcoeff(ri_at, ispin, iex)%array(:) = & - xas_atom_env%ri_dcoeff(ri_at, ispin, iex)%array(:) + work2(:) + DO i = 1, SIZE(work2) +!$OMP ATOMIC + xas_atom_env%ri_dcoeff(ri_at, ispin, iex)%array(i) = & + xas_atom_env%ri_dcoeff(ri_at, ispin, iex)%array(i) + work2(i) + END DO DEALLOCATE (work2) END DO !jnb @@ -1068,6 +1075,7 @@ CONTAINS DEALLOCATE (t_block) END DO !iter CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL !clean-up CALL dbcsr_release(ri_sinv) diff --git a/src/xas_tdp_correction.F b/src/xas_tdp_correction.F index 2dee44ecd7..722880ec09 100644 --- a/src/xas_tdp_correction.F +++ b/src/xas_tdp_correction.F @@ -985,6 +985,9 @@ CONTAINS CALL dbt_distribution_destroy(t_dist) !dist of 3c_ex and work match, can simply copy blocks over. Diagonal with factor 0.5 + +!$OMP PARALLEL DEFAULT(NONE) SHARED(xas_tdp_env,exat,work,orb_blk_size,nsgf_x) & +!$OMP PRIVATE(iter,ind,pblock,found) CALL dbt_iterator_start(iter, xas_tdp_env%ri_3c_ex) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind) @@ -999,6 +1002,7 @@ CONTAINS DEALLOCATE (pblock) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL dbt_finalize(work) !create (pq|X) based on work and copy over @@ -1331,6 +1335,9 @@ CONTAINS CALL dbt_distribution_destroy(t_dist) !copy block by block, dist match + +!$OMP PARALLEL DEFAULT(NONE) SHARED(t_work,t_pI_Y,nblk_ri,ri_blk_size,ao_blk_size) & +!$OMP PRIVATE(iter,ind,pblock,found,bo) CALL dbt_iterator_start(iter, t_work) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind) @@ -1346,6 +1353,7 @@ CONTAINS DEALLOCATE (pblock) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL dbt_finalize(t_pI_Y) !get optimal pgrid for (oI|Y) @@ -1430,6 +1438,10 @@ CONTAINS nocc = SIZE(occ_evals, 1) !Iterate over the tensors and sum. Both tensors have same dist + +!$OMP PARALLEL DEFAULT(NONE) REDUCTION(+:contrib,dev) & +!$OMP SHARED(ja_Ik_diff,occ_evals,virt_evals,omega,c_ss,nocc) & +!$OMP PRIVATE(iter,ind,boff,bsize,tensor_blk,found,idx1,idx2,idx3,j,A,k,denom,tmp) CALL dbt_iterator_start(iter, ja_Ik_diff) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind, blk_offset=boff, blk_size=bsize) @@ -1462,6 +1474,7 @@ CONTAINS DEALLOCATE (tensor_blk) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) @@ -1509,6 +1522,10 @@ CONTAINS ! of the 2 terms, but with a factor 2 !Iterate over the tensor and sum + +!$OMP PARALLEL DEFAULT(NONE) REDUCTION(+:contrib,dev) & +!$OMP SHARED(ja_Ik,j_evals,a_evals,k_evals,omega,c_os,a_offset) & +!$OMP PRIVATE(iter,ind,boff,bsize,ja_Ik_blk,found,idx1,idx2,idx3,j,A,k,denom,tmp) CALL dbt_iterator_start(iter, ja_Ik) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind, blk_offset=boff, blk_size=bsize) @@ -1542,6 +1559,7 @@ CONTAINS DEALLOCATE (ja_Ik_blk) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) @@ -1585,6 +1603,10 @@ CONTAINS nocc = SIZE(occ_evals, 1) !tensors have matching distributions, can do that safely + +!$OMP PARALLEL DEFAULT(NONE) REDUCTION(+:contrib,dev) & +!$OMP SHARED(aj_Ib_diff,occ_evals,virt_evals,omega,c_ss,nocc) & +!$OMP PRIVATE(iter,ind,boff,bsize,tensor_blk,found,idx1,idx2,idx3,j,A,b,denom,tmp) CALL dbt_iterator_start(iter, aj_Ib_diff) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind, blk_offset=boff, blk_size=bsize) @@ -1617,6 +1639,7 @@ CONTAINS DEALLOCATE (tensor_blk) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) @@ -1664,6 +1687,9 @@ CONTAINS ! MOs on one center have one spin and the 2 MOs on the other center have another spin ! In the end, the sum is such that can take one of those with a factor 2 +!$OMP PARALLEL DEFAULT(NONE) REDUCTION(+:contrib,dev) & +!$OMP SHARED(aj_Ib,a_evals,j_evals,b_evals,omega,c_os,a_offset,b_offset) & +!$OMP PRIVATE(iter,ind,boff,bsize,aj_Ib_blk,found,idx1,idx2,idx3,j,A,b,denom,tmp) CALL dbt_iterator_start(iter, aj_Ib) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind, blk_offset=boff, blk_size=bsize) @@ -1696,6 +1722,7 @@ CONTAINS DEALLOCATE (aj_Ib_blk) END DO CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle) diff --git a/src/xas_tdp_integrals.F b/src/xas_tdp_integrals.F index a29c227a7f..5518e893b6 100644 --- a/src/xas_tdp_integrals.F +++ b/src/xas_tdp_integrals.F @@ -236,7 +236,10 @@ CONTAINS CALL neighbor_list_iterator_release(ab_iter) CALL neighbor_list_iterator_release(ac_iter) +!TODO: Parallelize this properly. +!$OMP PARALLEL DEFAULT(NONE) SHARED(pq_X, idx1, idx2, idx3) CALL dbt_reserve_blocks(pq_X, idx1, idx2, idx3) +!$OMP END PARALLEL CALL dbt_finalize(pq_X) !clean-up diff --git a/src/xas_tdp_kernel.F b/src/xas_tdp_kernel.F index f6de17419e..e7482de76b 100644 --- a/src/xas_tdp_kernel.F +++ b/src/xas_tdp_kernel.F @@ -1105,8 +1105,7 @@ CONTAINS CHARACTER(len=*), PARAMETER :: routineN = 'contract3_RI_to_doMOs' - INTEGER :: handle, iatom, ik, ind(3), jatom, katom, & - nk, s1, s2 + INTEGER :: handle, i, iatom, ind(3), j, jatom, katom LOGICAL :: found, t_found REAL(dp) :: prefac REAL(dp), ALLOCATABLE, DIMENSION(:, :, :) :: iabc @@ -1118,8 +1117,9 @@ CONTAINS CALL timeset(routineN, handle) - nk = SIZE(vec) - +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(ab_Q,vec,mat_abIJ,atom_k) & +!$OMP PRIVATE(iter,ind,iatom,jatom,katom,prefac,iabc,t_found,found,pblock,i,j) CALL dbt_iterator_start(iter, ab_Q) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind) @@ -1138,14 +1138,17 @@ CONTAINS CALL dbcsr_get_block_p(mat_abIJ, iatom, jatom, pblock, found) IF ((.NOT. found) .OR. (.NOT. t_found)) CYCLE - s1 = SIZE(pblock, 1); s2 = SIZE(pblock, 2) - DO ik = 1, nk - CALL daxpy(s1*s2, prefac*vec(ik), iabc(:, :, ik), 1, pblock(:, :), 1) + DO i = 1, SIZE(pblock, 1) + DO j = 1, SIZE(pblock, 2) +!$OMP ATOMIC + pblock(i, j) = pblock(i, j) + prefac*DOT_PRODUCT(vec(:), iabc(i, j, :)) + END DO END DO DEALLOCATE (iabc) END DO !iter CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL !matrix only half filled => need to add its transpose CALL dbcsr_create(work, template=mat_abIJ) @@ -1183,7 +1186,7 @@ CONTAINS CHARACTER(LEN=*), PARAMETER :: routineN = 'contract2_AO_to_doMO_low' INTEGER :: handle, i, iatom, ind(3), j, jatom, & - katom, n, s1, s2 + katom, s1, s2 INTEGER, DIMENSION(:), POINTER :: atom_blk_size LOGICAL :: found, t_found REAL(dp), ALLOCATABLE, DIMENSION(:, :, :) :: iabc @@ -1196,6 +1199,9 @@ CONTAINS CALL dbcsr_get_info(mat_aIb, row_blk_size=atom_blk_size) +!$OMP PARALLEL DEFAULT(NONE) & +!$OMP SHARED(ab_Q,vec,mat_aIb,mat_bIa,atom_k,atom_blk_size) & +!$OMP PRIVATE(iter,ind,iatom,jatom,katom,iabc,t_found,found,s1,s2,j,i,pblock) CALL dbt_iterator_start(iter, ab_Q) DO WHILE (dbt_iterator_blocks_left(iter)) CALL dbt_iterator_next_block(iter, ind) @@ -1211,16 +1217,16 @@ CONTAINS ! Deal with mat_aIb IF (jatom == atom_k) THEN - n = SIZE(vec) s1 = atom_blk_size(iatom) s2 = SIZE(iabc, 3) CALL dbcsr_get_block_p(matrix=mat_aIb, row=iatom, col=jatom, BLOCK=pblock, found=found) IF (found) THEN - DO j = 1, s2 - DO i = 1, n - pblock(:, j) = pblock(:, j) + vec(i)*iabc(:, i, j) + DO i = 1, s1 + DO j = 1, s2 +!$OMP ATOMIC + pblock(i, j) = pblock(i, j) + DOT_PRODUCT(vec, iabc(i, :, j)) END DO END DO END IF @@ -1229,16 +1235,17 @@ CONTAINS ! Deal with mat_bIa, keep block diagonal empty IF (iatom == jatom) CYCLE IF (iatom == atom_k) THEN - - n = SIZE(vec) s1 = SIZE(iabc, 3) s2 = atom_blk_size(jatom) CALL dbcsr_get_block_p(matrix=mat_bIa, row=iatom, col=jatom, BLOCK=pblock, found=found) IF (found) THEN - DO i = 1, n - CALL daxpy(s1*s2, vec(i), TRANSPOSE(iabc(i, :, :)), 1, pblock(:, :), 1) + DO i = 1, s1 + DO j = 1, s2 +!$OMP ATOMIC + pblock(i, j) = pblock(i, j) + DOT_PRODUCT(vec, iabc(:, j, i)) + END DO END DO END IF END IF !iatom== atom_k @@ -1246,6 +1253,7 @@ CONTAINS DEALLOCATE (iabc) END DO !iter CALL dbt_iterator_stop(iter) +!$OMP END PARALLEL CALL timestop(handle)