from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
import os
kernel_dir = os.path.dirname(os.path.abspath(__file__))


setup(
    name="my_model_kernel",
    version="0.0.0",
    ext_modules=[
        CUDAExtension(
            name="my_model_kernel",
            sources = [
                "./my_model_kernel.cpp",
                "./trilinear/src/trilinear_kernel.cu",
                "./trilinear/src/trilinear_cuda.cu",
                "./trilinear/src/trilinear_cpu.cpp",
                "./quadrilinear/src/quadrilinear4d_cuda.cpp",
                "./quadrilinear/src/quadrilinear4d_kernel.cu",
                "./quadrilinear/src/quadrilinear4d.cpp"
            ],
            include_dirs = [
                os.path.join(kernel_dir, "trilinear", "include"),
                os.path.join(kernel_dir, "quadrilinear", "include"),
            ],
            extra_compile_args = {
                "cxx": [],
                "nvcc": [ "-O3",
                          "-std=c++17",
                          "-Wno-deprecated-gpu-targets",
                          "-U__CUDA_NO_HALF_OPERATORS__",
                          "-U__CUDA_NO_HALF_CONVERSIONS__",
                          "-U__CUDA_NO_BFLOAT16_OPERATORS__",
                          "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                          "-U__CUDA_NO_BFLOAT162_OPERATORS__",
                          "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
                          "--expt-relaxed-constexpr",
                          "--expt-extended-lambda",
                          "--use_fast_math",
                          "--ptxas-options=-v",
                          "-lineinfo",
                          "--threads", "16", "-allow-unsupported-compiler"],
            }
        ),
    ],
    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
)