Template/expval

.github/CHANGELOG.md

@@ -12,6 +12,9 @@
 
 ### Improvements
 
+* Refactor LKokkos `Measurements` class to use Kokkos `RangePolicy` together with special functors to obtain the expectation value of 1- to 4-wire generic unitary gates. For more than 4 wires, the general implementation using Kokkos `TeamPolicy` is employed to yield the best all-around performance. 
+  [(#489)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/489)
+
 * Add tests to increase LKokkos coverage.
   [(#485)](https://github.com/PennyLaneAI/pennylane-lightning/pull/485)
 

pennylane_lightning/core/_version.py

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

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/README.md

@@ -0,0 +1,75 @@
+# Implementation
+
+## Expval kernels
+
+The general multi-qubit operator kernel requires a private variable `coeffs_in` to store state vector coefficients.
+In the Kokkos framework, this variable cannot be a private member of the functor, say, because all threads will overwrite it.
+One must then use `TeamPolicy`s together with `scratch_memory_space` which allows creating and manipulating thread-local variables.
+This implementation however appears suboptimal compared with the straightforward `RangePolicy` with bit-injection one.
+
+The last being more verbose, it is only implemented for 1- and 2-qubit observables.
+It is however possible to generate the code automatically for higher qubit counts with the following Python script. 
+
+```python
+for n_wires in range(1, 6):
+    name = f"getExpVal{n_wires}QubitOpFunctor"
+    nU = 2**n_wires
+
+    print(
+        f"""template <class PrecisionT, std::size_t n_wires> struct {name} {{
+
+        using ComplexT = Kokkos::complex<PrecisionT>;
+        using KokkosComplexVector = Kokkos::View<ComplexT *>;
+        using KokkosIntVector = Kokkos::View<std::size_t *>;
+
+        KokkosComplexVector arr;
+        KokkosComplexVector matrix;
+        KokkosIntVector wires;
+        std::size_t dim;
+        std::size_t num_qubits;
+
+        {name}(const KokkosComplexVector &arr_,
+        const std::size_t num_qubits_,
+                                    const KokkosComplexVector &matrix_,
+                                    const KokkosIntVector &wires_) {{
+            wires = wires_;
+            arr = arr_;
+            matrix = matrix_;
+            num_qubits = num_qubits_;
+            dim = 1U << wires.size();
+        }}
+
+        KOKKOS_INLINE_FUNCTION
+        void operator()(const std::size_t k, PrecisionT &expval) const {{
+        const std::size_t kdim = k * dim;
+    """
+    )
+
+    for k in range(nU):
+        print(
+            f"""
+        std::size_t i{k:0{n_wires}b} = kdim | {k};
+        for (std::size_t pos = 0; pos < n_wires; pos++) {{
+            std::size_t x =
+                ((i{k:0{n_wires}b} >> (n_wires - pos - 1)) ^
+                    (i{k:0{n_wires}b} >> (num_qubits - wires(pos) - 1))) &
+                1U;
+            i{k:0{n_wires}b} = i{k:0{n_wires}b} ^ ((x << (n_wires - pos - 1)) |
+                            (x << (num_qubits - wires(pos) - 1)));
+        }}
+        """
+        )
+
+    # print("expval += real(")
+    for k in range(nU):
+        tmp = f"expval += real(conj(arr(i{k:0{n_wires}b})) * ("
+        tmp += f"matrix(0B{k:0{n_wires}b}{0:0{n_wires}b}) * arr(i{0:0{n_wires}b})"
+        for j in range(1, nU):
+            tmp += (
+                f" + matrix(0B{k:0{n_wires}b}{j:0{n_wires}b}) * arr(i{j:0{n_wires}b})"
+            )
+        print(tmp, end="")
+        print("));")
+    print("}")
+    print("};")
+```

pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/CMakeLists.txt

@@ -11,6 +11,7 @@ target_link_libraries(${PL_BACKEND}_measurements INTERFACE  lightning_compile_op
                                                             lightning_observables
                                                             lightning_utils
                                                             ${PL_BACKEND}
+                                                            ${PL_BACKEND}_observables
                                                             ${PL_BACKEND}_utils
                                                             )