submission 68318

import torch
from utils import DeterministicContext
from torch.utils.cpp_extension import load_inline
from typing import List
from task import input_t, output_t

cuda_histogram = """
#define BLOCK_SIZE 1024
#define MAX_VAL 256
#include <cstdint>

__global__ void histogram(const unsigned char* __restrict__ A, unsigned long long* __restrict__ output, int len) {
    __shared__ unsigned long long T[MAX_VAL];

    int tid = threadIdx.x;
    int lane_id = threadIdx.x & 31;  // Thread ID within warp (0-31)

    // Initialize shared memory histogram
    for (int i = tid; i < MAX_VAL; i += blockDim.x) T[i] = 0;
    __syncthreads();

    // Build histogram in shared memory with warp-level aggregation
    for (int i = blockIdx.x * blockDim.x + tid; i < len; i += blockDim.x * gridDim.x) {
        unsigned char bin = A[i];

#if __CUDA_ARCH__ >= 700
        // Use warp match intrinsic (Volta+) for efficient aggregation
        unsigned int match_mask = __match_any_sync(0xFFFFFFFF, bin);
        int leader = __ffs(match_mask) - 1;  // Find first set bit (lowest lane with this bin)

        if (lane_id == leader) {
            // Leader thread counts how many threads have the same bin
            int count = __popc(match_mask);
            atomicAdd(&T[bin], (unsigned long long)count);
        }
#else
        // Fallback: direct atomic add per thread
        atomicAdd(&T[bin], 1ULL);
#endif
    }
    __syncthreads();    // Accumulate shared memory histogram into global memory
    #pragma unroll
    for (int i = tid; i < MAX_VAL; i += blockDim.x) {
        if (T[i] > 0) {
            atomicAdd(&output[i], T[i]);
        }
    }
}

torch::Tensor& histogram(torch::Tensor& A, torch::Tensor& B) {
    int N = A.numel();

    int blocks = min((N + BLOCK_SIZE - 1) / BLOCK_SIZE, 1024);
    histogram<<<blocks, BLOCK_SIZE>>>(
        A.data_ptr<unsigned char>(),
        reinterpret_cast<unsigned long long*>(B.data_ptr<int64_t>()),
        N
    );
    return B;
}
"""

abi = torch._C._GLIBCXX_USE_CXX11_ABI
extra_compile_args={'cxx': [f"-O3", f"-D_GLIBCXX_USE_CXX11_ABI={abi}"],
                    'nvcc': ["-O3"]}

histogram_md = load_inline(
    name='sum_cuda_ext',
    cpp_sources="torch::Tensor& histogram(torch::Tensor& A, torch::Tensor& B);",
    cuda_sources=cuda_histogram,
    functions=['histogram'],
    extra_cflags=extra_compile_args['cxx'],
    extra_cuda_cflags=extra_compile_args['nvcc'],
    verbose=True,
)

def custom_kernel(data: input_t) -> output_t:
    """
    Custom implementation of vector addition using CUDA.
    Args:
        inputs: List of pairs of tensors [A, B] to be added.
    Returns:
        Tensor containing element-wise sum.
    """
    A, B = data
    B.zero_()  # Initialize output tensor to zero
    return histogram_md.histogram(A, B)