fix multiple reset()/setbasisstate() & resolve comments

.github/workflows/tests_lmps_tncuda_cpp.yml

@@ -70,7 +70,7 @@ jobs:
         pl_tensor_backend: ["cutensornet"]
         cuda_version: ["12"]
 
-    name: C++ tests (Lightning-Tensor-MPS-TNCuda)
+    name: C++ Tests (${{ matrix.pl_backend }}, method-${{ matrix.pl_tensor_method }}, backend-${{ matrix.pl_tensor_backend }}, cuda-${{ matrix.cuda_version }})  
     runs-on:
       - ${{ matrix.os }}
       - self-hosted
@@ -95,21 +95,27 @@ jobs:
         name: Install Python
         with:
           python-version: '3.9'
-
+
+      # Since the self-hosted runner can be re-used. It is best to set up all package
+      # installations in a virtual environment that gets cleaned at the end of each workflow run
       - name: Setup Python virtual environment
         id: setup_venv
         env:
           VENV_NAME: ${{ github.workspace }}/venv_${{ steps.setup_python.outputs.python-version }}_${{ github.sha }}
         run: |
           # Clear any pre-existing venvs
           rm -rf venv_*
+
           # Create new venv for this workflow_run
           python --version
           python -m venv ${{ env.VENV_NAME }}
+
           # Add the venv to PATH for subsequent steps
           echo ${{ env.VENV_NAME }}/bin >> $GITHUB_PATH
+
           # Adding venv name as an output for subsequent steps to reference if needed
           echo "venv_name=${{ env.VENV_NAME }}" >> $GITHUB_OUTPUT
+
       - name: Display Python-Path
         id: python_path
         run: |
@@ -185,8 +191,5 @@ jobs:
         if: always()
         run: |
           rm -rf ${{ steps.setup_venv.outputs.venv_name }}
-          rm -rf *
-          rm -rf .git
-          rm -rf .gitignore
-          rm -rf .github
+          rm -rf * .git .gitignore .github
           pip cache purge

pennylane_lightning/core/src/simulators/lightning_tensor/tensor/CMakeLists.txt

@@ -17,4 +17,4 @@ endif()
 target_include_directories(${PL_BACKEND}_tensor INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}/base)
 target_link_libraries(${PL_BACKEND}_tensor INTERFACE tensorBase lightning_utils lightning_compile_options lightning_external_libs)
 
-set_property(TARGET ${PL_BACKEND}_tensor PROPERTY POSITION_INDEPENDENT_CODE ON)
+set_property(TARGET ${PL_BACKEND}_tensor PROPERTY POSITION_INDEPENDENT_CODE ON)

pennylane_lightning/core/src/simulators/lightning_tensor/tensor/base/TensorBase.hpp

