submission 102250

import os
# Unset CUDA_VISIBLE_DEVICES before importing torch/triton
# This works around SLURM setting it to a value that prevents PyTorch from seeing GPUs
if 'CUDA_VISIBLE_DEVICES' in os.environ:
    del os.environ['CUDA_VISIBLE_DEVICES']

import triton.testing
import math
import logging

import triton
import triton.language as tl
import torch

torch.cuda.empty_cache()
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device count: {torch.cuda.device_count()}")
import torch.nn.functional as F

@triton.jit
def naive_prefixsum_kernel(
    input_ptr,
    output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

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

    # Single-thread accumulator
    # Only lane 0 computes the prefix sum sequentially.
    acc = tl.zeros((BLOCK_SIZE,), dtype=x.dtype)

    # get the lane id within the program
    lane_id = tl.arange(0, BLOCK_SIZE)

    # Run prefix sum sequentially on lane 0
    if tl.program_id(axis=1) == 0:  # always true (1D grid) but needed for Triton scoping
        # Only lane 0 executes the loop
        if lane_id == 0:
            running = 0.0
            for i in range(BLOCK_SIZE):
                running += x[i]
                acc[i] = running

    # Broadcast results: each lane loads its own index from acc
    out = acc[lane_id]

    tl.store(output_ptr + offsets, out, mask=mask)




# --------- kernel ----------------
@triton.jit
def prefixsum_kernel(
    input_ptr,      # pointer to input array
    output_ptr,     # pointer to output array
    n_elements,     # total number of elements
    BLOCK_SIZE: tl.constexpr,  # block size for processing
):
    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, with padding for elements out of bounds
    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)

    # Perform parallel prefix sum within the block using tl.cumsum.
    # cumsum is the built-in prefix sum operation in Triton.
    x = tl.cumsum(x, axis=0)

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

@triton.jit
def add_block_sums_kernel(
    output_ptr,
    block_sums_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """Add the cumulative sum from previous blocks to each element."""
    pid = tl.program_id(axis=0)
    
    if pid == 0:
        return  # First block doesn't need adjustment
    
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    # Load the cumulative sum from all previous blocks
    prev_sum = tl.load(block_sums_ptr + pid - 1)
    
    # Load current values, add previous sum, and store back
    x = tl.load(output_ptr + offsets, mask=mask)
    x = x + prev_sum
    tl.store(output_ptr + offsets, x, mask=mask)


# --------- wrapper function ----------------
def prefixsum(data):
    """
    Wrapper function to launch the prefix sum kernel.
    Args:
        data: Tuple of (input_tensor, output_tensor)
    Returns:
        output_tensor: Tensor containing the inclusive prefix sum
    """
    input_tensor, output_tensor = data
    n_elements = input_tensor.numel()
    
    # Choose block size
    BLOCK_SIZE = 4096
    
    # Calculate grid size
    num_blocks = triton.cdiv(n_elements, BLOCK_SIZE)
    grid = (num_blocks,)
    
    # Pass 1: Compute prefix sum within each block
    prefixsum_kernel[grid](
        input_tensor,
        output_tensor,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
    )

    # If we have multiple blocks, we need to propagate sums across blocks
    if num_blocks > 1:
        # Extract the last element of each block (the block sums)
        block_sums = torch.zeros(num_blocks, dtype=input_tensor.dtype, device=input_tensor.device)
        for i in range(num_blocks):
            last_idx = min((i + 1) * BLOCK_SIZE - 1, n_elements - 1)
            block_sums[i] = output_tensor[last_idx]

        # Compute prefix sum of block sums
        block_prefix = torch.cumsum(block_sums, dim=0)

        # Pass 2: Add accumulated sums from previous blocks
        add_block_sums_kernel[grid](
            output_tensor,
            block_prefix,
            n_elements,
            BLOCK_SIZE=BLOCK_SIZE,
        )

    return output_tensor


# Alias for submission system
custom_kernel = prefixsum