submission 539581

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

cuda_source = """
__device__ __forceinline__ double warp_reduce(double val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        int lo = __double2loint(val);
        int hi = __double2hiint(val);
        lo = __shfl_down_sync(0xFFFFFFFF, lo, offset);
        hi = __shfl_down_sync(0xFFFFFFFF, hi, offset);
        val += __hiloint2double(hi, lo);
    }
    return val;
}

__device__ __forceinline__ float4 ld_cs(const float4* addr) {
    float4 ret;
    asm("ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];" :
        "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) :
        "l"(addr));
    return ret;
}

static double* d_block_sums = nullptr;
static unsigned int* d_counter = nullptr;
static int alloc_blocks = 0;
static int cached_sms = 0;

__global__ void __launch_bounds__(512, 2)
vectorsum_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const int N,
    double* __restrict__ block_sums,
    unsigned int* __restrict__ block_counter
) {
    double s0 = 0.0, s1 = 0.0, s2 = 0.0, s3 = 0.0;
    double s4 = 0.0, s5 = 0.0, s6 = 0.0, s7 = 0.0;

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

    const int N4 = N >> 2;
    const float4* __restrict__ input4 = reinterpret_cast<const float4*>(input);

    const int stride8 = stride << 3;
    int i = tid;
    const int limit8 = N4 - (stride << 3) + stride;

    for (; i < limit8; i += stride8) {
        float4 v0 = ld_cs(&input4[i]);
        float4 v1 = ld_cs(&input4[i + stride]);
        float4 v2 = ld_cs(&input4[i + stride * 2]);
        float4 v3 = ld_cs(&input4[i + stride * 3]);
        float4 v4 = ld_cs(&input4[i + stride * 4]);
        float4 v5 = ld_cs(&input4[i + stride * 5]);
        float4 v6 = ld_cs(&input4[i + stride * 6]);
        float4 v7 = ld_cs(&input4[i + stride * 7]);

        s0 += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
        s1 += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
        s2 += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
        s3 += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.w;
        s4 += (double)v4.x + (double)v4.y + (double)v4.z + (double)v4.w;
        s5 += (double)v5.x + (double)v5.y + (double)v5.z + (double)v5.w;
        s6 += (double)v6.x + (double)v6.y + (double)v6.z + (double)v6.w;
        s7 += (double)v7.x + (double)v7.y + (double)v7.z + (double)v7.w;
    }

    for (; i < N4; i += stride) {
        float4 v = ld_cs(&input4[i]);
        s0 += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
    }

    const int tail = N4 << 2;
    for (int j = tail + tid; j < N; j += stride) {
        float ret;
        asm("ld.global.cs.f32 %0, [%1];" : "=f"(ret) : "l"(&input[j]));
        s0 += (double)ret;
    }

    double total = ((s0 + s1) + (s2 + s3)) + ((s4 + s5) + (s6 + s7));
    total = warp_reduce(total);

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

    if (lane == 0) shared[warp_id] = total;
    __syncthreads();

    if (warp_id == 0) {
        total = (lane < 16) ? shared[lane] : 0.0;
        total = warp_reduce(total);
    }

    __shared__ bool is_last;
    if (threadIdx.x == 0) {
        block_sums[blockIdx.x] = total;
        __threadfence();
        unsigned int count = atomicAdd(block_counter, 1u);
        is_last = (count == gridDim.x - 1);
    }
    __syncthreads();

    if (is_last) {
        double sum = 0.0;
        for (int k = threadIdx.x; k < (int)gridDim.x; k += blockDim.x) {
            sum += block_sums[k];
        }
        sum = warp_reduce(sum);
        if (lane == 0) shared[warp_id] = sum;
        __syncthreads();
        if (warp_id == 0) {
            sum = (lane < 16) ? shared[lane] : 0.0;
            sum = warp_reduce(sum);
            if (lane == 0) {
                output[0] = (float)sum;
                *block_counter = 0;
            }
        }
    }
}

torch::Tensor vectorsum_cuda(torch::Tensor input, torch::Tensor output) {
    const int N = input.numel();

    if (cached_sms == 0) {
        cudaDeviceGetAttribute(&cached_sms, cudaDevAttrMultiProcessorCount, 0);
    }

    const int grid = cached_sms * 2;

    if (grid > alloc_blocks) {
        if (d_block_sums) cudaFree(d_block_sums);
        if (d_counter) cudaFree(d_counter);
        cudaMalloc(&d_block_sums, grid * sizeof(double));
        cudaMalloc(&d_counter, sizeof(unsigned int));
        cudaMemset(d_counter, 0, sizeof(unsigned int));
        alloc_blocks = grid;
    }

    vectorsum_kernel<<<grid, 512>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        N,
        d_block_sums,
        d_counter
    );

    return output;
}
"""

cpp_source = """
#include <torch/extension.h>
torch::Tensor vectorsum_cuda(torch::Tensor input, torch::Tensor output);
"""

module = load_inline(
    name='vectorsum_fast',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['vectorsum_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math', '--expt-relaxed-constexpr'],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    input_tensor, output_tensor = data
    module.vectorsum_cuda(input_tensor, output_tensor)
    return output_tensor[0]