submission 489479

import triton
import triton.language as tl
import torch


@triton.jit
def _histogram_partial_kernel(
    data_ptr,
    partial_ptr,
    n_elements,
    stride_pr0,
    stride_pr1,
    BLOCK_SIZE: tl.constexpr,
    BINS: tl.constexpr,
):
    """
    Per-tile histogram builder (no atomics).
    Each program builds a histogram of its BLOCK_SIZE-sized tile and stores it to partial_ptr[pid, :].

    Args:
        data_ptr: pointer to uint8 input [n_elements]
        partial_ptr: pointer to int32 partial histograms [num_tiles, BINS]
        n_elements: size of input
        stride_pr0: stride of partial along the row (num_tiles) dimension
        stride_pr1: stride of partial along the bin dimension
        BLOCK_SIZE: number of elements per program tile
        BINS: number of histogram bins (256)
    """
    tl.static_assert(BINS == 256)
    pid = tl.program_id(axis=0)

    # Offsets for this tile
    block_start = pid * BLOCK_SIZE
    offs = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offs < n_elements

    # Load tile data (pad with 0 for masked lanes; mask will be applied in comparisons)
    vals_u8 = tl.load(data_ptr + offs, mask=mask, other=0)

    # Local histogram accumulator (int32 is sufficient per tile)
    idx_bins = tl.arange(0, BINS)
    local_hist = tl.zeros((BINS,), dtype=tl.int32)

    # For each possible bin value, count occurrences in this tile
    # We explicitly AND with 'mask' so padded zeros do not contribute to bin 0.
    for b in range(0, BINS):
        eq = vals_u8 == b
        count_b = tl.sum((eq & mask).to(tl.int32), axis=0)
        # Accumulate into the specific bin via vectorized mask to avoid scatter conflicts
        local_hist += (idx_bins == b).to(tl.int32) * count_b

    # Store this tile's histogram into the partial buffer
    row_ptrs = partial_ptr + pid * stride_pr0 + idx_bins * stride_pr1
    tl.store(row_ptrs, local_hist)


@triton.jit
def _histogram_reduce_kernel(
    partial_ptr,
    out_ptr,
    num_rows,
    stride_pr0,
    stride_pr1,
    stride_out,
    BLOCK_BINS: tl.constexpr,
    BINS: tl.constexpr,
):
    """
    Reduction kernel that sums partial histograms across rows into final int64 output.

    Args:
        partial_ptr: pointer to int32 partial histograms [num_rows, BINS]
        out_ptr: pointer to int64 output histogram [BINS]
        num_rows: number of partial rows (tiles)
        stride_pr0: stride of partial along the row dimension
        stride_pr1: stride of partial along the bin dimension
        stride_out: stride of output along the bin dimension (typically 1)
        BLOCK_BINS: number of bins handled per program
        BINS: total bins (256)
    """
    tl.static_assert(BINS == 256)

    pid = tl.program_id(axis=0)
    bin_start = pid * BLOCK_BINS
    offs_bins = bin_start + tl.arange(0, BLOCK_BINS)
    mask_bins = offs_bins < BINS

    acc = tl.zeros((BLOCK_BINS,), dtype=tl.int64)

    # Sum across all rows (tiles)
    for r in tl.range(0, num_rows, 1):
        ptrs = partial_ptr + r * stride_pr0 + offs_bins * stride_pr1
        vals = tl.load(ptrs, mask=mask_bins, other=0).to(tl.int64)
        acc += vals

    out_ptrs = out_ptr + offs_bins * stride_out
    tl.store(out_ptrs, acc, mask=mask_bins)


def kernel_function(data: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
    """
    Triton-based histogram kernel wrapper.

    Fused stages attempted:
    - We implement the histogram in two Triton kernels:
      1) Per-tile counting in SRAM without atomics
      2) Cross-tile reduction into the final int64 histogram

    Rationale for not fully fusing into a single kernel:
    - While a single-pass design using atomic_add to global int64 bins would conceptually be simpler,
      int64 atomics are not consistently supported across environments/architectures. To keep the
      implementation portable and deterministic under the provided test constraints, we perform a
      two-stage reduction without relying on atomics. This avoids heavy contention for the "evil"
      bin values and stays within the runtime restrictions (no PyTorch compute in execution path).

    Args:
        data: 1D CUDA tensor with dtype uint8, values in [0, 255]
        out: Optional 1D CUDA tensor with dtype int64 and shape (256,).
             If provided, it will be written to and returned. If None, a new output is allocated.

    Returns:
        1D CUDA tensor with dtype int64 and shape (256,) containing counts per bin.
    """
    # Basic checks and setup (allowed in wrapper)
    assert isinstance(data, torch.Tensor), "data must be a torch.Tensor"
    assert data.dtype == torch.uint8, "data must be torch.uint8"
    assert data.is_cuda, "data must be on CUDA device"
    data = data.contiguous()
    n = data.numel()

    if out is None:
        out = torch.empty(256, device=data.device, dtype=torch.int64)
    else:
        assert isinstance(out, torch.Tensor), "out must be a torch.Tensor"
        assert out.is_cuda, "out must be on CUDA device"
        assert out.dtype == torch.int64, "out must be torch.int64"
        assert out.numel() == 256, "out must have 256 elements"
        out = out.contiguous()

    # Kernel configuration
    BINS = 256
    BLOCK_SIZE = 4096  # tile size per program; power-of-two per guidelines
    BLOCK_BINS = 128   # bins handled per program in the reduction

    # Number of tiles/programs in partial pass
    num_tiles = triton.cdiv(n, BLOCK_SIZE)
    # Partial buffer to hold per-tile histograms in int32
    partial = torch.empty((num_tiles, BINS), device=data.device, dtype=torch.int32).contiguous()

    # Launch partial counting kernel
    grid_partial = (num_tiles,)
    _histogram_partial_kernel[grid_partial](
        data, partial,
        n,
        partial.stride(0), partial.stride(1),
        BLOCK_SIZE=BLOCK_SIZE,
        BINS=BINS,
    )

    # Launch reduction kernel over bins
    grid_reduce = (triton.cdiv(BINS, BLOCK_BINS),)
    _histogram_reduce_kernel[grid_reduce](
        partial, out,
        num_tiles,
        partial.stride(0), partial.stride(1),
        out.stride(0),
        BLOCK_BINS=BLOCK_BINS,
        BINS=BINS,
    )

    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"