diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp index 3d0000d08..946069cd5 100644 --- a/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp @@ -489,8 +489,75 @@ class StateVectorKokkos final PL_ABORT_IF(gate_matrix.empty(), std::string("Operation does not exist for ") + opName + std::string(" and no matrix provided.")); - PL_ABORT("Controlled matrix operation not yet supported"); - return; + return applyNCMultiQubitOp(vector2view(gate_matrix), + controlled_wires, controlled_values, + wires, inverse); + } + } + + /** + * @brief Apply a controlled-multi qubit operator to the state vector using + * a matrix + * + * @param matrix Kokkos gate matrix in the device space. + * @param controlled_wires Control wires. + * @param controlled_values Control values (true or false). + * @param wires Wires to apply gate to. + * @param inverse Indicates whether to use adjoint of gate. + */ + void applyNCMultiQubitOp(const KokkosVector matrix, + const std::vector &controlled_wires, + const std::vector &controlled_values, + const std::vector &wires, + bool inverse = false) { + + auto &&num_qubits = this->getNumQubits(); + std::size_t two2N = + std::exp2(num_qubits - wires.size() - controlled_wires.size()); + std::size_t dim = std::exp2(wires.size()); + KokkosVector matrix_trans("matrix_trans", matrix.size()); + + if (inverse) { + Kokkos::MDRangePolicy policy_2d({0, 0}, {dim, dim}); + Kokkos::parallel_for( + policy_2d, + KOKKOS_LAMBDA(const std::size_t i, const std::size_t j) { + matrix_trans(i + j * dim) = conj(matrix(i * dim + j)); + }); + } else { + matrix_trans = matrix; + } + + switch (wires.size()) { + case 1: + Kokkos::parallel_for(two2N, applyNC1QubitOpFunctor( + *data_, num_qubits, matrix_trans, + controlled_wires, controlled_values, + wires)); + break; + case 2: + Kokkos::parallel_for(two2N, applyNC2QubitOpFunctor( + *data_, num_qubits, matrix_trans, + controlled_wires, controlled_values, + wires)); + break; + case 3: + Kokkos::parallel_for(two2N, applyNC3QubitOpFunctor( + *data_, num_qubits, matrix_trans, + controlled_wires, controlled_values, + wires)); + break; + default: + std::size_t scratch_size = ScratchViewComplex::shmem_size(dim) + + ScratchViewSizeT::shmem_size(dim); + Kokkos::parallel_for( + "multiNCQubitOpFunctor", + TeamPolicy(two2N, Kokkos::AUTO, dim) + .set_scratch_size(0, Kokkos::PerTeam(scratch_size)), + NCMultiQubitOpFunctor( + *data_, num_qubits, matrix_trans, controlled_wires, + controlled_values, wires)); + break; } } @@ -546,7 +613,56 @@ class StateVectorKokkos final "number of wires"); applyMatrix(matrix.data(), wires, inverse); } + inline void applyControlledMatrix( + ComplexT *matrix, const std::vector &controlled_wires, + const std::vector &controlled_values, + const std::vector &wires, bool inverse = false) { + PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0"); + std::size_t n = static_cast(1U) << wires.size(); + KokkosVector matrix_(matrix, n * n); + applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires, + inverse); + } + inline void + applyControlledMatrix(const ComplexT *matrix, + const std::vector &controlled_wires, + const std::vector &controlled_values, + const std::vector &wires, + bool inverse = false) { + PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0"); + std::size_t n = static_cast(1U) << wires.size(); + std::size_t n2 = n * n; + KokkosVector matrix_("matrix_", n2); + Kokkos::deep_copy(matrix_, UnmanagedConstComplexHostView(matrix, n2)); + applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires, + inverse); + } + + /** + * @brief Apply a given controlled-matrix directly to the statevector. + * + * @param matrix Vector containing the statevector data (in row-major + * format). + * @param controlled_wires Control wires. + * @param controlled_values Control values (false or true). + * @param wires Wires to apply gate to. + * @param inverse Indicate whether inverse should be taken. + */ + inline void + applyControlledMatrix(const std::vector &matrix, + const std::vector &controlled_wires, + const std::vector &controlled_values, + const std::vector &wires, + bool inverse = false) { + PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0"); + PL_ABORT_IF(matrix.size() != exp2(2 * wires.size()), + "The size of matrix does not match with the given " + "number of wires"); + applyControlledMatrix(matrix.data(), controlled_wires, + controlled_values, wires, inverse); + } + /** * @brief Apply a single generator to the state vector using the given * kernel. diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp index 0fb6f5d86..d908a80a7 100644 --- a/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp @@ -47,7 +47,26 @@ namespace Pennylane::LightningKokkos { using StateVectorBackends = Pennylane::Util::TypeList, StateVectorKokkos, void>; +/** + * @brief Register controlled matrix kernel. + */ +template +void applyControlledMatrix( + StateVectorT &st, + const py::array_t, + py::array::c_style | py::array::forcecast> &matrix, + const std::vector &controlled_wires, + const std::vector &controlled_values, + const std::vector &wires, bool inverse = false) { + using ComplexT = typename StateVectorT::ComplexT; + st.applyControlledMatrix( + static_cast(matrix.request().ptr), controlled_wires, + controlled_values, wires, inverse); +} +/** + * @brief Register controlled gates. + */ template void registerControlledGate(PyClass &pyclass) { using PrecisionT = @@ -172,7 +191,9 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) { .def("collapse", &StateVectorT::collapse, "Collapse the statevector onto the 0 or 1 branch of a given wire.") .def("normalize", &StateVectorT::normalize, - "Normalize the statevector to norm 1."); + "Normalize the statevector to norm 1.") + .def("applyControlledMatrix", &applyControlledMatrix, + "Apply controlled operation"); } /** diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp index 575e215d2..3d7a7603b 100644 --- a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp @@ -23,7 +23,10 @@ namespace { using namespace Pennylane::Util; using Kokkos::Experimental::swap; +using Pennylane::LightningKokkos::Util::generateControlBitPatterns; using Pennylane::LightningKokkos::Util::one; +using Pennylane::LightningKokkos::Util::parity_2_offset; +using Pennylane::LightningKokkos::Util::reverseWires; using Pennylane::LightningKokkos::Util::vector2view; using Pennylane::LightningKokkos::Util::wires2Parity; using std::size_t; @@ -104,6 +107,67 @@ template struct multiQubitOpFunctor { } }; +template struct NCMultiQubitOpFunctor { + using KokkosComplexVector = Kokkos::View *>; + using KokkosIntVector = Kokkos::View; + using ScratchViewComplex = + Kokkos::View *, + Kokkos::DefaultExecutionSpace::scratch_memory_space, + Kokkos::MemoryTraits>; + using ScratchViewSizeT = + Kokkos::View>; + using MemberType = Kokkos::TeamPolicy<>::member_type; + + KokkosComplexVector arr; + KokkosComplexVector matrix; + KokkosIntVector indices; + KokkosIntVector parity; + KokkosIntVector rev_wires; + KokkosIntVector rev_wire_shifts; + std::size_t dim; + std::size_t num_qubits; + + NCMultiQubitOpFunctor(KokkosComplexVector arr_, std::size_t num_qubits_, + const KokkosComplexVector &matrix_, + const std::vector &controlled_wires_, + const std::vector &controlled_values_, + const std::vector &wires_) { + dim = one << wires_.size(); + arr = arr_; + matrix = matrix_; + num_qubits = num_qubits_; + std::tie(parity, rev_wires) = + reverseWires(num_qubits_, wires_, controlled_wires_); + indices = generateControlBitPatterns(num_qubits_, controlled_wires_, + controlled_values_, wires_); + } + + KOKKOS_INLINE_FUNCTION + void operator()(const MemberType &teamMember) const { + const std::size_t k = teamMember.league_rank(); + ScratchViewComplex coeffs_in(teamMember.team_scratch(0), dim); + const std::size_t offset = parity_2_offset(parity, k); + if (teamMember.team_rank() == 0) { + for (std::size_t i = 0; i < dim; i++) { + coeffs_in(i) = arr(indices(i) + offset); + } + } + + Kokkos::parallel_for( + Kokkos::TeamThreadRange(teamMember, dim), [&](const std::size_t i) { + const auto idx = indices(i) + offset; + arr(idx) = 0.0; + const std::size_t base_idx = i * dim; + + for (std::size_t j = 0; j < dim; j++) { + arr(idx) += matrix(base_idx + j) * coeffs_in(j); + } + }); + } +}; + template struct apply1QubitOpFunctor { using ComplexT = Kokkos::complex; using KokkosComplexVector = Kokkos::View; @@ -142,6 +206,47 @@ template struct apply1QubitOpFunctor { arr(i1) = matrix(0B10) * v0 + matrix(0B11) * v1; } }; + +template struct applyNC1QubitOpFunctor { + using ComplexT = Kokkos::complex; + using KokkosComplexVector = Kokkos::View; + using KokkosIntVector = Kokkos::View; + + KokkosComplexVector arr; + KokkosComplexVector matrix; + KokkosIntVector indices; + KokkosIntVector parity; + KokkosIntVector rev_wires; + KokkosIntVector rev_wire_shifts; + std::size_t num_qubits; + + applyNC1QubitOpFunctor(KokkosComplexVector arr_, std::size_t num_qubits_, + const KokkosComplexVector &matrix_, + const std::vector &controlled_wires_, + const std::vector &controlled_values_, + const std::vector &wires_) { + arr = arr_; + matrix = matrix_; + num_qubits = num_qubits_; + std::tie(parity, rev_wires) = + reverseWires(num_qubits_, wires_, controlled_wires_); + indices = generateControlBitPatterns(num_qubits_, controlled_wires_, + controlled_values_, wires_); + } + + KOKKOS_INLINE_FUNCTION + void operator()(const std::size_t k) const { + const std::size_t offset = parity_2_offset(parity, k); + std::size_t i0 = indices(0B00); + std::size_t i1 = indices(0B01); + ComplexT v0 = arr(i0 + offset); + ComplexT v1 = arr(i1 + offset); + + arr(i0 + offset) = matrix(0B00) * v0 + matrix(0B01) * v1; + arr(i1 + offset) = matrix(0B10) * v0 + matrix(0B11) * v1; + } +}; + template struct apply2QubitOpFunctor { using ComplexT = Kokkos::complex; using KokkosComplexVector = Kokkos::View; @@ -203,6 +308,56 @@ template struct apply2QubitOpFunctor { } }; +template struct applyNC2QubitOpFunctor { + using ComplexT = Kokkos::complex; + using KokkosComplexVector = Kokkos::View; + using KokkosIntVector = Kokkos::View; + + KokkosComplexVector arr; + KokkosComplexVector matrix; + KokkosIntVector indices; + KokkosIntVector parity; + KokkosIntVector rev_wires; + KokkosIntVector rev_wire_shifts; + std::size_t num_qubits; + + applyNC2QubitOpFunctor(KokkosComplexVector arr_, std::size_t num_qubits_, + const KokkosComplexVector &matrix_, + const std::vector &controlled_wires_, + const std::vector &controlled_values_, + const std::vector &wires_) { + arr = arr_; + matrix = matrix_; + num_qubits = num_qubits_; + std::tie(parity, rev_wires) = + reverseWires(num_qubits_, wires_, controlled_wires_); + indices = generateControlBitPatterns(num_qubits_, controlled_wires_, + controlled_values_, wires_); + } + + KOKKOS_INLINE_FUNCTION + void operator()(const std::size_t k) const { + const std::size_t offset = parity_2_offset(parity, k); + std::size_t i00 = indices(0B00); + std::size_t i01 = indices(0B01); + std::size_t i10 = indices(0B10); + std::size_t i11 = indices(0B11); + ComplexT v00 = arr(i00 + offset); + ComplexT v01 = arr(i01 + offset); + ComplexT v10 = arr(i10 + offset); + ComplexT v11 = arr(i11 + offset); + + arr(i00 + offset) = matrix(0B0000) * v00 + matrix(0B0001) * v01 + + matrix(0B0010) * v10 + matrix(0B0011) * v11; + arr(i01 + offset) = matrix(0B0100) * v00 + matrix(0B0101) * v01 + + matrix(0B0110) * v10 + matrix(0B0111) * v11; + arr(i10 + offset) = matrix(0B1000) * v00 + matrix(0B1001) * v01 + + matrix(0B1010) * v10 + matrix(0B1011) * v11; + arr(i11 + offset) = matrix(0B1100) * v00 + matrix(0B1101) * v01 + + matrix(0B1110) * v10 + matrix(0B1111) * v11; + } +}; + #define GATEENTRY3(xx, yy) xx << 3 | yy #define GATETERM3(xx, yy, vyy) matrix(GATEENTRY3(xx, yy)) * vyy #define GATESUM3(xx) \ @@ -265,6 +420,64 @@ template struct apply3QubitOpFunctor { } }; +template struct applyNC3QubitOpFunctor { + using ComplexT = Kokkos::complex; + using KokkosComplexVector = Kokkos::View; + using KokkosIntVector = Kokkos::View; + + KokkosComplexVector arr; + KokkosComplexVector matrix; + KokkosIntVector indices; + KokkosIntVector parity; + KokkosIntVector rev_wires; + KokkosIntVector rev_wire_shifts; + std::size_t num_qubits; + + applyNC3QubitOpFunctor(KokkosComplexVector arr_, std::size_t num_qubits_, + const KokkosComplexVector &matrix_, + const std::vector &controlled_wires_, + const std::vector &controlled_values_, + const std::vector &wires_) { + arr = arr_; + matrix = matrix_; + num_qubits = num_qubits_; + std::tie(parity, rev_wires) = + reverseWires(num_qubits_, wires_, controlled_wires_); + indices = generateControlBitPatterns(num_qubits_, controlled_wires_, + controlled_values_, wires_); + } + + KOKKOS_INLINE_FUNCTION + void operator()(const std::size_t k) const { + const std::size_t offset = parity_2_offset(parity, k); + std::size_t i000 = indices(0B000); + std::size_t i001 = indices(0B001); + std::size_t i010 = indices(0B010); + std::size_t i011 = indices(0B011); + std::size_t i100 = indices(0B100); + std::size_t i101 = indices(0B101); + std::size_t i110 = indices(0B110); + std::size_t i111 = indices(0B111); + ComplexT v000 = arr(i000 + offset); + ComplexT v001 = arr(i001 + offset); + ComplexT v010 = arr(i010 + offset); + ComplexT v011 = arr(i011 + offset); + ComplexT v100 = arr(i100 + offset); + ComplexT v101 = arr(i101 + offset); + ComplexT v110 = arr(i110 + offset); + ComplexT v111 = arr(i111 + offset); + + arr(i000 + offset) = GATESUM3(0B000); + arr(i001 + offset) = GATESUM3(0B001); + arr(i010 + offset) = GATESUM3(0B010); + arr(i011 + offset) = GATESUM3(0B011); + arr(i100 + offset) = GATESUM3(0B100); + arr(i101 + offset) = GATESUM3(0B101); + arr(i110 + offset) = GATESUM3(0B110); + arr(i111 + offset) = GATESUM3(0B111); + } +}; + #define GATEENTRY4(xx, yy) xx << 4 | yy #define GATETERM4(xx, yy, vyy) matrix(GATEENTRY4(xx, yy)) * vyy #define GATESUM4(xx) \ diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_NonParam.cpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_NonParam.cpp index d47be12c4..5f2b319b4 100644 --- a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_NonParam.cpp +++ b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_NonParam.cpp @@ -704,6 +704,62 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyCSWAP", } } + +TEMPLATE_TEST_CASE("StateVectorKokkos::applyMatrix/Controlled-Operation", + "[StateVectorKokkos_Nonparam]", float, double) { + using StateVectorT = StateVectorKokkos; + using PrecisionT = StateVectorT::PrecisionT; + + const std::size_t num_qubits = 5; + const TestType EP = 1e-4; + auto ini_st = createNonTrivialState(num_qubits); + + std::unordered_map str_to_gates_{}; + for (const auto &[gate_op, gate_name] : Constant::gate_names) { + str_to_gates_.emplace(gate_name, gate_op); + } + + std::unordered_map + str_to_controlled_gates_{}; + for (const auto &[gate_op, controlled_gate_name] : + Constant::controlled_gate_names) { + str_to_controlled_gates_.emplace(controlled_gate_name, gate_op); + } + + const bool inverse = GENERATE(false, true); + const std::string gate_name = + GENERATE("PauliX", "PauliY", "PauliZ", "Hadamard", "S", "T", "SWAP"); + DYNAMIC_SECTION("N-controlled Matrix - Gate = " + << gate_name << " Inverse = " << inverse) { + auto gate_matrix = getMatrix( + str_to_gates_.at(gate_name), {}, false); + + std::vector controlled_wires = {4}; + std::vector controlled_values = {true}; + + StateVectorT kokkos_sv_ops{ini_st.data(), ini_st.size()}; + StateVectorT kokkos_sv_mat{ini_st.data(), ini_st.size()}; + + const auto wires = + createWires(str_to_controlled_gates_.at(gate_name), num_qubits); + kokkos_sv_ops.applyOperation(gate_name, controlled_wires, + controlled_values, wires, inverse, {}); + kokkos_sv_mat.applyOperation("Matrix", controlled_wires, + controlled_values, wires, inverse, {}, + gate_matrix); + + auto result_ops = kokkos_sv_ops.getDataVector(); + auto result_mat = kokkos_sv_mat.getDataVector(); + + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(real(result_ops[j]) == + Approx(real(result_mat[j])).margin(EP)); + CHECK(imag(result_ops[j]) == + Approx(imag(result_mat[j])).margin(EP)); + } + } +} + TEMPLATE_TEST_CASE("StateVectorKokkos::applyMultiQubitOp", "[StateVectorKokkos_Nonparam][Inverse]", float, double) { const bool inverse = GENERATE(true, false); @@ -801,11 +857,86 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyMultiQubitOp", } } +TEMPLATE_TEST_CASE("StateVectorKokkos::applyNCMultiQubitOp","[StateVectorKokkos_Nonparam][Inverse]", float, double) { + const bool inverse = GENERATE(true, false); + std::size_t num_qubits = 3; + StateVectorKokkos sv_normal{num_qubits}; + StateVectorKokkos sv_mq{num_qubits}; + using UnmanagedComplexHostView = + Kokkos::View *, Kokkos::HostSpace, + Kokkos::MemoryTraits>; + + SECTION("0 Controlled Single Qubit via applyOperation") { + std::vector wires = {0}; + sv_normal.applyOperation("PauliX", wires, inverse); + auto sv_normal_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, sv_normal.getView()); + + std::vector controlled_wire = {}; + std::vector controlled_value = {}; + std::vector wire = {0}; + auto matrix = getPauliX(); + sv_mq.applyOperation("XXXXXXXX", controlled_wire, controlled_value, + wire, inverse, {}, matrix); + auto sv_mq_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, sv_mq.getView()); + + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(imag(sv_normal_host[j]) == Approx(imag(sv_mq_host[j]))); + CHECK(real(sv_normal_host[j]) == Approx(real(sv_mq_host[j]))); + } + } + + SECTION("Single Qubit via applyNCMultiQubitOp") { + std::vector wires = {0, 1}; + sv_normal.applyOperation("CNOT", wires, inverse); + auto sv_normal_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, sv_normal.getView()); + + std::vector controlled_wire = {0}; + std::vector controlled_value = {true}; + std::vector wire = {1}; + auto matrix = getPauliX(); + Kokkos::View *> device_matrix("device_matrix", + matrix.size()); + Kokkos::deep_copy(device_matrix, UnmanagedComplexHostView( + matrix.data(), matrix.size())); + sv_mq.applyNCMultiQubitOp(device_matrix, controlled_wire, + controlled_value, wire, inverse); + auto sv_mq_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, sv_mq.getView()); + + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(imag(sv_normal_host[j]) == Approx(imag(sv_mq_host[j]))); + CHECK(real(sv_normal_host[j]) == Approx(real(sv_mq_host[j]))); + } + } + + SECTION("Controlled Single Qubit via applyOperation") { + std::vector wires = {0, 1}; + sv_normal.applyOperation("CNOT", wires, inverse); + auto sv_normal_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, sv_normal.getView()); + + std::vector controlled_wire = {0}; + std::vector controlled_value = {true}; + std::vector wire = {1}; + auto matrix = getPauliX(); + sv_mq.applyOperation("XXXXXXXX", controlled_wire, controlled_value, + wire, inverse, {}, matrix); + auto sv_mq_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, sv_mq.getView()); + + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(imag(sv_normal_host[j]) == Approx(imag(sv_mq_host[j]))); + CHECK(real(sv_normal_host[j]) == Approx(real(sv_mq_host[j]))); + } + } +} + TEMPLATE_TEST_CASE("StateVectorKokkos::applyOperation non-param " "one-qubit with controls", "[StateVectorKokkos_NonParam]", float, double) { - // using PrecisionT = TestType; - // using ComplexT = StateVectorKokkos::ComplexT; const bool inverse = GENERATE(true, false); const std::size_t num_qubits = 4; const std::size_t control = GENERATE(0, 1, 2, 3); @@ -994,6 +1125,55 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyOperation non-param " } } } + + SECTION("N-controlled SWAP with matrix") { + if (control != wire0 && control != wire1 && wire0 != wire1) { + StateVectorT kokkos_sv0{ini_st.data(), ini_st.size()}; + StateVectorT kokkos_sv1{ini_st.data(), ini_st.size()}; + kokkos_sv0.applyOperation("CSWAP", {control, wire0, wire1}, + inverse); + auto matrix = getSWAP(); + kokkos_sv1.applyOperation( + "XXXXXXXX", std::vector{control}, + std::vector{true}, std::vector{wire0, wire1}, + inverse, {}, matrix); + auto result_sv0 = kokkos_sv0.getDataVector(); + auto result_sv1 = kokkos_sv1.getDataVector(); + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(real(result_sv0[j]) == + Approx(real(result_sv1[j])).margin(EP)); + CHECK(imag(result_sv0[j]) == + Approx(imag(result_sv1[j])).margin(EP)); + } + } + } +} + +TEMPLATE_TEST_CASE("StateVectorKokkos::applyOperation controlled Toffoli", + "[StateVectorKokkos_NonParam]", float, double) { + using StateVectorT = StateVectorKokkos; + + const TestType EP = 1e-4; + const std::size_t num_qubits = 6; + const bool inverse = GENERATE(true, false); + const std::size_t control = GENERATE(0, 1, 2); + + auto ini_st = createNonTrivialState(num_qubits); + StateVectorT kokkos_sv0{ini_st.data(), ini_st.size()}; + StateVectorT kokkos_sv1{ini_st.data(), ini_st.size()}; + auto matrix = getToffoli(); + kokkos_sv0.applyOperation( + "Matrix", std::vector{control}, std::vector{true}, + std::vector{3, 4, 5}, inverse, {}, matrix); + kokkos_sv1.applyOperation("PauliX", std::vector{control, 3, 4}, + std::vector{true, true, true}, + std::vector{5}, inverse); + auto result_sv0 = kokkos_sv0.getDataVector(); + auto result_sv1 = kokkos_sv1.getDataVector(); + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(real(result_sv0[j]) == Approx(real(result_sv1[j])).margin(EP)); + CHECK(imag(result_sv0[j]) == Approx(imag(result_sv1[j])).margin(EP)); + } } TEMPLATE_TEST_CASE("StateVectorKokkos::SetStateVector", diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Param.cpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Param.cpp index 829946332..2b79169f0 100644 --- a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Param.cpp +++ b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Param.cpp @@ -116,6 +116,70 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyMatrix/Operation", } } + +TEMPLATE_TEST_CASE("StateVectorKokkos::applyMatrix/Controlled Operation", + "[StateVectorKokkos_Operation]", float, double) { + using StateVectorT = StateVectorKokkos; + using PrecisionT = StateVectorT::PrecisionT; + + const std::size_t num_qubits = 5; + const TestType EP = 1e-4; + const TestType param = 0.12342; + auto ini_st = createNonTrivialState(num_qubits); + + std::unordered_map str_to_gates_{}; + for (const auto &[gate_op, gate_name] : Constant::gate_names) { + str_to_gates_.emplace(gate_name, gate_op); + } + + std::unordered_map + str_to_controlled_gates_{}; + for (const auto &[gate_op, controlled_gate_name] : + Constant::controlled_gate_names) { + str_to_controlled_gates_.emplace(controlled_gate_name, gate_op); + } + + const bool inverse = GENERATE(false, true); + const std::string gate_name = GENERATE( + "PhaseShift", "RX", "RY", "RZ", "Rot", "IsingXX", "IsingXY", "IsingYY", + "IsingZZ", "SingleExcitation", "SingleExcitationMinus", + "SingleExcitationPlus", "DoubleExcitation", "DoubleExcitationMinus", + "DoubleExcitationPlus"); + DYNAMIC_SECTION("N-controlled Matrix - Gate = " + << gate_name << " Inverse = " << inverse) { + auto gate_op = + reverse_lookup(Constant::gate_names, std::string_view{gate_name}); + auto num_params = lookup(Constant::gate_num_params, gate_op); + auto params = std::vector(num_params, param); + auto gate_matrix = getMatrix( + str_to_gates_.at(gate_name), params, inverse); + + StateVectorT kokkos_sv_ops{ini_st.data(), ini_st.size()}; + StateVectorT kokkos_sv_mat{ini_st.data(), ini_st.size()}; + + std::vector controlled_wires = {4}; + std::vector controlled_values = {true}; + + const auto wires = + createWires(str_to_controlled_gates_.at(gate_name), num_qubits); + kokkos_sv_ops.applyOperation(gate_name, controlled_wires, + controlled_values, wires, inverse, params); + kokkos_sv_mat.applyOperation("Matrix", controlled_wires, + controlled_values, wires, false, {}, + gate_matrix); + + auto result_ops = kokkos_sv_ops.getDataVector(); + auto result_mat = kokkos_sv_mat.getDataVector(); + + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(real(result_ops[j]) == + Approx(real(result_mat[j])).margin(EP)); + CHECK(imag(result_ops[j]) == + Approx(imag(result_mat[j])).margin(EP)); + } + } +} + TEMPLATE_TEST_CASE( "StateVectorKokkos::applyOperation param one-qubit with controls", "[StateVectorKokkos_Operation]", float, double) { @@ -568,6 +632,39 @@ TEMPLATE_TEST_CASE( } } } + + SECTION("N-controlled MultiRZ") { + std::vector wires = {control, wire0, wire1, wire2, wire3}; + std::sort(wires.begin(), wires.end()); + const ComplexT e = Kokkos::exp(ComplexT{0, -0.5} * param); + std::vector matrix(16, 0.0); + matrix[0] = e; + matrix[5] = conj(e); + matrix[10] = conj(e); + matrix[15] = e; + if (std::adjacent_find(wires.begin(), wires.end()) == wires.end()) { + StateVectorT kokkos_sv_mat{ini_st.data(), ini_st.size()}; + StateVectorT kokkos_sv_op{ini_st.data(), ini_st.size()}; + + kokkos_sv_mat.applyControlledMatrix( + matrix, {control, wire0, wire1}, + std::vector{true, false, true}, {wire2, wire3}, inverse); + kokkos_sv_op.applyOperation( + "MultiRZ", std::vector{control, wire0, wire1}, + std::vector{true, false, true}, + std::vector{wire2, wire3}, inverse, {param}); + + auto result_mat = kokkos_sv_mat.getDataVector(); + auto result_op = kokkos_sv_op.getDataVector(); + + for (std::size_t j = 0; j < exp2(num_qubits); j++) { + CHECK(real(result_op[j]) == + Approx(real(result_mat[j])).margin(EP)); + CHECK(imag(result_op[j]) == + Approx(imag(result_mat[j])).margin(EP)); + } + } + } } TEMPLATE_TEST_CASE("StateVectorKokkosManaged::applyIsingXY", diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py index 3888fc972..5d59dcf77 100644 --- a/pennylane_lightning/lightning_kokkos/_state_vector.py +++ b/pennylane_lightning/lightning_kokkos/_state_vector.py @@ -201,9 +201,14 @@ def _apply_lightning_controlled(self, operation): inv = False param = operation.parameters method(control_wires, control_values, target_wires, inv, param) - else: - raise qml.DeviceError( - "No gate operation supplied and controlled matrix not yet supported" + else: # apply gate as an n-controlled matrix + method = getattr(state, "applyControlledMatrix") + method( + qml.matrix(operation.base), + control_wires, + control_values, + target_wires, + False, ) def _apply_lightning_midmeasure( @@ -284,9 +289,7 @@ def _apply_lightning( elif method is not None: # apply specialized gate param = operation.parameters method(wires, invert_param, param) - elif isinstance(operation, qml.ops.Controlled) and not isinstance( - operation.base, (qml.QubitUnitary, qml.BlockEncode) - ): # apply n-controlled gate + elif isinstance(operation, qml.ops.Controlled): # apply n-controlled gate self._apply_lightning_controlled(operation) else: # apply gate as a matrix # Inverse can be set to False since qml.matrix(operation) is already in diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py index c1c156539..faa7e6d0b 100644 --- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py +++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py @@ -129,6 +129,7 @@ "C(DoubleExcitationMinus)", "C(DoubleExcitationPlus)", "C(MultiRZ)", + "C(QubitUnitary)", "CRot", "IsingXX", "IsingYY", diff --git a/tests/lightning_qubit/test_measurements_class.py b/tests/lightning_qubit/test_measurements_class.py index 9a606e66a..d7ec89802 100644 --- a/tests/lightning_qubit/test_measurements_class.py +++ b/tests/lightning_qubit/test_measurements_class.py @@ -799,8 +799,8 @@ def test_controlled_qubit_gates(self, operation, n_qubits, control_value, tol, l assert np.allclose(result, expected, tol * 10) @pytest.mark.skipif( - device_name not in ("lightning.qubit", "lightning.tensor"), - reason="N-controlled operations only implemented in lightning.qubit.", + device_name not in ("lightning.qubit", "lightning.tensor", "lightning.kokkos"), + reason="N-controlled operations only implemented in lightning.qubit, lightning.tensor, and lightning.kokkos.", ) def test_controlled_qubit_unitary_from_op(self, tol, lightning_sv): n_qubits = 10 diff --git a/tests/test_gates.py b/tests/test_gates.py index 2f288c020..d1c1840ad 100644 --- a/tests/test_gates.py +++ b/tests/test_gates.py @@ -372,8 +372,8 @@ def test_state_prep(n_targets, tol): @pytest.mark.skipif( - device_name != "lightning.qubit", - reason="N-controlled operations only implemented in lightning.qubit.", + device_name not in ("lightning.qubit", "lightning.kokkos"), + reason="N-controlled operations only implemented in lightning.qubit and lightning.kokkos.", ) @pytest.mark.parametrize("control_value", [False, True]) @pytest.mark.parametrize("n_qubits", list(range(2, 8))) @@ -497,8 +497,8 @@ def circuit(): @pytest.mark.skipif( - device_name != "lightning.qubit", - reason="N-controlled operations only implemented in lightning.qubit.", + device_name not in ("lightning.qubit", "lightning.kokkos"), + reason="N-controlled operations only implemented in lightning.qubit and lightning.kokkos.", ) def test_controlled_qubit_unitary_from_op(tol): n_qubits = 10 @@ -558,8 +558,8 @@ def test_paulirot(n_wires, n_targets, tol): @pytest.mark.skipif( - device_name not in ("lightning.qubit", "lightning.tensor"), - reason="N-controlled operations only implemented in lightning.qubit.", + device_name not in ("lightning.qubit", "lightning.tensor", "lightning.kokkos"), + reason="N-controlled operations only implemented in lightning.qubit, lightning.tensor, and lightning.kokkos.", ) @pytest.mark.parametrize("control_wires", range(4)) @pytest.mark.parametrize("target_wires", range(4))