submission 757213

# EVOLVE-BLOCK-START

import torch
import triton
import triton.language as tl
from typing import Tuple

# ----------------------------------------------------------------------
# Combine function for associative scan (simple addition)
# ----------------------------------------------------------------------
@triton.jit
def _add_combine(a, b):
    return a + b

# ----------------------------------------------------------------------
# Kernel 1: compute per‑block sums of the input tensor
# ----------------------------------------------------------------------
@triton.jit
def block_sum_kernel(
    x_ptr,                # *Pointer* to input tensor
    block_sum_ptr,        # *Pointer* to per‑block sum buffer
    n: tl.int64,          # total number of elements
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)                     # block index
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n

    # Load a chunk of the input (zero‑pad out‑of‑bounds)
    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)

    # Reduce to a single sum for this block
    block_sum = tl.sum(x)

    # Write the per‑block sum
    tl.store(block_sum_ptr + pid, block_sum)

# ----------------------------------------------------------------------
# Kernel 2: final inclusive scan + add exclusive per‑block offset
# ----------------------------------------------------------------------
@triton.jit
def final_scan_kernel(
    x_ptr,                # *Pointer* to input tensor
    out_ptr,              # *Pointer* to output tensor
    block_sum_ptr,        # *Pointer* to per‑block sum buffer (original)
    block_sum_scan_ptr,   # *Pointer* to inclusive scan of per‑block sums
    n: tl.int64,          # total number of elements
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)                     # block index
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n

    # Load the input slice for this block
    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)

    # Inclusive scan within the block
    local_scan = tl.associative_scan(x, axis=0, combine_fn=_add_combine)

    # Load per‑block sums (original) and its inclusive scan
    block_sum = tl.load(block_sum_ptr + pid)
    block_sum_scan = tl.load(block_sum_scan_ptr + pid)

    # Exclusive offset = sum of all previous blocks
    offset = block_sum_scan - block_sum

    # Apply the offset to the locally scanned values
    out = local_scan + offset

    # Write the final result
    tl.store(out_ptr + offsets, out, mask=mask)

# ----------------------------------------------------------------------
# Public API
# ----------------------------------------------------------------------
def custom_kernel(data: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
    """
    Inclusive prefix sum (cumsum) over a 1‑D float32 tensor using a two‑pass
    Triton implementation with a small‑size fallback to torch.cumsum.
    """
    x, output = data

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------
    if not (x.is_cuda and output.is_cuda):
        raise RuntimeError("custom_kernel expects CUDA tensors")
    if x.dtype != torch.float32 or output.dtype != torch.float32:
        raise RuntimeError("custom_kernel only supports float32 tensors")
    if x.shape != output.shape:
        raise ValueError("input and output must have the same shape")
    if x.ndim != 1:
        raise ValueError("custom_kernel only supports 1‑D tensors")

    n = x.numel()
    if n == 0:
        return output

    # ------------------------------------------------------------------
    # Small‑tensor fast path using native torch.cumsum
    # ------------------------------------------------------------------
    SMALL_THRESHOLD = 1 << 20  # 1,048,576 elements
    if n <= SMALL_THRESHOLD:
        torch.cumsum(x, dim=0, out=output)
        return output

    # ------------------------------------------------------------------
    # Triton launch configuration
    # ------------------------------------------------------------------
    BLOCK_SIZE = 32768                 # elements per block (multiple of 32)
    NUM_WARPS_SUM = 8                  # 8 warps → 256 threads per block (reduction)
    NUM_WARPS_SCAN = 8                 # 8 warps → 256 threads per block (scan)
    num_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE

    # Buffer for per‑block sums
    block_sums = torch.empty(num_blocks, dtype=x.dtype, device=x.device)

    # ------------------------------------------------------------------
    # Pass 1: compute per‑block sums
    # ------------------------------------------------------------------
    block_sum_kernel[(num_blocks,)](
        x,
        block_sums,
        n,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=NUM_WARPS_SUM,
        num_stages=3,          # modest pipelining for the reduction kernel
    )

    # ------------------------------------------------------------------
    # Scan the per‑block totals on the GPU (inclusive)
    # ------------------------------------------------------------------
    block_sums_scan = torch.cumsum(block_sums, dim=0)

    # ------------------------------------------------------------------
    # Pass 2: local inclusive scan + add exclusive per‑block offsets
    # ------------------------------------------------------------------
    final_scan_kernel[(num_blocks,)](
        x,
        output,
        block_sums,
        block_sums_scan,
        n,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=NUM_WARPS_SCAN,
        num_stages=6,          # deeper pipelining to hide memory latency
    )

    return output
# EVOLVE-BLOCK-END