static*/
| 81 | } // namespace |
| 82 | |
| 83 | /*static*/ xla::StatusOr<std::unique_ptr<XlaJitCompiledCpuFunction>> |
| 84 | XlaJitCompiledCpuFunction::Compile( |
| 85 | const GraphDef& graph_def, const tf2xla::Config& config, |
| 86 | const xla::ExecutableBuildOptions& build_options) { |
| 87 | // Convert the graph_def into an xla::XlaComputation. |
| 88 | TF_ASSIGN_OR_RETURN(se::Platform * platform, |
| 89 | xla::PlatformUtil::GetPlatform(kHostPlatform)); |
| 90 | TF_ASSIGN_OR_RETURN(xla::LocalClient * client, |
| 91 | xla::ClientLibrary::GetOrCreateLocalClient(platform)); |
| 92 | xla::XlaComputation computation; |
| 93 | TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToXla(graph_def, config, client, |
| 94 | &computation)); |
| 95 | |
| 96 | // Get and verify the program shape. |
| 97 | TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::ProgramShape> program_shape, |
| 98 | client->GetComputationShape(computation)); |
| 99 | if (program_shape->result().element_type() != xla::TUPLE) { |
| 100 | // The XlaCompiler we use to build the xla computation always generates a |
| 101 | // tuple result, and XlaCompiledCpuFunction relies on this for simpler |
| 102 | // calling semantics. |
| 103 | return errors::Internal( |
| 104 | "XlaJitCompiledCpuFunction requires the XLA result to be a tuple"); |
| 105 | } |
| 106 | // The parameter names are currently meaningless, and redundant with the rest |
| 107 | // of our metadata, so clear them out to avoid confusion and save space. |
| 108 | program_shape->clear_parameter_names(); |
| 109 | |
| 110 | // Compute arg shapes, needed to compile the executable. |
| 111 | std::vector<const xla::Shape*> arg_shapes; |
| 112 | arg_shapes.reserve(program_shape->parameters_size()); |
| 113 | for (int i = 0; i < program_shape->parameters_size(); ++i) { |
| 114 | arg_shapes.push_back(&program_shape->parameters(i)); |
| 115 | } |
| 116 | |
| 117 | // Compile the executable. The static_cast to the CpuExecutable subclass is |
| 118 | // necessary since the raw function and buffer assignments are only available |
| 119 | // there. |
| 120 | TF_ASSIGN_OR_RETURN(auto executables, |
| 121 | client->Compile(computation, arg_shapes, build_options)); |
| 122 | TF_RET_CHECK(executables.size() == 1); |
| 123 | std::unique_ptr<xla::LocalExecutable> executable = std::move(executables[0]); |
| 124 | const xla::cpu::CpuExecutable* cpu_executable = |
| 125 | static_cast<xla::cpu::CpuExecutable*>(executable->executable()); |
| 126 | XlaCompiledCpuFunction::RawFunction raw_function = |
| 127 | cpu_executable->compute_function(); |
| 128 | const xla::BufferAssignment& buffer_assignment = |
| 129 | cpu_executable->buffer_assignment(); |
| 130 | |
| 131 | // Compute buffer infos and the result index, needed to run the raw function. |
| 132 | std::vector<xla::cpu_function_runtime::BufferInfo> buffer_infos = |
| 133 | xla::cpu::CreateBufferInfosFromBufferAssignment(buffer_assignment); |
| 134 | std::vector<int32> arg_index_table = |
| 135 | xla::cpu::CreateArgIndexTableFromBufferInfos(buffer_infos); |
| 136 | TF_ASSIGN_OR_RETURN(size_t result_index, |
| 137 | ComputeResultIndex(buffer_assignment)); |
| 138 | |
| 139 | std::unique_ptr<XlaJitCompiledCpuFunction> jit_unique_ptr( |
| 140 | new XlaJitCompiledCpuFunction); |
nothing calls this directly
no test coverage detected