Update caffe package for cuda9 (#24831)
Add base CudaPackage, cuda_arch
This commit is contained in:
parent
f58b2e03ca
commit
81bad21d3a
1 changed files with 8 additions and 2 deletions
|
@ -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
|
||||
|
|
Loading…
Reference in a new issue