submission 68058

from typing import Any, Dict, Tuple

import torch
from task import input_t, output_t

import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
from torch.utils.cpp_extension import load_inline

_KERNEL_CACHE: Dict[Tuple[Tuple[int, ...], torch.dtype], Any] = {}
_VECADD_EXT = None


def _get_vecadd_extension():
    global _VECADD_EXT
    if _VECADD_EXT is not None:
        return _VECADD_EXT

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

namespace {

__device__ __forceinline__ ulonglong4 add_ulonglong4(ulonglong4 va, ulonglong4 vb) {
    ulonglong4 vc;
    __half2* ha = reinterpret_cast<__half2*>(&va);
    __half2* hb = reinterpret_cast<__half2*>(&vb);
    __half2* hc = reinterpret_cast<__half2*>(&vc);
#pragma unroll
    for (int i = 0; i < 8; ++i) {
        hc[i] = __hadd2(ha[i], hb[i]);
    }
    return vc;
}

__global__ void vecadd_fp16_kernel(
    const half* __restrict__ a,
    const half* __restrict__ b,
    half* __restrict__ out,
    size_t elements
) {
    constexpr int pack_elems = 16;  // 16 fp16 values per ulonglong4
    const size_t tid = threadIdx.x + blockIdx.x * blockDim.x;
    const size_t stride = static_cast<size_t>(blockDim.x) * gridDim.x;

    const size_t total_packs = elements / pack_elems;
    auto* a_vec = reinterpret_cast<const ulonglong4*>(a);
    auto* b_vec = reinterpret_cast<const ulonglong4*>(b);
    auto* o_vec = reinterpret_cast<ulonglong4*>(out);

    const size_t stride_packs = stride;
    auto* a_ptr = a_vec + tid;
    auto* b_ptr = b_vec + tid;
    auto* o_ptr = o_vec + tid;
    for (size_t pack = tid; pack < total_packs; pack += stride_packs * 2) {
        if (pack < total_packs) {
            ulonglong4 va0 = *a_ptr;
            ulonglong4 vb0 = *b_ptr;
            *o_ptr = add_ulonglong4(va0, vb0);
        }
        a_ptr += stride_packs;
        b_ptr += stride_packs;
        o_ptr += stride_packs;

        const size_t pack1 = pack + stride_packs;
        if (pack1 < total_packs) {
            ulonglong4 va1 = *a_ptr;
            ulonglong4 vb1 = *b_ptr;
            *o_ptr = add_ulonglong4(va1, vb1);
        }
        a_ptr += stride_packs;
        b_ptr += stride_packs;
        o_ptr += stride_packs;
    }

    const size_t tail_start = total_packs * pack_elems;
    for (size_t idx = tail_start + tid; idx < elements; idx += stride) {
        out[idx] = __hadd(a[idx], b[idx]);
    }
}

}  // namespace

