submission 490609

# kernel.py
"""
Histogram (uint8 -> int64[256]) implemented in Triton.

Pipeline / fusion note:
- This operation is inherently a global reduction into 256 bins using atomic adds.
- We additionally must *overwrite* the output (set bins to 0) before accumulation.
- These two stages cannot be safely fused into a single kernel because Triton provides no
  global barrier across programs; clearing + accumulating in one launch would race.
So we use two Triton kernels:
  (1) zero-out the 256-bin output
  (2) atomic accumulation over the input
"""

from __future__ import annotations

import torch
import triton
import triton.language as tl


@triton.jit
def _zero_256_i64_kernel(out_ptr,  # *i64
                         BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    # Only one program is expected, but keep pid in case of accidental larger grids.
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < 256
    tl.store(out_ptr + offs, tl.zeros([BLOCK], dtype=tl.int64), mask=mask)


@triton.jit
def _hist_u8_to_i64_kernel(x_ptr,  # *u8
                           out_ptr,  # *i64
                           n_elements: tl.int32,
                           BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < n_elements

    # Load uint8 values; cast to int32 for indexing.
    x = tl.load(x_ptr + offs, mask=mask, other=0).to(tl.int32)  # in [0, 255]

    # Atomic add into global histogram bins (int64 output).
    ones = tl.full([BLOCK_SIZE], 1, dtype=tl.int64)
    tl.atomic_add(out_ptr + x, ones, mask=mask)


def kernel_function(data, out: torch.Tensor | None = None) -> torch.Tensor:
    """
    Compute histogram of a CUDA uint8 1D tensor into 256 bins (int64), matching:
        out[...] = torch.bincount(data, minlength=256)

    Allowed wrapper work: argument checks, allocation, and Triton launches only.
    """
    # Support alternate calling convention kernel_function((data, out))
    if out is None and isinstance(data, (tuple, list)) and len(data) == 2:
        data, out = data

    if not isinstance(data, torch.Tensor):
        raise TypeError(f"data must be a torch.Tensor, got {type(data)}")
    if not data.is_cuda:
        raise ValueError("data must be a CUDA tensor")
    if data.dtype != torch.uint8:
        raise ValueError(f"data must be torch.uint8, got {data.dtype}")
    if data.dim() != 1:
        raise ValueError(f"data must be 1D, got shape={tuple(data.shape)}")
    if not data.is_contiguous():
        data = data.contiguous()

    if out is None:
        out = torch.empty((256,), device=data.device, dtype=torch.int64)
    if not isinstance(out, torch.Tensor):
        raise TypeError(f"out must be a torch.Tensor, got {type(out)}")
    if out.device != data.device:
        raise ValueError("out must be on the same device as data")
    if out.dtype != torch.int64:
        raise ValueError(f"out must be torch.int64, got {out.dtype}")
    if out.shape != (256,):
        raise ValueError(f"out must have shape (256,), got {tuple(out.shape)}")
    if not out.is_contiguous():
        raise ValueError("out must be contiguous")

    # Stage 1: overwrite output bins with zeros (in Triton, not PyTorch).
    _zero_256_i64_kernel[(1,)](out, BLOCK=256, num_warps=1)

    # Stage 2: atomic histogram accumulation.
    n_elements = data.numel()
    if n_elements == 0:
        return out

    BLOCK_SIZE = 1024
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
    _hist_u8_to_i64_kernel[grid](data, out, n_elements, BLOCK_SIZE=BLOCK_SIZE, num_warps=4)

    return out

import inspect

def custom_kernel(input):
    sig = inspect.signature(kernel_function)
    num_params = len(sig.parameters)

    if len(input) == num_params:
        return kernel_function(*input)
    return kernel_function(input)


# Ensure deterministic cuBLAS.
import os
if os.environ.get("CUBLAS_WORKSPACE_CONFIG", "") not in (":4096:8", ":16:8"):
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"