Add the supported CUDA compute capabilities into the toolchain created for manylinux2010 compatibility.

PiperOrigin-RevId: 262473631
This commit is contained in:
Amit Patankar 2019-08-08 18:14:44 -07:00 committed by Goldie Gadde
parent 8b21ae572e
commit e0b6dd6c7c
2 changed files with 14 additions and 2 deletions

View File

@ -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('.', '')

View File

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