[CINN] Fix block reduce CUDA template (#71471)

lshpku · web-flow · commit f0112fcac7c1 · 2025-03-07T11:17:03.000+08:00
diff --git a/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh b/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh
@@ -572,7 +572,10 @@ EXPAND_REDUCE_FP16_MACRO(CINN_WARP_SHUFFLE_INTERNAL_IMPL)
   }                                                                \
   __syncthreads();                                                 \
   if (threadIdx.x < (blockDim.x + 31) / 32) {                      \
-    shm[0] = cinn_warp_shuffle_internal(shm[threadIdx.x]);         \
+    tmp_val = cinn_warp_shuffle_internal(shm[threadIdx.x]);        \
+    if (threadIdx.x == 0) {                                        \
+      shm[0] = tmp_val;                                            \
+    }                                                              \
   }                                                                \
   __syncthreads();                                                 \
   return shm[0];
