2024-11-16 19:26:54 +08:00
|
|
|
# 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",
|
2024-11-18 19:54:12 +08:00
|
|
|
"csrc/md.cu",
|
2024-11-16 19:26:54 +08:00
|
|
|
]
|
|
|
|
|
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},
|
|
|
|
|
)
|
