submission 68301

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 <ATen/ATen.h>
#include <torch/extension.h>
#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;

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

    // Build histogram in shared memory
    for (int i = blockIdx.x * blockDim.x + tid; i < len; i += blockDim.x * gridDim.x) {
        atomicAdd(&T[A[i]], 1ULL);
    }
    __syncthreads();

    // Accumulate shared memory histogram into global memory
    for (int i = tid; i < MAX_VAL; i += blockDim.x) {
        atomicAdd(&output[i], T[i]);
    }
}

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

    int blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    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)