| 366 | } // namespace |
| 367 | |
| 368 | Status GenerateHeader(const CodegenOpts& opts, const tf2xla::Config& config, |
| 369 | const CompileResult& compile_result, |
| 370 | const MetadataResult& metadata_result, string* header) { |
| 371 | TF_RETURN_IF_ERROR(ValidateConfig(config)); |
| 372 | TF_RETURN_IF_ERROR(ValidateFeedFetchCppNames(config)); |
| 373 | const int64 result_index = compile_result.aot->result_buffer_index(); |
| 374 | const std::vector<BufferInfo>& buffer_infos = |
| 375 | compile_result.aot->buffer_infos(); |
| 376 | const std::vector<int32> arg_index_table = |
| 377 | ::xla::cpu::CreateArgIndexTableFromBufferInfos(buffer_infos); |
| 378 | std::vector<string> buffer_infos_as_strings = |
| 379 | BufferInfosToCppExpression(buffer_infos); |
| 380 | if (result_index < 0 || result_index >= buffer_infos.size()) { |
| 381 | return errors::InvalidArgument("result index: ", result_index, |
| 382 | " is outside the range of temp sizes: [0,", |
| 383 | buffer_infos.size(), ")"); |
| 384 | } |
| 385 | |
| 386 | // Compute sizes and generate methods. |
| 387 | std::vector<BufferInfo> buffer_infos_for_args = |
| 388 | ExtractEntryParamBufferInfos(buffer_infos); |
| 389 | std::vector<BufferInfo> buffer_infos_for_temps = |
| 390 | ExtractTempBufferInfos(buffer_infos); |
| 391 | const xla::ProgramShapeProto& ps = compile_result.program_shape; |
| 392 | string methods_arg, methods_result, methods_variable; |
| 393 | TF_RETURN_IF_ERROR(GenArgMethods(config, ps, compile_result, &methods_arg)); |
| 394 | TF_RETURN_IF_ERROR(GenResultMethods(config, ps, &methods_result)); |
| 395 | TF_RETURN_IF_ERROR(GenVariableMethods(config, ps, &methods_variable)); |
| 396 | const size_t arg_bytes_aligned = |
| 397 | xla::cpu_function_runtime::AlignedBufferBytes( |
| 398 | buffer_infos_for_args.data(), buffer_infos_for_args.size(), |
| 399 | /*allocate_entry_params=*/true); |
| 400 | const size_t arg_bytes_total = TotalBufferBytes(buffer_infos_for_args); |
| 401 | const size_t temp_bytes_aligned = |
| 402 | xla::cpu_function_runtime::AlignedBufferBytes( |
| 403 | buffer_infos_for_temps.data(), buffer_infos_for_temps.size(), |
| 404 | /*allocate_entry_params=*/true); |
| 405 | const size_t temp_bytes_total = TotalBufferBytes(buffer_infos_for_temps); |
| 406 | |
| 407 | // Create rewrite strings for namespace start and end. |
| 408 | string ns_start; |
| 409 | for (const string& n : opts.namespaces) { |
| 410 | ns_start += absl::StrCat("namespace ", n, " {\n"); |
| 411 | } |
| 412 | ns_start += "\n"; |
| 413 | string ns_end("\n"); |
| 414 | for (int i = opts.namespaces.size() - 1; i >= 0; --i) { |
| 415 | const string& n = opts.namespaces[i]; |
| 416 | ns_end += absl::StrCat("} // end namespace ", n, "\n"); |
| 417 | } |
| 418 | |
| 419 | // Generate metadata. |
| 420 | const string arg_names_code = |
| 421 | GenNameToIndexCode(config.feed(), opts.gen_name_to_index); |
| 422 | const string result_names_code = |
| 423 | GenNameToIndexCode(config.fetch(), opts.gen_name_to_index); |
| 424 | const string include_xla_data_proto = |
| 425 | opts.gen_program_shape |