submission 506500

"""
Histogram v2 — Direct global atomicAdd, block0 zeros output at start
@MemoryCoalesced

v18 bottleneck: scratch merge writes+reads 131KB extra per call.
For 1.3MB input that's 10% overhead → fixed 1.5µs cost.

Fix: eliminate scratch entirely.
  - Block 0 zeros output[0..255] immediately (256 stores, ~0.5µs)
  - Block 0 sets ready flag
  - All other blocks spin on ready flag at START (near-zero wait)
  - All blocks accumulate into smem
  - All blocks atomicAdd smem → output (2KB, 256 bins × num_blocks atomics)

Spin at START vs END:
  - END spin: wait for slowest block to finish processing its data chunk → 10-15µs
  - START spin: wait for block 0 to do 256 stores → ~0.1µs
  ∴ START spin is ~100x cheaper

Global atomicAdd contention at merge:
  - 128 blocks × 256 bins = 32K atomics, but spread across 256 independent addresses
  - 128 serialized atomics per bin × ~5ns = ~640ns = 0.6µs
  - Still cheaper than 131KB scratch write+read (~1.5µs)
"""

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 256
#define NUM_BINS      256
#define NUM_BLOCKS    116

__global__ void __launch_bounds__(BLOCK_THREADS, 4)
histogram_kernel(
    const uint4*        __restrict__ data16,
    unsigned long long* __restrict__ output,
    volatile int*       __restrict__ flag,  // [1], reset to 0 by block 0 at end
    int n16,
    int n)
{
    // Block 0: zero output immediately, then signal ready
    // All others: spin until ready, then start work
    if (blockIdx.x == 0) {
        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];
    smem[threadIdx.x] = 0u;
    __syncthreads();

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

    for (int i = tid; i < n16; i += stride) {
        uint4 v = __ldg(data16 + i);
        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);
    }
    const uint8_t* data1 = (const uint8_t*)data16;
    for (int i = n16 * 16 + tid; i < n; i += stride)
        atomicAdd(&smem[__ldg(data1 + i)], 1u);

    __syncthreads();

    // Merge smem → global (output already zeroed)
    atomicAdd(&output[threadIdx.x], (unsigned long long)smem[threadIdx.x]);

    // Block 0 resets flag for next call (it's guaranteed last to finish
    // since it did the most work: zeroing + full data processing)
    if (blockIdx.x == 0 && threadIdx.x == 0)
        *flag = 0;
}

void histogram_cuda(
    torch::Tensor data,
    torch::Tensor output,
    torch::Tensor flag)
{
    const int n   = data.numel();
    const int n16 = n / 16;
    histogram_kernel<<<NUM_BLOCKS, 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
    );
}
"""

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

# Single int flag, reset to 0 by kernel after each call
_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