Undo cuda aware mpi

Author: Flying-dragon-boxing

source/source_basis/module_pw/pw_basis.h

@@ -248,11 +248,11 @@ class PW_Basis
     ModuleBase::FFT_Bundle fft_bundle;
     //The position of pointer in and out can be equal(in-place transform) or different(out-of-place transform).
 
-    template <typename FPTYPE>
-    void real2recip(const FPTYPE* in,
-                    std::complex<FPTYPE>* out,
-                    const bool add = false,
-                    const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
+template <typename FPTYPE>
+void real2recip(const FPTYPE* in,
+                std::complex<FPTYPE>* out,
+                const bool add = false,
+                const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
     template <typename FPTYPE>
     void real2recip(const std::complex<FPTYPE>* in,
                     std::complex<FPTYPE>* out,
@@ -269,16 +269,16 @@ class PW_Basis
                     const bool add = false,
                     const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
 
+template <typename FPTYPE>
+void real2recip_gpu(const FPTYPE* in,
+                std::complex<FPTYPE>* out,
+                const bool add = false,
+                const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
     template <typename FPTYPE>
-    void real2recip_gpu(const FPTYPE* in,
-                    std::complex<FPTYPE>* out,
-                    const bool add = false,
-                    const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
-    template <typename FPTYPE>
-    void real2recip_gpu(const std::complex<FPTYPE>* in,
-                    std::complex<FPTYPE>* out,
-                    const bool add = false,
-                    const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
+void real2recip_gpu(const std::complex<FPTYPE>* in,
+                std::complex<FPTYPE>* out,
+                const bool add = false,
+                const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
     template <typename FPTYPE>
     void recip2real_gpu(const std::complex<FPTYPE>* in,
                     FPTYPE* out,

source/source_pw/module_pwdft/operator_pw/op_exx_pw.cpp

@@ -244,13 +244,26 @@ void OperatorEXXPW<T, Device>::act_op(const int nbands,
 
             // kpar: to real and bcast between same rank_in_pool
             // auto request = MPI_REQUEST_
+            T* psi_mq_real_cpu = new T[wfcpw->nrxx];
             if (iq_pool == GlobalV::MY_POOL)
             {
                 const T* psi_mq = get_pw(m_iband, iq_loc);
                 wfcpw->recip_to_real(ctx, psi_mq, psi_mq_real, iq_loc);
+                cudaMemcpy(psi_mq_real_cpu, psi_mq_real, wfcpw->nrxx, cudaMemcpyDeviceToHost);
                 // send
             }
             MPI_Bcast(psi_mq_real, wfcpw->nrxx, MPI_DOUBLE_COMPLEX, iq_pool, KP_WORLD);
+            if (iq_pool == GlobalV::MY_POOL)
+            {
+                const T* psi_mq = get_pw(m_iband, iq_loc);
+                wfcpw->recip_to_real(ctx, psi_mq, psi_mq_real, iq_loc);
+                // send
+            }
+            else
+            {
+                cudaMemcpy(psi_mq_real, psi_mq_real_cpu, wfcpw->nrxx, cudaMemcpyHostToDevice);
+            }
+            delete[] psi_mq_real_cpu;
             // std::cout << "psi_mq_real[0]: " << psi_mq_real[0] << std::endl;
             // if (GlobalV::RANK_IN_POOL == 1)
             // {