submission 612741

#!POPCORN leaderboard vectorsum_v2

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

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

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

__device__ __forceinline__ double warp_reduce_sum(double val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void __launch_bounds__(512)
vectorsum_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    double* __restrict__ partial_buf,
    unsigned int* __restrict__ retire_count,
    const int n,
    const int nblocks
) {
    double sum0 = 0.0;
    double sum1 = 0.0;

    const int tid = threadIdx.x + blockIdx.x * 512;
    const int grid_stride = 512 * nblocks;

    const int n4 = n >> 2;
    const float4* input4 = reinterpret_cast<const float4*>(input);

    int i = tid;
    for (; i + grid_stride < n4; i += grid_stride * 2) {
        float4 v0 = __ldg(&input4[i]);
        float4 v1 = __ldg(&input4[i + grid_stride]);
        sum0 += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
        sum1 += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
    }
    for (; i < n4; i += grid_stride) {
        float4 v = __ldg(&input4[i]);
        sum0 += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
    }

    for (int j = (n4 << 2) + tid; j < n; j += grid_stride) {
        sum0 += (double)__ldg(&input[j]);
    }

    double thread_sum = sum0 + sum1;
    thread_sum = warp_reduce_sum(thread_sum);

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

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

    if (warp_id == 0) {
        double val = (lane < 16) ? warp_sums[lane] : 0.0;
        val = warp_reduce_sum(val);
        if (lane == 0) {
            partial_buf[blockIdx.x] = val;
        }
    }

    __threadfence();

    __shared__ bool is_last;
    if (threadIdx.x == 0) {
        unsigned int ticket = atomicInc(retire_count, 0x7fffffff);
        is_last = (ticket == (unsigned int)(nblocks - 1));
    }
    __syncthreads();

    if (is_last) {
        double final_sum = 0.0;
        for (int j = threadIdx.x; j < nblocks; j += 512) {
            final_sum += partial_buf[j];
        }
        final_sum = warp_reduce_sum(final_sum);
        if (lane == 0) warp_sums[warp_id] = final_sum;
        __syncthreads();
        if (warp_id == 0) {
            double val = (lane < 16) ? warp_sums[lane] : 0.0;
            val = warp_reduce_sum(val);
            if (lane == 0) {
                *output = (float)val;
                *retire_count = 0;
            }
        }
    }
}

void vectorsum(torch::Tensor input, torch::Tensor output,
               torch::Tensor partial_buf, torch::Tensor retire_count,
               int nblocks) {
    const int n = input.numel();
    vectorsum_kernel<<<nblocks, 512>>>(
        input.data_ptr<float>(), output.data_ptr<float>(),
        partial_buf.data_ptr<double>(), (unsigned int*)retire_count.data_ptr<int>(),
        n, nblocks);
}
"""

cpp_source = """
void vectorsum(torch::Tensor input, torch::Tensor output,
               torch::Tensor partial_buf, torch::Tensor retire_count,
               int nblocks);
"""

module = load_inline(
    name="vectorsum_cuda",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["vectorsum"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
    verbose=False,
)

_partial_buf = torch.zeros(8192, dtype=torch.float64, device="cuda")
_retire_count = torch.zeros(1, dtype=torch.int32, device="cuda")


def custom_kernel(data: input_t) -> output_t:
    inp, out = data
    n = inp.numel()
    # Size-adaptive block count
    if n <= 2 * 1024 * 1024:
        nb = 296
    elif n <= 8 * 1024 * 1024:
        nb = 592
    elif n <= 32 * 1024 * 1024:
        nb = 592
    elif n <= 128 * 1024 * 1024:
        nb = 2048
    else:
        nb = 4096
    module.vectorsum(inp, out, _partial_buf, _retire_count, nb)
    return out[0]