submission 511990

  #@SUBMISSION@
   
  import os
-  
  import torch
  from torch.utils.cpp_extension import load_inline
-  
  from task import input_t, output_t
   
- os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "8.0+PTX")
+ os.environ["TORCH_CUDA_ARCH_LIST"] = "8.9"
   
  cuda_source = r"""
  #include <pybind11/pybind11.h>
  ⋯ 2 unchanged lines
  template <typename T>
  __device__ __forceinline__ T warpReduction(T v, unsigned mask = 0xffffffffu) {
      #pragma unroll
-     for (int offset = 16; offset > 0; offset >>= 1) {
-         v += __shfl_down_sync(mask, v, offset);
-     }
+     for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(mask, v, offset);
      return v;
  }
   
  template <unsigned blockSize>
- __global__ void reduction_doubleout(const float* __restrict__ data, double* out, size_t size) {
-     size_t index = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+ __global__ __launch_bounds__(blockSize, 2)
+ void reduction_doubleout(const float* __restrict__ data, double* out, size_t size) {
+     size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
      const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+     constexpr unsigned warps = (blockSize + 31) / 32;
+     constexpr unsigned mask = (1u << warps) - 1u;
+     __shared__ float warpSums[warps];
   
-     constexpr unsigned blockWarpCount = (blockSize + 31) / 32;
-     constexpr unsigned blockReductionThreadMask = (1u << blockWarpCount) - 1u;
-     __shared__ float warpSums[blockWarpCount];
-  
-     float threadSum = 0.0f;
-     while (index < size) {
-         threadSum += data[index];
-         index += stride;
+     float sum = 0.0f;
+     while (idx < size) {
+         sum += data[idx];
+         idx += stride;
      }
   
-     const float warpSum = warpReduction(threadSum);
-     const unsigned warpNumber = threadIdx.x >> 5;
-     if (!(threadIdx.x & 31)) warpSums[warpNumber] = warpSum;
+     const float w = warpReduction(sum);
+     if (!(threadIdx.x & 31)) warpSums[threadIdx.x >> 5] = w;
      __syncthreads();
   
-     if (threadIdx.x < blockWarpCount) {
-         const float blockSum = warpReduction(warpSums[threadIdx.x], blockReductionThreadMask);
-         if (!threadIdx.x) atomicAdd(out, static_cast<double>(blockSum));
+     if (threadIdx.x < warps) {
+         const float b = warpReduction(warpSums[threadIdx.x], mask);
+         if (!threadIdx.x) atomicAdd(out, static_cast<double>(b));
      }
  }
   
  ⋯ 2 unchanged lines
  }
   
  torch::Tensor kernelReduce(torch::Tensor x, torch::Tensor out) {
-     constexpr int threads = 256;
+     constexpr int threads = 448;
      const int blocks = (x.numel() + threads - 1) / threads;
-     reduction_doubleout<256><<<blocks, threads>>>(
+     reduction_doubleout<448><<<blocks, threads>>>(
          reinterpret_cast<const float*>(x.data_ptr<float>()),
          reinterpret_cast<double*>(out.data_ptr<double>()),
          x.numel());
  ⋯ 1 unchanged lines
  }
   
  torch::Tensor kernelCast(torch::Tensor in, torch::Tensor out) {
-     cast_one<<<1, 1>>>(
-         reinterpret_cast<const double*>(in.data_ptr<double>()),
-         reinterpret_cast<float*>(out.data_ptr<float>()));
+     cast_one<<<1, 1>>>(reinterpret_cast<const double*>(in.data_ptr<double>()), reinterpret_cast<float*>(out.data_ptr<float>()));
      return out;
  }
  """
  ⋯ 4 unchanged lines
  """
   
  module = load_inline(
-     name="vectorsum_fourcore_0227_top1_noalloc_cast",
+     name="vectorsum_fourcore_0227_combo_c_448_noalloc_sm89",
      cpp_sources=cpp_source,
      cuda_sources=cuda_source,
      functions=["kernelReduce", "kernelCast"],
      verbose=False,
-     extra_cuda_cflags=["-O3", "--use_fast_math"],
+     extra_cuda_cflags=["-O3", "--use_fast_math", "-maxrregcount=80"],
  )
   
  result = torch.tensor(0.0, dtype=torch.float64, device="cuda:0")
   
-  
  def custom_kernel(data: input_t) -> output_t:
-     x, out_buf = data
+     x, out = data
      result.fill_(0.0)
      module.kernelReduce(x, result)
-     module.kernelCast(result, out_buf)
-     return out_buf[0]
+     module.kernelCast(result, out)
+     return out[0]