submission 545126

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

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

# Two-level reduction: GPU blocks reduce locally, then CPU sums partials
# Use float64 accumulation for accuracy matching reference

@triton.jit
def sum_reduce_kernel(
    x_ptr, partial_ptr, n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    # Load as float32, accumulate in float64
    x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
    block_sum = tl.sum(x, axis=0)
    tl.store(partial_ptr + pid, block_sum)

@triton.jit
def sum_final_kernel(
    partial_ptr, out_ptr, n_partials,
    BLOCK_SIZE: tl.constexpr,
):
    offsets = tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_partials
    x = tl.load(partial_ptr + offsets, mask=mask, other=0.0)
    total = tl.sum(x, axis=0)
    tl.store(out_ptr, total.to(tl.float32))

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    n = data.numel()
    BLOCK_SIZE = 4096
    n_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
    partial = torch.empty(n_blocks, device=data.device, dtype=torch.float64)
    sum_reduce_kernel[(n_blocks,)](data, partial, n, BLOCK_SIZE=BLOCK_SIZE)
    
    # Second level reduction on GPU if too many partials
    if n_blocks <= 65536:
        # Use power-of-2 BLOCK_SIZE for final reduction
        FINAL_BS = 1
        while FINAL_BS < n_blocks:
            FINAL_BS *= 2
        if FINAL_BS > 65536:
            FINAL_BS = 65536
        out_scalar = torch.empty(1, device=data.device, dtype=torch.float32)
        sum_final_kernel[(1,)](partial, out_scalar, n_blocks, BLOCK_SIZE=FINAL_BS)
        return out_scalar[0]
    else:
        return partial.sum().to(torch.float32)