From 63bb4b1a93cebc4a08f1e253b74edfa7ccb91547 Mon Sep 17 00:00:00 2001
From: Taillefumier Mathieu <29380261+mtaillefumier@users.noreply.github.com>
Date: Tue, 23 Jun 2026 19:45:14 +0000
Subject: [PATCH] Fix a race condition in the grid backend (#5437)

Co-authored-by: Mathieu Taillefumier <mathieu.taillefumier@free.fr>
---
 src/grid/gpu/grid_gpu_collocate.cu |  1 +
 src/grid/gpu/grid_gpu_integrate.cu | 17 +++++++++++++----
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/grid/gpu/grid_gpu_collocate.cu b/src/grid/gpu/grid_gpu_collocate.cu
index 256d293017..06206baa45 100644
--- a/src/grid/gpu/grid_gpu_collocate.cu
+++ b/src/grid/gpu/grid_gpu_collocate.cu
@@ -110,6 +110,7 @@ __launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) {
     __syncthreads();
 
     cab_to_cxyz(task, smem_alpha, smem_cab, coef_);
+    __syncthreads();
   }
 }
 
diff --git a/src/grid/gpu/grid_gpu_integrate.cu b/src/grid/gpu/grid_gpu_integrate.cu
index 9f4de3985a..33446db940 100644
--- a/src/grid/gpu/grid_gpu_integrate.cu
+++ b/src/grid/gpu/grid_gpu_integrate.cu
@@ -86,6 +86,7 @@ __global__ __launch_bounds__(64) void compute_hab_v4(const kernel_params dev_) {
 
     compute_alpha(task, smem_alpha);
     __syncthreads();
+
     cxyz_to_cab(task, smem_alpha, coef_, smem_cab);
     __syncthreads();
 
@@ -112,6 +113,9 @@ __global__ __launch_bounds__(64) void compute_hab_v4(const kernel_params dev_) {
         }
       }
     }
+    // all warps need to be synchronized here before modifying the task
+    // information.
+    __syncthreads();
   }
 }
 
@@ -173,6 +177,9 @@ __global__ __launch_bounds__(64) void compute_hab_v2(const kernel_params dev_) {
         T *shared_virial = &shared_memory[6];
 
         if (CALCULATE_FORCES) {
+          // Important to synchronize the warps
+          __syncthreads();
+
           if (tid < 3) {
             shared_forces_a[tid] = get_force_a<COMPUTE_TAU, T>(
                 a, b, tid, task.zeta, task.zetb, task.n1, smem_cab);
@@ -187,8 +194,8 @@ __global__ __launch_bounds__(64) void compute_hab_v2(const kernel_params dev_) {
                                              task.zetb, task.rab, task.n1,
                                              smem_cab);
           }
+          __syncthreads();
         }
-        __syncthreads();
 
         T block_val1 = 0.0;
         // Try to use as many contiguous threads as possible and limit the
@@ -231,8 +238,12 @@ __global__ __launch_bounds__(64) void compute_hab_v2(const kernel_params dev_) {
         }
       }
     }
+    __syncthreads();
   }
 
+  // theoretically not needed
+  __syncthreads();
+
   if (CALCULATE_FORCES) {
     const int task_id = dev_.task_sorted_by_blocks_dev[offset];
     const auto &glb_task = dev_.tasks[task_id];
@@ -617,6 +628,7 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
     if (tid < min(length - ico, lbatch))
       dev_.ptr_dev[2][dev_.tasks[dev_.first_task + block_index()].coef_offset +
                       tid + ico] = sum[tid];
+    __syncthreads();
   }
 }
 
@@ -692,9 +704,6 @@ void context_info::compute_hab_coefficients() {
     compute_hab_v4<double, false>
         <<<this->nblocks, threads_per_block, smem_params.smem_per_block(),
            this->main_stream>>>(params);
-    // compute_hab_v2<double, double3, false, false>
-    //   <<<this->nblocks, threads_per_block, smem_params.smem_per_block(),
-    //   this->main_stream>>>(params);
     return;
   }