add stress_ewa_op for cpu

Author: jieli-matrix

source/module_hamilt_pw/hamilt_pwdft/kernels/stress_op.cpp

@@ -759,6 +759,65 @@ template struct cal_force_npw_op<double, base_device::DEVICE_CPU>;
 template struct cal_multi_dot_op<float, base_device::DEVICE_CPU>;
 template struct cal_multi_dot_op<double, base_device::DEVICE_CPU>;
 
+// CPU implementation of Ewald stress sincos operator
+template <typename FPTYPE>
+struct cal_stress_ewa_sincos_op<FPTYPE, base_device::DEVICE_CPU>
+{
+    void operator()(const base_device::DEVICE_CPU* ctx,
+                    const int& nat,
+                    const int& npw,
+                    const int& ig_gge0,
+                    const FPTYPE* gcar,
+                    const FPTYPE* tau,
+                    const FPTYPE* zv_facts,
+                    FPTYPE* rhostar_real,
+                    FPTYPE* rhostar_imag)
+    {
+        const FPTYPE TWO_PI = 2.0 * M_PI;
+
+        // Initialize output arrays
+        std::fill(rhostar_real, rhostar_real + npw, static_cast<FPTYPE>(0.0));
+        std::fill(rhostar_imag, rhostar_imag + npw, static_cast<FPTYPE>(0.0));
+
+#ifdef _OPENMP
+#pragma omp parallel for
+#endif
+        for (int ig = 0; ig < npw; ig++) {
+            if (ig == ig_gge0) continue; // Skip G=0
+
+            FPTYPE local_rhostar_real = 0.0;
+            FPTYPE local_rhostar_imag = 0.0;
+
+            // Double loop: iat -> ig (as requested)
+            for (int iat = 0; iat < nat; iat++) {
+                const FPTYPE tau_x = tau[iat * 3 + 0];
+                const FPTYPE tau_y = tau[iat * 3 + 1];
+                const FPTYPE tau_z = tau[iat * 3 + 2];
+                const FPTYPE zv = zv_facts[iat];
+
+                // Calculate phase: 2π * (G · τ) - similar to cal_force_ewa phase
+                const FPTYPE phase = TWO_PI * (gcar[ig * 3 + 0] * tau_x +
+                                               gcar[ig * 3 + 1] * tau_y +
+                                               gcar[ig * 3 + 2] * tau_z);
+
+                // Calculate sincos
+                FPTYPE sinp, cosp;
+                ModuleBase::libm::sincos(phase, &sinp, &cosp);
+
+                // Accumulate structure factor
+                local_rhostar_real += zv * cosp;
+                local_rhostar_imag += zv * sinp;
+            }
+
+            // Store results
+            rhostar_real[ig] = local_rhostar_real;
+            rhostar_imag[ig] = local_rhostar_imag;
+        }
+    }
+};
+
+template struct cal_stress_ewa_sincos_op<float, base_device::DEVICE_CPU>;
+template struct cal_stress_ewa_sincos_op<double, base_device::DEVICE_CPU>;
 
 // template struct prepare_vkb_deri_ptr_op<float, base_device::DEVICE_CPU>;
 // template struct prepare_vkb_deri_ptr_op<double, base_device::DEVICE_CPU>;

source/module_hamilt_pw/hamilt_pwdft/kernels/stress_op.h

@@ -260,6 +260,35 @@ struct cal_multi_dot_op{
                     const std::complex<FPTYPE>* psi);
 };
 
+template <typename FPTYPE, typename Device>
+struct cal_stress_ewa_sincos_op
+{
+    /// @brief Calculate Ewald stress sincos computation
+    /// Only computes rhostar, other calculations remain in original function
+    ///
+    /// Input Parameters
+    /// @param ctx - which device this function runs on
+    /// @param nat - total number of atoms
+    /// @param npw - number of plane waves
+    /// @param ig_gge0 - index of G=0 vector (-1 if not present)
+    /// @param gcar - G-vector Cartesian coordinates [npw * 3]
+    /// @param tau - atomic positions [nat * 3]
+    /// @param zv_facts - precomputed zv factors for each atom [nat]
+    ///
+    /// Output Parameters
+    /// @param rhostar_real - real part of structure factor [npw]
+    /// @param rhostar_imag - imaginary part of structure factor [npw]
+    void operator()(const Device* ctx,
+                    const int& nat,
+                    const int& npw,
+                    const int& ig_gge0,
+                    const FPTYPE* gcar,
+                    const FPTYPE* tau,
+                    const FPTYPE* zv_facts,
+                    FPTYPE* rhostar_real,
+                    FPTYPE* rhostar_imag);
+};
+
 
 #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
 template <typename FPTYPE>
@@ -434,6 +463,20 @@ struct cal_multi_dot_op<FPTYPE, base_device::DEVICE_GPU>{
                     const std::complex<FPTYPE>* psi);
 };
 
