submission 614486

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

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

cuda_source = r"""
#include <cuda_fp16.h>
#include <cuda_runtime.h>

__global__ __launch_bounds__(512, 2)
void vecadd_kernel(const half* __restrict__ A,
                   const half* __restrict__ B,
                   half* __restrict__ C,
                   const int N) {
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int idx = tid * 8;
    if (idx + 7 < N) {
        const float4 a = __ldg(reinterpret_cast<const float4*>(A + idx));
        const float4 b = __ldg(reinterpret_cast<const float4*>(B + idx));

        const half2* a_h2 = reinterpret_cast<const half2*>(&a);
        const half2* b_h2 = reinterpret_cast<const half2*>(&b);

        float4 c;
        half2* c_h2 = reinterpret_cast<half2*>(&c);

        c_h2[0] = __hadd2(a_h2[0], b_h2[0]);
        c_h2[1] = __hadd2(a_h2[1], b_h2[1]);
        c_h2[2] = __hadd2(a_h2[2], b_h2[2]);
        c_h2[3] = __hadd2(a_h2[3], b_h2[3]);

        *reinterpret_cast<float4*>(C + idx) = c;
    } else if (idx < N) {
        for (int i = idx; i < N && i < idx + 8; i++) {
            C[i] = __hadd(A[i], B[i]);
        }
    }
}

void vecadd_cuda(torch::Tensor A, torch::Tensor B, torch::Tensor C, int N) {
    const int threads = 512;
    const int elems_per_thread = 8;
    const int blocks = (N + threads * elems_per_thread - 1) / (threads * elems_per_thread);
    vecadd_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
    );
}
"""

cpp_source = r"""
#include <torch/extension.h>
void vecadd_cuda(torch::Tensor A, torch::Tensor B, torch::Tensor C, int N);
"""

module = load_inline(
    name="vecadd_v2",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["vecadd_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-maxrregcount=24", "--gpu-architecture=sm_100"],
    verbose=False,
)

def custom_kernel(data: input_t) -> output_t:
    A, B, output = data
    N = A.numel()
    module.vecadd_cuda(A, B, output, N)
    return output