1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
|
import sys
import torch.cuda
import os
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
from torch.testing._internal.common_utils import IS_WINDOWS
if sys.platform == 'win32':
vc_version = os.getenv('VCToolsVersion', '')
if vc_version.startswith('14.16.'):
CXX_FLAGS = ['/sdl']
else:
CXX_FLAGS = ['/sdl', '/permissive-']
else:
CXX_FLAGS = ['-g']
USE_NINJA = os.getenv('USE_NINJA') == '1'
ext_modules = [
CppExtension(
'torch_test_cpp_extension.cpp', ['extension.cpp'],
extra_compile_args=CXX_FLAGS),
CppExtension(
'torch_test_cpp_extension.ort', ['ort_extension.cpp'],
extra_compile_args=CXX_FLAGS),
CppExtension(
'torch_test_cpp_extension.rng', ['rng_extension.cpp'],
extra_compile_args=CXX_FLAGS),
]
if torch.cuda.is_available() and (CUDA_HOME is not None or ROCM_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)
if torch.cuda.is_available() and (CUDA_HOME is not None or ROCM_HOME is not None):
extension = CUDAExtension(
'torch_test_cpp_extension.torch_library', [
'torch_library.cu'
],
extra_compile_args={'cxx': CXX_FLAGS,
'nvcc': ['-O2']})
ext_modules.append(extension)
# todo(mkozuki): Figure out the root cause
if (not IS_WINDOWS) and torch.cuda.is_available() and CUDA_HOME is not None:
# malfet: One shoudl not assume that PyTorch re-exports CUDA dependencies
cublas_extension = CUDAExtension(
name='torch_test_cpp_extension.cublas_extension',
sources=['cublas_extension.cpp'],
libraries=['cublas'] if torch.version.hip is None else [],
)
ext_modules.append(cublas_extension)
cusolver_extension = CUDAExtension(
name='torch_test_cpp_extension.cusolver_extension',
sources=['cusolver_extension.cpp'],
libraries=['cusolver'] if torch.version.hip is None else [],
)
ext_modules.append(cusolver_extension)
if USE_NINJA and (not IS_WINDOWS) and torch.cuda.is_available() and CUDA_HOME is not None:
extension = CUDAExtension(
name='torch_test_cpp_extension.cuda_dlink',
sources=[
'cuda_dlink_extension.cpp',
'cuda_dlink_extension_kernel.cu',
'cuda_dlink_extension_add.cu',
],
dlink=True,
extra_compile_args={'cxx': CXX_FLAGS,
'nvcc': ['-O2', '-dc']})
ext_modules.append(extension)
setup(
name='torch_test_cpp_extension',
packages=['torch_test_cpp_extension'],
ext_modules=ext_modules,
include_dirs='self_compiler_include_dirs_test',
cmdclass={'build_ext': BuildExtension.with_options(use_ninja=USE_NINJA)})
|
