submission 512630

# prefixsum_v2 submission — two-pass inclusive prefix sum using Triton
#
# Algorithm:
#   Pass 1 (_block_reduce): each CTA sums its BLOCK elements → P block sums.
#   Intermediate: prefix-scan the P block sums with torch.cumsum.
#   Pass 2 (_final_scan): each CTA reloads its chunk, does a local inclusive
#       scan with tl.associative_scan, adds the inter-block prefix, stores.
#
# Data movement: 2×N reads + 1×N write — better than the 4×N of the naive
# store-partial-buffer approach.
#
# Target: NVIDIA B200 (CC 10.0), Triton 3.6.0

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


@triton.jit
def _add(a, b):
    return a + b


@triton.jit
def _block_reduce(in_ptr, sums_ptr, N: int, BLOCK: tl.constexpr):
    """Pass 1: reduce each block of BLOCK elements to a single sum."""
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    x = tl.load(in_ptr + offs, mask=mask, other=0.0)
    tl.store(sums_ptr + pid, tl.sum(x, axis=0))


@triton.jit
def _final_scan(in_ptr, out_ptr, sums_ptr, N: int, BLOCK: tl.constexpr):
    """Pass 2: local inclusive scan + add inter-block prefix."""
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    x = tl.load(in_ptr + offs, mask=mask, other=0.0)

    # Parallel inclusive prefix scan within this block
    x_scan = tl.associative_scan(x, 0, _add)

    # Inter-block prefix: sum of all elements in preceding blocks.
    # sums[i] holds the inclusive cumsum of block sums after the
    # torch.cumsum step, so sums[pid-1] = sum of blocks 0 .. pid-1.
    prev = tl.maximum(pid - 1, 0)
    prefix_raw = tl.load(sums_ptr + prev)
    prefix = tl.where(pid > 0, prefix_raw, 0.0)

    tl.store(out_ptr + offs, x_scan + prefix, mask=mask)


# Block size must be a power of 2 (required by tl.associative_scan).
# 4096 elements × 4 B = 16 KB SMEM per block — well within B200's 256 KB.
_BLOCK = 4096
_WARPS = 8  # 8×32 = 256 threads → 16 elements per thread


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

    P = triton.cdiv(N, _BLOCK)
    sums = torch.empty(P, dtype=torch.float32, device=x.device)

    # Pass 1: per-block reductions
    _block_reduce[(P,)](x, sums, N, BLOCK=_BLOCK, num_warps=_WARPS)

    # Prefix-scan the P block sums (tiny array — torch.cumsum is fine)
    if P > 1:
        torch.cumsum(sums, dim=0, out=sums)

    # Pass 2: local scan + inter-block prefix → final output
    _final_scan[(P,)](x, out, sums, N, BLOCK=_BLOCK, num_warps=_WARPS)

    return out