fix multi-GPU without direct access

brucefan1983 · Jan 30, 2025 · b6c977a · b6c977a
1 parent fdb9608
commit b6c977a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/src/main_nep/nep_charge.cu b/src/main_nep/nep_charge.cu
@@ -256,12 +256,6 @@ void NEP_Charge::find_k1k2k3()
       }
     }
   }
-  charge_para.k1_gpu.resize(charge_para.k1.size());
-  charge_para.k2_gpu.resize(charge_para.k2.size());
-  charge_para.k3_gpu.resize(charge_para.k3.size());
-  charge_para.k1_gpu.copy_from_host(charge_para.k1.data());
-  charge_para.k2_gpu.copy_from_host(charge_para.k2.data());
-  charge_para.k3_gpu.copy_from_host(charge_para.k3.data());
 }
 
 NEP_Charge::NEP_Charge(
@@ -348,6 +342,12 @@ NEP_Charge::NEP_Charge(
     nep_data[device_id].S_real.resize(Nc * charge_para.num_kpoints);
     nep_data[device_id].S_imag.resize(Nc * charge_para.num_kpoints);
     nep_data[device_id].D_real.resize(N);
+    nep_data[device_id].k1_gpu.resize(charge_para.k1.size());
+    nep_data[device_id].k2_gpu.resize(charge_para.k2.size());
+    nep_data[device_id].k3_gpu.resize(charge_para.k3.size());
+    nep_data[device_id].k1_gpu.copy_from_host(charge_para.k1.data());
+    nep_data[device_id].k2_gpu.copy_from_host(charge_para.k2.data());
+    nep_data[device_id].k3_gpu.copy_from_host(charge_para.k3.data());
   }
 }
 
@@ -1162,9 +1162,9 @@ void NEP_Charge::find_force(
       charge_para.num_kpoints,
       charge_para.alpha_factor,
       dataset[device_id].box_original.data(),
-      charge_para.k1_gpu.data(),
-      charge_para.k2_gpu.data(),
-      charge_para.k3_gpu.data(),
+      nep_data[device_id].k1_gpu.data(),
+      nep_data[device_id].k2_gpu.data(),
+      nep_data[device_id].k3_gpu.data(),
       nep_data[device_id].kx.data(),
       nep_data[device_id].ky.data(),
       nep_data[device_id].kz.data(),

diff --git a/src/main_nep/nep_charge.cuh b/src/main_nep/nep_charge.cuh
@@ -81,6 +81,9 @@ public:
     GPU_Vector<float> S_real;
     GPU_Vector<float> S_imag;
     GPU_Vector<float> D_real;
+    GPU_Vector<int> k1_gpu;
+    GPU_Vector<int> k2_gpu;
+    GPU_Vector<int> k3_gpu;
   };
 
   struct Charge_Para {
@@ -89,9 +92,6 @@ public:
     std::vector<int> k1;
     std::vector<int> k2;
     std::vector<int> k3;
-    GPU_Vector<int> k1_gpu;
-    GPU_Vector<int> k2_gpu;
-    GPU_Vector<int> k3_gpu;
     float alpha = 0.5f; // 1 / (2 Angstrom)
     float alpha_factor = 1.0f; // 1 / (4 * alpha * alpha)
     float two_alpha_over_sqrt_pi = 0.564189583547756f;