submission 66564

import torch
from torch.utils.cpp_extension import load_inline

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <stdint.h>

// Helpers --------------------------------------------------------------------
static inline int div_up_long(long long n, int d) {
    return (int)((n + d - 1) / d);
}

// Union to safely bit-cast between u32 and __half2
union u32_h2 {
    uint32_t u;
    __half2  h2;
};

// fast path: 128-bit vectorized (8 x fp16 per pack), requires 16B alignment
__global__ void add_fp16_vec8_kernel(const __half* __restrict__ a,
                                     const __half* __restrict__ b,
                                     __half* __restrict__ out,
                                     size_t n_packs /* number of 8-half packs */) {
    const uint4* __restrict__ a4 = reinterpret_cast<const uint4*>(a);
    const uint4* __restrict__ b4 = reinterpret_cast<const uint4*>(b);
    uint4* __restrict__ out4     = reinterpret_cast<uint4*>(out);

    const size_t tid    = blockIdx.x * blockDim.x + threadIdx.x;
    const size_t stride = (size_t)blockDim.x * gridDim.x;

    // ILP to increase memory-level parallelism per thread
    const int ILP = 4;
    size_t i = tid;

    for (; i + (size_t)(ILP - 1) * stride < n_packs; i += (size_t)ILP * stride) {
        #pragma unroll
        for (int k = 0; k < ILP; ++k) {
            size_t idx = i + (size_t)k * stride;
            uint4 av = a4[idx];
            uint4 bv = b4[idx];

            // Each uint4.{x,y,z,w} is 32-bit = __half2
            u32_h2 a0; a0.u = av.x; u32_h2 b0; b0.u = bv.x; u32_h2 r0; r0.h2 = __hadd2(a0.h2, b0.h2);
            u32_h2 a1; a1.u = av.y; u32_h2 b1; b1.u = bv.y; u32_h2 r1; r1.h2 = __hadd2(a1.h2, b1.h2);
            u32_h2 a2; a2.u = av.z; u32_h2 b2; b2.u = bv.z; u32_h2 r2; r2.h2 = __hadd2(a2.h2, b2.h2);
            u32_h2 a3; a3.u = av.w; u32_h2 b3; b3.u = bv.w; u32_h2 r3; r3.h2 = __hadd2(a3.h2, b3.h2);

            uint4 rv;
            rv.x = r0.u; rv.y = r1.u; rv.z = r2.u; rv.w = r3.u;
            out4[idx] = rv;
        }
    }
    for (; i < n_packs; i += stride) {
        uint4 av = a4[i];
        uint4 bv = b4[i];
        u32_h2 a0; a0.u = av.x; u32_h2 b0; b0.u = bv.x; u32_h2 r0; r0.h2 = __hadd2(a0.h2, b0.h2);
        u32_h2 a1; a1.u = av.y; u32_h2 b1; b1.u = bv.y; u32_h2 r1; r1.h2 = __hadd2(a1.h2, b1.h2);
        u32_h2 a2; a2.u = av.z; u32_h2 b2; b2.u = bv.z; u32_h2 r2; r2.h2 = __hadd2(a2.h2, b2.h2);
        u32_h2 a3; a3.u = av.w; u32_h2 b3; b3.u = bv.w; u32_h2 r3; r3.h2 = __hadd2(a3.h2, b3.h2);
        uint4 rv; rv.x = r0.u; rv.y = r1.u; rv.z = r2.u; rv.w = r3.u;
        out4[i] = rv;
    }
}

// medium path: __half2 (4B), requires 4B alignment; 2 x fp16 per element
__global__ void add_fp16_h2_kernel(const __half2* __restrict__ a,
                                   const __half2* __restrict__ b,
                                   __half2* __restrict__ out,
                                   size_t n_h2) {
    const size_t tid    = blockIdx.x * blockDim.x + threadIdx.x;
    const size_t stride = (size_t)blockDim.x * gridDim.x;
    for (size_t i = tid; i < n_h2; i += stride) {
        out[i] = __hadd2(a[i], b[i]);
    }
}

// tail / worst-case path: scalar half (handles misalignment + remainder)
__global__ void add_fp16_scalar_kernel(const __half* __restrict__ a,
                                       const __half* __restrict__ b,
                                       __half* __restrict__ out,
                                       size_t n) {
    const size_t tid    = blockIdx.x * blockDim.x + threadIdx.x;
    const size_t stride = (size_t)blockDim.x * gridDim.x;
    for (size_t i = tid; i < n; i += stride) {
        out[i] = __hadd(a[i], b[i]);
    }
}

