submission 779823

from utils import make_match_reference, DeterministicContext
import torch
from task import input_t, output_t
import sys

from torch.utils.cpp_extension import load_inline

N_ELEMENTS = 52428800

_CPP_SOURCE = r"""
#include <torch/extension.h>

torch::Tensor sum_reduce_cuda(std::vector<torch::Tensor> data);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("sum_reduce_cuda", &sum_reduce_cuda, "Sum reduction with custom CUDA kernel");
}
"""


_CUDA_SOURCE = r"""
#include <cuda_runtime.h>
#include <stdio.h>

// ================= 硬编码参数 (针对 52,428,800 精确计算) =================
constexpr int N_FLOATS          = 52428800;
constexpr int BLOCK_SIZE        = 256;
constexpr int GRID_SIZE         = 2048;
constexpr int FLOAT4_PER_THREAD = 25;
constexpr int STRIDE_FLOAT4     = BLOCK_SIZE * GRID_SIZE; // 524,288 (float4单位)

// ================= A100 极限优化内核 =================
__global__ void reduce_sum_50M_a100(const float* __restrict__ in, float* __restrict__ out) {
    // 仅需 8 个 slot (每个 warp 一个 leader)，32 Bytes，零 Bank Conflict
    __shared__ float sdata[8];
    
    int tid = threadIdx.x;
    // 在 float4 空间中的全局索引 (保证连续线程加载连续 float4，完美合并)
    int idx = blockIdx.x * BLOCK_SIZE + tid;

    float sum = 0.0f;
    const float4* in4 = reinterpret_cast<const float4*>(in);

    // 1. 向量化 + Grid-Stride + 完全展开的寄存器累加
    #pragma unroll
    for (int j = 0; j < FLOAT4_PER_THREAD; ++j) {
        float4 v = __ldg(&in4[idx + j * STRIDE_FLOAT4]); // 使用只读缓存加载，减少延迟
        sum += v.x + v.y + v.z + v.w;
    }

    // 2. Warp 级归约 (纯寄存器 Shuffle，无共享内存交互)
    sum += __shfl_xor_sync(0xffffffff, sum, 16);
    sum += __shfl_xor_sync(0xffffffff, sum, 8);
    sum += __shfl_xor_sync(0xffffffff, sum, 4);
    sum += __shfl_xor_sync(0xffffffff, sum, 2);
    sum += __shfl_xor_sync(0xffffffff, sum, 1);

    // 3. Warp Leader 写入共享内存
    if (tid % 32 == 0) {
        sdata[tid / 32] = sum;
    }
    __syncthreads();

    // 4. Warp 0 完成 Block 内最终归约 (8个值)
    if (tid < 32) {
        float wsum = (tid < 8) ? sdata[tid] : 0.0f;
        wsum += __shfl_xor_sync(0xffffffff, wsum, 4);
        wsum += __shfl_xor_sync(0xffffffff, wsum, 2);
        wsum += __shfl_xor_sync(0xffffffff, wsum, 1);

        if (tid == 0) {
            // A100 硬件加速 atomicAdd，2048 次原子操作开销 < 0.05%
            atomicAdd(out, wsum);
        }
    }
}


torch::Tensor sum_reduce_cuda(std::vector<torch::Tensor> data) {
    data[1].fill_(0); // 必须初始化为 0
    int n_elements = data[0].numel();
    if (n_elements == N_FLOATS) {
        reduce_sum_50M_a100<<<GRID_SIZE, BLOCK_SIZE>>>(
            data[0].data_ptr<float>(), 
            data[1].data_ptr<float>()
        );
        return data[1][0];
        
    } else { 
        // 退化到 PyTorch 内置实现，保证正确性
        return data[0].sum();
    }
}
"""


_EXT = load_inline(
        name="cuda_sum_reduce_000002_ext",
        cpp_sources=[_CPP_SOURCE],
        cuda_sources=[_CUDA_SOURCE],
        functions=None,
        extra_cflags=["-O3 -use_fast_math"],
        extra_cuda_cflags=["-O3 -use_fast_math"],
        with_cuda=True,
        verbose=False,
    )


def ref_kernel(data: input_t) -> output_t:
    """
    Reference implementation of vector sum reduction using PyTorch.
    Args:
        data: Input tensor to be reduced
    Returns:
        Tensor containing the sum of all elements
    """
    with DeterministicContext():
        data, output = data
        # Let's be on the safe side here, and do the reduction in 64 bit
        output = data.to(torch.float64).sum().to(torch.float32)
        return output

custom_kernel = _EXT.sum_reduce_cuda

def generate_input(size: int, seed: int) -> input_t:
    """
    Generates random input tensor of specified shape with random offset and scale.
    The data is first generated as standard normal, then scaled and offset
    to prevent trivial solutions.

    Returns:
        Tensor to be reduced
    """
    gen = torch.Generator(device="cuda")
    gen.manual_seed(seed)

    # Generate base random data
    data = torch.randn(
        size, device="cuda", dtype=torch.float32, generator=gen
    ).contiguous()

    # Generate random offset and scale (using different seeds to avoid correlation)
    offset_gen = torch.Generator(device="cuda")
    offset_gen.manual_seed(seed + 1)
    scale_gen = torch.Generator(device="cuda")
    scale_gen.manual_seed(seed + 2)

    # Generate random offset between -100 and 100
    offset = (torch.rand(1, device="cuda", generator=offset_gen) * 200 - 100).item()
    # Generate random scale between 0.1 and 10
    scale = (torch.rand(1, device="cuda", generator=scale_gen) * 9.9 + 0.1).item()

    # Apply scale and offset
    input_tensor = (data * scale + offset).contiguous()
    output_tensor = torch.empty(1, device="cuda", dtype=torch.float32)
    return input_tensor, output_tensor


check_implementation = make_match_reference(ref_kernel)
def warmup(fn, args, n_warmup=5):
    for _ in range(n_warmup):
        _ = fn(args)
        torch.cuda.synchronize()


# warmup(custom_kernel, generate_input(N_ELEMENTS, 42))