submission 66233

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

hip_source = """
#include <cuda_runtime.h>
#include <iostream>
#include <pybind11/buffer_info.h>
#include <pybind11/pybind11.h>
#include <stdexcept>
#include <torch/extension.h>
#include <cuda_fp16.h>
namespace py = pybind11;

__global__ void add_kernel(const __half *__restrict a, const __half *__restrict b,
                           __half *__restrict c, int N) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < N) {
    c[i] = __hadd(a[i],b[i]);
  }
}

void add_vectors(uintptr_t A, uintptr_t B, uintptr_t C, int N) {

  __half *A_ptr = reinterpret_cast<__half *>(A);
  __half *B_ptr = reinterpret_cast<__half *>(B);
  __half *C_ptr = reinterpret_cast<__half *>(C);

  int tpb = 256;
  dim3 threadPerBlock(tpb);
  dim3 numBlocks((N + tpb - 1) / tpb);
  add_kernel<<<numBlocks, threadPerBlock>>>(A_ptr, B_ptr, C_ptr, N);
  cudaDeviceSynchronize();
}

"""


# --- C++ Function Declaration (Header) ---
# This string provides the function prototype to the C++ compiler (for the bindings).
cpp_source_declaration = """
#include <torch/extension.h>
#include <pybind11/pybind11.h>

namespace py = pybind11;
void add_vectors(uintptr_t A, uintptr_t B, uintptr_t C, int N);
"""

# # Define paths and flags matching the setup
# extra_include_paths = [
#     "-I/opt/nvidia/hpc_sdk/Linux_x86_64/25.7/cuda/include/",
#     "-I/opt/nvidia/hpc_sdk/Linux_x86_64/25.7/comm_libs/nccl/include/",
# ]
# extra_library_paths = [
#     "-L/opt/nvidia/hpc_sdk/Linux_x86_64/25.7/cuda/lib64/",
#     "-L/opt/nvidia/hpc_sdk/Linux_x86_64/25.7/comm_libs/12.9/nccl/lib/",
#     "-L/opt/nvidia/hpc_sdk/Linux_x86_64/25.7/math_libs/12.9/lib64/",
#     "-lnccl",
#     "-lcublas",
#     "-lcudart",
# ]
extra_cflags = ["-O2"]
extra_cuda_cflags = ["-O2"]

# Load the inline extension
add_kernel = load_inline(
    name="add_kernel",
    verbose=True,
    cpp_sources=[cpp_source_declaration],
    cuda_sources=[hip_source],
    functions=[
        "add_vectors",
    ],
    extra_cflags=extra_cflags,
    extra_cuda_cflags=extra_cuda_cflags,
    # extra_include_paths=extra_include_paths,
    # extra_ldflags=extra_library_paths,
)


def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of vector addition using PyTorch.
    Args:
        data: Tuple of tensors [A, B] to be added.
    Returns:
        Tensor containing element-wise sums.
    """

    with DeterministicContext():
        A, B, output = data
        N = A.shape[0] ** 2
        add_kernel.add_vectors(A.data_ptr(), B.data_ptr(), output.data_ptr(), N)
        return output