Switch to using tf philox random. (tensorflow#551)

huckstar · May 3, 2021 · b9f15e1 · b9f15e1
1 parent ebe47ac
commit b9f15e1
Show file tree

Hide file tree

Showing 8 changed files with 156 additions and 88 deletions.
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
@@ -35,7 +35,10 @@ limitations under the License.
 #include "tensorflow/core/lib/core/error_codes.pb.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/lib/random/simple_philox.h"
 #include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/util/guarded_philox_random.h"
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
@@ -170,12 +173,17 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
     auto sv = ss.Create(largest_nq);
     auto scratch = ss.Create(largest_nq);
 
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
-    std::mt19937 gen(r_seed);
-    std::uniform_int_distribution<> distrib(1, 1 << 30);
+    tensorflow::GuardedPhiloxRandom random_gen;
+    int max_n_shots = 1;
+    for (int i = 0; i < num_samples.size(); i++) {
+      for (int j = 0; j < num_samples[i].size(); j++) {
+        max_n_shots = std::max(max_n_shots, num_samples[i][j]);
+      }
+    }
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
+    auto local_gen =
+        random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
+    tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // Simulate programs one by one. Parallelizing over state vectors
     // we no longer parallelize over circuits. Each time we encounter a
@@ -208,7 +216,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
       while (1) {
         ss.SetStateZero(sv);
 
-        QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
+        QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
                              scratch, sv, unused_stats);
 
         // Use this trajectory as a source for all expectation calculations.
@@ -270,10 +278,14 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
     output_tensor->setZero();
 
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
+    tensorflow::GuardedPhiloxRandom random_gen;
+    int max_n_shots = 1;
+    for (int i = 0; i < num_samples.size(); i++) {
+      for (int j = 0; j < num_samples[i].size(); j++) {
+        max_n_shots = std::max(max_n_shots, num_samples[i][j]);
+      }
+    }
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
 
     Status compute_status = Status::OK();
     auto c_lock = tensorflow::mutex();
@@ -286,8 +298,11 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
       auto sv = ss.Create(largest_nq);
       auto scratch = ss.Create(largest_nq);
 
-      std::mt19937 gen(r_seed + start);
-      std::uniform_int_distribution<> distrib(1, 1 << 30);
+      int n_rand = ncircuits.size() * max_n_shots + 1;
+      n_rand = (n_rand + num_threads) / num_threads;
+      auto local_gen =
+          random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
+      tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (int i = 0; i < ncircuits.size(); i++) {
         int nq = num_qubits[i];
@@ -318,8 +333,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
         while (1) {
           ss.SetStateZero(sv);
 
-          QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
-                               scratch, sv, unused_stats);
+          QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss,
+                               sim, scratch, sv, unused_stats);
 
           // Compute expectations across all ops using this trajectory.
           for (int j = 0; j < pauli_sums[i].size(); j++) {

diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
@@ -35,7 +35,10 @@ limitations under the License.
 #include "tensorflow/core/lib/core/error_codes.pb.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/lib/random/simple_philox.h"
 #include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/util/guarded_philox_random.h"
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
@@ -171,13 +174,20 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
     auto sv = ss.Create(largest_nq);
     auto scratch = ss.Create(largest_nq);
 
-    // time since epoch seeds random generator.
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
-    std::mt19937 gen(r_seed);
-    std::uniform_int_distribution<> distrib(1, 1 << 30);
+    tensorflow::GuardedPhiloxRandom random_gen;
+    int max_psum_length = 1;
+    int max_n_shots = 1;
+    for (int i = 0; i < pauli_sums.size(); i++) {
+      for (int j = 0; j < pauli_sums[i].size(); j++) {
+        max_psum_length =
+            std::max(max_psum_length, pauli_sums[i][j].terms().size());
+        max_n_shots = std::max(max_n_shots, num_samples[i][j]);
+      }
+    }
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
+    auto local_gen = random_gen.ReserveSamples128(
+        ncircuits.size() * (1 + max_psum_length) * max_n_shots);
+    tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // Simulate programs one by one. Parallelizing over state vectors
     // we no longer parallelize over circuits. Each time we encounter a
@@ -210,7 +220,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
       while (1) {
         ss.SetStateZero(sv);
 
-        QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
+        QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
                              scratch, sv, unused_stats);
 
         // Use this trajectory as a source for all expectation calculations.
@@ -221,7 +231,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
           float exp_v = 0.0;
           OP_REQUIRES_OK(context, ComputeSampledExpectationQsim(
                                       pauli_sums[i][j], sim, ss, sv, scratch, 1,
-                                      gen, &exp_v));
+                                      rand_source, &exp_v));
           rolling_sums[j] += static_cast<double>(exp_v);
           run_samples[j]++;
         }
