submission 512774

# submission_vectorsum_v2.py
# Sum reduction of a 1D float32 tensor → scalar.
# Interface: custom_kernel((input_tensor, output_tensor)) -> output_tensor
#   input_tensor:  (size,) float32
#   output_tensor: (1,)   float32 (pre-allocated, but reference returns 0-d)
#
# The reference returns data.to(float64).sum().to(float32) — a 0-d tensor.
# We must return a 0-d tensor too to pass the shape check in check_implementation.
# Float64 accumulators preserve accuracy for large arrays.

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

_BLOCK = 4096
_WARPS = 8


@triton.jit
def _partial_reduce(x_ptr, partial_ptr, N: int, BLOCK: tl.constexpr):
    """Reduce BLOCK float32 elements → one float64 partial sum."""
    pid  = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    x    = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float64)
    tl.store(partial_ptr + pid, tl.sum(x, axis=0))


def custom_kernel(data: input_t) -> output_t:
    x, out = data
    N = x.numel()

    if N == 0:
        result = torch.zeros((), dtype=torch.float32, device=x.device)
        out[0] = result
        return result

    P       = triton.cdiv(N, _BLOCK)
    partial = torch.empty(P, dtype=torch.float64, device=x.device)

    _partial_reduce[(P,)](x, partial, N, BLOCK=_BLOCK, num_warps=_WARPS)

    # Final reduction on P float64 partials → 0-d float32 scalar.
    # The reference also returns a 0-d tensor, so we match that shape.
    result = partial.sum().to(torch.float32)  # 0-d tensor
    out[0] = result
    return result