submission 612258

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

cuda_source = r"""
#include <torch/extension.h>
#include <cub/cub.cuh>
#include <cuda_runtime.h>

static void* d_temp_storage = nullptr;
static size_t temp_storage_bytes = 0;
static int last_n = 0;

void inclusive_scan(torch::Tensor input, torch::Tensor output) {
    int n = input.numel();
    float* d_in = input.data_ptr<float>();
    float* d_out = output.data_ptr<float>();

    if (n != last_n) {
        if (d_temp_storage) cudaFree(d_temp_storage);
        temp_storage_bytes = 0;
        cub::DeviceScan::InclusiveSum(nullptr, temp_storage_bytes, d_in, d_out, n);
        cudaMalloc(&d_temp_storage, temp_storage_bytes);
        last_n = n;
    }

    cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, d_in, d_out, n);
}
"""

cpp_source = r"""
void inclusive_scan(torch::Tensor input, torch::Tensor output);
"""

module = load_inline(
    name="cub_scan",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["inclusive_scan"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    data, output = data
    module.inclusive_scan(data, output)
    return output