submission 670022

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

cuda_source = """
#include <cuda_fp16.h>
#include <torch/extension.h>

// Shared memory for block-level reduction
__global__ void sum_reduce_kernel(
    const float* __restrict__ x,
    float* partial_sums,
    int N
) {
    extern __shared__ float sdata[];
    
    unsigned int tid = threadIdx.x;
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Load element (with boundary check)
    sdata[tid] = (idx < N) ? x[idx] : 0.0f;
    __syncthreads();
    
    // Parallel reduction in shared memory
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // Write partial sum for this block
    if (tid == 0) {
        partial_sums[blockIdx.x] = sdata[0];
    }
}

torch::Tensor sum_reduce_cuda(torch::Tensor x) {
    TORCH_CHECK(x.device().is_cuda(), "Must be CUDA tensor");
    TORCH_CHECK(x.dtype() == torch::kFloat32, "Must be float32");
    TORCH_CHECK(x.is_contiguous(), "Must be contiguous");
    
    int N = x.numel();
    const int BLOCK_SIZE = 1024;
    const int num_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    
    // Allocate partial sums buffer
    auto partial_sums = torch::empty(num_blocks, 
                                      torch::dtype(torch::kFloat32)
                                      .device(x.device()));
    
    // Launch kernel
    sum_reduce_kernel<<<num_blocks, BLOCK_SIZE, BLOCK_SIZE * sizeof(float)>>>(
        x.data_ptr<float>(),
        partial_sums.data_ptr<float>(),
        N
    );
    
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }
    
    // Final reduction on GPU
    auto final_sum = partial_sums.sum();
    
    return final_sum;
}
"""

cpp_source = """
#include <torch/extension.h>
torch::Tensor sum_reduce_cuda(torch::Tensor x);
"""

# Load module once at module level
sum_module = load_inline(
    name='sum_reduce_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['sum_reduce_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math'],
)


def custom_kernel(data: input_t) -> output_t:
    """
    Vector sum reduction using CUDA.
    """
    # Handle input format
    if isinstance(data, (tuple, list)):
        x = data[0]
    else:
        x = data
    
    # Ensure float32, contiguous, on GPU
    x = x.contiguous().cuda().to(torch.float32)
    
    return sum_module.sum_reduce_cuda(x)