submission 782559

- from utils import make_match_reference, DeterministicContext
  import torch
- from task import input_t, output_t
- import sys
-  
  from torch.utils.cpp_extension import load_inline
   
- N_ELEMENTS = 16384
+ N_ELEMENTS = 268435456
   
  _CPP_SOURCE = r"""
  #include <torch/extension.h>
  ⋯ 5 unchanged lines
  }
  """
   
-  
  _CUDA_SOURCE = r"""
- #include <cuda_fp16.h>
  #include <cuda_runtime.h>
  #include <torch/extension.h>
+ #include <cub/cub.cuh>
   
-  
  constexpr size_t N_SIZE = 268435456;
   
+  
+ float* d_temp_storage = nullptr;
+ size_t temp_storage_bytes = 0;
+  
  torch::Tensor cuda_prefixsum(std::vector<torch::Tensor> data) {
+     if (data[0].numel() == N_SIZE) {
+         if (d_temp_storage == nullptr) {
+             cub::DeviceScan::InclusiveSum(
+             d_temp_storage, temp_storage_bytes,
+             data[0].data_ptr<float>(), data[1].data_ptr<float>(), N_SIZE, 0);
+             cudaMalloc(&d_temp_storage, temp_storage_bytes);
+         }
+         cub::DeviceScan::InclusiveSum(
+         d_temp_storage, temp_storage_bytes,
+         data[0].data_ptr<float>(), data[1].data_ptr<float>(), N_SIZE, 0);
   
-     if (data[0].element_size() == N_SIZE) {
-         return torch::cumsum(data[0], 0);
+         return data[1];
      } else { 
-         // 退化到 PyTorch 内置实现，保证正确性
          return torch::cumsum(data[0], 0);
      }
  }
  """
   
  _EXT = load_inline(
-         name="cuda_prefixsum_extension_001",
+         name="cuda_prefixsum_extension_0023",
          cpp_sources=[_CPP_SOURCE],
          cuda_sources=[_CUDA_SOURCE],
          functions=None,
-         extra_cflags=["-O3 -use_fast_math"],
-         extra_cuda_cflags=["-O3 -use_fast_math -Xptxas=-v -maxrregcount=32"],
+         extra_cflags=["-O3"],
+         extra_cuda_cflags=[
+             "-O3", 
+             "-use_fast_math", 
+             "-Xptxas=-v", 
+             "-Xptxas=-dlcm=cg",       # L2 Cache 全局策略
+             "-Xptxas=-warn-spills",   # 监控寄存器溢出
+             '-gencode=arch=compute_80,code=sm_80'
+         ],
          with_cuda=True,
-         verbose=False,
+         verbose=True,
      )
   
- custom_kernel = _EXT.cuda_prefixsum
-  
- def ref_kernel(data: input_t) -> output_t:
-     """
-     Reference implementation of inclusive prefix sum using PyTorch.
-     Args:
-         data: Input tensor to compute prefix sum on
-     Returns:
-         Tensor containing the inclusive prefix sum
-     """
-     with DeterministicContext():
-         data, output = data
-         output = torch.cumsum(data.to(torch.float64), dim=0).to(torch.float64)
-         return output
-  
-  
- def generate_input(size: int, seed: int) -> input_t:
-     """
-     Generates random input tensor.
-     Returns:
-         Tensor to compute prefix sum on
-     """
-     gen = torch.Generator(device="cuda")
-     gen.manual_seed(seed)
-     x = torch.randn(
-         size, device="cuda", dtype=torch.float32, generator=gen
-     ).contiguous()
-     y = torch.empty(size, device="cuda", dtype=torch.float32).contiguous()
-     return x, y
-  
-  
- # This algorithm is very sensitive to the tolerance and the error is magnified by the input size
- # The tolerance is scaled by the square root of the input size
- def check_implementation(data: input_t, output: output_t) -> str:
-     # Then get the size for scaling the tolerance
-     n = data[0].numel()
-  
-     scale_factor = n ** 0.5  # Square root of input size
-     rtol = 1e-5 * scale_factor
-     atol = 1e-5 * scale_factor
-  
-     return match_reference(data, output, reference=ref_kernel, rtol=rtol, atol=atol)
-  
-  
- 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))
+ custom_kernel = _EXT.cuda_prefixsum
   No newline at end of file