submission 545317

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

# Per-block private histogram approach:
# Phase 1: Each block builds its own 256-bin histogram (low contention)
# Phase 2: Sum all per-block histograms

@triton.jit
def histogram_per_block_kernel(
    data_ptr, hist_ptr, n_elements, n_bins: tl.constexpr,
    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).to(tl.int32)
    # Each block writes to its own 256-bin region
    base = pid * n_bins
    tl.atomic_add(hist_ptr + base + vals, 1, mask=mask)

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    n = data.numel()
    BLOCK_SIZE = 4096
    n_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
    # Allocate per-block histograms
    per_block = torch.zeros(n_blocks, 256, device=data.device, dtype=torch.int32)
    histogram_per_block_kernel[(n_blocks,)](
        data, per_block, n, n_bins=256, BLOCK_SIZE=BLOCK_SIZE
    )
    # Reduce: sum across blocks
    output[...] = per_block.sum(dim=0)
    return output