submission 516170

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpu A100

import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t


THREADS = 256
SMALL_THRESHOLD = 8 * 1024 * 1024
LARGE_THRESHOLD = 32 * 1024 * 1024
SMALL_BLOCKS_PER_SM = 4
MEDIUM_BLOCKS_PER_SM = 8
LARGE_BLOCKS_PER_SM = 16


CPP_SRC = """
torch::Tensor vecsum_cuda(torch::Tensor input, torch::Tensor partial, torch::Tensor output);
"""

CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdint>

__device__ __forceinline__ float warp_reduce_sum(float v) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

__device__ __forceinline__ float block_reduce_sum(float v) {
    __shared__ float warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int warp_id = threadIdx.x >> 5;
    const int num_warps = blockDim.x >> 5;

    v = warp_reduce_sum(v);
    if (lane == 0) {
        warp_sums[warp_id] = v;
    }
    __syncthreads();

    float block_sum = (threadIdx.x < num_warps) ? warp_sums[lane] : 0.0f;
    if (warp_id == 0) {
        block_sum = warp_reduce_sum(block_sum);
    }
    return block_sum;
}

template <typename scalar_t, int ITEMS_PER_THREAD>
__global__ void partial_reduce_kernel(const scalar_t* __restrict__ input,
                                      float* __restrict__ partial,
                                      int64_t n) {
    float thread_sum = 0.0f;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;

    for (int64_t base = idx; base < n; base += stride * ITEMS_PER_THREAD) {
#pragma unroll
        for (int j = 0; j < ITEMS_PER_THREAD; ++j) {
            int64_t i = base + (int64_t)j * stride;
            if (i < n) {
                thread_sum += static_cast<float>(input[i]);
            }
        }
    }

    float block_sum = block_reduce_sum(thread_sum);
    if (threadIdx.x == 0) {
        partial[blockIdx.x] = block_sum;
    }
}

template <int ITEMS_PER_THREAD>
__global__ void partial_reduce_kernel_float4(const float* __restrict__ input,
                                             float* __restrict__ partial,
                                             int64_t n) {
    float thread_sum = 0.0f;
    const int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
    const int64_t n_vec = n / 4;
    const float4* input4 = reinterpret_cast<const float4*>(input);

    for (int64_t base = idx; base < n_vec; base += stride * ITEMS_PER_THREAD) {
#pragma unroll
        for (int j = 0; j < ITEMS_PER_THREAD; ++j) {
            const int64_t i = base + (int64_t)j * stride;
            if (i < n_vec) {
                const float4 v = input4[i];
                thread_sum += v.x + v.y + v.z + v.w;
            }
        }
    }

    for (int64_t i = n_vec * 4 + idx; i < n; i += stride) {
        thread_sum += input[i];
    }

    float block_sum = block_reduce_sum(thread_sum);
    if (threadIdx.x == 0) {
        partial[blockIdx.x] = block_sum;
    }
}

__global__ void final_reduce_kernel(const float* __restrict__ partial,
                                    float* __restrict__ output,
                                    int64_t n_partials) {
    float thread_sum = 0.0f;
    for (int64_t i = threadIdx.x; i < n_partials; i += blockDim.x) {
        thread_sum += partial[i];
    }

    float block_sum = block_reduce_sum(thread_sum);
    if (threadIdx.x == 0) {
        output[0] = block_sum;
    }
}

