Skip to content

Commit b9f15e1

Browse files
MichaelBroughtonMichaelBroughton
authored
Switch to using tf philox random. (#551)
1 parent ebe47ac commit b9f15e1

File tree

8 files changed

+156
-88
lines changed

8 files changed

+156
-88
lines changed

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

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,10 @@ limitations under the License.
3535
#include "tensorflow/core/lib/core/error_codes.pb.h"
3636
#include "tensorflow/core/lib/core/status.h"
3737
#include "tensorflow/core/lib/core/threadpool.h"
38+
#include "tensorflow/core/lib/random/random.h"
39+
#include "tensorflow/core/lib/random/simple_philox.h"
3840
#include "tensorflow/core/platform/mutex.h"
41+
#include "tensorflow/core/util/guarded_philox_random.h"
3942
#include "tensorflow_quantum/core/ops/parse_context.h"
4043
#include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
4144
#include "tensorflow_quantum/core/src/util_qsim.h"
@@ -170,12 +173,17 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
170173
auto sv = ss.Create(largest_nq);
171174
auto scratch = ss.Create(largest_nq);
172175

173-
unsigned long r_seed =
174-
std::chrono::duration_cast<std::chrono::milliseconds>(
175-
std::chrono::system_clock::now().time_since_epoch())
176-
.count();
177-
std::mt19937 gen(r_seed);
178-
std::uniform_int_distribution<> distrib(1, 1 << 30);
176+
tensorflow::GuardedPhiloxRandom random_gen;
177+
int max_n_shots = 1;
178+
for (int i = 0; i < num_samples.size(); i++) {
179+
for (int j = 0; j < num_samples[i].size(); j++) {
180+
max_n_shots = std::max(max_n_shots, num_samples[i][j]);
181+
}
182+
}
183+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
184+
auto local_gen =
185+
random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
186+
tensorflow::random::SimplePhilox rand_source(&local_gen);
179187

180188
// Simulate programs one by one. Parallelizing over state vectors
181189
// we no longer parallelize over circuits. Each time we encounter a
@@ -208,7 +216,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
208216
while (1) {
209217
ss.SetStateZero(sv);
210218

211-
QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
219+
QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
212220
scratch, sv, unused_stats);
213221

214222
// Use this trajectory as a source for all expectation calculations.
@@ -270,10 +278,14 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
270278

271279
output_tensor->setZero();
272280

273-
unsigned long r_seed =
274-
std::chrono::duration_cast<std::chrono::milliseconds>(
275-
std::chrono::system_clock::now().time_since_epoch())
276-
.count();
281+
tensorflow::GuardedPhiloxRandom random_gen;
282+
int max_n_shots = 1;
283+
for (int i = 0; i < num_samples.size(); i++) {
284+
for (int j = 0; j < num_samples[i].size(); j++) {
285+
max_n_shots = std::max(max_n_shots, num_samples[i][j]);
286+
}
287+
}
288+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
277289

278290
Status compute_status = Status::OK();
279291
auto c_lock = tensorflow::mutex();
@@ -286,8 +298,11 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
286298
auto sv = ss.Create(largest_nq);
287299
auto scratch = ss.Create(largest_nq);
288300

289-
std::mt19937 gen(r_seed + start);
290-
std::uniform_int_distribution<> distrib(1, 1 << 30);
301+
int n_rand = ncircuits.size() * max_n_shots + 1;
302+
n_rand = (n_rand + num_threads) / num_threads;
303+
auto local_gen =
304+
random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
305+
tensorflow::random::SimplePhilox rand_source(&local_gen);
291306

292307
for (int i = 0; i < ncircuits.size(); i++) {
293308
int nq = num_qubits[i];
@@ -318,8 +333,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
318333
while (1) {
319334
ss.SetStateZero(sv);
320335

321-
QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
322-
scratch, sv, unused_stats);
336+
QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss,
337+
sim, scratch, sv, unused_stats);
323338

324339
// Compute expectations across all ops using this trajectory.
325340
for (int j = 0; j < pauli_sums[i].size(); j++) {

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

Lines changed: 37 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,10 @@ limitations under the License.
3535
#include "tensorflow/core/lib/core/error_codes.pb.h"
3636
#include "tensorflow/core/lib/core/status.h"
3737
#include "tensorflow/core/lib/core/threadpool.h"
38+
#include "tensorflow/core/lib/random/random.h"
39+
#include "tensorflow/core/lib/random/simple_philox.h"
3840
#include "tensorflow/core/platform/mutex.h"
41+
#include "tensorflow/core/util/guarded_philox_random.h"
3942
#include "tensorflow_quantum/core/ops/parse_context.h"
4043
#include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
4144
#include "tensorflow_quantum/core/src/util_qsim.h"
@@ -171,13 +174,20 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
171174
auto sv = ss.Create(largest_nq);
172175
auto scratch = ss.Create(largest_nq);
173176

174-
// time since epoch seeds random generator.
175-
unsigned long r_seed =
176-
std::chrono::duration_cast<std::chrono::milliseconds>(
177-
std::chrono::system_clock::now().time_since_epoch())
178-
.count();
179-
std::mt19937 gen(r_seed);
180-
std::uniform_int_distribution<> distrib(1, 1 << 30);
177+
tensorflow::GuardedPhiloxRandom random_gen;
178+
int max_psum_length = 1;
179+
int max_n_shots = 1;
180+
for (int i = 0; i < pauli_sums.size(); i++) {
181+
for (int j = 0; j < pauli_sums[i].size(); j++) {
182+
max_psum_length =
183+
std::max(max_psum_length, pauli_sums[i][j].terms().size());
184+
max_n_shots = std::max(max_n_shots, num_samples[i][j]);
185+
}
186+
}
187+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
188+
auto local_gen = random_gen.ReserveSamples128(
189+
ncircuits.size() * (1 + max_psum_length) * max_n_shots);
190+
tensorflow::random::SimplePhilox rand_source(&local_gen);
181191

182192
// Simulate programs one by one. Parallelizing over state vectors
183193
// we no longer parallelize over circuits. Each time we encounter a
@@ -210,7 +220,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
210220
while (1) {
211221
ss.SetStateZero(sv);
212222

213-
QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
223+
QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
214224
scratch, sv, unused_stats);
215225

216226
// Use this trajectory as a source for all expectation calculations.
@@ -221,7 +231,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
221231
float exp_v = 0.0;
222232
OP_REQUIRES_OK(context, ComputeSampledExpectationQsim(
223233
pauli_sums[i][j], sim, ss, sv, scratch, 1,
224-
gen, &exp_v));
234+
rand_source, &exp_v));
225235
rolling_sums[j] += static_cast<double>(exp_v);
226236
run_samples[j]++;
227237
}
@@ -272,10 +282,17 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
272282

