[XLA] Improvements to replay_computation tool.

 * Reduce threshold at which we run fake-data generation on the device from 1gb to 1mb. At the old threshold, I observed cases where we'd spend many seconds, and >50% of our runtime, in logf(), used for computing random numbers. * Don't retrieve or print the result when running with fake data. Presumably this is uninteresting, because garbage in, garbage out. Retrieving this data can take as long as running the whole computation, and printing it can take many times longer. * Add a LOG(INFO) indicating how long execution took. * Add a --num_runs flag. This is particularly important on GPUs, where the first run does autotuning, and so isn't interesting from a performance perspective. PiperOrigin-RevId: 177185636

Author: Justin Lebar

tensorflow/compiler/xla/client/lib/testing.cc

@@ -51,7 +51,7 @@ std::unique_ptr<GlobalData> MakeFakeDataViaDeviceOrDie(const Shape& shape,
 
 std::unique_ptr<GlobalData> MakeFakeDataOrDie(const Shape& shape,
  Client* client) {
- if (ShapeUtil::ByteSizeOf(shape) < (1LL << 30)) {
+ if (ShapeUtil::ByteSizeOf(shape) < (1LL << 20)) {
  StatusOr<std::unique_ptr<Literal>> literal_status = MakeFakeLiteral(shape);
  if (!literal_status.ok()) {
  // If we got an Unimplemented error, fall back to making the fake data via

tensorflow/compiler/xla/tools/replay_computation.cc

@@ -65,8 +65,9 @@ namespace {
 // Similarly, infeeds fake data of shape fake_infeed_shape if it is provided;
 // otherwise, no infeed is performed.
 StatusOr<std::unique_ptr<Literal>> ReplayComputation(
- const SessionModule& module, tensorflow::StringPiece fake_infeed_shape,
- bool use_fake_data, Client* client) {
+ const SessionModule& module, int num_runs,
+ tensorflow::StringPiece fake_infeed_shape, bool use_fake_data,
+ Client* client) {
  TF_ASSIGN_OR_RETURN(Computation computation, client->LoadSnapshot(module));
 
  std::vector<std::unique_ptr<GlobalData>> arguments;
@@ -107,33 +108,58 @@ StatusOr<std::unique_ptr<Literal>> ReplayComputation(
  for (auto& argument : arguments) {
  execute_arguments.push_back(argument.get());
  }
- return client->ExecuteAndTransfer(computation, execute_arguments);
+
+ // Run the computation num_runs times, and return the result from the last
+ // execution.
+ std::unique_ptr<Literal> result;
+ for (int i = 0; i < num_runs; ++i) {
+ ExecutionProfile profile;
+ if (use_fake_data) {
+ // If using fake data, execute the computation but don't bother retrieving
+ // the result -- presumably it's uninteresting, since our data is fake.
+ TF_RETURN_IF_ERROR(client
+ ->Execute(computation, execute_arguments,
+ /*execution_options=*/nullptr, &profile)
+ .status());
+ } else {
+ TF_ASSIGN_OR_RETURN(result, client->ExecuteAndTransfer(
+ computation, execute_arguments,
+ /*execution_options=*/nullptr, &profile));
+ }
+ LOG(INFO) << "Execution took "
+ << static_cast<double>(profile.compute_time_ns()) / 1e9 << "s";
+ }
+
+ return std::move(result);
 }
 
-int RealMain(tensorflow::gtl::ArraySlice<char*> args,
+int RealMain(tensorflow::gtl::ArraySlice<char*> args, int num_runs,
  tensorflow::StringPiece fake_infeed_shape, bool use_fake_data) {
  Client* client = ClientLibrary::LocalClientOrDie();
  tensorflow::Env* env = tensorflow::Env::Default();
  int exit_status = EXIT_SUCCESS;
  for (char* arg : args) {
  SessionModule module;
  TF_CHECK_OK(tensorflow::ReadBinaryProto(env, arg, &module));
- StatusOr<std::unique_ptr<Literal>> result_status =
- ReplayComputation(module, fake_infeed_shape, use_fake_data, client);
+ StatusOr<std::unique_ptr<Literal>> result_status = ReplayComputation(
+ module, num_runs, fake_infeed_shape, use_fake_data, client);
  if (!result_status.ok()) {
  fprintf(stderr, "%s: error: %s\n", arg,
  result_status.status().ToString().c_str());
  exit_status = EXIT_FAILURE;
  continue;
  }
+
  std::unique_ptr<Literal> result = result_status.ConsumeValueOrDie();
- fprintf(stdout, "%s: %s :: %s:%s\n", arg, module.entry().name().c_str(),
- ShapeUtil::HumanString(result->shape()).c_str(),
- result->ToString().c_str());
- if (module.has_result()) {
- fprintf(stdout, "was %s:%s\n",
- ShapeUtil::HumanString(module.result().shape()).c_str(),
- Literal(module.result()).ToString().c_str());
+ if (result != nullptr) {
+ fprintf(stdout, "%s: %s :: %s:%s\n", arg, module.entry().name().c_str(),
+ ShapeUtil::HumanString(result->shape()).c_str(),
+ result->ToString().c_str());
+ if (module.has_result()) {
+ fprintf(stdout, "was %s:%s\n",
+ ShapeUtil::HumanString(module.result().shape()).c_str(),
+ Literal(module.result()).ToString().c_str());
+ }
  }
  }
  return exit_status;
@@ -147,9 +173,12 @@ int main(int argc, char** argv) {
  // Flags
  xla::string fake_infeed_shape;
  bool use_fake_data = false;
+ int num_runs = 1;
  const std::vector<tensorflow::Flag> flag_list = {
  tensorflow::Flag("use_fake_data", &use_fake_data,
  "Replay computation using fake data"),
+ tensorflow::Flag("num_runs", &num_runs,
+ "Number of times to run each computation"),
  tensorflow::Flag("fake_infeed_shape", &fake_infeed_shape,
  "Shape of fake data to construct for (infinite) infeed"),
  };
@@ -162,5 +191,5 @@ int main(int argc, char** argv) {
 
  tensorflow::gtl::ArraySlice<char*> args(argv, argc);
  args.pop_front(); // Pop off the binary name, argv[0]
- return xla::tools::RealMain(args, fake_infeed_shape, use_fake_data);
+ return xla::tools::RealMain(args, num_runs, fake_infeed_shape, use_fake_data);
 }