@@ -33,21 +33,20 @@ namespace Pennylane::LightningTensor {
  */
 template <class PrecisionT, class Derived> class TensorBase {
   private:
-    std::size_t rank_;                 // A rank N tensor has N modes
-    std::size_t length_;               // Number of elements
-    std::vector<std::size_t> modes_;   // modes for contraction identify
-    std::vector<std::size_t> extents_; // Number of elements in each mode
+    const std::size_t rank_;                 // A rank N tensor has N modes
+    std::size_t length_;                     // Number of elements
+    const std::vector<std::size_t> modes_;   // modes for contraction identify
+    const std::vector<std::size_t> extents_; // Number of elements in each mode
 
   public:
-    TensorBase(std::size_t rank, const std::vector<std::size_t> &modes,
-               const std::vector<std::size_t> &extents)
+    explicit TensorBase(std::size_t rank, const std::vector<std::size_t> &modes,
+                        const std::vector<std::size_t> &extents)
         : rank_(rank), modes_(modes), extents_(extents) {
         PL_ABORT_IF_NOT(rank_ == extents_.size(),
                         "Please check if rank or extents are set correctly.");
         length_ = 1;
-        for (auto extent : extents) {
-            length_ *= extent;
-        }
+        length_ = std::accumulate(extents.begin(), extents.end(),
+                                  std::size_t{1}, std::multiplies<>());
     }
 
     ~TensorBase() {}

pennylane_lightning/core/src/simulators/lightning_tensor/tensor/tncuda/TensorCuda.hpp

@@ -46,9 +46,10 @@ class TensorCuda final : public TensorBase<PrecisionT, TensorCuda<PrecisionT>> {
     using BaseType = TensorBase<PrecisionT, TensorCuda>;
     using CFP_t = decltype(cuUtil::getCudaType(PrecisionT{}));
 
-    TensorCuda(const std::size_t rank, const std::vector<std::size_t> &modes,
-               const std::vector<std::size_t> &extents,
-               const DevTag<int> &dev_tag, bool device_alloc = true)
+    explicit TensorCuda(const std::size_t rank,
+                        const std::vector<std::size_t> &modes,
+                        const std::vector<std::size_t> &extents,
+                        const DevTag<int> &dev_tag, bool device_alloc = true)
         : TensorBase<PrecisionT, TensorCuda<PrecisionT>>(rank, modes, extents),
           data_buffer_{std::make_shared<DataBuffer<CFP_t>>(
               BaseType::getLength(), dev_tag, device_alloc)} {}

pennylane_lightning/core/src/simulators/lightning_tensor/tensor/tncuda/tests/Tests_TensorCuda.cpp

@@ -67,4 +67,4 @@ TEMPLATE_TEST_CASE("TensorCuda::baseMethods", "[TensorCuda]", float, double) {
     SECTION("getExtents()") { CHECK(tensor.getExtents() == extents); }
 
     SECTION("getLength()") { CHECK(tensor.getLength() == length); }
-}
+}

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

@@ -61,7 +61,6 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     using BaseType = TNCudaBase<Precision, MPSTNCuda>;
 
     MPSStatus MPSInitialized_ = MPSStatus::MPSInitNotSet;
-    MPSStatus MPSFinalized_ = MPSStatus::MPSFinalizedNotSet;
 
     const std::size_t maxBondDim_;
 
@@ -161,11 +160,6 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
                         "Please ensure all elements of a basis state should be "
                         "either 0 or 1.");
 
-        PL_ABORT_IF(MPSInitialized_ == MPSStatus::MPSInitSet,
-                    "setBasisState() can be called only once.");
-
-        MPSInitialized_ = MPSStatus::MPSInitSet;
-
         CFP_t value_cu =
             Pennylane::LightningGPU::Util::complexToCu<ComplexT>({1.0, 0.0});
 
@@ -186,8 +180,10 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
                            &value_cu, sizeof(CFP_t), cudaMemcpyHostToDevice));
         }
 
-        updateQuantumStateMPS_(getSitesExtentsPtr().data(),
-                               getTensorsDataPtr().data());
+        if (MPSInitialized_ == MPSStatus::MPSInitNotSet) {
+            MPSInitialized_ = MPSStatus::MPSInitSet;
+            updateQuantumStateMPS_();
+        }
     };
 
     /**
@@ -201,13 +197,6 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
      * quantum state on host
      */
     auto getDataVector() -> std::vector<ComplexT> {
-        PL_ABORT_IF(MPSFinalized_ == MPSStatus::MPSFinalizedSet,
-                    "getDataVector() method to return the full state "
-                    "vector can't be called "
-                    "after cutensornetStateFinalizeMPS is called");
-
-        MPSFinalized_ = MPSStatus::MPSFinalizedSet;
-
         // 1D representation
         std::vector<std::size_t> output_modes(std::size_t{1}, std::size_t{1});
         std::vector<std::size_t> output_extent(
@@ -216,17 +205,10 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
                                             output_extent,
                                             BaseType::getDevTag());
 
-        std::vector<void *> output_tensorPtr(
-            std::size_t{1},
-            static_cast<void *>(output_tensor.getDataBuffer().getData()));
-
-        std::vector<int64_t *> output_extentsPtr;
-        std::vector<int64_t> extent_int64(
-            std::size_t{1},
-            static_cast<int64_t>(std::size_t{1} << BaseType::getNumQubits()));
-        output_extentsPtr.emplace_back(extent_int64.data());
+        void *output_tensorPtr[] = {
+            static_cast<void *>(output_tensor.getDataBuffer().getData())};
 
-        this->computeState(output_extentsPtr, output_tensorPtr);
+        this->computeState(output_tensorPtr);
 
         std::vector<ComplexT> results(output_extent.front());
         output_tensor.CopyGpuDataToHost(results.data(), results.size());
@@ -328,19 +310,17 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
      * @brief Update quantumState (cutensornetState_t) with data provided by a
      * user
      *
-     * @param extentsIn Extents of each sites
-     * @param tensorsIn Pointer to tensors provided by a user
      */
-    void updateQuantumStateMPS_(const int64_t *const *extentsIn,
-                                uint64_t **tensorsIn) {
+    void updateQuantumStateMPS_() {
         PL_CUTENSORNET_IS_SUCCESS(cutensornetStateInitializeMPS(
             /*const cutensornetHandle_t */ BaseType::getTNCudaHandle(),
             /*cutensornetState_t*/ BaseType::getQuantumState(),
             /*cutensornetBoundaryCondition_t */
             CUTENSORNET_BOUNDARY_CONDITION_OPEN,
-            /*const int64_t *const* */ extentsIn,
+            /*const int64_t *const* */ getSitesExtentsPtr().data(),
             /*const int64_t *const* */ nullptr,
-            /*void ** */ reinterpret_cast<void **>(tensorsIn)));
+            /*void ** */
+            reinterpret_cast<void **>(getTensorsDataPtr().data())));
     }
 };
 } // namespace Pennylane::LightningTensor::TNCuda

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

