submission 676874

#!POPCORN leaderboard histogram_v2
#!POPCORN gpu B200

import torch
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline

cuda_src = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>


__global__ void __launch_bounds__(256, 4)
histogram_kernel(const uint8_t* __restrict__ data, int64_t* __restrict__ output, int N) {
    extern __shared__ int smem_hist[];
    const int tid = threadIdx.x;

    // Zero shared memory
    smem_hist[tid] = 0;

    // Zero global output in-kernel (block 0 zeros all 256 bins)
    if (blockIdx.x == 0) {
        output[tid] = 0;
    }
    __syncthreads();

    // Grid-stride loop with vectorized loads (16 bytes = 16 uint8 elements per load)
    const int vec_n = N >> 4;
    const uint4* data_vec = reinterpret_cast<const uint4*>(data);

    int idx = blockIdx.x * blockDim.x + tid;
    const int stride = blockDim.x * gridDim.x;

    for (int i = idx; i < vec_n; i += stride) {
        uint4 val = data_vec[i];
        const uint8_t* b = reinterpret_cast<const uint8_t*>(&val);

        #pragma unroll
        for (int j = 0; j < 16; j++) {
            atomicAdd(&smem_hist[b[j]], 1);
        }
    }

    // Handle remaining elements
    int tail_start = vec_n * 16;
    for (int i = tail_start + idx; i < N; i += stride) {
        atomicAdd(&smem_hist[data[i]], 1);
    }

    __syncthreads();

    // Write back to global memory
    if (smem_hist[tid] > 0) {
        atomicAdd(reinterpret_cast<unsigned long long*>(&output[tid]),
                  static_cast<unsigned long long>(smem_hist[tid]));
    }
}

void histogram_inplace(torch::Tensor data, torch::Tensor output) {
    const int N = data.numel();
    int num_blocks = min(256, max(1, (N + 256*16 - 1) / (256*16)));
    histogram_kernel<<<num_blocks, 256, 256*sizeof(int)>>>(
        data.data_ptr<uint8_t>(), output.data_ptr<int64_t>(), N);
}
"""

cpp_src = r"""
void histogram_inplace(torch::Tensor data, torch::Tensor output);
"""

_ext = load_inline(
    name="histogram_block_private_v1",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["histogram_inplace"],
    with_cuda=True,
    extra_cflags=["-O3", "-std=c++17"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    data_tensor, output_tensor = data
    _ext.histogram_inplace(data_tensor, output_tensor)
    return output_tensor