submission 66907

+ 
  import os
  import torch
  from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
  
- # Optimize for L4 GPU architecture
- os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
+ # Multi-architecture support
+ os.environ['TORCH_CUDA_ARCH_LIST'] = '8.0;8.9;9.0'  # A100, L4, B100
  
  def custom_kernel(data: input_t) -> output_t:
      """
-     CUDA implementation of vectorSum with Kahan compensated summation.
-     
+     Optimized CUDA implementation of vectorSum
      Args:
          data: Tuple of (input_tensor, output_tensor)
-               - input_tensor: 1D tensor of shape (N,) with float32 values
-               - output_tensor: Pre-allocated output tensor of shape (1,)
-     
      Returns:
-         Scalar tensor containing the sum of all input elements
+         Scalar value equal to the sum of all elements.
      """
-     # Unpack the input tuple
-     input_tensor, output_tensor = data
+     A, output_buffer = data  # Unpack the tuple
+     A = A.contiguous()
      
-     # Ensure contiguous memory layout for optimal access
-     A = input_tensor.contiguous()
-     
-     # Validation checks
      assert A.is_cuda, "Input tensor must be on GPU"
      assert A.dim() == 1, "Input tensor must be 1-dimensional"
      assert A.dtype == torch.float32, "Input tensor must be float32"
      
      N = A.numel()
-     
-     # Handle edge case of empty input
      if N == 0:
-         output_tensor.fill_(0.0)
-         return output_tensor.squeeze()
+         return torch.zeros((), dtype=A.dtype, device=A.device)
      
