2023-02-16 07:47:03 +00:00
|
|
|
import setuptools
|
|
|
|
|
from torch.utils import cpp_extension
|
|
|
|
|
|
|
|
|
|
CXX_FLAGS = ['-g']
|
|
|
|
|
NVCC_FLAGS = ['-O2']
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
ext_modules = []
|
|
|
|
|
|
|
|
|
|
# Cache operations.
|
|
|
|
|
cache_extension = cpp_extension.CUDAExtension(
|
|
|
|
|
name='cacheflow.ops',
|
2023-02-16 20:05:45 +00:00
|
|
|
sources=['csrc/cache.cpp', 'csrc/cache_kernels.cu'],
|
2023-02-16 07:47:03 +00:00
|
|
|
extra_compile_args={'cxx': CXX_FLAGS, 'nvcc': NVCC_FLAGS},
|
|
|
|
|
)
|
|
|
|
|
ext_modules.append(cache_extension)
|
|
|
|
|
|
|
|
|
|
setuptools.setup(
|
|
|
|
|
name='cacheflow',
|
|
|
|
|
requires_python='>=3.9',
|
|
|
|
|
ext_modules=ext_modules,
|
|
|
|
|
cmdclass={'build_ext': cpp_extension.BuildExtension},
|
|
|
|
|
)
|