@@ -272,10 +282,17 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
 
     output_tensor->setZero();
 
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
+    tensorflow::GuardedPhiloxRandom random_gen;
+    int max_psum_length = 1;
+    int max_n_shots = 1;
+    for (int i = 0; i < pauli_sums.size(); i++) {
+      for (int j = 0; j < pauli_sums[i].size(); j++) {
+        max_psum_length =
+            std::max(max_psum_length, pauli_sums[i][j].terms().size());
+        max_n_shots = std::max(max_n_shots, num_samples[i][j]);
+      }
+    }
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
 
     Status compute_status = Status::OK();
     auto c_lock = tensorflow::mutex();
@@ -288,9 +305,10 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
       auto sv = ss.Create(largest_nq);
       auto scratch = ss.Create(largest_nq);
 
-      // time since epoch seeds random generator.
-      std::mt19937 gen(r_seed + start);
-      std::uniform_int_distribution<> distrib(1, 1 << 30);
+      int num_rand = ncircuits.size() * (1 + max_psum_length) * max_n_shots;
+      num_rand = (num_rand + num_threads) / num_threads + 1;
+      auto local_gen = random_gen.ReserveSamples128(num_rand);
+      tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (int i = 0; i < ncircuits.size(); i++) {
         int nq = num_qubits[i];
@@ -321,8 +339,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
         while (1) {
           ss.SetStateZero(sv);
 
-          QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
-                               scratch, sv, unused_stats);
+          QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss,
+                               sim, scratch, sv, unused_stats);
 
           // Compute expectations across all ops using this trajectory.
           for (int j = 0; j < pauli_sums[i].size(); j++) {
@@ -334,7 +352,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
             NESTED_FN_STATUS_SYNC(
                 compute_status,
                 ComputeSampledExpectationQsim(pauli_sums[i][j], sim, ss, sv,
-                                              scratch, 1, gen, &exp_v),
+                                              scratch, 1, rand_source, &exp_v),
                 c_lock);
             rolling_sums[j] += static_cast<double>(exp_v);
             run_samples[j]++;

diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
@@ -15,7 +15,6 @@ limitations under the License.
 
 #include <stdlib.h>
 
-#include <random>
 #include <string>
 
 #include "../qsim/lib/channel.h"
@@ -36,6 +35,9 @@ limitations under the License.
 #include "tensorflow/core/lib/core/error_codes.pb.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/lib/random/simple_philox.h"
+#include "tensorflow/core/util/guarded_philox_random.h"
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/src/circuit_parser_qsim.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
@@ -149,12 +151,11 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
     auto sv = ss.Create(largest_nq);
     auto scratch = ss.Create(largest_nq);
 
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
-    std::mt19937 gen(r_seed);
-    std::uniform_int_distribution<> distrib(1, 1 << 30);
+    tensorflow::GuardedPhiloxRandom random_gen;
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
+    auto local_gen =
+        random_gen.ReserveSamples32(2 * num_samples * ncircuits.size() + 2);
+    tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // Simulate programs one by one. Parallelizing over state vectors
     // we no longer parallelize over circuits. Each time we encounter a
@@ -180,7 +181,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
       for (int j = 0; j < num_samples; j++) {
         ss.SetStateZero(sv);
 
-        QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
+        QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
                              scratch, sv, gathered_samples);
         uint64_t q_ind = 0;
         uint64_t mask = 1;
