Merge pull request #372 from MichaelBroughton/mps

MPS Implementation p4: Two qubit gates.
quantumlib · Jul 14, 2021 · 0514a02 · 0514a02
2 parents cae9570 + 6129dc1
commit 0514a02
Show file tree

Hide file tree

Showing 5 changed files with 693 additions and 46 deletions.
diff --git a/WORKSPACE b/WORKSPACE
@@ -18,6 +18,11 @@ http_archive(
     ],
 )
 
+
+EIGEN_COMMIT = "12e8d57108c50d8a63605c6eb0144c838c128337"
+EIGEN_SHA256 = "f689246e342c3955af48d26ce74ac34d21b579a00675c341721a735937919b02"
+
+
 http_archive(
     name = "eigen",
     build_file_content = """
@@ -27,10 +32,10 @@ cc_library(
   visibility = ["//visibility:public"],
 )
     """,
-    sha256 = "a3c10a8c14f55e9f09f98b0a0ac6874c21bda91f65b7469d9b1f6925990e867b",  # SHARED_EIGEN_SHA
-        strip_prefix = "eigen-d10b27fe37736d2944630ecd7557cefa95cf87c9",
+    sha256 = EIGEN_SHA256,
+        strip_prefix = "eigen-{commit}".format(commit = EIGEN_COMMIT),
         urls = [
-            "https://storage.googleapis.com/mirror.tensorflow.org/gitlab.com/libeigen/eigen/-/archive/d10b27fe37736d2944630ecd7557cefa95cf87c9/eigen-d10b27fe37736d2944630ecd7557cefa95cf87c9.tar.gz",
-            "https://gitlab.com/libeigen/eigen/-/archive/d10b27fe37736d2944630ecd7557cefa95cf87c9/eigen-d10b27fe37736d2944630ecd7557cefa95cf87c9.tar.gz",
+            "https://storage.googleapis.com/mirror.tensorflow.org/gitlab.com/libeigen/eigen/-/archive/{commit}/eigen-{commit}.tar.gz".format(commit = EIGEN_COMMIT),
+            "https://gitlab.com/libeigen/eigen/-/archive/{commit}/eigen-{commit}.tar.gz".format(commit = EIGEN_COMMIT),
         ],
 )
