Integrate LGPU to Catalyst

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Integrate the Lightning-GPU with Catalyst.
+  [(#928)](https://github.com/PennyLaneAI/pennylane-lightning/pull/928)
+
 * Add `mid-circuit measurements` support to `lightning.gpu`'s single-GPU backend.
   [(#931)](https://github.com/PennyLaneAI/pennylane-lightning/pull/931)
 

pennylane_lightning/core/_version.py

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

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

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

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

@@ -0,0 +1,85 @@
+cmake_minimum_required(VERSION 3.20)
+
+project(lightning_gpu_catalyst LANGUAGES CXX)
+
+set(LGPU_CATALYST_FILES LightningGPUSimulator.cpp CACHE INTERNAL "")
+add_library(lightning_gpu_catalyst SHARED ${LGPU_CATALYST_FILES})
+
+include(FetchContent)
+
+if(LIGHTNING_CATALYST_SRC_PATH)
+    if(NOT IS_ABSOLUTE ${LIGHTNING_CATALYST_SRC_PATH})
+        message(FATAL_ERROR " LIGHTNING_CATALYST_SRC_PATH=${LIGHTNING_CATALYST_SRC_PATH} must be set to an absolute path")
+    endif()
+    if(CATALYST_GIT_TAG)
+        message(WARN " Setting `LIGHTNING_CATALYST_SRC_PATH=${LIGHTNING_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 ${LIGHTNING_CATALYST_SRC_PATH}
+        OUTPUT_VARIABLE CATALYST_GIT_TAG
+    )
+    message(INFO " Building against local Catalyst - path: ${LIGHTNING_CATALYST_SRC_PATH} - GIT TAG: ${CATALYST_GIT_TAG}")
+
+    target_include_directories(lightning_gpu_catalyst PUBLIC ${LIGHTNING_CATALYST_SRC_PATH}/runtime/lib/backend/common)
+    target_include_directories(lightning_gpu_catalyst PUBLIC ${LIGHTNING_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_gpu_catalyst PUBLIC ${CMAKE_CURRENT_BINARY_DIR}/include)
+
+endif()
+
+target_include_directories(lightning_gpu_catalyst INTERFACE  ${CMAKE_CURRENT_SOURCE_DIR})
+target_link_libraries(lightning_gpu_catalyst PUBLIC  lightning_compile_options
+                                                        lightning_gpu_algorithms
+                                                        lightning_gpu_measurements
+)
+
+if (BUILD_TESTS)
+    enable_testing()
+    add_subdirectory("tests")
+endif()

pennylane_lightning/core/src/simulators/lightning_gpu/catalyst/LightningGPUObsManager.hpp

@@ -0,0 +1,208 @@
+// Copyright 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 "ObservablesGPU.hpp"
+
+namespace Catalyst::Runtime::Simulator {
+
+/**
+ * @brief The LightningGPUObsManager 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 LightningGPUObsManager final {
+  private:
+    using StateVectorT =
+        Pennylane::LightningGPU::StateVectorCudaManaged<PrecisionT>;
+    using ObservableT = Pennylane::Observables::Observable<StateVectorT>;
+    using ObservablePairType = std::pair<std::shared_ptr<ObservableT>, ObsType>;
+    std::vector<ObservablePairType> observables_{};
+
+  public:
+    LightningGPUObsManager() = default;
+    ~LightningGPUObsManager() = default;
+
+    LightningGPUObsManager(const LightningGPUObsManager &) = delete;
+    LightningGPUObsManager &operator=(const LightningGPUObsManager &) = delete;
+    LightningGPUObsManager(LightningGPUObsManager &&) = delete;
+    LightningGPUObsManager &operator=(LightningGPUObsManager &&) = 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<std::size_t>(i) < this->observables_.size());
+        });
+    }
+
+    /**
+     * @brief Get the constructed observable instance.
+     *
+     * @param key The observable key
+     * @return std::shared_ptr<ObservableT>
+     */
+    [[nodiscard]] auto getObservable(ObsIdType key)
+        -> std::shared_ptr<ObservableT> {
+        RT_FAIL_IF(!this->isValidObservables({key}), "Invalid observable key");
+        return std::get<0>(this->observables_[key]);
+    }
+
+    /**
+     * @brief Get the number of observables.
+     *
+     * @return std::size_t
+     */
+    [[nodiscard]] auto numObservables() const -> std::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<std::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::LightningGPU::Observables::NamedObs<StateVectorT>>(
+                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<std::size_t> &wires) -> ObsIdType {
+        std::vector<std::complex<PrecisionT>> matrix_k;
+        matrix_k.reserve(matrix.size());
+        for (const auto &elem : matrix) {
+            matrix_k.push_back(static_cast<std::complex<PrecisionT>>(elem));
+        }
+
+        this->observables_.push_back(std::make_pair(
+            std::make_shared<Pennylane::LightningGPU::Observables::HermitianObs<
+                StateVectorT>>(
+                Pennylane::LightningGPU::Observables::HermitianObs<
+                    StateVectorT>{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<ObservableT>> obs_vec;
+        obs_vec.reserve(key_size);
+
+        for (const auto &key : obsKeys) {
+            RT_FAIL_IF(static_cast<std::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::LightningGPU::Observables::TensorProdObs<
+                               StateVectorT>::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<ObservableT>> obs_vec;
+        obs_vec.reserve(key_size);
+
+        for (auto key : obsKeys) {
+            RT_FAIL_IF(static_cast<std::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::LightningGPU::Observables::Hamiltonian<
+                StateVectorT>>(
+                Pennylane::LightningGPU::Observables::Hamiltonian<StateVectorT>(
+                    coeffs, std::move(obs_vec))),
+            ObsType::Hamiltonian));
+
+        return static_cast<ObsIdType>(obs_size);
+    }
+};
+} // namespace Catalyst::Runtime::Simulator