Add py-pennylane-lightning-kokkos package. (#36257)

* update python package

* change package inheritance

* small update

* enable cpp tests

* small update

* Add flaky package

* Restructure PennyLane deps and order

* Change Lightning defaults and add libomp support for MacOS

* Replace explicit git url with PyPI

* Add Flaky support

* Update PennyLane and PennyLane Lightning support

* fix format

* update packages versioning

* Add patching and default updates for lightning package

* Format

* fix patch version

* update py-flaky package

* update py-pennylane-lightning package

* update py-pennylane package

* remove explicity python dependence

* Remove redundant lines from patch-file

* Update SHA for new patch

* Initial commit for PLLKokkos.

* Comment verbose variant.

* Update develop commit version and restore verbose option.

* Add backends.

* Add mesa package dep (libxml2). Fix rocm install for py-pennylane-lightning-kokkos.

* Restore sycl backend.

* Revert mesa package.

* Make py-pe-li-kokkos into CudaPackage, ROCmPackage.

* Do not force kokkos+wrapper when +cuda

* Few mods following comments on py-pll.

* Update versions of py-pennylane*.

* Remove py-pennylane-lightning patch.

* Remove redundant preferred=True.

* Fix lint in py-pennylane-lightning-kokkos.

* Update var/spack/repos/builtin/packages/py-pennylane-lightning-kokkos/package.py

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

* Ninja and pip not required at runtime. Set lower bound on PL/PLL versions.

* Remove v0.29.0 from pennylane.

* Add AmintorDusko as maintainer.

---------

Co-authored-by: AmintorDusko <amintor_dusko@hotmail.com>
Co-authored-by: Lee J. O'Riordan <lee@xanadu.ai>
Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>
Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

2 files changed, 126 insertions, 3 deletions

diff --git a/var/spack/repos/builtin/packages/py-pennylane-lightning-kokkos/package.py b/var/spack/repos/builtin/packages/py-pennylane-lightning-kokkos/package.py
new file mode 100644
index 0000000000..aae4c5957a
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-pennylane-lightning-kokkos/package.py
@@ -0,0 +1,123 @@
+# 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 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"
+    url = "https://github.com/PennyLaneAI/pennylane-lightning-kokkos/archive/refs/tags/v0.29.1.tar.gz"
+
+    maintainers("AmintorDusko", "vincentmr")
+
+    version("main", branch="main")
+    version("0.29.1", sha256="f51ba7718defc7bb5064f690f381e04b2ec58cb09f22a171ae5f410860716e30")
+
+    # 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"],
+        "pthread": [False, "Whether to build Pthread 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(
+        "build_type",
+        default="Release",
+        description="CMake build type",
+        values=("Debug", "Release", "RelWithDebInfo", "MinSizeRel"),
+    )
+    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("py-setuptools", type="build")
+    depends_on("py-pybind11", type="build")
+    depends_on("py-pip", type="build")
+    depends_on("py-wheel", type="build")
+    depends_on("py-pennylane@0.28:", type=("build", "run"))
+    depends_on("py-pennylane-lightning@0.28:~kokkos", type=("build", "run"))
+
+    # 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")
+
+
+class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
+    build_directory = "build"
+
+    def cmake_args(self):
+        """
+        Here we specify all variant options that can be dynamically specified at build time
+        """
+        args = [
+            self.define_from_variant("PLKOKKOS_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
+        return args
+
+    def build(self, pkg, spec, prefix):
+        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")
diff --git a/var/spack/repos/builtin/packages/py-pennylane/package.py b/var/spack/repos/builtin/packages/py-pennylane/package.py
index 93c1a790e8..286343fc57 100644
--- a/var/spack/repos/builtin/packages/py-pennylane/package.py
+++ b/var/spack/repos/builtin/packages/py-pennylane/package.py
@@ -12,15 +12,15 @@ class PyPennylane(PythonPackage):
 
     homepage = "https://docs.pennylane.ai/"
     git = "https://github.com/PennyLaneAI/pennylane.git"
-    pypi = "PennyLane/PennyLane-0.28.0.tar.gz"
+    url = "https://github.com/PennyLaneAI/pennylane/archive/refs/tags/v0.29.1.tar.gz"
 
     maintainers("mlxd", "AmintorDusko", "marcodelapierre")
 
     version("master", branch="master")
-    version("0.29.1", sha256="c5d662994b741afa69e4fdadc79a1b75840275138a8b7e0cfc5fd64b66a12eef")
+    version("0.29.1", sha256="6ecfb305a3898347df8c539a89a67e748766941d159dbef9e34864872f13c45c")
     version(
         "0.28.0",
-        sha256="2a6100c00277c1eb59eab6856cdad7b1237e9d1fbda98b1e15020bd5a64b10a8",
+        sha256="7736a17dc972d918e3a737ce4360d16ac84b9f9f6fca440f167de579c926c114",
         deprecated=True,
     )