submission 102359

import os
# Unset CUDA_VISIBLE_DEVICES before importing torch/triton
if 'CUDA_VISIBLE_DEVICES' in os.environ:
    del os.environ['CUDA_VISIBLE_DEVICES']

import torch
import triton
import triton.language as tl


@triton.jit
def prefixsum_kernel(
    input_ptr,
    output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """Compute prefix sum within each block."""
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # Load data from input
    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)

    # Compute prefix sum within block
    x = tl.cumsum(x, axis=0)

    # Store the result
    tl.store(output_ptr + offsets, x, mask=mask)


@triton.jit
def extract_block_sums_kernel(
    data_ptr,
    block_sums_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """Extract the last element of each block (the block sum)."""
    pid = tl.program_id(axis=0)

    # Last index in this block
    last_idx = min((pid + 1) * BLOCK_SIZE - 1, n_elements - 1)

    # Load the last element (sum of entire block)
    block_sum = tl.load(data_ptr + last_idx)

    # Store to block_sums array
    tl.store(block_sums_ptr + pid, block_sum)


@triton.jit
def add_block_offsets_kernel(
    output_ptr,
    block_offsets_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """Add the offset from previous blocks to each element."""
    pid = tl.program_id(axis=0)

    if pid == 0:
        return  # First block has no offset

    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # Load the cumulative offset for this block
    offset = tl.load(block_offsets_ptr + pid - 1)

    # Load current values, add offset, and store back
    x = tl.load(output_ptr + offsets, mask=mask)
    x = x + offset
    tl.store(output_ptr + offsets, x, mask=mask)


def prefixsum_recursive(input_tensor, output_tensor, block_size=1024):
    """
    Recursive GPU-based prefix sum.

    Args:
        input_tensor: Input tensor (1D)
        output_tensor: Output tensor (1D, same size as input)
        block_size: Size of each block for parallel processing

    Returns:
        output_tensor: Result of prefix sum
    """
    n_elements = input_tensor.numel()
    num_blocks = triton.cdiv(n_elements, block_size)

    # Base case: single block
    if num_blocks == 1:
        grid = (1,)
        prefixsum_kernel[grid](
            input_tensor,
            output_tensor,
            n_elements,
            BLOCK_SIZE=block_size,
        )
        return output_tensor

    # Recursive case: multiple blocks
    grid = (num_blocks,)

    # Step 1: Compute prefix sum within each block
    prefixsum_kernel[grid](
        input_tensor,
        output_tensor,
        n_elements,
        BLOCK_SIZE=block_size,
    )

    # Step 2: Extract block sums (last element of each block)
    block_sums = torch.empty(num_blocks, dtype=input_tensor.dtype, device=input_tensor.device)
    extract_block_sums_kernel[grid](
        output_tensor,
        block_sums,
        n_elements,
        BLOCK_SIZE=block_size,
    )

    # Step 3: Recursively compute prefix sum of block sums
    block_offsets = torch.empty_like(block_sums)
    prefixsum_recursive(block_sums, block_offsets, block_size=block_size)

    # Step 4: Add block offsets to all elements
    add_block_offsets_kernel[grid](
        output_tensor,
        block_offsets,
        n_elements,
        BLOCK_SIZE=block_size,
    )

    return output_tensor


def custom_kernel(data, block_size=1024):
    """
    Optimized prefix sum using fully GPU-based recursive algorithm.

    Args:
        data: Tuple of (input_tensor, output_tensor)
        block_size: Block size for processing (default 1024)
                   Larger blocks = fewer recursion levels but less parallelism per block
                   Smaller blocks = more recursion but better intra-block parallelism

    Returns:
        output_tensor: Result of prefix sum
    """
    input_tensor, output_tensor = data

    # Auto-tune block size based on input size if desired
    n = input_tensor.numel()
    if block_size == 'auto':
        # Heuristic: use larger blocks for larger inputs
        if n < 1024:
            block_size = 256
        elif n < 16384:
            block_size = 1024
        elif n < 262144:
            block_size = 2048
        else:
            block_size = 4096

    return prefixsum_recursive(input_tensor, output_tensor, block_size=block_size)


# Backward compatibility
prefixsum = custom_kernel