submission 771487

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

import triton
import triton.language as tl


@triton.jit
def reduce_kernel(
    in_ptr,
    out_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # 加载数据块
    x = tl.load(in_ptr + offsets, mask=mask, other=0.0)

    partial = tl.sum(x, axis=0)
    tl.store(out_ptr + pid, partial)


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


def custom_kernel(data: input_t) -> output_t:
    input_tensor, _ = data
    n_elements = input_tensor.numel()

    # if n_elements == 0:
    #     return torch.zeros((), device=input_tensor.device, dtype=torch.float32)

    # x = input_tensor.contiguous()

    # A100 优化参数：每个 program 处理 4 个 1024 元素块（共 4096 元素）
    # 目标是减少 launch program 数，同时保持较高并行度和内存吞吐。
    BLOCK_SIZE = 2048
    num_warps = 8
    num_stages = 4

    n_partials = triton.cdiv(n_elements, BLOCK_SIZE)
    partials = torch.empty((n_partials,), device="cuda", dtype=torch.float32)
    reduce_kernel[(n_partials,)](
        input_tensor,
        partials,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
        num_stages=num_stages,
    )
    input = partials
    output = input_tensor

    while n_partials > 1:
        next_n = triton.cdiv(n_partials, BLOCK_SIZE)
        reduce_kernel[(next_n,)](
            input,
            output,  # 重用输入缓冲区作为中间结果存储
            n_partials,
            BLOCK_SIZE=BLOCK_SIZE,
            num_warps=num_warps,
            num_stages=num_stages,
        )
        # 交换输入输出缓冲区
        input, output = output, input
        n_partials = next_n

    return input[0].to(torch.float32)


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)