Optimize lightning.tensor by adding direct MPS sites data set

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"

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp

@@ -499,6 +499,16 @@ class TNCuda : public TNCudaBase<PrecisionT, Derived> {
                       projected_mode_values, numHyperSamples);
     }
 
+    /**
+     * @brief Get a const vector reference of sitesExtents_.
+     *
+     * @return const std::vector<std::vector<std::size_t>>
+     */
+    [[nodiscard]] auto getSitesExtents() const
+        -> const std::vector<std::vector<std::size_t>> & {
+        return sitesExtents_;
+    }
+
   protected:
     /**
      * @brief Get a vector of pointers to tensor data of each site.

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<np_arr_c> &tensors) {
+                // Extract the incoming MPS shape
+                std::vector<std::vector<std::size_t>> 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<std::size_t> 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<std::complex<PrecisionT> *>(

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp

@@ -323,4 +323,64 @@ TEMPLATE_TEST_CASE("MPOTNCuda::getBondDims()", "[MPOTNCuda]", float, double) {
 
         CHECK(bondDims == expected_bondDims);
     }
-}
+}
+
+TEMPLATE_TEST_CASE("MPSTNCuda::getSitesExtents()", "[MPSTNCuda]", float,
+                   double) {
+    SECTION("Check if sitesExtents retrun is correctly with 3 qubits") {
+        const std::size_t num_qubits = 3;
+        const std::size_t maxBondDim = 128;
+        const DevTag<int> dev_tag{0, 0};
+
+        const std::vector<std::vector<std::size_t>> reference{
+            {{2, 2}, {2, 2, 2}, {2, 2}}};
+
+        MPSTNCuda<TestType> mps{num_qubits, maxBondDim, dev_tag};
+
+        const auto &sitesExtents = mps.getSitesExtents();
+
+        CHECK(reference == sitesExtents);
+    }
+
+    SECTION("Check if sitesExtents retrun is correctly with 8 qubits") {
+        const std::size_t num_qubits = 8;
+        const std::size_t maxBondDim = 128;
+        const DevTag<int> dev_tag{0, 0};
+
+        const std::vector<std::vector<std::size_t>> reference{{{2, 2},
+                                                               {2, 2, 4},
+                                                               {4, 2, 8},
+                                                               {8, 2, 16},
+                                                               {16, 2, 8},
+                                                               {8, 2, 4},
+                                                               {4, 2, 2},
+                                                               {2, 2}}};
+
+        MPSTNCuda<TestType> mps{num_qubits, maxBondDim, dev_tag};
+
+        const auto &sitesExtents = mps.getSitesExtents();
+
+        CHECK(reference == sitesExtents);
+    }
+    SECTION("Check if sitesExtents retrun is correctly with 8 qubits and "
+            "maxBondDim=8") {
+        const std::size_t num_qubits = 8;
+        const std::size_t maxBondDim = 8;
+        const DevTag<int> dev_tag{0, 0};
+
+        const std::vector<std::vector<std::size_t>> reference{{{2, 2},
+                                                               {2, 2, 4},
+                                                               {4, 2, 8},
+                                                               {8, 2, 8},
+                                                               {8, 2, 8},
+                                                               {8, 2, 4},
+                                                               {4, 2, 2},
+                                                               {2, 2}}};
+
+        MPSTNCuda<TestType> mps{num_qubits, maxBondDim, dev_tag};
+
+        const auto &sitesExtents = mps.getSitesExtents();
+
+        CHECK(reference == sitesExtents);
+    }
+}

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<std::vector<std::size_t>> MPS_shape_dest{
+            {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}};
+
+        std::vector<std::vector<std::size_t>> 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<std::vector<std::size_t>> MPS_shape_dest{
+            {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}};
+
+        std::vector<std::vector<std::size_t>> 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<std::vector<std::size_t>> MPS_shape_dest{
+            {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}};
+
+        std::vector<std::vector<std::size_t>> 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<std::vector<std::size_t>> MPS_shape_dest{
+            {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}};
+
+        std::vector<std::vector<std::size_t>> 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"));
+    }
+}

pennylane_lightning/core/src/simulators/lightning_tensor/utils/tncuda_utils/tncuda_helpers.hpp

@@ -194,4 +194,28 @@ inline auto create_swap_wire_pair_queue(const std::vector<std::size_t> &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<std::vector<std::size_t>> &MPS_shape_dest,
+              const std::vector<std::vector<std::size_t>> &MPS_shape_source) {
+    if (MPS_shape_dest != MPS_shape_source) {
+        auto MPS_shape_source_str =
+            Pennylane::Util::vector2DToString<std::size_t>(MPS_shape_source);
+        auto MPS_shape_dest_str =
+            Pennylane::Util::vector2DToString<std::size_t>(MPS_shape_dest);
+
+        PL_ABORT_IF_NOT(MPS_shape_dest == MPS_shape_source,
+                        "The incoming MPS does not have the correct layout for "
+                        "lightning.tensor.\n    Incoming MPS: " +
+                            MPS_shape_source_str +
+                            "\n Destination MPS: " + MPS_shape_dest_str)
+    }
+}
+
 } // namespace Pennylane::LightningTensor::TNCuda::Util

pennylane_lightning/core/src/utils/Util.hpp

@@ -592,4 +592,34 @@ bool areVecsDisjoint(const std::vector<T> &v1, const std::vector<T> &v2) {
     }
     return true;
 }
+
+/**
+ * @brief Convert a 2D vector to string.
+ * @tparam T Data type.
+ * @param vec Vector to convert.
+ *
+ * @return std::string String with the vector values.
+ */
+template <typename T>
+std::string vector2DToString(const std::vector<std::vector<T>> &vec) {
+    std::ostringstream oss;
+    oss << "[";
+
+    for (std::size_t i = 0; i < vec.size(); ++i) {
+        oss << "[";
+        for (std::size_t j = 0; j < vec[i].size(); ++j) {
+            oss << vec[i][j];
+            if (j != vec[i].size() - 1) {
+                oss << ", "; // Add a comma between elements in the inner vector
+            }
+        }
+        oss << "]";
+        if (i != vec.size() - 1) {
+            oss << ", "; // Add a comma between inner vectors
+        }
+    }
+    oss << "]";
+    return oss.str(); // Return the resulting string
+}
+
 } // namespace Pennylane::Util

