submission 779738

  from utils import make_match_reference, DeterministicContext
  import torch
  from task import input_t, output_t
+ import sys
   
- import triton
- import triton.language as tl
+ from torch.utils.cpp_extension import load_inline
   
  N_ELEMENTS = 52428800
- BLOCK_SIZE = 1024
- # NUM_WARPS = 32
- NUM_STAGES = 3
- CHUNKS = 32
   
- N1 = triton.cdiv(N_ELEMENTS, BLOCK_SIZE * CHUNKS)
- _GLOBAL_REDUCE_BUF = torch.empty(N1, device="cuda", dtype=torch.float32)
- BLOCK_SIZE2=triton.next_power_of_2(N1)
+ _CPP_SOURCE = r"""
+ #include <torch/extension.h>
   
- @triton.jit
- def reduce_kernel(
-     in_ptr,
-     out_ptr,
-     n_elements,
-     BLOCK_SIZE: tl.constexpr,
-     CHUNKS_PER_BLOCK: tl.constexpr,
- ):
-     pid = tl.program_id(axis=0)
-     block_start = pid * BLOCK_SIZE * CHUNKS_PER_BLOCK
+ torch::Tensor sum_reduce_cuda(torch::Tensor input);
   
-     acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+     m.def("sum_reduce_cuda", &sum_reduce_cuda, "Sum reduction with custom CUDA kernel");
+ }
+ """
   
-     for c in range(CHUNKS_PER_BLOCK):
-         # 计算当前 chunk 的内存偏移量
-         offsets = block_start + c * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-         
-         # 边界保护
-         mask = offsets < n_elements
   
-         # 加载当前 chunk 的数据
-         x = tl.load(in_ptr + offsets, mask=mask, other=0.0)
-         
-         # 累加到 acc 中
-         acc += x
+ _CUDA_SOURCE = r"""
+ #include <cuda_runtime.h>
+ #include <stdio.h>
   
-     partial = tl.sum(acc, axis=0)
-     tl.store(out_ptr + pid, partial)
+ // ================= 硬编码参数 (针对 52,428,800 精确计算) =================
+ constexpr int N_FLOATS          = 52428800;
+ constexpr int BLOCK_SIZE        = 256;
+ constexpr int GRID_SIZE         = 2048;
+ constexpr int FLOAT4_PER_THREAD = 25;
+ constexpr int STRIDE_FLOAT4     = BLOCK_SIZE * GRID_SIZE; // 524,288 (float4单位)
   
+ // ================= A100 极限优化内核 =================
+ __global__ void reduce_sum_50M_a100(const float* __restrict__ in, float* __restrict__ out) {
+     // 仅需 8 个 slot (每个 warp 一个 leader)，32 Bytes，零 Bank Conflict
+     __shared__ float sdata[8];
+     
+     int tid = threadIdx.x;
+     // 在 float4 空间中的全局索引 (保证连续线程加载连续 float4，完美合并)
+     int idx = blockIdx.x * BLOCK_SIZE + tid;
   
+     float sum = 0.0f;
+     const float4* in4 = reinterpret_cast<const float4*>(in);
+  
+     // 1. 向量化 + Grid-Stride + 完全展开的寄存器累加
+     #pragma unroll
+     for (int j = 0; j < FLOAT4_PER_THREAD; ++j) {
+         float4 v = __ldg(&in4[idx + j * STRIDE_FLOAT4]); // 使用只读缓存加载，减少延迟
+         sum += v.x + v.y + v.z + v.w;
+     }
+  
+     // 2. Warp 级归约 (纯寄存器 Shuffle，无共享内存交互)
+     sum += __shfl_xor_sync(0xffffffff, sum, 16);
+     sum += __shfl_xor_sync(0xffffffff, sum, 8);
+     sum += __shfl_xor_sync(0xffffffff, sum, 4);
+     sum += __shfl_xor_sync(0xffffffff, sum, 2);
+     sum += __shfl_xor_sync(0xffffffff, sum, 1);
+  
+     // 3. Warp Leader 写入共享内存
+     if (tid % 32 == 0) {
+         sdata[tid / 32] = sum;
+     }
+     __syncthreads();
+  
+     // 4. Warp 0 完成 Block 内最终归约 (8个值)
+     if (tid < 32) {
+         float wsum = (tid < 8) ? sdata[tid] : 0.0f;
+         wsum += __shfl_xor_sync(0xffffffff, wsum, 4);
+         wsum += __shfl_xor_sync(0xffffffff, wsum, 2);
+         wsum += __shfl_xor_sync(0xffffffff, wsum, 1);
+  
+         if (tid == 0) {
+             // A100 硬件加速 atomicAdd，2048 次原子操作开销 < 0.05%
+             atomicAdd(out, wsum);
+         }
+     }
+ }
+  
+  
+ torch::Tensor sum_reduce_cuda(torch::Tensor input) {
+     auto final_output = torch::zeros({}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
+     reduce_sum_50M_a100<<<GRID_SIZE, BLOCK_SIZE>>>(
+         input.data_ptr<float>(), 
+         final_output.data_ptr<float>()
+     );
+     return final_output;
+ }
+ """
+  
+  
+ _EXT = load_inline(
+         name="cuda_sum_reduce_000002_ext",
+         cpp_sources=[_CPP_SOURCE],
+         cuda_sources=[_CUDA_SOURCE],
+         functions=None,
+         extra_cflags=["-O3 -use_fast_math"],
+         extra_cuda_cflags=["-O3 -use_fast_math"],
+         with_cuda=True,
+         verbose=False,
+     )
+  
+  
  def ref_kernel(data: input_t) -> output_t:
      """
      Reference implementation of vector sum reduction using PyTorch.
  ⋯ 14 unchanged lines
      n_elements = input_tensor.numel()
      if n_elements != N_ELEMENTS:
          return input_tensor.sum()
+     return _EXT.sum_reduce_cuda(input_tensor)
   
-     reduce_kernel[(N1,)](
-         input_tensor,
-         _GLOBAL_REDUCE_BUF,
-         n_elements,
-         BLOCK_SIZE=BLOCK_SIZE,
-         CHUNKS_PER_BLOCK=CHUNKS,
-         # num_warps=NUM_WARPS,
-         num_stages=NUM_STAGES,
-     )
   
-     reduce_kernel[(1,)](
-         _GLOBAL_REDUCE_BUF,
-         input_tensor,  # 重用输入缓冲区作为中间结果存储
-         N1,
-         BLOCK_SIZE=BLOCK_SIZE2,
-         CHUNKS_PER_BLOCK=1,  # 第二轮不需要分块了
-         # num_warps=NUM_WARPS,
-         num_stages=NUM_STAGES,
-     )
-     return input_tensor[0].to(torch.float32)
-  
-  
  def generate_input(size: int, seed: int) -> input_t:
      """
      Generates random input tensor of specified shape with random offset and scale.
  ⋯ 29 unchanged lines
   
   
  check_implementation = make_match_reference(ref_kernel)
-  
  def warmup(fn, args, n_warmup=5):
      for _ in range(n_warmup):
          _ = fn(args)
          torch.cuda.synchronize()
   
- # 使用
+  
  warmup(custom_kernel, generate_input(N_ELEMENTS, 42))
   No newline at end of file