Add Catalyst specific Lightning Kokkos class

.github/CHANGELOG.md

@@ -2,10 +2,10 @@
 
 ### New features since last release
 * Implement Python interface to the `lightning.tensor` device.
-  [(#748)](https://github.com/PennyLaneAI/pennylane-lightning/pull/748) 
+  [(#748)](https://github.com/PennyLaneAI/pennylane-lightning/pull/748)
 
 * Add `inverse` support for gate operations in `lightning.tensor` in the C++ layer.
-  [(#753)](https://github.com/PennyLaneAI/pennylane-lightning/pull/753) 
+  [(#753)](https://github.com/PennyLaneAI/pennylane-lightning/pull/753)
 
 * Add `observable` and `expval` support to `cutensornet` backed `lightning.tensor` C++ layer.
   [(#728)](https://github.com/PennyLaneAI/pennylane-lightning/pull/728)
@@ -19,6 +19,9 @@
 * Add support for `C(BlockEncode)` to Lightning devices.
   [(#743)](https://github.com/PennyLaneAI/pennylane-lightning/pull/743)
 
+* Add a Catalyst-specific wrapping class for Lightning Kokkos.
+  [(#770)](https://github.com/PennyLaneAI/pennylane-lightning/pull/770)
+
 ### Breaking changes
 
 * Removed the `QuimbMPS` class and the corresponding interface/backend from `lightning.tensor`.
@@ -45,17 +48,17 @@
   [(#763)](https://github.com/PennyLaneAI/pennylane-lightning/pull/763)
 
 * Change the type of tensor network objects passed to `ObservablesTNCuda` and `MeasurementsTNCuda` class from `StateTensorT` to `TensorNetT`.
-  [(#759)](https://github.com/PennyLaneAI/pennylane-lightning/pull/759) 
-  
+  [(#759)](https://github.com/PennyLaneAI/pennylane-lightning/pull/759)
+
 * Rationalize MCM tests, removing most end-to-end tests from the native MCM test file,
   but keeping one that validates multiple mid-circuit measurements with any allowed return.
   [(#754)](https://github.com/PennyLaneAI/pennylane/pull/754)
 
 * Refactor C++ library names for `lightning.tensor`.
-  [(#755)](https://github.com/PennyLaneAI/pennylane-lightning/pull/755) 
+  [(#755)](https://github.com/PennyLaneAI/pennylane-lightning/pull/755)
 
 * Set `state_tensor` as `const` for the `MeasurementTNCuda` class.
-  [(#753)](https://github.com/PennyLaneAI/pennylane-lightning/pull/753) 
+  [(#753)](https://github.com/PennyLaneAI/pennylane-lightning/pull/753)
 
 * Updated Kokkos version and support to 4.3.01.
   [(#725)](https://github.com/PennyLaneAI/pennylane-lightning/pull/725)

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.37.0-dev46"
+__version__ = "0.37.0-dev47"

pennylane_lightning/core/src/simulators/lightning_kokkos/CMakeLists.txt

@@ -70,6 +70,7 @@ endif()
 ###############################################################################
 set(COMPONENT_SUBDIRS   algorithms
                         bindings
+                        catalyst
                         gates
                         measurements
                         observables

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/CMakeLists.txt

@@ -0,0 +1,92 @@
+cmake_minimum_required(VERSION 3.20)
+
+project(lightning_kokkos_catalyst LANGUAGES CXX)
+
+set(CATALYST_FILES LightningKokkosSimulator.cpp CACHE INTERNAL "" FORCE)
+add_library(lightning_kokkos_catalyst SHARED ${CATALYST_FILES})
+
+include(FetchContent)
+
+if(CATALYST_SRC_PATH)
+    if(NOT IS_ABSOLUTE ${CATALYST_SRC_PATH})
+        message(FATAL_ERROR " CATALYST_SRC_PATH=${CATALYST_SRC_PATH} must be set to an absolute path")
+    endif()
+    if(CATALYST_GIT_TAG)
+        message(WARN " Setting `CATALYST_SRC_PATH=${CATALYST_SRC_PATH}` overrides `CATALYST_GIT_TAG=${CATALYST_GIT_TAG}`")
+    endif()
+
+    # Acquire local git hash and use for CATALYST_GIT_TAG
+    execute_process(COMMAND git rev-parse --short HEAD
+        WORKING_DIRECTORY ${CATALYST_SRC_PATH}
+        OUTPUT_VARIABLE CATALYST_GIT_TAG
+    )
+    message(INFO " Building against local Catalyst - path: ${CATALYST_SRC_PATH} - GIT TAG: ${CATALYST_GIT_TAG}")
+
+    target_include_directories(lightning_kokkos_catalyst PUBLIC ${CATALYST_SRC_PATH}/runtime/lib/backend/common)
+    target_include_directories(lightning_kokkos_catalyst PUBLIC ${CATALYST_SRC_PATH}/runtime/include)
+
+else()
+    if(NOT CATALYST_GIT_TAG)
+        set(CATALYST_GIT_TAG "main" CACHE STRING "GIT_TAG value to build Catalyst")
+    endif()
+    message(INFO " Building against Catalyst GIT TAG ${CATALYST_GIT_TAG}")
+
+    # Fetching /lib/backend/common hpp headers
+    set(LIB_BACKEND_COMMON_HEADERS  CacheManager.hpp
+                                    QubitManager.hpp
+                                    Utils.hpp
+    )
+
+    foreach(HEADER ${LIB_BACKEND_COMMON_HEADERS})
+        string(REGEX REPLACE "\\.[^.]*$" "" HEADER_NAME ${HEADER})
+        FetchContent_Declare(
+            ${HEADER_NAME}
+            URL                 https://raw.githubusercontent.com/PennyLaneAI/catalyst/${CATALYST_GIT_TAG}/runtime/lib/backend/common/${HEADER}
+            DOWNLOAD_NO_EXTRACT True
+            SOURCE_DIR          include
+        )
+
+        FetchContent_MakeAvailable(${HEADER_NAME})
+    endforeach()
+
+    # Fetching include hpp headers
+    set(INCLUDE_HEADERS DataView.hpp
+                        Exception.hpp
+                        QuantumDevice.hpp
+                        RuntimeCAPI.h
+                        Types.h
+    )
+
+    foreach(HEADER ${INCLUDE_HEADERS})
+        string(REGEX REPLACE "\\.[^.]*$" "" HEADER_NAME ${HEADER})
+        FetchContent_Declare(
+            ${HEADER_NAME}
+            URL                 https://raw.githubusercontent.com/PennyLaneAI/catalyst/${CATALYST_GIT_TAG}/runtime/include/${HEADER}
+            DOWNLOAD_NO_EXTRACT True
+            SOURCE_DIR          include
+        )
+
+        FetchContent_MakeAvailable(${HEADER_NAME})
+    endforeach()
+
+    target_include_directories(lightning_kokkos_catalyst PUBLIC ${CMAKE_CURRENT_BINARY_DIR}/include)
+
+endif()
+
+target_include_directories(lightning_kokkos_catalyst INTERFACE  ${CMAKE_CURRENT_SOURCE_DIR})
+target_link_libraries(lightning_kokkos_catalyst PUBLIC  lightning_compile_options
+                                                        lightning_external_libs
+                                                        lightning_base
+                                                        lightning_gates
+                                                        lightning_utils
+                                                        lightning_kokkos
+                                                        lightning_kokkos_algorithms
+                                                        lightning_kokkos_gates
+                                                        lightning_kokkos_measurements
+                                                        lightning_kokkos_utils
+)
+
+if (BUILD_TESTS)
+    enable_testing()
+    add_subdirectory("tests")
+endif()

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosObsManager.hpp

@@ -0,0 +1,210 @@
+// Copyright 2023-2024 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <array>
+#include <stdexcept>
+#include <tuple>
+#include <utility>
+
+#include "Types.h"
+#include "Utils.hpp"
+
+#include "ObservablesKokkos.hpp"
+
+namespace Catalyst::Runtime::Simulator {
+
+/**
+ * @brief The LightningKokkosObsManager caches observables of a program at
+ * runtime and maps each one to a const unique index (`int64_t`) in the scope of
+ * the global context manager.
+ */
+template <typename PrecisionT> class LightningKokkosObsManager {
+  private:
+    using VectorStateT =
+        Pennylane::LightningKokkos::StateVectorKokkos<PrecisionT>;
+    using ObservableClassName =
+        Pennylane::Observables::Observable<VectorStateT>;
+    using ObservablePairType =
+        std::pair<std::shared_ptr<ObservableClassName>, ObsType>;
+    std::vector<ObservablePairType> observables_{};
+
+  public:
+    LightningKokkosObsManager() = default;
+    ~LightningKokkosObsManager() = default;
+
+    LightningKokkosObsManager(const LightningKokkosObsManager &) = delete;
+    LightningKokkosObsManager &
+    operator=(const LightningKokkosObsManager &) = delete;
+    LightningKokkosObsManager(LightningKokkosObsManager &&) = delete;
+    LightningKokkosObsManager &operator=(LightningKokkosObsManager &&) = delete;
+
+    /**
+     * @brief A helper function to clear constructed observables in the program.
+     */
+    void clear() { this->observables_.clear(); }
+
+    /**
+     * @brief Check the validity of observable keys.
+     *
+     * @param obsKeys The vector of observable keys
+     * @return bool
+     */
+    [[nodiscard]] auto
+    isValidObservables(const std::vector<ObsIdType> &obsKeys) const -> bool {
+        return std::all_of(obsKeys.begin(), obsKeys.end(), [this](auto i) {
+            return (i >= 0 &&
+                    static_cast<size_t>(i) < this->observables_.size());
+        });
+    }
+
+    /**
+     * @brief Get the constructed observable instance.
+     *
+     * @param key The observable key
+     * @return std::shared_ptr<ObservableClassName>
+     */
+    [[nodiscard]] auto getObservable(ObsIdType key)
+        -> std::shared_ptr<ObservableClassName> {
+        RT_FAIL_IF(!this->isValidObservables({key}), "Invalid observable key");
+        return std::get<0>(this->observables_[key]);
+    }
+
+    /**
+     * @brief Get the number of observables.
+     *
+     * @return size_t
+     */
+    [[nodiscard]] auto numObservables() const -> size_t {
+        return this->observables_.size();
+    }
+
+    /**
+     * @brief Create and cache a new NamedObs instance.
+     *
+     * @param obsId The named observable id of type ObsId
+     * @param wires The vector of wires the observable acts on
+     * @return ObsIdType
+     */
+    [[nodiscard]] auto createNamedObs(ObsId obsId,
+                                      const std::vector<size_t> &wires)
+        -> ObsIdType {
+        auto &&obs_str = std::string(
+            Lightning::lookup_obs<Lightning::simulator_observable_support_size>(
+                Lightning::simulator_observable_support, obsId));
+
+        this->observables_.push_back(std::make_pair(
+            std::make_shared<Pennylane::LightningKokkos::Observables::NamedObs<
+                VectorStateT>>(obs_str, wires),
+            ObsType::Basic));
+        return static_cast<ObsIdType>(this->observables_.size() - 1);
+    }
+
+    /**
+     * @brief Create and cache a new HermitianObs instance.
+     *
+     * @param matrix The row-wise Hermitian matrix
+     * @param wires The vector of wires the observable acts on
+     * @return ObsIdType
+     */
+    [[nodiscard]] auto
+    createHermitianObs(const std::vector<std::complex<PrecisionT>> &matrix,
+                       const std::vector<size_t> &wires) -> ObsIdType {
+        std::vector<Kokkos::complex<PrecisionT>> matrix_k;
+        matrix_k.reserve(matrix.size());
+        for (const auto &elem : matrix) {
+            matrix_k.push_back(static_cast<Kokkos::complex<PrecisionT>>(elem));
+        }
+
+        this->observables_.push_back(std::make_pair(
+            std::make_shared<Pennylane::LightningKokkos::Observables::
+                                 HermitianObs<VectorStateT>>(
+                Pennylane::LightningKokkos::Observables::HermitianObs<
+                    VectorStateT>{matrix_k, wires}),
+            ObsType::Basic));
+
+        return static_cast<ObsIdType>(this->observables_.size() - 1);
+    }
+
+    /**
+     * @brief Create and cache a new TensorProd instance.
+     *
+     * @param obsKeys The vector of observable keys
+     * @return ObsIdType
+     */
+    [[nodiscard]] auto
+    createTensorProdObs(const std::vector<ObsIdType> &obsKeys) -> ObsIdType {
+        const auto key_size = obsKeys.size();
+        const auto obs_size = this->observables_.size();
+
+        std::vector<std::shared_ptr<ObservableClassName>> obs_vec;
+        obs_vec.reserve(key_size);
+
+        for (const auto &key : obsKeys) {
+            RT_FAIL_IF(static_cast<size_t>(key) >= obs_size || key < 0,
+                       "Invalid observable key");
+
+            auto &&[obs, type] = this->observables_[key];
+            obs_vec.push_back(obs);
+        }
+
+        this->observables_.push_back(std::make_pair(
+            Pennylane::LightningKokkos::Observables::TensorProdObs<
+                VectorStateT>::create(obs_vec),
+            ObsType::TensorProd));
+
+        return static_cast<ObsIdType>(obs_size);
+    }
+
+    /**
+     * @brief Create and cache a new HamiltonianObs instance.
+     *
+     * @param coeffs The vector of coefficients
+     * @param obsKeys The vector of observable keys
+     * @return ObsIdType
+     */
+    [[nodiscard]] auto
+    createHamiltonianObs(const std::vector<PrecisionT> &coeffs,
+                         const std::vector<ObsIdType> &obsKeys) -> ObsIdType {
+        const auto key_size = obsKeys.size();
+        const auto obs_size = this->observables_.size();
+
+        RT_FAIL_IF(
+            key_size != coeffs.size(),
+            "Incompatible list of observables and coefficients; "
+            "Number of observables and number of coefficients must be equal");
+
+        std::vector<std::shared_ptr<ObservableClassName>> obs_vec;
+        obs_vec.reserve(key_size);
+
+        for (auto key : obsKeys) {
+            RT_FAIL_IF(static_cast<size_t>(key) >= obs_size || key < 0,
+                       "Invalid observable key");
+
+            auto &&[obs, type] = this->observables_[key];
+            obs_vec.push_back(obs);
+        }
+
+        this->observables_.push_back(std::make_pair(
+            std::make_shared<Pennylane::LightningKokkos::Observables::
+                                 Hamiltonian<VectorStateT>>(
+                Pennylane::LightningKokkos::Observables::Hamiltonian<
+                    VectorStateT>(coeffs, std::move(obs_vec))),
+            ObsType::Hamiltonian));
+
+        return static_cast<ObsIdType>(obs_size);
+    }
+};
+} // namespace Catalyst::Runtime::Simulator