submission 612829

  ⋯ 18 unchanged lines
      return val;
  }
   
- __global__ void __launch_bounds__(512)
+ __global__ void __launch_bounds__(256)
  vectorsum_kernel(
      const float* __restrict__ input,
-     float* __restrict__ output,
-     double* __restrict__ partial_buf,
-     unsigned int* __restrict__ retire_count,
-     const int n,
-     const int nblocks
+     double* __restrict__ accumulator,
+     const int n
  ) {
-     double sum0 = 0.0;
-     double sum1 = 0.0;
+     double sum0 = 0.0, sum1 = 0.0, sum2 = 0.0, sum3 = 0.0;
   
-     const int tid = threadIdx.x + blockIdx.x * 512;
-     const int grid_stride = 512 * nblocks;
+     const int tid = threadIdx.x + blockIdx.x * 256;
+     const int grid_stride = 256 * gridDim.x;
   
      const int n4 = n >> 2;
      const float4* input4 = reinterpret_cast<const float4*>(input);
   
      int i = tid;
-     for (; i + grid_stride < n4; i += grid_stride * 2) {
+     const int gs4 = grid_stride * 4;
+     for (; i + grid_stride * 3 < n4; i += gs4) {
          float4 v0 = __ldg(&input4[i]);
          float4 v1 = __ldg(&input4[i + grid_stride]);
+         float4 v2 = __ldg(&input4[i + grid_stride * 2]);
+         float4 v3 = __ldg(&input4[i + grid_stride * 3]);
          sum0 += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
          sum1 += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
+         sum2 += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
+         sum3 += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.w;
      }
      for (; i < n4; i += grid_stride) {
          float4 v = __ldg(&input4[i]);
  ⋯ 4 unchanged lines
          sum0 += (double)__ldg(&input[j]);
      }
   
-     double thread_sum = sum0 + sum1;
+     double thread_sum = (sum0 + sum1) + (sum2 + sum3);
      thread_sum = warp_reduce_sum(thread_sum);
   
-     __shared__ double warp_sums[16];
+     __shared__ double warp_sums[8];
      const int lane = threadIdx.x & 31;
      const int warp_id = threadIdx.x >> 5;
   
  ⋯ 1 unchanged lines
      __syncthreads();
   
      if (warp_id == 0) {
-         double val = (lane < 16) ? warp_sums[lane] : 0.0;
+         double val = (lane < 8) ? warp_sums[lane] : 0.0;
          val = warp_reduce_sum(val);
          if (lane == 0) {
-             partial_buf[blockIdx.x] = val;
+             atomicAdd(accumulator, val);
          }
      }
+ }
   
-     __threadfence();
-  
-     __shared__ bool is_last;
-     if (threadIdx.x == 0) {
-         unsigned int ticket = atomicInc(retire_count, 0x7fffffff);
-         is_last = (ticket == (unsigned int)(nblocks - 1));
-     }
-     __syncthreads();
-  
-     if (is_last) {
-         double final_sum = 0.0;
-         for (int j = threadIdx.x; j < nblocks; j += 512) {
-             final_sum += partial_buf[j];
-         }
-         final_sum = warp_reduce_sum(final_sum);
-         if (lane == 0) warp_sums[warp_id] = final_sum;
-         __syncthreads();
-         if (warp_id == 0) {
-             double val = (lane < 16) ? warp_sums[lane] : 0.0;
-             val = warp_reduce_sum(val);
-             if (lane == 0) {
-                 *output = (float)val;
-                 *retire_count = 0;
-             }
-         }
-     }
+ __global__ void finalize_reset(double* __restrict__ acc, float* __restrict__ out) {
+     *out = (float)(*acc);
+     *acc = 0.0;
  }
   
  void vectorsum(torch::Tensor input, torch::Tensor output,
-                torch::Tensor partial_buf, torch::Tensor retire_count,
-                int nblocks) {
-     const int n = input.numel();
-     vectorsum_kernel<<<nblocks, 512>>>(
-         input.data_ptr<float>(), output.data_ptr<float>(),
-         partial_buf.data_ptr<double>(), (unsigned int*)retire_count.data_ptr<int>(),
-         n, nblocks);
+                torch::Tensor accumulator, int nblocks) {
+     vectorsum_kernel<<<nblocks, 256>>>(
+         input.data_ptr<float>(), accumulator.data_ptr<double>(), input.numel());
+     finalize_reset<<<1, 1>>>(accumulator.data_ptr<double>(), output.data_ptr<float>());
  }
  """
   
  cpp_source = """
  void vectorsum(torch::Tensor input, torch::Tensor output,
-                torch::Tensor partial_buf, torch::Tensor retire_count,
-                int nblocks);
+                torch::Tensor accumulator, int nblocks);
  """
   
  module = load_inline(
  ⋯ 5 unchanged lines
      verbose=False,
  )
   
- _partial_buf = torch.zeros(8192, dtype=torch.float64, device="cuda")
- _retire_count = torch.zeros(1, dtype=torch.int32, device="cuda")
+ _accumulator = torch.zeros(1, dtype=torch.float64, device="cuda")
   
   
+ def _autotune():
+     import sys
+     sys.path.insert(0, '.')
+     from reference import generate_input
+  
+     size = 52428800
+     best_time = 1e9
+     best_nb = 4096
+  
+     for nb in [2960, 4096, 5920, 11840]:
+         data = generate_input(size=size, seed=12345)
+         for _ in range(3):
+             module.vectorsum(data[0], data[1], _accumulator, nb)
+             torch.cuda.synchronize()
+         times = []
+         for _ in range(20):
+             data = generate_input(size=size, seed=12345)
+             torch.cuda.synchronize()
+             s = torch.cuda.Event(enable_timing=True)
+             e = torch.cuda.Event(enable_timing=True)
+             s.record()
+             module.vectorsum(data[0], data[1], _accumulator, nb)
+             e.record()
+             torch.cuda.synchronize()
+             times.append(s.elapsed_time(e))
+         avg = sum(times)/len(times)
+         if avg < best_time:
+             best_time = avg
+             best_nb = nb
+     return best_nb
+  
+ _best_nb = _autotune()
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      inp, out = data
      n = inp.numel()
-     # Size-adaptive block count
-     if n <= 2 * 1024 * 1024:
-         nb = 296
-     elif n <= 8 * 1024 * 1024:
-         nb = 592
-     elif n <= 32 * 1024 * 1024:
-         nb = 592
-     elif n <= 128 * 1024 * 1024:
-         nb = 2048
-     else:
-         nb = 4096
-     module.vectorsum(inp, out, _partial_buf, _retire_count, nb)
+     nb = max(1, min(_best_nb, n // 64))
+     module.vectorsum(inp, out, _accumulator, nb)
      return out[0]