submission 66718

  ⋯ 4 unchanged lines
  from task import input_t, output_t
   
  sum_cuda_source = """
- #define BLOCK_SIZE 1024
- template <typename scalar_t>
- __global__ void sum_kernel(const scalar_t* __restrict__ A,
-                            scalar_t* __restrict__ B,
-                            int N) {
+ #define BLOCK_SIZE 512
   
-     __shared__ scalar_t partial_sum[BLOCK_SIZE];
+ template <unsigned int blockSize>
+ __device__ void warpReduce(volatile float* sdata, unsigned int tid) {
+     if (blockSize >= 64) sdata[tid] += sdata[tid + 32];
+     if (blockSize >= 32) sdata[tid] += sdata[tid + 16];
+     if (blockSize >= 16) sdata[tid] += sdata[tid + 8];
+     if (blockSize >= 8) sdata[tid] += sdata[tid + 4];
+     if (blockSize >= 4) sdata[tid] += sdata[tid + 2];
+     if (blockSize >= 2) sdata[tid] += sdata[tid + 1];
+ }
   
+  
+ template <typename scalar_t, unsigned int blockSize>
+ __global__ void sum_kernel(scalar_t* __restrict__ A, int N) {
+     __shared__ scalar_t partial_sum[2 * blockSize];
+  
      unsigned int tid = threadIdx.x;
-     unsigned int i = blockIdx.x * (BLOCK_SIZE) + tid;
+     unsigned int i = blockIdx.x * (2 * blockSize) + tid;
   
-     if (i < N) {
+     if (i + blockSize < N) {
+         partial_sum[tid] = A[i] + A[i + blockSize];
+     }
+     else if (i < N) {
          partial_sum[tid] = A[i];
      }
      else {
          partial_sum[tid] = 0;
      }
-  
-     for (unsigned int stride = BLOCK_SIZE/2; stride >= 1; stride /= 2) {
-         __syncthreads();
-         if (tid < stride) {
-             partial_sum[tid] += partial_sum[tid + stride];
-         }
-     }
      __syncthreads();
   
+     if (blockSize >= 1024) { if (tid < 512) { partial_sum[tid] += partial_sum[tid + 512]; } __syncthreads(); }
+     if (blockSize >= 512) { if (tid < 256) { partial_sum[tid] += partial_sum[tid + 256]; } __syncthreads(); }
+     if (blockSize >= 256) { if (tid < 128) { partial_sum[tid] += partial_sum[tid + 128]; } __syncthreads(); }
+     if (blockSize >= 128) { if (tid < 64) { partial_sum[tid] += partial_sum[tid + 64]; } __syncthreads(); }
+     if (tid < 32) warpReduce<blockSize>(partial_sum, tid);
+  
      if (tid == 0) {
-         atomicAdd(B, partial_sum[0]);
+         A[blockIdx.x] = partial_sum[0];
      }
  }
   
- torch::Tensor sum_cuda(torch::Tensor A, torch::Tensor B) {
+ torch::Tensor sum_cuda(torch::Tensor A) {
      int N = A.numel();
   
-     int blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
+     int blocks = (N + (BLOCK_SIZE * 2) - 1) / (BLOCK_SIZE * 2);
+     do {
+         sum_kernel<float, BLOCK_SIZE><<<blocks, BLOCK_SIZE>>>(
+             A.data_ptr<float>(),
+             N
+         );
+         N = blocks;
+         blocks = (N + (BLOCK_SIZE * 2) - 1) / (BLOCK_SIZE * 2);
+     } while (N > 1);
   
-     sum_kernel<float><<<blocks, BLOCK_SIZE>>>(
-         A.data_ptr<float>(),
-         B.data_ptr<float>(),
-         N
-     );
-  
-     return B;
+     return A;
  }
  """
   
  sum_module = load_inline(
      name='sum_cuda_ext',
-     cpp_sources="torch::Tensor sum_cuda(torch::Tensor A, torch::Tensor B);",
+     cpp_sources="torch::Tensor sum_cuda(torch::Tensor A);",
      cuda_sources=sum_cuda_source,
      functions=['sum_cuda'],
      verbose=True,
  )
   
- def sum(A, B):
-     if not A.is_cuda or not B.is_cuda:
-         raise RuntimeError("Three tensors must be on GPU")
-     return sum_module.sum_cuda(A, B)
-  
  def custom_kernel(data: input_t) -> output_t:
      """
      Custom implementation of vector addition using CUDA.
  ⋯ 2 unchanged lines
      Returns:
          Tensor containing element-wise sum.
      """
-     A, B = data
-     B.zero_()
-     assert A.is_cuda and B.is_cuda, "Input tensors must be on GPU"
-  
-     # Simply reuse the existing add function we already defined
-     # This avoids the compilation issues with the inline kernel
-     return sum(A, B)[0]
+     A, _ = data
+     return sum_module.sum_cuda(A)[0]