Add native setStateVector support to lightning.gpu

.github/CHANGELOG.md

@@ -34,6 +34,9 @@
 
 ### Improvements
 
+* Add `setStateVector(state, wire)` support to the `lightning.gpu` C++ layer.
+  [(#930)](https://github.com/PennyLaneAI/pennylane-lightning/pull/930)
+
 * Always decompose `qml.QFT` in Lightning.
   [(#924)](https://github.com/PennyLaneAI/pennylane-lightning/pull/924)
 

pennylane_lightning/core/_version.py

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

pennylane_lightning/core/src/simulators/lightning_gpu/StateVectorCudaMPI.hpp

@@ -340,6 +340,46 @@ class StateVectorCudaMPI final
         mpi_manager_.Barrier();
     }
 
+    /**
+     * @brief Set values for a batch of elements of the state-vector.
+     *
+     * @param state_ptr Pointer to initial state data.
+     * @param num_states Length of initial state data.
+     * @param wires Wires.
+     */
+    void setStateVector(const ComplexT *state_ptr, const std::size_t num_states,
+                        const std::vector<std::size_t> &wires) {
+        PL_ABORT_IF_NOT(num_states == Pennylane::Util::exp2(wires.size()),
+                        "Inconsistent state and wires dimensions.");
+
+        const auto num_qubits = this->getTotalNumQubits();
+
+        PL_ABORT_IF_NOT(std::find_if(wires.begin(), wires.end(),
+                                     [&num_qubits](const auto i) {
+                                         return i >= num_qubits;
+                                     }) == wires.end(),
+                        "Invalid wire index.");
+
+        using index_type =
+            typename std::conditional<std::is_same<PrecisionT, float>::value,
+                                      int32_t, int64_t>::type;
+
+        // Calculate the indices of the state-vector to be set.
+        // TODO: Could move to GPU calculation if the state size is large.
+        std::vector<index_type> indices(num_states);
+        constexpr std::size_t one{1U};
+        for (std::size_t i = 0; i < num_states; i++) {
+            std::size_t index{0U};
+            for (std::size_t j = 0; j < wires.size(); j++) {
+                const std::size_t bit = (i & (one << j)) >> j;
+                index |= bit << (num_qubits - 1 - wires[wires.size() - 1 - j]);
+            }
+            indices[i] = index;
+        }
+        setStateVector<index_type>(num_states, state_ptr, indices.data());
+        mpi_manager_.Barrier();
+    }
+
     /**
      * @brief Apply a single gate to the state-vector. Offloads to custatevec
      * specific API calls if available. If unable, attempts to use prior cached

pennylane_lightning/core/src/simulators/lightning_gpu/StateVectorCudaManaged.hpp

@@ -217,6 +217,45 @@ class StateVectorCudaManaged
                             thread_per_block, stream_id);
     }
 
+    /**
+     * @brief Set values for a batch of elements of the state-vector.
+     *
+     * @param state_ptr Pointer to the initial state data.
+     * @param num_states Length of the initial state data.
+     * @param wires Wires.
+     */
+    void setStateVector(const ComplexT *state_ptr, const std::size_t num_states,
+                        const std::vector<std::size_t> &wires) {
+        PL_ABORT_IF_NOT(num_states == Pennylane::Util::exp2(wires.size()),
+                        "Inconsistent state and wires dimensions.");
+
+        const auto num_qubits = BaseType::getNumQubits();
+
+        PL_ABORT_IF_NOT(std::find_if(wires.begin(), wires.end(),
+                                     [&num_qubits](const auto i) {
+                                         return i >= num_qubits;
+                                     }) == wires.end(),
+                        "Invalid wire index.");
+
+        using index_type =
+            typename std::conditional<std::is_same<PrecisionT, float>::value,
+                                      int32_t, int64_t>::type;
+
+        // Calculate the indices of the state-vector to be set.
+        // TODO: Could move to GPU calculation if the state size is large.
+        std::vector<index_type> indices(num_states);
+        constexpr std::size_t one{1U};
+        for (std::size_t i = 0; i < num_states; i++) {
+            std::size_t index{0U};
+            for (std::size_t j = 0; j < wires.size(); j++) {
+                const std::size_t bit = (i & (one << j)) >> j;
+                index |= bit << (num_qubits - 1 - wires[wires.size() - 1 - j]);
+            }
+            indices[i] = index;
+        }
+        setStateVector<index_type>(num_states, state_ptr, indices.data());
+    }
+
     /**
      * @brief Multiplies the state-vector by a global phase.
      *

pennylane_lightning/core/src/simulators/lightning_gpu/bindings/LGPUBindings.hpp

@@ -105,6 +105,17 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
             },
             "Set State Vector on GPU with values and their corresponding "
             "indices for the state vector on device")
+        .def(
+            "setStateVector",
+            [](StateVectorT &sv, const np_arr_c &state,
+               const std::vector<std::size_t> &wires) {
+                const auto state_buffer = state.request();
+                const auto state_ptr =
+                    static_cast<const std::complex<ParamT> *>(state_buffer.ptr);
+                sv.setStateVector(state_ptr, state_buffer.size, wires);
+            },
+            "Set State Vector on GPU with values for the state vector and "
+            "wires on the host memory.")
         .def(
             "DeviceToDevice",
             [](StateVectorT &sv, const StateVectorT &other, bool async) {

pennylane_lightning/core/src/simulators/lightning_gpu/bindings/LGPUBindingsMPI.hpp

@@ -108,6 +108,17 @@ void registerBackendClassSpecificBindingsMPI(PyClass &pyclass) {
             },
             "Set State Vector on GPU with values and their corresponding "
             "indices for the state vector on device")
+        .def(
+            "setStateVector",
+            [](StateVectorT &sv, const np_arr_c &state,
+               const std::vector<std::size_t> &wires) {
+                const auto state_buffer = state.request();
+                const auto state_ptr =
+                    static_cast<const std::complex<ParamT> *>(state_buffer.ptr);
+                sv.setStateVector(state_ptr, state_buffer.size, wires);
+            },
+            "Set State Vector on GPU with values for the state vector and "
+            "wires on the host memory.")
         .def(
             "DeviceToDevice",
             [](StateVectorT &sv, const StateVectorT &other, bool async) {

pennylane_lightning/core/src/simulators/lightning_gpu/gates/tests/Test_StateVectorCudaManaged_NonParam.cpp

@@ -1103,6 +1103,13 @@ TEMPLATE_TEST_CASE("StateVectorCudaManaged::SetStateVector",
                                                indices.data(), false);
 
         CHECK(expected_state == Pennylane::Util::approx(sv.getDataVector()));
+
+        sv.initSV();
+        std::copy(init_state.begin(), init_state.end(),
+                  values.begin()); // copy the data to values
+        sv.setStateVector(values.data(), values.size(),
+                          std::vector<std::size_t>{0, 1, 2});
+        CHECK(init_state == Pennylane::Util::approx(sv.getDataVector()));
     }
 }
 // LCOV_EXCL_START

pennylane_lightning/core/src/simulators/lightning_gpu/gates/tests/mpi/Test_StateVectorCudaMPI_NonParam.cpp

@@ -15,6 +15,7 @@
 #include <complex>
 #include <iostream>
 #include <limits>
+#include <numeric>
 #include <type_traits>
 #include <utility>
 #include <vector>
@@ -173,6 +174,19 @@ TEMPLATE_TEST_CASE("StateVectorCudaMPI::SetStateVector",
         mpi_manager.Barrier();
 
         CHECK(expected_local_state == Pennylane::Util::approx(local_state));
+
+        sv.initSV();
+        std::vector<std::complex<PrecisionT>> values(init_state.size());
+        std::copy(init_state.begin(), init_state.end(),
+                  values.begin()); // copy the data to values
+        std::vector<std::size_t> wires(num_qubits);
+        std::iota(wires.begin(), wires.end(), 0);
+        sv.setStateVector(values.data(), values.size(), wires);
+
+        auto expected_local_state_vector = mpi_manager.scatter<cp_t>(values, 0);
+
+        CHECK(expected_local_state_vector ==
+              Pennylane::Util::approx(sv.getDataVector()));
     }
 }
 

pennylane_lightning/lightning_gpu/lightning_gpu.py

@@ -19,7 +19,7 @@
 
 from ctypes.util import find_library
 from importlib import util as imp_util
 from itertools import product
 from pathlib import Path
 from typing import List, Union
 from warnings import warn
@@ -429,7 +429,7 @@
         """
         self._gpu_state.setBasisState(index, use_async)
 
     def _apply_state_vector(self, state, device_wires, use_async=False):
         """Initialize the state vector on GPU with a specified state on host.
         Note that any use of this method will introduce host-overheads.
         Args:
@@ -456,20 +456,8 @@
             self.syncH2D(self._reshape(local_state, output_shape))
             return
 
-        # generate basis states on subset of qubits via the cartesian product
-        basis_states = np.array(list(product([0, 1], repeat=len(device_wires))))
-
-        # get basis states to alter on full set of qubits
-        unravelled_indices = np.zeros((2 ** len(device_wires), self.num_wires), dtype=int)
-        unravelled_indices[:, device_wires] = basis_states
-
-        # get indices for which the state is changed to input state vector elements
-        ravelled_indices = np.ravel_multi_index(unravelled_indices.T, [2] * self.num_wires)
-
-        # set the state vector on GPU with the unravelled_indices and their corresponding values
-        self._gpu_state.setStateVector(
-            ravelled_indices, state, use_async
-        )  # this operation on device
+        # set the state vector on GPU with the input values and wires
+        self._gpu_state.setStateVector(state, list(device_wires))
 
     def _apply_basis_state(self, state, wires):
         """Initialize the state vector in a specified computational basis state on GPU directly.