Add stateprep support to lightning.tensor - MPS (#849)

### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to
the
      [`tests`](../tests) directory!

- [ ] All new functions and code must be clearly commented and
documented.
If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [ ] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `.github/CHANGELOG.md` file, summarizing
the
      change, and including a link back to the PR.

- [x] Ensure that code is properly formatted by running `make format`. 

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.


------------------------------------------------------------------------------------------------------------

**Context:**

[SC-71278]

Add `qml.StatePrep()` and `qml.QubitStateVector()` support to
`lightning.tensor`. Note that `svd` decomposition is conduct by `numpy`
in the python layer based on the fact that `svd` operation with `numpy`
on cpu is faster than `cutensornet` on `A100` GPU if the `bond
dimension` is less than `256`. For more info, please visit
[here](https://developer.nvidia.com/cuquantum-sdk).

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

---------

Co-authored-by: ringo-but-quantum <[email protected]>
Co-authored-by: Vincent Michaud-Rioux <[email protected]>
Co-authored-by: Ali Asadi <[email protected]>
Co-authored-by: Lee James O'Riordan <[email protected]>
Co-authored-by: Josh Izaac <[email protected]>
Co-authored-by: Lee J. O'Riordan <[email protected]>
Co-authored-by: Pietropaolo Frisoni <[email protected]>
Co-authored-by: erick-xanadu <[email protected]>
Co-authored-by: Astral Cai <[email protected]>
Co-authored-by: Amintor Dusko <[email protected]>
Co-authored-by: Shiro-Raven <[email protected]>
Co-authored-by: albi3ro <[email protected]>

Author: 13 people

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add `qml.StatePrep()` and `qml.QubitStateVector()` support to `lightning.tensor`.
+  [(#849)](https://github.com/PennyLaneAI/pennylane-lightning/pull/849)
+
 * Add LightningGPU Linux (aarch64+GraceHopper) wheels to PyPI.
   [(#815)](https://github.com/PennyLaneAI/pennylane-lightning/pull/815)
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.38.0-dev43"
+__version__ = "0.38.0-dev44"

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

@@ -64,7 +64,7 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     MPSStatus MPSInitialized_ = MPSStatus::MPSInitNotSet;
 
     const std::size_t maxBondDim_;
-
+    const std::vector<std::size_t> bondDims_;
     const std::vector<std::vector<std::size_t>> sitesModes_;
     const std::vector<std::vector<std::size_t>> sitesExtents_;
     const std::vector<std::vector<int64_t>> sitesExtents_int64_;
@@ -86,19 +86,25 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     explicit MPSTNCuda(const std::size_t numQubits,
                        const std::size_t maxBondDim)
         : BaseType(numQubits), maxBondDim_(maxBondDim),
-          sitesModes_(setSitesModes_()), sitesExtents_(setSitesExtents_()),
+          bondDims_(setBondDims_()), sitesModes_(setSitesModes_()),
+          sitesExtents_(setSitesExtents_()),
           sitesExtents_int64_(setSitesExtents_int64_()) {
         initTensors_();
+        reset();
+        appendInitialMPSState_();
     }
 
     // TODO: Add method to the constructor to allow users to select methods at
     // runtime in the C++ layer
     explicit MPSTNCuda(const std::size_t numQubits,
                        const std::size_t maxBondDim, DevTag<int> dev_tag)
         : BaseType(numQubits, dev_tag), maxBondDim_(maxBondDim),
-          sitesModes_(setSitesModes_()), sitesExtents_(setSitesExtents_()),
+          bondDims_(setBondDims_()), sitesModes_(setSitesModes_()),
+          sitesExtents_(setSitesExtents_()),
           sitesExtents_int64_(setSitesExtents_int64_()) {
         initTensors_();
+        reset();
+        appendInitialMPSState_();
     }
 
     ~MPSTNCuda() = default;
@@ -162,6 +168,31 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
         setBasisState(zeroState);
     }
 
+    /**
+     * @brief Update the ith MPS site data.
+     *
+     * @param site_idx Index of the MPS site.
+     * @param host_data Pointer to the data on host.
+     * @param host_data_size Length of the data.
+     */
+    void updateMPSSiteData(const std::size_t site_idx,
+                           const ComplexT *host_data,
+                           std::size_t host_data_size) {
+        PL_ABORT_IF_NOT(
+            site_idx < BaseType::getNumQubits(),
+            "The site index should be less than the number of qubits.");
+
+        const std::size_t idx = BaseType::getNumQubits() - site_idx - 1;
+        PL_ABORT_IF_NOT(
+            host_data_size == tensors_[idx].getDataBuffer().getLength(),
+            "The length of the host data should match its copy on the device.");
+
+        tensors_[idx].getDataBuffer().zeroInit();
+
+        tensors_[idx].getDataBuffer().CopyHostDataToGpu(host_data,
+                                                        host_data_size);
+    }
+
     /**
      * @brief Update quantum state with a basis state.
      * NOTE: This API assumes the bond vector is a standard basis vector
@@ -192,18 +223,13 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
             if (i == 0) {
                 target = basisState[idx];
             } else {
-                target = basisState[idx] == 0 ? 0 : maxBondDim_;
+                target = basisState[idx] == 0 ? 0 : bondDims_[i - 1];
             }
 
             PL_CUDA_IS_SUCCESS(
                 cudaMemcpy(&tensors_[i].getDataBuffer().getData()[target],
                            &value_cu, sizeof(CFP_t), cudaMemcpyHostToDevice));
         }
-
-        if (MPSInitialized_ == MPSStatus::MPSInitNotSet) {
-            MPSInitialized_ = MPSStatus::MPSInitSet;
-            updateQuantumStateMPS_();
-        }
     };
 
     /**
@@ -307,6 +333,27 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     }
 
   private:
+    /**
+     * @brief Return bondDims to the member initializer
+     * NOTE: This method only works for the open boundary condition
+     * @return std::vector<std::size_t>
+     */
+    std::vector<std::size_t> setBondDims_() {
+        std::vector<std::size_t> localBondDims(BaseType::getNumQubits() - 1,
+                                               maxBondDim_);
+
+        const std::size_t ubDim = log2(maxBondDim_);
+        for (std::size_t i = 0; i < localBondDims.size(); i++) {
+            const std::size_t bondDim =
+                std::min(i + 1, BaseType::getNumQubits() - i - 1);
+
+            if (bondDim <= ubDim) {
+                localBondDims[i] = std::size_t{1} << bondDim;
+            }
+        }
+        return localBondDims;
+    }
+
     /**
      * @brief Return siteModes to the member initializer
      * NOTE: This method only works for the open boundary condition
@@ -348,15 +395,16 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
             if (i == 0) {
                 // Leftmost site (state mode, shared mode)
                 localSiteExtents = std::vector<std::size_t>(
-                    {BaseType::getQubitDims()[i], maxBondDim_});
+                    {BaseType::getQubitDims()[i], bondDims_[i]});
             } else if (i == BaseType::getNumQubits() - 1) {
                 // Rightmost site (shared mode, state mode)
                 localSiteExtents = std::vector<std::size_t>(
-                    {maxBondDim_, BaseType::getQubitDims()[i]});
+                    {bondDims_[i - 1], BaseType::getQubitDims()[i]});
             } else {
                 // Interior sites (state mode, state mode, shared mode)
                 localSiteExtents = std::vector<std::size_t>(
-                    {maxBondDim_, BaseType::getQubitDims()[i], maxBondDim_});
+                    {bondDims_[i - 1], BaseType::getQubitDims()[i],
+                     bondDims_[i]});
             }
             localSitesExtents.push_back(std::move(localSiteExtents));
         }
@@ -394,11 +442,11 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     }
 
     /**
-     * @brief Update quantumState (cutensornetState_t) with data provided by a
-     * user
+     * @brief Append initial MPS sites to the compute graph with data provided
+     * by a user
      *
      */
-    void updateQuantumStateMPS_() {
+    void appendInitialMPSState_() {
         PL_CUTENSORNET_IS_SUCCESS(cutensornetStateInitializeMPS(
             /*const cutensornetHandle_t */ BaseType::getTNCudaHandle(),
             /*cutensornetState_t*/ BaseType::getQuantumState(),

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp

@@ -74,6 +74,18 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
                 tensor_network.getData(data_ptr, state.size());
             },
             "Copy StateVector data into a Numpy array.")
+        .def(
+            "updateMPSSitesData",
+            [](TensorNet &tensor_network, std::vector<np_arr_c> &tensors) {
+                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> *>(
+                        numpyArrayInfo.ptr);
+                    tensor_network.updateMPSSiteData(idx, data_ptr,
+                                                     tensors[idx].size());
+                }
+            },
+            "Pass MPS site data to the C++ backend.")
         .def(
             "setBasisState",
             [](TensorNet &tensor_network,

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/measurements/tests/Test_MPSTNCuda_Expval.cpp

@@ -203,7 +203,7 @@ TEMPLATE_TEST_CASE("[PauliZ]", "[MPSTNCuda_Expval]", float, double) {
             PrecisionT ref = -0.2115276040475712;
 
             REQUIRE_THAT(res, Catch::Matchers::WithinRel(
-                                  ref, static_cast<PrecisionT>(cutoff)));
+                                  ref, static_cast<PrecisionT>(0.1)));
         }
     }
 }

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

@@ -58,6 +58,67 @@ TEMPLATE_PRODUCT_TEST_CASE("MPSTNCuda::Constructibility",
     }
 }
 
+TEMPLATE_TEST_CASE("MPSTNCuda::setIthMPSSite", "[MPSTNCuda]", float, double) {
+    SECTION("Set MPS site with wrong site index") {
+        const std::size_t num_qubits = 3;
+        const std::size_t maxBondDim = 3;
+        const std::size_t siteIdx = 3;
+
+        MPSTNCuda<TestType> mps_state{num_qubits, maxBondDim};
+
+        std::vector<std::complex<TestType>> site_data(1, {0.0, 0.0});
+
+        REQUIRE_THROWS_WITH(
+            mps_state.updateMPSSiteData(siteIdx, site_data.data(),
+                                        site_data.size()),
+            Catch::Matchers::Contains(
+                "The site index should be less than the number of qubits."));
+    }
+
+    SECTION("Set MPS site with wrong site data size") {
+        const std::size_t num_qubits = 3;
+        const std::size_t maxBondDim = 3;
+        const std::size_t siteIdx = 0;
+
+        MPSTNCuda<TestType> mps_state{num_qubits, maxBondDim};
+
+        std::vector<std::complex<TestType>> site_data(1, {0.0, 0.0});
+
+        REQUIRE_THROWS_WITH(
+            mps_state.updateMPSSiteData(siteIdx, site_data.data(),
+                                        site_data.size()),
+            Catch::Matchers::Contains("The length of the host data should "
+                                      "match its copy on the device."));
+    }
+
+    SECTION("Set MPS sites") {
+        const std::size_t num_qubits = 2;
+        const std::size_t maxBondDim = 3;
+
+        MPSTNCuda<TestType> mps_state{num_qubits, maxBondDim};
+
+        mps_state.reset(); // Reset the state to zero state
+
+        std::vector<std::complex<TestType>> site0_data(4, {0.0, 0.0}); // MSB
+        std::vector<std::complex<TestType>> site1_data(4, {0.0, 0.0}); // LSB
+
+        site0_data[2] = {1.0, 0.0};
+        site1_data[1] = {1.0, 0.0};
+
+        mps_state.updateMPSSiteData(0, site0_data.data(), site0_data.size());
+        mps_state.updateMPSSiteData(1, site1_data.data(), site1_data.size());
+
+        auto results = mps_state.getDataVector();
+
+        std::vector<std::complex<TestType>> expected_state(
+            std::size_t{1} << num_qubits, std::complex<TestType>({0.0, 0.0}));
+
+        expected_state[3] = {1.0, 0.0};
+
+        CHECK(expected_state == Pennylane::Util::approx(results));
+    }
+}
+
 TEMPLATE_TEST_CASE("MPSTNCuda::SetBasisStates() & reset()", "[MPSTNCuda]",
                    float, double) {
     std::vector<std::vector<std::size_t>> basisStates = {

pennylane_lightning/core/src/utils/cuda_utils/tests/Test_LinearAlgebra.cpp

@@ -1,4 +1,4 @@
-// Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+// Copyright 2018-2024 Xanadu Quantum Technologies Inc.
 
 // Licensed under the Apache License, Version 2.0 (the License);
 // you may not use this file except in compliance with the License.
@@ -14,6 +14,8 @@
 #include <algorithm>
 #include <complex>
 #include <cstdio>
+#include <cuComplex.h>
+#include <type_traits>
 #include <vector>
 
 #include <catch2/catch.hpp>