submission 506588

"""
Histogram v2 — Tuned for B200/H100
@MemoryCoalesced
"""

import torch
from torch.utils.cpp_extension import load_inline

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

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

#define BLOCK_THREADS 512
#define NUM_BINS      256

__device__ __forceinline__ void accumulate(unsigned int* smem, uint4 v) {
    atomicAdd(&smem[(v.x      ) & 0xff], 1u);
    atomicAdd(&smem[(v.x >>  8) & 0xff], 1u);
    atomicAdd(&smem[(v.x >> 16) & 0xff], 1u);
    atomicAdd(&smem[(v.x >> 24) & 0xff], 1u);
    atomicAdd(&smem[(v.y      ) & 0xff], 1u);
    atomicAdd(&smem[(v.y >>  8) & 0xff], 1u);
    atomicAdd(&smem[(v.y >> 16) & 0xff], 1u);
    atomicAdd(&smem[(v.y >> 24) & 0xff], 1u);
    atomicAdd(&smem[(v.z      ) & 0xff], 1u);
    atomicAdd(&smem[(v.z >>  8) & 0xff], 1u);
    atomicAdd(&smem[(v.z >> 16) & 0xff], 1u);
    atomicAdd(&smem[(v.z >> 24) & 0xff], 1u);
    atomicAdd(&smem[(v.w      ) & 0xff], 1u);
    atomicAdd(&smem[(v.w >>  8) & 0xff], 1u);
    atomicAdd(&smem[(v.w >> 16) & 0xff], 1u);
    atomicAdd(&smem[(v.w >> 24) & 0xff], 1u);
}

__global__ void __launch_bounds__(BLOCK_THREADS, 2)
histogram_kernel(
    const uint4*        __restrict__ data16,
    unsigned long long* __restrict__ output,
    volatile int*       __restrict__ flag,
    int n16,
    int n,
    int num_blocks)
{
    // Block 0 zeros output (256 entries, first 256 threads do one store each)
    if (blockIdx.x == 0) {
        if (threadIdx.x < NUM_BINS)
            output[threadIdx.x] = 0ULL;
        __threadfence();
        if (threadIdx.x == 0)
            *flag = 1;
    } else {
        if (threadIdx.x == 0)
            while (*flag == 0) {}
        __syncthreads();
    }

    __shared__ unsigned int smem[NUM_BINS];
    // 512 threads zero 256 bins: each thread zeros one entry (first 256 do it)
    if (threadIdx.x < NUM_BINS)
        smem[threadIdx.x] = 0u;
    __syncthreads();

    const int tid    = blockIdx.x * BLOCK_THREADS + threadIdx.x;
    const int stride = num_blocks * BLOCK_THREADS;

    // 2x unroll for better MLP
    int i = tid;
    for (; i + stride < n16; i += stride * 2) {
        accumulate(smem, __ldg(data16 + i));
        accumulate(smem, __ldg(data16 + i + stride));
    }
    if (i < n16)
        accumulate(smem, __ldg(data16 + i));

    // Tail bytes
    const uint8_t* data1 = (const uint8_t*)data16;
    for (int j = n16 * 16 + tid; j < n; j += stride)
        atomicAdd(&smem[__ldg(data1 + j)], 1u);

    __syncthreads();

    // Merge: 256 threads each do one atomicAdd
    if (threadIdx.x < NUM_BINS)
        atomicAdd(&output[threadIdx.x], (unsigned long long)smem[threadIdx.x]);

    if (blockIdx.x == 0 && threadIdx.x == 0)
        *flag = 0;
}

static int NUM_BLOCKS_CACHED = 0;

void histogram_cuda(
    torch::Tensor data,
    torch::Tensor output,
    torch::Tensor flag)
{
    if (NUM_BLOCKS_CACHED == 0) {
        int dev, sm_count;
        cudaGetDevice(&dev);
        cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
        NUM_BLOCKS_CACHED = sm_count;  // 1 block/SM, 512 threads = max work/thread
    }
    const int n   = data.numel();
    const int n16 = n / 16;
    histogram_kernel<<<NUM_BLOCKS_CACHED, BLOCK_THREADS>>>(
        (const uint4*)data.data_ptr<uint8_t>(),
        (unsigned long long*)output.data_ptr<int64_t>(),
        (volatile int*)flag.data_ptr<int>(),
        n16, n, NUM_BLOCKS_CACHED
    );
}
"""

module = load_inline(
    name='histogram_v27',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['histogram_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math', '-std=c++17'],
)

_flag = torch.zeros(1, device='cuda', dtype=torch.int32)

def custom_kernel(data: tuple) -> torch.Tensor:
    inp, out = data
    module.histogram_cuda(inp, out, _flag)
    return out