273283
output_tensor->setZero();
274284

275-
unsigned long r_seed =
276-
std::chrono::duration_cast<std::chrono::milliseconds>(
277-
std::chrono::system_clock::now().time_since_epoch())
278-
.count();
285+
tensorflow::GuardedPhiloxRandom random_gen;
286+
int max_psum_length = 1;
287+
int max_n_shots = 1;
288+
for (int i = 0; i < pauli_sums.size(); i++) {
289+
for (int j = 0; j < pauli_sums[i].size(); j++) {
290+
max_psum_length =
291+
std::max(max_psum_length, pauli_sums[i][j].terms().size());
292+
max_n_shots = std::max(max_n_shots, num_samples[i][j]);
293+
}
294+
}
295+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
279296

280297
Status compute_status = Status::OK();
281298
auto c_lock = tensorflow::mutex();
@@ -288,9 +305,10 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
288305
auto sv = ss.Create(largest_nq);
289306
auto scratch = ss.Create(largest_nq);
290307

291-
// time since epoch seeds random generator.
292-
std::mt19937 gen(r_seed + start);
293-
std::uniform_int_distribution<> distrib(1, 1 << 30);
308+
int num_rand = ncircuits.size() * (1 + max_psum_length) * max_n_shots;
309+
num_rand = (num_rand + num_threads) / num_threads + 1;
310+
auto local_gen = random_gen.ReserveSamples128(num_rand);
311+
tensorflow::random::SimplePhilox rand_source(&local_gen);
294312

295313
for (int i = 0; i < ncircuits.size(); i++) {
296314
int nq = num_qubits[i];
@@ -321,8 +339,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
321339
while (1) {
322340
ss.SetStateZero(sv);
323341

324-
QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
325-
scratch, sv, unused_stats);
342+
QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss,
343+
sim, scratch, sv, unused_stats);
326344

