Update caffe package for cuda9 (#24831)

Add base CudaPackage, cuda_arch
This commit is contained in:
jkelling 2021-07-13 00:31:38 +02:00 committed by GitHub
parent f58b2e03ca
commit 81bad21d3a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -6,7 +6,7 @@
from spack import *
class Caffe(CMakePackage):
class Caffe(CMakePackage, CudaPackage):
"""Caffe is a deep learning framework made with expression, speed, and
modularity in mind. It is developed by the Berkeley Vision and Learning
Center (BVLC) and by community contributors."""
@ -43,7 +43,7 @@ class Caffe(CMakePackage):
depends_on('hdf5 +hl +cxx')
# Optional dependencies
depends_on('opencv@3.2.0+core+highgui+imgproc', when='+opencv')
depends_on('opencv@3.2.0:3.4.12+core+highgui+imgproc+imgcodecs', when='+opencv')
depends_on('leveldb', when='+leveldb')
depends_on('lmdb', when='+lmdb')
depends_on('python@2.7:', when='+python')
@ -87,4 +87,10 @@ def cmake_args(self):
'-DCMAKE_CXX_COMPILER={0}'.format(self.spec['mpi'].mpicxx)
])
if '+cuda' in spec:
if spec.variants['cuda_arch'].value[0] != 'none':
cuda_arch = spec.variants['cuda_arch'].value
args.append(self.define('CUDA_ARCH_NAME', 'Manual'))
args.append(self.define('CUDA_ARCH_BIN', ' '.join(cuda_arch)))
return args