submission 66598

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

# Optimize for L4 GPU architecture
os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'

def custom_kernel(data: input_t) -> output_t:
    """
    CUDA implementation of vectorSum with Kahan compensated summation.
    
    Args:
        data: Tuple of (input_tensor, output_tensor)
              - input_tensor: 1D tensor of shape (N,) with float32 values
              - output_tensor: Pre-allocated output tensor of shape (1,)
    
    Returns:
        Scalar tensor containing the sum of all input elements
    """
    # Unpack the input tuple
    input_tensor, output_tensor = data
    
    # Ensure contiguous memory layout for optimal access
    A = input_tensor.contiguous()
    
    # Validation checks
    assert A.is_cuda, "Input tensor must be on GPU"
    assert A.dim() == 1, "Input tensor must be 1-dimensional"
    assert A.dtype == torch.float32, "Input tensor must be float32"
    
    N = A.numel()
    
    # Handle edge case of empty input
    if N == 0:
        output_tensor.fill_(0.0)
        return output_tensor.squeeze()
    
    # CUDA kernel source code with Kahan summation for better precision
    cuda_source = """
    #include <torch/extension.h>
    #include <cuda_runtime.h>

    // Kahan compensated summation for better numerical accuracy
    template <typename scalar_t>
    __inline__ __device__ void kahanSum(scalar_t& sum, scalar_t& c, scalar_t value) {
        scalar_t y = value - c;
        scalar_t t = sum + y;
        c = (t - sum) - y;
        sum = t;
    }

    // Warp-level reduction using shuffle instructions
    template <typename scalar_t>
    __inline__ __device__ scalar_t warpReduceSum(scalar_t val) {
        for (int offset = 16; offset > 0; offset /= 2) {
            val += __shfl_down_sync(0xffffffff, val, offset);
        }
        return val;
    }

    // Two-stage reduction: blocks -> final sum
    template <typename scalar_t>
    __global__ void vector_sum_kernel_stage1(
        const scalar_t* __restrict__ input,
        scalar_t* __restrict__ block_sums,
        const int N
    ) {
        extern __shared__ __align__(sizeof(scalar_t)) char shared_mem_char[];
        scalar_t* sharedMem = reinterpret_cast<scalar_t*>(shared_mem_char);

        const int tid = threadIdx.x;
        const int bid = blockIdx.x;
        const int globalIdx = bid * blockDim.x + tid;

        // Kahan summation for each thread
        scalar_t threadSum = 0;
        scalar_t compensation = 0;

        for (int i = globalIdx; i < N; i += gridDim.x * blockDim.x) {
            kahanSum(threadSum, compensation, input[i]);
        }

        // Warp-level reduction
        threadSum = warpReduceSum(threadSum);

        // Store warp results to shared memory
        if (tid % 32 == 0) {
            sharedMem[tid / 32] = threadSum;
        }
        __syncthreads();

        // First warp reduces all warp sums
        if (tid < 32) {
            scalar_t val = (tid < (blockDim.x + 31) / 32) ? sharedMem[tid] : 0;
            scalar_t blockSum = warpReduceSum(val);

            if (tid == 0) {
                block_sums[bid] = blockSum;
            }
        }
    }

    template <typename scalar_t>
    __global__ void vector_sum_kernel_stage2(
        const scalar_t* __restrict__ block_sums,
        scalar_t* __restrict__ output,
        const int num_blocks
    ) {
        extern __shared__ __align__(sizeof(scalar_t)) char shared_mem_char[];
        scalar_t* sharedMem = reinterpret_cast<scalar_t*>(shared_mem_char);

        const int tid = threadIdx.x;

        // Kahan summation for final reduction
        scalar_t threadSum = 0;
        scalar_t compensation = 0;

        for (int i = tid; i < num_blocks; i += blockDim.x) {
            kahanSum(threadSum, compensation, block_sums[i]);
        }

        // Warp-level reduction
        threadSum = warpReduceSum(threadSum);

        if (tid % 32 == 0) {
            sharedMem[tid / 32] = threadSum;
        }
        __syncthreads();

        if (tid < 32) {
            scalar_t val = (tid < (blockDim.x + 31) / 32) ? sharedMem[tid] : 0;
            scalar_t finalSum = warpReduceSum(val);

            if (tid == 0) {
                output[0] = finalSum;
            }
        }
    }

    torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output) {
        TORCH_CHECK(input.device().is_cuda(), "Input must be a CUDA tensor");
        TORCH_CHECK(input.dim() == 1, "Input must be 1-dimensional");
        
        const int N = input.numel();
        
        if (N == 0) {
            output.fill_(0.0);
            return output;
        }

        const int blockSize = 256;
        const int numBlocks = min((N + blockSize - 1) / blockSize, 4096);
        
        // Allocate intermediate storage for block sums
        auto block_sums = torch::empty({numBlocks}, input.options());
        
        // Stage 1: Reduce input to per-block sums
        dim3 grid1(numBlocks);
        dim3 block1(blockSize);
        size_t shared_mem_size1 = (blockSize / 32) * sizeof(float);
        
        AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "vector_sum_kernel_stage1", ([&] {
            vector_sum_kernel_stage1<scalar_t><<<grid1, block1, shared_mem_size1>>>(
                input.data_ptr<scalar_t>(),
                block_sums.data_ptr<scalar_t>(),
                N
            );
        }));
        
        // Stage 2: Reduce block sums to final result (deterministic, no atomics)
        dim3 grid2(1);
        dim3 block2(blockSize);
        size_t shared_mem_size2 = (blockSize / 32) * sizeof(float);
        
        AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "vector_sum_kernel_stage2", ([&] {
            vector_sum_kernel_stage2<scalar_t><<<grid2, block2, shared_mem_size2>>>(
                block_sums.data_ptr<scalar_t>(),
                output.data_ptr<scalar_t>(),
                numBlocks
            );
        }));

        cudaError_t err = cudaGetLastError();
        if (err != cudaSuccess) {
            throw std::runtime_error(cudaGetErrorString(err));
        }

        return output;
    }
    """

    # C++ binding code
    cpp_source = """
    #include <torch/extension.h>
    torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output);
    
    PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
        m.def("vector_sum_cuda", &vector_sum_cuda, "Vector sum reduction (CUDA)");
    }
    """

    # Compile and load the CUDA extension
    module = load_inline(
        name='vector_sum_optimized',
        cpp_sources=cpp_source,
        cuda_sources=cuda_source,
        extra_cuda_cflags=['-arch=sm_89', '-O3', '--use_fast_math'],
        verbose=False
    )

    # Execute the kernel and return the result
    result = module.vector_sum_cuda(A, output_tensor)
    return result.squeeze()