Add SetState and SetBasisState to kokkos

.github/CHANGELOG.md

@@ -55,6 +55,9 @@
 * The `setBasisState` and `setStateVector` methods of `StateVectorLQubit` and `StateVectorKokkos` are overloaded to support PennyLane-like parameters.
   [(#843)](https://github.com/PennyLaneAI/pennylane-lightning/pull/843)
 
+* Add `SetState` and `SetBasisState` to `LightningKokkosSimulator`.
+  [(#861)](https://github.com/PennyLaneAI/pennylane-lightning/pull/861)
+
 * `ENABLE_LAPACK` is off by default for all Lightning backends.
   [(#825)](https://github.com/PennyLaneAI/pennylane-lightning/pull/825)
 

pennylane_lightning/core/_version.py

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

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <cmath>
+
 #include <Kokkos_Complex.hpp>
 #include <Kokkos_Core.hpp>
 
@@ -137,6 +139,23 @@ void LightningKokkosSimulator::PrintState() {
 }
 /// LCOV_EXCL_STOP
 
+void LightningKokkosSimulator::SetState(DataView<std::complex<double>, 1> &data,
+                                        std::vector<QubitIdType> &wires) {
+    size_t expected_wires = static_cast<size_t>(log2(data.size()));
+    RT_ASSERT(expected_wires == wires.size());
+    std::vector<Kokkos::complex<double>> data_vector(data.begin(), data.end());
+    std::vector<std::size_t> wires_size_t(wires.begin(), wires.end());
+
+    this->device_sv->setStateVector(data_vector, wires_size_t);
+}
+
+void LightningKokkosSimulator::SetBasisState(DataView<int8_t, 1> &data,
+                                             std::vector<QubitIdType> &wires) {
+    std::vector<std::size_t> data_vector(data.begin(), data.end());
+    std::vector<std::size_t> wires_size_t(wires.begin(), wires.end());
+    this->device_sv->setBasisState(data_vector, wires_size_t);
+}
+
 auto LightningKokkosSimulator::Zero() const -> Result {
     return const_cast<Result>(&GLOBAL_RESULT_FALSE_CONST);
 }

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

@@ -120,6 +120,10 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
     void StopTapeRecording() override;
     void SetDeviceShots(size_t shots) override;
     void SetDevicePRNG(std::mt19937 *) override;
+    void SetState(DataView<std::complex<double>, 1> &,
+                  std::vector<QubitIdType> &) override;
+    void SetBasisState(DataView<int8_t, 1> &,
+                       std::vector<QubitIdType> &) override;
     [[nodiscard]] auto GetDeviceShots() const -> size_t override;
     void PrintState() override;
     [[nodiscard]] auto Zero() const -> Result override;

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/tests/Test_LightningKokkosSimulator.cpp

@@ -659,4 +659,46 @@ TEST_CASE("LightningKokkosSimulator::GateSet", "[GateSet]") {
             expected{3, 0, 1, {"Hadamard", "Hadamard", "IsingZZ"}, {}};
         REQUIRE(LKsim->CacheManagerInfo() == expected);
     }
+
+    SECTION("Test setStateVector") {
+        std::unique_ptr<LKSimulator> LKsim = std::make_unique<LKSimulator>();
+        constexpr std::size_t n_qubits = 2;
+        std::vector<intptr_t> Qs = LKsim->AllocateQubits(n_qubits);
+
+        std::vector<std::complex<double>> data = {{0.5, 0.5}, {0.0, 0.0}};
+        DataView<std::complex<double>, 1> data_view(data);
+        std::vector<QubitIdType> wires = {1};
+        LKsim->SetState(data_view, wires);
+
+        std::vector<std::complex<double>> state(1U << LKsim->GetNumQubits());
+        DataView<std::complex<double>, 1> view(state);
+        LKsim->State(view);
+
+        std::complex<double> c1{0.5, 0.5};
+        std::complex<double> c2{0.0, 0.0};
+        CHECK(state[0] == PLApproxComplex(c1).epsilon(1e-5));
+        CHECK(state[1] == PLApproxComplex(c2).epsilon(1e-5));
+        CHECK(state[2] == PLApproxComplex(c2).epsilon(1e-5));
+        CHECK(state[3] == PLApproxComplex(c2).epsilon(1e-5));
+    }
+
+    SECTION("Test setBasisState") {
+        std::unique_ptr<LKSimulator> LKsim = std::make_unique<LKSimulator>();
+        constexpr std::size_t n_qubits = 1;
+        std::vector<intptr_t> Qs = LKsim->AllocateQubits(n_qubits);
+
+        std::vector<int8_t> data = {0};
+        DataView<int8_t, 1> data_view(data);
+        std::vector<QubitIdType> wires = {0};
+        LKsim->SetBasisState(data_view, wires);
+
+        std::vector<std::complex<double>> state(1U << LKsim->GetNumQubits());
+        DataView<std::complex<double>, 1> view(state);
+        LKsim->State(view);
+
+        std::complex<double> c1{1.0, 0.0};
+        std::complex<double> c2{0.0, 0.0};
+        CHECK(state[0] == PLApproxComplex(c1).epsilon(1e-5));
+        CHECK(state[1] == PLApproxComplex(c2).epsilon(1e-5));
+    }
 }