submission 511952

+ #@SUBMISSION@
   
- import triton
- import triton.language as tl
+ import os
   
+ import torch
+ from torch.utils.cpp_extension import load_inline
+  
  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")
   
- @triton.jit
- def _zero_scalar(out_ptr):
-     tl.store(out_ptr, 0.0)
+ cuda_source = r"""
+ #include <pybind11/pybind11.h>
+ #include <torch/extension.h>
   
+ template <typename T>
+ __device__ T warpReduction(T threadSum, uint mask = 0xffffffff) {
+     T warpSum {threadSum};
+     #pragma unroll
+     for (int offset = 16; offset > 0; offset >>= 1) {
+         warpSum += __shfl_down_sync(mask, warpSum, offset);
+     }
+     return warpSum;
+ }
   
- @triton.jit
- def _sum_atomic_chunked(
-     x_ptr,
-     out_ptr,
-     n_elements,
-     BLOCK: tl.constexpr,
-     ITERS: tl.constexpr,
- ):
-     pid = tl.program_id(0)
-     base = pid * BLOCK * ITERS
-     tl.multiple_of(base, 256)
-     r = tl.arange(0, BLOCK)
+ 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;
   
-     acc = tl.zeros((), dtype=tl.float32)
-     for i in tl.static_range(0, ITERS):
-         offsets = base + i * BLOCK + r
-         x = tl.load(x_ptr + offsets, mask=offsets < n_elements, other=0.0, cache_modifier=".cg")
-         acc += tl.sum(x, axis=0)
+     constexpr uint blockWarpCount = (blockSize + 31) / 32;
+     constexpr uint blockReductionThreadMask = (1 << blockWarpCount) - 1;
+     __shared__ T warpSums[blockWarpCount];
   
-     tl.atomic_add(out_ptr, acc)
+     T threadSum {};
+     while (index < size) {
+         threadSum += data[index];
+         index += stride;
+     }
   
+     T warpSum = warpReduction(threadSum);
   
- def _pick_iters(n: int) -> int:
-     if n >= 20_000_000:
-         return 16
-     if n >= 2_000_000:
-         return 8
-     return 4
+     // First thread of each warp writes the warp sum to shared memory.
+     uint warpNumber = threadIdx.x >> 5;
+     if (!(threadIdx.x & 0b11111)) {
+         warpSums[warpNumber] = warpSum;
+     }
   
+     __syncthreads();
   
- def custom_kernel(data: input_t) -> output_t:
-     x, out = data
-     n = x.numel()
+     // 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);
+         }
+     }
+ }
   
-     _zero_scalar[(1,)](out, num_warps=1, num_stages=1)
+ torch::Tensor kernelCaller(torch::Tensor x, torch::Tensor result) {
+     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>()),
+         x.numel());
+     return result;
+ }
+ """
   
-     BLOCK = 1024
-     iters = _pick_iters(n)
-     grid = (triton.cdiv(n, BLOCK * iters),)
+ cpp_source = "torch::Tensor kernelCaller(torch::Tensor x, torch::Tensor result);"
   
-     _sum_atomic_chunked[grid](
-         x,
-         out,
-         n,
-         BLOCK=BLOCK,
-         ITERS=iters,
-         num_warps=4,
-         num_stages=4,
-     )
-     return out[0]
+ module = load_inline(
+     name="vectorsum_fourcore_pmpp_v2",
+     cpp_sources=cpp_source,
+     cuda_sources=cuda_source,
+     functions=["kernelCaller"],
+     verbose=True,
+     extra_cuda_cflags=["-O2"],
+ )
   
+ result = torch.tensor(0.0, dtype=torch.float64, device="cuda:0")
+  
+  
+ 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))
+     return out_buf[0]