submission 488446

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

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_fp16.h>

__global__ void vectorAdd_float4(
    const half* __restrict__ a,
    const half* __restrict__ b,
    half* __restrict__ c,
    int n
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Each thread processes 8 fp16 values (one float4 = 16 bytes = 8× half)
    int idx8 = tid * 8;
    
    if (idx8 + 7 < n) {
        // Load 16 bytes (8× fp16) as float4
        float4 a_vec = reinterpret_cast<const float4*>(a)[tid];
        float4 b_vec = reinterpret_cast<const float4*>(b)[tid];
        
        // Reinterpret as 4× half2
        half2* a_h2 = reinterpret_cast<half2*>(&a_vec);
        half2* b_h2 = reinterpret_cast<half2*>(&b_vec);
        
        half2 c_h2[4];
        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]);
        
        // Store back as float4
        reinterpret_cast<float4*>(c)[tid] = *reinterpret_cast<float4*>(c_h2);
    }
    
    // Handle tail elements (up to 7 remaining)
    if (tid == 0) {
        int n8 = (n >> 3) << 3;  // Round down to multiple of 8
        for (int i = n8; i < n; i++) {
            c[i] = __hadd(a[i], b[i]);
        }
    }
}

void vectoradd_kernel(
    torch::Tensor input1,
    torch::Tensor input2,
    torch::Tensor output
) {
    const half* a = reinterpret_cast<const half*>(input1.data_ptr<at::Half>());
    const half* b = reinterpret_cast<const half*>(input2.data_ptr<at::Half>());
    half* c = reinterpret_cast<half*>(output.data_ptr<at::Half>());
    
    int n = input1.numel();
    int n8 = (n + 7) >> 3;  // Number of float4 chunks
    
    int threads = 512;
    int blocks = (n8 + threads - 1) / threads;
    
    vectorAdd_float4<<<blocks, threads>>>(a, b, c, n);
}
"""

cpp_source = r"""
void vectoradd_kernel(torch::Tensor, torch::Tensor, torch::Tensor);
"""

module = load_inline(
    name="vectoradd_float4",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["vectoradd_kernel"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    verbose=False,
)

def ref_kernel(data: input_t) -> output_t:
    with DeterministicContext():
        A, B, output = data
        output[...] = A + B
        return output

def generate_input(size: int, seed: int) -> input_t:
    gen = torch.Generator(device="cuda")
    gen.manual_seed(seed)
    A = torch.randn(size, size, device="cuda", dtype=torch.float16, generator=gen).contiguous()
    B = torch.randn(size, size, device="cuda", dtype=torch.float16, generator=gen).contiguous()
    C = torch.empty_like(A)
    return A, B, C

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

check_implementation = make_match_reference(ref_kernel)