File size: 1,774 Bytes
89c0b51 |
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 |
# Copyright 2024 ByteDance and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
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,
)
|