pennylane_lightning/core/src/utils/tests/Test_Util.cpp

@@ -233,3 +233,15 @@ TEST_CASE("Util::areVecsDisjoint", "[Util][LinearAlgebra]") {
         REQUIRE(areVecsDisjoint(vec0, vec1) == false);
     }
 }
+
+TEST_CASE("Utils::vector2DToString", "[Utils]") {
+    SECTION("Test for convert 2Dvector to string") {
+        std::vector<std::vector<std::size_t>> vec{
+            {2, 2, 4}, {4, 2, 8}, {8, 2, 8}};
+        std::string ref_str{"[[2, 2, 4], [4, 2, 8], [8, 2, 8]]"};
+
+        std::string vec2str = vector2DToString<std::size_t>(vec);
+
+        REQUIRE(ref_str == vec2str);
+    }
+}

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
@@ -444,6 +444,14 @@
             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 == "tn":
+                    raise DeviceError("Exact Tensor Network does not support MPSPrep")
+
+                if self.method == "mps":
+                    mps = operations[0].mps
+                    self._tensornet.updateMPSSitesData(mps)
+                    operations = operations[1:]
 
         self._apply_lightning(operations)
 

pennylane_lightning/lightning_tensor/lightning_tensor.py

@@ -71,6 +71,7 @@
     {
         "Identity",
         "BasisState",
+        "MPSPrep",
         "QubitUnitary",
         "ControlledQubitUnitary",
         "DiagonalQubitUnitary",
@@ -169,6 +170,9 @@
     if isinstance(op, qml.ControlledQubitUnitary):
         return True
 
+    if isinstance(op, qml.MPSPrep):
+        return True
+
     return op.has_matrix and op.name in _operations