#!POPCORN leaderboard amd-identity
import torch
from torch.utils.cpp_extension import load_inline
from typing import List
from task import input_t, output_t

add_cuda_source = """
template<typename I, typename J>
auto cdiv(I x, J y){
    return (x + y - 1) / y;
}

template <typename scalar_t>
__global__ void copy_kern(const scalar_t* __restrict__ src, 
                           scalar_t* __restrict__ dst, 
                           int N) {


    for (uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < N; idx += blockDim.x * gridDim.x){
        dst[idx] = src[idx];
    }

}

void copy_f(torch::Tensor src, torch::Tensor dst) {
    
    int N = src.numel();  
    const int threads = 1024; 
    const int duplicity = 2;
    const int blocks = cdiv(N, threads * duplicity);
    
    AT_DISPATCH_FLOATING_TYPES_AND_HALF(src.scalar_type(), "copy_f", ([&] {
        copy_kern<scalar_t><<<blocks, threads>>>(
            src.data_ptr<scalar_t>(),
            dst.data_ptr<scalar_t>(),
            N
        );
    }));
}
"""

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

void copy_f(torch::Tensor src, torch::Tensor dst);
"""

torch_ops = load_inline(
    name="torch_ops",
    cpp_sources=add_cpp_source,
    cuda_sources=add_cuda_source,
    functions=["copy_f"],
    verbose=True,
)


def custom_kernel(data: input_t) -> output_t:
    input, output = data
    torch_ops.copy_f(input, output)
    return output


if __name__ == "__main__":
    x = torch.randn(123124, device="cuda")
    y = torch.randn(123124, device="cuda")
    custom_kernel((x, y))
    assert (x == y).all()

python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/torch_ops/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/torch_ops/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/torch_ops/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/torch_ops/hip.hip [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=torch_ops -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/torch_ops/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=torch_ops -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx900 --offload-arch=gfx906 --offload-arch=gfx908 --offload-arch=gfx90a --offload-arch=gfx942 --offload-arch=gfx1030 --offload-arch=gfx1100 --offload-arch=gfx1101 --offload-arch=gfx1102 --offload-arch=gfx1200 --offload-arch=gfx1201 -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/torch_ops/hip.hip -o hip.cuda.o 
[3/3] clang++ hip.cuda.o main.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o torch_ops.so

python eval.py test /tmp/tmp10208b7v
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/torch_ops/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/torch_ops/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/torch_ops/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/torch_ops/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
ninja: no work to do.
