From e0b6dd6c7ce5b3dc7cc9d02ad48a4548755b6ab0 Mon Sep 17 00:00:00 2001 From: Amit Patankar Date: Thu, 8 Aug 2019 18:14:44 -0700 Subject: [PATCH] Add the supported CUDA compute capabilities into the toolchain created for manylinux2010 compatibility. PiperOrigin-RevId: 262473631 --- .../clang/bin/crosstool_wrapper_driver_is_not_gcc | 7 ++++++- .../windows/msvc_wrapper_for_nvcc.py | 9 ++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/clang/bin/crosstool_wrapper_driver_is_not_gcc b/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/clang/bin/crosstool_wrapper_driver_is_not_gcc index 9800b7689a3..1243dbbefa0 100755 --- a/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/clang/bin/crosstool_wrapper_driver_is_not_gcc +++ b/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/clang/bin/crosstool_wrapper_driver_is_not_gcc @@ -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('.', '') diff --git a/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/windows/msvc_wrapper_for_nvcc.py b/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/windows/msvc_wrapper_for_nvcc.py index 79b98e587e3..a69d47f250f 100755 --- a/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/windows/msvc_wrapper_for_nvcc.py +++ b/third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.0/windows/msvc_wrapper_for_nvcc.py @@ -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))