submission 67752

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

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

__global__ void __launch_bounds__(256, 4) vectoradd_kernel(
    const half* __restrict__ a,
    const half* __restrict__ b,
    half* __restrict__ out,
    const int n)
{
    // Coalesced: consecutive threads access consecutive memory
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int base = idx * 16;  // 16 elements per thread for coalescing

    if (base + 15 < n) {
        uint4 a0, a1, b0, b1, c0, c1;

        const half* ap = a + base;
        const half* bp = b + base;
        half* cp = out + base;

        // Load 16 fp16 (2x 128-bit)
        asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(a0.x),"=r"(a0.y),"=r"(a0.z),"=r"(a0.w) : "l"(ap));
        asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(a1.x),"=r"(a1.y),"=r"(a1.z),"=r"(a1.w) : "l"(ap+8));
        asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(b0.x),"=r"(b0.y),"=r"(b0.z),"=r"(b0.w) : "l"(bp));
        asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(b1.x),"=r"(b1.y),"=r"(b1.z),"=r"(b1.w) : "l"(bp+8));

        // SIMD add
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.x) : "r"(a0.x), "r"(b0.x));
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.y) : "r"(a0.y), "r"(b0.y));
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.z) : "r"(a0.z), "r"(b0.z));
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.w) : "r"(a0.w), "r"(b0.w));
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.x) : "r"(a1.x), "r"(b1.x));
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.y) : "r"(a1.y), "r"(b1.y));
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.z) : "r"(a1.z), "r"(b1.z));
        asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.w) : "r"(a1.w), "r"(b1.w));

        // Store
        asm volatile ("st.global.cg.v4.u32 [%0],{%1,%2,%3,%4};" :: "l"(cp),"r"(c0.x),"r"(c0.y),"r"(c0.z),"r"(c0.w));
        asm volatile ("st.global.cg.v4.u32 [%0],{%1,%2,%3,%4};" :: "l"(cp+8),"r"(c1.x),"r"(c1.y),"r"(c1.z),"r"(c1.w));
    }
    else if (base < n) {
        for (int i = 0; i < 16 && base + i < n; i++) {
            out[base + i] = __hadd(a[base + i], b[base + i]);
        }
    }
}

torch::Tensor vectoradd_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor out) {
    const int n = a.numel();
    const half* a_ptr = reinterpret_cast<const half*>(a.data_ptr<at::Half>());
    const half* b_ptr = reinterpret_cast<const half*>(b.data_ptr<at::Half>());
    half* out_ptr = reinterpret_cast<half*>(out.data_ptr<at::Half>());

    const int threads = 256;
    const int blocks = (n + threads * 16 - 1) / (threads * 16);

    vectoradd_kernel<<<blocks, threads>>>(a_ptr, b_ptr, out_ptr, n);
    return out;
}
"""

cpp_source = "torch::Tensor vectoradd_cuda(torch::Tensor, torch::Tensor, torch::Tensor);"

_module = None

def _get_module():
    global _module
    if _module is None:
        _module = load_inline(
            name='vectoradd_coalesced',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['vectoradd_cuda'],
            extra_cuda_cflags=['-O3', '-use_fast_math', '-arch=sm_80', '-maxrregcount=64'],
            verbose=False
        )
    return _module

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