void vecadd_fp16(torch::Tensor a, torch::Tensor b, torch::Tensor out) {
    TORCH_CHECK(a.is_cuda(), "Input tensor A must be CUDA");
    TORCH_CHECK(b.is_cuda(), "Input tensor B must be CUDA");
    TORCH_CHECK(out.is_cuda(), "Output tensor must be CUDA");
    TORCH_CHECK(a.dtype() == torch::kFloat16, "A must be float16");
    TORCH_CHECK(b.dtype() == torch::kFloat16, "B must be float16");
    TORCH_CHECK(out.dtype() == torch::kFloat16, "output must be float16");
    TORCH_CHECK(a.numel() == b.numel(), "Input sizes must match");
    TORCH_CHECK(a.numel() == out.numel(), "Output size mismatch");

    auto stream = at::cuda::getCurrentCUDAStream();
    const auto elements = static_cast<size_t>(out.numel());

    constexpr int block = 320;
    const int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    constexpr int pack_elems = 16;
    const size_t total_packs = elements / pack_elems;
    int grid = static_cast<int>((total_packs + block - 1) / block);
    grid = std::max(1, grid);

    vecadd_fp16_kernel<<<grid, block, 0, stream>>>(
        reinterpret_cast<const half*>(a.data_ptr<at::Half>()),
        reinterpret_cast<const half*>(b.data_ptr<at::Half>()),
        reinterpret_cast<half*>(out.data_ptr<at::Half>()),
        elements
    );
    AT_CUDA_CHECK(cudaGetLastError());
}
"""

    _VECADD_EXT = load_inline(
        name="vectoradd_fp16_ext",
        cpp_sources="#include <torch/extension.h>\nvoid vecadd_fp16(torch::Tensor a, torch::Tensor b, torch::Tensor out);\n",
        cuda_sources=cuda_src,
        functions=["vecadd_fp16"],
        with_cuda=True,
        extra_cuda_cflags=[
            "-O3",
            "--use_fast_math",
            "-U__CUDA_NO_HALF_OPERATORS__",
            "-U__CUDA_NO_HALF_CONVERSIONS__",
        ],
    )
    return _VECADD_EXT


def _is_aligned(tensor: torch.Tensor, alignment: int = 16) -> bool:
    return tensor.data_ptr() % alignment == 0


@cute.kernel
def _tile_elementwise_add_kernel(
    gA: cute.Tensor,
    gB: cute.Tensor,
    gC: cute.Tensor,
    tv_layout,
) -> None:
    """Block-tiled elementwise add with per-thread vectorization."""
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    grid_x, _, _ = cute.arch.grid_dim()
    num_tiles = cute.size(gC, mode=[1])

    tile_idx = bidx
    while tile_idx < num_tiles:
        blk_coord = ((None, None), tile_idx)
        blkA = gA[blk_coord]
        blkB = gB[blk_coord]
        blkC = gC[blk_coord]

        tidfrgA = cute.composition(blkA, tv_layout)
        tidfrgB = cute.composition(blkB, tv_layout)
        tidfrgC = cute.composition(blkC, tv_layout)

        thr_coord = (tidx, None)
        thrA = tidfrgA[thr_coord]
        thrB = tidfrgB[thr_coord]
        thrC = tidfrgC[thr_coord]

        thrC[None] = thrA.load() + thrB.load()

        tile_idx += grid_x


@cute.jit
def _tiled_elementwise_add(
    mA: cute.Tensor,
    mB: cute.Tensor,
    mC: cute.Tensor,
) -> None:
    """Launch tiled kernel honoring remainder shapes."""
    element_type = mA.element_type
    bytes_per_element = element_type.width // 8
    coalesced_ldst_bytes = bytes_per_element

    thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0))
    val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0))
    val_layout = cute.recast_layout(element_type.width, 8, val_layout)
    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)

    gA = cute.zipped_divide(mA, tiler_mn)
    gB = cute.zipped_divide(mB, tiler_mn)
    gC = cute.zipped_divide(mC, tiler_mn)

    remap_block = cute.make_ordered_layout(
        cute.select(gA.shape[1], mode=[1, 0]), order=(1, 0)
    )
    gA = cute.composition(gA, (None, remap_block))
    gB = cute.composition(gB, (None, remap_block))
    gC = cute.composition(gC, (None, remap_block))

    total_tiles = int(cute.size(gC, mode=[1]))
    block_dim_x = int(cute.size(tv_layout, mode=[0]))
    max_ctas = 128
    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(torch.cuda.current_device())
        max_ctas = max(1, props.multi_processor_count * 2)
    grid_dim_x = min(total_tiles, max_ctas) if total_tiles > 0 else 1

    _tile_elementwise_add_kernel(gA, gB, gC, tv_layout).launch(
        grid=(grid_dim_x, 1, 1),
        block=(block_dim_x, 1, 1),
    )


def _ensure_contiguous(tensor: torch.Tensor) -> torch.Tensor:
    if tensor.is_contiguous():
        return tensor
    return tensor.contiguous()


def _torch_to_cute(tensor: torch.Tensor) -> cute.Tensor:
    return from_dlpack(tensor, assumed_align=16)


def _get_compiled_kernel(
    key: Tuple[Tuple[int, ...], torch.dtype],
    a_: cute.Tensor,
    b_: cute.Tensor,
    c_: cute.Tensor,
):
    compiled = _KERNEL_CACHE.get(key)
    if compiled is None:
        compiled = cute.compile(_tiled_elementwise_add, a_, b_, c_)
        _KERNEL_CACHE[key] = compiled
    return compiled


def custom_kernel(data: input_t) -> output_t:
    """Run vector addition via a tiled Cutlass (CuTe) kernel."""
    A, B, output = data

    A_contig = _ensure_contiguous(A)
    B_contig = _ensure_contiguous(B)
    output_contig = output if output.is_contiguous() else output.contiguous()

    a_cute = _torch_to_cute(A_contig)
    b_cute = _torch_to_cute(B_contig)
    c_cute = _torch_to_cute(output_contig)

    if (
        A_contig.is_cuda
        and A_contig.dtype == torch.float16
        and _is_aligned(A_contig, 32)
        and _is_aligned(B_contig, 32)
        and _is_aligned(output_contig, 32)
    ):
        ext = _get_vecadd_extension()
        ext.vecadd_fp16(A_contig, B_contig, output_contig)
    else:
        kernel_key = (tuple(A_contig.shape), A_contig.dtype)
        compiled_kernel = _get_compiled_kernel(
            kernel_key, a_cute, b_cute, c_cute
        )
        compiled_kernel(a_cute, b_cute, c_cute)

    if output_contig is not output:
        output.copy_(output_contig)

    return output