submission 780560

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

sum_cuda_source = """
__global__ void first_kernel(const float * input, int N, int warpsPerBlock, float * resultsBlocks) {
    int bid = blockIdx.x;
    int tid = threadIdx.x;
    int warpId = tid / 32;
    int warpLane = tid % 32;
    int gridStride = gridDim.x * blockDim.x;
    
    extern __shared__ float shmem[];

    float localValue = 0.0f;

    int globalId = (bid * blockDim.x + tid);

    for (int i = globalId * 4; i<N-3; i+= gridStride * 4) {
        float4 vals = reinterpret_cast<const float4 *>(&input[i])[0];
        localValue += vals.x + vals.y + vals.z + vals.w;
    }

    int nVec = (N / 4) * 4;
    for (int i = nVec + globalId; i < N; i += gridStride) {
        localValue += input[i];
    }

    // warp lvl reduction
    for (int offset = 16; offset > 0; offset /= 2) {
        localValue += __shfl_down_sync(0xFFFFFFFF, localValue, offset);
    }

    if (warpLane == 0) {
        shmem[warpId] = localValue;
    }
    __syncthreads();

    
    // block lvl reduction, made by a single warp to avoid synchronizations
    if (warpId == 0) {
        float warpVal = 0.0f;
        if (warpLane < warpsPerBlock) {
            warpVal = shmem[warpLane];
        }
        for (int offset = 16; offset > 0; offset /= 2) {
            warpVal += __shfl_down_sync(0xFFFFFFFF, warpVal, offset);
        }

        if (warpLane == 0) {
            resultsBlocks[bid] = warpVal;
        }
    }
    

  
}

__global__ void gmem_red_single_block(float * resultsBlocks, float * output, int nbBlocks) {
    // this kernel needs to be started on a single block
    float localValue = 0.0f;
    if (threadIdx.x < nbBlocks) {
        localValue = resultsBlocks[threadIdx.x];
    }
    int tid = threadIdx.x;
    extern __shared__ float shmem[];
    int warpsPerBlock = blockDim.x / 32;

    int warpId = threadIdx.x / 32;
    int warpLane = threadIdx.x % 32;
    for (int offset = 16; offset > 0; offset /=2) {
        localValue += __shfl_down_sync(0xFFFFFFFF, localValue, offset);
    }

    if (warpLane == 0) {
        shmem[warpId] = localValue;
    }
    __syncthreads();
    // block lvl reduction
    for (int offset = warpsPerBlock / 2; offset > 0; offset /= 2) {
        if (tid < offset) {
            shmem[tid] += shmem[tid + offset];
        } 
        __syncthreads();
    }

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




torch::Tensor sum_cuda(torch::Tensor input, torch::Tensor output) {
    TORCH_CHECK(input.device().is_cuda(), "Tensor input must be a CUDA tensor");
    TORCH_CHECK(output.device().is_cuda(), "Tensor output must be a CUDA tensor");
    
    int N = input.numel();  

    int warpsPerBlock = 8;
    int threadsPerBlock = warpsPerBlock * 32;
    int nbBlocks = 512;
    size_t shmemSize = warpsPerBlock * sizeof(float);

    auto resultsBlocks = torch::empty({nbBlocks}, input.options());

    first_kernel<<<nbBlocks, threadsPerBlock, shmemSize>>>(input.data_ptr<float>(), N, warpsPerBlock, resultsBlocks.data_ptr<float>());

    int threadsKernel2 = nbBlocks;
    size_t shmemKernel2 = (nbBlocks / 32) * sizeof(float);
    gmem_red_single_block<<<1,threadsKernel2,shmemKernel2>>>(resultsBlocks.data_ptr<float>(), output.data_ptr<float>(), nbBlocks);
       
    
    

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

    return output;
}
"""

sum_cpp_source = """
#include <torch/extension.h>

torch::Tensor sum_cuda(torch::Tensor input, torch::Tensor output);
"""

sum_module = load_inline(
    name='sum_cuda',
    cpp_sources=sum_cpp_source,
    cuda_sources=sum_cuda_source,
    functions=['sum_cuda'],
    verbose=True,
)

def sum(input,output):
    if not input.is_cuda or not output.is_cuda:
        raise RuntimeError("Tensor must be on GPU")
    return sum_module.sum_cuda(input,output)

def custom_kernel(data: input_t) -> output_t:
    """
    Custom implementation of vector sum reduction using CUDA.
    Args:
        inputs: List of pairs of tensors [A, B] to be added.
    Returns:
        Tensor containing element-wise sum.
    """
    input, output = data
    assert input.is_cuda, "Input tensor must be on GPU"

    # Simply reuse the existing add function we already defined
    # This avoids the compilation issues with the inline kernel
    res =  sum(input, output)
    return res.squeeze()