@@ -236,10 +237,8 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
       }
     }
 
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
+    tensorflow::GuardedPhiloxRandom random_gen;
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
 
     auto DoWork = [&](int start, int end) {
       // Begin simulation.
@@ -250,8 +249,11 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
       auto sv = ss.Create(largest_nq);
       auto scratch = ss.Create(largest_nq);
 
-      std::mt19937 gen(r_seed + start);
-      std::uniform_int_distribution<> distrib(1, 1 << 30);
+      int needed_random =
+          4 * (num_samples * ncircuits.size() + num_threads) / num_threads;
+      needed_random += 4;
+      auto local_gen = random_gen.ReserveSamples32(needed_random);
+      tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (int i = 0; i < ncircuits.size(); i++) {
         int nq = num_qubits[i];
@@ -277,8 +279,8 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
         while (1) {
           ss.SetStateZero(sv);
-          QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
-                               scratch, sv, gathered_samples);
+          QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss,
+                               sim, scratch, sv, gathered_samples);
 
           uint64_t q_ind = 0;
           uint64_t mask = 1;

diff --git a/tensorflow_quantum/core/ops/tfq_simulate_sampled_expectation_op.cc b/tensorflow_quantum/core/ops/tfq_simulate_sampled_expectation_op.cc
@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 
 #include <memory>
-#include <random>
 #include <vector>
 
 #include "../qsim/lib/circuit.h"
@@ -29,7 +28,10 @@ limitations under the License.
 #include "tensorflow/core/lib/core/error_codes.pb.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/lib/random/simple_philox.h"
 #include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/util/guarded_philox_random.h"
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
@@ -158,11 +160,17 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
     auto sv = ss.Create(largest_nq);
     auto scratch = ss.Create(largest_nq);
 
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
-    std::mt19937 gen(r_seed);
+    tensorflow::GuardedPhiloxRandom random_gen;
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
+    int largest_sum = -1;
+    for (const auto& sums : pauli_sums) {
+      for (const auto& sum : sums) {
+        largest_sum = std::max(largest_sum, sum.terms().size());
+      }
+    }
+    auto local_gen = random_gen.ReserveSamples32(
+        largest_sum * pauli_sums[0].size() * fused_circuits.size() + 1);
+    tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // Simulate programs one by one. Parallelizing over state vectors
     // we no longer parallelize over circuits. Each time we encounter a
@@ -192,7 +200,7 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
         float exp_v = 0.0;
         OP_REQUIRES_OK(context, ComputeSampledExpectationQsim(
                                     pauli_sums[i][j], sim, ss, sv, scratch,
-                                    num_samples[i][j], gen, &exp_v));
+                                    num_samples[i][j], rand_source, &exp_v));
         (*output_tensor)(i, j) = exp_v;
       }
     }
@@ -211,10 +219,17 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
 
     const int output_dim_op_size = output_tensor->dimension(1);
 
-    unsigned long r_seed =
-        std::chrono::duration_cast<std::chrono::milliseconds>(
-            std::chrono::system_clock::now().time_since_epoch())
-            .count();
+    tensorflow::GuardedPhiloxRandom random_gen;
+    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
+    int largest_sum = -1;
+    for (const auto& sums : pauli_sums) {
+      for (const auto& sum : sums) {
+        largest_sum = std::max(largest_sum, sum.terms().size());
+      }
+    }
+    const int num_threads = context->device()
+                                ->tensorflow_cpu_worker_threads()
+                                ->workers->NumThreads();
 
     Status compute_status = Status::OK();
     auto c_lock = tensorflow::mutex();
@@ -229,7 +244,11 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
       auto sv = ss.Create(largest_nq);
       auto scratch = ss.Create(largest_nq);
 
-      std::mt19937 gen(r_seed + start);
+      int n_random = largest_sum * output_dim_op_size * fused_circuits.size();
+      n_random /= num_threads;
+      n_random += 1;
+      auto local_gen = random_gen.ReserveSamples32(n_random);
+      tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (int i = start; i < end; i++) {
         cur_batch_index = i / output_dim_op_size;
@@ -264,7 +283,8 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
             compute_status,
             ComputeSampledExpectationQsim(
                 pauli_sums[cur_batch_index][cur_op_index], sim, ss, sv, scratch,
-                num_samples[cur_batch_index][cur_op_index], gen, &exp_v),
+                num_samples[cur_batch_index][cur_op_index], rand_source,
+                &exp_v),
             c_lock);
 
         (*output_tensor)(cur_batch_index, cur_op_index) = exp_v;