submission 67857

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

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_fp16.h>
#include <stdint.h>

template<int VEC>
__global__ void vec_add_half2_ptx(const __half* __restrict__ A,
                                  const __half* __restrict__ B,
                                  __half* __restrict__ C,
                                  int64_t N) {
    const int64_t idx0 = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * VEC;
    const int64_t stride = int64_t(blockDim.x) * gridDim.x * VEC;

    const uintptr_t baseA = reinterpret_cast<uintptr_t>(A);
    const uintptr_t baseB = reinterpret_cast<uintptr_t>(B);
    const uintptr_t baseC = reinterpret_cast<uintptr_t>(C);
    const bool aligned16 = (((baseA | baseB | baseC) & 0xF) == 0);

    for (int64_t idx = idx0; idx < N; idx += stride) {
        int64_t remaining = N - idx;
        if (remaining >= VEC) {
            const char* a_ptr = reinterpret_cast<const char*>(A) + idx * 2;
            const char* b_ptr = reinterpret_cast<const char*>(B) + idx * 2;
            char*       c_ptr = reinterpret_cast<char*>(C) + idx * 2;

            if (aligned16) {
                // 128-bit vectorized path: load/store 16 bytes per thread, do 4x add.f16x2
                unsigned rA0, rA1, rA2, rA3;
                unsigned rB0, rB1, rB2, rB3;
                unsigned rC0, rC1, rC2, rC3;
                asm volatile(
                    "{\n\t"
                    "ld.global.v4.b32 {%0,%1,%2,%3}, [%12];\n\t"
                    "ld.global.v4.b32 {%4,%5,%6,%7}, [%13];\n\t"
                    "add.f16x2 %8,  %0, %4;\n\t"
                    "add.f16x2 %9,  %1, %5;\n\t"
                    "add.f16x2 %10, %2, %6;\n\t"
                    "add.f16x2 %11, %3, %7;\n\t"
                    "st.global.v4.b32 [%14], {%8,%9,%10,%11};\n\t"
                    "}\n"
                    : "=&r"(rA0), "=&r"(rA1), "=&r"(rA2), "=&r"(rA3),
                      "=&r"(rB0), "=&r"(rB1), "=&r"(rB2), "=&r"(rB3),
                      "=&r"(rC0), "=&r"(rC1), "=&r"(rC2), "=&r"(rC3)
                    : "l"(a_ptr), "l"(b_ptr), "l"(c_ptr)
                    : "memory");
            } else {
                // 32-bit path: tolerate only 4B alignment using scalar b32 ops, still using add.f16x2
                unsigned rA0, rA1, rA2, rA3;
                unsigned rB0, rB1, rB2, rB3;
                unsigned rC0, rC1, rC2, rC3;

                const char* a0 = a_ptr + 0;
                const char* a1 = a_ptr + 4;
                const char* a2 = a_ptr + 8;
                const char* a3 = a_ptr + 12;

                const char* b0 = b_ptr + 0;
                const char* b1 = b_ptr + 4;
                const char* b2 = b_ptr + 8;
                const char* b3 = b_ptr + 12;

                char* c0 = c_ptr + 0;
                char* c1 = c_ptr + 4;
                char* c2 = c_ptr + 8;
                char* c3 = c_ptr + 12;

                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA0) : "l"(a0) : "memory");
                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA1) : "l"(a1) : "memory");
                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA2) : "l"(a2) : "memory");
                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA3) : "l"(a3) : "memory");

                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB0) : "l"(b0) : "memory");
                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB1) : "l"(b1) : "memory");
                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB2) : "l"(b2) : "memory");
                asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB3) : "l"(b3) : "memory");

                asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC0) : "r"(rA0), "r"(rB0));
                asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC1) : "r"(rA1), "r"(rB1));
                asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC2) : "r"(rA2), "r"(rB2));
                asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC3) : "r"(rA3), "r"(rB3));

                asm volatile("st.global.b32 [%0], %1;" :: "l"(c0), "r"(rC0) : "memory");
                asm volatile("st.global.b32 [%0], %1;" :: "l"(c1), "r"(rC1) : "memory");
                asm volatile("st.global.b32 [%0], %1;" :: "l"(c2), "r"(rC2) : "memory");
                asm volatile("st.global.b32 [%0], %1;" :: "l"(c3), "r"(rC3) : "memory");
            }
        } else {
            // Tail: scalar half add for remaining < VEC elements
            for (int64_t t = 0; t < remaining; ++t) {
                C[idx + t] = __hadd(A[idx + t], B[idx + t]);
            }
        }
    }
}

void vec_add_ptx(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda() && C.is_cuda(), "Tensors must be CUDA");
    TORCH_CHECK(A.scalar_type() == at::kHalf && B.scalar_type() == at::kHalf && C.scalar_type() == at::kHalf, "Tensors must be float16 (half)");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && C.is_contiguous(), "Tensors must be contiguous");
    TORCH_CHECK(A.numel() == B.numel() && A.numel() == C.numel(), "Size mismatch");

    c10::cuda::CUDAGuard guard(A.device());
    const int64_t N = A.numel();
    if (N == 0) return;

    const int threads = 256;      // good balance for SM80
    const int VEC = 8;            // 8 fp16 elems (16B) per thread for 128-bit transactions
    const int64_t elems_per_block = int64_t(threads) * VEC;
    int blocks = static_cast<int>((N + elems_per_block - 1) / elems_per_block);
    blocks = std::min(blocks, 65535); // cap to keep scheduler happy

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    vec_add_half2_ptx<VEC><<<blocks, threads, 0, stream>>>(
        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
    );
    AT_CUDA_CHECK(cudaGetLastError());
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("vec_add_ptx", &vec_add_ptx, "Vector add using inline PTX (half)");
}
'''

_ext_mod = None

def _get_ext_mod():
    global _ext_mod
    if _ext_mod is None:
        _ext_mod = load_inline(
            name="vecadd_half_ptx_sm80",
            cpp_sources="",
            cuda_sources=_cuda_src,
            extra_cuda_cflags=[
                "-O3",
                "--use_fast_math",
                "-U__CUDA_NO_HALF_OPERATORS__",
                "-U__CUDA_NO_HALF_CONVERSIONS__",
                "-gencode=arch=compute_80,code=sm_80",
            ],
            verbose=False,
        )
    return _ext_mod

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

    # Fast path: CUDA half, contiguous, same numel
    if (
        A.is_cuda and B.is_cuda and output.is_cuda and
        A.dtype == torch.float16 and B.dtype == torch.float16 and output.dtype == torch.float16 and
        A.is_contiguous() and B.is_contiguous() and output.is_contiguous() and
        A.numel() == B.numel() == output.numel()
    ):
        mod = _get_ext_mod()
        mod.vec_add_ptx(A, B, output)
        return output

    # Safe fallback without allocating a new output tensor
    output.copy_(A)
    output.add_(B)
    return output