-
Notifications
You must be signed in to change notification settings - Fork 36
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add Catalyst specific Lightning Kokkos class (#770)
**Context:** Catalyst needs to build a class wrapping the Lightning Kokkos class. Here we are moving the logic from Catalyst to the Lightning repository. **Description of the Change:** Moving code from Catalyst to this repository and updating the build system to build against Catalyst. **Benefits:** Catalyst wheels will build faster. **Possible Drawbacks:** Our build system now relies on headers coming from Catalyst. It is unclear if this may cause a deadlock in the future because of our cyclic dependencies (PennyLane, PennyLane Lightning, Catalyst). Chances are small because we don't rely on Catalyst wheels or build it from scratch. **Usage:** This PR offers three ways to configure the CMake building system: 1. Providing a local source path for Catalyst. Example: ```bash cmake -BBuildTests -G Ninja \ -DCMAKE_INSTALL_PREFIX=$path_to_kokkos \ -DCATALYST_SRC_PATH=$path_to_catalyst \ -DCMAKE_BUILD_TYPE=Debug \ -DBUILD_TESTS=ON \ -DENABLE_WARNINGS=ON \ -DPL_BACKEND=lightning_kokkos ``` 3. Providing a GIT TAG. Example with the main branch: ```bash cmake -BBuildTests -G Ninja \ -DCMAKE_INSTALL_PREFIX=$path_to_kokkos \ -DCATALYST_GIT_TAG=main \ -DCMAKE_BUILD_TYPE=Debug \ -DBUILD_TESTS=ON \ -DENABLE_WARNINGS=ON \ -DPL_BACKEND=lightning_kokkos ``` 4. If the dev/user doesn't say anything the configuration will default to get headers from Catalyst GitHub main branch. Example: ```bash cmake -BBuildTests -G Ninja \ -DCMAKE_INSTALL_PREFIX=$path_to_kokkos \ -DCMAKE_BUILD_TYPE=Debug \ -DBUILD_TESTS=ON \ -DENABLE_WARNINGS=ON \ -DPL_BACKEND=lightning_kokkos ``` Check out [ADR 76](PennyLaneAI/adrs#76) for alternative approaches. [sc-59312] --------- Co-authored-by: ringo-but-quantum <[email protected]> Co-authored-by: Ali Asadi <[email protected]>
- Loading branch information
1 parent
91241a0
commit b9b6e1c
Showing
12 changed files
with
3,809 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -16,4 +16,4 @@ | |
Version number (major.minor.patch[-label]) | ||
""" | ||
|
||
__version__ = "0.38.0-dev0" | ||
__version__ = "0.38.0-dev1" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
92 changes: 92 additions & 0 deletions
92
pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/CMakeLists.txt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,92 @@ | ||
cmake_minimum_required(VERSION 3.20) | ||
|
||
project(lightning_kokkos_catalyst LANGUAGES CXX) | ||
|
||
set(LK_CATALYST_FILES LightningKokkosSimulator.cpp CACHE INTERNAL "") | ||
add_library(lightning_kokkos_catalyst SHARED ${LK_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_kokkos_catalyst PUBLIC ${LIGHTNING_CATALYST_SRC_PATH}/runtime/lib/backend/common) | ||
target_include_directories(lightning_kokkos_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_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() |
208 changes: 208 additions & 0 deletions
208
...ane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosObsManager.hpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,208 @@ | ||
// 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 final { | ||
private: | ||
using StateVectorT = | ||
Pennylane::LightningKokkos::StateVectorKokkos<PrecisionT>; | ||
using ObservableT = Pennylane::Observables::Observable<StateVectorT>; | ||
using ObservablePairType = std::pair<std::shared_ptr<ObservableT>, 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<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 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< | ||
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<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<StateVectorT>>( | ||
Pennylane::LightningKokkos::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<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< | ||
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<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<StateVectorT>>( | ||
Pennylane::LightningKokkos::Observables::Hamiltonian< | ||
StateVectorT>(coeffs, std::move(obs_vec))), | ||
ObsType::Hamiltonian)); | ||
|
||
return static_cast<ObsIdType>(obs_size); | ||
} | ||
}; | ||
} // namespace Catalyst::Runtime::Simulator |
Oops, something went wrong.