| 91 | } // namespace |
| 92 | |
| 93 | Status CompileGraph(const GraphDef& graph_def, const tf2xla::Config& config, |
| 94 | const MainFlags& flags, CompileResult* compile_result) { |
| 95 | // Converts the graph into an XLA computation, and compiles the |
| 96 | // computation. |
| 97 | // TODO(toddw): Should we let the user pick the XLA cpu vs. gpu client? |
| 98 | se::Platform* cpu_platform = |
| 99 | se::MultiPlatformManager::PlatformWithName("Host").ValueOrDie(); |
| 100 | xla::CompileOnlyClient* client = |
| 101 | xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform) |
| 102 | .ValueOrDie(); |
| 103 | xla::XlaComputation computation; |
| 104 | TF_RETURN_IF_ERROR( |
| 105 | ConvertGraphDefToXla(graph_def, config, client, &computation)); |
| 106 | if (!flags.out_session_module.empty()) { |
| 107 | TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::HloSnapshot> module, |
| 108 | computation.Snapshot()); |
| 109 | // Serialize the HloSnapshot deterministically so that all the outputs of a |
| 110 | // tf_library genrule are deterministic. |
| 111 | const size_t size = module->ByteSizeLong(); |
| 112 | auto serialized = absl::make_unique<char[]>(size); |
| 113 | TF_RET_CHECK( |
| 114 | SerializeToBufferDeterministic(*module, serialized.get(), size)); |
| 115 | TF_RETURN_IF_ERROR( |
| 116 | WriteStringToFile(Env::Default(), flags.out_session_module, |
| 117 | absl::string_view(serialized.get(), size))); |
| 118 | } |
| 119 | xla::cpu::CpuAotCompilationOptions aot_opts( |
| 120 | flags.target_triple, flags.target_cpu, flags.target_features, |
| 121 | flags.entry_point, |
| 122 | xla::cpu::CpuAotCompilationOptions::RelocationModel::BigPic); |
| 123 | |
| 124 | return CompileXla(client, computation, aot_opts, compile_result); |
| 125 | } |
| 126 | |
| 127 | } // namespace tfcompile |
| 128 | } // namespace tensorflow |
no test coverage detected