| 426 | } |
| 427 | |
| 428 | int RealMain(absl::Span<char* const> args, const Options& opts) { |
| 429 | LocalClient* client = ClientLibrary::LocalClientOrDie(); |
| 430 | int exit_status = EXIT_SUCCESS; |
| 431 | |
| 432 | std::vector<HloSnapshot> snapshots; |
| 433 | for (char* arg : args) { |
| 434 | StatusOr<std::vector<HloSnapshot>> maybe_snapshot = |
| 435 | ParseInputFile(arg, opts); |
| 436 | if (maybe_snapshot.ok()) { |
| 437 | auto new_snapshots = std::move(maybe_snapshot).ValueOrDie(); |
| 438 | snapshots.insert(snapshots.end(), |
| 439 | std::make_move_iterator(new_snapshots.begin()), |
| 440 | std::make_move_iterator(new_snapshots.end())); |
| 441 | } else { |
| 442 | LOG(ERROR) << maybe_snapshot.status(); |
| 443 | } |
| 444 | } |
| 445 | |
| 446 | // Compile all the modules in parallel. |
| 447 | LOG(INFO) << "Compiling " << snapshots.size() << " modules in parallel."; |
| 448 | std::vector<StatusOr<std::unique_ptr<LocalExecutable>>> executables; |
| 449 | { |
| 450 | constexpr size_t kThreadLimits = 100; |
| 451 | // ThreadPool CHECK-fails if we give it 0 threads. |
| 452 | tensorflow::thread::ThreadPool thread_pool( |
| 453 | tensorflow::Env::Default(), tensorflow::ThreadOptions(), |
| 454 | "compile_modules", |
| 455 | std::min<size_t>(std::max(kThreadLimits, snapshots.size()), 1), |
| 456 | /*low_latency_hint=*/false); |
| 457 | executables.resize(snapshots.size()); |
| 458 | for (int64 i = 0; i < snapshots.size(); ++i) { |
| 459 | thread_pool.Schedule([&snapshots, &executables, client, i] { |
| 460 | executables[i] = CompileExecutable(snapshots[i], client); |
| 461 | }); |
| 462 | } |
| 463 | } |
| 464 | LOG(INFO) << "Done compiling; now running the modules."; |
| 465 | |
| 466 | for (int64 i = 0; i < executables.size(); ++i) { |
| 467 | if (!executables[i].ok()) { |
| 468 | LOG(ERROR) << "Compilation failed: " << executables[i].status() << ": " |
| 469 | << snapshots[i].ShortDebugString(); |
| 470 | exit_status = EXIT_FAILURE; |
| 471 | continue; |
| 472 | } |
| 473 | |
| 474 | if (opts.compile_only) { |
| 475 | continue; |
| 476 | } |
| 477 | |
| 478 | LocalExecutable* executable = executables[i].ValueOrDie().get(); |
| 479 | LOG(ERROR) << "Running iteration " << i; |
| 480 | StatusOr<Literal> result_status = |
| 481 | ReplayComputation(snapshots[i], executable, client, opts); |
| 482 | LOG(ERROR) << "iteration complete."; |
| 483 | if (!result_status.ok()) { |
| 484 | fprintf(stderr, "%s: error: %s\n", args[i], |
| 485 | result_status.status().ToString().c_str()); |
no test coverage detected