submission 780027

  ⋯ 9 unchanged lines
  #include <cuda.h>
  #include <cuda_runtime.h>
   
- // ---------------------------------------------------------------------------
- // Warp-level reduction using shuffle (double precision)
- // ---------------------------------------------------------------------------
  __device__ __forceinline__ double warp_reduce_sum(double val) {
      #pragma unroll
      for (int offset = 16; offset > 0; offset >>= 1)
  ⋯ 1 unchanged lines
      return val;
  }
   
- // ---------------------------------------------------------------------------
- // Phase 1: Each block computes a partial sum into partials[blockIdx.x]
- //   - 4x unrolled float4 loads for ILP
- //   - Warp shuffle + shared mem reduction
- //   - No atomics – just a store per block
- // ---------------------------------------------------------------------------
- __global__ void __launch_bounds__(256, 4)
+ __global__ void __launch_bounds__(256, 8)
  vector_sum_phase1(
      const float* __restrict__ x,
      double*      __restrict__ partials,
  ⋯ 7 unchanged lines
   
      double acc = 0.0;
   
-     // --- Vectorized float4 loads with 4x unroll for ILP ---
-     const int n_vec = n >> 2;  // n / 4
+     const int n_vec = n >> 2;
      const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x);
   
-     // 4x unrolled loop
      const int stride4 = stride * 4;
      int i = gid;
      int limit = n_vec - 3 * stride;
  ⋯ 7 unchanged lines
          acc += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
          acc += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.w;
      }
-     // Cleanup remaining float4s
      for (; i < n_vec; i += stride) {
          float4 v = __ldg(&x4[i]);
          acc += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
      }
   
-     // --- Tail elements ---
      const int tail_base = n_vec << 2;
      for (int j = tail_base + gid; j < n; j += stride)
          acc += (double)__ldg(&x[j]);
   
-     // --- Warp-level reduce ---
      acc = warp_reduce_sum(acc);
   
      const int lane    = tid & 31;
  ⋯ 3 unchanged lines
          sdata[warp_id] = acc;
      __syncthreads();
   
-     // --- Block-level reduce (first warp) ---
-     const int n_warps = blockDim.x >> 5;  // 256/32 = 8
+     const int n_warps = blockDim.x >> 5;
      acc = (tid < n_warps) ? sdata[tid] : 0.0;
      if (warp_id == 0)
          acc = warp_reduce_sum(acc);
  ⋯ 2 unchanged lines
          partials[blockIdx.x] = acc;
  }
   
- // ---------------------------------------------------------------------------
- // Phase 2: Reduce partials array (small – just num_blocks doubles)
- // Single block, single warp is enough for up to ~1024 partials
- // ---------------------------------------------------------------------------
  __global__ void vector_sum_phase2(
      const double* __restrict__ partials,
      float*        __restrict__ out,
  ⋯ 5 unchanged lines
   
      acc = warp_reduce_sum(acc);
   
-     // If more than 1 warp, use shared mem
      __shared__ double sdata[8];
      const int lane = threadIdx.x & 31;
      const int warp_id = threadIdx.x >> 5;
  ⋯ 11 unchanged lines
      }
  }
   
- // ---------------------------------------------------------------------------
- // Host launcher
- // ---------------------------------------------------------------------------
  torch::Tensor vector_sum_cuda(torch::Tensor x) {
      TORCH_CHECK(x.is_cuda() && x.is_contiguous());
      TORCH_CHECK(x.scalar_type() == torch::kFloat32);
   
      const int n = x.numel();
   
-     // Persistent thread config: fill A100's 108 SMs with 4 blocks/SM
      const int THREADS = 256;
-     const int BLOCKS  = 432;  // 108 SMs * 4
-     const int n_warps = THREADS / 32;  // 8
+     const int BLOCKS  = 864;
+     const int n_warps = THREADS / 32;
      const size_t smem = n_warps * sizeof(double);
   
-     // Allocate partials buffer and output
      auto partials = torch::empty({BLOCKS}, x.options().dtype(torch::kFloat64));
      auto out = torch::empty({1}, x.options().dtype(torch::kFloat32));
   
  ⋯ 2 unchanged lines
          partials.data_ptr<double>(),
          n);
   
-     // Phase 2: reduce 432 doubles with 256 threads (1 block)
      vector_sum_phase2<<<1, 256, 8 * sizeof(double)>>>(
          partials.data_ptr<double>(),
          out.data_ptr<float>(),