Skip to content

Commit d3b8bb0

Browse files
tomlqcgithub-actions[bot]vincentmrmaliasadiAmintorDusko
authored and
multiphaseCFD
committed
Add MidMeasureMP Kokkos (#658)
* collapse() and normalized() with Kokkos * Auto update version * CodeFactor * normalize() abort if norm is zero * unittest for normalize() * Auto update version * unittest for normalize() in LighningQubit * Auto update version * Add MCM bindings and tests for L-Kokkos. (#672) * Add MCM bindings and tests for L-Kokkos. * Auto update version * Refactor collapse. * Try parallelizing tests. * Only run test_native_mcm once amongst all workflows since it takes quite a bit of time. * trigger ci --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> * Do not run test_native_mcm with LK-GPU. * Auto update version * trigger ci * Update pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp Co-authored-by: Ali Asadi <[email protected]> * Auto update version * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <[email protected]> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <[email protected]> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <[email protected]> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <[email protected]> * Increase shotsfor flaky test_composite_mcm_single_measure_obs * Create separate device_allowed_operations list for ops supported by the device but not Catalyst. * Use isinstance instead of name comparison. * Check projector. * Update operations at device init. * Modify stopping condition in LK. * Cannot use instance with Hamiltonian check. * Update pennylane_lightning/core/src/simulators/lightning_kokkos/tests/Test_StateVectorLKokkos.cpp Co-authored-by: Amintor Dusko <[email protected]> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Amintor Dusko <[email protected]> * Remove obsolete comment. * Auto update version * trigger ci --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: Ali Asadi <[email protected]> Co-authored-by: Amintor Dusko <[email protected]>
1 parent b76db37 commit d3b8bb0

File tree

14 files changed

+262
-38
lines changed

14 files changed

+262
-38
lines changed

‎.github/CHANGELOG.md

+3
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22

33
### New features since last release
44

5+
* `lightning.kokkos` supports mid-circuit measurements.
6+
[(#672)](https://github.com/PennyLaneAI/pennylane-lightning/pull/672)
7+
58
* Add dynamic linking to LAPACK/OpenBlas shared objects in scipy.libs for both C++ and Python layer.
69
[(#653)](https://github.com/PennyLaneAI/pennylane-lightning/pull/653)
710

‎.github/workflows/tests_gpu_cuda.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ jobs:
311311
OMP_PROC_BIND: false
312312
run: |
313313
cd main/
314-
PL_DEVICE=lightning.qubit python -m pytest tests/ $COVERAGE_FLAGS
314+
PL_DEVICE=lightning.qubit python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
315315
pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
316316
pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
317317
PL_DEVICE=lightning.gpu python -m pytest tests/ $COVERAGE_FLAGS

‎.github/workflows/tests_gpu_kokkos.yml

+3-3
Original file line numberDiff line numberDiff line change
@@ -303,9 +303,9 @@ jobs:
303303
run: |
304304
cd main/
305305
DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
306+
PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS 2> /dev/null || echo Something went wrong with Pytest
306307
pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append 2> /dev/null || echo Something went wrong with pl-device-test shot=20000
307308
pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append 2> /dev/null || echo Something went wrong with pl-device-test shot=None
308-
PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS 2> /dev/null || echo Something went wrong with Pytest
309309
mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
310310
311311
- name: Install all backend devices
@@ -325,10 +325,10 @@ jobs:
325325
OMP_PROC_BIND: false
326326
run: |
327327
cd main/
328-
PL_DEVICE=lightning.qubit python -m pytest tests/ $COVERAGE_FLAGS
328+
PL_DEVICE=lightning.qubit python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
329329
pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
330330
pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
331-
PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS
331+
PL_DEVICE=lightning.kokkos python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS --cov-append
332332
pl-device-test --device lightning.kokkos --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
333333
pl-device-test --device lightning.kokkos --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
334334
mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml

‎.github/workflows/tests_linux.yml

+8-6
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ env:
2121
COVERAGE_FLAGS: "--cov=pennylane_lightning --cov-report=term-missing --no-flaky-report -p no:warnings --tb=native"
2222
GCC_VERSION: 11
2323
OMP_NUM_THREADS: "2"
24+
OMP_PROC_BIND: "false"
2425

2526
concurrency:
2627
group: tests_linux-${{ github.ref }}-${{ inputs.lightning-version }}-${{ inputs.pennylane-version }}
@@ -353,8 +354,8 @@ jobs:
353354
python -m pip install pytest-xdist
354355
cd main/
355356
DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
356-
OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "not unitary_correct" $COVERAGE_FLAGS
357-
PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "unitary_correct" $COVERAGE_FLAGS --cov-append
357+
OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "not unitary_correct and not test_native_mcm" $COVERAGE_FLAGS
358+
PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "unitary_correct and not test_native_mcm" $COVERAGE_FLAGS --cov-append
358359
pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
359360
pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
360361
mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -558,7 +559,8 @@ jobs:
558559
python -m pip install pytest-xdist
559560
cd main/
560561
DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
561-
PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
562+
PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
563+
OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "test_native_mcm" $COVERAGE_FLAGS --cov-append
562564
pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
563565
pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
564566
mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -579,11 +581,11 @@ jobs:
579581
# TODO: Remove installing pytest-xdist with release v0.36.0
580582
python -m pip install pytest-xdist
581583
cd main/
582-
OMP_NUM_THREADS=1 PL_DEVICE=lightning.qubit python -m pytest -n auto tests/ -k "not unitary_correct" $COVERAGE_FLAGS
583-
PL_DEVICE=lightning.qubit python -m pytest tests/ -k "unitary_correct" $COVERAGE_FLAGS --cov-append
584+
OMP_NUM_THREADS=1 PL_DEVICE=lightning.qubit python -m pytest -n auto tests/ -k "not unitary_correct and not test_native_mcm" $COVERAGE_FLAGS
585+
PL_DEVICE=lightning.qubit python -m pytest tests/ -k "unitary_correct and not test_native_mcm" $COVERAGE_FLAGS --cov-append
584586
pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
585587
pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
586-
PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS --cov-append
588+
PL_DEVICE=lightning.kokkos python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS --cov-append
587589
pl-device-test --device lightning.kokkos --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
588590
pl-device-test --device lightning.kokkos --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
589591
mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}

‎pennylane_lightning/core/lightning_base.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -94,9 +94,9 @@ def stopping_condition(self):
9494
and observable) and returns ``True`` if supported by the device."""
9595

9696
def accepts_obj(obj):
97-
if obj.name == "QFT":
97+
if isinstance(obj, qml.QFT):
9898
return len(obj.wires) < 10
99-
if obj.name == "GroverOperator":
99+
if isinstance(obj, qml.GroverOperator):
100100
return len(obj.wires) < 13
101101
is_not_tape = not isinstance(obj, qml.tape.QuantumTape)
102102
is_supported = getattr(self, "supports_operation", lambda name: False)(obj.name)

‎pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+47
Original file line numberDiff line numberDiff line change
@@ -757,6 +757,53 @@ class StateVectorKokkos final
757757
return -static_cast<fp_t>(0.5);
758758
}
759759

760+
/**
761+
* @brief Collapse the state vector after having measured one of the
762+
* qubits.
763+
*
764+
* The branch parameter imposes the measurement result on the given wire.
765+
*
766+
* @param wire Wire to collapse.
767+
* @param branch Branch 0 or 1.
768+
*/
769+
void collapse(std::size_t wire, bool branch) {
770+
KokkosVector matrix("gate_matrix", 4);
771+
Kokkos::parallel_for(
772+
matrix.size(), KOKKOS_LAMBDA(std::size_t k) {
773+
matrix(k) = ((k == 0 && branch == 0) || (k == 3 && branch == 1))
774+
? ComplexT{1.0, 0.0}
775+
: ComplexT{0.0, 0.0};
776+
});
777+
applyMultiQubitOp(matrix, {wire}, false);
778+
normalize();
779+
}
780+
781+
/**
782+
* @brief Normalize vector (to have norm 1).
783+
*/
784+
void normalize() {
785+
auto sv_view = getView();
786+
787+
PrecisionT squaredNorm = 0.0;
788+
Kokkos::parallel_reduce(
789+
sv_view.size(),
790+
KOKKOS_LAMBDA(std::size_t i, PrecisionT & sum) {
791+
const PrecisionT norm = Kokkos::abs(sv_view(i));
792+
sum += norm * norm;
793+
},
794+
squaredNorm);
795+
796+
PL_ABORT_IF(squaredNorm <
797+
std::numeric_limits<PrecisionT>::epsilon() * 1e2,
798+
"vector has norm close to zero and can't be normalized");
799+
800+
const std::complex<PrecisionT> inv_norm =
801+
1. / Kokkos::sqrt(squaredNorm);
802+
Kokkos::parallel_for(
803+
sv_view.size(),
804+
KOKKOS_LAMBDA(std::size_t i) { sv_view(i) *= inv_norm; });
805+
}
806+
760807
/**
761808
* @brief Update data of the class
762809
*

‎pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp

+5-1
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,11 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
137137
sv.applyOperation(str, wires, inv, std::vector<ParamT>{},
138138
conv_matrix);
139139
},
140-
"Apply operation via the gate matrix");
140+
"Apply operation via the gate matrix")
141+
.def("collapse", &StateVectorT::collapse,
142+
"Collapse the statevector onto the 0 or 1 branch of a given wire.")
143+
.def("normalize", &StateVectorT::normalize,
144+
"Normalize the statevector to norm 1.");
141145
}
142146

143147
/**

‎pennylane_lightning/core/src/simulators/lightning_kokkos/gates/GateFunctorsGenerator.hpp

+1
Original file line numberDiff line numberDiff line change
@@ -972,4 +972,5 @@ template <class PrecisionT, bool adj = false> struct generatorMultiRZFunctor {
972972
1 - 2 * int(Kokkos::Impl::bit_count(k & wires_parity) % 2));
973973
}
974974
};
975+
975976
} // namespace Pennylane::LightningKokkos::Functors

‎pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp

+5-1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414
#pragma once
15+
#include <chrono>
1516

1617
#include <Kokkos_Core.hpp>
1718
#include <Kokkos_Random.hpp>
@@ -682,7 +683,10 @@ class Measurements final
682683
});
683684

684685
// Sampling using Random_XorShift64_Pool
685-
Kokkos::Random_XorShift64_Pool<> rand_pool(5374857);
686+
Kokkos::Random_XorShift64_Pool<> rand_pool(
687+
std::chrono::high_resolution_clock::now()
688+
.time_since_epoch()
689+
.count());
686690

687691
Kokkos::parallel_for(
688692
Kokkos::RangePolicy<KokkosExecSpace>(0, num_samples),

‎pennylane_lightning/core/src/simulators/lightning_kokkos/tests/Test_StateVectorLKokkos.cpp

+69-1
Original file line numberDiff line numberDiff line change
@@ -308,4 +308,72 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::StateVectorKokkos",
308308
REQUIRE(sv.getDataVector() == data_);
309309
// REQUIRE(sv.getDataVector() == approx(st_data));
310310
}
311-
}
311+
}
312+
313+
TEMPLATE_TEST_CASE("StateVectorKokkos::collapse", "[StateVectorKokkos]", float,
314+
double) {
315+
using PrecisionT = TestType;
316+
using ComplexT = typename StateVectorKokkos<PrecisionT>::ComplexT;
317+
using TestVectorT = TestVector<ComplexT>;
318+
319+
const std::size_t num_qubits = 3;
320+
321+
// TODO @tomlqc use same template for testing all Lightning flavours?
322+
323+
SECTION("Collapse the state vector after having measured one of the "
324+
"qubits.") {
325+
TestVectorT init_state = createPlusState_<ComplexT>(num_qubits);
326+
327+
const ComplexT coef{0.5, PrecisionT{0.0}};
328+
const ComplexT zero{PrecisionT{0.0}, PrecisionT{0.0}};
329+
330+
std::vector<std::vector<std::vector<ComplexT>>> expected_state = {
331+
{{coef, coef, coef, coef, zero, zero, zero, zero},
332+
{coef, coef, zero, zero, coef, coef, zero, zero},
333+
{coef, zero, coef, zero, coef, zero, coef, zero}},
334+
{{zero, zero, zero, zero, coef, coef, coef, coef},
335+
{zero, zero, coef, coef, zero, zero, coef, coef},
336+
{zero, coef, zero, coef, zero, coef, zero, coef}},
337+
};
338+
339+
std::size_t wire = GENERATE(0, 1, 2);
340+
std::size_t branch = GENERATE(0, 1);
341+
StateVectorKokkos<PrecisionT> sv(
342+
reinterpret_cast<ComplexT *>(init_state.data()), init_state.size());
343+
sv.collapse(wire, branch);
344+
345+
PrecisionT eps = std::numeric_limits<PrecisionT>::epsilon() * 10e3;
346+
REQUIRE(isApproxEqual(sv.getData(), sv.getDataVector().size(),
347+
expected_state[branch][wire].data(),
348+
expected_state[branch][wire].size(), eps));
349+
}
350+
}
351+
352+
TEMPLATE_TEST_CASE("StateVectorKokkos::normalize", "[StateVectorKokkos]", float,
353+
double) {
354+
using PrecisionT = TestType;
355+
using ComplexT = typename StateVectorKokkos<PrecisionT>::ComplexT;
356+
357+
// TODO @tomlqc use same template for testing all Lightning flavours?
358+
359+
SECTION("Normalize state vector.") {
360+
const ComplexT init{1.0, PrecisionT{0.0}};
361+
const ComplexT half{0.5, PrecisionT{0.0}};
362+
const ComplexT zero{PrecisionT{0.0}, PrecisionT{0.0}};
363+
364+
std::vector<ComplexT> init_state = {init, zero, init, init,
365+
zero, zero, init, zero};
366+
367+
std::vector<ComplexT> expected_state = {half, zero, half, half,
368+
zero, zero, half, zero};
369+
370+
StateVectorKokkos<PrecisionT> sv(
371+
reinterpret_cast<ComplexT *>(init_state.data()), init_state.size());
372+
sv.normalize();
373+
374+
PrecisionT eps = std::numeric_limits<PrecisionT>::epsilon() * 1e3;
375+
REQUIRE(isApproxEqual(sv.getData(), sv.getDataVector().size(),
376+
expected_state.data(), expected_state.size(),
377+
eps));
378+
}
379+
}

‎pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubit.cpp

+27-1
Original file line numberDiff line numberDiff line change
@@ -245,9 +245,35 @@ TEMPLATE_PRODUCT_TEST_CASE("StateVectorLQubit::collapse", "[StateVectorLQubit]",
245245

246246
std::size_t wire = GENERATE(0, 1, 2);
247247
std::size_t branch = GENERATE(0, 1);
248-
StateVectorLQubitManaged<PrecisionT> sv(init_state);
248+
StateVectorT sv(init_state.data(), init_state.size());
249249
sv.collapse(wire, branch);
250250

251251
REQUIRE(sv.getDataVector() == approx(expected_state[branch][wire]));
252252
}
253253
}
254+
255+
TEMPLATE_PRODUCT_TEST_CASE("StateVectorLQubit::normalize",
256+
"[StateVectorLQubit]",
257+
(StateVectorLQubitManaged, StateVectorLQubitRaw),
258+
(float, double)) {
259+
using StateVectorT = TestType;
260+
using PrecisionT = typename StateVectorT::PrecisionT;
261+
using ComplexT = typename StateVectorT::ComplexT;
262+
263+
SECTION("Normalize state vector.") {
264+
const ComplexT init{1.0, PrecisionT{0.0}};
265+
const ComplexT half{0.5, PrecisionT{0.0}};
266+
const ComplexT zero{PrecisionT{0.0}, PrecisionT{0.0}};
267+
268+
std::vector<ComplexT> init_state = {init, zero, init, init,
269+
zero, zero, init, zero};
270+
271+
std::vector<ComplexT> expected_state = {half, zero, half, half,
272+
zero, zero, half, zero};
273+
274+
StateVectorT sv(init_state.data(), init_state.size());
275+
sv.normalize();
276+
277+
REQUIRE(sv.getDataVector() == approx(expected_state));
278+
}
279+
}

‎pennylane_lightning/core/src/utils/TestHelpers.hpp

+13
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,19 @@ auto createZeroState(size_t num_qubits) -> TestVector<ComplexT> {
288288
return res;
289289
}
290290

291+
/**
292+
* @brief create |+>^N
293+
*/
294+
template <typename ComplexT>
295+
auto createPlusState_(size_t num_qubits) -> TestVector<ComplexT> {
296+
TestVector<ComplexT> res(size_t{1U} << num_qubits, 1.0,
297+
getBestAllocator<ComplexT>());
298+
for (auto &elem : res) {
299+
elem /= std::sqrt(1U << num_qubits);
300+
}
301+
return res;
302+
}
303+
291304
/**
292305
* @brief create |+>^N
293306
*/

0 commit comments

Comments
 (0)