-     # CUDA kernel source code with Kahan summation for better precision
      cuda_source = """
      #include <torch/extension.h>
      #include <cuda_runtime.h>
- 
-     // Kahan compensated summation for better numerical accuracy
-     template <typename scalar_t>
-     __inline__ __device__ void kahanSum(scalar_t& sum, scalar_t& c, scalar_t value) {
-         scalar_t y = value - c;
-         scalar_t t = sum + y;
-         c = (t - sum) - y;
-         sum = t;
-     }
- 
-     // Warp-level reduction using shuffle instructions
-     template <typename scalar_t>
-     __inline__ __device__ scalar_t warpReduceSum(scalar_t val) {
-         for (int offset = 16; offset > 0; offset /= 2) {
-             val += __shfl_down_sync(0xffffffff, val, offset);
+     #include <cub/cub.cuh>
+     #include <cuda/atomic>
+     
+     // Optimized warp reduction
+     __inline__ __device__ float warpReduceSum(float val) {
+         #pragma unroll
+         for (int mask = 16; mask > 0; mask >>= 1) {
+             val += __shfl_xor_sync(0xffffffff, val, mask);
          }
          return val;
      }
- 
-     // Two-stage reduction: blocks -> final sum
-     template <typename scalar_t>
-     __global__ void vector_sum_kernel_stage1(
-         const scalar_t* __restrict__ input,
-         scalar_t* __restrict__ block_sums,
+     
+     // Single-pass reduction kernel with aggressive optimization
+     template<int BLOCK_SIZE, int ITEMS_PER_THREAD>
+     __global__ void __launch_bounds__(BLOCK_SIZE)
+     vector_sum_kernel_fast(
+         const float* __restrict__ input,
+         float* __restrict__ output,
          const int N
      ) {
-         extern __shared__ __align__(sizeof(scalar_t)) char shared_mem_char[];
-         scalar_t* sharedMem = reinterpret_cast<scalar_t*>(shared_mem_char);
- 
+         // Shared memory for warp-level reductions
+         __shared__ float shared_data[BLOCK_SIZE / 32];
+         
          const int tid = threadIdx.x;
-         const int bid = blockIdx.x;
-         const int globalIdx = bid * blockDim.x + tid;
- 
-         // Kahan summation for each thread
-         scalar_t threadSum = 0;
-         scalar_t compensation = 0;
- 
-         for (int i = globalIdx; i < N; i += gridDim.x * blockDim.x) {
-             kahanSum(threadSum, compensation, input[i]);
-         }
- 
-         // Warp-level reduction
-         threadSum = warpReduceSum(threadSum);
- 
-         // Store warp results to shared memory
-         if (tid % 32 == 0) {
-             sharedMem[tid / 32] = threadSum;
-         }
-         __syncthreads();
- 
-         // First warp reduces all warp sums
-         if (tid < 32) {
-             scalar_t val = (tid < (blockDim.x + 31) / 32) ? sharedMem[tid] : 0;
-             scalar_t blockSum = warpReduceSum(val);
- 
-             if (tid == 0) {
-                 block_sums[bid] = blockSum;
+         const int gid = blockIdx.x * BLOCK_SIZE + tid;
+         const int grid_size = gridDim.x * BLOCK_SIZE;
+         const int warp_id = tid >> 5;
+         const int lane_id = tid & 31;
+         
+         float sum = 0.0f;
+         
+         // Aggressive unrolling - each thread processes ITEMS_PER_THREAD elements
+         int idx = gid;
+         
+         // Main loop - vectorized loads when possible
+         if (N >= grid_size * 4) {
+             // Fast path for large arrays - no boundary checks in main loop
+             const int limit = N - (N % (grid_size * 4));
+             
+             while (idx < limit) {
+                 float4 v1 = *reinterpret_cast<const float4*>(&input[idx]);
+                 idx += grid_size;
+                 float4 v2 = *reinterpret_cast<const float4*>(&input[idx]);
+                 idx += grid_size;
+                 float4 v3 = *reinterpret_cast<const float4*>(&input[idx]);
+                 idx += grid_size;
+                 float4 v4 = *reinterpret_cast<const float4*>(&input[idx]);
+                 idx += grid_size;
+                 
+                 sum += (v1.x + v1.y + v1.z + v1.w) + 
+                        (v2.x + v2.y + v2.z + v2.w) + 
+                        (v3.x + v3.y + v3.z + v3.w) + 
+                        (v4.x + v4.y + v4.z + v4.w);
              }
+             
+             // Handle remainder
+             while (idx < N) {
+                 sum += input[idx];
+                 idx += grid_size;
+             }
+         } else {
+             // Fallback for smaller arrays
+             while (idx < N) {
+                 sum += input[idx];
+                 idx += grid_size;
+             }
          }
-     }
- 
-     template <typename scalar_t>
-     __global__ void vector_sum_kernel_stage2(
-         const scalar_t* __restrict__ block_sums,
-         scalar_t* __restrict__ output,
-         const int num_blocks
-     ) {
-         extern __shared__ __align__(sizeof(scalar_t)) char shared_mem_char[];
-         scalar_t* sharedMem = reinterpret_cast<scalar_t*>(shared_mem_char);
- 
-         const int tid = threadIdx.x;
- 
-         // Kahan summation for final reduction
-         scalar_t threadSum = 0;
-         scalar_t compensation = 0;
- 
-         for (int i = tid; i < num_blocks; i += blockDim.x) {
-             kahanSum(threadSum, compensation, block_sums[i]);
-         }
- 
+         
          // Warp-level reduction
-         threadSum = warpReduceSum(threadSum);
- 
-         if (tid % 32 == 0) {
-             sharedMem[tid / 32] = threadSum;
+         sum = warpReduceSum(sum);
+         
+         // Write warp sums to shared memory
+         if (lane_id == 0) {
+             shared_data[warp_id] = sum;
          }
          __syncthreads();
- 
+         
+         // Final reduction by first warp
          if (tid < 32) {
-             scalar_t val = (tid < (blockDim.x + 31) / 32) ? sharedMem[tid] : 0;
-             scalar_t finalSum = warpReduceSum(val);
- 
+             sum = (tid < (BLOCK_SIZE / 32)) ? shared_data[tid] : 0.0f;
+             sum = warpReduceSum(sum);
+             
+             // Use warp-aggregated atomic to reduce contention
              if (tid == 0) {
-                 output[0] = finalSum;
+                 atomicAdd(output, sum);
              }
          }
      }
- 
-     torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output) {
-         TORCH_CHECK(input.device().is_cuda(), "Input must be a CUDA tensor");
-         TORCH_CHECK(input.dim() == 1, "Input must be 1-dimensional");
+     
+     // CUB-based implementation for maximum performance
+     torch::Tensor vector_sum_cuda_cub(torch::Tensor input) {
+         const int N = input.numel();
+         auto output = torch::zeros({}, input.options());
          
+         if (N == 0) return output;
+         
+         // Use CUB for optimal performance
+         size_t temp_storage_bytes = 0;
+         cub::DeviceReduce::Sum(
+             nullptr, temp_storage_bytes,
+             input.data_ptr<float>(), output.data_ptr<float>(), N
+         );
+         
+         auto temp_storage = torch::empty({(long)temp_storage_bytes}, 
+                                         torch::TensorOptions()
+                                         .dtype(torch::kUInt8)
+                                         .device(input.device()));
+         
+         cub::DeviceReduce::Sum(
+             temp_storage.data_ptr(), temp_storage_bytes,
+             input.data_ptr<float>(), output.data_ptr<float>(), N
+         );
+         
+         return output;
+     }
+     
+     // Custom implementation for cases where CUB might not be optimal
+     torch::Tensor vector_sum_cuda_custom(torch::Tensor input) {
          const int N = input.numel();
+         auto output = torch::zeros({}, input.options());
          
-         if (N == 0) {
-             output.fill_(0.0);
-             return output;
-         }
- 
-         const int blockSize = 256;
-         const int numBlocks = min((N + blockSize - 1) / blockSize, 4096);
+         if (N == 0) return output;
          
-         // Allocate intermediate storage for block sums
-         auto block_sums = torch::empty({numBlocks}, input.options());
+         // Aggressive configuration for maximum occupancy
+         const int BLOCK_SIZE = 256;
+         const int ITEMS_PER_THREAD = 4;
          
-         // Stage 1: Reduce input to per-block sums
-         dim3 grid1(numBlocks);
-         dim3 block1(blockSize);
-         size_t shared_mem_size1 = (blockSize / 32) * sizeof(float);
+         // Launch MANY blocks to saturate memory bandwidth
+         int num_blocks;
+         if (N < 1024) {
+             num_blocks = 1;
+         } else if (N < 1024 * 1024) {
+             num_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
+         } else {
+             // For large arrays, use enough blocks to saturate bandwidth
+             // Target: each block processes ~2K-4K elements
+             num_blocks = std::min(65536, (N + 2048 - 1) / 2048);
+             
+             // Ensure we have enough blocks to saturate the GPU
+             cudaDeviceProp prop;
+             cudaGetDeviceProperties(&prop, 0);
+             int min_blocks = prop.multiProcessorCount * 16;  // 16 blocks per SM
+             num_blocks = std::max(num_blocks, min_blocks);
+         }
          
-         AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "vector_sum_kernel_stage1", ([&] {
-             vector_sum_kernel_stage1<scalar_t><<<grid1, block1, shared_mem_size1>>>(
-                 input.data_ptr<scalar_t>(),
-                 block_sums.data_ptr<scalar_t>(),
+         vector_sum_kernel_fast<BLOCK_SIZE, ITEMS_PER_THREAD>
+             <<<num_blocks, BLOCK_SIZE>>>(
+                 input.data_ptr<float>(),
+                 output.data_ptr<float>(),
                  N
              );
-         }));
          
-         // Stage 2: Reduce block sums to final result (deterministic, no atomics)
-         dim3 grid2(1);
-         dim3 block2(blockSize);
-         size_t shared_mem_size2 = (blockSize / 32) * sizeof(float);
-         
-         AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "vector_sum_kernel_stage2", ([&] {
-             vector_sum_kernel_stage2<scalar_t><<<grid2, block2, shared_mem_size2>>>(
-                 block_sums.data_ptr<scalar_t>(),
-                 output.data_ptr<scalar_t>(),
-                 numBlocks
-             );
-         }));
- 
          cudaError_t err = cudaGetLastError();
          if (err != cudaSuccess) {
              throw std::runtime_error(cudaGetErrorString(err));
          }
- 
+         
          return output;
      }
+     
+     torch::Tensor vector_sum_cuda(torch::Tensor input) {
+         // For best performance, use CUB for large arrays
+         const int N = input.numel();
+         
+         // Use CUB for large arrays (it's highly optimized by NVIDIA)
+         if (N > 1024 * 1024) {
+             return vector_sum_cuda_cub(input);
+         } else {
+             // Use custom kernel for smaller arrays (less overhead)
+             return vector_sum_cuda_custom(input);
+         }
+     }
      """
