submission 513376

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


@triton.jit
def _reduce_kernel(x_ptr, block_sums_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offset < n
    x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32)
    block_sum = tl.sum(x, axis=0)
    tl.store(block_sums_ptr + pid, block_sum)


@triton.jit
def _scan_and_add_kernel(x_ptr, out_ptr, block_sums_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offset < n
    x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32)
    scan = tl.cumsum(x, axis=0)
    if pid > 0:
        prefix = tl.load(block_sums_ptr + pid - 1)
        scan = scan + prefix
    tl.store(out_ptr + offset, scan, mask=mask)


def _recursive_scan(block_sums, num_blocks, BLOCK_SIZE=8192):
    if num_blocks <= 65536:
        torch.cumsum(block_sums[:num_blocks], dim=0, out=block_sums[:num_blocks])
    else:
        nb2 = (num_blocks + BLOCK_SIZE - 1) // BLOCK_SIZE
        block_sums2 = torch.empty(nb2, device=block_sums.device, dtype=torch.float32)
        _reduce_kernel[(nb2,)](block_sums, block_sums2, num_blocks, BLOCK_SIZE=BLOCK_SIZE)
        _recursive_scan(block_sums2, nb2, BLOCK_SIZE)
        _scan_and_add_kernel[(nb2,)](block_sums, block_sums, block_sums2, num_blocks, BLOCK_SIZE=BLOCK_SIZE)


def custom_kernel(data: input_t) -> output_t:
    x, output = data
    n = x.numel()

    if n <= 0:
        return output

    BLOCK_SIZE = 8192

    if n <= BLOCK_SIZE:
        torch.cumsum(x, dim=0, out=output)
        return output

    num_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
    block_sums = torch.empty(num_blocks, device=x.device, dtype=torch.float32)

    # Pass 1: Reduce each block
    _reduce_kernel[(num_blocks,)](x, block_sums, n, BLOCK_SIZE=BLOCK_SIZE)

    # Pass 2: Scan block sums
    _recursive_scan(block_sums, num_blocks, BLOCK_SIZE)

    # Pass 3: Local scan + add prefix
    _scan_and_add_kernel[(num_blocks,)](x, output, block_sums, n, BLOCK_SIZE=BLOCK_SIZE)

    return output