327345
// Compute expectations across all ops using this trajectory.
328346
for (int j = 0; j < pauli_sums[i].size(); j++) {
@@ -334,7 +352,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
334352
NESTED_FN_STATUS_SYNC(
335353
compute_status,
336354
ComputeSampledExpectationQsim(pauli_sums[i][j], sim, ss, sv,
337-
scratch, 1, gen, &exp_v),
355+
scratch, 1, rand_source, &exp_v),
338356
c_lock);
339357
rolling_sums[j] += static_cast<double>(exp_v);
340358
run_samples[j]++;

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

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ limitations under the License.
1515

1616
#include <stdlib.h>
1717

18-
#include <random>
1918
#include <string>
2019

2120
#include "../qsim/lib/channel.h"
@@ -36,6 +35,9 @@ limitations under the License.
3635
#include "tensorflow/core/lib/core/error_codes.pb.h"
3736
#include "tensorflow/core/lib/core/status.h"
3837
#include "tensorflow/core/lib/core/threadpool.h"
38+
#include "tensorflow/core/lib/random/random.h"
39+
#include "tensorflow/core/lib/random/simple_philox.h"
40+
#include "tensorflow/core/util/guarded_philox_random.h"
3941
#include "tensorflow_quantum/core/ops/parse_context.h"
4042
#include "tensorflow_quantum/core/src/circuit_parser_qsim.h"
4143
#include "tensorflow_quantum/core/src/util_qsim.h"
@@ -149,12 +151,11 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
149151
auto sv = ss.Create(largest_nq);
150152
auto scratch = ss.Create(largest_nq);
151153

152-
unsigned long r_seed =
153-
std::chrono::duration_cast<std::chrono::milliseconds>(
154-
std::chrono::system_clock::now().time_since_epoch())
155-
.count();
156-
std::mt19937 gen(r_seed);
157-
std::uniform_int_distribution<> distrib(1, 1 << 30);
154+
tensorflow::GuardedPhiloxRandom random_gen;
155+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
156+
auto local_gen =
157+
random_gen.ReserveSamples32(2 * num_samples * ncircuits.size() + 2);
158+
tensorflow::random::SimplePhilox rand_source(&local_gen);
158159

159160
// Simulate programs one by one. Parallelizing over state vectors
160161
// we no longer parallelize over circuits. Each time we encounter a
@@ -180,7 +181,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
180181
for (int j = 0; j < num_samples; j++) {
181182
ss.SetStateZero(sv);
182183

183-
QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
184+
QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
184185
scratch, sv, gathered_samples);
185186
uint64_t q_ind = 0;
186187
uint64_t mask = 1;
@@ -236,10 +237,8 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
236237
}
237238
}
238239

239-
unsigned long r_seed =
240-
std::chrono::duration_cast<std::chrono::milliseconds>(
241-
std::chrono::system_clock::now().time_since_epoch())
242-
.count();
240+
tensorflow::GuardedPhiloxRandom random_gen;
241+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
243242

