submission 523190

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 <cuda.h>
#include <cuda_runtime.h>
#include <cub/cub.cuh>
#include <cub/agent/single_pass_scan_operators.cuh>

// Custom inclusive scan using CUB building blocks with optimized tile size
// and CUB's warp-cooperative decoupled lookback

constexpr int BT = 128;
constexpr int IPT = 36;
constexpr int TILE = BT * IPT;  // 4608

// cub::Sum was removed in CUDA 13 — define our own
struct SumOp {
    __host__ __device__ __forceinline__
    float operator()(const float& a, const float& b) const { return a + b; }
};

typedef cub::ScanTileState<float> TileStateT;
typedef cub::BlockLoad<float, BT, IPT, cub::BLOCK_LOAD_WARP_TRANSPOSE> BlockLoadT;
typedef cub::BlockStore<float, BT, IPT, cub::BLOCK_STORE_WARP_TRANSPOSE> BlockStoreT;
typedef cub::BlockScan<float, BT, cub::BLOCK_SCAN_WARP_SCANS> BlockScanT;
typedef cub::TilePrefixCallbackOp<float, SumOp, TileStateT> PrefixOpT;

// cub::DeviceScanInitKernel is an internal API removed in CUDA 13.
// Replicate it: call InitializeStatus on each tile state entry.
__global__ void init_tile_state_kernel(TileStateT tile_state, int num_tiles) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < num_tiles)
        tile_state.InitializeStatus(num_tiles);
}

__global__ __launch_bounds__(BT)
void scan_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    TileStateT tile_state,
    const int n
) {
    const int tid = threadIdx.x;
    const int tile_id = blockIdx.x;
    const int tile_offset = tile_id * TILE;
    const int valid = min(TILE, n - tile_offset);

    __shared__ union {
        typename BlockLoadT::TempStorage load;
        typename BlockStoreT::TempStorage store;
        struct {
            typename BlockScanT::TempStorage scan;
            typename PrefixOpT::TempStorage prefix;
        };
    } temp;

    // Load tile data
    float items[IPT];
    if (valid == TILE)
        BlockLoadT(temp.load).Load(input + tile_offset, items);
    else
        BlockLoadT(temp.load).Load(input + tile_offset, items, valid, 0.0f);
    __syncthreads();

    // Block-level inclusive scan + inter-block lookback
    if (tile_id == 0) {
        float block_agg;
        BlockScanT(temp.scan).InclusiveSum(items, items, block_agg);
        if (tid == 0)
            tile_state.SetInclusive(0, block_agg);
    } else {
        PrefixOpT prefix_op(tile_state, temp.prefix, SumOp(), tile_id);
        BlockScanT(temp.scan).InclusiveSum(items, items, prefix_op);
    }
    __syncthreads();

    // Store results
    if (valid == TILE)
        BlockStoreT(temp.store).Store(output + tile_offset, items);
    else
        BlockStoreT(temp.store).Store(output + tile_offset, items, valid);
}

// Pre-allocated tile state storage
static void* d_tile_state_storage = nullptr;
static size_t d_tile_state_bytes = 0;
static int alloc_max_tiles = 0;

torch::Tensor inclusive_scan(torch::Tensor input, torch::Tensor output) {
    const int n = input.size(0);
    const int num_tiles = (n + TILE - 1) / TILE;

    // Grow tile state storage if needed
    if (num_tiles > alloc_max_tiles) {
        if (d_tile_state_storage) cudaFree(d_tile_state_storage);
        size_t needed;
        TileStateT::AllocationSize(num_tiles, needed);
        d_tile_state_bytes = needed * 2;
        cudaMalloc(&d_tile_state_storage, d_tile_state_bytes);
        alloc_max_tiles = num_tiles;
    }

    // Initialize tile state
    TileStateT tile_state;
    tile_state.Init(num_tiles, d_tile_state_storage, d_tile_state_bytes);
    init_tile_state_kernel<<<(num_tiles + 255) / 256, 256>>>(tile_state, num_tiles);

    // Launch scan kernel
    scan_kernel<<<num_tiles, BT>>>(
        input.data_ptr<float>(), output.data_ptr<float>(),
        tile_state, n);

    return output;
}
"""

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

print("Compiling submission kernel...")
_module = load_inline(
    name='prefix_sum_submission_v2',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['inclusive_scan'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math', '--threads', '4']
)
print("Done.")


def custom_kernel(data: input_t) -> output_t:
    input_tensor, output_tensor = data
    _module.inclusive_scan(input_tensor, output_tensor)
    return output_tensor