+template <typename FPTYPE>
+struct cal_stress_ewa_sincos_op<FPTYPE, base_device::DEVICE_GPU>
+{
+    void operator()(const base_device::DEVICE_GPU* ctx,
+                    const int& nat,
+                    const int& npw,
+                    const int& ig_gge0,
+                    const FPTYPE* gcar,
+                    const FPTYPE* tau,
+                    const FPTYPE* zv_facts,
+                    FPTYPE* rhostar_real,
+                    FPTYPE* rhostar_imag);
+};
+
 /**
  * The operator is used to compute the auxiliary amount of stress /force 
  * in parallel on the GPU. They identify type with the type provided and 

source/module_hamilt_pw/hamilt_pwdft/stress_func_ewa.cpp

@@ -4,6 +4,7 @@
 #include "module_base/tool_threading.h"
 #include "module_base/libm/libm.h"
 #include "module_hamilt_pw/hamilt_pwdft/global.h"
+#include "kernels/stress_op.h"
 
 #ifdef _OPENMP
 #include <omp.h>
@@ -56,6 +57,45 @@ void Stress_Func<FPTYPE, Device>::stress_ewa(const UnitCell& ucell,
     if (PARAM.globalv.gamma_only_pw && is_pw) fact=2.0;
 //    else fact=1.0;
 
+    // Prepare data for the sincos op
+    std::vector<FPTYPE> zv_facts_host(ucell.nat);
+    std::vector<FPTYPE> tau_flat(ucell.nat * 3);
+
+    for (int iat = 0; iat < ucell.nat; iat++) {
+        int it = ucell.iat2it[iat];
+        int ia = ucell.iat2ia[iat];
+
+        zv_facts_host[iat] = static_cast<FPTYPE>(ucell.atoms[it].ncpp.zv);
+
+        tau_flat[iat * 3 + 0] = static_cast<FPTYPE>(ucell.atoms[it].tau[ia][0]);
+        tau_flat[iat * 3 + 1] = static_cast<FPTYPE>(ucell.atoms[it].tau[ia][1]);
+        tau_flat[iat * 3 + 2] = static_cast<FPTYPE>(ucell.atoms[it].tau[ia][2]);
+    }
+
+    std::vector<FPTYPE> gcar_flat(rho_basis->npw * 3);
+    for (int ig = 0; ig < rho_basis->npw; ig++) {
+        gcar_flat[ig * 3 + 0] = static_cast<FPTYPE>(rho_basis->gcar[ig][0]);
+        gcar_flat[ig * 3 + 1] = static_cast<FPTYPE>(rho_basis->gcar[ig][1]);
+        gcar_flat[ig * 3 + 2] = static_cast<FPTYPE>(rho_basis->gcar[ig][2]);
+    }
+
+    // Allocate result arrays
+    std::vector<FPTYPE> rhostar_real_host(rho_basis->npw);
+    std::vector<FPTYPE> rhostar_imag_host(rho_basis->npw);
+
+    // Call sincos op (outside OpenMP parallel region, op has its own parallelization)
+    hamilt::cal_stress_ewa_sincos_op<FPTYPE, Device>()(
+        this->ctx,
+        ucell.nat,
+        rho_basis->npw,
+        rho_basis->ig_gge0,
+        gcar_flat.data(),
+        tau_flat.data(),
+        zv_facts_host.data(),
+        rhostar_real_host.data(),
+        rhostar_imag_host.data()
+    );
+
 #ifdef _OPENMP
 #pragma omp parallel
 {
@@ -76,27 +116,21 @@ void Stress_Func<FPTYPE, Device>::stress_ewa(const UnitCell& ucell,
 	ig_end = ig + ig_end;
 
     FPTYPE g2,g2a;
-    FPTYPE arg;
-    std::complex<FPTYPE> rhostar;
     FPTYPE sewald;
     for(; ig < ig_end; ig++)
 	{
 		if(ig == ig0)  continue;
 		g2 = rho_basis->gg[ig]* ucell.tpiba2;
 		g2a = g2 /4.0/alpha;
-		rhostar=std::complex<FPTYPE>(0.0,0.0);
-		for(int it=0; it < ucell.ntype; it++)
-		{
-			for(int i=0; i<ucell.atoms[it].na; i++)
-			{
-				arg = (rho_basis->gcar[ig] * ucell.atoms[it].tau[i]) * (ModuleBase::TWO_PI);
-				FPTYPE sinp, cosp;
-                ModuleBase::libm::sincos(arg, &sinp, &cosp);
-				rhostar = rhostar + std::complex<FPTYPE>(ucell.atoms[it].ncpp.zv * cosp,ucell.atoms[it].ncpp.zv * sinp);
-			}
-		}
-		rhostar /= ucell.omega;
-		sewald = fact* (ModuleBase::TWO_PI) * ModuleBase::e2 * ModuleBase::libm::exp(-g2a) / g2 * pow(std::abs(rhostar),2);
+
+		// Use precomputed rhostar values
+		FPTYPE rhostar_real = rhostar_real_host[ig] / ucell.omega;
+		FPTYPE rhostar_imag = rhostar_imag_host[ig] / ucell.omega;
+
+		// Calculate |rhostar|² - mathematically equivalent to pow(std::abs(rhostar), 2)
+		FPTYPE rhostar_abs2 = rhostar_real * rhostar_real + rhostar_imag * rhostar_imag;
+
+		sewald = fact* (ModuleBase::TWO_PI) * ModuleBase::e2 * ModuleBase::libm::exp(-g2a) / g2 * rhostar_abs2;
 		local_sdewald -= sewald;
 		for(int l=0;l<3;l++)
 		{