Reallocate energyvector reduction storage when it increases

GOMC-WSU · Jul 27, 2024 · 7151453 · 7151453
1 parent 09ef528
commit 7151453
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 15 deletions.
diff --git a/src/GPU/CalculateEnergyCUDAKernel.cu b/src/GPU/CalculateEnergyCUDAKernel.cu
@@ -104,29 +104,18 @@ void CallBoxInterGPU(VariablesCUDA *vars, const std::vector<int> &cellVector,
 #endif
 
   // ReduceSum
-  void *d_temp_storage = NULL;
-  size_t temp_storage_bytes = 0;
-  // LJ ReduceSum
-  DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, vars->gpu_LJEn,
+  DeviceReduce::Sum(vars->cub_energyVec_storage, vars->cub_energyVec_storage_size, vars->gpu_LJEn,
                     vars->gpu_finalVal, energyVectorLen);
-  CubDebugExit(CUMALLOC(&d_temp_storage, temp_storage_bytes));
-  DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, vars->gpu_LJEn,
-                    vars->gpu_finalVal, energyVectorLen);
-  // Copy back the result to CPU ! :)
-  CubDebugExit(cudaMemcpy(&LJEn, vars->gpu_finalVal, sizeof(double),
-                          cudaMemcpyDeviceToHost));
+  cudaMemcpy(&LJEn, vars->gpu_finalVal, sizeof(double), cudaMemcpyDeviceToHost);
   if (electrostatic) {
     // Real Term ReduceSum
-    DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, vars->gpu_REn,
+    DeviceReduce::Sum(vars->cub_energyVec_storage, vars->cub_energyVec_storage_size, vars->gpu_REn,
                       vars->gpu_finalVal, energyVectorLen);
-    // Copy back the result to CPU ! :)
-    CubDebugExit(cudaMemcpy(&REn, vars->gpu_finalVal, sizeof(double),
-                            cudaMemcpyDeviceToHost));
+    cudaMemcpy(&REn, vars->gpu_finalVal, sizeof(double), cudaMemcpyDeviceToHost);
   } else {
     REn = 0.0;
   }
 
-  CUFREE(d_temp_storage);
   CUFREE(gpu_particleKind);
   CUFREE(gpu_particleMol);
   CUFREE(gpu_neighborList);

diff --git a/src/GPU/ConstantDefinitionsCUDAKernel.cu b/src/GPU/ConstantDefinitionsCUDAKernel.cu
@@ -321,6 +321,21 @@ void UpdateEnergyVecs(VariablesCUDA *vars, int newVecLen, bool electrostatic) {
   if (electrostatic) {
     CUMALLOC((void **)&vars->gpu_REn, vars->gpu_energyVecLen * sizeof(double));
   }
+
+  // Check if more temporary storage is needed for this larger reduction size.
+  // If so, free and malloc a new array for the additional space.
+  void *d_temp_storage = nullptr;
+  size_t temp_storage_bytes = 0;
+  cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, vars->gpu_LJEn,
+                         vars->gpu_finalVal, vars->gpu_energyVecLen);
+  if (temp_storage_bytes > vars->cub_energyVec_storage_size) {
+    // Free the current allocation if this isn't the first allocation
+    if (vars->cub_energyVec_storage_size > 0) {
+      CUFREE(vars->cub_energyVec_storage);
+    }
+    vars->cub_energyVec_storage_size = temp_storage_bytes;
+    CUMALLOC(&(vars->cub_energyVec_storage), vars->cub_energyVec_storage_size);
+  }
 }
 
 void DestroyEwaldCUDAVars(VariablesCUDA *vars) {
@@ -366,6 +381,7 @@ void DestroyExp6CUDAVars(VariablesCUDA *vars) {
 }
 
 void DestroyCUDAVars(VariablesCUDA *vars) {
+  CUFREE(vars->cub_energyVec_storage);
   CUFREE(vars->gpu_sigmaSq);
   CUFREE(vars->gpu_epsilon_Cn);
   CUFREE(vars->gpu_n);

diff --git a/src/GPU/VariablesCUDA.cuh b/src/GPU/VariablesCUDA.cuh
@@ -64,6 +64,8 @@ public:
   VariablesCUDA() {
     cub_reduce_storage_size = 0;
     cub_reduce_storage = nullptr;
+    cub_energyVec_storage_size = 0;
+    cub_energyVec_storage = nullptr;
     gpu_sigmaSq = nullptr;
     gpu_epsilon_Cn = nullptr;
     gpu_n = nullptr;
@@ -99,6 +101,8 @@ public:
 
   size_t cub_reduce_storage_size;
   void *cub_reduce_storage;
+  size_t cub_energyVec_storage_size;
+  void *cub_energyVec_storage;
   double *gpu_sigmaSq;
   double *gpu_epsilon_Cn;
   double *gpu_n;