submission 598236

import triton
import triton.language as tl
import torch


@triton.jit
def _sum_reduction_kernel(
    input_ptr,
    partial_sums_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """
    First stage: Each block computes partial sum of BLOCK_SIZE elements.
    Uses float32 accumulation for numerical stability.
    """
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    
    # Create offsets for this block
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    # Load elements with masking for out-of-bounds
    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
    
    # Convert to float32 for accumulation to improve precision
    x_f32 = x.to(tl.float32)
    
    # Compute sum within this block using tl.sum reduction
    block_sum = tl.sum(x_f32, axis=0)
    
    # Store partial sum - only one value per block
    tl.store(partial_sums_ptr + pid, block_sum)


@triton.jit
def _final_reduction_kernel(
    partial_sums_ptr,
    output_ptr,
    n_partials,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Second stage: Sum all partial sums into final result.
    Handles case where number of partials fits in one block.
    """
    # For the final reduction, we process all partials in one block
    offsets = tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_partials
    
    # Load partial sums
    partials = tl.load(partial_sums_ptr + offsets, mask=mask, other=0.0)
    
    # Sum all partials
    total_sum = tl.sum(partials, axis=0)
    
    # Store final result
    tl.store(output_ptr, total_sum)


def kernel_function(input_tensor, output_tensor=None):
    """
    Computes the sum of all elements in the input tensor.
    
    This is a fused two-stage reduction:
    - Stage 1: Parallel block-wise partial sums
    - Stage 2: Final reduction of partial sums
    
    All computation happens in Triton kernels with float32 accumulation
    for numerical stability.
    
    Args:
        input_tensor: Input tensor of shape (N,), float32
        output_tensor: Optional output tensor (ignored, for compatibility)
    
    Returns:
        Scalar tensor with the sum (shape ())
    """
    assert input_tensor.is_cuda, "Input must be on CUDA"
    
    n_elements = input_tensor.numel()
    
    # Choose block size - power of 2 for efficiency
    BLOCK_SIZE = 1024
    
    # Calculate number of blocks needed for first stage
    n_blocks = triton.cdiv(n_elements, BLOCK_SIZE)
    
    # Allocate temporary storage for partial sums
    partial_sums = torch.empty(n_blocks, device=input_tensor.device, dtype=torch.float32)
    
    # Stage 1: Compute partial sums per block
    grid_stage1 = (n_blocks,)
    _sum_reduction_kernel[grid_stage1](
        input_tensor,
        partial_sums,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    
    # Stage 2: Final reduction of partial sums
    # Allocate scalar output tensor with shape ()
    result = torch.empty((), device=input_tensor.device, dtype=torch.float32)
    
    # Use a block size that's a power of 2 and >= n_blocks
    FINAL_BLOCK_SIZE = 1024  # Can handle up to 1024 partial sums
    
    if n_blocks <= FINAL_BLOCK_SIZE:
        # Single block can handle all partials
        grid_stage2 = (1,)
        _final_reduction_kernel[grid_stage2](
            partial_sums,
            result,
            n_blocks,
            BLOCK_SIZE=FINAL_BLOCK_SIZE,
        )
    else:
        # Need recursive reduction for very large inputs
        # Keep reducing until we have <= FINAL_BLOCK_SIZE partials
        current_partials = partial_sums
        current_n = n_blocks
        
        while current_n > FINAL_BLOCK_SIZE:
            new_n_blocks = triton.cdiv(current_n, BLOCK_SIZE)
            new_partial_sums = torch.empty(new_n_blocks, device=input_tensor.device, dtype=torch.float32)
            
            grid = (new_n_blocks,)
            _sum_reduction_kernel[grid](
                current_partials,
                new_partial_sums,
                current_n,
                BLOCK_SIZE=BLOCK_SIZE,
            )
            
            current_partials = new_partial_sums
            current_n = new_n_blocks
        
        # Final reduction
        grid_stage2 = (1,)
        _final_reduction_kernel[grid_stage2](
            current_partials,
            result,
            current_n,
            BLOCK_SIZE=FINAL_BLOCK_SIZE,
        )
    
    # Return scalar tensor with shape ()
    return result

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"