Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

MPS Implementation p4: Two qubit gates. #372

Merged
merged 12 commits into from
Jul 14, 2021
Merged

MPS Implementation p4: Two qubit gates. #372

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
5742c86
MPS Implementation p4: Two qubit gates.
MichaelBroughton Jul 2, 2021
92125db
U + V now share scratch space better.
MichaelBroughton Jul 2, 2021
c8021cf
make s_vector a map.
MichaelBroughton Jul 2, 2021
d4981bc
Added equivalent Cirq code for fuzz test.
MichaelBroughton Jul 2, 2021
928ba7e
more comments.
MichaelBroughton Jul 2, 2021
9dd669a
removed use of CXX11-Tensor.
MichaelBroughton Jul 11, 2021
20436da
Sergei feedback.
MichaelBroughton Jul 13, 2021
6a999fc
missed renames.
MichaelBroughton Jul 13, 2021
f35e74e
Merge branch 'master' into mps
95-martin-orion Jul 14, 2021
2678fa0
winfix attempt 1.
MichaelBroughton Jul 14, 2021
41cccc3
Merge branch 'mps' of github.com:MichaelBroughton/qsim into mps
MichaelBroughton Jul 14, 2021
6129dc1
win fix 2.
MichaelBroughton Jul 14, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions WORKSPACE
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,11 @@ http_archive(
],
)


EIGEN_COMMIT = "12e8d57108c50d8a63605c6eb0144c838c128337"
EIGEN_SHA256 = "f689246e342c3955af48d26ce74ac34d21b579a00675c341721a735937919b02"


http_archive(
name = "eigen",
build_file_content = """
Expand All @@ -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),
],
)
116 changes: 94 additions & 22 deletions lib/mps_simulator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <vector>

#include "../eigen/Eigen/Dense"
#include "../eigen/Eigen/SVD"
#include "mps_statespace.h"

namespace qsim {
Expand All @@ -38,16 +39,15 @@ 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 OneQBMatrix = Eigen::Matrix<Complex, 2, 2, Eigen::RowMajor>;
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved
using ConstOneQBMap = Eigen::Map<const OneQBMatrix>;

// Note: ForArgs are currently unused.
Expand All @@ -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;
Expand Down Expand Up @@ -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.
}

Expand All @@ -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 {
Expand All @@ -130,39 +130,111 @@ 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();
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved
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);
ConstOneQBMap B((Complex*)matrix);
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved
MatrixMap C((Complex*)(raw_state + end), 2, bd);

for (unsigned block_sep = l_offset; block_sep < r_offset;
block_sep += 4 * bd) {
fp_type* cur_block = raw_state + block_sep;
ConstMatrixMap A =
ConstMatrixMap((Complex*)(cur_block), 2, bd);
ConstMatrixMap A((Complex*)(cur_block), 2, bd);
C.noalias() = B * A;
memcpy(cur_block, raw_state + end, sizeof(fp_type) * bd * 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();
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved
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);
ConstOneQBMap B((Complex*)matrix);
ConstMatrixMap A((Complex*)(raw_state + offset), bd, 2);
MatrixMap C((Complex*)(raw_state + end), bd, 2);
C.noalias() = A * B.transpose();
memcpy(raw_state + offset, raw_state + end, sizeof(fp_type) * bd * 4);
}

void ApplyGate2(const std::vector<unsigned>& qs, const fp_type* matrix,
State& state) const {
// TODO: micro-benchmark this function and improve performance.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this TODO resolved? If not, please open an issue for it and reference that issue here.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is not resolved. Opened issue here: #373

const auto bd = state.bond_dim();
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved
const auto nq = state.num_qubits();
fp_type* raw_state = state.get();
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved

const auto i_dim = (qs[0] == 0) ? 1 : bd;
const auto j_dim = 2;
const auto k_dim = bd;
const auto l_dim = 2;
const auto m_dim = (qs[1] == nq - 1) ? 1 : bd;

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 B_0((Complex*)(raw_state + B_0_offset), i_dim * j_dim, k_dim);
MatrixMap B_1((Complex*)(raw_state + B_1_offset), k_dim, l_dim * m_dim);

// Merge both blocks into scratch space.
MatrixMap C((Complex*)(raw_state + end), i_dim * j_dim, l_dim * m_dim);
C.noalias() = B_0 * B_1;

// Transpose inner dims in-place.
MatrixMap 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){
C_t.row(i + 1).swap(C_t.row(i + 2));
}

// Transpose gate matrix and place in 3rd (last) scratch block.
const auto scratch3_offset = end + 8 * bd * bd;
ConstMatrixMap G_mat((Complex*)matrix, 4, 4);
MatrixMap G_t_mat((Complex*)(raw_state + scratch3_offset), 4, 4);
G_t_mat = G_mat.transpose();
G_t_mat.col(1).swap(G_t_mat.col(2));

// Contract gate and merged block tensors, placing result in B0B1.
for (unsigned i = 0; i < i_dim; i++) {
MichaelBroughton marked this conversation as resolved.
Show resolved Hide resolved
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 K_i((Complex*)dest_block, 4, m_dim);
ConstMatrixMap C_i((Complex*)src_block, 4, m_dim);
// [i, np, m] = [np, lj] * [i, lj, m]
K_i.noalias() = G_t_mat * C_i;
}

// SVD B0B1.
MatrixMap K((Complex*)(raw_state + B_0_offset), 2 * i_dim, 2 * m_dim);
Eigen::BDCSVD<Matrix> svd(K, 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 U((Complex*)(raw_state + end), 2 * i_dim, p);
U.noalias() = svd.matrixU();
B_0.fill(Complex(0, 0));
const auto keep_cols = (U.cols() > bd) ? bd : U.cols();
B_0.block(0, 0, U.rows(), keep_cols).noalias() =
U(Eigen::all, Eigen::seq(0, keep_cols - 1));

// Place row product of S V into scratch to truncate and then B1.
MatrixMap V((Complex*)(raw_state + end), p, 2 * m_dim);
MatrixMap s_vector((Complex*)(raw_state + end + 8 * bd * bd), p, 1);
V.noalias() = svd.matrixV().adjoint();
s_vector.noalias() = svd.singularValues();
B_1.fill(Complex(0, 0));
const auto keep_rows = (V.rows() > bd) ? bd : V.rows();
const auto row_seq = Eigen::seq(0, keep_rows - 1);
for (unsigned i = 0; i < keep_rows; i++) {
V.row(i) *= s_vector(i);
}
B_1.block(0, 0, keep_rows, V.cols()).noalias() = V(row_seq, Eigen::all);
}

For for_;
};

Expand Down
4 changes: 2 additions & 2 deletions lib/mps_statespace.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,8 +93,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);
Expand Down
2 changes: 2 additions & 0 deletions tests/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -581,6 +581,8 @@ cc_test(
}),
deps = [
"@com_google_googletest//:gtest_main",
"//lib:gate_appl",
"//lib:gates_cirq",
"//lib:mps_simulator",
"//lib:formux",
],
Expand Down
Loading