Fix GPU memory MPI communication bugs for DeltaSpin/DFT+U

Three critical bugs fixed where GPU memory pointers were passed directly
to MPI functions (MPI_Allreduce/MPI_Bcast), causing segfaults:

1. op_pw_nl.cpp: becp GPU memory in Nonlocal::act()
   - Add CPU buffer becp_h, d2h/h2d copy around MPI_Allreduce

2. diago_dav_subspace.cpp: hcc/scc GPU memory in cal_elem()
   - Add CPU buffers hcc_h/scc_h, d2h/h2d copy around MPI_Allreduce

3. diago_dav_subspace.cpp: vcc GPU memory in diag_zhegvx()
   - Add CPU buffer vcc_h for MPI_Bcast, d2h/h2d copy

4. elecstate_pw.cpp: becp GPU memory in cal_becsum()
   - Add CPU buffer becp_h, d2h/h2d copy around MPI_Allreduce

All fixes use if constexpr to only apply GPU-specific code paths
when Device is DEVICE_GPU, preserving CPU path behavior.

Additionally, diag_zhegvx GPU path now falls back to CPU hegvx_op
for robustness against ill-conditioned overlap matrices that caused
'heevd: failed to invert matrix' errors with GPU cusolver hegvd.

Author: dyzheng

source/source_estate/elecstate_pw.cpp

@@ -1,5 +1,7 @@
 #include "elecstate_pw.h"
 
+#include <type_traits>
+
 #include "source_base/constants.h"
 #include "source_base/libm/libm.h"
 #include "source_base/math_ylmreal.h"
@@ -55,6 +57,10 @@ ElecStatePW<T, Device>::~ElecStatePW()
     }
     delmem_complex_op()(this->wfcr);
     delmem_complex_op()(this->wfcr_another_spin);
+    if (this->becp_h != nullptr)
+    {
+        delmem_complex_h_op()(this->becp_h);
+    }
 }
 
 template<typename T, typename Device>
@@ -339,7 +345,26 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
                       becp,
                       this->ppcell->nkb);
         }
-        Parallel_Reduce::reduce_pool(becp, this->ppcell->nkb * nbands);
+        // Copy becp from GPU to CPU for MPI reduction (only needed for GPU)
+        if constexpr (std::is_same<Device, base_device::DEVICE_GPU>::value)
+        {
+            if (this->becp_h_size < this->ppcell->nkb * nbands)
+            {
+                if (this->becp_h != nullptr)
+                {
+                    delmem_complex_h_op()(this->becp_h);
+                }
+                resmem_complex_h_op()(this->becp_h, this->ppcell->nkb * nbands, "ElecStatePW::becp_h");
+                this->becp_h_size = this->ppcell->nkb * nbands;
+            }
+            syncmem_complex_d2h_op()(this->becp_h, becp, this->ppcell->nkb * nbands);
+            Parallel_Reduce::reduce_pool(this->becp_h, this->ppcell->nkb * nbands);
+            syncmem_complex_h2d_op()(becp, this->becp_h, this->ppcell->nkb * nbands);
+        }
+        else
+        {
+            Parallel_Reduce::reduce_pool(becp, this->ppcell->nkb * nbands);
+        }
 
         // sum over bands: \sum_i <psi_i|beta_l><beta_m|psi_i> w_i
         for (int it = 0; it < ucell->ntype; it++)

source/source_estate/elecstate_pw.h

@@ -100,9 +100,16 @@ class ElecStatePW : public ElecState
     using setmem_complex_op = base_device::memory::set_memory_op<T, Device>;
     using resmem_complex_op = base_device::memory::resize_memory_op<T, Device>;
     using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;
+    using resmem_complex_h_op = base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>;
+    using delmem_complex_h_op = base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>;
+    using syncmem_complex_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;
+    using syncmem_complex_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
 
     using gemv_op = ModuleBase::gemv_op<T, Device>;
     using gemm_op = ModuleBase::gemm_op<T, Device>;
+
+    T* becp_h = nullptr;
+    size_t becp_h_size = 0;
 };
 
 } // namespace elecstate

source/source_hsolver/diago_dav_subspace.cpp

@@ -1,5 +1,7 @@
 #include "diago_dav_subspace.h"
 
+#include <type_traits>
+
 #include "diago_iter_assist.h"
 
 #include "source_base/module_device/device.h"
