Fix: Backport selected fixes from develop to LTS

source/module_base/parallel_reduce.cpp

@@ -4,6 +4,15 @@
 
 #include <vector>
 
+#ifdef __MPI
+const MPI_Datatype Parallel_Reduce::MPI_Type<int>::value = MPI_INT;
+const MPI_Datatype Parallel_Reduce::MPI_Type<double>::value = MPI_DOUBLE;
+const MPI_Datatype Parallel_Reduce::MPI_Type<float>::value = MPI_FLOAT;
+const MPI_Datatype Parallel_Reduce::MPI_Type<std::complex<double>>::value = MPI_DOUBLE_COMPLEX;
+const MPI_Datatype Parallel_Reduce::MPI_Type<std::complex<float>>::value = MPI_C_FLOAT_COMPLEX;
+const MPI_Datatype Parallel_Reduce::MPI_Type<long long>::value = MPI_LONG_LONG;
+#endif
+
 template <>
 void Parallel_Reduce::reduce_all<int>(int& object)
 {

source/module_base/parallel_reduce.h

@@ -12,6 +12,47 @@ using std::complex;
 
 namespace Parallel_Reduce
 {
+#ifdef __MPI
+template <typename T>
+struct MPI_Type;
+
+template <>
+struct MPI_Type<int>
+{
+    static const MPI_Datatype value;
+};
+
+template <>
+struct MPI_Type<double>
+{
+    static const MPI_Datatype value;
+};
+
+template <>
+struct MPI_Type<float>
+{
+    static const MPI_Datatype value;
+};
+
+template <>
+struct MPI_Type<std::complex<double>>
+{
+    static const MPI_Datatype value;
+};
+
+template <>
+struct MPI_Type<std::complex<float>>
+{
+    static const MPI_Datatype value;
+};
+
+template <>
+struct MPI_Type<long long>
+{
+    static const MPI_Datatype value;
+};
+#endif
+
 /// reduce in all process
 template <typename T>
 void reduce_all(T& object);

source/module_basis/module_pw/module_fft/fft_bundle.cpp

@@ -100,6 +100,18 @@ void FFT_Bundle::initfft(int nx_in,
             fft_double = make_unique<FFT_CUDA<double>>();
             fft_double->initfft(nx_in,ny_in,nz_in);
         #endif
+        // Host-memory interface: always provide a double CPU backend, and a
+        // float one when float FFTs are enabled. Single precision still needs
+        // a double host backend for host-only post-processing.
+        if (float_flag)
+        {
+            host_fft_float = make_unique<FFT_CPU<float>>(this->fft_mode);
+            host_fft_float->initfft(nx_in, ny_in, nz_in, lixy_in, rixy_in, ns_in,
+                                    nplane_in, nproc_in, gamma_only_in, xprime_in);
+        }
+        host_fft_double = make_unique<FFT_CPU<double>>(this->fft_mode);
+        host_fft_double->initfft(nx_in, ny_in, nz_in, lixy_in, rixy_in, ns_in,
+                                 nplane_in, nproc_in, gamma_only_in, xprime_in);
     }
 
 }
@@ -108,74 +120,80 @@ void FFT_Bundle::setupFFT()
 {
     if (double_flag){fft_double->setupFFT();}
     if (float_flag) {fft_float->setupFFT();}
+    if (host_fft_double != nullptr) { host_fft_double->setupFFT(); }
+    if (host_fft_float != nullptr) { host_fft_float->setupFFT(); }
 }
 
 void FFT_Bundle::clearFFT()
 {
     if (double_flag){fft_double->cleanFFT();}
     if (float_flag) {fft_float->cleanFFT();}
+    if (host_fft_double != nullptr) { host_fft_double->cleanFFT(); }
+    if (host_fft_float != nullptr) { host_fft_float->cleanFFT(); }
 }
 void FFT_Bundle::clear()
 {
     this->clearFFT();
     if (double_flag){fft_double->clear();}
     if (float_flag) {fft_float->clear();}
+    if (host_fft_double != nullptr) { host_fft_double->clear(); }
+    if (host_fft_float != nullptr) { host_fft_float->clear(); }
 }
 
 template <> void 
 FFT_Bundle::fftxyfor(std::complex<float>* in, 
                      std::complex<float>* out) 
-const {fft_float->fftxyfor(in,out);}
+const {get_host_fft_float()->fftxyfor(in,out);}
 template <> void 
 FFT_Bundle::fftxyfor(std::complex<double>* in,
                      std::complex<double>* out) 
-const {fft_double->fftxyfor(in,out);}
+const {get_host_fft_double()->fftxyfor(in,out);}
 
 
 template <> void 
 FFT_Bundle::fftzfor(std::complex<float>* in, 
                     std::complex<float>* out) 
-const {fft_float->fftzfor(in,out);}
+const {get_host_fft_float()->fftzfor(in,out);}
 template <> void 
 FFT_Bundle::fftzfor(std::complex<double>* in, 
                     std::complex<double>* out) 
-const {fft_double->fftzfor(in,out);}
+const {get_host_fft_double()->fftzfor(in,out);}
 
 template <> void 
 FFT_Bundle::fftxybac(std::complex<float>* in, 
                      std::complex<float>* out) 
-const {fft_float->fftxybac(in,out);}
+const {get_host_fft_float()->fftxybac(in,out);}
 template <> void 
 FFT_Bundle::fftxybac(std::complex<double>* in, 
                      std::complex<double>* out) 
-const {fft_double->fftxybac(in,out);}
+const {get_host_fft_double()->fftxybac(in,out);}
 
 template <> void 
 FFT_Bundle::fftzbac(std::complex<float>* in, 
                     std::complex<float>* out) 
-const {fft_float->fftzbac(in,out);}
+const {get_host_fft_float()->fftzbac(in,out);}
 template <> void 
 FFT_Bundle::fftzbac(std::complex<double>* in, 
                     std::complex<double>* out) 
-const {fft_double->fftzbac(in,out);}
+const {get_host_fft_double()->fftzbac(in,out);}
 
 template <> void 
 FFT_Bundle::fftxyr2c(float* in, 
                      std::complex<float>* out) 
-const {fft_float->fftxyr2c(in,out);}
+const {get_host_fft_float()->fftxyr2c(in,out);}
 template <> void 
 FFT_Bundle::fftxyr2c(double* in, 
                     std::complex<double>* out) 
-const {fft_double->fftxyr2c(in,out);}
+const {get_host_fft_double()->fftxyr2c(in,out);}
 
 template <> void 
 FFT_Bundle::fftxyc2r(std::complex<float>* in, 
                     float* out) 
-const {fft_float->fftxyc2r(in,out);}
+const {get_host_fft_float()->fftxyc2r(in,out);}
 template <> void 
 FFT_Bundle::fftxyc2r(std::complex<double>* in, 
                     double* out) 
-const {fft_double->fftxyc2r(in,out);}
+const {get_host_fft_double()->fftxyc2r(in,out);}
 
 template <> void  
 FFT_Bundle::fft3D_forward(const base_device::DEVICE_GPU* ctx, 
@@ -201,19 +219,19 @@ const {fft_double->fft3D_backward(in, out);}
 
 // access the real space data
 template <> float* 
-FFT_Bundle::get_rspace_data()  const {return fft_float->get_rspace_data();}
+FFT_Bundle::get_rspace_data()  const {return get_host_fft_float()->get_rspace_data();}
 template <> double* 
-FFT_Bundle::get_rspace_data()  const {return fft_double->get_rspace_data();}
+FFT_Bundle::get_rspace_data()  const {return get_host_fft_double()->get_rspace_data();}
 
 template <> std::complex<float>* 
-FFT_Bundle::get_auxr_data()    const {return fft_float->get_auxr_data();}
+FFT_Bundle::get_auxr_data()    const {return get_host_fft_float()->get_auxr_data();}
 template <> std::complex<double>* 
-FFT_Bundle::get_auxr_data()    const {return fft_double->get_auxr_data();}
+FFT_Bundle::get_auxr_data()    const {return get_host_fft_double()->get_auxr_data();}
 
 template <> std::complex<float>* 
-FFT_Bundle::get_auxg_data()    const {return fft_float->get_auxg_data();}
+FFT_Bundle::get_auxg_data()    const {return get_host_fft_float()->get_auxg_data();}
 template <> std::complex<double>* 
-FFT_Bundle::get_auxg_data()    const {return fft_double->get_auxg_data();}
+FFT_Bundle::get_auxg_data()    const {return get_host_fft_double()->get_auxg_data();}
 
 template <> std::complex<float>* 
 FFT_Bundle::get_auxr_3d_data() const {return fft_float->get_auxr_3d_data();}

source/module_basis/module_pw/module_fft/fft_bundle.h

@@ -206,7 +206,21 @@ class FFT_Bundle
         bool double_flag=false;
         std::shared_ptr<FFT_BASE<float>> fft_float=nullptr;
         std::shared_ptr<FFT_BASE<double>> fft_double=nullptr;
-
+        // Extra CPU backends serving the host-memory ("cpu-like") interface
+        // when the primary backend is a device FFT. They keep host
+        // post-processing (e.g. Wannier90/Berry phase) working on GPU runs.
+        std::shared_ptr<FFT_BASE<float>> host_fft_float=nullptr;
+        std::shared_ptr<FFT_BASE<double>> host_fft_double=nullptr;
+
+        FFT_BASE<float>* get_host_fft_float() const
+        {
+            return (this->device == "cpu") ? this->fft_float.get() : this->host_fft_float.get();
+        }
+        FFT_BASE<double>* get_host_fft_double() const
+        {
+            return (this->device == "cpu") ? this->fft_double.get() : this->host_fft_double.get();
+        }
+
         std::string device = "cpu";
         std::string precision = "double";
 };   

source/module_basis/module_pw/pw_basis_k.cpp

@@ -412,24 +412,6 @@ std::vector<int> PW_Basis_K::get_ig2iz(const int ik) const
     return ig_to_iz;
 }
 