torch::Tensor vecsum_cuda(torch::Tensor input, torch::Tensor partial, torch::Tensor output) {
    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(partial.is_cuda(), "partial must be CUDA");
    TORCH_CHECK(output.is_cuda(), "output must be CUDA");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
    TORCH_CHECK(partial.is_contiguous(), "partial must be contiguous");
    TORCH_CHECK(partial.scalar_type() == torch::kFloat32, "partial must be float32");
    TORCH_CHECK(output.scalar_type() == torch::kFloat32, "output must be float32");
    TORCH_CHECK(partial.numel() >= 1, "partial must have at least one element");
    TORCH_CHECK(output.numel() >= 1, "output must have at least one element");

    const auto n = input.numel();
    if (n == 0) {
        output.zero_();
        return output;
    }

    const int threads = 256;
    const int blocks = static_cast<int>(partial.numel());
    if (input.scalar_type() == torch::kFloat32 &&
        (reinterpret_cast<std::uintptr_t>(input.data_ptr<float>()) & 0xF) == 0) {
        if (n >= 32 * 1024 * 1024) {
            partial_reduce_kernel_float4<8><<<blocks, threads>>>(
                input.data_ptr<float>(),
                partial.data_ptr<float>(),
                n
            );
        } else {
            partial_reduce_kernel_float4<4><<<blocks, threads>>>(
                input.data_ptr<float>(),
                partial.data_ptr<float>(),
                n
            );
        }
    } else {
        AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "vecsum_partial", ([&] {
            if (n >= 32 * 1024 * 1024) {
                partial_reduce_kernel<scalar_t, 8><<<blocks, threads>>>(
                    input.data_ptr<scalar_t>(),
                    partial.data_ptr<float>(),
                    n
                );
            } else {
                partial_reduce_kernel<scalar_t, 4><<<blocks, threads>>>(
                    input.data_ptr<scalar_t>(),
                    partial.data_ptr<float>(),
                    n
                );
            }
        }));
    }
    final_reduce_kernel<<<1, threads>>>(
        partial.data_ptr<float>(),
        output.data_ptr<float>(),
        partial.numel()
    );

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err));
    return output;
}
"""

_module = load_inline(
    name="vectorsum_inline_cuda_v8",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["vecsum_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
    verbose=False,
)

_TMP_OUTPUT = {}
_PARTIAL = {}
_SM_COUNT = {}


def _device_key(device: torch.device) -> tuple[str, int | None]:
    return (device.type, device.index)


def _sm_count(device: torch.device) -> int:
    key = _device_key(device)
    sm = _SM_COUNT.get(key)
    if sm is None:
        sm = torch.cuda.get_device_properties(device).multi_processor_count
        _SM_COUNT[key] = sm
    return sm


def _blocks_per_sm(n: int) -> int:
    if n < SMALL_THRESHOLD:
        return SMALL_BLOCKS_PER_SM
    if n >= LARGE_THRESHOLD:
        return LARGE_BLOCKS_PER_SM
    return MEDIUM_BLOCKS_PER_SM


def _items_per_thread(n: int) -> int:
    return 8 if n >= LARGE_THRESHOLD else 4


def _choose_partial_blocks(n: int, device: torch.device) -> int:
    sm = _sm_count(device)
    target_blocks = sm * _blocks_per_sm(n)
    items = _items_per_thread(n)
    useful_blocks = (n + (THREADS * items) - 1) // (THREADS * items)
    return max(1, min(target_blocks, useful_blocks))


def _tmp_output(device: torch.device) -> torch.Tensor:
    key = _device_key(device)
    buf = _TMP_OUTPUT.get(key)
    if buf is None:
        buf = torch.empty(1, device=device, dtype=torch.float32)
        _TMP_OUTPUT[key] = buf
    return buf


def _partial_buffer(device: torch.device, blocks: int) -> torch.Tensor:
    key = _device_key(device)
    buf = _PARTIAL.get(key)
    if buf is None or buf.numel() < blocks:
        buf = torch.empty(blocks, device=device, dtype=torch.float32)
        _PARTIAL[key] = buf
    return buf[:blocks]


def custom_kernel(data: input_t) -> output_t:
    inp, output = data
    x = inp if inp.is_contiguous() else inp.contiguous()

    blocks = _choose_partial_blocks(x.numel(), x.device)
    partial = _partial_buffer(x.device, blocks)
    tmp = _tmp_output(x.device)

    _module.vecsum_cuda(x, partial, tmp)
    output[0] = tmp[0].to(output.dtype)
    return output[0]