submission 763183

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

sum_cuda_source = """
#include <cuda_fp16.h>
#include <cuda_runtime.h>

#define MAX_BLOCKS 4096

static double* d_partial = nullptr;
__device__ unsigned int g_retirement_count = 0;

__global__ void __launch_bounds__(512)
sum_kernel(const float* __restrict__ input,
           double* __restrict__ partial_sums,
           float* __restrict__ output,
           int N) {
    double thread_sum = 0.0;

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    int N4 = N / 4;
    const float4* input4 = reinterpret_cast<const float4*>(input);

    int vec_idx = idx;
    for (; vec_idx + stride < N4; vec_idx += 2 * stride) {
        float4 v0 = __ldg(&input4[vec_idx]);
        float4 v1 = __ldg(&input4[vec_idx + stride]);
        thread_sum += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
        thread_sum += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
    }
    for (; vec_idx < N4; vec_idx += stride) {
        float4 v = __ldg(&input4[vec_idx]);
        thread_sum += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
    }

    int scalar_start = N4 * 4;
    for (int i = scalar_start + threadIdx.x; i < N; i += blockDim.x) {
        thread_sum += (double)__ldg(&input[i]);
    }

    unsigned mask = 0xffffffff;
    for (int offset = 16; offset > 0; offset >>= 1) {
        thread_sum += __shfl_down_sync(mask, thread_sum, offset);
    }

    __shared__ double warp_sums[16];
    int lane = threadIdx.x & 31;
    int warp_id = threadIdx.x >> 5;

    if (lane == 0) {
        warp_sums[warp_id] = thread_sum;
    }
    __syncthreads();

    double block_sum = 0.0;
    if (warp_id == 0) {
        int num_warps = blockDim.x >> 5;
        thread_sum = (lane < num_warps) ? warp_sums[lane] : 0.0;
        for (int offset = 16; offset > 0; offset >>= 1) {
            thread_sum += __shfl_down_sync(mask, thread_sum, offset);
        }
        block_sum = thread_sum;
    }

    __shared__ bool s_is_last;
    if (threadIdx.x == 0) {
        partial_sums[blockIdx.x] = block_sum;
        __threadfence();
        unsigned int ticket = atomicAdd(&g_retirement_count, 1);
        s_is_last = (ticket == gridDim.x - 1);
    }
    __syncthreads();

    if (s_is_last) {
        double final_sum = 0.0;
        for (int i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
            final_sum += partial_sums[i];
        }

        for (int offset = 16; offset > 0; offset >>= 1) {
            final_sum += __shfl_down_sync(mask, final_sum, offset);
        }

        if (lane == 0) {
            warp_sums[warp_id] = final_sum;
        }
        __syncthreads();

        if (warp_id == 0) {
            int num_warps = blockDim.x >> 5;
            final_sum = (lane < num_warps) ? warp_sums[lane] : 0.0;
            for (int offset = 16; offset > 0; offset >>= 1) {
                final_sum += __shfl_down_sync(mask, final_sum, offset);
            }
            if (lane == 0) {
                output[0] = (float)final_sum;
                g_retirement_count = 0;
            }
        }
    }
}

void sum_cuda(torch::Tensor input, torch::Tensor output) {
    int N = input.numel();
    const int threads = 512;
    int blocks = min((N / 4 + threads - 1) / threads, MAX_BLOCKS);
    if (blocks < 1) blocks = 1;

    if (!d_partial) {
        cudaMalloc(&d_partial, MAX_BLOCKS * sizeof(double));
    }

    sum_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        d_partial,
        output.data_ptr<float>(),
        N
    );
}
"""

sum_cpp_source = """
#include <torch/extension.h>
void 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,
    extra_cuda_cflags=['-O3', '--use_fast_math'],
)

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    sum_module.sum_cuda(data, output)
    return output[0]