submission 763182

  ⋯ 5 unchanged lines
  #include <cuda_fp16.h>
  #include <cuda_runtime.h>
   
+ #define MAX_BLOCKS 2048
+  
  static double* d_partial = nullptr;
+ __device__ unsigned int g_retirement_count = 0;
   
  __global__ void __launch_bounds__(512)
- sum_reduce_kernel(const float* __restrict__ input,
-                   double* __restrict__ partial_sums,
-                   int N) {
+ sum_kernel(const float* __restrict__ input,
+            double* __restrict__ partial_sums,
+            float* __restrict__ output,
+            int N) {
      double thread_sum = 0.0;
   
      int idx = blockIdx.x * blockDim.x + threadIdx.x;
  ⋯ 1 unchanged lines
   
      int N4 = N / 4;
      const float4* input4 = reinterpret_cast<const float4*>(input);
-     for (int vec_idx = idx; vec_idx < N4; vec_idx += stride) {
+  
+     int vec_idx = idx;
+     for (; vec_idx + stride < N4; vec_idx += 2 * stride) {
+         float4 v0 = __ldg(&input4[vec_idx]);
+         float4 v1 = __ldg(&input4[vec_idx + stride]);
+         thread_sum += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
+         thread_sum += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
+     }
+     for (; vec_idx < N4; vec_idx += stride) {
          float4 v = __ldg(&input4[vec_idx]);
          thread_sum += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
      }
  ⋯ 17 unchanged lines
      }
      __syncthreads();
   
+     double block_sum = 0.0;
      if (warp_id == 0) {
          int num_warps = blockDim.x >> 5;
          thread_sum = (lane < num_warps) ? warp_sums[lane] : 0.0;
          for (int offset = 16; offset > 0; offset >>= 1) {
              thread_sum += __shfl_down_sync(mask, thread_sum, offset);
          }
-         if (lane == 0) {
-             partial_sums[blockIdx.x] = thread_sum;
-         }
+         block_sum = thread_sum;
      }
- }
   
- __global__ void final_reduce_kernel(const double* __restrict__ partial_sums,
-                                      float* __restrict__ output,
-                                      int num_blocks) {
-     double thread_sum = 0.0;
-     for (int i = threadIdx.x; i < num_blocks; i += blockDim.x) {
-         thread_sum += partial_sums[i];
+     __shared__ bool s_is_last;
+     if (threadIdx.x == 0) {
+         partial_sums[blockIdx.x] = block_sum;
+         __threadfence();
+         unsigned int ticket = atomicAdd(&g_retirement_count, 1);
+         s_is_last = (ticket == gridDim.x - 1);
      }
-  
-     unsigned mask = 0xffffffff;
-     for (int offset = 16; offset > 0; offset >>= 1) {
-         thread_sum += __shfl_down_sync(mask, thread_sum, offset);
-     }
-  
-     __shared__ double warp_sums[8];
-     int lane = threadIdx.x & 31;
-     int warp_id = threadIdx.x >> 5;
-  
-     if (lane == 0) {
-         warp_sums[warp_id] = thread_sum;
-     }
      __syncthreads();
   
-     if (warp_id == 0) {
-         int num_warps = blockDim.x >> 5;
-         thread_sum = (lane < num_warps) ? warp_sums[lane] : 0.0;
+     if (s_is_last) {
+         double final_sum = 0.0;
+         for (int i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+             final_sum += partial_sums[i];
+         }
+  
          for (int offset = 16; offset > 0; offset >>= 1) {
-             thread_sum += __shfl_down_sync(mask, thread_sum, offset);
+             final_sum += __shfl_down_sync(mask, final_sum, offset);
          }
+  
          if (lane == 0) {
-             output[0] = (float)thread_sum;
+             warp_sums[warp_id] = final_sum;
          }
+         __syncthreads();
+  
+         if (warp_id == 0) {
+             int num_warps = blockDim.x >> 5;
+             final_sum = (lane < num_warps) ? warp_sums[lane] : 0.0;
+             for (int offset = 16; offset > 0; offset >>= 1) {
+                 final_sum += __shfl_down_sync(mask, final_sum, offset);
+             }
+             if (lane == 0) {
+                 output[0] = (float)final_sum;
+                 g_retirement_count = 0;
+             }
+         }
      }
  }
   
  void sum_cuda(torch::Tensor input, torch::Tensor output) {
      int N = input.numel();
      const int threads = 512;
-     int blocks = min((N / 4 + threads - 1) / threads, 1024);
+     int blocks = min((N / 4 + threads - 1) / threads, MAX_BLOCKS);
      if (blocks < 1) blocks = 1;
   
      if (!d_partial) {
-         cudaMalloc(&d_partial, 1024 * sizeof(double));
+         cudaMalloc(&d_partial, MAX_BLOCKS * sizeof(double));
      }
   
-     sum_reduce_kernel<<<blocks, threads>>>(
+     sum_kernel<<<blocks, threads>>>(
          input.data_ptr<float>(),
          d_partial,
+         output.data_ptr<float>(),
          N
      );
-  
-     final_reduce_kernel<<<1, 256>>>(
-         d_partial,
-         output.data_ptr<float>(),
-         blocks
-     );
  }
  """