Add tag dispatch for memory location with LQ classes (#492)

* Add tag dispatch for memory location with LQ classes

* Auto update version

* Update changelog

* Trigger CI

* Move MemoryStorageLocation namespace to Memory.hpp

* Exclude PLABORT unhittable from coverage

---------

Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com>

Author: mlxd

.github/CHANGELOG.md

@@ -6,6 +6,9 @@
 
 ### Improvements
 
+* Add memory locality tag reporting and adjoint diff dispatch for `lightning.qubit` statevector classes.
+  [(#492)](https://github.com/PennyLaneAI/pennylane-lightning/pull/492)
+
 ### Documentation
 
 ### Bug fixes

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.33.0-dev1"
+__version__ = "0.33.0-dev2"

pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubit.hpp

@@ -57,6 +57,7 @@ template <class PrecisionT, class Derived>
 class StateVectorLQubit : public StateVectorBase<PrecisionT, Derived> {
   public:
     using ComplexT = std::complex<PrecisionT>;
+    using MemoryStorageT = Pennylane::Util::MemoryStorageLocation::Undefined;
 
   protected:
     const Threading threading_;

pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp

@@ -54,6 +54,7 @@ class StateVectorLQubitManaged final
   public:
     using PrecisionT = fp_t;
     using ComplexT = std::complex<PrecisionT>;
+    using MemoryStorageT = Pennylane::Util::MemoryStorageLocation::Internal;
 
   private:
     using BaseType =

pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitRaw.hpp

@@ -56,6 +56,7 @@ class StateVectorLQubitRaw final
   public:
     using PrecisionT = fp_t;
     using ComplexT = std::complex<PrecisionT>;
+    using MemoryStorageT = Pennylane::Util::MemoryStorageLocation::External;
 
   private:
     using BaseType =

pennylane_lightning/core/src/simulators/lightning_qubit/algorithms/AdjointJacobianLQubit.hpp

@@ -32,6 +32,8 @@
 /// @cond DEV
 namespace {
 using namespace Pennylane::Algorithms;
+using namespace Pennylane::Util::MemoryStorageLocation;
+
 using Pennylane::LightningQubit::Util::innerProdC;
 using Pennylane::LightningQubit::Util::Transpose;
 
@@ -256,12 +258,13 @@ class AdjointJacobian final
         // Pointer to data storage for StateVectorLQubitRaw<PrecisionT>:
         std::unique_ptr<std::vector<std::vector<ComplexT>>> H_lambda_storage;
         size_t lambda_qubits = lambda.getNumQubits();
-        if constexpr (std::is_same_v<StateVectorLQubitManaged<PrecisionT>,
-                                     StateVectorT>) {
+        if constexpr (std::is_same_v<typename StateVectorT::MemoryStorageT,
+                                     MemoryStorageLocation::Internal>) {
             H_lambda = std::make_unique<std::vector<StateVectorT>>(
                 num_observables, StateVectorT{lambda_qubits});
-        } else if constexpr (std::is_same_v<StateVectorLQubitRaw<PrecisionT>,
-                                            StateVectorT>) {
+        } else if constexpr (std::is_same_v<
+                                 typename StateVectorT::MemoryStorageT,
+                                 MemoryStorageLocation::External>) {
             H_lambda_storage =
                 std::make_unique<std::vector<std::vector<ComplexT>>>(
                     num_observables, std::vector<ComplexT>(lambda.getLength()));
@@ -273,6 +276,10 @@ class AdjointJacobian final
                                 (*H_lambda_storage)[ind].size());
                 H_lambda->push_back(sv);
             }
+        } else {
+            /// LCOV_EXCL_START
+            PL_ABORT("Undefined memory storage location for StateVectorT.");
+            /// LCOV_EXCL_STOP
         }
 
         StateVectorLQubitManaged<PrecisionT> mu(lambda_qubits);

pennylane_lightning/core/src/utils/Memory.hpp

@@ -172,6 +172,29 @@ bool operator!=([[maybe_unused]] const AlignedAllocator<T> &lhs,
     return lhs.alignment() != rhs.alignment();
 }
 
+/**
+ * @brief The following namespace holds compile-time tags for indicating where
+ * statevector memory storage lives.
+ */
+namespace MemoryStorageLocation {
+/**
+ * @brief Tag to indicate internal memory storage for compile-time dispatch.
+ *
+ */
+struct Internal {};
+
+/**
+ * @brief Tag to indicate external memory storage for compile-time dispatch.
+ *
+ */
+struct External {};
+/**
+ * @brief Tag to indicate undefined memory storage for compile-time dispatch.
+ *
+ */
+struct Undefined {};
+} // namespace MemoryStorageLocation
+
 ///@cond DEV
 template <class PrecisionT, class TypeList> struct commonAlignmentHelper {
     constexpr static size_t value = std::max(

pennylane_lightning/core/src/utils/Util.hpp

@@ -509,5 +509,4 @@ auto transpose_state_tensor(const std::vector<T> &tensor,
     }
     return transposed_tensor;
 }
-
 } // namespace Pennylane::Util