submission 150025

#!POPCORN leaderboard sort_v2

from task import input_t, output_t
import torch
import triton
import triton.language as tl


@triton.jit
def simple_sort_kernel(
    input_ptr,
    output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    """Simple sort kernel using tl.sort."""
    offsets = tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    data = tl.load(input_ptr + offsets, mask=mask, other=float('inf'))

    # Use Triton's built-in sort
    sorted_data = tl.sort(data)

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


def triton_custom_kernel(data: input_t) -> output_t:
    inp, output = data

    n_elements = inp.numel()

    # Round up to next power of 2 for sort
    block_size = 1
    while block_size < n_elements:
        block_size *= 2

    # Cap block size to reasonable maximum
    block_size = min(block_size, 1024)

    # Launch kernel
    simple_sort_kernel[(1,)](
        inp,
        output,
        n_elements,
        BLOCK_SIZE=block_size,
    )

    return output

def naive_custom_kernel(data):
    inp, out = data

    out.copy_(torch.sort(inp)[0])
    return out


custom_kernel = naive_custom_kernel