submission 513510

import torch
import triton
import triton.language as tl

@triton.jit
def sum_kernel(x_ptr, output_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
    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
    tl.atomic_add(output_ptr, tl.sum(x))

def custom_kernel(data):
    x = data[0].flatten()

    # Kahan summation for numerical stability
    total = torch.tensor(0.0, dtype=torch.float64, device=x.device)
    c = torch.tensor(0.0, dtype=torch.float64, device=x.device)

    # Process in chunks to reduce error accumulation
    chunk_size = 65536
    for i in range(0, x.numel(), chunk_size):
        chunk = x[i:i+chunk_size].to(torch.float64)
        y = chunk.sum() - c
        t = total + y
        c = (t - total) - y
        total = t
    return total.to(x.dtype)