submission 68322

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

# Tune occupancy vs. atomic contention. 

# Additionally:
# -- factor in SMEM throughput
# -- factor in LSU throughput

cpp_src = """
#include <torch/extension.h>

void histogram_kernel(torch::Tensor input, torch::Tensor output);
"""

cuda_src = """
#include <torch/extension.h>
#include <cooperative_groups.h>

namespace cg = cooperative_groups;

constexpr int threadSize = 1;
constexpr int threadsPerBlock = 256;

constexpr int blockSize = threadsPerBlock * threadSize;

__global__ void histogram_block(uint8_t* input, int64_t* output, int n) {
   __shared__ unsigned long long local[256];

    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    cg::thread_block block = cg::this_thread_block();

    // 0. Zero out local histogram.
    if (block.thread_rank() < 256) {
        local[block.thread_rank()] = 0;
    }
    block.sync();

    // 1. Compute local histogram.
    if (idx < n) {
        uint8_t item = input[idx];

        atomicAdd(&local[item], 1);
    } 

    // 2. Synchronize.
    block.sync(); 

    // 3. Writeback to global.
    if (block.thread_rank() < 256) {
        atomicAdd(reinterpret_cast<unsigned long long*>(&output[block.thread_rank()]), local[block.thread_rank()]);
    }
}

void histogram_kernel(torch::Tensor input, torch::Tensor output) {
    int n = input.numel();

    cudaMemset(output.data_ptr<int64_t>(), 0, sizeof(int64_t) * 256);
    
    histogram_block<<<(n + blockSize - 1) / blockSize, threadsPerBlock>>>(
        input.data_ptr<uint8_t>(), 
        output.data_ptr<int64_t>(), 
        n
    );
}
"""

module = load_inline(
    name="histogram_kernel_module",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["histogram_kernel"],
    verbose=False
)

# Note: input/output are GPU tensors!
def custom_kernel(input: input_t) -> output_t:
    inp_t, out_t = input  # Unpack the input tuple
    
    # Call the kernel directly
    module.histogram_kernel(inp_t, out_t)
    
    return out_t