Add the supported CUDA compute capabilities into the toolchain created for manylinux2010 compatibility.
PiperOrigin-RevId: 262473631
This commit is contained in:
parent
8b21ae572e
commit
e0b6dd6c7c
@ -53,6 +53,11 @@ NVCC_PATH = '/usr/local/cuda-10.0/bin/nvcc'
|
||||
PREFIX_DIR = os.path.dirname(GCC_HOST_COMPILER_PATH)
|
||||
NVCC_VERSION = '10.0'
|
||||
|
||||
# Environment variable for supported TF CUDA Compute Capabilities
|
||||
# eg. export TF_CUDA_COMPUTE_CAPABILITIES=3.5,3.7,5.2,6.0,6.1,7.0
|
||||
CUDA_COMPUTE_ENV_VAR = 'TF_CUDA_COMPUTE_CAPABILITIES'
|
||||
DEFAULT_CUDA_COMPUTE_CAPABILITIES = '3.5,6.0'
|
||||
|
||||
def Log(s):
|
||||
print('gpus/crosstool: {0}'.format(s))
|
||||
|
||||
@ -202,7 +207,7 @@ def InvokeNvcc(argv, log=False):
|
||||
srcs = ' '.join(src_files)
|
||||
out = ' -o ' + out_file[0]
|
||||
|
||||
supported_cuda_compute_capabilities = [ "3.0", "6.0" ]
|
||||
supported_cuda_compute_capabilities = os.environ.get(CUDA_COMPUTE_ENV_VAR, DEFAULT_CUDA_COMPUTE_CAPABILITIES).split(',')
|
||||
nvccopts = '-D_FORCE_INLINES '
|
||||
for capability in supported_cuda_compute_capabilities:
|
||||
capability = capability.replace('.', '')
|
||||
|
@ -36,7 +36,14 @@ GCC_HOST_COMPILER_PATH = ('/dt7/usr/bin/gcc')
|
||||
NVCC_PATH = '/usr/local/cuda-10.0/bin/nvcc'
|
||||
NVCC_VERSION = '10.0'
|
||||
NVCC_TEMP_DIR = "C:\\Windows\\Temp\\nvcc_inter_files_tmp_dir"
|
||||
supported_cuda_compute_capabilities = [ "3.0", "6.0" ]
|
||||
DEFAULT_CUDA_COMPUTE_CAPABILITIES = '3.5,6.0'
|
||||
|
||||
# Taken from environment variable for supported TF CUDA Compute Capabilities
|
||||
# eg. export TF_CUDA_COMPUTE_CAPABILITIES=3.5,3.7,5.2,6.0,6.1,7.0
|
||||
supported_cuda_compute_capabilities = os.environ.get(
|
||||
'TF_CUDA_COMPUTE_CAPABILITIES',
|
||||
DEFAULT_CUDA_COMPUTE_CAPABILITIES).split(',')
|
||||
|
||||
|
||||
def Log(s):
|
||||
print('gpus/crosstool: {0}'.format(s))
|
||||
|
Loading…
x
Reference in New Issue
Block a user