# coding=utf-8

from torch.utils.cpp_extension import BuildExtension, CUDAExtension
from setuptools import setup


cuda_exts = []
files = [
    "csrc/core.cu",
    "csrc/matrix.cu",
    "csrc/core_bind.cpp",
    "csrc/max.cu",
    "csrc/md.cu",
]
extension = CUDAExtension(
    name="torch_cuda_ext.core",
    sources=files,
    extra_compile_args={"cxx": ["-g"], "nvcc": ["-O2"]},
    include_dirs=["/home/squall/program/cutlass/include"],
)

cuda_exts.append(extension)
setup(
    name="torch_cuda_ext",
    version="1.0",
    ext_modules=cuda_exts,
    cmdclass={"build_ext": BuildExtension},
)