apply features changes brought by PR#927

pennylane_lightning/core/src/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

@@ -309,13 +309,20 @@ void LightningGPUSimulator::PartialProbs(
     std::move(dv_probs.begin(), dv_probs.end(), probs.begin());
 }
 
-void LightningGPUSimulator::Sample(DataView<double, 2> &samples,
-                                   std::size_t shots) {
+std::vector<size_t> LightningGPUSimulator::GenerateSamples(size_t shots) {
+    // generate_samples is a member function of the Measures class.
     Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
-    // PL-Lightning-GPU generates samples using the alias method.
-    // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+
+    if (this->gen) {
+        return m.generate_samples(shots, (*(this->gen))());
+    }
+    return m.generate_samples(shots);
+}
+
+void LightningGPUSimulator::Sample(DataView<double, 2> &samples,
+                                   std::size_t shots) {
+    auto li_samples = this->GenerateSamples(shots);
 
     RT_FAIL_IF(samples.size() != li_samples.size(),
                "Invalid size for the pre-allocated samples");
@@ -348,13 +355,7 @@ void LightningGPUSimulator::PartialSample(DataView<double, 2> &samples,
     // get device wires
     auto &&dev_wires = getDeviceWires(wires);
 
-    // generate_samples is a member function of the MeasuresGPU class.
-    Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
-        *(this->device_sv)};
-
-    // PL-Lightning-GPU generates samples using the alias method.
-    // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+    auto li_samples = this->GenerateSamples(shots);
 
     // The lightning samples are layed out as a single vector of size
     // shots*qubits, where each element represents a single bit. The
@@ -378,13 +379,7 @@ void LightningGPUSimulator::Counts(DataView<double, 1> &eigvals,
     RT_FAIL_IF(eigvals.size() != numElements || counts.size() != numElements,
                "Invalid size for the pre-allocated counts");
 
-    // generate_samples is a member function of the MeasuresGPU class.
-    Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
-        *(this->device_sv)};
-
-    // PL-Lightning-GPU generates samples using the alias method.
-    // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+    auto li_samples = this->GenerateSamples(shots);
 
     // Fill the eigenvalues with the integer representation of the corresponding
     // computational basis bitstring. In the future, eigenvalues can also be
@@ -423,13 +418,7 @@ void LightningGPUSimulator::PartialCounts(DataView<double, 1> &eigvals,
     // get device wires
     auto &&dev_wires = getDeviceWires(wires);
 
-    // generate_samples is a member function of the MeasuresGPU class.
-    Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
-        *(this->device_sv)};
-
-    // PL-Lightning-GPU generates samples using the alias method.
-    // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+    auto li_samples = this->GenerateSamples(shots);
 
     // Fill the eigenvalues with the integer representation of the corresponding
     // computational basis bitstring. In the future, eigenvalues can also be

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

@@ -97,6 +97,8 @@ class LightningGPUSimulator final : public Catalyst::Runtime::QuantumDevice {
         return res;
     }
 
+    auto GenerateSamples(size_t shots) -> std::vector<size_t>;
+
   public:
     explicit LightningGPUSimulator(const std::string &kwargs = "{}") {
         auto &&args = Catalyst::Runtime::parse_kwargs(kwargs);

pennylane_lightning/core/src/simulators/lightning_gpu/catalyst/tests/Test_LightningGPUMeasures.cpp

@@ -1754,26 +1754,71 @@ TEST_CASE("Counts and PartialCounts tests with numWires=0-4 shots=100",
 }
 
 TEST_CASE("Measurement with a seeded device", "[Measures]") {
-    for (std::size_t _ = 0; _ < 5; _++) {
-        std::unique_ptr<LGPUSimulator> sim = std::make_unique<LGPUSimulator>();
-        std::unique_ptr<LGPUSimulator> sim1 = std::make_unique<LGPUSimulator>();
+    std::array<std::unique_ptr<LGPUSimulator>, 2> sims;
+    std::vector<std::mt19937> gens{std::mt19937{37}, std::mt19937{37}};
 
-        std::mt19937 gen(37);
-        sim->SetDevicePRNG(&gen);
+    auto circuit = [](LGPUSimulator &sim, std::mt19937 &gen) {
+        sim.SetDevicePRNG(&gen);
         std::vector<intptr_t> Qs;
         Qs.reserve(1);
-        Qs.push_back(sim->AllocateQubit());
-        sim->NamedOperation("Hadamard", {}, {Qs[0]}, false);
-        auto m = sim->Measure(Qs[0]);
-
-        std::mt19937 gen1(37);
-        sim1->SetDevicePRNG(&gen1);
-        std::vector<intptr_t> Qs1;
-        Qs1.reserve(1);
-        Qs1.push_back(sim1->AllocateQubit());
-        sim1->NamedOperation("Hadamard", {}, {Qs1[0]}, false);
-        auto m1 = sim1->Measure(Qs1[0]);
-
-        CHECK(*m == *m1);
+        Qs.push_back(sim.AllocateQubit());
+        sim.NamedOperation("Hadamard", {}, {Qs[0]}, false);
+        auto m = sim.Measure(Qs[0]);
+        return m;
+    };
+
+    for (std::size_t trial = 0; trial < 5; trial++) {
+        sims[0] = std::make_unique<LGPUSimulator>();
+        sims[1] = std::make_unique<LGPUSimulator>();
+
+        auto m0 = circuit(*(sims[0]), gens[0]);
+        auto m1 = circuit(*(sims[1]), gens[1]);
+
+        CHECK(*m0 == *m1);
+    }
+}
+
+TEST_CASE("Sample with a seeded device", "[Measures]") {
+    std::size_t shots = 100;
+    std::array<std::unique_ptr<LGPUSimulator>, 2> sims;
+    std::vector<std::vector<double>> sample_vec(2,
+                                                std::vector<double>(shots * 4));
+
+    std::vector<MemRefT<double, 2>> buffers{
+        MemRefT<double, 2>{
+            sample_vec[0].data(), sample_vec[0].data(), 0, {shots, 1}, {1, 1}},
+        MemRefT<double, 2>{
+            sample_vec[1].data(), sample_vec[1].data(), 0, {shots, 1}, {1, 1}},
+    };
+    std::vector<DataView<double, 2>> views{
+        DataView<double, 2>(buffers[0].data_aligned, buffers[0].offset,
+                            buffers[0].sizes, buffers[0].strides),
+        DataView<double, 2>(buffers[1].data_aligned, buffers[1].offset,
+                            buffers[1].sizes, buffers[1].strides)};
+
+    std::vector<std::mt19937> gens{std::mt19937{37}, std::mt19937{37}};
+
+    auto circuit = [shots](LGPUSimulator &sim, DataView<double, 2> &view,
+                           std::mt19937 &gen) {
+        sim.SetDevicePRNG(&gen);
+        std::vector<intptr_t> Qs;
+        Qs.reserve(1);
+        Qs.push_back(sim.AllocateQubit());
+        sim.NamedOperation("Hadamard", {}, {Qs[0]}, false);
+        sim.NamedOperation("RX", {0.5}, {Qs[0]}, false);
+        sim.Sample(view, shots);
+    };
+
+    for (std::size_t trial = 0; trial < 5; trial++) {
+        sims[0] = std::make_unique<LGPUSimulator>();
+        sims[1] = std::make_unique<LGPUSimulator>();
+
+        for (std::size_t sim_idx = 0; sim_idx < sims.size(); sim_idx++) {
+            circuit(*(sims[sim_idx]), views[sim_idx], gens[sim_idx]);
+        }
+
+        for (std::size_t i = 0; i < sample_vec[0].size(); i++) {
+            CHECK((sample_vec[0][i] == sample_vec[1][i]));
+        }
     }
 }