244243
auto DoWork = [&](int start, int end) {
245244
// Begin simulation.
@@ -250,8 +249,11 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
250249
auto sv = ss.Create(largest_nq);
251250
auto scratch = ss.Create(largest_nq);
252251

253-
std::mt19937 gen(r_seed + start);
254-
std::uniform_int_distribution<> distrib(1, 1 << 30);
252+
int needed_random =
253+
4 * (num_samples * ncircuits.size() + num_threads) / num_threads;
254+
needed_random += 4;
255+
auto local_gen = random_gen.ReserveSamples32(needed_random);
256+
tensorflow::random::SimplePhilox rand_source(&local_gen);
255257

256258
for (int i = 0; i < ncircuits.size(); i++) {
257259
int nq = num_qubits[i];
@@ -277,8 +279,8 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
277279

278280
while (1) {
279281
ss.SetStateZero(sv);
280-
QTSimulator::RunOnce(param, ncircuits[i], distrib(gen), ss, sim,
281-
scratch, sv, gathered_samples);
282+
QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss,
283+
sim, scratch, sv, gathered_samples);
282284

283285
uint64_t q_ind = 0;
284286
uint64_t mask = 1;

‎tensorflow_quantum/core/ops/tfq_simulate_sampled_expectation_op.cc‎

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ limitations under the License.
1414
==============================================================================*/
1515

1616
#include <memory>
17-
#include <random>
1817
#include <vector>
1918

2019
#include "../qsim/lib/circuit.h"
@@ -29,7 +28,10 @@ limitations under the License.
2928
#include "tensorflow/core/lib/core/error_codes.pb.h"
3029
#include "tensorflow/core/lib/core/status.h"
3130
#include "tensorflow/core/lib/core/threadpool.h"
31+
#include "tensorflow/core/lib/random/random.h"
32+
#include "tensorflow/core/lib/random/simple_philox.h"
3233
#include "tensorflow/core/platform/mutex.h"
34+
#include "tensorflow/core/util/guarded_philox_random.h"
3335
#include "tensorflow_quantum/core/ops/parse_context.h"
3436
#include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
3537
#include "tensorflow_quantum/core/src/util_qsim.h"
@@ -158,11 +160,17 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
158160
auto sv = ss.Create(largest_nq);
159161
auto scratch = ss.Create(largest_nq);
160162

161-
unsigned long r_seed =
162-
std::chrono::duration_cast<std::chrono::milliseconds>(
163-
std::chrono::system_clock::now().time_since_epoch())
164-
.count();
165-
std::mt19937 gen(r_seed);
163+
tensorflow::GuardedPhiloxRandom random_gen;
164+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
165+
int largest_sum = -1;
166+
for (const auto& sums : pauli_sums) {
167+
for (const auto& sum : sums) {
168+
largest_sum = std::max(largest_sum, sum.terms().size());
169+
}
170+
}
171+
auto local_gen = random_gen.ReserveSamples32(
172+
largest_sum * pauli_sums[0].size() * fused_circuits.size() + 1);
173+
tensorflow::random::SimplePhilox rand_source(&local_gen);
166174

167175
// Simulate programs one by one. Parallelizing over state vectors
168176
// we no longer parallelize over circuits. Each time we encounter a
@@ -192,7 +200,7 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
192200
float exp_v = 0.0;
193201
OP_REQUIRES_OK(context, ComputeSampledExpectationQsim(
194202
pauli_sums[i][j], sim, ss, sv, scratch,
195-
num_samples[i][j], gen, &exp_v));
203+
num_samples[i][j], rand_source, &exp_v));
196204
(*output_tensor)(i, j) = exp_v;
197205
}
198206
}
@@ -211,10 +219,17 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
211219

212220
const int output_dim_op_size = output_tensor->dimension(1);
213221

214-
unsigned long r_seed =
215-
std::chrono::duration_cast<std::chrono::milliseconds>(
216-
std::chrono::system_clock::now().time_since_epoch())
217-
.count();
222+
tensorflow::GuardedPhiloxRandom random_gen;
223+
random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
224+
int largest_sum = -1;
225+
for (const auto& sums : pauli_sums) {
226+
for (const auto& sum : sums) {
227+
largest_sum = std::max(largest_sum, sum.terms().size());
228+
}
229+
}
230+
const int num_threads = context->device()
231+
->tensorflow_cpu_worker_threads()
232+
->workers->NumThreads();
218233

219234
Status compute_status = Status::OK();
220235
auto c_lock = tensorflow::mutex();
@@ -229,7 +244,11 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
229244
auto sv = ss.Create(largest_nq);
230245
auto scratch = ss.Create(largest_nq);
231246

232-
std::mt19937 gen(r_seed + start);
247+
int n_random = largest_sum * output_dim_op_size * fused_circuits.size();
248+
n_random /= num_threads;
249+
n_random += 1;
250+
auto local_gen = random_gen.ReserveSamples32(n_random);
251+
tensorflow::random::SimplePhilox rand_source(&local_gen);
233252

234253
for (int i = start; i < end; i++) {
235254
cur_batch_index = i / output_dim_op_size;
@@ -264,7 +283,8 @@ class TfqSimulateSampledExpectationOp : public tensorflow::OpKernel {
264283
compute_status,
265284
ComputeSampledExpectationQsim(
266285
pauli_sums[cur_batch_index][cur_op_index], sim, ss, sv, scratch,
267-
num_samples[cur_batch_index][cur_op_index], gen, &exp_v),
286+
num_samples[cur_batch_index][cur_op_index], rand_source,
287+
&exp_v),
268288
c_lock);
269289

270290
(*output_tensor)(cur_batch_index, cur_op_index) = exp_v;

0 commit comments

Comments
 (0)