submission 610921

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

add_cuda_source = """
#include <cuda_fp16.h>

__global__ void add_kernel(const half* __restrict__ A, 
                           const half* __restrict__ B, 
                           half* __restrict__ C, 
                           int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    // Each thread handles 8 fp16 elements via one float4 (16 bytes) load
    const int elems_per_thread = 8;
    int base = idx * elems_per_thread;

    if (base + elems_per_thread <= N) {
        float4 a = *reinterpret_cast<const float4*>(A + base);
        float4 b = *reinterpret_cast<const float4*>(B + base);

        half2* a_h = reinterpret_cast<half2*>(&a);
        half2* b_h = reinterpret_cast<half2*>(&b);
        half2 c_h[4];

        #pragma unroll
        for (int i = 0; i < 4; i++) {
            c_h[i] = __hadd2(a_h[i], b_h[i]);
        }

        *reinterpret_cast<float4*>(C + base) = *reinterpret_cast<float4*>(c_h);
    } else {
        // Tail: handle remaining elements one by one
        for (int i = base; i < N; i++) {
            C[i] = __hadd(A[i], B[i]);
        }
    }
}

torch::Tensor add_cuda(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
    TORCH_CHECK(A.device().is_cuda(), "Tensor A must be a CUDA tensor");
    TORCH_CHECK(B.device().is_cuda(), "Tensor B must be a CUDA tensor");
    TORCH_CHECK(C.device().is_cuda(), "Tensor C must be a CUDA tensor");
    TORCH_CHECK(A.sizes() == B.sizes(), "Input tensors must have the same size");
    
    int N = A.numel();  

    const int threads = 256; 
    const int elems_per_thread = 8;
    const int blocks = (N + threads * elems_per_thread - 1) / (threads * elems_per_thread);  
    
    add_kernel<<<blocks, threads>>>(
        reinterpret_cast<const half*>(A.data_ptr<at::Half>()),
        reinterpret_cast<const half*>(B.data_ptr<at::Half>()),
        reinterpret_cast<half*>(C.data_ptr<at::Half>()),
        N
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }

    return C;
}
"""

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

torch::Tensor add_cuda(torch::Tensor A, torch::Tensor B, torch::Tensor C);
"""

add_module = load_inline(
    name='add_cuda',
    cpp_sources=add_cpp_source,
    cuda_sources=add_cuda_source,
    functions=['add_cuda'],
    verbose=True,
)

def custom_kernel(data: input_t) -> output_t:
    A, B, C = data

    assert A.is_cuda and B.is_cuda, "Input tensors must be on GPU"
    assert A.shape == B.shape, "Input tensors must have the same shape"
    assert A.dtype == torch.float16 and B.dtype == torch.float16, "Input tensors must be float16"

    return add_module.add_cuda(A, B, C)