-template <>
-float * PW_Basis_K::get_kvec_c_data() const {
-    return this->s_kvec_c;
-}
-template <>
-double * PW_Basis_K::get_kvec_c_data() const {
-    return this->d_kvec_c;
-}
-
-template <>
-float * PW_Basis_K::get_gcar_data() const {
-    return this->s_gcar;
-}
-template <>
-double * PW_Basis_K::get_gcar_data() const {
-    return this->d_gcar;
-}
-
 template <>
 float * PW_Basis_K::get_gk2_data() const {
     return this->s_gk2;

source/module_basis/module_pw/pw_basis_k.h

@@ -183,6 +183,23 @@ class PW_Basis_K : public PW_Basis
 
 };
 
+// Trivial accessors are kept inline so that unit tests and other small
+// translation units that call grad_wfc do not need to link pw_basis_k.cpp.
+template <>
+inline float* PW_Basis_K::get_gcar_data() const { return this->s_gcar; }
+template <>
+inline double* PW_Basis_K::get_gcar_data() const
+{
+    return this->d_gcar != nullptr ? this->d_gcar : reinterpret_cast<double*>(&this->gcar[0][0]);
+}
+template <>
+inline float* PW_Basis_K::get_kvec_c_data() const { return this->s_kvec_c; }
+template <>
+inline double* PW_Basis_K::get_kvec_c_data() const
+{
+    return this->d_kvec_c != nullptr ? this->d_kvec_c : reinterpret_cast<double*>(&this->kvec_c[0][0]);
+}
+
 }
 #endif //PlaneWave_K class
 