void add_fp16_out(torch::Tensor a, torch::Tensor b, torch::Tensor out) {
    TORCH_CHECK(a.is_cuda() && b.is_cuda() && out.is_cuda(), "Tensors must be CUDA");
    TORCH_CHECK(a.scalar_type() == torch::kHalf && b.scalar_type() == torch::kHalf && out.scalar_type() == torch::kHalf,
                "Tensors must be torch.float16");
    TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && out.is_contiguous(), "Tensors must be contiguous");
    TORCH_CHECK(a.sizes() == b.sizes() && a.sizes() == out.sizes(), "Shapes must match");

    c10::cuda::CUDAGuard device_guard(a.get_device());
    auto* a_ptr = reinterpret_cast<const __half*>(a.data_ptr<at::Half>());
    auto* b_ptr = reinterpret_cast<const __half*>(b.data_ptr<at::Half>());
    auto* o_ptr = reinterpret_cast<__half*>(out.data_ptr<at::Half>());

    const size_t n = (size_t)out.numel();
    if (n == 0) return;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Choose launch geometry
    const int threads = 256;
    const int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    int blocks;

    // Check alignments
    uintptr_t a_addr = reinterpret_cast<uintptr_t>(a_ptr);
    uintptr_t b_addr = reinterpret_cast<uintptr_t>(b_ptr);
    uintptr_t o_addr = reinterpret_cast<uintptr_t>(o_ptr);

    const bool aligned16 = ((a_addr | b_addr | o_addr) & 0xF) == 0; // 16B
    const bool aligned4  = ((a_addr | b_addr | o_addr) & 0x3) == 0; // 4B

    if (aligned16) {
        const size_t n_packs = n / 8;  // 8 half per 128-bit pack
        if (n_packs) {
            blocks = std::min(div_up_long((long long)n_packs, threads), sm * 32);
            add_fp16_vec8_kernel<<<blocks, threads, 0, stream>>>(a_ptr, b_ptr, o_ptr, n_packs);
        }
        const size_t tail = n - n_packs * 8;
        if (tail) {
            const __half* a_tail = a_ptr + n_packs * 8;
            const __half* b_tail = b_ptr + n_packs * 8;
            __half*       o_tail = o_ptr + n_packs * 8;
            blocks = std::min(div_up_long((long long)tail, threads), sm * 4);
            add_fp16_scalar_kernel<<<blocks, threads, 0, stream>>>(a_tail, b_tail, o_tail, tail);
        }
    } else if (aligned4) {
        const size_t n_h2 = n / 2;
        if (n_h2) {
            blocks = std::min(div_up_long((long long)n_h2, threads), sm * 32);
            const __half2* a2 = reinterpret_cast<const __half2*>(a_ptr);
            const __half2* b2 = reinterpret_cast<const __half2*>(b_ptr);
            __half2*       o2 = reinterpret_cast<__half2*>(o_ptr);
            add_fp16_h2_kernel<<<blocks, threads, 0, stream>>>(a2, b2, o2, n_h2);
        }
        if (n & 1) {
            const __half* a_tail = a_ptr + (n_h2 * 2);
            const __half* b_tail = b_ptr + (n_h2 * 2);
            __half*       o_tail = o_ptr + (n_h2 * 2);
            add_fp16_scalar_kernel<<<1, 32, 0, stream>>>(a_tail, b_tail, o_tail, 1);
        }
    } else {
        blocks = std::min(div_up_long((long long)n, threads), sm * 32);
        add_fp16_scalar_kernel<<<blocks, threads, 0, stream>>>(a_ptr, b_ptr, o_ptr, n);
    }

    TORCH_CHECK(cudaGetLastError() == cudaSuccess, "CUDA kernel launch failed");
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("add_fp16_out", &add_fp16_out, "Vector add (fp16) into out (CUDA)");
}
"""

ext = load_inline(
    name="vecadd_fp16_inline",
    cpp_sources="",
    cuda_sources=cuda_src,
    extra_cflags=["-O3"],
    extra_cuda_cflags=[
        "-O3",
        "--use_fast_math",
        # Enable half/half2 intrinsics explicitly:
        "-U__CUDA_NO_HALF_OPERATORS__",
        "-U__CUDA_NO_HALF_CONVERSIONS__",
        "-U__CUDA_NO_HALF2_OPERATORS__",   # <-- added
        # "-maxrregcount=64",  # optional tuning
    ],
    verbose=False,
    with_cuda=True,
)

def add_cuda_(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
    assert A.is_cuda and B.is_cuda and out.is_cuda
    assert A.dtype == torch.float16 and B.dtype == torch.float16 and out.dtype == torch.float16
    assert A.shape == B.shape == out.shape
    assert A.is_contiguous() and B.is_contiguous() and out.is_contiguous()
    ext.add_fp16_out(A, B, out)
    return out

def add_cuda(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(A)
    return add_cuda_(A, B, out)


def custom_kernel(data):
    A, B, output = data
    add_cuda_(A, B, output)   # writes directly into 'output'
    return output