zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

py-cupy allow customizing architecture and threads (#37072)

Browse Source 
* py-cupy allow customizing architecture and threads

* Update var/spack/repos/builtin/packages/py-cupy/package.py

Co-authored-by: Massimiliano Culpo <massimiliano.culpo@gmail.com>

* add missing self

---------

Co-authored-by: Robert Underwood <runderwood@anl.gov>
Co-authored-by: Massimiliano Culpo <massimiliano.culpo@gmail.com>
This commit is contained in:
Robert Underwood 2023-04-21 10:52:47 -04:00 committed by GitHub
parent aaac8b0545
commit ff689be250
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
var/spack/repos/builtin/packages/py-cupy/package.py
View File

@ -6,7 +6,7 @@
from spack.package import *
class PyCupy(PythonPackage):
class PyCupy(PythonPackage, CudaPackage):
"""CuPy is an open-source array library accelerated with
NVIDIA CUDA. CuPy provides GPU accelerated computing with
Python. CuPy uses CUDA-related libraries including cuBLAS,
@ -32,3 +32,12 @@ class PyCupy(PythonPackage):
depends_on("nccl")
depends_on("cudnn")
depends_on("cutensor")
conflicts("~cuda")
def setup_build_environment(self, env):
env.set("CUPY_NUM_BUILD_JOBS", make_jobs)
if not self.spec.satisfies("cuda_arch=none"):
cuda_arch = self.spec.variants["cuda_arch"].value
arch_str = ";".join("arch=compute_{0},code=sm_{0}".format(i) for i in cuda_arch)
env.set("CUPY_NVCC_GENERATE_CODE", arch_str)