submission 68310

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

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

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

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

constexpr int threadSize = 1;
constexpr int threadsPerBlock = 128;

constexpr int blockSize = threadsPerBlock * threadSize;

__global__ void histogram_block(uint8_t* input, int64_t* output, int n) {
    // Kernel implementation here

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

    if (idx < n) {
        uint8_t item = input[idx];

        atomicAdd(reinterpret_cast<unsigned long long int*>(&output[item]), 1);
    } 
}

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