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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
|
import os
import sys
from setuptools import setup
import torch.cuda
from torch.testing._internal.common_utils import IS_WINDOWS
from torch.utils.cpp_extension import (
BuildExtension,
CppExtension,
CUDA_HOME,
CUDAExtension,
ROCM_HOME,
)
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.maia",
["maia_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)
if torch.backends.mps.is_available():
extension = CppExtension(
"torch_test_cpp_extension.mps",
["mps_extension.mm"],
extra_compile_args=CXX_FLAGS,
)
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 should 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)},
entry_points={
"torch.backends": [
"device_backend = torch_test_cpp_extension:_autoload",
],
},
)
|
