submission 613124

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>

// Fused kernel: zeros output atomically (via exchange), then does histogram.
// The first thread to reach each output bin sets it to 0 using atomicExch.
// Other threads just atomicAdd.
// We use a flag to synchronize: first block zeros, then signals others.
// But we can't do grid sync without cooperative launch.
// Instead: use atomicExch(0) + atomicAdd approach.
// Since the output starts with garbage, the first atomicAdd will be wrong.
// We need to ensure output is zeroed before any atomicAdd.
//
// Simpler: just use cudaMemset in C++ before kernel launch.
// The overhead is that cudaMemsetAsync issues a DMA operation.
// Alternative: output.zero_() in Python.
//
// Actually, let's try torch.zeros in Python instead:

__global__ __launch_bounds__(512)
void histogram_custom(
    const uint8_t* __restrict__ data,
    int64_t* __restrict__ output,
    const int64_t n
) {
    constexpr int WARPS = 16;
    __shared__ uint32_t smem[WARPS * 256];

    const int warp_id = threadIdx.x >> 5;

    #pragma unroll
    for (int i = threadIdx.x; i < WARPS * 256; i += 512) {
        smem[i] = 0;
    }
    __syncthreads();

    uint32_t* my_hist = smem + warp_id * 256;

    const int64_t tid = blockIdx.x * 512 + threadIdx.x;
    const int64_t grid_stride = 512LL * gridDim.x;

    const int64_t n16 = n >> 4;
    const uint4* data_vec = (const uint4*)data;

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

        atomicAdd(&my_hist[(vals.x      ) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.x >>  8) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.x >> 16) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.x >> 24)       ], 1u);
        atomicAdd(&my_hist[(vals.y      ) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.y >>  8) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.y >> 16) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.y >> 24)       ], 1u);
        atomicAdd(&my_hist[(vals.z      ) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.z >>  8) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.z >> 16) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.z >> 24)       ], 1u);
        atomicAdd(&my_hist[(vals.w      ) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.w >>  8) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.w >> 16) & 0xFF], 1u);
        atomicAdd(&my_hist[(vals.w >> 24)       ], 1u);
    }

    {
        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();

    if (threadIdx.x < 256) {
        uint32_t total = 0;
        #pragma unroll
        for (int w = 0; w < WARPS; w++) {
            total += smem[w * 256 + threadIdx.x];
        }
        if (total > 0) {
            atomicAdd((unsigned long long*)&output[threadIdx.x], (unsigned long long)total);
        }
    }
}

torch::Tensor histogram_cuda(torch::Tensor data, torch::Tensor output) {
    const int64_t n = data.numel();
    cudaMemsetAsync(output.data_ptr<int64_t>(), 0, 256 * sizeof(int64_t));

    // Auto-detect SM count for optimal block config
    static int sm_count = -1;
    if (sm_count < 0) {
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, 0);
        sm_count = prop.multiProcessorCount;
    }

    // 2 blocks per SM is optimal for 512 threads/block
    int blocks = sm_count * 2;

    if (n < 512 * 16 * blocks) {
        blocks = (n + 512 * 16 - 1) / (512 * 16);
        if (blocks < 1) blocks = 1;
    }

    histogram_custom<<<blocks, 512>>>(
        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_512_auto",
    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