submission 539398

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 volatile("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;
}

__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 volatile("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,
                             torch::Tensor block_sums, torch::Tensor counter) {
    const int N = input.numel();
    const int grid = block_sums.size(0);

    vectorsum_kernel<<<grid, 512>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        N,
        block_sums.data_ptr<double>(),
        (unsigned int*)counter.data_ptr<int>()
    );

    return output;
}
"""

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

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=True,
)

_block_sums = None
_counter = None


def custom_kernel(data: input_t) -> output_t:
    global _block_sums, _counter
    input_tensor, output_tensor = data
    if _block_sums is None:
        num_sms = torch.cuda.get_device_properties(0).multi_processor_count
        grid = num_sms * 2
        _block_sums = torch.zeros(grid, dtype=torch.float64, device=input_tensor.device)
        _counter = torch.zeros(1, dtype=torch.int32, device=input_tensor.device)
    module.vectorsum_cuda(input_tensor, output_tensor, _block_sums, _counter)
    return output_tensor[0]