mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-08 07:39:33 +01:00
e870a9a870
Summary: The flags `/sdl` and `/permissive-` are switched on automatically when using the VS GUI. Adding those checks will ensure that those annoying errors won't appear when users use the VS GUI to build their project. More info: https://docs.microsoft.com/en-us/cpp/build/reference/sdl-enable-additional-security-checks?view=vs-2017 https://docs.microsoft.com/en-us/cpp/build/reference/permissive-standards-conformance?view=vs-2017 Pull Request resolved: https://github.com/pytorch/pytorch/pull/29709 Differential Revision: D18473888 Pulled By: bddppq fbshipit-source-id: 21156b0232a5dc3b566d14491d00bacb11493254
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
import sys
|
|
import torch.cuda
|
|
from setuptools import setup
|
|
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
|
|
from torch.utils.cpp_extension import CUDA_HOME
|
|
|
|
CXX_FLAGS = ['/sdl', '/permissive-'] if sys.platform == 'win32' else ['-g', '-Werror']
|
|
|
|
ext_modules = [
|
|
CppExtension(
|
|
'torch_test_cpp_extension.cpp', ['extension.cpp'],
|
|
extra_compile_args=CXX_FLAGS),
|
|
CppExtension(
|
|
'torch_test_cpp_extension.msnpu', ['msnpu_extension.cpp'],
|
|
extra_compile_args=CXX_FLAGS),
|
|
]
|
|
|
|
if torch.cuda.is_available() and CUDA_HOME is not None:
|
|
extension = CUDAExtension(
|
|
'torch_test_cpp_extension.cuda', [
|
|
'cuda_extension.cpp',
|
|
'cuda_extension_kernel.cu',
|
|
'cuda_extension_kernel2.cu',
|
|
],
|
|
extra_compile_args={'cxx': CXX_FLAGS,
|
|
'nvcc': ['-O2']})
|
|
ext_modules.append(extension)
|
|
|
|
setup(
|
|
name='torch_test_cpp_extension',
|
|
packages=['torch_test_cpp_extension'],
|
|
ext_modules=ext_modules,
|
|
cmdclass={'build_ext': BuildExtension})
|