submission 113003

import os
import sys

# --- optional: guard for local eval + Popcorn messing with sys.stdout/stderr ---
if sys.stdout is None:
    if getattr(sys, "__stdout__", None) is not None:
        sys.stdout = sys.__stdout__
    else:
        sys.stdout = open(os.devnull, "w")

if sys.stderr is None:
    if getattr(sys, "__stderr__", None) is not None:
        sys.stderr = sys.__stderr__
    else:
        sys.stderr = open(os.devnull, "w")

# Make sure we target Blackwell in this container
# I do not need other archs .. os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "10.0")
os.environ['TORCH_CUDA_ARCH_LIST'] = '10.0+PTX'  # so I can test on my local sm120

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

# Single .cu file containing both kernel + host wrapper.
add_cuda_source = r"""
#include <torch/extension.h>

template <typename scalar_t>
__global__ void add_kernel(const scalar_t* __restrict__ A,
                           const scalar_t* __restrict__ B,
                           scalar_t* __restrict__ C,
                           int64_t N) {
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host wrapper called from Python. Signature MUST be (Tensor, Tensor)
// because load_inline(functions=["add_cuda"]) will declare it that way.
torch::Tensor add_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.device().is_cuda(), "A must be a CUDA tensor");
    TORCH_CHECK(B.device().is_cuda(), "B must be a CUDA tensor");
    TORCH_CHECK(A.sizes() == B.sizes(), "A and B must have the same size");

    auto C = torch::empty_like(A);
    int64_t N = A.numel();

    const int threads = 256;
    const int blocks = (N + threads - 1) / threads;

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(A.scalar_type(), "add_kernel", ([&] {
        add_kernel<scalar_t><<<blocks, threads>>>(
            A.data_ptr<scalar_t>(),
            B.data_ptr<scalar_t>(),
            C.data_ptr<scalar_t>(),
            N
        );
    }));

    auto err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }

    return C;
}
"""

add_cpp_source = """
#include <torch/extension.h>

torch::Tensor add_cuda(torch::Tensor A, torch::Tensor B);
"""

add_module = load_inline(
    name="add_cuda",
    cpp_sources=[add_cpp_source],
    cuda_sources=[add_cuda_source],
    functions=["add_cuda"],      # expects add_cuda(Tensor, Tensor) -> Tensor; must be here
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=True,               # keep True while debugging; can be False later
)


def custom_kernel(data: input_t) -> output_t:
    """
    Vector add kernel, CUDA inline version.

    Eval passes 3 tensors: (A, B, C), like in the Triton reference.
    We only *need* A and B: we compute C_out = A + B and return it.
    """
    # Accept both (A,B) and (A,B,C) just to be defensive
    if len(data) == 2:
        A, B = data
    else:
        A, B, _C_in = data

    # Same semantics as Triton: work on CUDA tensors in place-ish,
    # but our implementation just returns the result.
    assert A.is_cuda and B.is_cuda, "Input tensors must live on GPU"
    assert A.shape == B.shape, "A and B must have the same shape"
    # Typicky float16 pro challenge, ale necháme to flexibilní
    # assert A.dtype == torch.float16 and B.dtype == torch.float16

    C_out = add_module.add_cuda(A, B)
    return C_out