submission 191358

#!POPCORN leaderboard vectoradd_v2

# This is a submission template for popcorn leaderboard 'vectoradd_v2'.
# Your task is as follows:
# > Implement a float16 vector addition kernel.
# > 
# > Input: tuple(torch.Tensor, torch.Tensor) with tensors of shape (N, N) and type torch.float16. These tensors are from
# > a normal distribution with mean 0 and variance 1.
# > Output: torch.Tensor of shape (N, N) and type torch.float16
# The deadline for this leaderboard is 2025-12-30 00:00:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

import subprocess
import torch

try:
    import cuda.compute
except ImportError:
    if "L4" in torch.cuda.get_device_name():
        # Prefetch seems to be better on L4, so I'm using the potential-submissions branch
        import os
        subprocess.check_call(["rm", "-rf", "cccl"])
        subprocess.check_call(["git", "clone", "--depth", "1", "--branch", "potential-submissions", "https://github.com/NaderAlAwar/cccl.git"])
        subprocess.check_call(["git", "checkout", "83ad53b"], cwd="cccl")

        env = os.environ.copy()
        env["CC"] = "gcc"
        env["CXX"] = "g++"
        env["CMAKE_ARGS"] = "-DCMAKE_CXX_STANDARD=20 -DCMAKE_CUDA_STANDARD=20"

        subprocess.check_call(["pip", "install", ".[cu12]"], cwd="cccl/python/cuda_cccl", env=env)
    elif "H100" in torch.cuda.get_device_name() or "B200" in torch.cuda.get_device_name() or "A100" in torch.cuda.get_device_name():
        import os
        # subprocess.check_call(["rm", "-rf", "cccl"])
        subprocess.run(["git", "clone", "--depth", "2", "--branch", "potential-submissions", "https://github.com/NaderAlAwar/cccl.git"])
        subprocess.check_call(["git", "checkout", "e063657650987d055d1765a81395b0a4f2302b4e"], cwd="cccl")

        env = os.environ.copy()
        env["CC"] = "gcc"
        env["CXX"] = "g++"
        env["CMAKE_ARGS"] = "-DCMAKE_CXX_STANDARD=20 -DCMAKE_CUDA_STANDARD=20"

        subprocess.check_call(["pip", "install", "-e", ".[cu12]"], cwd="cccl/python/cuda_cccl", env=env)
    else:
        subprocess.run(["pip", "install", "cuda-cccl[cu12]==0.4.3"])

from task import input_t, output_t

import cuda.compute
from cuda.compute import OpKind

build_A = torch.empty(2, 2, dtype=torch.float16, device="cuda")
build_B = torch.empty(2, 2, dtype=torch.float16, device="cuda")
build_output = torch.empty(2, 2, dtype=torch.float16, device="cuda")

transformer = cuda.compute.make_binary_transform(build_A, build_B, build_output, OpKind.PLUS)

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