submission 511988

  ⋯ 6 unchanged lines
   
  from task import input_t, output_t
   
- # Keep default arch list broad enough for environments that may not support sm_89 directly.
  os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "8.0+PTX")
   
  cuda_source = r"""
  ⋯ 1 unchanged lines
  #include <torch/extension.h>
   
  template <typename T>
- __device__ T warpReduction(T threadSum, uint mask = 0xffffffff) {
-     T warpSum {threadSum};
+ __device__ __forceinline__ T warpReduction(T v, unsigned mask = 0xffffffffu) {
      #pragma unroll
      for (int offset = 16; offset > 0; offset >>= 1) {
-         warpSum += __shfl_down_sync(mask, warpSum, offset);
+         v += __shfl_down_sync(mask, v, offset);
      }
-     return warpSum;
+     return v;
  }
   
- template <typename T, typename T2, uint blockSize>
- __global__ void reduction_basic(T* data, T2* out, size_t size) {
-     uint index = blockIdx.x * blockDim.x + threadIdx.x;
-     const uint stride = gridDim.x * blockDim.x;
+ 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;
+     const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
   
-     constexpr uint blockWarpCount = (blockSize + 31) / 32;
-     constexpr uint blockReductionThreadMask = (1 << blockWarpCount) - 1;
-     __shared__ T warpSums[blockWarpCount];
+     constexpr unsigned blockWarpCount = (blockSize + 31) / 32;
+     constexpr unsigned blockReductionThreadMask = (1u << blockWarpCount) - 1u;
+     __shared__ float warpSums[blockWarpCount];
   
-     T threadSum {};
+     float threadSum = 0.0f;
      while (index < size) {
          threadSum += data[index];
          index += stride;
      }
   
-     T warpSum = warpReduction(threadSum);
-  
-     // First thread of each warp writes the warp sum to shared memory.
-     uint warpNumber = threadIdx.x >> 5;
-     if (!(threadIdx.x & 0b11111)) {
-         warpSums[warpNumber] = warpSum;
-     }
-  
+     const float warpSum = warpReduction(threadSum);
+     const unsigned warpNumber = threadIdx.x >> 5;
+     if (!(threadIdx.x & 31)) warpSums[warpNumber] = warpSum;
      __syncthreads();
   
-     // First warp reduces the block warp sums.
      if (threadIdx.x < blockWarpCount) {
-         T blockSum = warpReduction(warpSums[threadIdx.x], blockReductionThreadMask);
-         if (!threadIdx.x) {
-             atomicAdd(out, (T2)blockSum);
-         }
+         const float blockSum = warpReduction(warpSums[threadIdx.x], blockReductionThreadMask);
+         if (!threadIdx.x) atomicAdd(out, static_cast<double>(blockSum));
      }
  }
   
- torch::Tensor kernelCaller(torch::Tensor x, torch::Tensor result) {
+ __global__ void cast_one(const double* __restrict__ in, float* out) {
+     if (!threadIdx.x && !blockIdx.x) out[0] = static_cast<float>(in[0]);
+ }
+  
+ torch::Tensor kernelReduce(torch::Tensor x, torch::Tensor out) {
      constexpr int threads = 256;
      const int blocks = (x.numel() + threads - 1) / threads;
-     reduction_basic<float, double, 256><<<blocks, threads>>>(
-         reinterpret_cast<float*>(x.data_ptr<float>()),
-         reinterpret_cast<double*>(result.data_ptr<double>()),
+     reduction_doubleout<256><<<blocks, threads>>>(
+         reinterpret_cast<const float*>(x.data_ptr<float>()),
+         reinterpret_cast<double*>(out.data_ptr<double>()),
          x.numel());
-     return result;
+     return out;
  }
+  
+ 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>()));
+     return out;
+ }
  """
   
- cpp_source = "torch::Tensor kernelCaller(torch::Tensor x, torch::Tensor result);"
+ cpp_source = """
+ torch::Tensor kernelReduce(torch::Tensor x, torch::Tensor out);
+ torch::Tensor kernelCast(torch::Tensor in, torch::Tensor out);
+ """
   
  module = load_inline(
-     name="vectorsum_fourcore_pmpp_v2",
+     name="vectorsum_fourcore_0227_top1_noalloc_cast",
      cpp_sources=cpp_source,
      cuda_sources=cuda_source,
-     functions=["kernelCaller"],
-     verbose=True,
-     extra_cuda_cflags=["-O2"],
+     functions=["kernelReduce", "kernelCast"],
+     verbose=False,
+     extra_cuda_cflags=["-O3", "--use_fast_math"],
  )
   
  result = torch.tensor(0.0, dtype=torch.float64, device="cuda:0")
  ⋯ 2 unchanged lines
  def custom_kernel(data: input_t) -> output_t:
      x, out_buf = data
      result.fill_(0.0)
-     module.kernelCaller(x, result)
-     out_buf.copy_(result.to(torch.float32))
+     module.kernelReduce(x, result)
+     module.kernelCast(result, out_buf)
      return out_buf[0]