submission 614166

#!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(double val) {
    val += __shfl_down_sync(0xffffffff, val, 16);
    val += __shfl_down_sync(0xffffffff, val, 8);
    val += __shfl_down_sync(0xffffffff, val, 4);
    val += __shfl_down_sync(0xffffffff, val, 2);
    val += __shfl_down_sync(0xffffffff, val, 1);
    return val;
}

__device__ __forceinline__ float4 load_cs(const float4* ptr) {
    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"(ptr));
    return ret;
}

__global__ void __launch_bounds__(256)
vectorsum_kernel(
    const float* __restrict__ input,
    double* __restrict__ acc,
    const int n
) {
    double s0 = 0.0, s1 = 0.0, s2 = 0.0, s3 = 0.0;
    const int tid = threadIdx.x + blockIdx.x * 256;
    const int gs = 256 * gridDim.x;
    const int n4 = n >> 2;
    const float4* in4 = reinterpret_cast<const float4*>(input);

    int i = tid;
    for (; i + gs * 3 < n4; i += gs * 4) {
        float4 v0 = load_cs(&in4[i]);
        float4 v1 = load_cs(&in4[i + gs]);
        float4 v2 = load_cs(&in4[i + gs * 2]);
        float4 v3 = load_cs(&in4[i + gs * 3]);
        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;
    }
    for (; i < n4; i += gs) {
        float4 v = load_cs(&in4[i]);
        s0 += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
    }
    for (int j = (n4 << 2) + tid; j < n; j += gs)
        s0 += (double)__ldg(&input[j]);

    double ts = (s0 + s1) + (s2 + s3);
    ts = warp_reduce(ts);

    __shared__ double ws[8];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) ws[wid] = ts;
    __syncthreads();
    if (wid == 0) {
        double v = (lane < 8) ? ws[lane] : 0.0;
        v = warp_reduce(v);
        if (lane == 0) atomicAdd(acc, v);
    }
}

__global__ void finalize(double* __restrict__ acc, float* __restrict__ out) {
    *out = (float)(*acc);
    *acc = 0.0;
}

void vectorsum(torch::Tensor input, torch::Tensor output,
               torch::Tensor acc, int nb) {
    vectorsum_kernel<<<nb, 256>>>(
        input.data_ptr<float>(), acc.data_ptr<double>(), input.numel());
    finalize<<<1, 1>>>(acc.data_ptr<double>(), output.data_ptr<float>());
}
"""

cpp_source = "void vectorsum(torch::Tensor, torch::Tensor, torch::Tensor, int);"

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

_acc = torch.zeros(1, dtype=torch.float64, device="cuda")


def _tune():
    import sys
    sys.path.insert(0, '.')
    from reference import generate_input
    size = 52428800
    best_t, best_nb = 1e9, 5920
    for nb in [2960, 4096, 5920]:
        data = generate_input(size=size, seed=12345)
        for _ in range(3):
            module.vectorsum(data[0], data[1], _acc, nb)
            torch.cuda.synchronize()
        ts = []
        for _ in range(15):
            data = generate_input(size=size, seed=12345)
            torch.cuda.synchronize()
            s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
            s.record(); module.vectorsum(data[0], data[1], _acc, nb); e.record()
            torch.cuda.synchronize(); ts.append(s.elapsed_time(e))
        a = sum(ts)/len(ts)
        if a < best_t: best_t, best_nb = a, nb
    return best_nb

_nb = _tune()


def custom_kernel(data: input_t) -> output_t:
    inp, out = data
    n = inp.numel()
    nb = max(1, min(_nb, n // 64))
    module.vectorsum(inp, out, _acc, nb)
    return out[0]