Skip to content

Commit 5fecece

Browse files
mhuckamhucka
authored
Merge pull request #837 from mhucka/mhucka-move-balancetrajectory
Refactor to move BalanceTrajectory to separate file
2 parents cc4136a + 7bddbee commit 5fecece

File tree

10 files changed

+311
-210
lines changed

10 files changed

+311
-210
lines changed

‎tensorflow_quantum/core/ops/noise/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ cc_binary(
6060
"//tensorflow_quantum/core/ops:parse_context",
6161
"//tensorflow_quantum/core/ops:tfq_simulate_utils",
6262
"//tensorflow_quantum/core/src:circuit_parser_qsim",
63+
"//tensorflow_quantum/core/src:util_balance_trajectory",
6364
"//tensorflow_quantum/core/src:util_qsim",
6465
"@qsim//lib:qsim_lib",
6566
# tensorflow core framework

‎tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ limitations under the License.
4141
#include "tensorflow_quantum/core/ops/parse_context.h"
4242
#include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
4343
#include "tensorflow_quantum/core/proto/program.pb.h"
44+
#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
4445
#include "tensorflow_quantum/core/src/util_qsim.h"
4546

4647
namespace tfq {

‎tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ limitations under the License.
4141
#include "tensorflow_quantum/core/ops/parse_context.h"
4242
#include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
4343
#include "tensorflow_quantum/core/proto/program.pb.h"
44+
#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
4445
#include "tensorflow_quantum/core/src/util_qsim.h"
4546

4647
namespace tfq {

‎tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ limitations under the License.
4040
#include "tensorflow_quantum/core/ops/parse_context.h"
4141
#include "tensorflow_quantum/core/proto/program.pb.h"
4242
#include "tensorflow_quantum/core/src/circuit_parser_qsim.h"
43+
#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
4344
#include "tensorflow_quantum/core/src/util_qsim.h"
4445

4546
namespace tfq {

‎tensorflow_quantum/core/src/BUILD‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,12 +93,12 @@ cc_test(
9393

9494
cc_library(
9595
name = "util_qsim",
96-
srcs = [],
9796
hdrs = ["util_qsim.h"],
9897
deps = [
9998
":circuit_parser_qsim",
10099
"//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
101100
"//tensorflow_quantum/core/proto:projector_sum_cc_proto",
101+
"@com_google_absl//absl/functional:any_invocable",
102102
"@com_google_absl//absl/container:inlined_vector", # unclear why needed.
103103
"@com_google_absl//absl/status",
104104
"@com_google_absl//absl/status:statusor",
@@ -126,6 +126,25 @@ cc_test(
126126
],
127127
)
128128

129+
cc_library(
130+
name = "util_balance_trajectory",
131+
srcs = ["util_balance_trajectory.cc"],
132+
hdrs = ["util_balance_trajectory.h"],
133+
deps = [
134+
],
135+
)
136+
137+
cc_test(
138+
name = "util_balance_trajectory_test",
139+
size = "small",
140+
srcs = ["util_balance_trajectory_test.cc"],
141+
linkstatic = 0,
142+
deps = [
143+
":util_balance_trajectory",
144+
"@com_google_googletest//:gtest_main",
145+
],
146+
)
147+
129148
cc_library(
130149
name = "program_resolution",
131150
srcs = ["program_resolution.cc"],

‎tensorflow_quantum/core/src/util_balance_trajectory.cc‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
17+
18+
namespace tfq {
19+
20+
// Balance the number of trajectory computations done between
21+
// threads. num_samples is a 2d vector containing the number of reps
22+
// requested for each pauli_sum[i,j]. After running thread_offsets
23+
// contains 0/-1 values that will offset the work for each thread.
24+
// to make it as close to uniform as possible. **Assumes circuits
25+
// have roughly equal simulation cost**
26+
void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
27+
const int& num_threads,
28+
std::vector<std::vector<int>>* thread_offsets) {
29+
std::vector<int> rep_limits(num_samples.size(), -1);
30+
std::vector<int> height(num_threads, 0);
31+
32+
for (int i = 0; i < num_samples.size(); i++) {
33+
for (int j = 0; j < num_samples[i].size(); j++) {
34+
rep_limits[i] = std::max(rep_limits[i], num_samples[i][j]);
35+
}
36+
}
37+
int prev_max_height = -1;
38+
for (int j = 0; j < num_samples.size(); j++) {
39+
int run_ceiling = ((rep_limits[j] + num_threads - 1) / num_threads);
40+
int num_lo = num_threads * run_ceiling - rep_limits[j];
41+
int num_hi = num_threads - num_lo;
42+
int cur_max = prev_max_height;
43+
for (int i = 0; i < num_threads; i++) {
44+
if (height[i] == cur_max && num_lo) {
45+
// previously had extra work on this thread and
46+
// have remaining low budget to give.
47+
height[i]++;
48+
(*thread_offsets)[i][j] = -1;
49+
num_lo--;
50+
} else if (height[i] == cur_max - 1 && num_hi) {
51+
// previously had less work on this thread and
52+
// remaining high budget to give.
53+
height[i] += 2;
54+
(*thread_offsets)[i][j] = 0;
55+
num_hi--;
56+
} else if (num_hi) {
57+
height[i] += 2;
58+
(*thread_offsets)[i][j] = 0;
59+
num_hi--;
60+
} else {
61+
height[i]++;
62+
(*thread_offsets)[i][j] = -1;
63+
num_lo--;
64+
}
65+
prev_max_height = std::max(height[i], prev_max_height);
66+
}
67+
}
68+
}
69+
70+
// Simpler case of TrajectoryBalance where num_samples is fixed
71+
// across all circuits.
72+
void BalanceTrajectory(const int& num_samples, const int& num_threads,
73+
std::vector<std::vector<int>>* thread_offsets) {
74+
std::vector<int> height(num_threads, 0);
75+
76+
int prev_max_height = -1;
77+
for (int j = 0; j < (*thread_offsets)[0].size(); j++) {
78+
int run_ceiling = ((num_samples + num_threads - 1) / num_threads);
79+
int num_lo = num_threads * run_ceiling - num_samples;
80+
int num_hi = num_threads - num_lo;
81+
int cur_max = prev_max_height;
82+
for (int i = 0; i < num_threads; i++) {
83+
if (height[i] == cur_max && num_lo) {
84+
// previously had extra work on this thread and
85+
// have remaining low budget to give.
86+
height[i]++;
87+
(*thread_offsets)[i][j] = -1;
88+
num_lo--;
89+
} else if (height[i] == cur_max - 1 && num_hi) {
90+
// previously had less work on this thread and
91+
// remaining high budget to give.
92+
height[i] += 2;
93+
(*thread_offsets)[i][j] = 0;
94+
num_hi--;
95+
} else if (num_hi) {
96+
height[i] += 2;
97+
(*thread_offsets)[i][j] = 0;
98+
num_hi--;
99+
} else {
100+
height[i]++;
101+
(*thread_offsets)[i][j] = -1;
102+
num_lo--;
103+
}
104+
prev_max_height = std::max(height[i], prev_max_height);
105+
}
106+
}
107+
}
108+
109+
} // namespace tfq

‎tensorflow_quantum/core/src/util_balance_trajectory.h‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#ifndef UTIL_BALANCE_TRAJECTORY_H_
17+
#define UTIL_BALANCE_TRAJECTORY_H_
18+
19+
#include <algorithm>
20+
#include <cstdint>
21+
#include <vector>
22+
23+
namespace tfq {
24+
25+
void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
26+
const int& num_threads,
27+
std::vector<std::vector<int>>* thread_offsets);
28+
29+
void BalanceTrajectory(const int& num_samples, const int& num_threads,
30+
std::vector<std::vector<int>>* thread_offsets);
31+
32+
} // namespace tfq
33+
34+
#endif // UTIL_BALANCE_TRAJECTORY_H_

‎tensorflow_quantum/core/src/util_balance_trajectory_test.cc‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
17+
18+
#include "gtest/gtest.h"
19+
20+
namespace tfq {
21+
namespace {
22+
23+
static void AssertWellBalanced(const std::vector<std::vector<int>>& n_reps,
24+
const int& num_threads,
25+
const std::vector<std::vector<int>>& offsets) {
26+
auto max_work = std::vector<int>(n_reps.size(), -1);
27+
for (int i = 0; i < n_reps.size(); i++) {
28+
for (int j = 0; j < n_reps[0].size(); j++) {
29+
max_work[i] = std::max(max_work[i], n_reps[i][j]);
30+
}
31+
}
32+
33+
for (int i = 0; i < n_reps.size(); i++) {
34+
int sum = 0;
35+
int prev_local_work = 0;
36+
for (int k = 0; k < num_threads; k++) {
37+
int local_work = (max_work[i] + num_threads - 1) / num_threads;
38+
local_work += offsets[k][i];
39+
sum += local_work;
40+
if (k > 0) {
41+
EXPECT_LT(abs(local_work - prev_local_work), 2);
42+
}
43+
prev_local_work = local_work;
44+
}
45+
EXPECT_EQ(sum, max_work[i]);
46+
}
47+
}
48+
49+
TEST(UtilQsimTest, BalanceTrajectorySimple) {
50+
std::vector<std::vector<int>> n_reps = {{1, 3, 5, 10, 15},
51+
{1, 10, 20, 30, 40},
52+
{50, 70, 100, 100, 100},
53+
{100, 200, 200, 200, 200}};
54+
const int num_threads = 3;
55+
// [num_threads, n_reps.size()]
56+
std::vector<std::vector<int>> offsets = {
57+
{0, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}};
58+
59+
BalanceTrajectory(n_reps, num_threads, &offsets);
60+
AssertWellBalanced(n_reps, num_threads, offsets);
61+
}
62+
63+
TEST(UtilQsimTest, BalanceTrajectoryPreventIdle) {
64+
std::vector<std::vector<int>> n_reps = {{1, 1, 1, 1, 11},
65+
{1, 1, 1, 11, 1},
66+
{1, 1, 11, 1, 1},
67+
{1, 11, 1, 1, 1},
68+
{11, 1, 1, 1, 1}};
69+
const int num_threads = 10;
70+
// [num_threads, n_reps.size()]
71+
std::vector<std::vector<int>> offsets = {
72+
{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
73+
{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
74+
{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
75+
76+
BalanceTrajectory(n_reps, num_threads, &offsets);
77+
AssertWellBalanced(n_reps, num_threads, offsets);
78+
}
79+
80+
TEST(UtilQsimTest, BalanceTrajectoryLowRep) {
81+
std::vector<std::vector<int>> n_reps = {
82+
{1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1},
83+
{1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}};
84+
const int num_threads = 5;
85+
// [num_threads, n_reps.size()]
86+
std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
87+
{0, 0, 0, 0, 0, 0, 0},
88+
{0, 0, 0, 0, 0, 0, 0},
89+
{0, 0, 0, 0, 0, 0, 0},
90+
{0, 0, 0, 0, 0, 0, 0}};
91+
92+
BalanceTrajectory(n_reps, num_threads, &offsets);
93+
AssertWellBalanced(n_reps, num_threads, offsets);
94+
}
95+
96+
TEST(UtilQsimTest, BalanceTrajectoryFewHigh) {
97+
std::vector<std::vector<int>> n_reps = {
98+
{1, 100, 1, 1, 1}, {1, 1, 1, 1, 1000}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1},
99+
{1, 1, 1, 1, 1}, {1, 10, 1, 1, 1}, {1, 1, 1, 1, 1000}};
100+
const int num_threads = 5;
101+
// [num_threads, n_reps.size()]
102+
std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
103+
{0, 0, 0, 0, 0, 0, 0},
104+
{0, 0, 0, 0, 0, 0, 0},
105+
{0, 0, 0, 0, 0, 0, 0},
106+
{0, 0, 0, 0, 0, 0, 0}};
107+
108+
BalanceTrajectory(n_reps, num_threads, &offsets);
109+
AssertWellBalanced(n_reps, num_threads, offsets);
110+
}
111+
112+
TEST(UtilQsimTest, BalanceTrajectory1D) {
113+
const int n_reps = 100;
114+
const int num_threads = 5;
115+
// [num_threads, batch_size]
116+
std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
117+
{0, 0, 0, 0, 0, 0, 0},
118+
{0, 0, 0, 0, 0, 0, 0},
119+
{0, 0, 0, 0, 0, 0, 0},
120+
{0, 0, 0, 0, 0, 0, 0}};
121+
122+
std::vector<std::vector<int>> tmp(offsets[0].size(),
123+
std::vector<int>(2, n_reps));
124+
BalanceTrajectory(n_reps, num_threads, &offsets);
125+
AssertWellBalanced(tmp, num_threads, offsets);
126+
}
127+
128+
TEST(UtilQsimTest, BalanceTrajectory1D_2) {
129+
const int n_reps = 11;
130+
const int num_threads = 10;
131+
// [num_threads, batch_size]
132+
std::vector<std::vector<int>> offsets = {
133+
{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
134+
{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
135+
{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
136+
137+
std::vector<std::vector<int>> tmp(offsets[0].size(),
138+
std::vector<int>(2, n_reps));
139+
BalanceTrajectory(n_reps, num_threads, &offsets);
140+
AssertWellBalanced(tmp, num_threads, offsets);
141+
}
142+
143+
} // namespace
144+
} // namespace tfq

0 commit comments

Comments
 (0)