Spaces:
Running
Running
import os | |
from setuptools import setup | |
from torch.utils.cpp_extension import BuildExtension, CUDAExtension | |
_src_path = os.path.dirname(os.path.abspath(__file__)) | |
nvcc_flags = [ | |
'-O3', | |
# '-std=c++14', | |
'-std=c++17', | |
'-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__', | |
] | |
if os.name == "posix": | |
# c_flags = ['-O3', '-std=c++14'] | |
c_flags = ['-O3', '-std=c++17'] | |
elif os.name == "nt": | |
c_flags = ['/O2', '/std:c++17'] | |
# find cl.exe | |
def find_cl_path(): | |
import glob | |
for edition in ["Enterprise", "Professional", "BuildTools", "Community"]: | |
paths = sorted(glob.glob(r"C:\\Program Files (x86)\\Microsoft Visual Studio\\*\\%s\\VC\\Tools\\MSVC\\*\\bin\\Hostx64\\x64" % edition), reverse=True) | |
if paths: | |
return paths[0] | |
# If cl.exe is not on path, try to find it. | |
if os.system("where cl.exe >nul 2>nul") != 0: | |
cl_path = find_cl_path() | |
if cl_path is None: | |
raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation") | |
os.environ["PATH"] += ";" + cl_path | |
setup( | |
name='shencoder', # package name, import this to use python API | |
ext_modules=[ | |
CUDAExtension( | |
name='_shencoder', # extension name, import this to use CUDA API | |
sources=[os.path.join(_src_path, 'src', f) for f in [ | |
'shencoder.cu', | |
'bindings.cpp', | |
]], | |
extra_compile_args={ | |
'cxx': c_flags, | |
'nvcc': nvcc_flags, | |
} | |
), | |
], | |
cmdclass={ | |
'build_ext': BuildExtension, | |
} | |
) |