SUBROUTINE test_complete_redistribute() INTEGER, POINTER, CONTIGUOUS :: rbs(:), cbs(:), rd(:), cd(:) INTEGER, POINTER, CONTIGUOUS :: new_rbs(:), new_cbs(:), new_rd(:), new_cd(:) INTEGER :: itype, icase, r, c, i, j, k, n, p, ro, co, offset INTEGER :: rows(16), cols(16) INTEGER, ALLOCATABLE :: index_before(:) TYPE(dbcsr_distribution_obj) :: dist, new_dist TYPE(dbcsr_type) :: source, target, converted TYPE(dbcsr_data_obj) :: values, actual COMPLEX(dp) :: reference(10, 10), expected(10, 10) LOGICAL :: keep ALLOCATE (rbs(4), cbs(4), rd(4), cd(4), new_rbs(3), new_cbs(3), new_rd(3), new_cd(3)) rbs = [2, 3, 1, 4]; cbs = [3, 1, 2, 4] new_rbs = [3, 4, 3]; new_cbs = [2, 5, 3] rd = [(MOD(i - 1, npdims(1)), i=1, 4)] cd = [(MOD(i - 1, npdims(2)), i=1, 4)] new_rd = [(MOD(i, npdims(1)), i=1, 3)] new_cd = [(MOD(i, npdims(2)), i=1, 3)] CALL dbcsr_distribution_new(dist, mp_env, rd, cd) CALL dbcsr_distribution_new(new_dist, mp_env, new_rd, new_cd) CALL dbcsr_data_init(values) CALL dbcsr_data_init(actual) CALL dbcsr_data_new(values, dbcsr_type_complex_8, data_size=100) CALL dbcsr_data_new(actual, dbcsr_type_complex_8_2d, 10, 10) DO itype = 1, SIZE(data_types) reference = CMPLX(0.0_dp, 0.0_dp, dp) CALL dbcsr_create(source, "reblock source", dist, dbcsr_type_no_symmetry, & row_blk_size=rbs, col_blk_size=cbs, data_type=data_types(itype)) n = 0 DO r = 1, SIZE(rbs) ro = SUM(rbs(:r - 1)) DO c = 1, SIZE(cbs) IF (MOD(r + c, 3) == 0) CYCLE co = SUM(cbs(:c - 1)) DO j = 1, cbs(c) DO i = 1, rbs(r) reference(ro + i, co + j) = CMPLX(100*(ro + i) + co + j, 0, dp) END DO END DO CALL mp_cart_rank(group, [rd(r), cd(c)], p) IF (p /= mynode) CYCLE n = n + 1; rows(n) = r; cols(n) = c END DO END DO IF (itype > 2) reference = reference*CMPLX(1.0_dp, -1.0_dp, dp) CALL dbcsr_reserve_blocks(source, rows(:n), cols(:n)) DO k = 1, n r = rows(k); c = cols(k) ro = SUM(rbs(:r - 1)); co = SUM(cbs(:c - 1)) offset = source%blk_p(k) values%d%c_dp(offset:offset + rbs(r)*cbs(c) - 1) = & RESHAPE(reference(ro + 1:ro + rbs(r), co + 1:co + cbs(c)), [rbs(r)*cbs(c)]) END DO CALL dbcsr_data_set_size_referenced(values, source%nze) CALL dbcsr_data_convert(source%data_area, values) DO icase = 1, 4 ! Empty target, populated target, retained sparsity, and empty source. keep = icase == 3 CALL dbcsr_create(target, "reblock target", new_dist, dbcsr_type_no_symmetry, & row_blk_size=new_rbs, col_blk_size=new_cbs, data_type=data_types(itype)) IF (icase > 1) THEN n = 0 DO r = 1, SIZE(new_rbs) DO c = 1, SIZE(new_cbs) IF (MOD(r + c, 2) == 0) CYCLE CALL mp_cart_rank(group, [new_rd(r), new_cd(c)], p) IF (p /= mynode) CYCLE n = n + 1; rows(n) = r; cols(n) = c END DO END DO CALL dbcsr_reserve_blocks(target, rows(:n), cols(:n)) values%d%c_dp(:) = CMPLX(-7.0_dp, 0.0_dp, dp) CALL dbcsr_data_set_size_referenced(values, target%nze) CALL dbcsr_data_convert(target%data_area, values) END IF IF (icase == 4) THEN CALL dbcsr_clear(source) reference = CMPLX(0.0_dp, 0.0_dp, dp) END IF ALLOCATE (index_before, SOURCE=target%index) CALL dbcsr_complete_redistribute(source, target, keep_sparsity=keep) CALL dbcsr_verify_matrix(target) IF (.NOT. ASSOCIATED(target%dist%d, new_dist%d)) & DBCSR_ABORT("Reblocking changed distribution") IF (keep .OR. icase == 4) THEN IF (ANY(target%index /= index_before)) & DBCSR_ABORT("Reblocking rebuilt an unchanged index") END IF DEALLOCATE (index_before) IF (icase == 1) THEN ! CSR conversion requires newly built matrices to have packed row-wise data. offset = 1 DO r = 1, SIZE(new_rbs) DO k = target%row_p(r) + 1, target%row_p(r + 1) IF (target%blk_p(k) /= offset) & DBCSR_ABORT("Reblocking changed data order") offset = offset + new_rbs(r)*new_cbs(target%col_i(k)) END DO END DO END IF CALL dbcsr_create(converted, template=target, data_type=dbcsr_type_complex_8) CALL dbcsr_copy(converted, target) CALL dbcsr_to_dense_local(converted, actual) CALL dbcsr_release(converted) expected = reference ! Compare locally, including zeros outside this rank's ownership. DO r = 1, SIZE(new_rbs) ro = SUM(new_rbs(:r - 1)) DO c = 1, SIZE(new_cbs) co = SUM(new_cbs(:c - 1)) CALL mp_cart_rank(group, [new_rd(r), new_cd(c)], p) IF (p /= mynode .OR. (keep .AND. MOD(r + c, 2) == 0)) & expected(ro + 1:ro + new_rbs(r), co + 1:co + new_cbs(c)) = 0.0_dp END DO END DO IF (ANY(ABS(actual%d%c2_dp - expected) > 1.0E-6_dp)) & DBCSR_ABORT("Reblocking changed matrix values") CALL dbcsr_release(target) END DO CALL dbcsr_release(source) END DO CALL dbcsr_data_release(values) CALL dbcsr_data_release(actual) CALL dbcsr_distribution_release(dist) CALL dbcsr_distribution_release(new_dist) DEALLOCATE (rbs, cbs, rd, cd, new_rbs, new_cbs, new_rd, new_cd) END SUBROUTINE test_complete_redistribute