| 64 | } |
| 65 | |
| 66 | Status Main(const MainFlags& flags) { |
| 67 | // Process config. |
| 68 | tf2xla::Config config; |
| 69 | if (flags.config.empty()) { |
| 70 | return errors::InvalidArgument("Must specify --config"); |
| 71 | } |
| 72 | TF_RETURN_IF_ERROR(ReadProtoFile(flags.config, &config)); |
| 73 | TF_RETURN_IF_ERROR(ValidateConfig(config)); |
| 74 | if (flags.dump_fetch_nodes) { |
| 75 | std::set<string> nodes; |
| 76 | for (const tf2xla::Fetch& fetch : config.fetch()) { |
| 77 | nodes.insert(fetch.id().node_name()); |
| 78 | } |
| 79 | std::cout << absl::StrJoin(nodes, ","); |
| 80 | return Status::OK(); |
| 81 | } |
| 82 | |
| 83 | // Read and initialize the graph. |
| 84 | if (flags.graph.empty()) { |
| 85 | return errors::InvalidArgument("Must specify --graph"); |
| 86 | } |
| 87 | GraphDef graph_def; |
| 88 | TF_RETURN_IF_ERROR(ReadProtoFile(flags.graph, &graph_def)); |
| 89 | CompileResult compile_result; |
| 90 | TF_RETURN_IF_ERROR(CompileGraph(graph_def, config, flags, &compile_result)); |
| 91 | |
| 92 | // Write output files. |
| 93 | Env* env = Env::Default(); |
| 94 | const std::vector<char>& obj = compile_result.aot->object_file_data(); |
| 95 | TF_RETURN_IF_ERROR( |
| 96 | WriteStringToFile(env, flags.out_function_object, |
| 97 | absl::string_view(obj.data(), obj.size()))); |
| 98 | CodegenOpts codegen_opts; |
| 99 | codegen_opts.gen_name_to_index = flags.gen_name_to_index; |
| 100 | codegen_opts.gen_program_shape = flags.gen_program_shape; |
| 101 | codegen_opts.target_triple = flags.target_triple; |
| 102 | if (flags.cpp_class.empty()) { |
| 103 | return errors::InvalidArgument("Must specify --cpp_class"); |
| 104 | } |
| 105 | codegen_opts.gen_hlo_profile_printer_data = |
| 106 | xla::GetDebugOptionsFromFlags().xla_hlo_profile(); |
| 107 | TF_RETURN_IF_ERROR(ParseCppClass(flags.cpp_class, &codegen_opts.class_name, |
| 108 | &codegen_opts.namespaces)); |
| 109 | |
| 110 | MetadataResult metadata_result; |
| 111 | TF_RETURN_IF_ERROR( |
| 112 | GenerateMetadata(codegen_opts, compile_result, &metadata_result)); |
| 113 | TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_metadata_object, |
| 114 | metadata_result.object_file_data)); |
| 115 | string header; |
| 116 | TF_RETURN_IF_ERROR(GenerateHeader(codegen_opts, config, compile_result, |
| 117 | metadata_result, &header)); |
| 118 | TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_header, header)); |
| 119 | return Status::OK(); |
| 120 | } |
| 121 | |
| 122 | } // end namespace tfcompile |
| 123 | } // end namespace tensorflow |
no test coverage detected