submission 623314

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpu B200

import torch
import triton
import triton.language as tl
from task import input_t, output_t


BLOCK_SIZE = 1024
NUM_WARPS = 8
PIPELINE_STAGES = 4
PROGRAMS_PER_SM = 4
_WORKSPACE: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
_NUM_SMS: dict[int, int] = {}


@triton.jit
def reduce_persistent_kernel(
    x_ptr,
    partial_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
    PIPELINE_STAGES: tl.constexpr,
):
    pid = tl.program_id(0)
    program_count = tl.num_programs(0)
    offsets = tl.arange(0, BLOCK_SIZE)
    block_start = pid * BLOCK_SIZE
    block_offsets = block_start + offsets
    mask = block_offsets < n_elements

    x = tl.load(x_ptr + block_offsets, mask=mask, other=0.0)
    acc = tl.sum(x, axis=0, dtype=tl.float32)
    block_step = program_count * BLOCK_SIZE

    for block_start in tl.range(block_start + block_step, n_elements, block_step, num_stages=PIPELINE_STAGES):
        block_offsets = block_start + offsets
        mask = block_offsets < n_elements
        x = tl.load(x_ptr + block_offsets, mask=mask, other=0.0)
        acc += tl.sum(x, axis=0, dtype=tl.float32)

    tl.store(partial_ptr + pid, acc)


def _get_workspace(
    device: torch.device,
    partial_count: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    key = device.index
    buffers = _WORKSPACE.get(key)
    if buffers is None or buffers[0].numel() < partial_count:
        buffers = (
            torch.empty(partial_count, device=device, dtype=torch.float32),
            torch.empty(partial_count, device=device, dtype=torch.float32),
        )
        _WORKSPACE[key] = buffers
    return buffers


def _get_num_sms(device: torch.device) -> int:
    key = device.index
    num_sms = _NUM_SMS.get(key)
    if num_sms is None:
        num_sms = torch.cuda.get_device_properties(key).multi_processor_count
        _NUM_SMS[key] = num_sms
    return num_sms


def custom_kernel(data: input_t) -> output_t:
    x, output = data
    x = x.reshape(-1)

    if x.numel() == 0:
        output.zero_()
        return output[0]

    if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
        output.view(()).copy_(torch.sum(x))
        return output[0]

    max_programs = _get_num_sms(x.device) * PROGRAMS_PER_SM
    partial_count = min(triton.cdiv(x.numel(), BLOCK_SIZE), max_programs)
    buf_a, buf_b = _get_workspace(x.device, partial_count)

    current = x
    current_size = x.numel()
    next_buf = buf_a

    # reduce until there is only one left
    while current_size > 1:
        next_size = min(triton.cdiv(current_size, BLOCK_SIZE), max_programs)
        reduce_persistent_kernel[(next_size,)](
            current,
            next_buf,
            current_size,
            BLOCK_SIZE=BLOCK_SIZE,
            PIPELINE_STAGES=PIPELINE_STAGES,
            num_warps=NUM_WARPS,
        )
        current = next_buf
        current_size = next_size
        next_buf = buf_b if next_buf is buf_a else buf_a

    output.view(()).copy_(current[0].to(output.dtype))
    return output[0]