@@ -103,6 +105,8 @@ Diago_DavSubspace<T, Device>::~Diago_DavSubspace()
         delmem_real_op()(this->d_precondition);
         delmem_complex_op()(this->d_scc);
         delmem_real_op()(this->d_eigenvalue);
+        delmem_complex_h_op()(this->hcc_h);
+        delmem_complex_h_op()(this->scc_h);
     }
 #endif
 }
@@ -585,8 +589,30 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
         mtfunc::dsp_dav_subspace_reduce(hcc, scc, nbase, this->nbase_x, this->notconv, this->diag_comm.comm);
 #else
         assert(this->diag_comm.comm == POOL_WORLD);
-        Parallel_Reduce::reduce_pool(hcc + nbase * this->nbase_x, notconv * this->nbase_x);
-        Parallel_Reduce::reduce_pool(scc + nbase * this->nbase_x, notconv * this->nbase_x);
+        // For GPU, need to copy to CPU before MPI reduction
+        if constexpr (std::is_same<Device, base_device::DEVICE_GPU>::value)
+        {
+            // Allocate CPU buffers if needed
+            if (this->hcc_h == nullptr)
+            {
+                resmem_complex_h_op()(this->hcc_h, this->nbase_x * this->nbase_x, "DAV::hcc_h");
+                resmem_complex_h_op()(this->scc_h, this->nbase_x * this->nbase_x, "DAV::scc_h");
+            }
+            // Copy from GPU to CPU
+            syncmem_d2h_op()(this->hcc_h, hcc + nbase * this->nbase_x, notconv * this->nbase_x);
+            syncmem_d2h_op()(this->scc_h, scc + nbase * this->nbase_x, notconv * this->nbase_x);
+            // Reduce on CPU
+            Parallel_Reduce::reduce_pool(this->hcc_h, notconv * this->nbase_x);
+            Parallel_Reduce::reduce_pool(this->scc_h, notconv * this->nbase_x);
+            // Copy back to GPU
+            syncmem_h2d_op()(hcc + nbase * this->nbase_x, this->hcc_h, notconv * this->nbase_x);
+            syncmem_h2d_op()(scc + nbase * this->nbase_x, this->scc_h, notconv * this->nbase_x);
+        }
+        else
+        {
+            Parallel_Reduce::reduce_pool(hcc + nbase * this->nbase_x, notconv * this->nbase_x);
+            Parallel_Reduce::reduce_pool(scc + nbase * this->nbase_x, notconv * this->nbase_x);
+        }
 #endif
     }
 #endif
@@ -615,9 +641,24 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
 #if defined(__CUDA) || defined(__ROCM)
         if (this->diag_comm.rank == 0)
         {
-            syncmem_complex_op()(this->d_scc, scc, nbase * this->nbase_x);
-            ct::kernels::lapack_hegvd<T, ct_Device>()(nbase, this->nbase_x, this->hcc, this->d_scc, this->d_eigenvalue, this->vcc);
-            syncmem_var_d2h_op()((*eigenvalue_iter).data(), this->d_eigenvalue, this->nbase_x);
+            // Copy hcc and scc from GPU to CPU for robust eigensolving
+            std::vector<T> hcc_h(nbase * this->nbase_x, *this->zero);
+            std::vector<T> scc_h(nbase * this->nbase_x, *this->zero);
+            syncmem_d2h_op()(hcc_h.data(), this->hcc, nbase * this->nbase_x);
+            syncmem_d2h_op()(scc_h.data(), scc, nbase * this->nbase_x);
+
+            // Solve on CPU using hegvx (more robust than hegvd for ill-conditioned overlap)
+            std::vector<T> vcc_h(nbase * this->nbase_x, *this->zero);
+            hegvx_op<T, base_device::DEVICE_CPU>()(this->cpu_ctx,
+                                                    nbase,
+                                                    this->nbase_x,
+                                                    hcc_h.data(),
+                                                    scc_h.data(),
+                                                    nband,
+                                                    (*eigenvalue_iter).data(),
+                                                    vcc_h.data());
+            // Copy eigenvectors back to GPU
+            syncmem_h2d_op()(this->vcc, vcc_h.data(), nbase * this->nbase_x);
         }
 #endif
     }