@@ -187,11 +187,9 @@ class TNCudaBase : public TensornetBase<Precision, Derived> {
     /**
      * @brief Save quantumState information to data provided by a user
      *
-     * @param extentsPtr Extents of each sites
      * @param tensorPtr Pointer to tensors provided by a user
      */
-    void computeState(std::vector<int64_t *> &extentsPtr,
-                      std::vector<void *> &tensorPtr) {
+    void computeState(void **tensorPtr) {
         cutensornetWorkspaceDescriptor_t workDesc;
         PL_CUTENSORNET_IS_SUCCESS(
             cutensornetCreateWorkspaceDescriptor(getTNCudaHandle(), &workDesc));
@@ -223,9 +221,9 @@ class TNCudaBase : public TensornetBase<Precision, Derived> {
             /* const cutensornetHandle_t */ getTNCudaHandle(),
             /* cutensornetState_t */ getQuantumState(),
             /* cutensornetWorkspaceDescriptor_t */ workDesc,
-            /* int64_t * */ extentsPtr.data(),
+            /* int64_t * */ nullptr,
             /* int64_t *stridesOut */ nullptr,
-            /* void * */ tensorPtr.data(),
+            /* void * */ tensorPtr,
             /* cudaStream_t */ getDevTag().getStreamID()));
 
         PL_CUTENSORNET_IS_SUCCESS(

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

@@ -135,10 +135,33 @@ TEMPLATE_TEST_CASE("MPSTNCuda::SetBasisStates() & reset()", "[MPSTNCuda]",
         CHECK(expected_state ==
               Pennylane::Util::approx(mps_state.getDataVector()));
     }
+
+    SECTION("Test different bondDim and different basisstate & reset()") {
+        std::size_t bondDim = GENERATE(2, 3, 4, 5);
+        std::size_t stateIdx = GENERATE(0, 1, 2, 3, 4, 5, 6, 7);
+        std::size_t num_qubits = 3;
+        std::size_t maxBondDim = bondDim;
+
+        MPSTNCuda<TestType> mps_state{num_qubits, maxBondDim};
+
+        mps_state.setBasisState(basisStates[stateIdx]);
+
+        mps_state.reset();
+
+        std::vector<std::complex<TestType>> expected_state(
+            size_t{1} << num_qubits, std::complex<TestType>({0.0, 0.0}));
+
+        std::size_t index = 0;
+
+        expected_state[index] = {1.0, 0.0};
+
+        CHECK(expected_state ==
+              Pennylane::Util::approx(mps_state.getDataVector()));
+    }
 }
 
 TEMPLATE_TEST_CASE("MPSTNCuda::getDataVector()", "[MPSTNCuda]", float, double) {
-    std::size_t num_qubits = 3;
+    std::size_t num_qubits = 10;
     std::size_t maxBondDim = 2;
     DevTag<int> dev_tag{0, 0};
 

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

@@ -29,8 +29,6 @@ enum class MPSStatus : uint32_t {
     BEGIN = 0,
     MPSInitNotSet = 0,
     MPSInitSet,
-    MPSFinalizedNotSet,
-    MPSFinalizedSet,
     END
 };