submission 68391

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

# NOTE: All sizes are multiples of 16 !!!

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;

template <size_t Width>
struct uint8v {
    uint8_t data[Width];
};

template <size_t ThreadWidth, size_t JoinWidth>
__global__ void histogram_block(uint8_t* input, int64_t* output, int n) {
    static_assert(ThreadWidth % JoinWidth == 0);

   __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. Local Histogram Calculation.
    if (idx * ThreadWidth < n) {
        uint8v<ThreadWidth> items = reinterpret_cast<uint8v<ThreadWidth>*>(input)[idx];

        for (int joinTile = 0; joinTile < ThreadWidth / JoinWidth; joinTile++) {
            uint8v<JoinWidth> tile = reinterpret_cast<uint8v<JoinWidth>*>(items.data)[joinTile];

            bool equal = true;
            int first  = tile.data[0];
            for (int tileLane = 1; tileLane < JoinWidth; tileLane++) {
                if (first != tile.data[tileLane]) {
                    equal = false;
                }
            }

            if (equal) {
                atomicAdd(&local[first], JoinWidth);
            } else {
                for (int tileLane = 0; tileLane < JoinWidth; tileLane++) {
                    atomicAdd(&local[tile.data[tileLane]], 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);

    uint8_t* d_in  = input.data_ptr<uint8_t>();
    int64_t* d_out = output.data_ptr<int64_t>();

    if (n < 16384) {
        constexpr int threadSize = 1;
        constexpr int joinWidth  = 1; 
        int threadsPerBlock = 512;
        int workPerBlock    = threadsPerBlock * threadSize;

        int blockCount = (n + workPerBlock - 1) / workPerBlock;

        histogram_block<threadSize, joinWidth><<<blockCount, threadsPerBlock>>> (d_in, d_out, n);
    } else if (n < 65536) {
        constexpr int threadSize = 4;
        constexpr int joinWidth  = 2; 
        int threadsPerBlock = 512;
        int workPerBlock    = threadsPerBlock * threadSize;

        int blockCount = (n + workPerBlock - 1) / workPerBlock;

        histogram_block<threadSize, joinWidth><<<blockCount, threadsPerBlock>>> (d_in, d_out, n);
    } else {
        constexpr int threadSize = 8;
        constexpr int joinWidth  = 2; 
        int threadsPerBlock = 512;
        int workPerBlock    = threadsPerBlock * threadSize;

        int blockCount = (n + workPerBlock - 1) / workPerBlock;

        histogram_block<threadSize, joinWidth><<<blockCount, threadsPerBlock>>> (d_in, d_out, 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