Optimize diagonal FDM density rotations

Author: rokzitko

c++/dmnrg.hpp

@@ -10,6 +10,9 @@
 #include <stdexcept>
 #include <string>
 #include <cmath>
+#include <chrono>
+#include <limits>
+#include <utility>
 
 #include "operators.hpp"
 #include "symmetry.hpp"
@@ -35,6 +38,98 @@ void check_trace_rho(const DensMatElements<S> &m, MF mult, const double ref_valu
     throw std::runtime_error(fmt::format("check_trace_rho() failed, tr={}, ref_value={}", tr, ref_value));
 }
 
+class DmnrgProfile {
+ public:
+  using clock = std::chrono::steady_clock;
+  using time_point = clock::time_point;
+  enum class Contribution { generic, kk, dd };
+
+  DmnrgProfile(std::string label, const size_t N, const bool enabled) : label_(std::move(label)), N_(N), enabled_(enabled) {}
+
+  [[nodiscard]] bool enabled() const noexcept { return enabled_; }
+  [[nodiscard]] static time_point now() { return clock::now(); }
+  [[nodiscard]] static double elapsed(const time_point start) { return std::chrono::duration<double>(clock::now() - start).count(); }
+
+  void add_diag_load(const double seconds) { diag_load_seconds_ += seconds; }
+  void add_init(const double seconds) { init_seconds_ += seconds; }
+  void add_trace(const double seconds) { trace_seconds_ += seconds; }
+  void add_save(const double seconds) { save_seconds_ += seconds; }
+  void add_missing_rho() { ++missing_rho_; }
+  void add_missing_diag() { ++missing_diag_; }
+  void add_zero_dim() { ++zero_dim_; }
+  void add_zero_omega() { ++zero_omega_; }
+  void add_zero_rmax() { ++zero_rmax_; }
+
+  void add_rotate(const size_t q, const size_t r, const size_t element_size, const double seconds,
+                  const Contribution contribution = Contribution::generic, const bool diagonal = false) {
+    ++contributions_;
+    rotate_seconds_ += seconds;
+    if (contribution == Contribution::kk) {
+      ++kk_contributions_;
+      kk_rotate_seconds_ += seconds;
+    } else if (contribution == Contribution::dd) {
+      ++dd_contributions_;
+      dd_rotate_seconds_ += seconds;
+    }
+    q_min_ = std::min(q_min_, q);
+    q_max_ = std::max(q_max_, q);
+    r_min_ = std::min(r_min_, r);
+    r_max_ = std::max(r_max_, r);
+    if (diagonal)
+      multiply_adds_ += static_cast<long double>(q) * static_cast<long double>(r)
+                     + static_cast<long double>(q) * static_cast<long double>(r) * static_cast<long double>(r);
+    else
+      multiply_adds_ += static_cast<long double>(q) * static_cast<long double>(q) * static_cast<long double>(r)
+                     + static_cast<long double>(q) * static_cast<long double>(r) * static_cast<long double>(r);
+    u_bytes_ += static_cast<long double>(q) * static_cast<long double>(r) * static_cast<long double>(element_size);
+    rho_bytes_ += static_cast<long double>(diagonal ? q : q * q) * static_cast<long double>(element_size);
+    out_bytes_ += static_cast<long double>(r) * static_cast<long double>(r) * static_cast<long double>(element_size);
+  }
+
+  void report() const {
+    if (!enabled_) return;
+    const auto q_min = contributions_ ? q_min_ : 0;
+    const auto r_min = contributions_ ? r_min_ : 0;
+    std::cout << fmt::format(
+        "[{} profile] N={} contrib={} rotate={:.6g}s kk={} kk_rotate={:.6g}s dd={} dd_rotate={:.6g}s diag_load={:.6g}s init={:.6g}s trace={:.6g}s save={:.6g}s "
+        "q={}..{} r={}..{} madd={:.6g} U={:.6g}MB rho={:.6g}MB out={:.6g}MB missing_rho={} missing_diag={} zero_dim={} zero_omega={} zero_rmax={}",
+        label_, N_, contributions_, rotate_seconds_, kk_contributions_, kk_rotate_seconds_, dd_contributions_, dd_rotate_seconds_,
+        diag_load_seconds_, init_seconds_, trace_seconds_, save_seconds_,
+        q_min, q_max_, r_min, r_max_, static_cast<double>(multiply_adds_), megabytes(u_bytes_), megabytes(rho_bytes_), megabytes(out_bytes_),
+        missing_rho_, missing_diag_, zero_dim_, zero_omega_, zero_rmax_) << std::endl;
+  }
+
+ private:
+  [[nodiscard]] static double megabytes(const long double bytes) { return static_cast<double>(bytes / (1024.0L * 1024.0L)); }
+
+  std::string label_;
+  size_t N_{};
+  bool enabled_{};
+  size_t contributions_{};
+  size_t kk_contributions_{};
+  size_t dd_contributions_{};
+  size_t missing_rho_{};
+  size_t missing_diag_{};
+  size_t zero_dim_{};
+  size_t zero_omega_{};
+  size_t zero_rmax_{};
+  size_t q_min_ = std::numeric_limits<size_t>::max();
+  size_t q_max_{};
+  size_t r_min_ = std::numeric_limits<size_t>::max();
+  size_t r_max_{};
+  long double multiply_adds_{};
+  long double u_bytes_{};
+  long double rho_bytes_{};
+  long double out_bytes_{};
+  double rotate_seconds_{};
+  double kk_rotate_seconds_{};
+  double dd_rotate_seconds_{};
+  double diag_load_seconds_{};
+  double init_seconds_{};
+  double trace_seconds_{};
+  double save_seconds_{};
+};
+
 // Calculate rho_N, the density matrix at the last NRG iteration. It is
 // normalized to 1. Note: in CFS approach, we consider all states in the
 // last iteration to be "discarded".
