submission 782561

import torch
from torch.utils.cpp_extension import load_inline

N_ELEMENTS = 268435456

_CPP_SOURCE = r"""
#include <torch/extension.h>

torch::Tensor cuda_prefixsum(std::vector<torch::Tensor> data);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("cuda_prefixsum", &cuda_prefixsum, "prefixsum with custom CUDA kernel");
}
"""

_CUDA_SOURCE = r"""
#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);

        return data[1];
    } else { 
        return torch::cumsum(data[0], 0);
    }
}
"""

_EXT = load_inline(
        name="cuda_prefixsum_extension_0023",
        cpp_sources=[_CPP_SOURCE],
        cuda_sources=[_CUDA_SOURCE],
        functions=None,
        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=True,
    )

custom_kernel = _EXT.cuda_prefixsum