Template/expval (#489)

* M pennylane_lightning/core/src/bindings/Bindings.hpp; hack `JacobianData` to work with devices. M pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp; `applyMatrix` bugfix: use intermediate hostview to copy matrix data; same bugfix for `getDataVector`. M pennylane_lightning/core/src/simulators/lightning_kokkos/algorithms/AdjointJacobianKokkos.hpp; use copy constructor. M pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp; use copy constructor. M pennylane_lightning/core/src/simulators/lightning_kokkos/observables/ObservablesKokkos.hpp; use copy constructor. M requirements-dev.txt; add clang-format-14. * Auto update version * Update changelog. * Auto update version * Auto update version * Add an argument to adjointJacobian to avoid syncing and copying state vector data in adjoint-diff. * Reformat * trigger CI * [skip ci] Update changelog. * Introduce std::unordered_map<std::string, ExpValFunc> expval_funcs_. * Introduce applyExpectationValueFunctor. * Add binding to LKokkos expval(matrix, wires). Combine expval functor calls into two templated methods. Call specialized expval methods when possible. Remove obsolete 'Apply directly' tests. * Update changelog. * Add test for arbitrary expval(Hermitian). * Add getExpectationValueMultiQubitOpFunctor. * Add typename hint for macos. * Add typename macos. * Use Kokkos::ThreadVectorRange policy for innerloop in getExpectationValueMultiQubitOpFunctor. * Auto update version * Auto update version * Couple fix for HIP. * WIP * Add specialized 3-5 qubit expval functors. * Add implementation notes. * [skip ci] Polish getExpValMatrix and limit getExpValNQubitOpFunctor to N == 4 (seems optimal on OPENMP/HIP backends although CUDA can go further). * Fix whitespace warning. * Auto update version * Use simplified bitshifts in 1- & 2-qubit expval kernels. * Update tests. * Update changelog. * Bump pennylane version. * Auto update version * Reimplement expval functors with macros. * Auto update version * Update CHANGELOG.md * Auto update version * trigger CI * trigger CI * Bump Kokkos to 4.1.00 in CI. * Revert kokkos ver. * Add tests for macroed expval functors. * Remove redundant black lines. * Use matrix interface to get expval of HermitianObs in LKokkos. * Cover kokkos_args error. --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>
PennyLaneAI · Sep 7, 2023 · e96a53f · e96a53f
1 parent 8eed6f2
commit e96a53f
Show file tree

Hide file tree

Showing 11 changed files with 709 additions and 117 deletions.
diff --git a/.github/CHANGELOG.md b/.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)
 

diff --git a/pennylane_lightning/core/_version.py b/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"
diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/README.md b/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("};")
+```
diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/CMakeLists.txt b/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
                                                             )