submission 510383

import triton
import triton.language as tl
import torch


@triton.jit
def _prefix_sum_single_block_kernel(
    input_ptr,
    output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Single-block inclusive prefix sum using a work-efficient parallel scan.
    This kernel handles arrays that fit within a single block.
    """
    # Load all elements into shared memory
    offsets = tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    # Load input values
    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
    
    # Perform inclusive scan using Hillis-Steele algorithm
    # This is a simple parallel prefix sum that works well for moderate sizes
    stride = 1
    while stride < BLOCK_SIZE:
        # Get value from 'stride' positions before
        shifted_offsets = offsets - stride
        shifted_mask = shifted_offsets >= 0
        
        # Load the shifted values (from the same array x)
        x_shifted = tl.where(shifted_mask, 
                             tl.load(input_ptr + shifted_offsets, mask=shifted_mask & (shifted_offsets < n_elements), other=0.0),
                             0.0)
        
        # This approach won't work directly - we need to be more careful
        stride *= 2
    
    # Store results
    tl.store(output_ptr + offsets, x, mask=mask)


@triton.jit
def _local_scan_kernel(
    input_ptr,
    output_ptr,
    block_sums_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """
    First pass: compute local prefix sums within each block and store block totals.
    """
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    # Load input values
    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
    
    # Compute inclusive prefix sum within block using cumsum
    # Triton doesn't have built-in cumsum, so we use associative scan
    result = tl.cumsum(x, axis=0)
    
    # Store local prefix sums
    tl.store(output_ptr + offsets, result, mask=mask)
    
    # Store the block total (last valid element's prefix sum)
    # We need the sum of all elements in this block
    block_sum = tl.sum(x, axis=0)
    tl.store(block_sums_ptr + pid, block_sum)


@triton.jit  
def _add_block_prefix_kernel(
    output_ptr,
    block_prefix_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Third pass: add the prefix sum of previous blocks to each element.
    """
    pid = tl.program_id(0)
    
    # Skip first block - it doesn't need adjustment
    if pid == 0:
        return
    
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    # Load the prefix sum of all previous blocks
    prefix = tl.load(block_prefix_ptr + pid - 1)
    
    # Load current values and add prefix
    current = tl.load(output_ptr + offsets, mask=mask, other=0.0)
    result = current + prefix
    
    # Store updated values
    tl.store(output_ptr + offsets, result, mask=mask)


def kernel_function(data: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
    """
    Computes inclusive prefix sum (cumulative sum) of the input tensor.
    
    This is a fused multi-pass implementation:
    1. Local scan within blocks + compute block totals
    2. Scan of block totals (done sequentially for small number of blocks)
    3. Add block prefixes to local results
    
    Args:
        data: Input 1D tensor
        output: Output tensor (pre-allocated)
    
    Returns:
        Output tensor containing inclusive prefix sum
    """
    n_elements = data.numel()
    
    # Choose block size
    BLOCK_SIZE = 1024
    
    # Number of blocks needed
    n_blocks = triton.cdiv(n_elements, BLOCK_SIZE)
    
    # Allocate temporary storage for block sums
    block_sums = torch.empty(n_blocks, device=data.device, dtype=torch.float32)
    
    # Pass 1: Local prefix sums within each block
    grid = (n_blocks,)
    _local_scan_kernel[grid](
        data, output, block_sums,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    
    # Pass 2: Compute prefix sum of block sums
    # For small number of blocks, we can do this with a single block
    if n_blocks > 1:
        block_prefix = torch.empty(n_blocks, device=data.device, dtype=torch.float32)
        
        # Use a simple kernel to compute prefix sum of block sums
        _compute_block_prefix[(1,)](
            block_sums, block_prefix, n_blocks, 
            BLOCK_SIZE=triton.next_power_of_2(n_blocks),
        )
        
        # Pass 3: Add block prefixes to local results
        _add_block_prefix_kernel[grid](
            output, block_prefix,
            n_elements,
            BLOCK_SIZE=BLOCK_SIZE,
        )
    
    return output


@triton.jit
def _compute_block_prefix(
    block_sums_ptr,
    block_prefix_ptr,
    n_blocks,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Compute inclusive prefix sum of block sums.
    Single block kernel for small arrays.
    """
    offsets = tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_blocks
    
    # Load block sums
    x = tl.load(block_sums_ptr + offsets, mask=mask, other=0.0)
    
    # Compute inclusive prefix sum
    result = tl.cumsum(x, axis=0)
    
    # Store results
    tl.store(block_prefix_ptr + offsets, result, mask=mask)

import inspect

def custom_kernel(input):
    sig = inspect.signature(kernel_function)
    num_params = len(sig.parameters)

    if len(input) == num_params:
        return kernel_function(*input)
    return kernel_function(input)


# Ensure deterministic cuBLAS.
import os
if os.environ.get("CUBLAS_WORKSPACE_CONFIG", "") not in (":4096:8", ":16:8"):
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"