submission 611747

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

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

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

// Per-warp private histograms in shared memory to minimize atomic contention.
// Each warp gets its own 256-bin histogram. At the end, we reduce across warps
// and atomicAdd to global memory.

__global__ void histogram_kernel(
    const uint8_t* __restrict__ data,
    int64_t* __restrict__ output,
    const int64_t n
) {
    const int WARPS_PER_BLOCK = blockDim.x / 32;
    extern __shared__ uint32_t smem[];

    const int warp_id = threadIdx.x / 32;
    const int lane_id = threadIdx.x & 31;

    // Zero shared memory - each warp zeros its own histogram
    uint32_t* my_hist = smem + warp_id * 256;
    #pragma unroll
    for (int i = lane_id; i < 256; i += 32) {
        my_hist[i] = 0;
    }
    __syncthreads();

    // Global index and stride
    const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t grid_stride = (int64_t)blockDim.x * gridDim.x;

    // Process 16 bytes (16 uint8 values) at a time using uint4 loads
    const int64_t n16 = n / 16;
    const uint4* data_vec = (const uint4*)data;

    for (int64_t i = tid; i < n16; i += grid_stride) {
        uint4 vals = __ldg(&data_vec[i]);

        #pragma unroll
        for (int shift = 0; shift < 32; shift += 8) {
            atomicAdd(&my_hist[(vals.x >> shift) & 0xFF], 1u);
            atomicAdd(&my_hist[(vals.y >> shift) & 0xFF], 1u);
            atomicAdd(&my_hist[(vals.z >> shift) & 0xFF], 1u);
            atomicAdd(&my_hist[(vals.w >> shift) & 0xFF], 1u);
        }
    }

    // Handle remaining elements
    {
        int64_t tail_start = n16 * 16;
        for (int64_t i = tail_start + tid; i < n; i += grid_stride) {
            atomicAdd(&my_hist[__ldg(&data[i])], 1u);
        }
    }

    __syncthreads();

    // Reduce across warps and write to global memory
    for (int bin = threadIdx.x; bin < 256; bin += blockDim.x) {
        uint32_t total = 0;
        for (int w = 0; w < WARPS_PER_BLOCK; w++) {
            total += smem[w * 256 + bin];
        }
        if (total > 0) {
            atomicAdd((unsigned long long*)&output[bin], (unsigned long long)total);
        }
    }
}

torch::Tensor histogram_cuda(torch::Tensor data, torch::Tensor output) {
    const int64_t n = data.numel();
    output.zero_();

    const int threads = 256;
    const int max_blocks = 160 * 4;
    const int64_t elements_per_block = threads * 16;
    int blocks = (n + elements_per_block - 1) / elements_per_block;
    if (blocks > max_blocks) blocks = max_blocks;
    if (blocks < 1) blocks = 1;

    const int warps_per_block = threads / 32;
    const int smem_size = warps_per_block * 256 * sizeof(uint32_t);

    histogram_kernel<<<blocks, threads, smem_size>>>(
        data.data_ptr<uint8_t>(),
        output.data_ptr<int64_t>(),
        n
    );

    return output;
}
""";

cpp_source = r"""
torch::Tensor histogram_cuda(torch::Tensor data, torch::Tensor output);
""";

module = load_inline(
    name="histogram_kernel",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["histogram_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


def custom_kernel(data: input_t) -> output_t:
    data_tensor, output = data
    module.histogram_cuda(data_tensor, output)
    return output