Merge pull request #837 from mhucka/mhucka-move-balancetrajectory

Refactor to move BalanceTrajectory to separate file

tensorflow_quantum/core/ops/noise/BUILD

@@ -60,6 +60,7 @@ cc_binary(
         "//tensorflow_quantum/core/ops:parse_context",
         "//tensorflow_quantum/core/ops:tfq_simulate_utils",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
+        "//tensorflow_quantum/core/src:util_balance_trajectory",
         "//tensorflow_quantum/core/src:util_qsim",
         "@qsim//lib:qsim_lib",
         # tensorflow core framework

tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {

tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {

tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc

@@ -40,6 +40,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
 #include "tensorflow_quantum/core/src/circuit_parser_qsim.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {

tensorflow_quantum/core/src/BUILD

@@ -93,12 +93,12 @@ cc_test(
 
 cc_library(
     name = "util_qsim",
-    srcs = [],
     hdrs = ["util_qsim.h"],
     deps = [
         ":circuit_parser_qsim",
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/container:inlined_vector",  # unclear why needed.
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -126,6 +126,25 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "util_balance_trajectory",
+    srcs = ["util_balance_trajectory.cc"],
+    hdrs = ["util_balance_trajectory.h"],
+    deps = [
+    ],
+)
+
+cc_test(
+    name = "util_balance_trajectory_test",
+    size = "small",
+    srcs = ["util_balance_trajectory_test.cc"],
+    linkstatic = 0,
+    deps = [
+        ":util_balance_trajectory",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "program_resolution",
     srcs = ["program_resolution.cc"],

tensorflow_quantum/core/src/util_balance_trajectory.cc

@@ -0,0 +1,109 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
+
+namespace tfq {
+
+// Balance the number of trajectory computations done between
+// threads. num_samples is a 2d vector containing the number of reps
+// requested for each pauli_sum[i,j]. After running thread_offsets
+// contains 0/-1 values that will offset the work for each thread.
+// to make it as close to uniform as possible. **Assumes circuits
+// have roughly equal simulation cost**
+void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
+                       const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets) {
+  std::vector<int> rep_limits(num_samples.size(), -1);
+  std::vector<int> height(num_threads, 0);
+
+  for (int i = 0; i < num_samples.size(); i++) {
+    for (int j = 0; j < num_samples[i].size(); j++) {
+      rep_limits[i] = std::max(rep_limits[i], num_samples[i][j]);
+    }
+  }
+  int prev_max_height = -1;
+  for (int j = 0; j < num_samples.size(); j++) {
+    int run_ceiling = ((rep_limits[j] + num_threads - 1) / num_threads);
+    int num_lo = num_threads * run_ceiling - rep_limits[j];
+    int num_hi = num_threads - num_lo;
+    int cur_max = prev_max_height;
+    for (int i = 0; i < num_threads; i++) {
+      if (height[i] == cur_max && num_lo) {
+        // previously had extra work on this thread and
+        // have remaining low budget to give.
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      } else if (height[i] == cur_max - 1 && num_hi) {
+        // previously had less work on this thread and
+        // remaining high budget to give.
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else if (num_hi) {
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else {
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      }
+      prev_max_height = std::max(height[i], prev_max_height);
+    }
+  }
+}
+
+// Simpler case of TrajectoryBalance where num_samples is fixed
+// across all circuits.
+void BalanceTrajectory(const int& num_samples, const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets) {
+  std::vector<int> height(num_threads, 0);
+
+  int prev_max_height = -1;
+  for (int j = 0; j < (*thread_offsets)[0].size(); j++) {
+    int run_ceiling = ((num_samples + num_threads - 1) / num_threads);
+    int num_lo = num_threads * run_ceiling - num_samples;
+    int num_hi = num_threads - num_lo;
+    int cur_max = prev_max_height;
+    for (int i = 0; i < num_threads; i++) {
+      if (height[i] == cur_max && num_lo) {
+        // previously had extra work on this thread and
+        // have remaining low budget to give.
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      } else if (height[i] == cur_max - 1 && num_hi) {
+        // previously had less work on this thread and
+        // remaining high budget to give.
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else if (num_hi) {
+        height[i] += 2;
+        (*thread_offsets)[i][j] = 0;
+        num_hi--;
+      } else {
+        height[i]++;
+        (*thread_offsets)[i][j] = -1;
+        num_lo--;
+      }
+      prev_max_height = std::max(height[i], prev_max_height);
+    }
+  }
+}
+
+}  // namespace tfq

tensorflow_quantum/core/src/util_balance_trajectory.h

@@ -0,0 +1,34 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef UTIL_BALANCE_TRAJECTORY_H_
+#define UTIL_BALANCE_TRAJECTORY_H_
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+namespace tfq {
+
+void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
+                       const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets);
+
+void BalanceTrajectory(const int& num_samples, const int& num_threads,
+                       std::vector<std::vector<int>>* thread_offsets);
+
+}  // namespace tfq
+
+#endif  // UTIL_BALANCE_TRAJECTORY_H_

tensorflow_quantum/core/src/util_balance_trajectory_test.cc

@@ -0,0 +1,144 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
+
+#include "gtest/gtest.h"
+
+namespace tfq {
+namespace {
+
+static void AssertWellBalanced(const std::vector<std::vector<int>>& n_reps,
+                               const int& num_threads,
+                               const std::vector<std::vector<int>>& offsets) {
+  auto max_work = std::vector<int>(n_reps.size(), -1);
+  for (int i = 0; i < n_reps.size(); i++) {
+    for (int j = 0; j < n_reps[0].size(); j++) {
+      max_work[i] = std::max(max_work[i], n_reps[i][j]);
+    }
+  }
+
+  for (int i = 0; i < n_reps.size(); i++) {
+    int sum = 0;
+    int prev_local_work = 0;
+    for (int k = 0; k < num_threads; k++) {
+      int local_work = (max_work[i] + num_threads - 1) / num_threads;
+      local_work += offsets[k][i];
+      sum += local_work;
+      if (k > 0) {
+        EXPECT_LT(abs(local_work - prev_local_work), 2);
+      }
+      prev_local_work = local_work;
+    }
+    EXPECT_EQ(sum, max_work[i]);
+  }
+}
+
+TEST(UtilQsimTest, BalanceTrajectorySimple) {
+  std::vector<std::vector<int>> n_reps = {{1, 3, 5, 10, 15},
+                                          {1, 10, 20, 30, 40},
+                                          {50, 70, 100, 100, 100},
+                                          {100, 200, 200, 200, 200}};
+  const int num_threads = 3;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {
+      {0, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectoryPreventIdle) {
+  std::vector<std::vector<int>> n_reps = {{1, 1, 1, 1, 11},
+                                          {1, 1, 1, 11, 1},
+                                          {1, 1, 11, 1, 1},
+                                          {1, 11, 1, 1, 1},
+                                          {11, 1, 1, 1, 1}};
+  const int num_threads = 10;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectoryLowRep) {
+  std::vector<std::vector<int>> n_reps = {
+      {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1},
+      {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}, {1, 1, 1, 1, 1}};
+  const int num_threads = 5;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectoryFewHigh) {
+  std::vector<std::vector<int>> n_reps = {
+      {1, 100, 1, 1, 1}, {1, 1, 1, 1, 1000}, {1, 1, 1, 1, 1},   {1, 1, 1, 1, 1},
+      {1, 1, 1, 1, 1},   {1, 10, 1, 1, 1},   {1, 1, 1, 1, 1000}};
+  const int num_threads = 5;
+  // [num_threads, n_reps.size()]
+  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(n_reps, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectory1D) {
+  const int n_reps = 100;
+  const int num_threads = 5;
+  // [num_threads, batch_size]
+  std::vector<std::vector<int>> offsets = {{0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0},
+                                           {0, 0, 0, 0, 0, 0, 0}};
+
+  std::vector<std::vector<int>> tmp(offsets[0].size(),
+                                    std::vector<int>(2, n_reps));
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(tmp, num_threads, offsets);
+}
+
+TEST(UtilQsimTest, BalanceTrajectory1D_2) {
+  const int n_reps = 11;
+  const int num_threads = 10;
+  // [num_threads, batch_size]
+  std::vector<std::vector<int>> offsets = {
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0},
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
+
+  std::vector<std::vector<int>> tmp(offsets[0].size(),
+                                    std::vector<int>(2, n_reps));
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+  AssertWellBalanced(tmp, num_threads, offsets);
+}
+
+}  // namespace
+}  // namespace tfq