submission 614161

  ⋯ 10 unchanged lines
  #include <cuda_runtime.h>
  #include <torch/extension.h>
   
- __device__ __forceinline__ double warp_reduce_sum(double val) {
-     #pragma unroll
-     for (int offset = 16; offset > 0; offset >>= 1) {
-         val += __shfl_down_sync(0xffffffff, val, offset);
-     }
+ __device__ __forceinline__ double warp_reduce(double val) {
+     val += __shfl_down_sync(0xffffffff, val, 16);
+     val += __shfl_down_sync(0xffffffff, val, 8);
+     val += __shfl_down_sync(0xffffffff, val, 4);
+     val += __shfl_down_sync(0xffffffff, val, 2);
+     val += __shfl_down_sync(0xffffffff, val, 1);
      return val;
  }
   
  ⋯ 8 unchanged lines
  __global__ void __launch_bounds__(256)
  vectorsum_kernel(
      const float* __restrict__ input,
-     double* __restrict__ accumulator,
+     double* __restrict__ acc,
      const int n
  ) {
-     double sum0 = 0.0, sum1 = 0.0, sum2 = 0.0, sum3 = 0.0;
-  
+     double s0 = 0.0, s1 = 0.0, s2 = 0.0, s3 = 0.0;
      const int tid = threadIdx.x + blockIdx.x * 256;
-     const int grid_stride = 256 * gridDim.x;
-  
+     const int gs = 256 * gridDim.x;
      const int n4 = n >> 2;
-     const float4* input4 = reinterpret_cast<const float4*>(input);
+     const float4* in4 = reinterpret_cast<const float4*>(input);
   
      int i = tid;
-     const int gs4 = grid_stride * 4;
-     for (; i + grid_stride * 3 < n4; i += gs4) {
-         float4 v0 = load_cs(&input4[i]);
-         float4 v1 = load_cs(&input4[i + grid_stride]);
-         float4 v2 = load_cs(&input4[i + grid_stride * 2]);
-         float4 v3 = load_cs(&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 + gs * 3 < n4; i += gs * 4) {
+         float4 v0 = load_cs(&in4[i]);
+         float4 v1 = load_cs(&in4[i + gs]);
+         float4 v2 = load_cs(&in4[i + gs * 2]);
+         float4 v3 = load_cs(&in4[i + gs * 3]);
+         s0 += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
+         s1 += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
+         s2 += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
+         s3 += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.w;
      }
-     for (; i < n4; i += grid_stride) {
-         float4 v = load_cs(&input4[i]);
-         sum0 += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
+     for (; i < n4; i += gs) {
+         float4 v = load_cs(&in4[i]);
+         s0 += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
      }
+     for (int j = (n4 << 2) + tid; j < n; j += gs)
+         s0 += (double)__ldg(&input[j]);
   
-     for (int j = (n4 << 2) + tid; j < n; j += grid_stride) {
-         sum0 += (double)__ldg(&input[j]);
-     }
+     double ts = (s0 + s1) + (s2 + s3);
+     ts = warp_reduce(ts);
   
-     double thread_sum = (sum0 + sum1) + (sum2 + sum3);
-     thread_sum = warp_reduce_sum(thread_sum);
-  
-     __shared__ double warp_sums[8];
+     __shared__ double ws[8];
      const int lane = threadIdx.x & 31;
-     const int warp_id = threadIdx.x >> 5;
-  
-     if (lane == 0) warp_sums[warp_id] = thread_sum;
+     const int wid = threadIdx.x >> 5;
+     if (lane == 0) ws[wid] = ts;
      __syncthreads();
-  
-     if (warp_id == 0) {
-         double val = (lane < 8) ? warp_sums[lane] : 0.0;
-         val = warp_reduce_sum(val);
-         if (lane == 0) {
-             atomicAdd(accumulator, val);
-         }
+     if (wid == 0) {
+         double v = (lane < 8) ? ws[lane] : 0.0;
+         v = warp_reduce(v);
+         if (lane == 0) atomicAdd(acc, v);
      }
  }
   
- __global__ void finalize_reset(double* __restrict__ acc, float* __restrict__ out) {
+ __global__ void finalize(double* __restrict__ acc, float* __restrict__ out) {
      *out = (float)(*acc);
      *acc = 0.0;
  }
   
  void vectorsum(torch::Tensor input, torch::Tensor output,
-                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>());
+                torch::Tensor acc, int nb) {
+     vectorsum_kernel<<<nb, 256>>>(
+         input.data_ptr<float>(), acc.data_ptr<double>(), input.numel());
+     finalize<<<1, 1>>>(acc.data_ptr<double>(), output.data_ptr<float>());
  }
  """
   
- cpp_source = """
- void vectorsum(torch::Tensor input, torch::Tensor output,
-                torch::Tensor accumulator, int nblocks);
- """
+ cpp_source = "void vectorsum(torch::Tensor, torch::Tensor, torch::Tensor, int);"
   
  module = load_inline(
      name="vectorsum_cuda",
      cpp_sources=[cpp_source],
      cuda_sources=[cuda_source],
      functions=["vectorsum"],
-     extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
+     extra_cuda_cflags=["-O3", "--use_fast_math"],
      verbose=False,
  )
   
- _accumulator = torch.zeros(1, dtype=torch.float64, device="cuda")
+ _acc = torch.zeros(1, dtype=torch.float64, device="cuda")
   
   
- def _autotune():
+ def _tune():
      import sys
      sys.path.insert(0, '.')
      from reference import generate_input
-  
      size = 52428800
-     best_time = 1e9
-     best_nb = 5920
-  
-     for nb in [2960, 4096, 5920, 8192]:
+     best_t, best_nb = 1e9, 5920
+     for nb in [2960, 4096, 5920]:
          data = generate_input(size=size, seed=12345)
          for _ in range(3):
-             module.vectorsum(data[0], data[1], _accumulator, nb)
+             module.vectorsum(data[0], data[1], _acc, nb)
              torch.cuda.synchronize()
-         times = []
-         for _ in range(20):
+         ts = []
+         for _ in range(15):
              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
+             s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+             s.record(); module.vectorsum(data[0], data[1], _acc, nb); e.record()
+             torch.cuda.synchronize(); ts.append(s.elapsed_time(e))
+         a = sum(ts)/len(ts)
+         if a < best_t: best_t, best_nb = a, nb
      return best_nb
   
- _best_nb = _autotune()
+ _nb = _tune()
   
   
  def custom_kernel(data: input_t) -> output_t:
      inp, out = data
      n = inp.numel()
-     nb = max(1, min(_best_nb, n // 64))
-     module.vectorsum(inp, out, _accumulator, nb)
+     nb = max(1, min(_nb, n // 64))
+     module.vectorsum(inp, out, _acc, nb)
      return out[0]