156 lines
7.0 KiB
Python
156 lines
7.0 KiB
Python
# Copyright 2013-2024 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 PyPennylaneLightningKokkos(CMakePackage, PythonExtension, CudaPackage, ROCmPackage):
|
|
"""The PennyLane-Lightning-Kokkos plugin provides
|
|
a fast state-vector simulator with Kokkos kernels."""
|
|
|
|
homepage = "https://docs.pennylane.ai/projects/lightning-kokkos"
|
|
git = "https://github.com/PennyLaneAI/pennylane-lightning-kokkos.git"
|
|
|
|
maintainers("AmintorDusko", "vincentmr")
|
|
|
|
license("Apache-2.0")
|
|
|
|
version("main", branch="main")
|
|
version("master", branch="master")
|
|
version("0.37.0", sha256="3f70e3e3b7e4d0f6a679919c0c83e451e129666b021bb529dd02eb915d0666a0")
|
|
version("0.36.0", sha256="c5fb24bdaf2ebdeaf614bfb3a8bcc07ee83c2c7251a3893399bb0c189d2d1d01")
|
|
version("0.35.1", sha256="d39a2749d08ef2ba336ed2d6f77b3bd5f6d1b25292263a41b97943ae7538b7da")
|
|
version("0.35.0", sha256="1a16fd3dbf03788e4f8dd510bbb668e7a7073ca62be4d9414e2c32e0166e8bda")
|
|
version("0.34.0", sha256="398c3a1d4450a9f3e146204c22329da9adc3f83a1685ae69187f3b25f47824c0")
|
|
version("0.33.1", sha256="878f63cd1afadd52386b1aca9c0e3fb0a097b64ce8e347b325ebc7cac722e5e0")
|
|
version("0.33.0", sha256="c4cab4a8a1a53edc0990a2a429805dca1c6a46a7ffcc6f77c985b88cd6ff6247")
|
|
version("0.32.0", sha256="06f19dfb1073387ef9ee30c38ea44884844a771373256b694a0e1ceb87195bb2")
|
|
version("0.31.0", sha256="fe10322fee0fa7df45cd3a81d6c229a79c7dfa7f20ff7d67c65c9a28f494dc89")
|
|
version("0.30.0", sha256="7c8f0e0431f8052993cd8033a316f53590c7bf5419445d0725e214b93cbc661b")
|
|
version("0.29.1", sha256="f51ba7718defc7bb5064f690f381e04b2ec58cb09f22a171ae5f410860716e30")
|
|
|
|
depends_on("cxx", type="build") # generated
|
|
|
|
depends_on("kokkos@:3.7.2", when="@:0.30", type=("run", "build"))
|
|
depends_on("kokkos@4:4.1", when="@0.31", type=("run", "build"))
|
|
depends_on("kokkos@4:4.2", when="@0.32:0.36", type=("run", "build"))
|
|
depends_on("kokkos@4.3:", when="@0.37:", type=("run", "build"))
|
|
|
|
# kokkos backends
|
|
backends = {
|
|
"cuda": [False, "Whether to build CUDA backend"],
|
|
"openmp": [False, "Whether to build OpenMP backend"],
|
|
"openmptarget": [False, "Whether to build the OpenMPTarget backend"],
|
|
"threads": [False, "Whether to build the C++ threads backend"],
|
|
"rocm": [False, "Whether to build HIP backend"],
|
|
"serial": [True, "Whether to build serial backend"],
|
|
"sycl": [False, "Whether to build the SYCL backend"],
|
|
}
|
|
|
|
for backend in backends:
|
|
deflt_bool, descr = backends[backend]
|
|
variant(backend.lower(), default=deflt_bool, description=descr)
|
|
depends_on(f"kokkos+{backend.lower()}", when=f"+{backend.lower()}", type=("run", "build"))
|
|
# CUDA
|
|
for val in CudaPackage.cuda_arch_values:
|
|
depends_on("kokkos cuda_arch={0}".format(val), when="cuda_arch={0}".format(val))
|
|
# Use +wrapper when not %clang %cce
|
|
depends_on("kokkos+wrapper", when="%gcc+cuda")
|
|
|
|
# ROCm
|
|
for val in ROCmPackage.amdgpu_targets:
|
|
depends_on("kokkos amdgpu_target={0}".format(val), when="amdgpu_target={0}".format(val))
|
|
|
|
conflicts(
|
|
"+cuda",
|
|
when="+rocm",
|
|
msg="CUDA and ROCm are not compatible in PennyLane-Lightning-Kokkos.",
|
|
)
|
|
|
|
# build options
|
|
extends("python")
|
|
variant("cpptests", default=False, description="Build CPP tests")
|
|
variant("native", default=False, description="Build natively for given hardware")
|
|
variant("sanitize", default=False, description="Build with address sanitization")
|
|
|
|
# hard dependencies
|
|
depends_on("cmake@3.20:", type="build")
|
|
depends_on("ninja", type="build")
|
|
depends_on("python@3.8:", type=("build", "run"))
|
|
depends_on("python@3.9:", type=("build", "run"), when="@0.32:")
|
|
depends_on("py-setuptools", type="build")
|
|
depends_on("py-pybind11", type="link")
|
|
depends_on("py-pip", type="build")
|
|
depends_on("py-wheel", type="build")
|
|
depends_on("py-pennylane@0.28:0.30", type=("build", "run"), when="@:0.30")
|
|
depends_on("py-pennylane@0.30:", type=("build", "run"), when="@0.31")
|
|
# This requirement differs from `pennylane>=0.30` in `setup.py`,
|
|
# but the introduction of `StatePrep` demands `pennylane>=0.32`
|
|
depends_on("py-pennylane@0.32:", type=("build", "run"), when="@0.32")
|
|
depends_on("py-pennylane-lightning~kokkos", type=("build", "run"), when="@:0.31")
|
|
for v in range(33, 38):
|
|
depends_on(f"py-pennylane@0.{v}:", type="run", when=f"@0.{v}")
|
|
depends_on(f"py-pennylane-lightning@0.{v}", type=("build", "run"), when=f"@0.{v}")
|
|
|
|
# variant defined dependencies
|
|
depends_on("llvm-openmp", when="+openmp %apple-clang")
|
|
|
|
# Test deps
|
|
depends_on("py-pytest", type="test")
|
|
depends_on("py-pytest-mock", type="test")
|
|
depends_on("py-flaky", type="test")
|
|
|
|
def url_for_version(self, version):
|
|
extra = "-kokkos" if version <= Version("0.32.0") else ""
|
|
return f"https://github.com/PennyLaneAI/pennylane-lightning{extra}/archive/refs/tags/v{version}.tar.gz"
|
|
|
|
|
|
class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
|
|
build_directory = "build"
|
|
|
|
def setup_build_environment(self, env):
|
|
env.set("PL_BACKEND", "lightning_kokkos")
|
|
cm_args = " ".join([s[2:] for s in self.cmake_args()])
|
|
env.set("CMAKE_ARGS", f"{cm_args}")
|
|
|
|
def cmake_args(self):
|
|
"""
|
|
Here we specify all variant options that can be dynamically specified at build time
|
|
"""
|
|
args_prefix = "PLKOKKOS_" if self.spec.version < Version("0.33") else ""
|
|
|
|
args = [
|
|
self.define_from_variant(f"{args_prefix}BUILD_TESTS", "cpptests"),
|
|
self.define_from_variant("PLKOKKOS_ENABLE_NATIVE", "native"),
|
|
self.define_from_variant("PLKOKKOS_ENABLE_SANITIZER", "sanitize"),
|
|
]
|
|
args.append("-DCMAKE_PREFIX_PATH=" + self.spec["kokkos"].prefix)
|
|
if "+rocm" in self.spec:
|
|
args.append(self.define("CMAKE_CXX_COMPILER", self.spec["hip"].hipcc))
|
|
args.append(
|
|
"-DPLKOKKOS_ENABLE_WARNINGS=OFF"
|
|
) # otherwise build might fail due to Kokkos::InitArguments deprecated
|
|
if self.spec.version >= Version("0.33"):
|
|
args.append("-DPL_BACKEND=lightning_kokkos")
|
|
return args
|
|
|
|
def build(self, pkg, spec, prefix):
|
|
if self.spec.version < Version("0.32"):
|
|
super().build(pkg, spec, prefix)
|
|
cm_args = ";".join([s[2:] for s in self.cmake_args()])
|
|
args = ["-i", f"--define={cm_args}"]
|
|
python("setup.py", "build_ext", *args)
|
|
|
|
def install(self, pkg, spec, prefix):
|
|
pip_args = std_pip_args + [f"--prefix={prefix}", "."]
|
|
pip(*pip_args)
|
|
super().install(pkg, spec, prefix)
|
|
|
|
@run_after("install")
|
|
@on_package_attributes(run_tests=True)
|
|
def install_test(self):
|
|
with working_dir(self.stage.source_path):
|
|
pl_device_test = Executable(join_path(self.prefix, "bin", "pl-device-test"))
|
|
pl_device_test("--device", "lightning.kokkos", "--shots", "None", "--skip-ops")
|