diff --git a/src/methods/GW/thc_gw.icc b/src/methods/GW/thc_gw.icc index f0f89daf..4f7081e2 100644 --- a/src/methods/GW/thc_gw.icc +++ b/src/methods/GW/thc_gw.icc @@ -260,6 +260,10 @@ namespace methods { int np_Q_2 = (t_intra_comm.size()/qpools_2) / np_P_2; auto dSigma_skPQ_2 = make_distributed_array( t_intra_comm, {1, qpools_2, np_P_2, np_Q_2}, {ns, nkpts_ibz, NP, NQ}); + // The two layouts are fixed for every tau point and symmetry operation. + // Build their collective metadata once instead of repeating the layout + // all-gather in every bounded redistribution. + auto Sigma_redistribution_plan = math::nda::make_redistribution_plan(dSigma_skPQ, dSigma_skPQ_2); app_log(2, " Self-energy from THC to primary basis:"); app_log(2, " - processor grid for Sigma: (t, q, P, Q) = ({}, {}, {}, {})\n", tpools, qpools_2, np_P_2, np_Q_2); @@ -291,7 +295,8 @@ namespace methods { sSigma_tskab.set_zero(); _Timer.start("SIGMA_AUX_TO_PRIM"); // redistribute Sigma to improve locality of (P,Q) indices - math::nda::redistribute(dSigma_skPQ, dSigma_skPQ_2); + math::nda::redistribute( + dSigma_skPQ, dSigma_skPQ_2, Sigma_redistribution_plan); thc_solver_comm::aux_to_primary(0,0,tpool_id, ComplexType(1.0),dSigma_skPQ_2, sSigma_tskab, thc, MF->ks_to_k(isym)); _Timer.stop("SIGMA_AUX_TO_PRIM"); @@ -491,32 +496,79 @@ namespace methods { T_skia.win().fence(); T_skia.all_reduce(); - nda::array buffer_ib(nbnd, nbnd); - - auto sDelta_tskij = make_shared_array(*mpi, sSigma_tskij.shape()); - auto Delta_loc = sDelta_tskij.local(); - sDelta_tskij.win().fence(); - for (size_t tsk = rank; tsk < nts * ns * nkpts; tsk += size) { - // tsk = it * ns*nkpts + is * nkpts + ik - size_t it = tsk / (ns * nkpts); - size_t is = (tsk / nkpts) % ns; - size_t ik = tsk % nkpts; - - size_t it_pos = (it < nt_half)? it : nts - it - 1; - RealType factor = -1.0 * MF->madelung() * eps_inv_head(it_pos).real(); - - auto G_ab = G_tskij(it, is, ik, nda::ellipsis{}); - auto Delta_ij = Delta_loc(it, is, ik, nda::ellipsis{}); - auto T_ia = T_skia_loc(is, ik, nda::ellipsis{}); - nda::blas::gemm(T_ia, G_ab, buffer_ib); - nda::blas::gemm(ComplexType(factor), buffer_ib, nda::dagger(T_skia_loc(is, ik, nda::ellipsis{})), - ComplexType(0.0), Delta_ij); + utils::check(nbnd == nbnd2, "Sigma_div_correction expects square band matrices: {} != {}", nbnd, nbnd2); + + // Reduce the correction in bounded tiles and apply each reduced tile + // directly to Sigma. The old full-size sDelta_tskij duplicated the + // already resident G and Sigma tensors (6.6 GiB in the production + // case). Keep a full band matrix whenever it fits so the usual path + // still uses only two GEMMs per (tau, spin, k) point; band tiling is a + // safety fallback for exceptionally large band spaces. + constexpr size_t delta_tile_bytes = size_t{32} * 1024 * 1024; + size_t max_tile_elements = std::max(1, delta_tile_bytes / sizeof(ComplexType)); + long band_tile = std::min(nbnd, static_cast(std::sqrt(static_cast(max_tile_elements)))); + band_tile = std::max(1, band_tile); + size_t band_tile_elements = static_cast(band_tile) * static_cast(band_tile); + long tsk_tile = static_cast(std::max(1, max_tile_elements / band_tile_elements)); + size_t ntsk = static_cast(nts) * static_cast(ns) * static_cast(nkpts); + + using Array_view_3D_t = nda::array_view; + auto sDelta_tile = make_shared_array( + *mpi, {tsk_tile, band_tile, band_tile}); + nda::array buffer_ib(band_tile, nbnd); + decltype(nda::range::all) all; + + for (size_t tsk0 = 0; tsk0 < ntsk; tsk0 += static_cast(tsk_tile)) { + size_t tsk1 = std::min(ntsk, tsk0 + static_cast(tsk_tile)); + long tsk_count = static_cast(tsk1 - tsk0); + for (long i0 = 0; i0 < nbnd; i0 += band_tile) { + long i1 = std::min(nbnd, i0 + band_tile); + long ni = i1 - i0; + for (long j0 = 0; j0 < nbnd; j0 += band_tile) { + long j1 = std::min(nbnd, j0 + band_tile); + long nj = j1 - j0; + + sDelta_tile.set_zero(); + auto Delta_tile_loc = sDelta_tile.local(); + size_t rank_offset = (static_cast(rank) + static_cast(size) - + tsk0 % static_cast(size)) % static_cast(size); + for (size_t tsk = tsk0 + rank_offset; tsk < tsk1; tsk += static_cast(size)) { + + // tsk = it * ns*nkpts + is * nkpts + ik + size_t it = tsk / (static_cast(ns) * static_cast(nkpts)); + size_t is = (tsk / static_cast(nkpts)) % static_cast(ns); + size_t ik = tsk % static_cast(nkpts); + size_t it_pos = (it < static_cast(nt_half))? it : static_cast(nts) - it - 1; + RealType factor = -1.0 * MF->madelung() * eps_inv_head(it_pos).real(); + + auto G_ab = G_tskij(it, is, ik, nda::ellipsis{}); + auto T_ia = T_skia_loc(is, ik, nda::range(i0, i1), all); + auto T_ja = T_skia_loc(is, ik, nda::range(j0, j1), all); + auto buffer_i = buffer_ib(nda::range(0, ni), all); + auto Delta_ij = Delta_tile_loc( + static_cast(tsk - tsk0), nda::range(0, ni), nda::range(0, nj)); + nda::blas::gemm(T_ia, G_ab, buffer_i); + nda::blas::gemm(ComplexType(factor), buffer_i, nda::dagger(T_ja), + ComplexType(0.0), Delta_ij); + } + sDelta_tile.all_reduce(delta_tile_bytes); + + if (sSigma_tskij.node_comm()->root()) { + auto Sigma_loc = sSigma_tskij.local(); + auto Delta_tile_loc_reduced = sDelta_tile.local(); + for (long itsk = 0; itsk < tsk_count; ++itsk) { + size_t tsk = tsk0 + static_cast(itsk); + size_t it = tsk / (static_cast(ns) * static_cast(nkpts)); + size_t is = (tsk / static_cast(nkpts)) % static_cast(ns); + size_t ik = tsk % static_cast(nkpts); + Sigma_loc(it, is, ik, nda::range(i0, i1), nda::range(j0, j1)) += + Delta_tile_loc_reduced(itsk, nda::range(0, ni), nda::range(0, nj)); + } + } + } + } } - sDelta_tskij.win().fence(); - sDelta_tskij.all_reduce(); - - if (sSigma_tskij.node_comm()->root()) - sSigma_tskij.local() += sDelta_tskij.local(); + sSigma_tskij.node_sync(); sSigma_tskij.communicator()->barrier(); } else { utils::check(false, "Unsupported divergence treatment: {}", div_treatment); diff --git a/src/methods/SCF/tests/test_scf_common.cpp b/src/methods/SCF/tests/test_scf_common.cpp index b01afe16..f607a280 100644 --- a/src/methods/SCF/tests/test_scf_common.cpp +++ b/src/methods/SCF/tests/test_scf_common.cpp @@ -43,6 +43,7 @@ #include "methods/ERI/eri_utils.hpp" #include "hamiltonian/pseudo/pseudopot.h" #include "numerics/iter_scf/iter_scf_utils.hpp" +#include "numerics/iter_scf/diis/vspace_fock_sigma.hpp" namespace bdft_tests { @@ -186,6 +187,69 @@ namespace bdft_tests { } } + TEST_CASE("dyson_diis_streamed_commutator", "[methods_scf][diis]") { + decltype(nda::range::all) all; + imag_axes_ft::IAFT ft(50.0, 1.5, imag_axes_ft::ir_basis, "medium"); + const long nt = ft.nt_f(); + const long nw = ft.nw_f(); + constexpr long ns = 1; + constexpr long nk = 2; + constexpr long nao = 3; + constexpr double mu = 0.17; + + nda::array G_t(nt, ns, nk, nao, nao); + nda::array Sigma_t(nt, ns, nk, nao, nao); + nda::array F(ns, nk, nao, nao); + nda::array S(ns, nk, nao, nao); + nda::array H0(ns, nk, nao, nao); + + for (long it = 0; it < nt; ++it) + for (long k = 0; k < nk; ++k) + for (long i = 0; i < nao; ++i) + for (long j = 0; j < nao; ++j) { + const double x = 1.0 + it + 3*k + 5*i + 7*j; + G_t(it, 0, k, i, j) = ComplexType(0.002*x, -0.001*(x + i - j)); + Sigma_t(it, 0, k, i, j) = ComplexType(-0.0007*(x + j), 0.0003*(x + i)); + } + for (long k = 0; k < nk; ++k) + for (long i = 0; i < nao; ++i) + for (long j = 0; j < nao; ++j) { + const double x = 1.0 + 2*k + 3*i + 5*j; + F(0, k, i, j) = ComplexType(0.01*x, 0.002*(i - j)); + H0(0, k, i, j) = ComplexType(-0.02*x, 0.001*(i + j)); + S(0, k, i, j) = ComplexType((i == j ? 1.0 : 0.01*x), 0.0005*(i - j)); + } + + iter_scf::FockSigma fs(F, Sigma_t, mu); + nda::array C_stream; + iter_scf::commutator_t(C_stream, &ft, G_t, fs, mu, S, H0); + + // Dense formulation retained in the test as an independent numerical + // oracle for the bounded-memory frequency-streaming implementation. + nda::array G_w(nw, ns, nk, nao, nao); + nda::array Sigma_w(nw, ns, nk, nao, nao); + nda::array C_w(nw, ns, nk, nao, nao); + nda::array C_reference(nt, ns, nk, nao, nao); + ft.tau_to_w(G_t, G_w, imag_axes_ft::fermion); + ft.tau_to_w(Sigma_t, Sigma_w, imag_axes_ft::fermion); + for (long iw = 0; iw < nw; ++iw) + for (long k = 0; k < nk; ++k) { + const auto omega_mu = ft.omega(ft.wn_mesh()(iw)) + mu; + auto G_wsk = G_w(iw, 0, k, all, all); + auto Sigma_wsk = Sigma_w(iw, 0, k, all, all); + auto M = nda::make_regular(omega_mu*S(0, k, all, all) - H0(0, k, all, all) + - F(0, k, all, all) - Sigma_wsk); + nda::array left(nao, nao); + nda::array right(nao, nao); + nda::blas::gemm(G_wsk, M, left); + nda::blas::gemm(M, G_wsk, right); + C_w(iw, 0, k, all, all) = left - right; + } + ft.w_to_tau(C_w, C_reference, imag_axes_ft::fermion); + + ARRAY_EQUAL(C_stream, C_reference, 5e-11); + } + TEST_CASE("dyson_scf_gw_diis_vs_damping", "[methods_scf]") { auto& mpi_context = utils::make_unit_test_mpi_context(); diff --git a/src/methods/embedding/downfold_1e.cpp b/src/methods/embedding/downfold_1e.cpp index c1e13ef4..295385c3 100644 --- a/src/methods/embedding/downfold_1e.cpp +++ b/src/methods/embedding/downfold_1e.cpp @@ -730,9 +730,7 @@ namespace methods { mpi->comm.barrier(); // Vhf_skij = Vhf_skij + Vcorr_skij - Vcorr_dc_skij - proj.upfold(sVcorr_skij, Vcorr_dc_sIab); - if (sVhf_skij.node_comm()->root()) sVhf_skij.local() -= sVcorr_skij.local(); - mpi->comm.barrier(); + proj.upfold_add(sVhf_skij, Vcorr_dc_sIab, ComplexType(-1.0)); // If weiss_f_iter==-1, we are in the 1st iteration of embedding and there is // no impurity self-energy. diff --git a/src/methods/embedding/embed_eri_t.cpp b/src/methods/embedding/embed_eri_t.cpp index f648be61..f3eb7ba1 100644 --- a/src/methods/embedding/embed_eri_t.cpp +++ b/src/methods/embedding/embed_eri_t.cpp @@ -21,6 +21,8 @@ #include "mpi3/communicator.hpp" +#include +#include #include #include "nda/nda.hpp" #include "nda/blas.hpp" @@ -37,6 +39,47 @@ #include "cholesky.hpp" namespace methods { + namespace { + // Keep MPI collectives below both the legacy int-count limit and a modest + // byte count. Some MPI implementations still multiply count by the + // datatype extent in an int internally even when the C++ wrapper accepts + // a wider Size type. + template + constexpr size_t mpi_chunk_elements() { + constexpr size_t max_chunk_bytes = size_t{32} * 1024 * 1024; + return std::max( + 1, std::min(static_cast(std::numeric_limits::max()), + max_chunk_bytes / sizeof(T))); + } + + template + void all_reduce_in_place_chunked(communicator_t &comm, T *data, size_t count, Op op) { + auto const chunk_size = mpi_chunk_elements(); + for (size_t offset = 0; offset < count; offset += chunk_size) { + auto n = std::min(chunk_size, count - offset); + comm.all_reduce_in_place_n(data + offset, static_cast(n), op); + } + } + + template + void reduce_in_place_chunked(communicator_t &comm, T *data, size_t count, Op op, int root = 0) { + auto const chunk_size = mpi_chunk_elements(); + for (size_t offset = 0; offset < count; offset += chunk_size) { + auto n = std::min(chunk_size, count - offset); + comm.reduce_in_place_n(data + offset, static_cast(n), op, root); + } + } + + template + void broadcast_chunked(communicator_t &comm, T *data, size_t count, int root = 0) { + auto const chunk_size = mpi_chunk_elements(); + for (size_t offset = 0; offset < count; offset += chunk_size) { + auto n = std::min(chunk_size, count - offset); + comm.broadcast_n(data + offset, static_cast(n), root); + } + } + } // namespace + template auto embed_eri_t::compute_collation_impurity_basis( thc_t &thc, const projector_boson_t &proj_boson, nda::range u_rng) @@ -1457,7 +1500,7 @@ namespace methods { // Bare interactions app_log(1, "Downfolding the bare Coulomb interactions...\n"); - auto V_qabcd = downfold_Vq(thc, B_qIPab); + auto V_qabcd = downfold_Vq_root(thc, B_qIPab); // Dynamical screened interactions app_log(1, "Downfolding the dynamic screened interactions with screening type = {}.\n", screen_type); @@ -1472,17 +1515,29 @@ namespace methods { // FIXME We assume particle-hole symmetry. This may not always be the case! scr_coulomb.dyson_W_in_place(dW_wqPQ, thc); auto [eps_inv_head_wq, eps_inv_head_w] = solvers::div_utils::eps_inv_head_w(dW_wqPQ, thc, *_MF, _div_treatment); - auto W_wqabcd = downfold_Wq(thc, dW_wqPQ, B_qIPab); + auto W_wqabcd = downfold_Wq_root(thc, dW_wqPQ, B_qIPab); - // projection for local quantities + // Projection for local quantities. The q-resolved tensors are assembled + // only on rank 0; broadcast only their much smaller q averages, which are + // needed by the common downstream path on every rank. nda::array V_abcd(nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); - nda::array W_wabcd(W_wqabcd.shape(0), nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); - for (size_t iq_full=0; iq_full < _MF->nqpts(); ++iq_full) { - V_abcd += V_qabcd(iq_full, nda::ellipsis{}); - W_wabcd += W_wqabcd(nda::range::all, iq_full, nda::ellipsis{}); + nda::array W_wabcd(eps_inv_head_w.shape(0), nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + V_abcd() = ComplexType(0.0); + W_wabcd() = ComplexType(0.0); + if (mpi->comm.root()) { + utils::check(V_qabcd.shape()[0] == _MF->nqpts() and + W_wqabcd.shape()[0] == eps_inv_head_w.shape(0) and + W_wqabcd.shape()[1] == _MF->nqpts(), + "embed_eri_t::rpa_q_eri_impl: Inconsistent root-owned q tensor shapes."); + for (size_t iq_full=0; iq_full < _MF->nqpts(); ++iq_full) { + V_abcd += V_qabcd(iq_full, nda::ellipsis{}); + W_wabcd += W_wqabcd(nda::range::all, iq_full, nda::ellipsis{}); + } } V_abcd() /= _MF->nqpts(); W_wabcd() /= _MF->nqpts(); + broadcast_chunked(mpi->comm, V_abcd.data(), V_abcd.size()); + broadcast_chunked(mpi->comm, W_wabcd.data(), W_wabcd.size()); // finite-size correction to V_abcd V_div_correction(V_abcd, B_qIPab, thc); @@ -1523,17 +1578,38 @@ namespace methods { auto dV_qPQ = thc.dZ({1, 1, mpi->comm.size()}); // dV_qPQ = dV_qPQ + dW_wqPQ[0,...] { - // lazy for now, add redistribute routine that can operate on a submatrix, or eval_W_selected_frequencies - math::nda::redistribute_in_place(dW_wqPQ,{1,1,1,mpi->comm.size()}, - {dW_wqPQ.global_shape()[0],dW_wqPQ.global_shape()[1],dW_wqPQ.global_shape()[2],dV_qPQ.block_size()[2]}); - utils::check(dV_qPQ.local_shape()[0] == dW_wqPQ.local_shape()[1] and - dV_qPQ.local_shape()[1] == dW_wqPQ.local_shape()[2] and - dV_qPQ.local_shape()[2] == dW_wqPQ.local_shape()[3] and - dV_qPQ.origin()[0] == dW_wqPQ.origin()[1] and - dV_qPQ.origin()[1] == dW_wqPQ.origin()[2] and - dV_qPQ.origin()[2] == dW_wqPQ.origin()[3], - "Error in rpa_chol_eri_impl: Inconsistent data distribution, should not happen. \n"); - dV_qPQ.local() += dW_wqPQ.local()(0,nda::ellipsis{}); + // Materialize only the local part of w=0. Ranks in other frequency + // pools contribute an empty block. The streaming redistribution uses + // the source/destination block metadata to transfer overlaps directly + // to dV_qPQ, so the complete (w,q,P,Q) tensor is never redistributed or + // duplicated merely to select one frequency. + using local_Array_3D_t = memory::array; + auto const w0 = 0L; + auto const w_rng = dW_wqPQ.local_range(0); + bool const owns_w0 = w_rng.first() <= w0 and w0 < w_rng.last(); + auto const dW_shape = dW_wqPQ.local_shape(); + std::array w0_local_shape = owns_w0 ? + std::array{dW_shape[1], dW_shape[2], dW_shape[3]} : + std::array{0, 0, 0}; + local_Array_3D_t W0_local(w0_local_shape); + if (owns_w0) + W0_local = dW_wqPQ.local()(w0 - dW_wqPQ.origin()[0], nda::ellipsis{}); + + auto const dW_global_shape = dW_wqPQ.global_shape(); + auto const dW_origin = dW_wqPQ.origin(); + memory::irregular_block_darray_t dW0_qPQ( + std::addressof(mpi->comm), + {dW_global_shape[1], dW_global_shape[2], dW_global_shape[3]}, + {dW_origin[1], dW_origin[2], dW_origin[3]}, + std::move(W0_local)); + + utils::check(dW0_qPQ.global_shape() == dV_qPQ.global_shape(), + "Error in rpa_chol_eri_impl: w=0 and V global shapes differ."); + auto transfer_plan = math::nda::make_redistribution_plan(dW0_qPQ, dV_qPQ); + transfer_plan.validate_source_coverage(); + transfer_plan.validate_destination_coverage(); + math::nda::redistribute_streaming(dW0_qPQ, dV_qPQ, transfer_plan, + ComplexType(1.0), ComplexType(1.0)); } dW_wqPQ.reset(); app_log(1, "Treatment of long-wavelength divergence in V (bare): {}", _bare_div_treatment); @@ -1642,7 +1718,7 @@ namespace methods { nda::blas::gemm(ComplexType(1.0), B_cd_P_conj, T_P_ab, ComplexType(1.0), V_cd_ab); } } - comm.all_reduce_in_place_n(V_cdab.data(), V_cdab.size(), std::plus<>{}); + all_reduce_in_place_chunked(comm, V_cdab.data(), V_cdab.size(), std::plus<>{}); V_cdab() /= (nqpts); // finite-size correction to V_cdab @@ -1653,6 +1729,19 @@ namespace methods { template B_t> auto embed_eri_t::downfold_Vq(thc_t &thc, const B_t &B_qIPab) + -> nda::array { + auto V_qcdab = downfold_Vq_root(thc, B_qIPab); + auto mpi = thc.mpi(); + auto &comm = mpi->comm; + auto [nqpts, nImps, NP, nImpOrbs, nImpOrbs2] = B_qIPab.shape(); + if (not comm.root()) + V_qcdab.resize(nqpts, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + broadcast_chunked(comm, V_qcdab.data(), V_qcdab.size()); + return V_qcdab; + } + + template B_t> + auto embed_eri_t::downfold_Vq_root(thc_t &thc, const B_t &B_qIPab) -> nda::array { // B_qIPab lives in the full MP mesh @@ -1674,7 +1763,14 @@ namespace methods { auto Q_rng = dV_qPQ.local_range(2); auto [q_origin, P_origin, Q_origin] = dV_qPQ.origin(); - nda::array V_qcdab(nqpts, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + // The complete q-resolved result is an output artifact consumed on rank 0 + // by the serial HDF5 writer. Other ranks remain empty. + nda::array V_qcdab; + if (comm.root()) + V_qcdab.resize(nqpts, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + size_t orbital_slab_elements = static_cast(nImpOrbs) * nImpOrbs * nImpOrbs * nImpOrbs; + size_t q_batch_size = std::max(1, mpi_chunk_elements() / + std::max(1, orbital_slab_elements)); // V_qcdab = conj(B_qPdc) * [ V_qPQ ] * B_qQab nda::array V_PQ_loc(NP_loc, NQ_loc); @@ -1682,33 +1778,38 @@ namespace methods { nda::array B_cdP_conj(nImpOrbs, nImpOrbs, NP_loc); auto B_cd_P_conj = nda::reshape(B_cdP_conj, shape_t<2>{nImpOrbs*nImpOrbs, NP_loc}); - // Bare interactions - for (size_t iq_loc = 0; iq_loc < nq_loc; ++iq_loc) { - // iq lives in IBZ - size_t iq = q_origin + iq_loc; - - // loop over all symmetry-related q-points for iq - for (size_t iq_full=0; iq_fullqp_to_ibz(iq_full)!=iq) continue; - auto B_Q_ab = nda::reshape(B_qIPab(iq_full, 0, Q_rng, nda::ellipsis{}), - shape_t<2>{NQ_loc, nImpOrbs*nImpOrbs}); - if (_MF->qp_trev(iq_full)) { - V_PQ_loc = nda::conj( dV_qPQ.local()(iq_loc, nda::ellipsis{}) ); - nda::blas::gemm(V_PQ_loc, B_Q_ab, T_P_ab); - } else { - V_PQ_loc = dV_qPQ.local()(iq_loc, nda::ellipsis{}); - nda::blas::gemm(V_PQ_loc, B_Q_ab, T_P_ab); - } + // Bare interactions, reduced in bounded batches of full-BZ q slabs. + for (size_t q_begin=0; q_begin(nqpts), q_begin + q_batch_size); + nda::array V_batch(q_end - q_begin, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + V_batch() = ComplexType(0.0); + for (size_t iq_full=q_begin; iq_fullqp_to_ibz(iq_full); + if (q_origin <= iq and iq < q_origin + nq_loc) { + long iq_loc = iq - q_origin; + auto B_Q_ab = nda::reshape(B_qIPab(iq_full, 0, Q_rng, nda::ellipsis{}), + shape_t<2>{NQ_loc, nImpOrbs*nImpOrbs}); + if (_MF->qp_trev(iq_full)) { + V_PQ_loc = nda::conj( dV_qPQ.local()(iq_loc, nda::ellipsis{}) ); + nda::blas::gemm(V_PQ_loc, B_Q_ab, T_P_ab); + } else { + V_PQ_loc = dV_qPQ.local()(iq_loc, nda::ellipsis{}); + nda::blas::gemm(V_PQ_loc, B_Q_ab, T_P_ab); + } - for (size_t P = 0; P < NP_loc; ++P) { - auto B_dc = B_qIPab(iq_full, 0, P_origin+P, nda::ellipsis{}); - B_cdP_conj(nda::range::all, nda::range::all, P) = nda::conj(nda::transpose(B_dc)); + for (size_t P = 0; P < NP_loc; ++P) { + auto B_dc = B_qIPab(iq_full, 0, P_origin+P, nda::ellipsis{}); + B_cdP_conj(nda::range::all, nda::range::all, P) = nda::conj(nda::transpose(B_dc)); + } + auto V_2D = nda::reshape(V_batch(iq_full - q_begin, nda::ellipsis{}), + shape_t<2>{nImpOrbs*nImpOrbs, nImpOrbs*nImpOrbs}); + nda::blas::gemm(ComplexType(1.0), B_cd_P_conj, T_P_ab, ComplexType(1.0), V_2D); } - auto Vq_2D = nda::reshape(V_qcdab(iq_full, nda::ellipsis{}), shape_t<2>{nImpOrbs*nImpOrbs, nImpOrbs*nImpOrbs}); - nda::blas::gemm(ComplexType(1.0), B_cd_P_conj, T_P_ab, ComplexType(1.0), Vq_2D); } + reduce_in_place_chunked(comm, V_batch.data(), V_batch.size(), std::plus<>{}); + if (comm.root()) + V_qcdab(nda::range(q_begin, q_end), nda::ellipsis{}) = V_batch; } - comm.all_reduce_in_place_n(V_qcdab.data(), V_qcdab.size(), std::plus<>{}); return V_qcdab; } @@ -1741,7 +1842,7 @@ namespace methods { for (long P = 0; P < NQ; ++P) nda::blas::gerc(ComplexType(e0)*chi_head(P),T_skIPa(0, is, ik, I, P, nda::range::all), T_skIPa(0, is, ik, I, P, nda::range::all), BB_ab); - mpi->comm.reduce_in_place_n(BB_ab.data(),BB_ab.size(),std::plus<>{},0); + reduce_in_place_chunked(mpi->comm, BB_ab.data(), BB_ab.size(), std::plus<>{}); if(not root) BB_ab = nda::array(0,0); } @@ -1819,7 +1920,7 @@ namespace methods { nda::transpose(Vloc(iq, nda::ellipsis{})), value_type(0.0), T); } - mpi->comm.all_reduce_in_place_n(T.data(),T.size(),std::plus<>{}); + all_reduce_in_place_chunked(mpi->comm, T.data(), T.size(), std::plus<>{}); for( long ab=0; abcomm.all_reduce_in_place_n(B.data(),B.size(),std::plus<>{}); + all_reduce_in_place_chunked(mpi->comm, B.data(), B.size(), std::plus<>{}); nda::blas::gemm(value_type(1.0), B, Vloc(iq, nda::ellipsis{}), value_type(0.0), T); T = nda::conj(T); @@ -1918,7 +2019,7 @@ namespace methods { B() = ComplexType(0.0); for( auto [in,n] : itertools::enumerate(index) ) B(in,Q_rng) = nda::conj(Bq(all,n)); - mpi->comm.all_reduce_in_place_n(B.data(),B.size(),std::plus<>{}); + all_reduce_in_place_chunked(mpi->comm, B.data(), B.size(), std::plus<>{}); nda::blas::gemm(value_type(1.0), B, Vloc(iq, nda::ellipsis{}), value_type(0.0), T); nda::blas::gemm(norm, T, Bq, value_type(1.0), D); @@ -2142,7 +2243,7 @@ namespace methods { } } } - dW_wqPQ.communicator()->all_reduce_in_place_n(W_wcdab.data(), W_wcdab.size(), std::plus<>{}); + all_reduce_in_place_chunked(*dW_wqPQ.communicator(), W_wcdab.data(), W_wcdab.size(), std::plus<>{}); W_wcdab() /= nqpts; // finite-size correction to W_cdab(w) @@ -2154,6 +2255,21 @@ namespace methods { template B_t> auto embed_eri_t::downfold_Wq([[maybe_unused]] thc_t &thc, memory::darray_t &dW_wqPQ, const B_t &B_qIPab) + -> nda::array { + auto W_wqcdab = downfold_Wq_root(thc, dW_wqPQ, B_qIPab); + auto &comm = *dW_wqPQ.communicator(); + auto [nqpts, nImps, NP, nImpOrbs, nImpOrbs2] = B_qIPab.shape(); + long nw_half = dW_wqPQ.global_shape()[0]; + if (not comm.root()) + W_wqcdab.resize(nw_half, nqpts, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + broadcast_chunked(comm, W_wqcdab.data(), W_wqcdab.size()); + return W_wqcdab; + } + + template B_t> + auto embed_eri_t::downfold_Wq_root([[maybe_unused]] thc_t &thc, + memory::darray_t &dW_wqPQ, + const B_t &B_qIPab) -> nda::array { auto [nqpts, nImps, NP, nImpOrbs, nImpOrbs2] = B_qIPab.shape(); auto nw_half = dW_wqPQ.global_shape()[0]; @@ -2163,8 +2279,14 @@ namespace methods { auto q_origin = dW_wqPQ.origin()[1]; auto P_origin = dW_wqPQ.origin()[2]; - nda::array W_wqcdab(nw_half, nqpts, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); - W_wqcdab() = ComplexType(0.0); + auto &comm = *dW_wqPQ.communicator(); + // Root owns the complete HDF5 output; other ranks remain empty. + nda::array W_wqcdab; + if (comm.root()) + W_wqcdab.resize(nw_half, nqpts, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + size_t orbital_slab_elements = static_cast(nImpOrbs) * nImpOrbs * nImpOrbs * nImpOrbs; + size_t wq_batch_size = std::max(1, mpi_chunk_elements() / + std::max(1, orbital_slab_elements)); // W_wqcdab = conj(B_qPdc) * [ W_wqPQ ] * B_qQab nda::array T_P_ab(NP_loc, nImpOrbs*nImpOrbs); @@ -2173,15 +2295,21 @@ namespace methods { auto B_cd_P_conj = nda::reshape(B_cdP_conj, shape_t<2>{nImpOrbs*nImpOrbs, NP_loc}); auto W_loc = dW_wqPQ.local(); - for (size_t iw_loc = 0; iw_loc < nw_loc; ++iw_loc) { - size_t iw = w_origin + iw_loc; - //auto W_cd_ab = nda::reshape(W_wcdab(iw, nda::ellipsis()), shape_t<2>{nImpOrbs*nImpOrbs, nImpOrbs*nImpOrbs}); - for (size_t iq_loc = 0; iq_loc < nq_loc; ++iq_loc) { - size_t iq = q_origin + iq_loc; // iq lives inside IBZ - - // loop over all symmetry-related q-points - for (size_t iq_full=0; iq_fullqp_to_ibz(iq_full)!=iq) continue; + size_t nwq = static_cast(nw_half) * nqpts; + for (size_t wq_begin=0; wq_begin W_batch(wq_end - wq_begin, nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + W_batch() = ComplexType(0.0); + for (size_t iwq=wq_begin; iwqqp_to_ibz(iq_full); + bool owns_wq = w_origin <= static_cast(iw) and + static_cast(iw) < w_origin + nw_loc and + q_origin <= iq and iq < q_origin + nq_loc; + if (owns_wq) { + long iw_loc = static_cast(iw) - w_origin; + long iq_loc = iq - q_origin; auto B_Q_ab = nda::reshape(B_qIPab(iq_full, 0, Q_rng, nda::ellipsis{}), shape_t<2>{NQ_loc, nImpOrbs*nImpOrbs}); if (_MF->qp_trev(iq_full)) W_PQ = nda::conj( W_loc(iw_loc, iq_loc, nda::ellipsis()) ); @@ -2193,12 +2321,20 @@ namespace methods { auto B_dc = B_qIPab(iq_full, 0, P_origin+P, nda::ellipsis{}); B_cdP_conj(nda::range::all, nda::range::all, P) = nda::conj(nda::transpose(B_dc)); } - auto W_wq_2D = nda::reshape(W_wqcdab(iw, iq_full, nda::ellipsis{}), shape_t<2>{nImpOrbs*nImpOrbs, nImpOrbs*nImpOrbs}); - nda::blas::gemm(ComplexType(1.0), B_cd_P_conj, T_P_ab, ComplexType(1.0), W_wq_2D); + auto W_2D = nda::reshape(W_batch(iwq - wq_begin, nda::ellipsis{}), + shape_t<2>{nImpOrbs*nImpOrbs, nImpOrbs*nImpOrbs}); + nda::blas::gemm(ComplexType(1.0), B_cd_P_conj, T_P_ab, ComplexType(1.0), W_2D); + } + } + reduce_in_place_chunked(comm, W_batch.data(), W_batch.size(), std::plus<>{}); + if (comm.root()) { + for (size_t iwq=wq_begin; iwqall_reduce_in_place_n(W_wqcdab.data(), W_wqcdab.size(), std::plus<>{}); return W_wqcdab; } diff --git a/src/methods/embedding/embed_eri_t.h b/src/methods/embedding/embed_eri_t.h index 3ec6f382..5eaadebc 100644 --- a/src/methods/embedding/embed_eri_t.h +++ b/src/methods/embedding/embed_eri_t.h @@ -252,9 +252,14 @@ namespace methods { -> nda::array; template B_t> + // Replicated compatibility wrapper. auto downfold_Vq(thc_t &thc, const B_t &B_qIPab) -> nda::array; + template B_t> + auto downfold_Vq_root(thc_t &thc, const B_t &B_qIPab) + -> nda::array; + auto downfold_cholesky(THC_ERI auto &thc, const projector_boson_t &proj_boson, math::nda::DistributedArrayOfRank<3> auto &dV_qPQ, ComplexType div_correction_factor, double thresh = 1e-6); @@ -274,10 +279,16 @@ namespace methods { -> nda::array; template B_t> + // Replicated compatibility wrapper. auto downfold_Wq(thc_t &thc, memory::darray_t &dW_wqPQ, const B_t &B_qIPab) -> nda::array; + template B_t> + auto downfold_Wq_root(thc_t &thc, memory::darray_t &dW_wqPQ, + const B_t &B_qIPab) + -> nda::array; + template B_t> void V_div_correction(nda::array &V_cdab, const B_t &B_qIPab, thc_t &thc); diff --git a/src/methods/embedding/embed_t.cpp b/src/methods/embedding/embed_t.cpp index dd474534..33da6734 100644 --- a/src/methods/embedding/embed_t.cpp +++ b/src/methods/embedding/embed_t.cpp @@ -468,14 +468,8 @@ namespace methods { auto& sVhf_skij = mb_state.sF_skij.value(); auto& proj = mb_state.proj_boson.value().proj_fermi(); - auto sVhf_correction_upfold = sVhf_skij; - sVhf_correction_upfold.set_zero(); - nda::array Vhf_correction = mb_state.Vhf_imp_sIab.value() - mb_state.Vhf_dc_sIab.value(); - proj.upfold(sVhf_correction_upfold, Vhf_correction); - - if (sVhf_skij.node_comm()->root()) sVhf_skij.local() += sVhf_correction_upfold.local(); - sVhf_skij.communicator()->barrier(); + proj.upfold_add(sVhf_skij, Vhf_correction); } void embed_t::add_Sigma_dyn_correction(MBState &mb_state, bool subtract_dc) { @@ -485,9 +479,6 @@ namespace methods { auto nImpOrbs = proj.nImpOrbs(); auto& sSigma_tskij = mb_state.sSigma_tskij.value(); - auto sSigma_correction_upfold = sSigma_tskij; - sSigma_correction_upfold.set_zero(); - nda::array Sigma_imp_tsIab(mb_state.ft->nt_f(), _MF->nspin(), nImps, nImpOrbs, nImpOrbs); mb_state.ft->w_to_tau(mb_state.Sigma_imp_wsIab.value(), Sigma_imp_tsIab, imag_axes_ft::fermion); mb_state.ft->check_leakage(Sigma_imp_tsIab, imag_axes_ft::fermion, sSigma_tskij.communicator(), "impurity self-energy"); @@ -499,9 +490,7 @@ namespace methods { Sigma_imp_tsIab -= Sigma_dc_tsIab; } - proj.upfold(sSigma_correction_upfold, Sigma_imp_tsIab); - if (sSigma_tskij.node_comm()->root()) sSigma_tskij.local() += sSigma_correction_upfold.local(); - sSigma_tskij.communicator()->barrier(); + proj.upfold_add(sSigma_tskij, Sigma_imp_tsIab); } template Array_base_t> @@ -509,12 +498,7 @@ namespace methods { MBState &mb_state) { auto& proj = mb_state.proj_boson.value().proj_fermi(); - auto sVcorr_skij_upfold = sVcorr_skij; - sVcorr_skij_upfold.set_zero(); - - proj.upfold(sVcorr_skij_upfold, mb_state.Vcorr_dc_sIab.value()); - if (sVcorr_skij.node_comm()->root()) sVcorr_skij.local() -= sVcorr_skij_upfold.local(); - sVcorr_skij.communicator()->barrier(); + proj.upfold_add(sVcorr_skij, mb_state.Vcorr_dc_sIab.value(), ComplexType(-1.0)); } } // methods diff --git a/src/methods/embedding/projector_boson_t.cpp b/src/methods/embedding/projector_boson_t.cpp index b7026746..6b39df86 100644 --- a/src/methods/embedding/projector_boson_t.cpp +++ b/src/methods/embedding/projector_boson_t.cpp @@ -19,10 +19,33 @@ */ +#include +#include + #include "numerics/nda_functions.hpp" #include "methods/ERI/thc_reader_t.hpp" #include "methods/embedding/projector_boson_t.h" +namespace { + + constexpr long max_product_block_size = 64; + constexpr size_t max_projector_tile_bytes = size_t{32} * 1024 * 1024; + + // Keep enough independent blocks to occupy the MPI ranks when possible, + // while bounding both the product-basis dimension and the T tile bytes. + long product_block_size(long Np, long n_task_groups, long mpi_size, + size_t elements_per_product) { + n_task_groups = std::max(1L, n_task_groups); + auto blocks_per_group = std::max(1L, (mpi_size + n_task_groups - 1) / n_task_groups); + auto balanced_block_size = std::max(1L, Np / blocks_per_group); + auto max_tile_elements = std::max(1, max_projector_tile_bytes / sizeof(ComplexType)); + auto memory_block_size = std::max( + 1, static_cast(max_tile_elements / std::max(1, elements_per_product))); + return std::min({max_product_block_size, balanced_block_size, memory_block_size}); + } + +} // anonymous namespace + namespace methods { auto projector_boson_t::calc_bosonic_projector(THC_ERI auto &thc) const @@ -32,31 +55,48 @@ namespace methods { auto W_rng = _proj_fermi.W_rng(); auto nqpts = _MF->nqpts(); auto [ns, nkpts, nImps, nImpOrbs, nOrbs_W] = C_skIai.shape(); - - nda::array T_skIPa(ns, nkpts, nImps, thc.Np(), nImpOrbs); - for (long isk = 0; isk < ns*nkpts; ++isk) { - long is = isk / nkpts; // isk = is * nkpts + ik - long ik = isk % nkpts; - for (long I = 0; I < nImps; ++I) { - nda::blas::gemm(thc.X(is, 0, ik)(nda::range::all, W_rng[I]), nda::dagger(C_skIai(is, ik, I, nda::ellipsis{})), - T_skIPa(is, ik, I, nda::ellipsis{})); - } - } + long NP = thc.Np(); auto sB_qIPab = math::shm::make_shared_array( - *mpi, {nqpts, nImps, thc.Np(), nImpOrbs, nImpOrbs}); + *mpi, {nqpts, nImps, NP, nImpOrbs, nImpOrbs}); auto B_qIPab = sB_qIPab.local(); sB_qIPab.win().fence(); - for (size_t iqIP = mpi->comm.rank(); iqIP < nkpts*nImps*thc.Np(); iqIP += mpi->comm.size()) { - size_t iq = iqIP / (nImps * thc.Np()); // iqIP = iq * nImps * Np + I * Np + p - size_t I = (iqIP / thc.Np()) % nImps; - size_t P = iqIP % thc.Np(); + + // T is independent of q in the no-symmetry path. Distribute disjoint + // (impurity, product-basis block) pairs across ranks, then reuse each tile + // for every q. Thus every B(q,I,P,:,:) has exactly one writer before the + // existing all-reduce, without replicating the full T tensor on each rank. + size_t elements_per_product = static_cast(ns) * nkpts * nImpOrbs; + auto P_block_size = product_block_size( + NP, nImps, mpi->comm.size(), elements_per_product); + auto nP_blocks = (NP + P_block_size - 1) / P_block_size; + nda::array T_skPa(ns, nkpts, P_block_size, nImpOrbs); + for (long iIP = mpi->comm.rank(); iIP < nImps*nP_blocks; iIP += mpi->comm.size()) { + long I = iIP / nP_blocks; + long iP_block = iIP % nP_blocks; + long P0 = iP_block * P_block_size; + long P1 = std::min(P0 + P_block_size, NP); + auto P_rng = nda::range(P0, P1); + auto p_rng = nda::range(0, P1-P0); + for (long isk = 0; isk < ns*nkpts; ++isk) { long is = isk / nkpts; long ik = isk % nkpts; - long ikmq = _MF->qk_to_k2(iq, ik); - nda::blas::gerc(ComplexType(1.0),T_skIPa(is, ikmq, I, P, nda::range::all), - T_skIPa(is, ik, I, P, nda::range::all), B_qIPab(iq, I, P, nda::ellipsis{})); + nda::blas::gemm(thc.X(is, 0, ik)(P_rng, W_rng[I]), + nda::dagger(C_skIai(is, ik, I, nda::ellipsis{})), + T_skPa(is, ik, p_rng, nda::range::all)); + } + + for (long iq = 0; iq < nqpts; ++iq) { + for (long isk = 0; isk < ns*nkpts; ++isk) { + long is = isk / nkpts; + long ik = isk % nkpts; + long ikmq = _MF->qk_to_k2(iq, ik); + for (long p = 0; p < P1-P0; ++p) + nda::blas::gerc(ComplexType(1.0), T_skPa(is, ikmq, p, nda::range::all), + T_skPa(is, ik, p, nda::range::all), + B_qIPab(iq, I, P0+p, nda::ellipsis{})); + } } } sB_qIPab.win().fence(); @@ -76,61 +116,99 @@ namespace methods { auto [ns, nkpts, nImps, nImpOrbs, nOrbs_W] = C_skIai.shape(); auto nqpts = _MF->nqpts(); auto qsymms = _MF->qsymms(); + long NP = thc.Np(); auto sB_qIPab = math::shm::make_shared_array( - *mpi, {nqpts, nImps, thc.Np(), nImpOrbs, nImpOrbs}); + *mpi, {nqpts, nImps, NP, nImpOrbs, nImpOrbs}); auto B_qIPab = sB_qIPab.local(); sB_qIPab.win().fence(); // intermediate objects - nda::array T_skIPb(ns, nkpts, nImps, thc.Np(), nImpOrbs); + size_t elements_per_product = static_cast(ns) * nkpts * nImpOrbs; + auto P_block_size = product_block_size( + NP, nqpts*nImps, mpi->comm.size(), elements_per_product); + auto nP_blocks = (NP + P_block_size - 1) / P_block_size; + nda::array T_skPb(ns, nkpts, P_block_size, nImpOrbs); nda::array Cfull_jb(_MF->nbnd(), nImpOrbs); nda::array tmp_ib(_MF->nbnd(), nImpOrbs); + std::optional> Crot_skib; using math::sparse::T; using math::sparse::csrmm; sB_qIPab.win().fence(); - for (size_t iq=mpi->comm.rank(); iqcomm.size()) { + // T is q-dependent. Give every (q,I) at least one worker, adding only the + // minimum extra workers needed for MPI occupancy. Each worker owns a + // disjoint subset of P blocks and reuses one cached symmetry rotation. + long nqI = nqpts*nImps; + long max_worker_tasks = nqI*nP_blocks; + long nworker_tasks = std::min(max_worker_tasks, std::max(nqI, long(mpi->comm.size()))); + long workers_per_qI = nworker_tasks / nqI; + long qI_with_extra_worker = nworker_tasks % nqI; + for (long qI = 0; qI < nqI; ++qI) { + long iq = qI / nImps; + long I = qI % nImps; + long nworkers = workers_per_qI + (qI < qI_with_extra_worker ? 1 : 0); + long task0 = qI*workers_per_qI + std::min(qI, qI_with_extra_worker); + // symmetry index auto sym_it = std::find(qsymms.begin(), qsymms.end(), _MF->qp_symm(iq)); auto isym = std::distance(qsymms.begin(), sym_it); - // calculate T_skIPb - for (size_t isk=0; iskks_to_k(isym, ik); - for (size_t I=0; Icomm.size() != mpi->comm.rank()) continue; + + if (isym != 0) { + if (not Crot_skib) + Crot_skib.emplace(ns, nkpts, _MF->nbnd(), nImpOrbs); + for (long isk = 0; isk < ns*nkpts; ++isk) { + long is = isk / nkpts; + long ik = isk % nkpts; auto [cjg, D_ij] = _MF->symmetry_rotation(isym, ik); - // D_ij * Cfull_jb = tmp_ib Cfull_jb() = 0.0; - if(not cjg) { - Cfull_jb(W_rng[I], nda::range::all) = nda::conj(nda::transpose(C_skIai(is, ik, I, nda::ellipsis{}))); + if (not cjg) { + Cfull_jb(W_rng[I], nda::range::all) = + nda::conj(nda::transpose(C_skIai(is, ik, I, nda::ellipsis{}))); csrmm(ComplexType(1.0), *D_ij, Cfull_jb, ComplexType(0.0), tmp_ib); } else { - Cfull_jb(W_rng[I], nda::range::all) = nda::transpose(C_skIai(is, ik, I, nda::ellipsis{})); + Cfull_jb(W_rng[I], nda::range::all) = + nda::transpose(C_skIai(is, ik, I, nda::ellipsis{})); csrmm(ComplexType(1.0), *D_ij, Cfull_jb, ComplexType(0.0), tmp_ib); tmp_ib = nda::conj(tmp_ib); } - // X_Pi * tmp_ib = TskI_Pb - nda::blas::gemm(thc.X(is, 0, ikR), tmp_ib, T_skIPb(is,ik,I,nda::ellipsis{})); + Crot_skib.value()(is, ik, nda::ellipsis{}) = tmp_ib; } } - } - for (long I = 0; I < nImps; ++I) { - for (long isk = 0; isk < ns*nkpts; ++isk) { - long is = isk / nkpts; - long ik = isk % nkpts; - long ikmq = _MF->qk_to_k2(iq, ik); + for (long iP_block = iworker; iP_block < nP_blocks; iP_block += nworkers) { + long P0 = iP_block * P_block_size; + long P1 = std::min(P0 + P_block_size, NP); + auto P_rng = nda::range(P0, P1); + auto p_rng = nda::range(0, P1-P0); + + for (long isk = 0; isk < ns*nkpts; ++isk) { + long is = isk / nkpts; + long ik = isk % nkpts; + auto ikR = _MF->ks_to_k(isym, ik); + if (isym == 0) { + nda::blas::gemm(thc.X(is, 0, ik)(P_rng, W_rng[I]), + nda::dagger(C_skIai(is, ik, I, nda::ellipsis{})), + T_skPb(is, ik, p_rng, nda::range::all)); + } else { + nda::blas::gemm(thc.X(is, 0, ikR)(P_rng, nda::range::all), + Crot_skib.value()(is, ik, nda::ellipsis{}), + T_skPb(is, ik, p_rng, nda::range::all)); + } + } - for (long P = 0; P < thc.Np(); ++P) - nda::blas::gerc(ComplexType(1.0),T_skIPb(is, ikmq, I, P, nda::range::all), - T_skIPb(is, ik, I, P, nda::range::all), B_qIPab(iq, I, P,nda::ellipsis{})); + for (long isk = 0; isk < ns*nkpts; ++isk) { + long is = isk / nkpts; + long ik = isk % nkpts; + long ikmq = _MF->qk_to_k2(iq, ik); + for (long p = 0; p < P1-P0; ++p) + nda::blas::gerc(ComplexType(1.0), T_skPb(is, ikmq, p, nda::range::all), + T_skPb(is, ik, p, nda::range::all), + B_qIPab(iq, I, P0+p, nda::ellipsis{})); + } } } } @@ -155,4 +233,4 @@ namespace methods { template sArray_t projector_boson_t::calc_bosonic_projector(thc_reader_t &thc) const; template sArray_t projector_boson_t::calc_bosonic_projector_symm(thc_reader_t &thc) const; -} // \ No newline at end of file +} // diff --git a/src/methods/embedding/projector_t.cpp b/src/methods/embedding/projector_t.cpp index 482fe898..8e1cf8f7 100644 --- a/src/methods/embedding/projector_t.cpp +++ b/src/methods/embedding/projector_t.cpp @@ -19,12 +19,31 @@ */ +#include +#include + #include "methods/SCF/scf_common.hpp" #include "numerics/nda_functions.hpp" #include "methods/embedding/projector_t.h" namespace methods { + namespace { + // Each node owns one shared reduction tile. Keep its peak size bounded + // independently of the number of time/frequency and k points. + // A very small tile starves broad-window calculations because one rank + // owns each (tau,spin,k) matrix. 256 MiB still removes the multi-GiB + // full correction buffer while retaining useful concurrency on large + // communicators (for nOrbs_W=300 it holds 186 matrices, not only five). + constexpr std::size_t upfold_tile_bytes = 256 * 1024 * 1024; + + long upfold_tile_size(long nmat, long matrix_size) { + auto matrices_per_tile = static_cast( + upfold_tile_bytes / (sizeof(ComplexType) * static_cast(matrix_size))); + return std::min(nmat, std::max(1L, matrices_per_tile)); + } + } + void projector_t::print_metadata() { app_log(1, " Projector Information"); app_log(1, " ---------------------"); @@ -37,6 +56,14 @@ namespace methods { template Array_base_t, nda::ArrayOfRank<4> Oloc_t> void projector_t::upfold(sArray_t &O_skij, const Oloc_t &Oloc_sIab) const { + O_skij.set_zero(); + upfold_add(O_skij, Oloc_sIab); + } + + template Array_base_t, nda::ArrayOfRank<4> Oloc_t> + void projector_t::upfold_add(sArray_t &O_skij, const Oloc_t &Oloc_sIab, + ComplexType alpha) const { + utils::check(O_skij.shape()[0] == Oloc_sIab.shape(0) and O_skij.shape()[0] == _C_skIai.shape(0), "embed_t::upfold: ns mismatches.{}, {}, {}", O_skij.shape()[0], Oloc_sIab.shape(0), _C_skIai.shape(0)); @@ -45,44 +72,67 @@ namespace methods { utils::check(O_skij.shape()[1]==_MF->nkpts_ibz(), "embed_t::upfold: O_skij.shape[1]({})!=nkpts_ibz({}).", O_skij.shape()[1], _MF->nkpts_ibz()); - O_skij.set_zero(); auto [ns, nkpts_ibz, nbnd, nbnd_b] = O_skij.shape(); + utils::check(nbnd == nbnd_b, "embed_t::upfold: crystal-basis matrix is not square. {}, {}", nbnd, nbnd_b); nda::array buffer_ib(_nOrbs_W, _nImpOrbs); - - auto O_buffer = sArray_t( + long nsk = ns * nkpts_ibz; + long tile_size = upfold_tile_size(nsk, _nOrbs_W * _nOrbs_W); + auto O_tile = math::shm::shared_array>( O_skij.communicator(), O_skij.internode_comm(), O_skij.node_comm(), - {ns, nkpts_ibz, _nOrbs_W, _nOrbs_W}); + {tile_size, _nOrbs_W, _nOrbs_W}); + auto O_buf = O_tile.local(); + int rank = O_skij.communicator()->rank(); + int size = O_skij.communicator()->size(); - auto O_buf_loc = O_buffer.local(); - int rank = O_buffer.communicator()->rank(); - int size = O_buffer.communicator()->size(); + O_skij.node_sync(); + O_skij.win().fence(); for (long imp = 0; imp < _nImps; ++imp) { - O_buffer.set_zero(); - O_buffer.win().fence(); - for (long sk = rank; sk < ns * nkpts_ibz; sk += size) { - long is = sk / nkpts_ibz; - long ik = sk % nkpts_ibz; - - nda::blas::gemm(ComplexType(1.0), - nda::dagger(_C_skIai(is,ik,imp,nda::ellipsis{})), - Oloc_sIab(is,imp,nda::ellipsis{}), - ComplexType(0.0), - buffer_ib); - nda::blas::gemm(buffer_ib, _C_skIai(is,ik,imp,nda::ellipsis{}), O_buf_loc(is,ik,nda::ellipsis{})); - } - O_buffer.win().fence(); - O_buffer.all_reduce(); - if (O_skij.node_comm()->root()) { - O_skij.local()(nda::range::all, nda::range::all, _W_rng[imp], _W_rng[imp]) += O_buf_loc; + for (long tile_begin = 0; tile_begin < nsk; tile_begin += tile_size) { + long tile_count = std::min(tile_size, nsk - tile_begin); + O_tile.set_zero(); + O_tile.win().fence(); + long rank_offset = (rank + size - tile_begin % size) % size; + for (long offset = rank_offset; offset < tile_count; offset += size) { + long sk = tile_begin + offset; + long is = sk / nkpts_ibz; + long ik = sk % nkpts_ibz; + + nda::blas::gemm(ComplexType(1.0), + nda::dagger(_C_skIai(is,ik,imp,nda::ellipsis{})), + Oloc_sIab(is,imp,nda::ellipsis{}), + ComplexType(0.0), buffer_ib); + nda::blas::gemm(alpha, buffer_ib, _C_skIai(is,ik,imp,nda::ellipsis{}), + ComplexType(0.0), O_buf(offset,nda::ellipsis{})); + } + O_tile.win().fence(); + O_tile.all_reduce(); + if (O_skij.node_comm()->root()) { + auto O_loc = O_skij.local(); + for (long offset = 0; offset < tile_count; ++offset) { + long sk = tile_begin + offset; + long is = sk / nkpts_ibz; + long ik = sk % nkpts_ibz; + O_loc(is,ik,_W_rng[imp],_W_rng[imp]) += O_buf(offset,nda::ellipsis{}); + } + } } } - O_skij.communicator()->barrier(); + O_skij.win().fence(); + O_skij.node_sync(); } template Array_base_t, nda::ArrayOfRank<5> Oloc_t> void projector_t::upfold(sArray_t &O_tskij, const Oloc_t &Oloc_tsIab) const { + O_tskij.set_zero(); + upfold_add(O_tskij, Oloc_tsIab); + } + + template Array_base_t, nda::ArrayOfRank<5> Oloc_t> + void projector_t::upfold_add(sArray_t &O_tskij, const Oloc_t &Oloc_tsIab, + ComplexType alpha) const { + utils::check(O_tskij.shape()[0] == Oloc_tsIab.shape(0), "embed_t::upfold: nts mismatches. {}, {}", O_tskij.shape()[0], Oloc_tsIab.shape(0)); utils::check(O_tskij.shape()[1] == Oloc_tsIab.shape(1) and @@ -93,48 +143,69 @@ namespace methods { utils::check(O_tskij.shape()[2]==_MF->nkpts_ibz(), "embed_t::upfold: O_tskij.shape[2]({})!=nkpts_ibz({}).", O_tskij.shape()[2], _MF->nkpts_ibz()); - O_tskij.set_zero(); auto [nts, ns, nkpts_ibz, nbnd, nbnd_b] = O_tskij.shape(); + utils::check(nbnd == nbnd_b, "embed_t::upfold: crystal-basis matrix is not square. {}, {}", nbnd, nbnd_b); nda::array buffer_ib(_nOrbs_W, _nImpOrbs); - - auto O_buffer = sArray_t( + long ntsk = nts * ns * nkpts_ibz; + long tile_size = upfold_tile_size(ntsk, _nOrbs_W * _nOrbs_W); + auto O_tile = math::shm::shared_array>( O_tskij.communicator(), O_tskij.internode_comm(), O_tskij.node_comm(), - {nts, ns, nkpts_ibz, _nOrbs_W, _nOrbs_W}); + {tile_size, _nOrbs_W, _nOrbs_W}); + auto O_buf = O_tile.local(); + int rank = O_tskij.communicator()->rank(); + int size = O_tskij.communicator()->size(); - auto O_buf = O_buffer.local(); - int rank = O_buffer.communicator()->rank(); - int size = O_buffer.communicator()->size(); + O_tskij.node_sync(); + O_tskij.win().fence(); for (long imp = 0; imp < _nImps; ++imp) { - O_buffer.set_zero(); - O_buffer.win().fence(); - for (long tsk = rank; tsk < nts*ns*nkpts_ibz; tsk += size) { - long it = tsk / (ns*nkpts_ibz); // tsk = it * ns*nkpts_ibz + is * nkpts_ibz + ik - long is = (tsk / nkpts_ibz) % ns; - long ik = tsk % nkpts_ibz; - - nda::blas::gemm(ComplexType(1.0), - nda::dagger(_C_skIai(is,ik,imp,nda::ellipsis{})), - Oloc_tsIab(it,is,imp,nda::ellipsis{}), - ComplexType(0.0), - buffer_ib); - - nda::blas::gemm(buffer_ib, _C_skIai(is,ik,imp,nda::ellipsis{}), O_buf(it,is,ik,nda::ellipsis{})); + for (long tile_begin = 0; tile_begin < ntsk; tile_begin += tile_size) { + long tile_count = std::min(tile_size, ntsk - tile_begin); + O_tile.set_zero(); + O_tile.win().fence(); + long rank_offset = (rank + size - tile_begin % size) % size; + for (long offset = rank_offset; offset < tile_count; offset += size) { + long tsk = tile_begin + offset; + long it = tsk / (ns*nkpts_ibz); + long is = (tsk / nkpts_ibz) % ns; + long ik = tsk % nkpts_ibz; + + nda::blas::gemm(ComplexType(1.0), + nda::dagger(_C_skIai(is,ik,imp,nda::ellipsis{})), + Oloc_tsIab(it,is,imp,nda::ellipsis{}), + ComplexType(0.0), buffer_ib); + nda::blas::gemm(alpha, buffer_ib, _C_skIai(is,ik,imp,nda::ellipsis{}), + ComplexType(0.0), O_buf(offset,nda::ellipsis{})); + } + O_tile.win().fence(); + O_tile.all_reduce(); + if (O_tskij.node_comm()->root()) { + auto O_loc = O_tskij.local(); + for (long offset = 0; offset < tile_count; ++offset) { + long tsk = tile_begin + offset; + long it = tsk / (ns*nkpts_ibz); + long is = (tsk / nkpts_ibz) % ns; + long ik = tsk % nkpts_ibz; + O_loc(it,is,ik,_W_rng[imp],_W_rng[imp]) += O_buf(offset,nda::ellipsis{}); + } + } } - O_buffer.win().fence(); - O_buffer.all_reduce(); - - O_tskij.win().fence(); - if (O_tskij.node_comm()->root()) { - O_tskij.local()(nda::range::all, nda::range::all, nda::range::all, _W_rng[imp], _W_rng[imp]) += O_buf; - } - O_tskij.win().fence(); } + O_tskij.win().fence(); + O_tskij.node_sync(); } template Array_base_t, nda::ArrayOfRank<6> Ac_t> void projector_t::upfold(sArray_t &O_tskij, const Ac_t &O_tskIab) const { + O_tskij.set_zero(); + upfold_add(O_tskij, O_tskIab); + } + + template Array_base_t, nda::ArrayOfRank<6> Ac_t> + void projector_t::upfold_add(sArray_t &O_tskij, const Ac_t &O_tskIab, + ComplexType alpha) const { + utils::check(O_tskij.shape()[0] == O_tskIab.shape(0), "embed_t::upfold: nts mismatches. {}, {}", O_tskij.shape()[0], O_tskIab.shape(0)); utils::check(O_tskij.shape()[1] == O_tskIab.shape(1) and @@ -144,41 +215,55 @@ namespace methods { utils::check(O_tskIab.shape(4) == _nImpOrbs, "embed_t::upfold: nImpOrbs mismatches. {}, {}", O_tskIab.shape(4), _nImpOrbs); - O_tskij.set_zero(); auto [nts, ns, nkpts, nbnd, nbnd_b] = O_tskij.shape(); + utils::check(nbnd == nbnd_b, "embed_t::upfold: crystal-basis matrix is not square. {}, {}", nbnd, nbnd_b); nda::array buffer_ib(_nOrbs_W, _nImpOrbs); - - auto O_buffer = sArray_t( + long ntsk = nts * ns * nkpts; + long tile_size = upfold_tile_size(ntsk, _nOrbs_W * _nOrbs_W); + auto O_tile = math::shm::shared_array>( O_tskij.communicator(), O_tskij.internode_comm(), O_tskij.node_comm(), - {nts, ns, nkpts, _nOrbs_W, _nOrbs_W}); + {tile_size, _nOrbs_W, _nOrbs_W}); + auto O_buf = O_tile.local(); + int rank = O_tskij.communicator()->rank(); + int size = O_tskij.communicator()->size(); - auto O_buf = O_buffer.local(); - int rank = O_buffer.communicator()->rank(); - int size = O_buffer.communicator()->size(); + O_tskij.node_sync(); + O_tskij.win().fence(); for (long imp = 0; imp < _nImps; ++imp) { - O_buffer.set_zero(); - O_buffer.win().fence(); - for (long tsk = rank; tsk < nts * ns * nkpts; tsk += size) { - long it = tsk / (ns * nkpts); // tsk = it * ns*nkpts + is * nkpts + ik - long is = (tsk / nkpts) % ns; - long ik = tsk % nkpts; - - nda::blas::gemm(ComplexType(1.0), nda::dagger(_C_skIai(is, ik, imp, nda::ellipsis{})), - O_tskIab(it, is, ik, imp, nda::ellipsis{}), - ComplexType(0.0), buffer_ib); - - nda::blas::gemm(buffer_ib, _C_skIai(is,ik,imp,nda::ellipsis{}), O_buf(it,is,ik,nda::ellipsis{})); - } - O_buffer.win().fence(); - O_buffer.all_reduce(); - - O_tskij.win().fence(); - if (O_tskij.node_comm()->root()) { - O_tskij.local()(nda::range::all, nda::range::all, nda::range::all, _W_rng[imp], _W_rng[imp]) += O_buf; + for (long tile_begin = 0; tile_begin < ntsk; tile_begin += tile_size) { + long tile_count = std::min(tile_size, ntsk - tile_begin); + O_tile.set_zero(); + O_tile.win().fence(); + long rank_offset = (rank + size - tile_begin % size) % size; + for (long offset = rank_offset; offset < tile_count; offset += size) { + long tsk = tile_begin + offset; + long it = tsk / (ns * nkpts); + long is = (tsk / nkpts) % ns; + long ik = tsk % nkpts; + + nda::blas::gemm(ComplexType(1.0), nda::dagger(_C_skIai(is, ik, imp, nda::ellipsis{})), + O_tskIab(it, is, ik, imp,nda::ellipsis{}), + ComplexType(0.0), buffer_ib); + nda::blas::gemm(alpha, buffer_ib, _C_skIai(is,ik,imp,nda::ellipsis{}), + ComplexType(0.0), O_buf(offset,nda::ellipsis{})); + } + O_tile.win().fence(); + O_tile.all_reduce(); + if (O_tskij.node_comm()->root()) { + auto O_loc = O_tskij.local(); + for (long offset = 0; offset < tile_count; ++offset) { + long tsk = tile_begin + offset; + long it = tsk / (ns * nkpts); + long is = (tsk / nkpts) % ns; + long ik = tsk % nkpts; + O_loc(it,is,ik,_W_rng[imp],_W_rng[imp]) += O_buf(offset,nda::ellipsis{}); + } + } } - O_tskij.win().fence(); } + O_tskij.win().fence(); + O_tskij.node_sync(); } template @@ -497,6 +582,9 @@ namespace methods { template void projector_t::upfold(sArray_t>&, const nda::array&) const; template void projector_t::upfold(sArray_t>&, const nda::array&) const; template void projector_t::upfold(sArray_t>&, const nda::array&) const; + template void projector_t::upfold_add(sArray_t>&, const nda::array&, ComplexType) const; + template void projector_t::upfold_add(sArray_t>&, const nda::array&, ComplexType) const; + template void projector_t::upfold_add(sArray_t>&, const nda::array&, ComplexType) const; template nda::array projector_t::downfold_k_fbz(const sArray_t>&) const; diff --git a/src/methods/embedding/projector_t.h b/src/methods/embedding/projector_t.h index 1fc01832..818ce131 100644 --- a/src/methods/embedding/projector_t.h +++ b/src/methods/embedding/projector_t.h @@ -229,6 +229,17 @@ namespace methods { template Array_base_t, nda::ArrayOfRank<6> Ac_t> void upfold(sArray_t &O_tskij, const Ac_t &O_tskIab) const; + /** Upfold Oloc and accumulate alpha * Oloc in an existing crystal-basis array. */ + template Array_base_t, nda::ArrayOfRank<4> Oloc_t> + void upfold_add(sArray_t &O_skij, const Oloc_t &Oloc_sIab, + ComplexType alpha = ComplexType(1.0)) const; + template Array_base_t, nda::ArrayOfRank<5> Oloc_t> + void upfold_add(sArray_t &O_tskij, const Oloc_t &Oloc_tsIab, + ComplexType alpha = ComplexType(1.0)) const; + template Array_base_t, nda::ArrayOfRank<6> Ac_t> + void upfold_add(sArray_t &O_tskij, const Ac_t &O_tskIab, + ComplexType alpha = ComplexType(1.0)) const; + template Array_base_t> auto downfold_k_fbz(const sArray_t &O_tskij) const -> nda::array; diff --git a/src/methods/embedding/tests/test_embed.cpp b/src/methods/embedding/tests/test_embed.cpp index d4bcf3b0..3f8dc454 100644 --- a/src/methods/embedding/tests/test_embed.cpp +++ b/src/methods/embedding/tests/test_embed.cpp @@ -50,6 +50,123 @@ namespace bdft_tests { namespace mpi3 = boost::mpi3; using namespace methods; + TEST_CASE("projector upfold_add matches direct Cdagger O C", "[methods][embed][projector]") { + auto& mpi = utils::make_unit_test_mpi_context(); + auto mf = mf::default_MF(mpi, "qe_lih222"); + + auto ns = mf.nspin(); + auto nk = mf.nkpts_ibz(); + auto nbnd = mf.nbnd(); + constexpr long nImps = 1; + constexpr long nImpOrbs = 2; + constexpr long nW = 2; + constexpr long W0 = 1; + REQUIRE(nbnd >= W0+nW); + + // Keep the coefficients in the test so the reference does not call either + // upfold implementation. The in-memory constructor reorders no k points + // when the MF k-point array itself is supplied. + nda::array C_ksIai(nk, ns, nImps, nImpOrbs, nW); + for (long k = 0; k < nk; ++k) + for (long s = 0; s < ns; ++s) + for (long a = 0; a < nImpOrbs; ++a) + for (long i = 0; i < nW; ++i) + C_ksIai(k,s,0,a,i) = ComplexType( + 0.15*(1+a+i) + 0.01*(1+k+s), 0.025*(1+k)*(a-i)); + + nda::array band_window(nImps, nk, 2); + for (long k = 0; k < nk; ++k) { + band_window(0,k,0) = W0+1; // 1-based inclusive lower bound + band_window(0,k,1) = W0+nW; // 1-based inclusive upper bound + } + nda::array kpts_crys = mf.kpts_crystal(); + projector_t proj(mf, C_ksIai, band_window, kpts_crys, false, false); + + nda::array Oloc_sIab(ns, nImps, nImpOrbs, nImpOrbs); + for (long s = 0; s < ns; ++s) + for (long a = 0; a < nImpOrbs; ++a) + for (long b = 0; b < nImpOrbs; ++b) + Oloc_sIab(s,0,a,b) = ComplexType(0.1*(1+s+a+2*b), 0.03*(1+a-b)); + + nda::array seed(ns, nk, nbnd, nbnd); + seed() = ComplexType(0.375, -0.125); + auto static_oracle = [&](ComplexType alpha) { + nda::array expected = seed; + for (long s = 0; s < ns; ++s) + for (long k = 0; k < nk; ++k) + for (long i = 0; i < nW; ++i) + for (long j = 0; j < nW; ++j) + for (long a = 0; a < nImpOrbs; ++a) + for (long b = 0; b < nImpOrbs; ++b) + expected(s,k,W0+i,W0+j) += alpha * std::conj(C_ksIai(k,s,0,a,i)) + * Oloc_sIab(s,0,a,b) * C_ksIai(k,s,0,b,j); + return expected; + }; + + auto sTarget = math::shm::make_shared_array( + *mpi, {ns, nk, nbnd, nbnd}); + for (ComplexType alpha : {ComplexType(1.0), ComplexType(-1.0)}) { + if (sTarget.node_comm()->root()) sTarget.local() = seed; + sTarget.node_sync(); + proj.upfold_add(sTarget, Oloc_sIab, alpha); + ARRAY_EQUAL(sTarget.local(), static_oracle(alpha), 1e-12); + } + + constexpr long nt = 2; + nda::array Oloc_tsIab(nt, ns, nImps, nImpOrbs, nImpOrbs); + for (long t = 0; t < nt; ++t) + Oloc_tsIab(t,nda::ellipsis{}) = ComplexType(t+1.0, -0.125*t) * Oloc_sIab; + + nda::array seed_t(nt, ns, nk, nbnd, nbnd); + seed_t() = ComplexType(-0.25, 0.0625); + auto time_oracle = [&] { + nda::array expected = seed_t; + for (long t = 0; t < nt; ++t) + for (long s = 0; s < ns; ++s) + for (long k = 0; k < nk; ++k) + for (long i = 0; i < nW; ++i) + for (long j = 0; j < nW; ++j) + for (long a = 0; a < nImpOrbs; ++a) + for (long b = 0; b < nImpOrbs; ++b) + expected(t,s,k,W0+i,W0+j) += std::conj(C_ksIai(k,s,0,a,i)) + * Oloc_tsIab(t,s,0,a,b) * C_ksIai(k,s,0,b,j); + return expected; + }; + + auto sTarget_t = math::shm::make_shared_array( + *mpi, {nt, ns, nk, nbnd, nbnd}); + if (sTarget_t.node_comm()->root()) sTarget_t.local() = seed_t; + sTarget_t.node_sync(); + proj.upfold_add(sTarget_t, Oloc_tsIab); + ARRAY_EQUAL(sTarget_t.local(), time_oracle(), 1e-12); + + nda::array Oloc_tskIab(nt, ns, nk, nImps, nImpOrbs, nImpOrbs); + for (long t = 0; t < nt; ++t) + for (long s = 0; s < ns; ++s) + for (long k = 0; k < nk; ++k) + Oloc_tskIab(t,s,k,nda::ellipsis{}) = ComplexType(1.0+0.2*k, 0.1*t) + * Oloc_sIab(s,nda::ellipsis{}); + + auto k_time_oracle = [&] { + nda::array expected = seed_t; + for (long t = 0; t < nt; ++t) + for (long s = 0; s < ns; ++s) + for (long k = 0; k < nk; ++k) + for (long i = 0; i < nW; ++i) + for (long j = 0; j < nW; ++j) + for (long a = 0; a < nImpOrbs; ++a) + for (long b = 0; b < nImpOrbs; ++b) + expected(t,s,k,W0+i,W0+j) += std::conj(C_ksIai(k,s,0,a,i)) + * Oloc_tskIab(t,s,k,0,a,b) * C_ksIai(k,s,0,b,j); + return expected; + }; + + if (sTarget_t.node_comm()->root()) sTarget_t.local() = seed_t; + sTarget_t.node_sync(); + proj.upfold_add(sTarget_t, Oloc_tskIab); + ARRAY_EQUAL(sTarget_t.local(), k_time_oracle(), 1e-12); + } + TEST_CASE("downfold_1e_mb", "[methods][embed][df_1e]") { auto& mpi = utils::make_unit_test_mpi_context(); @@ -525,7 +642,8 @@ TEST_CASE("downfold_1e_mb_qp", "[methods][embed][df_1e]") { TEST_CASE("compute_downfolded_coulomb_tensors", "[methods][embed][df_2e]") { auto& mpi = utils::make_unit_test_mpi_context(); - auto test_compute = [&](std::shared_ptr &mf, std::string wannier_file) { + auto test_compute = [&](std::shared_ptr &mf, std::string wannier_file, + bool test_q_dependent_output) { int nIpts = mf->nbnd() * 20; std::string cd_dir = ""; std::string storage = "incore"; @@ -613,6 +731,46 @@ TEST_CASE("downfold_1e_mb_qp", "[methods][embed][df_1e]") { VALUE_EQUAL(Wloc_crpa(0,1,1,1,1), -0.220910992415, 1e-5); VALUE_EQUAL(Wloc_crpa(0,0,0,1,1), -0.115140041097, 1e-5); + if (test_q_dependent_output) { + // The q-resolved implementation keeps the full tensor only on the + // HDF5-writing rank. Its public local tensors and on-disk schema must + // remain identical. Ignore both divergence corrections here so the + // local result is exactly the arithmetic average of the stored q slabs. + embed_eri_t embed_q(*mf, "ignore_g0", "ignore_g0"); + auto [Vloc_crpa_q, Wloc_crpa_q] = embed_q.compute_downfolded_coulomb_tensors( + thc, mb_state, screen_type_crpa, false, false, &ft, + greens_func_source, greens_func_iteration, true, true); + + if (mpi->comm.root()) { + nda::array V_qabcd; + nda::array U_qwabcd; + h5::file file(prefix + ".mbpt.h5", 'r'); + auto grp = h5::group(file).open_group("scf/iter0/downfolded_model"); + nda::h5_read(grp, "V_qabcd", V_qabcd); + nda::h5_read(grp, "U_qwabcd", U_qwabcd); + long nImpOrbs = Vloc_crpa_q.shape()[0]; + REQUIRE(V_qabcd.shape() == std::array{mf->nqpts(), nImpOrbs, + nImpOrbs, nImpOrbs, nImpOrbs}); + REQUIRE(U_qwabcd.shape() == std::array{Wloc_crpa_q.shape()[0], mf->nqpts(), + nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs}); + nda::array V_from_q(nImpOrbs, nImpOrbs, nImpOrbs, nImpOrbs); + nda::array W_from_q(Wloc_crpa_q.shape()); + V_from_q() = ComplexType(0.0); + W_from_q() = ComplexType(0.0); + for (long iq = 0; iq < mf->nqpts(); ++iq) { + V_from_q += V_qabcd(iq, nda::ellipsis{}); + W_from_q += U_qwabcd(nda::range::all, iq, nda::ellipsis{}); + } + V_from_q() /= mf->nqpts(); + W_from_q() /= mf->nqpts(); + ARRAY_EQUAL(Vloc_crpa_q, V_from_q, 1e-11); + ARRAY_EQUAL(Wloc_crpa_q, W_from_q, 1e-11); + REQUIRE(nda::sum(nda::abs(V_qabcd(0, nda::ellipsis{}))) > 0.0); + REQUIRE(nda::sum(nda::abs(U_qwabcd(0, 0, nda::ellipsis{}))) > 0.0); + } + mpi->comm.barrier(); + } + mpi->comm.barrier(); if (mpi->comm.root()) { @@ -626,14 +784,14 @@ TEST_CASE("downfold_1e_mb_qp", "[methods][embed][df_1e]") { auto [outdir, prefix] = utils::utest_filename("qe_lih222"); auto mf = std::make_shared(mf::default_MF(mpi, "qe_lih222")); std::string wannier_file = outdir + "/lih_wan.h5"; - test_compute(mf, wannier_file); + test_compute(mf, wannier_file, true); } SECTION("sym_qe") { auto [outdir, prefix] = utils::utest_filename("qe_lih222_sym"); auto mf = std::make_shared(mf::default_MF(mpi, "qe_lih222_sym")); std::string wannier_file = outdir + "/lih_wan.h5"; - test_compute(mf, wannier_file); + test_compute(mf, wannier_file, true); } } diff --git a/src/numerics/distributed_array/detail/concepts.hpp b/src/numerics/distributed_array/detail/concepts.hpp index 3b97e0bb..95959f2c 100644 --- a/src/numerics/distributed_array/detail/concepts.hpp +++ b/src/numerics/distributed_array/detail/concepts.hpp @@ -34,20 +34,27 @@ using get_value_t = typename std::decay_t::value_type; /* * Some concepts */ +// A collection of rank-local rectangular blocks. Unlike DistributedArray, +// this concept does not claim that the blocks originate from a regular +// processor grid or have a uniform algorithmic block size. template -concept DistributedArray = requires(A const& a) { +concept BlockDistributedArray = requires(A const& a) { { ::nda::MemoryArray::Array_t> }; { std::decay_t::rank > 0 }; { std::is_scalar>::value }; { a.communicator() }; - { a.grid() }; { a.local() }; { a.local_shape() } -> ::nda::StdArrayOfLong; { a.global_shape() } -> ::nda::StdArrayOfLong; - { a.block_size() } -> ::nda::StdArrayOfLong; { a.origin() } -> ::nda::StdArrayOfLong; }; +template +concept DistributedArray = BlockDistributedArray and requires(A const& a) { + { a.grid() }; + { a.block_size() } -> ::nda::StdArrayOfLong; +}; + template concept DistributedArrayView = DistributedArray and requires(A const& a) { { ::nda::MemoryArray::Array_view_t> }; diff --git a/src/numerics/distributed_array/nda.hpp b/src/numerics/distributed_array/nda.hpp index d125e444..7c280d51 100644 --- a/src/numerics/distributed_array/nda.hpp +++ b/src/numerics/distributed_array/nda.hpp @@ -36,7 +36,9 @@ using darray_t = math::nda::distributed_array; template<::nda::MemoryArray local_Array_t,class comm> using darray_view_t = math::nda::distributed_array_view; + +template<::nda::MemoryArray local_Array_t,class comm> +using irregular_block_darray_t = math::nda::irregular_block_distributed_array; } #endif - diff --git a/src/numerics/distributed_array/nda_matrix.hpp b/src/numerics/distributed_array/nda_matrix.hpp index 2922f1c9..9a1ae36e 100644 --- a/src/numerics/distributed_array/nda_matrix.hpp +++ b/src/numerics/distributed_array/nda_matrix.hpp @@ -351,6 +351,63 @@ class distributed_array Array_t A; }; +/* + * Owning collection of irregular rank-local rectangular blocks. This type is + * deliberately not a DistributedArray: it exposes no synthetic processor grid + * or algorithmic block size. It is intended for metadata-driven transfers of + * selected slices and other sparse-in-rank block layouts. + */ +template<::nda::Array Array_base_t, typename communicator_t> +class irregular_block_distributed_array { +public: + using Array_t = typename std::decay_t::regular_type; + static constexpr int rank = ::nda::get_rank; + using value_type = typename Array_t::value_type; + + template<::nda::MemoryArray Arr> + requires (::nda::get_rank> == rank) + irregular_block_distributed_array(communicator_t *comm, + std::array global_shape, + std::array origin, + Arr &&local) : + comm_(comm), global_shape_(global_shape), origin_(origin), + local_(std::forward(local)) { + utils::check(comm_ != nullptr, "irregular_block_distributed_array: Null communicator."); + for (int dim = 0; dim < rank; ++dim) { + utils::check(global_shape_[dim] >= 0 and origin_[dim] >= 0 and + local_.shape()[dim] >= 0 and + origin_[dim] <= global_shape_[dim] - local_.shape()[dim], + "irregular_block_distributed_array: Invalid block on axis {}: origin {}, size {}, global {}.", + dim, origin_[dim], local_.shape()[dim], global_shape_[dim]); + } +#if defined(SYNCHRONIZE_DISTRIBUTED_ARRAY) + auto root_shape = global_shape_; + comm_->broadcast_n(root_shape.data(), rank, 0); + utils::check(root_shape == global_shape_, + "irregular_block_distributed_array: Inconsistent global shape."); +#endif + } + + irregular_block_distributed_array() = default; + irregular_block_distributed_array(irregular_block_distributed_array const&) = default; + irregular_block_distributed_array(irregular_block_distributed_array&&) = default; + irregular_block_distributed_array& operator=(irregular_block_distributed_array const&) = default; + irregular_block_distributed_array& operator=(irregular_block_distributed_array&&) = default; + + auto local() { return local_(); } + auto local() const { return local_(); } + auto const& local_shape() const { return local_.shape(); } + auto const& global_shape() const { return global_shape_; } + auto const& origin() const { return origin_; } + communicator_t *communicator() const { return comm_; } + +private: + communicator_t *comm_ = nullptr; + std::array global_shape_{}; + std::array origin_{}; + Array_t local_; +}; + /* * Non-owning version of distributed array */ diff --git a/src/numerics/distributed_array/nda_utils.hpp b/src/numerics/distributed_array/nda_utils.hpp index fdd2739b..3ffda66c 100644 --- a/src/numerics/distributed_array/nda_utils.hpp +++ b/src/numerics/distributed_array/nda_utils.hpp @@ -24,6 +24,9 @@ #include #include +#include +#include +#include #include "configuration.hpp" #include "utilities/check.hpp" #include "nda/nda.hpp" @@ -862,10 +865,436 @@ void redistribute_alltoallv(Arr1_t& A, Arr2_t& B, get_value_t a = 1, get } } -template +// Layout metadata shared by repeated bounded-memory redistributions. Building +// the plan is collective, while executing it does not repeat the layout +// all-gather. This matters for hot paths (notably Sigma) whose distributions +// stay fixed across many tau points and symmetry operations. +template +class redistribution_plan { +public: + static constexpr long rank = Rank; + + template + redistribution_plan(Arr1_t const& A, Arr2_t const& B, size_t max_chunk_elements = 0) { + using value_t = typename std::decay_t::Array_t::value_type; + static_assert(get_rank == Rank and get_rank == Rank, "Rank mismatch."); + utils::check(A.global_shape() == B.global_shape(), "Size mismatch."); + utils::check(*A.communicator() == *B.communicator(), "Communicator mismatch."); + + global_shape_ = A.global_shape(); + A_origin_ = A.origin(); + A_local_shape_ = A.local_shape(); + B_origin_ = B.origin(); + B_local_shape_ = B.local_shape(); + mpi_size_ = A.communicator()->size(); + mpi_rank_ = A.communicator()->rank(); + + std::copy_n(A_origin_.data(), Rank, local_blocks_.data()); + std::copy_n(A_local_shape_.data(), Rank, local_blocks_.data() + Rank); + std::copy_n(B_origin_.data(), Rank, local_blocks_.data() + 2 * Rank); + std::copy_n(B_local_shape_.data(), Rank, local_blocks_.data() + 3 * Rank); + blocks_.resize(static_cast(mpi_size_) * 4 * Rank); + A.communicator()->all_gather_n(local_blocks_.data(), 4 * Rank, blocks_.data(), 4 * Rank); + + destination_element_size_ = sizeof(value_t); + size_t local_pack_elements = static_cast(A.local().size()); + utils::check(local_pack_elements <= std::numeric_limits::max() - + static_cast(B.local().size()), + "Local pack-size overflow in redistribution_plan."); + local_pack_elements += static_cast(B.local().size()); + size_t max_pack_elements = A.communicator()->all_reduce_value( + local_pack_elements, boost::mpi3::max<>{}); + constexpr size_t collective_pack_budget = size_t{64} * 1024 * 1024; + collective_pack_too_large_ = + max_pack_elements > collective_pack_budget / destination_element_size_; + + constexpr size_t default_chunk_bytes = size_t{32} * 1024 * 1024; + size_t mpi_count_limit = static_cast(std::numeric_limits::max()); + if (max_chunk_elements == 0) + max_chunk_elements = std::max(1, default_chunk_bytes / sizeof(value_t)); + max_chunk_elements_ = std::max(1, std::min(max_chunk_elements, mpi_count_limit)); + } + + template + void validate(Arr1_t const& A, Arr2_t const& B) const { + static_assert(get_rank == Rank and get_rank == Rank, "Rank mismatch."); + utils::check(A.global_shape() == global_shape_ and B.global_shape() == global_shape_, + "redistribution_plan global shape changed."); + utils::check(A.origin() == A_origin_ and A.local_shape() == A_local_shape_, + "redistribution_plan source layout changed."); + utils::check(B.origin() == B_origin_ and B.local_shape() == B_local_shape_, + "redistribution_plan destination layout changed."); + utils::check(*A.communicator() == *B.communicator() and + A.communicator()->size() == mpi_size_ and A.communicator()->rank() == mpi_rank_, + "redistribution_plan communicator changed."); + using destination_value_t = typename std::decay_t::Array_t::value_type; + utils::check(sizeof(destination_value_t) == destination_element_size_, + "redistribution_plan destination element size changed: {} != {}.", + sizeof(destination_value_t), destination_element_size_); + } + + void validate_source_coverage() const { validate_exact_coverage(0, "source"); } + void validate_destination_coverage() const { validate_exact_coverage(2, "destination"); } + + long mpi_size() const { return mpi_size_; } + long mpi_rank() const { return mpi_rank_; } + size_t max_chunk_elements() const { return max_chunk_elements_; } + bool collective_pack_too_large() const { return collective_pack_too_large_; } + long local_block(long slot, long axis) const { + return local_blocks_[static_cast(slot * Rank + axis)]; + } + long block(long peer, long slot, long axis) const { + return blocks_[static_cast((peer * 4 + slot) * Rank + axis)]; + } + +private: + size_t checked_block_volume(long peer, long slot, const char *label) const { + size_t volume = 1; + bool empty = false; + for (long axis = 0; axis < Rank; ++axis) { + long origin = block(peer, slot, axis); + long extent = block(peer, slot + 1, axis); + utils::check(origin >= 0 and extent >= 0 and origin <= global_shape_[axis] - extent, + "redistribution_plan {} block {} is out of bounds on axis {}: origin {}, size {}, global {}.", + label, peer, axis, origin, extent, global_shape_[axis]); + empty = empty or extent == 0; + if (not empty) { + utils::check(static_cast(extent) <= std::numeric_limits::max() / volume, + "redistribution_plan {} block-volume overflow on rank {}.", label, peer); + volume *= static_cast(extent); + } + } + return empty ? 0 : volume; + } + + void validate_exact_coverage(long slot, const char *label) const { + size_t global_volume = 1; + for (long axis = 0; axis < Rank; ++axis) { + utils::check(global_shape_[axis] >= 0, + "redistribution_plan global shape is negative on axis {}.", axis); + if (global_shape_[axis] == 0) { + global_volume = 0; + break; + } + utils::check(static_cast(global_shape_[axis]) <= + std::numeric_limits::max() / global_volume, + "redistribution_plan global-volume overflow."); + global_volume *= static_cast(global_shape_[axis]); + } + + size_t covered_volume = 0; + for (long peer = 0; peer < mpi_size_; ++peer) { + size_t volume = checked_block_volume(peer, slot, label); + utils::check(volume <= std::numeric_limits::max() - covered_volume, + "redistribution_plan {} covered-volume overflow.", label); + covered_volume += volume; + } + + for (long left = 0; left < mpi_size_; ++left) { + if (checked_block_volume(left, slot, label) == 0) continue; + for (long right = left + 1; right < mpi_size_; ++right) { + if (checked_block_volume(right, slot, label) == 0) continue; + bool overlaps = true; + for (long axis = 0; axis < Rank; ++axis) { + long left_begin = block(left, slot, axis); + long left_end = left_begin + block(left, slot + 1, axis); + long right_begin = block(right, slot, axis); + long right_end = right_begin + block(right, slot + 1, axis); + overlaps = overlaps and std::max(left_begin, right_begin) < std::min(left_end, right_end); + } + utils::check(not overlaps, + "redistribution_plan {} blocks overlap on ranks {} and {}.", + label, left, right); + } + } + utils::check(covered_volume == global_volume, + "redistribution_plan {} blocks leave holes: covered {} of {} elements.", + label, covered_volume, global_volume); + } + + std::array global_shape_{}; + std::array A_origin_{}; + std::array A_local_shape_{}; + std::array B_origin_{}; + std::array B_local_shape_{}; + std::array local_blocks_{}; + std::vector blocks_; + long mpi_size_ = 0; + long mpi_rank_ = 0; + size_t max_chunk_elements_ = 0; + size_t destination_element_size_ = 0; + bool collective_pack_too_large_ = true; +}; + +template +auto make_redistribution_plan(Arr1_t const& A, Arr2_t const& B, size_t max_chunk_elements = 0) { + static_assert(get_rank == get_rank, "Rank mismatch."); + return redistribution_plan>(A, B, max_chunk_elements); +} + +// Bounded-memory redistribution. Each ring round communicates with at most one +// source and one destination rank, and each rectangular overlap is tiled before +// it is packed. This avoids the full-size send and receive pack buffers used by +// redistribute_alltoallv for large distributed tensors. +template +void redistribute_streaming(Arr1_t& A, Arr2_t& B, + redistribution_plan> const& plan, + get_value_t a = 1, + get_value_t b = 0) { + using local_Arr1_t = typename std::decay_t::Array_t::regular_type; + using local_Arr2_t = typename std::decay_t::Array_t::regular_type; + static_assert(get_rank == get_rank, "Rank mismatch."); + plan.validate(A, B); + + constexpr long rank = get_rank; + auto b_one = get_value_t{1}; + auto comm = A.communicator(); + long mpi_size = plan.mpi_size(); + long mpi_rank = plan.mpi_rank(); + size_t max_chunk_elements = plan.max_chunk_elements(); + auto Aloc = A.local(); + auto Bloc = B.local(); + + if (b == get_value_t(0)) { + if (Bloc.size() > 0) ::nda::tensor::set(get_value_t(0), Bloc); + } else if (b != get_value_t(1)) { + if (Bloc.size() > 0) ::nda::tensor::scale(b, Bloc); + } + if (a == get_value_t(0)) return; + + if (mpi_size == 1) { + if constexpr (::nda::mem::have_device_compatible_addr_space) { + ::nda::tensor::add(a, Aloc, b_one, Bloc); + } else { + static_assert(::nda::mem::have_host_compatible_addr_space, "oh oh."); + Bloc += a * Aloc; + } + return; + } + + auto overlap_global = [&](bool local_A, long peer) { + std::vector<::nda::range> overlap(rank, ::nda::range(0)); + bool nonempty = true; + for (long r = 0; r < rank; ++r) { + long local_slot = local_A ? 0 : 2; + long peer_slot = local_A ? 2 : 0; + long i0 = plan.local_block(local_slot, r); + long i1 = i0 + plan.local_block(local_slot + 1, r); + long j0 = plan.block(peer, peer_slot, r); + long j1 = j0 + plan.block(peer, peer_slot + 1, r); + if (j1 > i0 and j0 < i1) { + overlap[r] = ::nda::range(std::max(i0, j0), std::min(i1, j1)); + } else { + nonempty = false; + break; + } + } + return std::make_pair(nonempty, std::move(overlap)); + }; + + auto make_tile_extent = [&](std::vector<::nda::range> const& overlap) { + std::vector extent(rank), tile_extent(rank, 1); + for (long r = 0; r < rank; ++r) + extent[r] = static_cast(overlap[r].last() - overlap[r].first()); + + size_t capacity = max_chunk_elements; + auto set_tile_extent = [&](long r) { + size_t take = std::min(static_cast(extent[r]), capacity); + tile_extent[r] = static_cast(std::max(1, take)); + capacity = std::max(1, capacity / static_cast(tile_extent[r])); + }; + if constexpr (local_Arr2_t::layout_t::is_stride_order_Fortran()) { + for (long r = 0; r < rank; ++r) set_tile_extent(r); + } else { + for (long r = rank - 1; r >= 0; --r) set_tile_extent(r); + } + return tile_extent; + }; + + auto number_of_tiles = [&](std::vector<::nda::range> const& overlap, + std::vector const& tile_extent) { + size_t ntiles = 1; + for (long r = 0; r < rank; ++r) { + size_t extent = static_cast(overlap[r].last() - overlap[r].first()); + size_t block = static_cast(tile_extent[r]); + size_t nblocks = extent / block + (extent % block != 0); + utils::check(nblocks == 0 or ntiles <= std::numeric_limits::max() / nblocks, + "Tile-count overflow in redistribute_streaming."); + ntiles *= nblocks; + } + return ntiles; + }; + + // Generate only the current tile. Keeping all tile ranges would make the + // metadata itself scale with tensor size when a small memory cap is used. + auto tile_at = [&](std::vector<::nda::range> const& overlap, + std::vector const& tile_extent, size_t tile_index) { + std::vector<::nda::range> tile(rank, ::nda::range(0)); + for (long r = rank - 1; r >= 0; --r) { + size_t extent = static_cast(overlap[r].last() - overlap[r].first()); + size_t block = static_cast(tile_extent[r]); + size_t nblocks = extent / block + (extent % block != 0); + size_t iblock = tile_index % nblocks; + tile_index /= nblocks; + long first = static_cast(overlap[r].first()) + static_cast(iblock * block); + tile[r] = ::nda::range(first, std::min(first + tile_extent[r], static_cast(overlap[r].last()))); + } + utils::check(tile_index == 0, "Tile index overflow in redistribute_streaming."); + return tile; + }; + + auto to_local_ranges = [&](std::vector<::nda::range> const& global_ranges, bool local_A) { + std::vector<::nda::range> local_ranges(rank, ::nda::range(0)); + long slot = local_A ? 0 : 2; + for (long r = 0; r < rank; ++r) { + long origin = plan.local_block(slot, r); + local_ranges[r] = ::nda::range(static_cast(global_ranges[r].first()) - origin, + static_cast(global_ranges[r].last()) - origin); + } + return local_ranges; + }; + + size_t sent_elements = 0; + size_t received_elements = 0; + + // Same-rank overlap needs no MPI scratch. + { + auto [has_self_A, self_global_A] = overlap_global(true, mpi_rank); + auto [has_self_B, self_global_B] = overlap_global(false, mpi_rank); + utils::check(has_self_A == has_self_B, "Logic error in redistribute_streaming self overlap."); + if (has_self_A) { + auto A_self = detail::get_sub_matrix(Aloc, to_local_ranges(self_global_A, true)); + auto B_self = detail::get_sub_matrix(Bloc, to_local_ranges(self_global_B, false)); + utils::check(A_self.size() == B_self.size(), "Self-overlap size mismatch in redistribute_streaming."); + if constexpr (::nda::mem::have_device_compatible_addr_space) { + ::nda::tensor::add(a, A_self, b_one, B_self); + } else { + static_assert(::nda::mem::have_host_compatible_addr_space, "oh oh."); + B_self += a * A_self; + } + sent_elements += A_self.size(); + received_elements += B_self.size(); + } + } + + constexpr int redistribute_tag = 0; + for (long step = 1; step < mpi_size; ++step) { + long destination = (mpi_rank + step) % mpi_size; + long source = (mpi_rank - step + mpi_size) % mpi_size; + auto [has_send, send_global] = overlap_global(true, destination); + auto [has_recv, recv_global] = overlap_global(false, source); + auto send_tile_extent = has_send ? make_tile_extent(send_global) : std::vector{}; + auto recv_tile_extent = has_recv ? make_tile_extent(recv_global) : std::vector{}; + size_t send_ntiles = has_send ? number_of_tiles(send_global, send_tile_extent) : 0; + size_t recv_ntiles = has_recv ? number_of_tiles(recv_global, recv_tile_extent) : 0; + size_t nrounds = std::max(send_ntiles, recv_ntiles); + + for (size_t itile = 0; itile < nrounds; ++itile) { + std::optional send_buffer; + std::optional recv_buffer; + std::vector<::nda::range> send_tile; + std::vector<::nda::range> recv_tile; + if (itile < send_ntiles) { + send_tile = tile_at(send_global, send_tile_extent, itile); + auto A_tile = detail::get_sub_matrix(Aloc, to_local_ranges(send_tile, true)); + send_buffer.emplace(A_tile); + utils::check(send_buffer->size() <= max_chunk_elements, "Send tile exceeds redistribute_streaming chunk limit."); + } + if (itile < recv_ntiles) { + recv_tile = tile_at(recv_global, recv_tile_extent, itile); + auto B_tile = detail::get_sub_matrix(Bloc, to_local_ranges(recv_tile, false)); + recv_buffer.emplace(B_tile.shape()); + utils::check(recv_buffer->size() <= max_chunk_elements, "Receive tile exceeds redistribute_streaming chunk limit."); + } + + boost::mpi3::request recv_request; + boost::mpi3::request send_request; + if (recv_buffer) + recv_request = comm->ireceive_n(recv_buffer->data(), static_cast(recv_buffer->size()), source, redistribute_tag); + if (send_buffer) + send_request = comm->isend_n(send_buffer->data(), static_cast(send_buffer->size()), destination, redistribute_tag); + + if (recv_buffer) { + recv_request.wait(); + auto B_tile = detail::get_sub_matrix(Bloc, to_local_ranges(recv_tile, false)); + if constexpr (::nda::mem::have_device_compatible_addr_space) { + ::nda::tensor::add(a, *recv_buffer, b_one, B_tile); + } else { + static_assert(::nda::mem::have_host_compatible_addr_space, "oh oh."); + B_tile += a * (*recv_buffer); + } + received_elements += recv_buffer->size(); + } + if (send_buffer) { + send_request.wait(); + sent_elements += send_buffer->size(); + } + } + } + + utils::check(sent_elements == static_cast(Aloc.size()), + "redistribute_streaming did not cover the full local source: {} != {}", sent_elements, Aloc.size()); + utils::check(received_elements == static_cast(Bloc.size()), + "redistribute_streaming did not cover the full local destination: {} != {}", received_elements, Bloc.size()); +} + +// Backward-compatible one-shot entry point. Repeated callers should construct +// one redistribution_plan and use the overload above. +template +void redistribute_streaming(Arr1_t& A, Arr2_t& B, + get_value_t a = 1, + get_value_t b = 0, + size_t max_chunk_elements = 0) { + auto plan = make_redistribution_plan(A, B, max_chunk_elements); + redistribute_streaming(A, B, plan, a, b); +} + +// Reuse cached layout metadata while retaining the small-layout all-to-all-v +// fast path for regular arrays. Irregular block collections always use the +// metadata-driven streaming path because they intentionally expose no regular +// grid contract. +template +void redistribute(Arr1_t& A, Arr2_t& B, + redistribution_plan> const& plan, + get_value_t a = 1, + get_value_t b = 0) { + plan.validate(A, B); + if constexpr (DistributedArray and DistributedArray) { + if (a == get_value_t(1) and b == get_value_t(0) and + not plan.collective_pack_too_large()) { + redistribute_alltoallv(A, B, a, b); + return; + } + } + redistribute_streaming(A, B, plan, a, b); +} + +template void redistribute(Arr1_t& A, Arr2_t& B, get_value_t a = 1, get_value_t b = 0) { switch(Alg) { + case 0: { + // Keep the optimized collective path for small arrays. For large arrays, + // cap explicit communication scratch instead of allocating full local + // send and receive packs. Select collectively from the largest actual + // local pack requirement so every rank takes the same path. + using value_B_t = typename std::decay_t::Array_t::value_type; + size_t local_pack_elements = static_cast(A.local().size()); + utils::check(local_pack_elements <= std::numeric_limits::max() - static_cast(B.local().size()), + "Local pack-size overflow in redistribute."); + local_pack_elements += static_cast(B.local().size()); + size_t max_pack_elements = A.communicator()->all_reduce_value(local_pack_elements, boost::mpi3::max<>{}); + constexpr size_t collective_pack_budget = size_t{64} * 1024 * 1024; + bool collective_pack_too_large = max_pack_elements > collective_pack_budget / sizeof(value_B_t); + // redistribute_alltoallv currently only implements B = A. Route scaled + // redistributions through the streaming implementation as well. + if (a != get_value_t(1) or b != get_value_t(0) or + collective_pack_too_large) + redistribute_streaming(A, B, a, b); + else + redistribute_alltoallv(A, B, a, b); + break; + } case 1: redistribute_standard(A, B, a, b); break; @@ -875,6 +1304,11 @@ void redistribute(Arr1_t& A, Arr2_t& B, get_value_t a = 1, get_value_t; + auto A = make_distributed_array(world, {size, 1}, {size + 5, 3 * size + 7}); + auto B = make_distributed_array(world, {1, size}, {size + 5, 3 * size + 7}); + auto Aloc = A.local(); + auto origin = A.origin(); + auto gshape = A.global_shape(); + for (long i = 0; i < Aloc.shape(0); ++i) + for (long j = 0; j < Aloc.shape(1); ++j) + Aloc(i, j) = static_cast((origin[0] + i) * gshape[1] + origin[1] + j); + + redistribute_streaming(A, B, 1.0, 0.0, 7); + + auto Bloc = B.local(); + origin = B.origin(); + for (long i = 0; i < Bloc.shape(0); ++i) + for (long j = 0; j < Bloc.shape(1); ++j) + REQUIRE(Bloc(i, j) == static_cast((origin[0] + i) * gshape[1] + origin[1] + j)); + + // A reusable plan must observe new source values while retaining the fixed + // layouts. The adaptive plan overload retains the collective fast path for + // this small regular layout. + auto plan = make_redistribution_plan(A, B, 7); + Aloc += 1000.0; + Bloc = -1.0; + redistribute(A, B, plan); + for (long i = 0; i < Bloc.shape(0); ++i) + for (long j = 0; j < Bloc.shape(1); ++j) + REQUIRE(Bloc(i, j) == 1000.0 + + static_cast((origin[0] + i) * gshape[1] + origin[1] + j)); + } + + // Exercise the truthful irregular-block layout used for a selected-axis + // transfer: only an active subset of ranks owns source blocks, while the + // destination is regularly distributed over every rank. + { + using larray = nda::array; + long active_ranks = std::max(1L, (size + 1) / 2); + long global_rows = 2 * active_ranks + 1; + bool active = world.rank() < active_ranks; + long base_rows = global_rows / active_ranks; + long remainder = global_rows % active_ranks; + long local_rows = active ? base_rows + (world.rank() < remainder ? 1 : 0) : 0; + long row_origin = active ? world.rank() * base_rows + std::min(long(world.rank()), remainder) : 0; + std::array local_shape = active ? std::array{local_rows, 5} : + std::array{0, 0}; + larray local(local_shape); + for (long i = 0; i < local_rows; ++i) + for (long j = 0; j < 5; ++j) + local(i, j) = static_cast((row_origin + i) * 5 + j); + + memory::irregular_block_darray_t selected( + std::addressof(world), {global_rows, 5}, {row_origin, 0}, std::move(local)); + auto B = make_distributed_array(world, {size, 1}, {global_rows, 5}); + auto plan = make_redistribution_plan(selected, B, 7); + plan.validate_source_coverage(); + plan.validate_destination_coverage(); + redistribute_streaming(selected, B, plan); + + auto Bloc = B.local(); + auto origin = B.origin(); + for (long i = 0; i < Bloc.shape(0); ++i) + for (long j = 0; j < Bloc.shape(1); ++j) + REQUIRE(Bloc(i, j) == static_cast((origin[0] + i) * 5 + origin[1] + j)); + } + + // Exercise a higher-rank C-layout tensor and the general B = a*A + b*B path. + { + using larray = nda::array; + auto A = make_distributed_array(world, {size, 1, 1, 1}, {size + 3, size + 5, 3, 2}); + auto B = make_distributed_array(world, {1, size, 1, 1}, {size + 3, size + 5, 3, 2}); + auto Aloc = A.local(); + auto Aorigin = A.origin(); + auto gshape = A.global_shape(); + for (long i = 0; i < Aloc.shape(0); ++i) + for (long j = 0; j < Aloc.shape(1); ++j) + for (long k = 0; k < Aloc.shape(2); ++k) + for (long l = 0; l < Aloc.shape(3); ++l) { + double x = static_cast((((Aorigin[0] + i) * gshape[1] + Aorigin[1] + j) * + gshape[2] + Aorigin[2] + k) * gshape[3] + Aorigin[3] + l); + Aloc(i, j, k, l) = ComplexType{x, -0.5 * x}; + } + + auto Bloc = B.local(); + auto Borigin = B.origin(); + for (long i = 0; i < Bloc.shape(0); ++i) + for (long j = 0; j < Bloc.shape(1); ++j) + for (long k = 0; k < Bloc.shape(2); ++k) + for (long l = 0; l < Bloc.shape(3); ++l) { + double x = static_cast((((Borigin[0] + i) * gshape[1] + Borigin[1] + j) * + gshape[2] + Borigin[2] + k) * gshape[3] + Borigin[3] + l); + Bloc(i, j, k, l) = ComplexType{1.0 + 0.25 * x, 2.0 - 0.125 * x}; + } + + ComplexType a{0.5, -1.0}; + ComplexType b{-0.25, 0.5}; + // Non-default coefficients must make the default dispatcher select the + // streaming implementation, even though this test tensor is small. + redistribute(A, B, a, b); + + for (long i = 0; i < Bloc.shape(0); ++i) + for (long j = 0; j < Bloc.shape(1); ++j) + for (long k = 0; k < Bloc.shape(2); ++k) + for (long l = 0; l < Bloc.shape(3); ++l) { + double x = static_cast((((Borigin[0] + i) * gshape[1] + Borigin[1] + j) * + gshape[2] + Borigin[2] + k) * gshape[3] + Borigin[3] + l); + ComplexType expected = a * ComplexType{x, -0.5 * x} + + b * ComplexType{1.0 + 0.25 * x, 2.0 - 0.125 * x}; + REQUIRE(std::abs(Bloc(i, j, k, l) - expected) < 1.0e-12); + } + } +} + } // bdft_tests diff --git a/src/numerics/imag_axes_ft/IAFT.icc b/src/numerics/imag_axes_ft/IAFT.icc index 5fbeee7c..5b6cc6fc 100644 --- a/src/numerics/imag_axes_ft/IAFT.icc +++ b/src/numerics/imag_axes_ft/IAFT.icc @@ -174,7 +174,10 @@ void IAFT::w_to_tau_partial(ndaArray_A&& Xw_i, ndaArray_B&& X_ti, stats_e stats, auto Xw_i_1D = nda::reshape(Xw_i, shape_t<1>{dim1}); auto T_t = (stats == fermion)? nda::make_regular(Ttw_ff()(nda::range::all, iwn)) : nda::make_regular(Ttw_bb()(nda::range::all, iwn)); - X_ti_2D += nda::blas::outer_product(T_t, Xw_i_1D); + // Accumulate the rank-1 update directly into the output. Constructing an + // outer_product here materializes a full (nt, dim1) temporary, which defeats + // the bounded-memory purpose of the partial transform for large tensors. + nda::blas::ger(ComplexType{1.0, 0.0}, T_t, Xw_i_1D, X_ti_2D); } /** diff --git a/src/numerics/iter_scf/diis/com_diis_residual.h b/src/numerics/iter_scf/diis/com_diis_residual.h index 94c7408c..52211310 100644 --- a/src/numerics/iter_scf/diis/com_diis_residual.h +++ b/src/numerics/iter_scf/diis/com_diis_residual.h @@ -110,14 +110,17 @@ class com_diis_residual : public diis_residual { utils::check(com_initialized, "DIIS commutator residual is not initialized"); upload_g_mu(); // TODO if it hasn't been supplied externally // Warning! Sigma here is in tau! - FockSigma x_last = current_state->get(); + const FockSigma& x_last = current_state->get_ref(); Array_5D C_t; commutator_t(C_t, FT, G_incoming, x_last, mu, _S, _H0); + G_incoming = Array_5D{}; + iter = -1; auto Fz = x_last.get_fock(); Fz() = 0; - res.set_fock_sigma(Fz, C_t); + res.set_fock_sigma(std::move(Fz), std::move(C_t)); + res.set_mu(0.0); return true; } diff --git a/src/numerics/iter_scf/diis/diis_alg.hpp b/src/numerics/iter_scf/diis/diis_alg.hpp index cb698e83..294cd0eb 100644 --- a/src/numerics/iter_scf/diis/diis_alg.hpp +++ b/src/numerics/iter_scf/diis/diis_alg.hpp @@ -73,11 +73,11 @@ template // (no extrapolation as well) bool grow_xvsp_only; - auto get_extrapolated_state() { + const Vector& get_extrapolated_state() const { if(extrapolated_state == nullptr) { APP_ABORT("DIIS state is not initialized! ABORT!"); } - return extrapolated_state->get(); + return extrapolated_state->get_ref(); } private: @@ -158,7 +158,7 @@ template * return 1 if extrapolation was performed * 0 if no extrapolation (just growing the subspace) */ - int next_step(const Vector& new_vec) { + int next_step(Vector new_vec) { if (x_vsp->size() == 0 || grow_xvsp_only) { app_log(2, diis_str + "Growing vector subspace only. No extrapolation.\n"); x_vsp->add_to_vspace(new_vec); // growing vector space @@ -170,15 +170,16 @@ template // Normal execution app_log(2, diis_str + "Growing vector and residual subspaces for DIIS\n"); // Fill the extrapolated state with the current vector for residual computation - extrapolated_state->put(new_vec); - Vector res; - if(! residual->get_diis_residual(res) ) { - APP_ABORT(diis_str + "Could not get residual!!! ABORT!"); + extrapolated_state->put(std::move(new_vec)); + { + Vector res; + if(! residual->get_diis_residual(res) ) { + APP_ABORT(diis_str + "Could not get residual!!! ABORT!"); + } + update_overlaps(res); // the overlap with res is added in any case... + res_vsp->add_to_vspace(res); // growing residual space } - update_overlaps(res); // the overlap with res is added in any case... - - res_vsp->add_to_vspace(res); // growing residual space - x_vsp->add_to_vspace(new_vec); // growing vector space + x_vsp->add_to_vspace(extrapolated_state->get_ref()); // growing vector space } else { // The subspace is already of the maximum size app_log(2, diis_str + "Reached maximum subspace -> the first vector will be kicked out of the subspace.\n"); @@ -188,14 +189,16 @@ template purge_overlap(0); // purge overlap matrix of residuals // Fill the extrapolated state with the current vector for residual computation - extrapolated_state->put(new_vec); - Vector res; - if(! residual->get_diis_residual(res) ) { - APP_ABORT(diis_str + "Could not get residual!!! ABORT!"); + extrapolated_state->put(std::move(new_vec)); + { + Vector res; + if(! residual->get_diis_residual(res) ) { + APP_ABORT(diis_str + "Could not get residual!!! ABORT!"); + } + update_overlaps(res); + res_vsp->add_to_vspace(res); // growing residual space } - update_overlaps(res); - res_vsp->add_to_vspace(res); // growing residual space - x_vsp->add_to_vspace(new_vec); // growing vector space + x_vsp->add_to_vspace(extrapolated_state->get_ref()); // growing vector space } if (extrap && (res_vsp->size() > 1) ) { @@ -213,7 +216,7 @@ template // build extrapolated vector Vector result = x_vsp->make_linear_comb(m_C); app_log(2, ""); - extrapolated_state->put(result); // update extrapolated state + extrapolated_state->put(std::move(result)); // update extrapolated state return 1; } else { diff --git a/src/numerics/iter_scf/diis/diis_t.hpp b/src/numerics/iter_scf/diis/diis_t.hpp index 19c290b2..62ec6739 100644 --- a/src/numerics/iter_scf/diis/diis_t.hpp +++ b/src/numerics/iter_scf/diis/diis_t.hpp @@ -107,7 +107,7 @@ namespace iter_scf { comFS_residual.initialize(&extrapolated_state, S, H0, FT, mbpt_output); // providing non-owning pointers to DIIS kernel as well as the starting state d_alg.init(&extrapolated_state, &comFS_residual, &x_vsp, &res_vsp, - max_subsp_size, true, FockSigma(F, Sigma, mu)); + max_subsp_size, true, extrapolated_state.get_ref()); initialized = true; } @@ -237,19 +237,22 @@ namespace iter_scf { // DO DIIS d_alg.extrap = true; d_alg.grow_xvsp_only = false; - FockSigma fs(F, Sigma, get_mu()); int is_extrapolated = d_alg.next_step(FockSigma(F, Sigma, get_mu())); if(is_extrapolated != 0) { - auto Fdiff = nda::make_regular(F - d_alg.get_extrapolated_state().get_fock()); - auto Sdiff = nda::make_regular(Sigma - d_alg.get_extrapolated_state().get_sigma()); + const auto& extrapolated = d_alg.get_extrapolated_state(); + const auto& extrapolated_F = extrapolated.get_fock(); + const auto& extrapolated_Sigma = extrapolated.get_sigma(); + auto Fdiff = nda::make_regular(F - extrapolated_F); auto Fmax_iter = max_element(Fdiff.data(), Fdiff.data()+Fdiff.size(), [](auto a, auto b) { return std::abs(a) < std::abs(b); }); - auto Smax_iter = max_element(Sdiff.data(), Sdiff.data()+Sdiff.size(), - [](auto a, auto b) { return std::abs(a) < std::abs(b); }); - F = d_alg.get_extrapolated_state().get_fock(); - Sigma = d_alg.get_extrapolated_state().get_sigma(); - - return std::array{std::abs(*Fmax_iter), std::abs(*Smax_iter)}; + double Smax_iter = 0.0; + nda::for_each(Sigma.shape(), [&](auto... i) { + Smax_iter = std::max(Smax_iter, std::abs(Sigma(i...) - extrapolated_Sigma(i...))); + }); + F = extrapolated_F; + Sigma = extrapolated_Sigma; + + return std::array{std::abs(*Fmax_iter), Smax_iter}; } else { // No DIIS extrapolation has been applied diff --git a/src/numerics/iter_scf/diis/state.h b/src/numerics/iter_scf/diis/state.h index 9b279151..cfe54f79 100644 --- a/src/numerics/iter_scf/diis/state.h +++ b/src/numerics/iter_scf/diis/state.h @@ -57,6 +57,11 @@ class opt_state { return x; } + const Vector& get_ref() const { + utils::check(inited, "State is not initialized"); + return x; + } + void set(const Vector x_) {x = x_; inited = true;} void set(const Vector& x_) {x = x_; inited = true;} void set(Vector&& x_) noexcept {x = std::move(x_); inited = true;} diff --git a/src/numerics/iter_scf/diis/vspace_fock_sigma.hpp b/src/numerics/iter_scf/diis/vspace_fock_sigma.hpp index 1085a9f0..51edc5ea 100644 --- a/src/numerics/iter_scf/diis/vspace_fock_sigma.hpp +++ b/src/numerics/iter_scf/diis/vspace_fock_sigma.hpp @@ -46,6 +46,7 @@ class FockSigma { inited_S = true; inited_mu = true; } + FockSigma(FockSigma&&) noexcept = default; FockSigma(const Array_4D& Fock_, const Array_5D& Sigma_, const double mu_) : _Fock(Fock_), _Sigma(Sigma_), _mu(mu_) { @@ -63,6 +64,7 @@ class FockSigma { inited_mu = true; return *this; } + FockSigma& operator=(FockSigma&&) noexcept = default; ComplexType dot_prod(const FockSigma& rhs) const { utils::check(inited_F, "FockSigma: Fock matrix is not initialized"); @@ -73,25 +75,15 @@ class FockSigma { auto vec_F= nda::reshape(_Fock, std::array{Fdim}); auto vec_S= nda::reshape(_Sigma, std::array{Sdim}); */ - auto matvec_F= nda::reshape(_Fock, std::array{Fdim, 1}); - auto matvec_S= nda::reshape(_Sigma, std::array{Sdim, 1}); - - auto rFock = rhs.get_fock(); - auto rSigma = rhs.get_sigma(); + const auto& rFock = rhs.get_fock(); + const auto& rSigma = rhs.get_sigma(); size_t rFdim = std::reduce(rFock.shape().begin(), rFock.shape().end(), 1, std::multiplies()); size_t rSdim = std::reduce(rSigma.shape().begin(), rSigma.shape().end(), 1, std::multiplies()); -/* - auto vec_rF= nda::reshape(rFock, std::array{rFdim}); - auto vec_rS= nda::reshape(rSigma, std::array{rSdim}); -*/ - auto matvec_rF= nda::reshape(rFock, std::array{rFdim, 1}); - auto matvec_rS= nda::reshape(rSigma, std::array{rSdim, 1}); - //return nda::blas::dotc(vec_F,vec_rF) + nda::blas::dotc(vec_S,vec_rS); - nda::array res1(1,1); - nda::array res2(1,1); - nda::blas::gemm(nda::make_regular(nda::conj(nda::transpose(matvec_F))), matvec_rF, res1); - nda::blas::gemm(nda::make_regular(nda::conj(nda::transpose(matvec_S))), matvec_rS, res2); - return res1(0,0) + res2(0,0); + auto vec_rF = nda::reshape(rFock, std::array{static_cast(rFdim)}); + auto vec_rS = nda::reshape(rSigma, std::array{static_cast(rSdim)}); + auto vec_F = nda::reshape(_Fock, std::array{static_cast(Fdim)}); + auto vec_S = nda::reshape(_Sigma, std::array{static_cast(Sdim)}); + return nda::blas::dotc(vec_F, vec_rF) + nda::blas::dotc(vec_S, vec_rS); } const Array_4D& get_fock() const { @@ -112,14 +104,26 @@ class FockSigma { _Fock = F_; inited_F = true; } + void set_fock(Array_4D&& F_) { + _Fock = std::move(F_); + inited_F = true; + } void set_sigma(Array_5D& S_) { _Sigma = S_; inited_S = true; } + void set_sigma(Array_5D&& S_) { + _Sigma = std::move(S_); + inited_S = true; + } void set_fock_sigma(Array_4D& F_, Array_5D& S_) { set_fock(F_); set_sigma(S_); } + void set_fock_sigma(Array_4D&& F_, Array_5D&& S_) { + set_fock(std::move(F_)); + set_sigma(std::move(S_)); + } void set_zero() { _Fock() = 0; @@ -157,8 +161,12 @@ class FockSigma { void add(FockSigma&& a, ComplexType c) { utils::check(inited_F, "FockSigma: Fock matrix is not initialized"); utils::check(inited_S, "FockSigma: Sigma is not initialized"); - _Fock += c * a.get_fock(); - _Sigma += c * a.get_sigma(); + const auto& aFock = a.get_fock(); + const auto& aSigma = a.get_sigma(); + utils::check(_Fock.shape() == aFock.shape(), "FockSigma::add: incompatible Fock shapes"); + utils::check(_Sigma.shape() == aSigma.shape(), "FockSigma::add: incompatible Sigma shapes"); + for (size_t i = 0; i < _Fock.size(); ++i) _Fock.data()[i] += c * aFock.data()[i]; + for (size_t i = 0; i < _Sigma.size(); ++i) _Sigma.data()[i] += c * aSigma.data()[i]; } void read_from_file(std::string filename, const size_t vec_number) { @@ -215,48 +223,44 @@ void commutator_t(Array_G& C_t, const imag_axes_ft::IAFT *FT, size_t nk = G_t.shape()[2]; size_t nao = G_t.shape()[3]; size_t nw = FT->nw_f(); - nda::array G_w(nw,ns,nk,nao,nao); - nda::array Sigma_w(nw,ns,nk,nao,nao); - // G_w is filled - FT->tau_to_w(G_t, G_w, imag_axes_ft::fermion); - // Sigma_t is filled - auto Sigma_t = FS_t.get_sigma(); - auto Fock = FS_t.get_fock(); - // Sigma_w is filled - FT->tau_to_w(Sigma_t, Sigma_w, imag_axes_ft::fermion); - - nda::array Dm(ns,nk,nao,nao); - FT->tau_to_beta(G_t, Dm); - - nda::array C_w(nw, ns, nk, nao, nao); - C_w () = 0; + const auto& Sigma_t = FS_t.get_sigma(); + const auto& Fock = FS_t.get_fock(); C_t = nda::array(nt,ns,nk,nao,nao); // To make sure an array of appropriate size is ready C_t () = 0; + // Stream the frequency axis. The former implementation materialized full + // G(w), Sigma(w), and C(w) arrays simultaneously on the global root. + nda::array G_wskij(ns,nk,nao,nao); + nda::array Sigma_wskij(ns,nk,nao,nao); + nda::array C_wskij(ns,nk,nao,nao); + nda::array I1(nao, nao); nda::array I2(nao, nao); - for(size_t iw = 0; iw < nw; iw++) - for(size_t s = 0; s < ns; s++) - for(size_t k = 0; k < nk; k++) { - long wn = FT->wn_mesh()(iw); - ComplexType omega_mu = FT->omega(wn) + mu; - auto S_sk = S(s,k,all,all); - auto F_sk = Fock(s,k,all,all); - auto H0_sk = H0(s,k,all,all); - auto G_wsk = G_w(iw,s,k,all,all); - auto Sigma_wsk = Sigma_w(iw,s,k,all,all); + app_log(2, "DIIS: Streaming commutator residual over {} frequencies", nw); + for(size_t iw = 0; iw < nw; iw++) { + FT->tau_to_w(G_t, G_wskij, imag_axes_ft::fermion, iw); + FT->tau_to_w(Sigma_t, Sigma_wskij, imag_axes_ft::fermion, iw); + for(size_t s = 0; s < ns; s++) + for(size_t k = 0; k < nk; k++) { + long wn = FT->wn_mesh()(iw); + ComplexType omega_mu = FT->omega(wn) + mu; + auto S_sk = S(s,k,all,all); + auto F_sk = Fock(s,k,all,all); + auto H0_sk = H0(s,k,all,all); + auto G_wsk = G_wskij(s,k,all,all); + auto Sigma_wsk = Sigma_wskij(s,k,all,all); - nda::array G0inv_Sigma_wsk = nda::make_regular(omega_mu * S_sk - H0_sk - F_sk - Sigma_wsk); - nda::array_view C_wsk = C_w(iw,s,k,all,all); - I1() = 0; - I2() = 0; - nda::blas::gemm(G_wsk, G0inv_Sigma_wsk, I1); - nda::blas::gemm(G0inv_Sigma_wsk, G_wsk, I2); - C_wsk = nda::make_regular(I1 - I2); + nda::array G0inv_Sigma_wsk = nda::make_regular(omega_mu * S_sk - H0_sk - F_sk - Sigma_wsk); + auto C_wsk = C_wskij(s,k,all,all); + nda::blas::gemm(G_wsk, G0inv_Sigma_wsk, I1); + nda::blas::gemm(G0inv_Sigma_wsk, G_wsk, I2); + C_wsk = I1 - I2; + } + FT->w_to_tau_partial(C_wskij, C_t, imag_axes_ft::fermion, iw); + if ((iw + 1) == nw || (iw + 1) % std::max(size_t{1}, nw / 10) == 0) + app_log(2, "DIIS: Commutator frequencies {}/{}", iw + 1, nw); } - - FT->w_to_tau(C_w, C_t, imag_axes_ft::fermion); } diff --git a/src/numerics/shared_array/nda.hpp b/src/numerics/shared_array/nda.hpp index 28caed9b..2e911d8e 100644 --- a/src/numerics/shared_array/nda.hpp +++ b/src/numerics/shared_array/nda.hpp @@ -22,6 +22,8 @@ #ifndef NUMERICS_SHARED_ARRAY_NDA_HPP #define NUMERICS_SHARED_ARRAY_NDA_HPP +#include + #include "configuration.hpp" #include "mpi3/communicator.hpp" #include "mpi3/shared_window.hpp" @@ -151,14 +153,20 @@ namespace math { node_sync(); } - void all_reduce() { + static constexpr size_t default_all_reduce_chunk_bytes = size_t{256} * 1024 * 1024; + + void all_reduce(size_t max_chunk_bytes = default_all_reduce_chunk_bytes) { node_sync(); if (_node_comm->root()) { - // split all_reduce() to avoid mpi count overflow - for (size_t shift=0; shift<_size; shift+=size_t(1e9)) { + // Bound both the MPI int count and the message size that may drive + // implementation-internal collective scratch. + size_t max_count = std::max( + 1, std::min(max_chunk_bytes / sizeof(value_type), + static_cast(std::numeric_limits::max()))); + for (size_t shift = 0; shift < _size; shift += max_count) { value_type *start = (value_type*)_win->base(0) + shift; - size_t count = (shift+size_t(1e9) < _size)? size_t(1e9) : _size-shift; - _internode_comm->all_reduce_in_place_n(start, count, std::plus<>{}); + size_t count = std::min(max_count, static_cast(_size) - shift); + _internode_comm->all_reduce_in_place_n(start, static_cast(count), std::plus<>{}); } } node_sync(); diff --git a/src/numerics/shared_array/tests/test_shared.cpp b/src/numerics/shared_array/tests/test_shared.cpp index b778c576..2cd4bd83 100644 --- a/src/numerics/shared_array/tests/test_shared.cpp +++ b/src/numerics/shared_array/tests/test_shared.cpp @@ -83,4 +83,34 @@ TEST_CASE("distributed_shared_nda", "[math]") { array.node_sync(); } +TEST_CASE("shared_nda_chunked_all_reduce", "[math][shared][all_reduce]") { + auto world = mpi3::environment::get_world_instance(); + auto node_comm = world.split_shared(); + int node_size = node_comm.size(); + int color = world.rank() % node_size; + int key = world.rank() / node_size; + auto internode_comm = world.split(color, key); + + auto check_type = [&]() { + using Array_view_t = nda::array_view; + auto array = make_shared_array(world, internode_comm, node_comm, shape_t<1>{17}); + if (node_comm.root()) { + T value = T(internode_comm.rank() + 1); + array.local() = value; + } + array.node_sync(); + + // Force the 17 elements through a 3,3,3,3,3,2 chunk sequence. + array.all_reduce(3 * sizeof(T)); + + double expected = 0.5 * internode_comm.size() * (internode_comm.size() + 1); + auto local = array.local(); + for (long i = 0; i < local.size(); ++i) + REQUIRE(std::abs(local(i) - T(expected)) < 1.0e-12); + }; + + check_type.template operator()(); + check_type.template operator()(); +} + } // bdft_tests