diff --git a/lib/mps_simulator.h b/lib/mps_simulator.h
@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "../eigen/Eigen/Dense"
+#include "../eigen/Eigen/SVD"
 #include "mps_statespace.h"
 
 namespace qsim {
@@ -38,17 +39,16 @@ template <typename For, typename fp_type = float>
 class MPSSimulator final {
  public:
   using MPSStateSpace_ = MPSStateSpace<For, fp_type>;
-  using MPS = typename MPSStateSpace_::MPS;
+  using State = typename MPSStateSpace_::MPS;
 
   using Complex = std::complex<fp_type>;
-  using Matrix = Eigen::Matrix<Complex, Eigen::Dynamic,
-                               Eigen::Dynamic, Eigen::RowMajor>;
+  using Matrix =
+      Eigen::Matrix<Complex, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
   using ConstMatrixMap = Eigen::Map<const Matrix>;
   using MatrixMap = Eigen::Map<Matrix>;
 
-  using OneQBMatrix =
-      Eigen::Matrix<Complex, 2, 2, Eigen::RowMajor>;
-  using ConstOneQBMap = Eigen::Map<const OneQBMatrix>;
+  using OneQubitMatrix = Eigen::Matrix<Complex, 2, 2, Eigen::RowMajor>;
+  using ConstOneQubitMap = Eigen::Map<const OneQubitMatrix>;
 
   // Note: ForArgs are currently unused.
   template <typename... ForArgs>
@@ -61,16 +61,16 @@ class MPSSimulator final {
    * @param state The state of the system, to be updated by this method.
    */
   void ApplyGate(const std::vector<unsigned>& qs, const fp_type* matrix,
-                 MPS& state) const {
+                 State& state) const {
     // Assume qs[0] < qs[1] < qs[2] < ... .
 
     switch (qs.size()) {
       case 1:
         ApplyGate1(qs, matrix, state);
         break;
-      // case 2:
-      //   ApplyGate2(qs, matrix, state);
-      //   break;
+      case 2:
+        ApplyGate2(qs, matrix, state);
+        break;
       // case 3:
       //   ApplyGate3(qs, matrix, state);
       //   break;
@@ -99,7 +99,7 @@ class MPSSimulator final {
    */
   void ApplyControlledGate(const std::vector<unsigned>& qs,
                            const std::vector<unsigned>& cqs, uint64_t cmask,
-                           const fp_type* matrix, MPS& state) const {
+                           const fp_type* matrix, State& state) const {
     // TODO.
   }
 
@@ -113,15 +113,15 @@ class MPSSimulator final {
    */
   std::complex<double> ExpectationValue(const std::vector<unsigned>& qs,
                                         const fp_type* matrix,
-                                        const MPS& state) const {
+                                        const State& state) const {
     // Assume qs[0] < qs[1] < qs[2] < ... .
     // TODO.
     return std::complex<double>(-10., -10.);
   }
 
  private:
   void ApplyGate1(const std::vector<unsigned>& qs, const fp_type* matrix,
-                  MPS& state) const {
+                  State& state) const {
     if (qs[0] == state.num_qubits() - 1) {
       Apply1Right(qs, matrix, state);
     } else {
@@ -130,37 +130,109 @@ class MPSSimulator final {
   }
 
   void Apply1LeftOrInterior(const std::vector<unsigned>& qs,
-                            const fp_type* matrix, MPS& state) const {
+                            const fp_type* matrix, State& state) const {
     fp_type* raw_state = state.get();
-    const auto bd = state.bond_dim();
+    const auto bond_dim = state.bond_dim();
     const auto l_offset = MPSStateSpace_::GetBlockOffset(state, qs[0]);
     const auto r_offset = MPSStateSpace_::GetBlockOffset(state, qs[0] + 1);
     const auto end = MPSStateSpace_::Size(state);
-    ConstOneQBMap B = ConstOneQBMap((Complex*)matrix);
-    MatrixMap C = MatrixMap((Complex*)(raw_state + end), 2, bd);
+    ConstOneQubitMap gate_matrix((Complex*) matrix);
+    MatrixMap scratch_block((Complex*)(raw_state + end), 2, bond_dim);
 
     for (unsigned block_sep = l_offset; block_sep < r_offset;
-         block_sep += 4 * bd) {
+         block_sep += 4 * bond_dim) {
       fp_type* cur_block = raw_state + block_sep;
-      ConstMatrixMap A =
-          ConstMatrixMap((Complex*)(cur_block), 2, bd);
-      C.noalias() = B * A;
-      memcpy(cur_block, raw_state + end, sizeof(fp_type) * bd * 4);
+      ConstMatrixMap mps_block((Complex*) cur_block, 2, bond_dim);
+      scratch_block.noalias() = gate_matrix * mps_block;
+      memcpy(cur_block, raw_state + end, sizeof(fp_type) * bond_dim * 4);
     }
   }
 
   void Apply1Right(const std::vector<unsigned>& qs, const fp_type* matrix,
-                   MPS& state) const {
+                   State& state) const {
     fp_type* raw_state = state.get();
-    const auto bd = state.bond_dim();
+    const auto bond_dim = state.bond_dim();
     const auto offset = MPSStateSpace_::GetBlockOffset(state, qs[0]);
     const auto end = MPSStateSpace_::Size(state);
-    ConstOneQBMap B = ConstOneQBMap((Complex*)matrix);
-    ConstMatrixMap A =
-        ConstMatrixMap((Complex*)(raw_state + offset), bd, 2);
-    MatrixMap C = MatrixMap((Complex*)(raw_state + end), bd, 2);
-    C.noalias() = A * B.transpose();
-    memcpy(raw_state + offset, raw_state + end, sizeof(fp_type) * bd * 4);
+    ConstOneQubitMap gate_matrix((Complex*) matrix);
+    ConstMatrixMap mps_block((Complex*)(raw_state + offset), bond_dim, 2);
+    MatrixMap scratch_block((Complex*)(raw_state + end), bond_dim, 2);
+    scratch_block.noalias() = mps_block * gate_matrix.transpose();
+    memcpy(raw_state + offset, raw_state + end, sizeof(fp_type) * bond_dim * 4);
+  }
+
+  void ApplyGate2(const std::vector<unsigned>& qs, const fp_type* matrix,
+                  State& state) const {
+    // TODO: micro-benchmark this function and improve performance.
+    const auto bond_dim = state.bond_dim();
+    const auto num_qubits = state.num_qubits();
+    fp_type* raw_state = state.get();
+
+    const auto i_dim = (qs[0] == 0) ? 1 : bond_dim;
+    const auto j_dim = 2;
+    const auto k_dim = bond_dim;
+    const auto l_dim = 2;
+    const auto m_dim = (qs[1] == num_qubits - 1) ? 1 : bond_dim;
+
+    const auto b_0_offset = MPSStateSpace_::GetBlockOffset(state, qs[0]);
+    const auto b_1_offset = MPSStateSpace_::GetBlockOffset(state, qs[1]);
+    const auto end = MPSStateSpace_::Size(state);
+
+    MatrixMap block_0((Complex*)(raw_state + b_0_offset), i_dim * j_dim, k_dim);
+    MatrixMap block_1((Complex*)(raw_state + b_1_offset), k_dim, l_dim * m_dim);
+
+    // Merge both blocks into scratch space.
+    MatrixMap scratch_c((Complex*)(raw_state + end), i_dim * j_dim, l_dim * m_dim);
+    scratch_c.noalias() = block_0 * block_1;
+
+    // Transpose inner dims in-place.
+    MatrixMap scratch_c_t((Complex*)(raw_state + end), i_dim * j_dim * l_dim, m_dim);
+    for (unsigned i = 0; i < i_dim * j_dim * l_dim; i += 4) {
+      scratch_c_t.row(i + 1).swap(scratch_c_t.row(i + 2));
+    }
+
+    // Transpose gate matrix and place in 3rd (last) scratch block.
+    const auto scratch3_offset = end + 8 * bond_dim * bond_dim;
+    ConstMatrixMap gate_matrix((Complex*) matrix, 4, 4);
+    MatrixMap gate_matrix_transpose((Complex*)(raw_state + scratch3_offset), 4, 4);
+    gate_matrix_transpose = gate_matrix.transpose();
+    gate_matrix_transpose.col(1).swap(gate_matrix_transpose.col(2));
+
+    // Contract gate and merged block tensors, placing result in B0B1.
+    for (unsigned i = 0; i < i_dim; ++i) {
+      fp_type* src_block = raw_state + end + i * 8 * m_dim;
+      fp_type* dest_block = raw_state + b_0_offset + i * 8 * m_dim;
+      MatrixMap block_b0b1((Complex*) dest_block, 4, m_dim);
+      ConstMatrixMap scratch_c_i((Complex*) src_block, 4, m_dim);
+      // [i, np, m] = [np, lj] * [i, lj, m]
+      block_b0b1.noalias() = gate_matrix_transpose * scratch_c_i;
+    }
+
+    // SVD B0B1.
+    MatrixMap full_b0b1((Complex*)(raw_state + b_0_offset), 2 * i_dim, 2 * m_dim);
+    Eigen::BDCSVD<Matrix> svd(full_b0b1, Eigen::ComputeThinU | Eigen::ComputeThinV);
+    const auto p = std::min(2 * i_dim, 2 * m_dim);
+
+    // Place U in scratch to truncate and then B0.
+    MatrixMap svd_u((Complex*)(raw_state + end), 2 * i_dim, p);
+    svd_u.noalias() = svd.matrixU();
+    block_0.fill(Complex(0, 0));
+    const auto keep_cols = (svd_u.cols() > bond_dim) ? bond_dim : svd_u.cols();
+    block_0.block(0, 0, svd_u.rows(), keep_cols).noalias() =
+        svd_u(Eigen::all, Eigen::seq(0, keep_cols - 1));
+
+    // Place row product of S V into scratch to truncate and then B1.
+    MatrixMap svd_v((Complex*)(raw_state + end), p, 2 * m_dim);
+    MatrixMap s_vector((Complex*)(raw_state + end + 8 * bond_dim * bond_dim), p, 1);
+    svd_v.noalias() = svd.matrixV().adjoint();
+    s_vector.noalias() = svd.singularValues();
+    block_1.fill(Complex(0, 0));
+    const auto keep_rows = (svd_v.rows() > bond_dim) ? bond_dim : svd_v.rows();
+    const auto row_seq = Eigen::seq(0, keep_rows - 1);
+    for (unsigned i = 0; i < keep_rows; ++i) {
+      svd_v.row(i) *= s_vector(i);
+    }
+    block_1.block(0, 0, keep_rows, svd_v.cols()).noalias() = svd_v(row_seq, Eigen::all);
   }
 
   For for_;

diff --git a/lib/mps_statespace.h b/lib/mps_statespace.h
@@ -18,13 +18,21 @@
 // For templates will take care of parallelization.
 #define EIGEN_DONT_PARALLELIZE 1
 
+#ifdef _WIN32
+  #include <malloc.h>
+#endif
+
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
 #include <memory>
 
 #include "../eigen/Eigen/Dense"
 
+namespace qsim {
+
+namespace mps {
+
 namespace detail {
 
 inline void do_not_free(void*) {}
@@ -39,9 +47,6 @@ inline void free(void* ptr) {
 
 }  // namespace detail
 
-namespace qsim {
-
-namespace mps {
 /**
  * Class containing context and routines for fixed bond dimension
  * truncated Matrix Product State (MPS) simulation.
@@ -93,8 +98,8 @@ class MPSStateSpace {
   // Requires num_qubits >= 2 and bond_dim >= 2.
   static MPS CreateMPS(unsigned num_qubits, unsigned bond_dim) {
     auto end_sizes = 2 * 4 * bond_dim;
-    auto internal_sizes = 4 * bond_dim * bond_dim * num_qubits;
-    // Use two extra "internal style" blocks past the end of the
+    auto internal_sizes = 4 * bond_dim * bond_dim * (num_qubits + 1);
+    // Use three extra "internal style" blocks past the end of the
     //   working allocation for scratch space. Needed for gate
     //   application.
     auto size = sizeof(fp_type) * (end_sizes + internal_sizes);

diff --git a/tests/BUILD b/tests/BUILD
@@ -581,6 +581,8 @@ cc_test(
     }),
     deps = [
         "@com_google_googletest//:gtest_main",
+        "//lib:gate_appl",
+        "//lib:gates_cirq",
         "//lib:mps_simulator",
         "//lib:formux",
     ],