update py-cupy to enable ROCm builds and add variant to control optional dependencies (#38919)

* py-cupy updates: add +rocm and +all

* rocm deps are link only

* set parallelism for both +rocm and +cuda

* add missing deps; remove unnecessary deps; uncomment maintainers; get hipcc properly

3 files changed, 125 insertions, 15 deletions

diff --git a/var/spack/repos/builtin/packages/py-cmaes/package.py b/var/spack/repos/builtin/packages/py-cmaes/package.py
new file mode 100644
index 0000000000..26730ea43d
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-cmaes/package.py
@@ -0,0 +1,22 @@
+# Copyright 2013-2023 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack.package import *
+
+
+class PyCmaes(PythonPackage):
+    """Lightweight Covariance Matrix Adaptation Evolution Strategy (CMA-ES) implementation."""
+
+    homepage = "https://github.com/CyberAgentAILab/cmaes"
+    pypi = "cmaes/cmaes-0.10.0.tar.gz"
+
+    maintainers("eugeneswalker")
+
+    version("0.10.0", sha256="48afc70df027114739872b50489ae6b32461c307b92d084a63c7090a9742faf9")
+
+    depends_on("py-setuptools@61:", type="build")
+
+    depends_on("py-numpy", type=("build", "run"))
diff --git a/var/spack/repos/builtin/packages/py-cupy/package.py b/var/spack/repos/builtin/packages/py-cupy/package.py
index 2074632c23..feffd1dcf6 100644
--- a/var/spack/repos/builtin/packages/py-cupy/package.py
+++ b/var/spack/repos/builtin/packages/py-cupy/package.py
@@ -3,10 +3,12 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+import os
+
 from spack.package import *
 
 
-class PyCupy(PythonPackage, CudaPackage):
+class PyCupy(PythonPackage, CudaPackage, ROCmPackage):
     """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,
@@ -16,13 +18,15 @@ class PyCupy(PythonPackage, CudaPackage):
     homepage = "https://cupy.dev/"
     pypi = "cupy/cupy-8.0.0.tar.gz"
 
-    version("11.2.0", sha256="c33361f117a347a63f6996ea97446d17f1c038f1a1f533e502464235076923e2")
-    version("11.3.0", sha256="d057cc2f73ecca06fae8b9c270d9e14116203abfd211a704810cc50a453b4c9e")
-    version("11.4.0", sha256="03d52b2626e02a3a2b46d714c1cd03e702c8fe33915fcca6ed8de5c539964f49")
-    version("11.5.0", sha256="4bc8565bded22cc89b210fd9fb48a5d5316f30701e12bb23852a60314e1f9f6e")
-    version("11.6.0", sha256="53dbb840072bb32d4bfbaa6bfa072365a30c98b1fcd1f43e48969071ad98f1a7")
-    version("12.0.0", sha256="61ddbbef73d50d606bd5087570645f3c91ec9176c2566784c1d486d6a3404545")
     version("12.1.0", sha256="f6d31989cdb2d96581da12822e28b102f29e254427195c2017eac327869b7320")
+    version("12.0.0", sha256="61ddbbef73d50d606bd5087570645f3c91ec9176c2566784c1d486d6a3404545")
+    version("11.6.0", sha256="53dbb840072bb32d4bfbaa6bfa072365a30c98b1fcd1f43e48969071ad98f1a7")
+    version("11.5.0", sha256="4bc8565bded22cc89b210fd9fb48a5d5316f30701e12bb23852a60314e1f9f6e")
+    version("11.4.0", sha256="03d52b2626e02a3a2b46d714c1cd03e702c8fe33915fcca6ed8de5c539964f49")
+    version("11.3.0", sha256="d057cc2f73ecca06fae8b9c270d9e14116203abfd211a704810cc50a453b4c9e")
+    version("11.2.0", sha256="c33361f117a347a63f6996ea97446d17f1c038f1a1f533e502464235076923e2")
+
+    variant("all", default=False, description="Enable optional py-scipy, optuna, and cython")
 
     depends_on("python@3.7:", when="@:11", type=("build", "run"))
     depends_on("python@3.8:", when="@12:", type=("build", "run"))
@@ -31,18 +35,69 @@ class PyCupy(PythonPackage, CudaPackage):
     depends_on("py-fastrlock@0.5:", type=("build", "run"))
     depends_on("py-numpy@1.20:1.25", when="@:11", type=("build", "run"))
     depends_on("py-numpy@1.20:1.26", when="@12:", type=("build", "run"))
-    depends_on("py-scipy@1.6:1.12", type=("build", "run"))
-    depends_on("cuda@:11.9", when="@:11")
-    depends_on("cuda@:12.1", when="@12:")
-    depends_on("nccl")
-    depends_on("cudnn")
-    depends_on("cutensor")
 
-    conflicts("~cuda")
+    depends_on("py-scipy@1.6:1.12", when="+all", type=("build", "run"))
+    depends_on("py-cython@0.29.22:2", when="+all", type=("build", "run"))
+    depends_on("py-optuna@2:", when="+all", type=("build", "run"))
+
+    # Based on https://github.com/cupy/cupy/releases
+    depends_on("cuda@:11.9", when="@:11 +cuda")
+    depends_on("cuda@:12.1", when="@12: +cuda")
+
+    for a in CudaPackage.cuda_arch_values:
+        depends_on("nccl +cuda cuda_arch={0}".format(a), when="+cuda cuda_arch={0}".format(a))
+
+    depends_on("cudnn", when="+cuda")
+    depends_on("cutensor", when="+cuda")
+
+    for _arch in ROCmPackage.amdgpu_targets:
+        arch_str = "amdgpu_target={0}".format(_arch)
+        rocm_str = "+rocm {0}".format(arch_str)
+        depends_on("rocprim {0}".format(arch_str), when=rocm_str, type=("link"))
+        depends_on("rocsolver {0}".format(arch_str), when=rocm_str, type=("link"))
+        depends_on("rocthrust {0}".format(arch_str), when=rocm_str, type=("link"))
+        depends_on("rocrand {0}".format(arch_str), when=rocm_str, type=("link"))
+        depends_on("hipcub {0}".format(rocm_str), when=rocm_str, type=("link"))
+        depends_on("hipblas {0}".format(rocm_str), when=rocm_str, type=("link"))
+        depends_on("hiprand {0}".format(rocm_str), when=rocm_str, type=("link"))
+        depends_on("hipsparse {0}".format(rocm_str), when=rocm_str, type=("link"))
+        depends_on("hipfft {0}".format(rocm_str), when=rocm_str, type=("link"))
+
+    depends_on("rccl", when="+rocm", type=("link"))
+    depends_on("roctracer-dev", when="+rocm", type=("link"))
+    depends_on("rocprofiler-dev", when="+rocm", type=("link"))
+
+    conflicts("~cuda ~rocm")
+    conflicts("+cuda +rocm")
+    conflicts("+cuda cuda_arch=none")
 
     def setup_build_environment(self, env):
         env.set("CUPY_NUM_BUILD_JOBS", make_jobs)
-        if not self.spec.satisfies("cuda_arch=none"):
+        if self.spec.satisfies("+cuda"):
             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)
+        elif self.spec.satisfies("+rocm"):
+            spec = self.spec
+
+            incs = {
+                "roctracer-dev": ["include/roctracer"],
+                "hiprand": ["include"],
+                "rocrand": ["include"],
+                "rocthrust": ["include"],
+                "rocprim": ["include"],
+                "hip": ["include", "include/hip"],
+            }
+
+            inc_dirs = []
+            for pkg, ds in incs.items():
+                for d in ds:
+                    p = os.path.join(spec[pkg].prefix, d)
+                    if os.path.exists(p):
+                        inc_dirs.append(p)
+
+            env.set("CUPY_INCLUDE_PATH", ":".join(inc_dirs))
+
+            env.set("HIPCC", self.spec["hip"].hipcc)
+            env.set("ROCM_HOME", self.spec["hipcub"].prefix)
+            env.set("CUPY_INSTALL_USE_HIP", 1)
diff --git a/var/spack/repos/builtin/packages/py-optuna/package.py b/var/spack/repos/builtin/packages/py-optuna/package.py
new file mode 100644
index 0000000000..de74a405ec
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-optuna/package.py
@@ -0,0 +1,33 @@
+# Copyright 2013-2023 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack.package import *
+
+
+class PyOptuna(PythonPackage):
+    """Optuna is an automatic hyperparameter optimization software framework,
+    particularly designed for machine learning. It features an imperative,
+    define-by-run style user API. Thanks to our define-by-run API, the code
+    written with Optuna enjoys high modularity, and the user of Optuna can
+    dynamically construct the search spaces for the hyperparameters."""
+
+    homepage = "https://optuna.org/"
+    pypi = "optuna/optuna-3.2.0.tar.gz"
+
+    maintainers("elliottslaughter", "eugeneswalker")
+
+    version("3.2.0", sha256="683d8693643a761a41d251a6b8e13263b24acacf9fc46a9233d5f6aa3ce5c683")
+
+    depends_on("py-setuptools@61.1:", type="build")
+
+    depends_on("py-alembic@1.5:", type=("build", "run"))
+    depends_on("py-cmaes@0.9.1:", type=("build", "run"))
+    depends_on("py-colorlog", type=("build", "run"))
+    depends_on("py-numpy", type=("build", "run"))
+    depends_on("py-packaging@20:", type=("build", "run"))
+    depends_on("py-pyyaml", type=("build", "run"))
+    depends_on("py-sqlalchemy@1.3:", type=("build", "run"))
+    depends_on("py-tqdm", type=("build", "run"))