submission 545265

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

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

@triton.jit
def histogram_kernel(
    data_ptr, output_ptr, n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    vals = tl.load(data_ptr + offsets, mask=mask, other=0)
    # Convert to int32 for indexing
    vals = vals.to(tl.int32)
    # Atomic add to global histogram
    for i in range(BLOCK_SIZE):
        if pid * BLOCK_SIZE + i < n_elements:
            bin_idx = tl.load(data_ptr + pid * BLOCK_SIZE + i).to(tl.int32)
            tl.atomic_add(output_ptr + bin_idx, 1)

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    # torch.bincount is already very fast, let's just use it
    output[...] = torch.bincount(data, minlength=256)
    return output