submission 757228

# EVOLVE-BLOCK-START

import torch
import triton
import triton.language as tl
from typing import TypeVar

input_t = TypeVar("input_t", bound=tuple[torch.Tensor, torch.Tensor])
output_t = TypeVar("output_t", bound=torch.Tensor)


@triton.jit
def histogram_kernel(
    data_ptr,
    tmp_hist_ptr,
    N,
    BLOCK_SIZE: tl.constexpr,
    CHUNK: tl.constexpr,
):
    """
    Build a per‑block 256‑bin histogram in a temporary global buffer.
    Each block writes to its own slice of `tmp_hist` (256 int32 entries).
    """
    pid = tl.program_id(0)

    # Pointer to the slice belonging to this block (256 int32 bins)
    tmp_ptr = tmp_hist_ptr + pid * 256

    # Thread index inside the block
    thread_idx = tl.arange(0, BLOCK_SIZE)

    # Base offset for this thread: each thread processes CHUNK elements spaced by BLOCK_SIZE
    base = pid * BLOCK_SIZE * CHUNK + thread_idx

    # Loop over the CHUNK elements assigned to this thread
    for i in range(CHUNK):
        idx = base + i * BLOCK_SIZE
        mask = idx < N

        # Load a uint8 value (0‑255) and cast to int32 for indexing
        val = tl.load(data_ptr + idx, mask=mask, other=0).to(tl.int32)

        # Increment the per‑block histogram atomically
        tl.atomic_add(tmp_ptr + val, 1, mask=mask)


@triton.jit
def reduce_histogram_kernel(
    tmp_hist_ptr,
    out_ptr,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Reduce all per‑block histograms into the final output.
    One thread per bin (0‑255) adds its bin count from this block to `out`.
    """
    pid = tl.program_id(0)

    # One thread per histogram bin
    bin_idx = tl.arange(0, BLOCK_SIZE)
    mask = bin_idx < 256

    # Load the count for this bin from this block's temporary histogram (int32 → int64)
    cnt = tl.load(tmp_hist_ptr + pid * 256 + bin_idx, mask=mask).to(tl.int64)

    # Atomically add to the final output (int64)
    tl.atomic_add(out_ptr + bin_idx, cnt, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    """
    Compute a 256‑bin histogram of a uint8 1‑D tensor.
    This implementation builds per‑block sub‑histograms in a temporary buffer
    to dramatically reduce atomic contention, then reduces them into the final
    output.
    """
    data_tensor, output_tensor = data

    # Ensure inputs are contiguous and on the same device
    data_tensor = data_tensor.contiguous()
    output_tensor = output_tensor.contiguous()
    device = data_tensor.device

    # Zero the output tensor (required for atomic accumulation)
    output_tensor.zero_()

    N = data_tensor.numel()

    # Kernel launch configuration
    BLOCK_SIZE = 256          # one thread per histogram bin (256 threads)
    CHUNK = 16                # elements processed per thread

    # Number of blocks needed to cover the input
    num_blocks = max(1, triton.cdiv(N, BLOCK_SIZE * CHUNK))
    grid = (num_blocks,)

    # Temporary buffer for per‑block histograms (int32)
    tmp_hist = torch.zeros((num_blocks, 256), dtype=torch.int32, device=device)

    # -----------------------------------------------------------------
    # Phase 1: build per‑block histograms
    histogram_kernel[grid](
        data_tensor,
        tmp_hist,
        N,
        BLOCK_SIZE=BLOCK_SIZE,
        CHUNK=CHUNK,
        num_warps=8,
        num_stages=4,
    )
    # -----------------------------------------------------------------
    # Phase 2: reduce per‑block histograms into the final output
    reduce_histogram_kernel[grid](
        tmp_hist,
        output_tensor,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=8,
        num_stages=2,
    )
    # -----------------------------------------------------------------

    return output_tensor
# EVOLVE-BLOCK-END