Compiles the XLA computation into executable code.
| 45 | |
| 46 | // Compiles the XLA computation into executable code. |
| 47 | Status CompileXla(xla::CompileOnlyClient* client, |
| 48 | const xla::XlaComputation& computation, |
| 49 | const xla::cpu::CpuAotCompilationOptions& aot_opts, |
| 50 | CompileResult* compile_result) { |
| 51 | // Retrieves arg and result layouts from the computation. |
| 52 | // TODO(toddw): Should we let the user choose the major/minor ordering? |
| 53 | xla::StatusOr<std::unique_ptr<xla::ProgramShape>> pshape_or = |
| 54 | client->GetComputationShape(computation); |
| 55 | if (!pshape_or.ok()) { |
| 56 | return errors::Unknown("Couldn't get XLA program shape: ", |
| 57 | pshape_or.status().error_message()); |
| 58 | } |
| 59 | compile_result->program_shape = pshape_or.ValueOrDie()->ToProto(); |
| 60 | xla::ProgramShapeProto* pshape = &compile_result->program_shape; |
| 61 | |
| 62 | // AotXlaComputationInstance::argument_layouts is a vector of Shape |
| 63 | // pointers. Accumulate the Shape objects themselves in a separate vector |
| 64 | // while building the vector of pointers. |
| 65 | std::vector<const xla::Shape*> arg_layout_ptrs(pshape->parameters_size()); |
| 66 | std::vector<xla::Shape> arg_layouts(pshape->parameters_size()); |
| 67 | for (int i = 0; i < pshape->parameters_size(); ++i) { |
| 68 | arg_layouts[i] = xla::Shape(*pshape->mutable_parameters(i)); |
| 69 | arg_layout_ptrs[i] = &arg_layouts[i]; |
| 70 | } |
| 71 | xla::CompileOnlyClient::AotXlaComputationInstance instance; |
| 72 | instance.computation = &computation; |
| 73 | instance.argument_layouts = std::move(arg_layout_ptrs); |
| 74 | xla::Shape result_shape(pshape->result()); |
| 75 | instance.result_layout = &result_shape; |
| 76 | xla::StatusOr<std::vector<std::unique_ptr<xla::AotCompilationResult>>> |
| 77 | aot_or = client->CompileAheadOfTime({instance}, aot_opts); |
| 78 | if (!aot_or.ok()) { |
| 79 | return errors::Unknown("XLA compilation failed: ", |
| 80 | aot_or.status().error_message()); |
| 81 | } |
| 82 | compile_result->aot = |
| 83 | xla::unique_ptr_static_cast<xla::cpu::CpuAotCompilationResult>( |
| 84 | std::move(aot_or.ValueOrDie().back())); |
| 85 | compile_result->entry_point = aot_opts.entry_point_name(); |
| 86 | compile_result->pointer_size = |
| 87 | xla::CompileOnlyClient::PointerSizeForTriple(aot_opts.triple()); |
| 88 | return Status::OK(); |
| 89 | } |
| 90 | |
| 91 | } // namespace |
| 92 |
no test coverage detected