submission 612853

#!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;
}

__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__ accumulator,
    const int n
) {
    double sum0 = 0.0, sum1 = 0.0, sum2 = 0.0, sum3 = 0.0;

    const int tid = threadIdx.x + blockIdx.x * 256;
    const int grid_stride = 256 * gridDim.x;

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

    int i = tid;
    const int gs4 = grid_stride * 4;
    for (; i + grid_stride * 3 < n4; i += gs4) {
        float4 v0 = load_cs(&input4[i]);
        float4 v1 = load_cs(&input4[i + grid_stride]);
        float4 v2 = load_cs(&input4[i + grid_stride * 2]);
        float4 v3 = load_cs(&input4[i + grid_stride * 3]);
        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;
        sum2 += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
        sum3 += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.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) + (sum2 + sum3);
    thread_sum = warp_reduce_sum(thread_sum);

    __shared__ double warp_sums[8];
    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 < 8) ? warp_sums[lane] : 0.0;
        val = warp_reduce_sum(val);
        if (lane == 0) {
            atomicAdd(accumulator, val);
        }
    }
}

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

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

cpp_source = """
void vectorsum(torch::Tensor input, torch::Tensor output,
               torch::Tensor accumulator, 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,
)

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


def _autotune():
    import sys
    sys.path.insert(0, '.')
    from reference import generate_input

    size = 52428800
    best_time = 1e9
    best_nb = 5920

    for nb in [1480, 2960, 4096, 5920, 8192, 11840]:
        data = generate_input(size=size, seed=12345)
        for _ in range(5):
            module.vectorsum(data[0], data[1], _accumulator, nb)
            torch.cuda.synchronize()
        times = []
        for _ in range(30):
            data = generate_input(size=size, seed=12345)
            torch.cuda.synchronize()
            s = torch.cuda.Event(enable_timing=True)
            e = torch.cuda.Event(enable_timing=True)
            s.record()
            module.vectorsum(data[0], data[1], _accumulator, nb)
            e.record()
            torch.cuda.synchronize()
            times.append(s.elapsed_time(e))
        avg = sum(times)/len(times)
        best = min(times)
        bw = size * 4 / (best * 1e-3) / 1e12
        print(f'nb={nb}: avg={avg*1000:.1f}us best={best*1000:.1f}us bw={bw:.2f}TB/s')
        if avg < best_time:
            best_time = avg
            best_nb = nb

    print(f'Best: nb={best_nb} avg={best_time*1000:.1f}us')
    return best_nb

_best_nb = _autotune()


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