File size: 1,176 Bytes
48b5986 | 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 | import os
from torch.utils.cpp_extension import load
def compile(name, sources, extra_include_paths, build_directory):
os.environ["TORCH_CUDA_ARCH_LIST"] = "7.0;8.0"
return load(
name=name,
sources=sources,
extra_include_paths=extra_include_paths,
extra_cflags=[
"-O3",
"-DVERSION_GE_1_1",
"-DVERSION_GE_1_3",
"-DVERSION_GE_1_5",
],
extra_cuda_cflags=[
"-O3",
"--use_fast_math",
"-DVERSION_GE_1_1",
"-DVERSION_GE_1_3",
"-DVERSION_GE_1_5",
"-std=c++17",
"-maxrregcount=50",
"-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__",
"--expt-relaxed-constexpr",
"--expt-extended-lambda",
"-gencode",
"arch=compute_70,code=sm_70",
"-gencode",
"arch=compute_80,code=sm_80",
"-gencode",
"arch=compute_86,code=sm_86",
"-gencode",
"arch=compute_90,code=sm_90",
],
verbose=True,
build_directory=build_directory,
)
|