source/module_basis/module_pw/pw_distributeg.cpp

@@ -48,7 +48,6 @@ void PW_Basis::count_pw_st(
 {
     ModuleBase::GlobalFunc::ZEROS(st_length2D, this->fftnxy);
     ModuleBase::GlobalFunc::ZEROS(st_bottom2D, this->fftnxy);
-    ModuleBase::Vector3<double> f;
 
     // determine the scaning area along x-direct, if gamma-only && xprime, only positive axis is used.
     int ix_end = int(this->nx / 2) + 1;
@@ -86,10 +85,18 @@ void PW_Basis::count_pw_st(
         }
     }
 
-    this->liy = this->riy = 0;
-    this->lix = this->rix = 0;
-    this->npwtot = 0;
-    this->nstot = 0;
+    int npwtot_local = 0;
+    int nstot_local = 0;
+    int lix_local = 0, rix_local = 0;
+    int liy_local = 0, riy_local = 0;
+
+#ifdef _OPENMP
+#pragma omp parallel for collapse(1) \
+    shared(st_length2D, st_bottom2D) \
+    reduction(+: npwtot_local, nstot_local) \
+    reduction(min: rix_local, riy_local) \
+    reduction(max: lix_local, liy_local)
+#endif
     for (int ix = ix_start; ix <= ix_end; ++ix)
     {
         for (int iy = iy_start; iy <= iy_end; ++iy)
@@ -101,37 +108,53 @@ void PW_Basis::count_pw_st(
             // so that its index in st_length and st_bottom is 9 * 10 + 2 = 92.
             int x = ix;
             int y = iy;
-            if (x < 0) x += this->nx;
-            if (y < 0) y += this->ny;
+            if (x < 0) { x += this->nx; }
+            if (y < 0) { y += this->ny; }
             int index = x * this->fftny + y;
 
             int length = 0; // number of planewave on stick (x, y).
+
+            ModuleBase::Vector3<double> f;
+            f.x = ix;
+            f.y = iy;
+
             for (int iz = iz_start; iz <= iz_end; ++iz)
             {
-                f.x = ix;
-                f.y = iy;
                 f.z = iz;
                 double modulus = f * (this->GGT * f);
                 if (modulus <= this->ggecut || this->full_pw)
                 {
-                    if (length == 0) st_bottom2D[index] = iz; // length == 0 means this point is the bottom of stick (x, y).
-                    ++this->npwtot;
+                    if (length == 0)
+                    {
+                        st_bottom2D[index] = iz; // length == 0 means this point is the bottom of stick (x, y).
+                    }
+                    ++npwtot_local;
                     ++length;
-                    if(iy < this->riy) this->riy = iy;
-                    if(iy > this->liy) this->liy = iy;
-                    if(ix < this->rix) this->rix = ix;
-                    if(ix > this->lix) this->lix = ix;
+                    if(ix > lix_local) { lix_local = ix; }
+                    if(ix < rix_local) { rix_local = ix; }
+                    if(iy > liy_local) { liy_local = iy; }
+                    if(iy < riy_local) { riy_local = iy; }
                 }
             }
             if (length > 0)
             {
                 st_length2D[index] = length;
-                ++this->nstot;
+                ++nstot_local;
             }
         }
     }
-    riy += this->ny;
-    rix += this->nx;
+
+    if (npwtot_local == 0)  // no planewave
+    {
+        lix_local = rix_local = liy_local = riy_local = 0;
+    }
+
+    this->npwtot = npwtot_local;
+    this->nstot = nstot_local;
+    this->lix = lix_local;
+    this->rix = rix_local + this->nx;
+    this->liy = liy_local;
+    this->riy = riy_local + this->ny;
     return;
 }