submission 512857

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

cuda_src = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>
#include <ATen/ATen.h>

// Per-warp private histograms, larger blocks to reduce global atomic pressure
// Block of 512 threads = 16 warps, each warp gets its own 256-bin histogram
// Total smem: 16 * 256 * 4 = 16KB per block
__global__ void histogram_warp_private_kernel(
    const uint8_t* __restrict__ data,
    int64_t* __restrict__ output,
    int size
) {
    const int WARP_SIZE = 32;
    const int WARPS_PER_BLOCK = blockDim.x / WARP_SIZE;

    extern __shared__ uint32_t smem_hist[];  // [WARPS_PER_BLOCK][256]

    int tid = threadIdx.x;
    int warp_id = tid / WARP_SIZE;
    int lane_id = tid % WARP_SIZE;

    uint32_t* my_hist = smem_hist + warp_id * 256;

    // Initialize
    for (int i = lane_id; i < 256; i += WARP_SIZE) {
        my_hist[i] = 0;
    }
    __syncthreads();

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

    // Process 16 bytes at a time using uint4 loads
    const uint4* data128 = reinterpret_cast<const uint4*>(data);
    int vec_size = size / 16;

    for (int i = global_tid; i < vec_size; i += stride) {
        uint4 val = __ldg(&data128[i]);
        uint32_t b0, b1, b2, b3;

        b0 = val.x & 0xFF; b1 = (val.x >> 8) & 0xFF; b2 = (val.x >> 16) & 0xFF; b3 = val.x >> 24;
        atomicAdd(&my_hist[b0], 1u);
        atomicAdd(&my_hist[b1], 1u);
        atomicAdd(&my_hist[b2], 1u);
        atomicAdd(&my_hist[b3], 1u);

        b0 = val.y & 0xFF; b1 = (val.y >> 8) & 0xFF; b2 = (val.y >> 16) & 0xFF; b3 = val.y >> 24;
        atomicAdd(&my_hist[b0], 1u);
        atomicAdd(&my_hist[b1], 1u);
        atomicAdd(&my_hist[b2], 1u);
        atomicAdd(&my_hist[b3], 1u);

        b0 = val.z & 0xFF; b1 = (val.z >> 8) & 0xFF; b2 = (val.z >> 16) & 0xFF; b3 = val.z >> 24;
        atomicAdd(&my_hist[b0], 1u);
        atomicAdd(&my_hist[b1], 1u);
        atomicAdd(&my_hist[b2], 1u);
        atomicAdd(&my_hist[b3], 1u);

        b0 = val.w & 0xFF; b1 = (val.w >> 8) & 0xFF; b2 = (val.w >> 16) & 0xFF; b3 = val.w >> 24;
        atomicAdd(&my_hist[b0], 1u);
        atomicAdd(&my_hist[b1], 1u);
        atomicAdd(&my_hist[b2], 1u);
        atomicAdd(&my_hist[b3], 1u);
    }

    __syncthreads();

    // Reduce all warp histograms and write to global output
    for (int i = tid; i < 256; i += blockDim.x) {
        uint32_t sum = 0;
        for (int w = 0; w < WARPS_PER_BLOCK; w++) {
            sum += smem_hist[w * 256 + i];
        }
        if (sum > 0) {
            atomicAdd(reinterpret_cast<unsigned long long*>(&output[i]),
                      (unsigned long long)sum);
        }
    }
}

void run_histogram(at::Tensor data, at::Tensor output) {
    int size = (int)data.size(0);
    const uint8_t* data_ptr = data.data_ptr<uint8_t>();
    int64_t* output_ptr = output.data_ptr<int64_t>();
    int device = data.device().index();
    int num_sms;
    cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device);

    // 512 threads = 16 warps, 16KB smem per block
    int block_size = 512;
    int warps_per_block = block_size / 32;
    size_t smem_size = warps_per_block * 256 * sizeof(uint32_t);

    // Use 2 blocks per SM
    int num_blocks = num_sms * 2;

    histogram_warp_private_kernel<<<num_blocks, block_size, smem_size>>>(
        data_ptr, output_ptr, size);
}
"""

cpp_src = """
#include <ATen/ATen.h>
void run_histogram(at::Tensor data, at::Tensor output);
"""

_module = None

def get_module():
    global _module
    if _module is None:
        _module = load_inline(
            name='histogram_warp_private_v5',
            cpp_sources=cpp_src,
            cuda_sources=cuda_src,
            functions=['run_histogram'],
            verbose=False,
            extra_cuda_cflags=['-O3', '--use_fast_math'],
        )
    return _module

def custom_kernel(data: input_t) -> output_t:
    inp, output = data
    output.zero_()
    get_module().run_histogram(inp, output)
    return output