| 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, | |
| ) | |