submission 34557

#!POPCORN leaderboard histogram_v2

from typing import TypedDict, TypeVar
import torch

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


class TestSpec(TypedDict):
    size: int
    seed: int
    contention: int


from torch.utils.cpp_extension import load_inline

cuda_src = r"""

__global__ void zero_kernel(int64_t* hist, int bins) {
    int idx = threadIdx.x;
    if (idx < bins) {
        hist[idx] = 0;
    }
}

__global__ void histogram_kernel(const uint8_t* data, int64_t* global_hist, int64_t n) {
    __shared__ unsigned int local_hist[256];

    int tid = threadIdx.x;
    // init shared histogram
    for (int i = tid; i < 256; i += blockDim.x)
        local_hist[i] = 0;
    __syncthreads();

    // stride loop over input
    for (size_t i = blockIdx.x * blockDim.x + tid; i < n; i += gridDim.x * blockDim.x) {
        atomicAdd(&local_hist[data[i]], 1);
    }
    __syncthreads();

    // merge into global histogram
    for (int i = tid; i < 256; i += blockDim.x)
        atomicAdd((unsigned long long int*)&global_hist[i], (unsigned long long)local_hist[i]);
}

void histogram(torch::Tensor x, torch::Tensor hist) {
    auto n = x.numel();
    int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    zero_kernel<<<1, 256>>>(hist.data_ptr<int64_t>(), 256);
    histogram_kernel<<<blocks, threads>>>(x.data_ptr<uint8_t>(), hist.data_ptr<int64_t>(), n);
    cudaDeviceSynchronize();
}

"""

cpp_src = "void histogram(torch::Tensor x, torch::Tensor hist);"

module = load_inline(
    name="histogram_cuda_module",
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=["histogram"],
    with_cuda=True,
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


def inline_kernel(data: input_t) -> output_t:
    x, output = data
    module.histogram(x, output)
    return output

import torch
def baseline_kernel(data: input_t) -> output_t:
    """
    Reference implementation of histogram using PyTorch.
    Args:
        data: tensor of shape (size,)
    Returns:
        Tensor containing bin counts
    """
    data, output = data
    # Count values in each bin
    output[...] = torch.bincount(data, minlength=256)
    return output


custom_kernel = inline_kernel


def test():
    # Sanity check
    import torch
    N = 32
    gen = torch.Generator(device='cuda')
    gen.manual_seed(42)

    # Generate integer values between 0 and 256
    a = torch.randint(0, 256, (N,), device='cuda', dtype=torch.uint8, generator=gen)
    output = torch.empty(256, device='cuda', dtype=torch.int64).contiguous()
    module.histogram(a, output)   # launches the kernel on the current CUDA stream

    # correctness check
    # torch.testing.assert_close(out, a + b)
    print("Input:      ", a)
    print("Histogram:  ", output)

# test()