[WIP] Reduce extra memory usage in scGW and GW+EDMFT workflows

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<local_Array_4D_t>(
           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<ComplexType, 2> buffer_ib(nbnd, nbnd);
-
-        auto sDelta_tskij = make_shared_array<Array_view_5D_t>(*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<size_t>(1, delta_tile_bytes / sizeof(ComplexType));
+        long band_tile = std::min<long>(nbnd, static_cast<long>(std::sqrt(static_cast<double>(max_tile_elements))));
+        band_tile = std::max<long>(1, band_tile);
+        size_t band_tile_elements = static_cast<size_t>(band_tile) * static_cast<size_t>(band_tile);
+        long tsk_tile = static_cast<long>(std::max<size_t>(1, max_tile_elements / band_tile_elements));
+        size_t ntsk = static_cast<size_t>(nts) * static_cast<size_t>(ns) * static_cast<size_t>(nkpts);
+
+        using Array_view_3D_t = nda::array_view<ComplexType, 3>;
+        auto sDelta_tile = make_shared_array<Array_view_3D_t>(
+            *mpi, {tsk_tile, band_tile, band_tile});
+        nda::array<ComplexType, 2> buffer_ib(band_tile, nbnd);
+        decltype(nda::range::all) all;
+
+        for (size_t tsk0 = 0; tsk0 < ntsk; tsk0 += static_cast<size_t>(tsk_tile)) {
+          size_t tsk1 = std::min(ntsk, tsk0 + static_cast<size_t>(tsk_tile));
+          long tsk_count = static_cast<long>(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<size_t>(rank) + static_cast<size_t>(size) -
+                                    tsk0 % static_cast<size_t>(size)) % static_cast<size_t>(size);
+              for (size_t tsk = tsk0 + rank_offset; tsk < tsk1; tsk += static_cast<size_t>(size)) {
+
+                // tsk = it * ns*nkpts + is * nkpts + ik
+                size_t it = tsk / (static_cast<size_t>(ns) * static_cast<size_t>(nkpts));
+                size_t is = (tsk / static_cast<size_t>(nkpts)) % static_cast<size_t>(ns);
+                size_t ik = tsk % static_cast<size_t>(nkpts);
+                size_t it_pos = (it < static_cast<size_t>(nt_half))? it : static_cast<size_t>(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<long>(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<size_t>(itsk);
+                  size_t it = tsk / (static_cast<size_t>(ns) * static_cast<size_t>(nkpts));
+                  size_t is = (tsk / static_cast<size_t>(nkpts)) % static_cast<size_t>(ns);
+                  size_t ik = tsk % static_cast<size_t>(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);

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<ComplexType, 5> G_t(nt, ns, nk, nao, nao);
+    nda::array<ComplexType, 5> Sigma_t(nt, ns, nk, nao, nao);
+    nda::array<ComplexType, 4> F(ns, nk, nao, nao);
+    nda::array<ComplexType, 4> S(ns, nk, nao, nao);
+    nda::array<ComplexType, 4> 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<ComplexType, 5> 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<ComplexType, 5> G_w(nw, ns, nk, nao, nao);
+    nda::array<ComplexType, 5> Sigma_w(nw, ns, nk, nao, nao);
+    nda::array<ComplexType, 5> C_w(nw, ns, nk, nao, nao);
+    nda::array<ComplexType, 5> 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<ComplexType, 2> left(nao, nao);
+      nda::array<ComplexType, 2> 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();
 

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.