From 87bcd1083ae365ad61baa173865e56e6b795896d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?= Date: Fri, 20 Dec 2024 13:49:25 -0500 Subject: [PATCH] Optimize `lightning.tensor` by adding direct MPS sites data set (#983) **Context:** Optimize `lightning.tensor` by adding direct MPS sites data set **Description of the Change:** Adding the `MPSPrep` gate to be able to pass an MPS directly to the Tensor Network. The `MPSPrep` gate frontend was developed on this [PR](https://github.com/PennyLaneAI/pennylane/pull/6431) **Benefits:** Avoid the decomposition from state vector to MPS sites which are expensive and inefficient **Possible Drawbacks:** **Related GitHub Issues:** [sc-74709] --------- Co-authored-by: ringo-but-quantum --- .github/CHANGELOG.md | 3 + pennylane_lightning/core/_version.py | 2 +- .../lightning_tensor/tncuda/TNCuda.hpp | 60 +++++++------ .../tncuda/bindings/LTensorTNCudaBindings.hpp | 17 +++- .../tncuda/tests/Tests_MPSTNCuda.cpp | 64 +++++++++++++- .../tncuda_utils/tests/Test_TNCuda_utils.cpp | 49 +++++++++++ .../utils/tncuda_utils/tncuda_helpers.hpp | 15 ++++ pennylane_lightning/core/src/utils/Util.hpp | 1 + .../lightning_tensor/_tensornet.py | 15 +++- .../lightning_tensor/lightning_tensor.py | 4 + .../lightning_tensor/test_lightning_tensor.py | 85 ++++++++++++++++++- .../lightning_tensor/test_tensornet_class.py | 2 +- 12 files changed, 281 insertions(+), 36 deletions(-) diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md index e265b3424e..fb008b06a2 100644 --- a/.github/CHANGELOG.md +++ b/.github/CHANGELOG.md @@ -29,6 +29,9 @@ ### Improvements +* Optimize lightning.tensor by adding direct MPS sites data set with `qml.MPSPrep`. + [(#983)](https://github.com/PennyLaneAI/pennylane-lightning/pull/983) + * Replace the `dummy_tensor_update` method with the `cutensornetStateCaptureMPS`API to ensure that further gates apply is allowed after the `cutensornetStateCompute` call. [(#1028)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1028/) diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py index f811d3d60e..30f24595a9 100644 --- a/pennylane_lightning/core/_version.py +++ b/pennylane_lightning/core/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.40.0-dev41" +__version__ = "0.40.0-dev42" diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp index fb249f0519..6352bcc481 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp @@ -60,32 +60,6 @@ class TNCuda : public TNCudaBase { using ComplexT = std::complex; using BaseType = TNCudaBase; - protected: - // Note both maxBondDim_ and bondDims_ are used for both MPS and Exact - // Tensor Network. Per Exact Tensor Network, maxBondDim_ is 1 and bondDims_ - // is {1}. Per Exact Tensor Network, setting bondDims_ allows call to - // appendInitialMPSState_() to append the initial state to the Exact Tensor - // Network state. - const std::size_t - maxBondDim_; // maxBondDim_ default is 1 for Exact Tensor Network - const std::vector - bondDims_; // bondDims_ default is {1} for Exact Tensor Network - - private: - const std::vector> sitesModes_; - const std::vector> sitesExtents_; - const std::vector> sitesExtents_int64_; - - SharedCublasCaller cublascaller_; - - std::shared_ptr> gate_cache_; - std::set gate_ids_; - - std::vector identiy_gate_ids_; - - std::vector> tensors_; - std::vector> tensors_out_; - public: TNCuda() = delete; @@ -499,7 +473,27 @@ class TNCuda : public TNCudaBase { projected_mode_values, numHyperSamples); } + /** + * @brief Get a const vector reference of sitesExtents_. + * + * @return const std::vector> + */ + [[nodiscard]] auto getSitesExtents() const + -> const std::vector> & { + return sitesExtents_; + } + protected: + // Note both maxBondDim_ and bondDims_ are used for both MPS and Exact + // Tensor Network. For Exact Tensor Network, maxBondDim_ is 1 and bondDims_ + // is {1}. For Exact Tensor Network, setting bondDims_ allows call to + // appendInitialMPSState_() to append the initial state to the Exact Tensor + // Network state. + const std::size_t + maxBondDim_; // maxBondDim_ default is 1 for Exact Tensor Network + const std::vector + bondDims_; // bondDims_ default is {1} for Exact Tensor Network + /** * @brief Get a vector of pointers to tensor data of each site. * @@ -578,6 +572,20 @@ class TNCuda : public TNCudaBase { } private: + const std::vector> sitesModes_; + const std::vector> sitesExtents_; + const std::vector> sitesExtents_int64_; + + SharedCublasCaller cublascaller_; + + std::shared_ptr> gate_cache_; + std::set gate_ids_; + + std::vector identiy_gate_ids_; + + std::vector> tensors_; + std::vector> tensors_out_; + /** * @brief Get accessor of a state tensor * diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp index 1522305a82..9892a95b6d 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp @@ -33,13 +33,14 @@ #include "TypeList.hpp" #include "Util.hpp" #include "cuda_helpers.hpp" +#include "tncuda_helpers.hpp" /// @cond DEV namespace { using namespace Pennylane; using namespace Pennylane::Bindings; using namespace Pennylane::LightningGPU::Util; -using Pennylane::LightningTensor::TNCuda::MPSTNCuda; +using namespace Pennylane::LightningTensor::TNCuda::Util; } // namespace /// @endcond @@ -137,6 +138,20 @@ void registerBackendClassSpecificBindingsMPS(PyClass &pyclass) { .def( "updateMPSSitesData", [](TensorNet &tensor_network, std::vector &tensors) { + // Extract the incoming MPS shape + std::vector> MPS_shape_source; + for (std::size_t idx = 0; idx < tensors.size(); idx++) { + py::buffer_info numpyArrayInfo = tensors[idx].request(); + auto MPS_site_source_shape = numpyArrayInfo.shape; + std::vector MPS_site_source( + MPS_site_source_shape.begin(), + MPS_site_source_shape.end()); + MPS_shape_source.emplace_back(std::move(MPS_site_source)); + } + + const auto &MPS_shape_dest = tensor_network.getSitesExtents(); + MPSShapeCheck(MPS_shape_dest, MPS_shape_source); + for (std::size_t idx = 0; idx < tensors.size(); idx++) { py::buffer_info numpyArrayInfo = tensors[idx].request(); auto *data_ptr = static_cast *>( diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp index b75c272697..f63ef70847 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp @@ -291,7 +291,7 @@ TEMPLATE_TEST_CASE("MPSTNCuda::getDataVector()", "[MPSTNCuda]", float, double) { TEMPLATE_TEST_CASE("MPOTNCuda::getBondDims()", "[MPOTNCuda]", float, double) { using cp_t = std::complex; - SECTION("Check if bondDims is correctly set") { + SECTION("Check if bondDims is correct set") { const std::size_t num_qubits = 3; const std::size_t maxBondDim = 128; const DevTag dev_tag{0, 0}; @@ -323,4 +323,64 @@ TEMPLATE_TEST_CASE("MPOTNCuda::getBondDims()", "[MPOTNCuda]", float, double) { CHECK(bondDims == expected_bondDims); } -} \ No newline at end of file +} + +TEMPLATE_TEST_CASE("MPSTNCuda::getSitesExtents()", "[MPSTNCuda]", float, + double) { + SECTION("Check if sitesExtents retrun is correct with 3 qubits") { + const std::size_t num_qubits = 3; + const std::size_t maxBondDim = 128; + const DevTag dev_tag{0, 0}; + + const std::vector> reference{ + {{2, 2}, {2, 2, 2}, {2, 2}}}; + + MPSTNCuda mps{num_qubits, maxBondDim, dev_tag}; + + const auto &sitesExtents = mps.getSitesExtents(); + + CHECK(reference == sitesExtents); + } + + SECTION("Check if sitesExtents retrun is correct with 8 qubits") { + const std::size_t num_qubits = 8; + const std::size_t maxBondDim = 128; + const DevTag dev_tag{0, 0}; + + const std::vector> reference{{{2, 2}, + {2, 2, 4}, + {4, 2, 8}, + {8, 2, 16}, + {16, 2, 8}, + {8, 2, 4}, + {4, 2, 2}, + {2, 2}}}; + + MPSTNCuda mps{num_qubits, maxBondDim, dev_tag}; + + const auto &sitesExtents = mps.getSitesExtents(); + + CHECK(reference == sitesExtents); + } + SECTION("Check if sitesExtents retrun is correct with 8 qubits and " + "maxBondDim=8") { + const std::size_t num_qubits = 8; + const std::size_t maxBondDim = 8; + const DevTag dev_tag{0, 0}; + + const std::vector> reference{{{2, 2}, + {2, 2, 4}, + {4, 2, 8}, + {8, 2, 8}, + {8, 2, 8}, + {8, 2, 4}, + {4, 2, 2}, + {2, 2}}}; + + MPSTNCuda mps{num_qubits, maxBondDim, dev_tag}; + + const auto &sitesExtents = mps.getSitesExtents(); + + CHECK(reference == sitesExtents); + } +} diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tests/Test_TNCuda_utils.cpp b/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tests/Test_TNCuda_utils.cpp index 63c399b1ba..7d4d5f2767 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tests/Test_TNCuda_utils.cpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tests/Test_TNCuda_utils.cpp @@ -84,3 +84,52 @@ TEST_CASE("swap_op_wires_queue", "[TNCuda_utils]") { REQUIRE(swap_wires_queue[1] == swap_wires_queue_ref1); } } + +TEST_CASE("MPSShapeCheck", "[TNCuda_utils]") { + SECTION("Correct incoming MPS shape") { + std::vector> MPS_shape_dest{ + {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}}; + + std::vector> MPS_shape_source{ + {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}}; + + REQUIRE_NOTHROW(MPSShapeCheck(MPS_shape_dest, MPS_shape_source)); + } + + SECTION("Incorrect incoming MPS shape, bond dimension") { + std::vector> MPS_shape_dest{ + {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}}; + + std::vector> incorrect_MPS_shape{ + {2, 2}, {2, 2, 2}, {2, 2, 2}, {2, 2}}; + + REQUIRE_THROWS_WITH( + MPSShapeCheck(MPS_shape_dest, incorrect_MPS_shape), + Catch::Matchers::Contains("The incoming MPS does not have the " + "correct layout for lightning.tensor")); + } + SECTION("Incorrect incoming MPS shape, physical dimension") { + std::vector> MPS_shape_dest{ + {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}}; + + std::vector> incorrect_shape{ + {4, 2}, {2, 4, 4}, {4, 4, 2}, {2, 4}}; + + REQUIRE_THROWS_WITH( + MPSShapeCheck(MPS_shape_dest, incorrect_shape), + Catch::Matchers::Contains("The incoming MPS does not have the " + "correct layout for lightning.tensor")); + } + SECTION("Incorrect incoming MPS shape, number sites") { + std::vector> MPS_shape_dest{ + {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}}; + + std::vector> incorrect_shape{ + {2, 2}, {2, 2, 2}, {2, 2}}; + + REQUIRE_THROWS_WITH( + MPSShapeCheck(MPS_shape_dest, incorrect_shape), + Catch::Matchers::Contains("The incoming MPS does not have the " + "correct layout for lightning.tensor")); + } +} diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tncuda_helpers.hpp b/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tncuda_helpers.hpp index 4fadfadb40..67e5a3d8ce 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tncuda_helpers.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tncuda_helpers.hpp @@ -194,4 +194,19 @@ inline auto create_swap_wire_pair_queue(const std::vector &wires) return {local_wires, swap_wires_queue}; } +/** + * @brief Check if the provided MPS has the correct dimension for C++ + * backend. + * + * @param MPS_shape_dest Dimension list of destination MPS. + * @param MPS_shape_source Dimension list of incoming MPS. + */ +inline void +MPSShapeCheck(const std::vector> &MPS_shape_dest, + const std::vector> &MPS_shape_source) { + PL_ABORT_IF_NOT(MPS_shape_dest == MPS_shape_source, + "The incoming MPS does not have the correct layout for " + "lightning.tensor.") +} + } // namespace Pennylane::LightningTensor::TNCuda::Util diff --git a/pennylane_lightning/core/src/utils/Util.hpp b/pennylane_lightning/core/src/utils/Util.hpp index 6f512ac40d..70ec857aaa 100644 --- a/pennylane_lightning/core/src/utils/Util.hpp +++ b/pennylane_lightning/core/src/utils/Util.hpp @@ -592,4 +592,5 @@ bool areVecsDisjoint(const std::vector &v1, const std::vector &v2) { } return true; } + } // namespace Pennylane::Util diff --git a/pennylane_lightning/lightning_tensor/_tensornet.py b/pennylane_lightning/lightning_tensor/_tensornet.py index 571c0ac6f0..c74e664adb 100644 --- a/pennylane_lightning/lightning_tensor/_tensornet.py +++ b/pennylane_lightning/lightning_tensor/_tensornet.py @@ -28,7 +28,7 @@ import numpy as np import pennylane as qml -from pennylane import BasisState, DeviceError, StatePrep +from pennylane import BasisState, DeviceError, MPSPrep, StatePrep from pennylane.ops.op_math import Adjoint from pennylane.tape import QuantumScript from pennylane.wires import Wires @@ -433,17 +433,24 @@ def apply_operations(self, operations): # State preparation is currently done in Python if operations: # make sure operations[0] exists if isinstance(operations[0], StatePrep): - if self.method == "tn": - raise DeviceError("Exact Tensor Network does not support StatePrep") - if self.method == "mps": self._apply_state_vector( operations[0].parameters[0].copy(), operations[0].wires ) operations = operations[1:] + if self.method == "tn": + raise DeviceError("Exact Tensor Network does not support StatePrep") elif isinstance(operations[0], BasisState): self._apply_basis_state(operations[0].parameters[0], operations[0].wires) operations = operations[1:] + elif isinstance(operations[0], MPSPrep): + if self.method == "mps": + mps = operations[0].mps + self._tensornet.updateMPSSitesData(mps) + operations = operations[1:] + + if self.method == "tn": + raise DeviceError("Exact Tensor Network does not support MPSPrep") self._apply_lightning(operations) diff --git a/pennylane_lightning/lightning_tensor/lightning_tensor.py b/pennylane_lightning/lightning_tensor/lightning_tensor.py index 523f6bd40b..386ae661f3 100644 --- a/pennylane_lightning/lightning_tensor/lightning_tensor.py +++ b/pennylane_lightning/lightning_tensor/lightning_tensor.py @@ -71,6 +71,7 @@ { "Identity", "BasisState", + "MPSPrep", "QubitUnitary", "ControlledQubitUnitary", "DiagonalQubitUnitary", @@ -169,6 +170,9 @@ def stopping_condition(op: Operator) -> bool: if isinstance(op, qml.ControlledQubitUnitary): return True + if isinstance(op, qml.MPSPrep): + return True + return op.has_matrix and op.name in _operations diff --git a/tests/lightning_tensor/test_lightning_tensor.py b/tests/lightning_tensor/test_lightning_tensor.py index ce1cb5b4ae..984ec0031b 100644 --- a/tests/lightning_tensor/test_lightning_tensor.py +++ b/tests/lightning_tensor/test_lightning_tensor.py @@ -18,7 +18,7 @@ import numpy as np import pennylane as qml import pytest -from conftest import LightningDevice, device_name # tested device +from conftest import LightningDevice, LightningException, device_name from pennylane.tape import QuantumScript from pennylane.wires import Wires @@ -27,6 +27,7 @@ else: from pennylane_lightning.lightning_tensor import LightningTensor + if not LightningDevice._CPP_BINARY_AVAILABLE: # pylint: disable=protected-access pytest.skip("Device doesn't have C++ support yet.", allow_module_level=True) @@ -157,3 +158,85 @@ def test_execute_and_compute_vjp(self, method): match="The computation of vector-Jacobian product has yet to be implemented for the lightning.tensor device.", ): dev.execute_and_compute_vjp(circuits=None, cotangents=None) + + +@pytest.mark.parametrize( + "wires,max_bond,MPS_shape", + [ + (2, 128, [[2, 2], [2, 2]]), + ( + 8, + 128, + [[2, 2], [2, 2, 4], [4, 2, 8], [8, 2, 16], [16, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]], + ), + (8, 8, [[2, 2], [2, 2, 4], [4, 2, 8], [8, 2, 8], [8, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]]), + (15, 2, [[2, 2]] + [[2, 2, 2] for _ in range(13)] + [[2, 2]]), + ], +) +def test_MPSPrep_check_pass(wires, max_bond, MPS_shape): + """Test the correct behavior regarding MPS shape of MPSPrep.""" + MPS = [np.zeros(i) for i in MPS_shape] + dev = LightningTensor(wires=wires, method="mps", max_bond_dim=max_bond) + dev_wires = dev.wires.tolist() + + def circuit(MPS): + qml.MPSPrep(mps=MPS, wires=dev_wires) + return qml.state() + + qnode_ltensor = qml.QNode(circuit, dev) + + _ = qnode_ltensor(MPS) + + +@pytest.mark.parametrize( + "wires,max_bond,MPS_shape", + [ + ( + 8, + 8, + [[2, 2], [2, 2, 4], [4, 2, 8], [8, 2, 16], [16, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]], + ), # Incorrect max bond dim. + (15, 2, [[2, 2]] + [[2, 2, 2] for _ in range(14)] + [[2, 2]]), # Incorrect amount of sites + ], +) +def test_MPSPrep_check_fail(wires, max_bond, MPS_shape): + """Test the exceptions regarding MPS shape of MPSPrep.""" + + MPS = [np.zeros(i) for i in MPS_shape] + dev = LightningTensor(wires=wires, method="mps", max_bond_dim=max_bond) + dev_wires = dev.wires.tolist() + + def circuit(MPS): + qml.MPSPrep(mps=MPS, wires=dev_wires) + return qml.state() + + qnode_ltensor = qml.QNode(circuit, dev) + + with pytest.raises( + LightningException, + match="The incoming MPS does not have the correct layout for lightning.tensor", + ): + _ = qnode_ltensor(MPS) + + +@pytest.mark.parametrize( + "wires, MPS_shape", + [ + (2, [[2, 2], [2, 2]]), + ], +) +def test_MPSPrep_with_tn(wires, MPS_shape): + """Test the exception of MPSPrep with the method exact tensor network (tn).""" + + MPS = [np.zeros(i) for i in MPS_shape] + dev = LightningTensor(wires=wires, method="tn") + dev_wires = dev.wires.tolist() + + def circuit(MPS): + qml.MPSPrep(mps=MPS, wires=dev_wires) + return qml.state() + + qnode_ltensor = qml.QNode(circuit, dev) + + with pytest.raises(qml.DeviceError, match="Exact Tensor Network does not support MPSPrep"): + _ = qnode_ltensor(MPS) diff --git a/tests/lightning_tensor/test_tensornet_class.py b/tests/lightning_tensor/test_tensornet_class.py index 62ef974749..64f367ed7a 100644 --- a/tests/lightning_tensor/test_tensornet_class.py +++ b/tests/lightning_tensor/test_tensornet_class.py @@ -71,7 +71,7 @@ def test_wrong_device_name(): def test_wrong_method_name(): - """Test an invalid device name""" + """Test an invalid method name""" with pytest.raises(qml.DeviceError, match="The method "): LightningTensorNet(3, max_bond_dim=5, device_name="lightning.tensor", method="spider_web")