- 
-     # C++ binding code
+     
      cpp_source = """
      #include <torch/extension.h>
-     torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output);
-     
+     torch::Tensor vector_sum_cuda(torch::Tensor input);
      PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-         m.def("vector_sum_cuda", &vector_sum_cuda, "Vector sum reduction (CUDA)");
+         m.def("vector_sum_cuda", &vector_sum_cuda, "Optimized vector sum");
      }
      """
- 
-     # Compile and load the CUDA extension
+     
      module = load_inline(
-         name='vector_sum_optimized',
+         name='vector_sum_ultra_optimized',
          cpp_sources=cpp_source,
          cuda_sources=cuda_source,
-         extra_cuda_cflags=['-arch=sm_89', '-O3', '--use_fast_math'],
+         extra_cuda_cflags=[
+             '-O3',
+             '--use_fast_math',
+             '--expt-relaxed-constexpr',
+             '-Xptxas="-v"',
+             '-lineinfo'
+         ],
+         extra_include_paths=['/usr/local/cuda/include'],  # For CUB
          verbose=False
      )
- 
-     # Execute the kernel and return the result
-     result = module.vector_sum_cuda(A, output_tensor)
-     return result.squeeze()
   No newline at end of file
+     
+     result = module.vector_sum_cuda(A)
+     return result
   No newline at end of file