@@ -715,9 +756,25 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
     if (this->diag_comm.nproc > 1)
     {
         // vcc: nbase * nband
-        for (int i = 0; i < nband; i++)
+        if constexpr (std::is_same<Device, base_device::DEVICE_GPU>::value)
         {
-            MPI_Bcast(&vcc[i * this->nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, this->diag_comm.comm);
+            // Copy vcc from GPU to CPU for MPI broadcast
+            T* vcc_h = nullptr;
+            resmem_complex_h_op()(vcc_h, nband * this->nbase_x, "DAV::vcc_h");
+            syncmem_d2h_op()(vcc_h, vcc, nband * this->nbase_x);
+            for (int i = 0; i < nband; i++)
+            {
+                MPI_Bcast(&vcc_h[i * this->nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, this->diag_comm.comm);
+            }
+            syncmem_h2d_op()(vcc, vcc_h, nband * this->nbase_x);
+            delmem_complex_h_op()(vcc_h);
+        }
+        else
+        {
+            for (int i = 0; i < nband; i++)
+            {
+                MPI_Bcast(&vcc[i * this->nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, this->diag_comm.comm);
+            }
         }
         MPI_Bcast((*eigenvalue_iter).data(), nband, MPI_DOUBLE, 0, this->diag_comm.comm);
     }

source/source_hsolver/diago_dav_subspace.h

@@ -190,12 +190,18 @@ class Diago_DavSubspace
     using syncmem_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
     using syncmem_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;
 
+    using resmem_complex_h_op = base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>;
+    using delmem_complex_h_op = base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>;
+
     // Note that ct_Device is different from base_device!
     using ct_Device = typename ct::PsiToContainer<Device>::type;
     // using hegvd_op = container::kernels::lapack_hegvd<T, ct_Device>;
 
     const T *one = nullptr, *zero = nullptr, *neg_one = nullptr;
     const T one_ = static_cast<T>(1.0), zero_ = static_cast<T>(0.0), neg_one_ = static_cast<T>(-1.0);
+
+    T* hcc_h = nullptr;
+    T* scc_h = nullptr;
 };
 
 } // namespace hsolver

source/source_pw/module_pwdft/op_pw_nl.cpp

@@ -1,5 +1,7 @@
 #include "op_pw_nl.h"
 
+#include <type_traits>
+
 #include "source_io/module_parameter/parameter.h"
 #include "source_base/timer.h"
 #include "source_base/parallel_reduce.h"
@@ -34,6 +36,7 @@ template<typename T, typename Device>
 Nonlocal<OperatorPW<T, Device>>::~Nonlocal() {
     delmem_complex_op()(this->ps);
     delmem_complex_op()(this->becp);
+    delmem_complex_h_op()(this->becp_h);
 }
 
 template<typename T, typename Device>
@@ -285,7 +288,25 @@ void Nonlocal<OperatorPW<T, Device>>::act(
             );
         }
 
-        Parallel_Reduce::reduce_pool(becp, nkb * nbands);
+        // Copy becp from GPU to CPU for MPI reduction (only needed for GPU)
+        if constexpr (std::is_same<Device, base_device::DEVICE_GPU>::value)
+        {
+            if (this->nkb_m_h < nkb * nbands)
+            {
+                resmem_complex_h_op()(this->becp_h, nkb * nbands, "Nonlocal<PW>::becp_h");
+                this->nkb_m_h = nkb * nbands;
+            }
+            syncmem_complex_d2h_op()(this->becp_h, this->becp, nkb * nbands);
+
+            Parallel_Reduce::reduce_pool(becp_h, nkb * nbands);
+
+            // Copy back to GPU
+            syncmem_complex_h2d_op()(this->becp, this->becp_h, nkb * nbands);
+        }
+        else
+        {
+            Parallel_Reduce::reduce_pool(becp, nkb * nbands);
+        }
 
         this->add_nonlocal_pp(tmhpsi, becp, nbands);
     }

source/source_pw/module_pwdft/op_pw_nl.h

@@ -98,9 +98,15 @@ class Nonlocal<OperatorPW<T, Device>> : public OperatorPW<T, Device>
     using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;
 #endif
     using syncmem_complex_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
+    using syncmem_complex_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;
+    using resmem_complex_h_op = base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>;
+    using delmem_complex_h_op = base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>;
 
     T one{1, 0};
     T zero{0, 0};
+
+    mutable T *becp_h = nullptr;
+    mutable size_t nkb_m_h = 0;
 };
 
 } // namespace hamilt