submission 515731

#!POPCORN leaderboard vectoradd_v2
#!POPCORN gpu A100
import subprocess
import sys
subprocess.check_call([sys.executable, "-m", "pip", "install", "numba"])

import numba
from numba import cuda
import math
import torch

@cuda.jit
def vec_add_fp16_kernel(A, B, C, SIZE):
    idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    if idx < SIZE:
        C[idx] = A[idx] + B[idx]

def custom_kernel(data):
    A, B, C = data
    SIZE = A.numel()
    A_flat = cuda.as_cuda_array(A.view(-1))
    B_flat = cuda.as_cuda_array(B.view(-1))
    C_flat = cuda.as_cuda_array(C.view(-1))
    threads_per_block = 1024
    blocks = math.ceil(SIZE / threads_per_block)
    vec_add_fp16_kernel[blocks, threads_per_block](A_flat, B_flat, C_flat, SIZE)
    return C