hlrs-software-stack/spack
7
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

TensorFlow: must specify CUDA compute capabilities (#14118)

Browse source 
* TensorFlow: must specify CUDA compute capabilities

* Correctly restrict cuda_arch version ranges
This commit is contained in:
Adam J. Stewart 2019-12-11 20:16:52 -06:00 committed by GitHub
parent 7ec1917d9b
commit 1f9d37c346
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
var/spack/repos/builtin/packages/py-tensorflow/package.py
View file

@ -205,6 +205,11 @@ class PyTensorflow(Package, CudaPackage):
conflicts('+computecpp', when='~opencl')
conflicts('+rocm', when='@:1.11')
conflicts('+cuda', when='platform=darwin', msg='There is no GPU support for macOS')
conflicts('cuda_arch=none', when='+cuda', msg='Must specify CUDA compute capabilities of your GPU, see https://developer.nvidia.com/cuda-gpus')
conflicts('cuda_arch=20', when='@1.12.1,1.14:', msg='TensorFlow only supports compute capabilities >= 3.5')
conflicts('cuda_arch=30', when='@1.12.1,1.14:', msg='TensorFlow only supports compute capabilities >= 3.5')
conflicts('cuda_arch=32', when='@1.12.1,1.14:', msg='TensorFlow only supports compute capabilities >= 3.5')
conflicts('cuda_arch=20', when='@1.4:1.12.0,1.12.2:1.12.3', msg='Only compute capabilities 3.0 or higher are supported')
conflicts('+tensorrt', when='@:1.5')
conflicts('+tensorrt', when='~cuda')
conflicts('+tensorrt', when='platform=darwin', msg='Currently TensorRT is only supported on Linux platform')
@ -423,10 +428,9 @@ def setup_build_environment(self, env):
# Please note that each additional compute capability significantly
# increases your build time and binary size, and that TensorFlow
# only supports compute capabilities >= 3.5
if spec.variants['cuda_arch'].value != 'none':
capabilities = ','.join('{0:.1f}'.format(
float(i) / 10.0) for i in spec.variants['cuda_arch'].value)
env.set('TF_CUDA_COMPUTE_CAPABILITIES', capabilities)
capabilities = ','.join('{0:.1f}'.format(
float(i) / 10.0) for i in spec.variants['cuda_arch'].value)
env.set('TF_CUDA_COMPUTE_CAPABILITIES', capabilities)
else:
env.set('TF_NEED_CUDA', '0')