@@ -73,43 +168,55 @@ void cdmI(const size_t i,        // Subspace index
           const size_t N,
           const t_coef factor, // multiplicative factor that accounts for multiplicity
           const BackiterStore &store_all,
-          const Params &P)
+          const Params &P,
+          DmnrgProfile *profile = nullptr,
+          const DmnrgProfile::Contribution contribution = DmnrgProfile::Contribution::generic,
+          const bool diagonal_rho = false)
 {
   my_assert(i < P.combs);
   nrglog('D', "cdmI i=" << i << " I1=" << I1 << " factor=" << factor);
   // Range of indexes r and r' in matrix C^{QS,N}_{r,r'}, cf. Eq. (3.55) in my dissertation.
   const auto dim = size2(rhoNEW);
   // number of states taken into account in the density-matrix at *current* (Nth) stage (in subspace I1)
   const auto nromega = size2(rhoN);
-  if (nromega == 0 || dim == 0) return;   // continue only if connection exists
+  if (nromega == 0) { if (profile != nullptr) profile->add_zero_omega(); return; }
+  if (dim == 0) { if (profile != nullptr) profile->add_zero_dim(); return; } // continue only if connection exists
   // rmax (info[I1].rmax[i]) is the range of r in U^N_I1(omega|ri), only those states that we actually kept..
   const auto rmax = store_all[N].at(I1).rmax.rmax(i);
-  if (rmax == 0) return;    // rmax can be zero in the case a subspace has been completely truncated
+  if (rmax == 0) { if (profile != nullptr) profile->add_zero_rmax(); return; } // rmax can be zero in the case a subspace has been completely truncated
   my_assert(rmax == dim);   // Otherwise, rmax must equal dim
   // Check range of omega: do the dimensions of C^N_I1(omega omega') and U^N_I1(omega|r1) match?
   my_assert(nromega <= diagI1.getnrstored());
   const auto &U0 = diagI1.U.get(i);
   const auto U = NRG::submatrix_const(U0, {0, nromega}, {0, size2(U0)});
-  rotate<S>(rhoNEW, factor, U, rhoN);
+  const auto rotate_start = profile != nullptr && profile->enabled() ? DmnrgProfile::now() : DmnrgProfile::time_point{};
+  if (diagonal_rho)
+    rotate_diagonal<S>(rhoNEW, factor, U, rhoN);
+  else
+    rotate<S>(rhoNEW, factor, U, rhoN);
+  if (profile != nullptr && profile->enabled())
+    profile->add_rotate(nromega, dim, sizeof(S), DmnrgProfile::elapsed(rotate_start), contribution, diagonal_rho);
 }
 
 // Calculation of the shell-N REDUCED DENSITY MATRICES: Calculate rho at previous iteration (N-1) from rho
 // at the current iteration (N, rho)
 template<scalar S>
 auto calc_densitymatrix_iterN(const DiagInfo<S> &diag, const DensMatElements<S> &rho,
-                              const size_t N, const BackiterStore &store_all, const Symmetry<S> *Sym, const Params &P) {
+                              const size_t N, const BackiterStore &store_all, const Symmetry<S> *Sym, const Params &P,
+                              DmnrgProfile *profile = nullptr) {
   nrglog('D', "calc_densitymatrix_iterN N=" << N);
   DensMatElements<S> rhoPrev;
   for (const auto &[I, dimsub] : store_all[N - 1]) { // loop over all subspaces at *previous* iteration
     const auto dim  = dimsub.kept();
     rhoPrev[I]      = zero_matrix<S>(dim);
-    if (dim == 0) continue;
+    if (dim == 0) { if (profile != nullptr) profile->add_zero_dim(); continue; }
     const auto ns = Sym->new_subspaces(I);
     for (const auto &[i, sub] : ns | ranges::views::enumerate) {
       const auto x = rho.find(sub);
       const auto y = diag.find(sub);
-      if (x != rho.end() && y != diag.end())
-        cdmI(i, sub, x->second, y->second, rhoPrev[I], N, double(Sym->mult(sub)) / double(Sym->mult(I)), store_all, P);
+      if (x == rho.end()) { if (profile != nullptr) profile->add_missing_rho(); continue; }
+      if (y == diag.end()) { if (profile != nullptr) profile->add_missing_diag(); continue; }
+      cdmI(i, sub, x->second, y->second, rhoPrev[I], N, double(Sym->mult(sub)) / double(Sym->mult(I)), store_all, P, profile);
     }
   }
   return rhoPrev;
@@ -137,10 +244,20 @@ void calc_densitymatrix(DensMatElements<S> &rho, const BackiterStore &store_all,
   const auto section_timing = mt.time_it("DM");
   for (size_t N = P.Nmax - 1; N > P.Ninit; N--) {
     std::cout << "[DM] " << N << std::endl;
+    DmnrgProfile profile("DM", N, P.logletter('Y'));
+    auto timer = DmnrgProfile::now();
     const DiagInfo<S> diag_loaded(N, P);
-    auto rhoPrev = calc_densitymatrix_iterN(diag_loaded, rho, N, store_all, Sym, P); // need store_all for backiteration!
-    if (P.checkrho) check_trace_rho(rhoPrev, Sym->multfnc()); // Make sure rho is normalized to 1.
+    profile.add_diag_load(DmnrgProfile::elapsed(timer));
+    auto rhoPrev = calc_densitymatrix_iterN(diag_loaded, rho, N, store_all, Sym, P, &profile); // need store_all for backiteration!
+    if (P.checkrho) {
+      timer = DmnrgProfile::now();
+      check_trace_rho(rhoPrev, Sym->multfnc()); // Make sure rho is normalized to 1.
+      profile.add_trace(DmnrgProfile::elapsed(timer));
+    }
+    timer = DmnrgProfile::now();
     rhoPrev.save(N-1, P, filename);
+    profile.add_save(DmnrgProfile::elapsed(timer));
+    profile.report();
     rho.swap(rhoPrev);
   }
 }
@@ -179,30 +296,45 @@ auto calc_fulldensitymatrix_iterN(const Step &step, // only required for step::l
                                   const DiagInfo<S> &diag,
                                   const DensMatElements<S> &rhoFDM, // input
                                   const size_t N, const ThermoStore<S> &store, const BackiterStore &store_all, const Stats<S> &stats,
-                                  const Symmetry<S> *Sym, const Params &P) {
+                                  const Symmetry<S> *Sym, const Params &P,
+                                  DmnrgProfile *profile = nullptr) {
   nrglog('D', "calc_fulldensitymatrix_iterN N=" << N);
   DensMatElements<S> rhoDD;
   DensMatElements<S> rhoFDMPrev;
-  if (!step.last(N))
+  if (!step.last(N)) {
+    const auto timer = profile != nullptr && profile->enabled() ? DmnrgProfile::now() : DmnrgProfile::time_point{};
     rhoDD = init_rho_FDM(N, store, stats, Sym->multfnc(), P.T, P.checkrho); // store here!
+    if (profile != nullptr && profile->enabled()) profile->add_init(DmnrgProfile::elapsed(timer));
+  }
   for (const auto &[I, ds] : store_all[N - 1]) { // loop over all subspaces at *previous* iteration, hence store_all here
     const auto subs = Sym->new_subspaces(I);
     const auto dim  = ds.kept();
     rhoFDMPrev[I]   = zero_matrix<S>(dim);
-    if (!dim) continue;
+    if (!dim) { if (profile != nullptr) profile->add_zero_dim(); continue; }
     for (const auto i : Sym->combs()) {
       const auto sub = subs[i];
       // DM construction for non-Abelian symmetries: must include the ratio of multiplicities as a coefficient.
       const auto coef = double(Sym->mult(sub)) / double(Sym->mult(I));
       // Contribution from the KK sector.
       const auto x1 = rhoFDM.find(sub);
       const auto y = diag.find(sub);
-      if (x1 != rhoFDM.end() && y != diag.end())
-        cdmI(i, sub, x1->second, y->second, rhoFDMPrev[I], N, coef, store_all, P);
+      if (x1 == rhoFDM.end()) {
+        if (profile != nullptr) profile->add_missing_rho();
+      } else if (y == diag.end()) {
+        if (profile != nullptr) profile->add_missing_diag();
+      } else {
+        cdmI(i, sub, x1->second, y->second, rhoFDMPrev[I], N, coef, store_all, P, profile, DmnrgProfile::Contribution::kk);
+      }
       // Contribution from the DD sector. rhoDD -> rhoFDMPrev
-      if (!step.last(N))
-        if (const auto x2 = rhoDD.find(sub); x2 !=rhoDD.end() && y != diag.end())
-          cdmI(i, sub, x2->second, y->second, rhoFDMPrev[I], N, coef, store_all, P);
+      if (!step.last(N)) {
+        if (const auto x2 = rhoDD.find(sub); x2 == rhoDD.end()) {
+          if (profile != nullptr) profile->add_missing_rho();
+        } else if (y == diag.end()) {
+          if (profile != nullptr) profile->add_missing_diag();
+        } else {
+          cdmI(i, sub, x2->second, y->second, rhoFDMPrev[I], N, coef, store_all, P, profile, DmnrgProfile::Contribution::dd, true);
+        }
+      }
       // (Exception: for the N-1 iteration, the rhoPrev is already initialized with the DD sector of the last iteration.) }
     } // over combinations
   } // over subspaces
@@ -216,16 +348,24 @@ void calc_fulldensitymatrix(const Step &step, DensMatElements<S> &rhoFDM, const
   const auto section_timing = mt.time_it("FDM");
   for (size_t N = P.Nmax - 1; N > P.Ninit; N--) {
     std::cout << "[FDM] " << N << std::endl;
+    DmnrgProfile profile("FDM", N, P.logletter('Y'));
+    auto timer = DmnrgProfile::now();
     const DiagInfo<S> diag_loaded(N, P); // = load_and_project(N, Sym, P);
-    auto rhoFDMPrev        = calc_fulldensitymatrix_iterN(step, diag_loaded, rhoFDM, N, store, store_all, stats, Sym, P);
+    profile.add_diag_load(DmnrgProfile::elapsed(timer));
+    auto rhoFDMPrev        = calc_fulldensitymatrix_iterN(step, diag_loaded, rhoFDM, N, store, store_all, stats, Sym, P, &profile);
     if (P.checkrho) {
+      timer = DmnrgProfile::now();
       const auto tr          = rhoFDMPrev.trace(Sym->multfnc());
       const auto expected    = std::accumulate(stats.wn.begin() + N, stats.wn.begin() + P.Nmax, 0.0);
       const auto diff        = (tr - expected) / expected;
       nrglog('w', "tr[rhoFDM(" << N << ")]=" << tr << " sum(wn)=" << expected << " diff=" << diff);
       my_assert(num_equal(diff, 0.0));
+      profile.add_trace(DmnrgProfile::elapsed(timer));
     }
+    timer = DmnrgProfile::now();
     rhoFDMPrev.save(N-1, P, filename);
+    profile.add_save(DmnrgProfile::elapsed(timer));
+    profile.report();
     rhoFDM.swap(rhoFDMPrev);
   }
 }

c++/numerics_Eigen.hpp

@@ -117,6 +117,15 @@ void rotate(EM &M, const t_coef factor, const U_type &U, const EM &O) {
   }
 }
 
+template<scalar S, typename U_type, Eigen_matrix EM, typename t_coef = coef_traits<S>> // XXX: U_type
+void rotate_diagonal(EM &M, const t_coef factor, const U_type &U, const EM &O) {
+  if (finite_size(U)) {
+    assert(size1(M) == size2(U) && size1(U) == size1(O) && size2(O) == size1(U) && size2(U) == size2(M));
+    assert(my_isfinite(factor));
+    M.noalias() += factor * U.adjoint() * O.diagonal().asDiagonal() * U;
+  }
+}
+
 template<scalar S>
 Eigen::Block<const EigenMatrix<S>> submatrix_const(const EigenMatrix<S> &M, const std::pair<size_t,size_t> &r1, const std::pair<size_t,size_t> &r2)
 {

c++/params.hpp

@@ -450,8 +450,9 @@ class Params {
    w - calculation of weights w_n
    M - MPI parallelization details
    ! - debug internal variables
-   D - DMNRG calculation details
-   Z - report the values of different partition functions
+    D - DMNRG calculation details
+    Y - DMNRG profiling summary
+    Z - report the values of different partition functions
    Useful combinations:
     @0 - high-level calculation flow
     is - debug matrix construction

test/unit/numerics/numerics.cpp

@@ -697,6 +697,18 @@ TEST(numerics, rotate) {
   EXPECT_TRUE(r.isApprox(ref));
 }
 
+TEST(numerics, rotate_diagonal) {
+  auto u = generate_matrix<double>(3,2);
+  auto o = NRG::zero_matrix<double>(3,3);
+  u(0,0) = 1; u(0,1) = 2; u(1,0) = 3; u(1,1) = 4; u(2,0) = 5; u(2,1) = 6;
+  o(0,0) = 2; o(1,1) = 3; o(2,2) = 4;
+  auto diagonal = NRG::zero_matrix<double>(2,2);
+  auto dense = NRG::zero_matrix<double>(2,2);
+  rotate_diagonal<double>(diagonal, 1.0, u, o);
+  rotate<double>(dense, 1.0, u, o);
+  EXPECT_TRUE(diagonal.isApprox(dense));
+}
+
 TEST(numerics, matrix_prod) {
   auto a = generate_matrix<double>(2,2);
   auto b = generate_matrix<double>(2,2);