fix GPU ut error (deepmodeling#3186)

YuLiu98 · Nov 10, 2023 · f63b98e · f63b98e
1 parent bf64879
commit f63b98e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/source/module_hamilt_pw/hamilt_pwdft/kernels/test/force_op_test.cpp b/source/module_hamilt_pw/hamilt_pwdft/kernels/test/force_op_test.cpp
@@ -241,12 +241,17 @@ TEST_F(TestSrcPWForceMultiDevice, cal_force_nl_op_gpu)
 {
     std::vector<double> res(expected_force.size(), 0);
     double * d_res = nullptr, * d_wg = nullptr, * d_deeq = nullptr;
+    double * d_ekb = nullptr, * d_qq_nt = nullptr;
     resmem_var_op()(gpu_ctx, d_wg, wg.size());
     resmem_var_op()(gpu_ctx, d_res, res.size());
     resmem_var_op()(gpu_ctx, d_deeq, deeq.size());
+    resmem_var_op()(gpu_ctx, d_ekb, ekb.size());
+    resmem_var_op()(gpu_ctx, d_qq_nt, qq_nt.size());
     syncmem_var_h2d_op()(gpu_ctx, cpu_ctx, d_wg, wg.data(), wg.size());
     syncmem_var_h2d_op()(gpu_ctx, cpu_ctx, d_res, res.data(), res.size());
     syncmem_var_h2d_op()(gpu_ctx, cpu_ctx, d_deeq, deeq.data(), deeq.size());
+    syncmem_var_h2d_op()(gpu_ctx, cpu_ctx, d_ekb, ekb.data(), ekb.size());
+    syncmem_var_h2d_op()(gpu_ctx, cpu_ctx, d_qq_nt, qq_nt.data(), qq_nt.size());
 
     int * d_atom_nh = nullptr, * d_atom_na = nullptr;
     resmem_int_op()(gpu_ctx, d_atom_nh, atom_nh.size());
@@ -279,6 +284,8 @@ TEST_F(TestSrcPWForceMultiDevice, cal_force_nl_op_gpu)
         d_atom_na,
         tpiba,
         d_wg,
+        d_ekb,
+        d_qq_nt,
         d_deeq,
         d_becp,
         d_dbecp,
@@ -292,6 +299,8 @@ TEST_F(TestSrcPWForceMultiDevice, cal_force_nl_op_gpu)
     delmem_var_op()(gpu_ctx, d_wg);
     delmem_var_op()(gpu_ctx, d_res);
     delmem_var_op()(gpu_ctx, d_deeq);
+    delmem_var_op()(gpu_ctx, d_ekb);
+    delmem_var_op()(gpu_ctx, d_qq_nt);
 
     delmem_int_op()(gpu_ctx, d_atom_nh);
     delmem_int_op()(gpu_ctx, d_atom_na);

diff --git a/source/module_hamilt_pw/hamilt_pwdft/kernels/test/stress_op_test.cpp b/source/module_hamilt_pw/hamilt_pwdft/kernels/test/stress_op_test.cpp
@@ -235,6 +235,7 @@ TEST(TestSrcPWStressMultiDevice, cal_stress_nl_op_gpu)
     std::vector<double> stress(expected_stress.size(), 0);
     std::complex<double> * d_becp = nullptr, * d_dbecp = nullptr;
     double * d_wg = nullptr, * d_deeq = nullptr, * d_stress = nullptr;
+    double * d_ekb = nullptr, * d_qq_nt = nullptr;
     int * d_atom_nh = nullptr, * d_atom_na = nullptr;
     resmem_zd_op()(gpu_ctx, d_becp, becp.size());
     resmem_zd_op()(gpu_ctx, d_dbecp, dbecp.size());
@@ -244,9 +245,13 @@ TEST(TestSrcPWStressMultiDevice, cal_stress_nl_op_gpu)
     resmem_dd_op()(gpu_ctx, d_wg, wg.size());
     resmem_dd_op()(gpu_ctx, d_deeq, deeq.size());
     resmem_dd_op()(gpu_ctx, d_stress, stress.size());
+    resmem_dd_op()(gpu_ctx, d_ekb, ekb.size());
+    resmem_dd_op()(gpu_ctx, d_qq_nt, qq_nt.size());
     syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, d_wg, wg.data(), wg.size());
     syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, d_deeq, deeq.data(), deeq.size());
     syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, d_stress, stress.data(), stress.size());
+    syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, d_ekb, ekb.data(), ekb.size());
+    syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, d_qq_nt, qq_nt.data(), qq_nt.size());
 
     using delmem_int_op = psi::memory::delete_memory_op<int, psi::DEVICE_GPU>;
     using resmem_int_op = psi::memory::resize_memory_op<int, psi::DEVICE_GPU>;
@@ -274,6 +279,8 @@ TEST(TestSrcPWStressMultiDevice, cal_stress_nl_op_gpu)
             d_atom_nh,
             d_atom_na,
             d_wg,
+            d_ekb,
+            d_qq_nt,
             d_deeq,
             d_becp,
             d_dbecp,
@@ -291,6 +298,8 @@ TEST(TestSrcPWStressMultiDevice, cal_stress_nl_op_gpu)
     delmem_dd_op()(gpu_ctx, d_wg);
     delmem_dd_op()(gpu_ctx, d_deeq);
     delmem_dd_op()(gpu_ctx, d_stress);
+    delmem_dd_op()(gpu_ctx, d_ekb);
+    delmem_dd_op()(gpu_ctx, d_qq_nt);
 
     delmem_int_op()(gpu_ctx, d_atom_nh);
     delmem_int_op()(gpu_ctx, d_atom_na);