submission 545313

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
from task import input_t, output_t

# Sort with stable=False for potential speedup
# The reference uses torch.sort(data)[0] which defaults to stable=False
# Pre-allocate both values and indices buffers

_vals_cache = {}
_idx_cache = {}

def _warmup():
    for size in [1024, 4096, 16384, 65536, 262144]:
        x = torch.randn(size, device='cuda', dtype=torch.float32)
        torch.sort(x, stable=False)
    torch.cuda.synchronize()

_warmup()

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    n = data.numel()
    if n not in _idx_cache or _idx_cache[n].device != data.device:
        _idx_cache[n] = torch.empty(n, device=data.device, dtype=torch.int64)
    torch.sort(data, stable=False, out=(output, _idx_cache[n]))
    return output