Feature/lkokkos mpi

cmake/support_kokkos.cmake

@@ -5,7 +5,7 @@
 # Include this file only once
 include_guard()
 
-set(KOKKOS_VERSION 4.2.00)
+set(KOKKOS_VERSION 4.2.01)
 
 # Macro to aid in finding Kokkos with 3 potential install options:
 # 1. Fully integrated Kokkos packages and CMake module files

mpitests/test_adjoint_jacobian.py

@@ -1396,9 +1396,7 @@ def test_qubit_unitary(dev, n_targets):
 
     np.random.seed(1337)
     par = 2 * np.pi * np.random.rand(n_wires)
-    U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(
-        2**n_targets, 2**n_targets
-    )
+    U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(2**n_targets, 2**n_targets)
     U, _ = np.linalg.qr(U)
     init_state = np.random.rand(2**n_wires) + 1j * np.random.rand(2**n_wires)
     init_state /= np.sqrt(np.dot(np.conj(init_state), init_state))
@@ -1446,9 +1444,7 @@ def test_diff_qubit_unitary(dev, n_targets):
 
     np.random.seed(1337)
     par = 2 * np.pi * np.random.rand(n_wires)
-    U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(
-        2**n_targets, 2**n_targets
-    )
+    U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(2**n_targets, 2**n_targets)
     U, _ = np.linalg.qr(U)
     init_state = np.random.rand(2**n_wires) + 1j * np.random.rand(2**n_wires)
     init_state /= np.sqrt(np.dot(np.conj(init_state), init_state))

pennylane_lightning/core/_serialize.py

@@ -298,9 +298,7 @@ def serialize_observables(self, tape: QuantumTape, wires_map: dict = None) -> Li
                 offset_indices.append(offset_indices[-1] + 1)
         return serialized_obs, offset_indices
 
-    def serialize_ops(
-        self, tape: QuantumTape, wires_map: dict = None
-    ) -> Tuple[
+    def serialize_ops(self, tape: QuantumTape, wires_map: dict = None) -> Tuple[
         List[List[str]],
         List[np.ndarray],
         List[List[int]],

pennylane_lightning/core/src/CMakeLists.txt

@@ -23,4 +23,14 @@ if (BUILD_TESTS)
     FetchAndIncludeCatch()
 
     include(CTest)
+endif()
+
+##########################
+## Create project target #
+##########################
+
+if (ENABLE_MPI)
+    add_executable(lightning lightning.cpp)
+    target_link_libraries(lightning pennylane_lightning)
+    set_property(TARGET lightning PROPERTY CXX_STANDARD 20)
 endif()

pennylane_lightning/core/src/bindings/Bindings.hpp

@@ -59,8 +59,12 @@ using namespace Pennylane::LightningQubit::Measures;
 
 #include "AdjointJacobianKokkos.hpp"
 #include "LKokkosBindings.hpp" // StateVectorBackends, registerBackendClassSpecificBindings, registerBackendSpecificMeasurements, registerBackendSpecificAlgorithms
-#include "MeasurementsKokkos.hpp"
 #include "ObservablesKokkos.hpp"
+#if _ENABLE_MPI == 1
+#include "MeasurementsKokkosMPI.hpp"
+#else
+#include "MeasurementsKokkos.hpp"
+#endif
 
 /// @cond DEV
 namespace {
@@ -97,6 +101,11 @@ namespace {
 using Pennylane::Util::bestCPUMemoryModel;
 using Pennylane::Util::CPUMemoryModel;
 using Pennylane::Util::getMemoryModel;
+#if _ENABLE_PLKOKKOS == 1 && _ENABLE_MPI == 1
+template <class T> using measure = MeasurementsMPI<T>;
+#else
+template <class T> using measure = Measurements<T>;
+#endif
 } // namespace
 /// @endcond
 
@@ -426,48 +435,50 @@ void registerBackendAgnosticMeasurements(PyClass &pyclass) {
     using PrecisionT =
         typename StateVectorT::PrecisionT; // Statevector's precision.
     using ParamT = PrecisionT;             // Parameter's data precision
-
+    using measureT = measure<StateVectorT>;
     pyclass
-        .def("probs",
-             [](Measurements<StateVectorT> &M,
-                const std::vector<size_t> &wires) {
-                 return py::array_t<ParamT>(py::cast(M.probs(wires)));
-             })
-        .def("probs",
-             [](Measurements<StateVectorT> &M) {
-                 return py::array_t<ParamT>(py::cast(M.probs()));
-             })
         .def(
             "expval",
-            [](Measurements<StateVectorT> &M,
+            [](measureT &M,
                const std::shared_ptr<Observable<StateVectorT>> &ob) {
                 return M.expval(*ob);
             },
             "Expected value of an observable object.")
         .def(
             "var",
-            [](Measurements<StateVectorT> &M,
+            [](measureT &M,
                const std::shared_ptr<Observable<StateVectorT>> &ob) {
                 return M.var(*ob);
             },
             "Variance of an observable object.")
-        .def("generate_samples", [](Measurements<StateVectorT> &M,
-                                    size_t num_wires, size_t num_shots) {
-            auto &&result = M.generate_samples(num_shots);
-            const size_t ndim = 2;
-            const std::vector<size_t> shape{num_shots, num_wires};
-            constexpr auto sz = sizeof(size_t);
-            const std::vector<size_t> strides{sz * num_wires, sz};
-            // return 2-D NumPy array
-            return py::array(py::buffer_info(
-                result.data(), /* data as contiguous array  */
-                sz,            /* size of one scalar        */
-                py::format_descriptor<size_t>::format(), /* data type */
-                ndim,   /* number of dimensions      */
-                shape,  /* shape of the matrix       */
-                strides /* strides for each axis     */
-                ));
-        });
+#if _ENABLE_MPI != 1 || _ENABLE_PLKOKKOS != 1
+        .def("probs",
+             [](measureT &M, const std::vector<size_t> &wires) {
+                 return py::array_t<ParamT>(py::cast(M.probs(wires)));
+             })
+        .def("probs",
+             [](measureT &M) {
+                 return py::array_t<ParamT>(py::cast(M.probs()));
+             })
+        .def("generate_samples",
+             [](measureT &M, size_t num_wires, size_t num_shots) {
+                 auto &&result = M.generate_samples(num_shots);
+                 const size_t ndim = 2;
+                 const std::vector<size_t> shape{num_shots, num_wires};
+                 constexpr auto sz = sizeof(size_t);
+                 const std::vector<size_t> strides{sz * num_wires, sz};
+                 // return 2-D NumPy array
+                 return py::array(py::buffer_info(
+                     result.data(), /* data as contiguous array  */
+                     sz,            /* size of one scalar        */
+                     py::format_descriptor<size_t>::format(), /* data type */
+                     ndim,   /* number of dimensions      */
+                     shape,  /* shape of the matrix       */
+                     strides /* strides for each axis     */
+                     ));
+             })
+#endif
+        ;
 }
 
 /**
@@ -659,7 +670,7 @@ template <class StateVectorT> void lightningClassBindings(py::module_ &m) {
     //                             Measurements
     //***********************************************************************//
     class_name = "MeasurementsC" + bitsize;
-    auto pyclass_measurements = py::class_<Measurements<StateVectorT>>(
+    auto pyclass_measurements = py::class_<measure<StateVectorT>>(
         m, class_name.c_str(), py::module_local());
 
 #ifdef _ENABLE_PLGPU