submission 150472

import math
from typing import Tuple

import torch
import triton
import triton.language as tl

from task import input_t, output_t


@triton.jit
def _bitonic_step(x_ptr, n_elements, j, k, BLOCK: tl.constexpr):
    """
    One stage of the bitonic sorting network.

    Args:
        x_ptr: Pointer to the data being sorted (float tensor).
        n_elements: Number of valid elements in the array.
        j: Inner distance within the current bitonic merge step.
        k: Current size of the bitonic sequence.
        BLOCK: Number of elements handled per program instance.
    """
    pid = tl.program_id(0)
    base = pid * BLOCK  # start offset
    idx = base + tl.arange(0, BLOCK)
    in_bounds = idx < n_elements

    partner = idx ^ j
    partner_in_bounds = partner < n_elements

    # Load values; pad out-of-range lanes with +inf so they never win a swap.
    val = tl.load(x_ptr + idx, mask=in_bounds, other=0.0)
    partner_val = tl.load(x_ptr + partner, mask=partner_in_bounds, other=0.0)
    # Mask out-of-range lanes with +inf so they never win a swap.
    inf_const = float("inf")
    val = tl.where(in_bounds, val, inf_const)
    partner_val = tl.where(partner_in_bounds, partner_val, inf_const)

    ascending_region = (idx & k) == 0
    should_swap = tl.where(
        ascending_region, val > partner_val, val < partner_val
    )

    # Only let the lower index of each pair write both elements to avoid races.
    pair_mask = in_bounds & partner_in_bounds & (partner > idx)
    do_swap = pair_mask & should_swap

    new_val = tl.where(do_swap, partner_val, val)
    new_partner_val = tl.where(do_swap, val, partner_val)

    tl.store(x_ptr + idx, new_val, mask=pair_mask)
    tl.store(x_ptr + partner, new_partner_val, mask=pair_mask)


def custom_kernel(
    input: input_t,
    block_size: int = 256,
    num_warps: int = 8,
) -> output_t:
    """
    Sort a 1D CUDA tensor in ascending order using a bitonic sorting network.

    Args:
        input: Tuple of (data, output_buffer) where output_buffer has the same
            shape and dtype as data.
        block_size: Number of elements processed per Triton program.
        num_warps: Triton launch hint for how many hardware warps to use.
    """
    data, output = input
    if not data.is_cuda:
        raise ValueError("custom_kernel requires a CUDA tensor")
    if data.dim() != 1:
        raise ValueError("custom_kernel currently supports 1D tensors only")
    if not data.is_floating_point():
        raise ValueError("custom_kernel expects floating-point data")
    if output.numel() != data.numel():
        raise ValueError("output buffer must match input shape")

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

    # Pad to next power of two into a temporary buffer to keep the bitonic
    # network correct when the length is not a power of two.
    padded = 1 << (n - 1).bit_length()
    tmp = torch.full((padded,), float("inf"), device=data.device, dtype=data.dtype)
    tmp[:n].copy_(data)

    grid = (triton.cdiv(padded, block_size),)
    k = 2
    while k <= padded:
        j = k >> 1
        while j > 0:
            _bitonic_step[grid](
                tmp, padded, j, k, BLOCK=block_size, num_warps=num_warps
            )
            j >>= 1
        k <<= 1

    output.copy_(tmp[:n])
    return output