| 112 | } |
| 113 | |
| 114 | Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr<Graph> graph, |
| 115 | XlaCompilationDevice* device, FunctionLibraryRuntime* flib, |
| 116 | int64 step_id) { |
| 117 | // Resource cleanup is a bit messy. XlaContext is a ref-countd resource; the |
| 118 | // resource manager takes ownership via Create, and unrefs via Cleanup. We |
| 119 | // explicitly add a reference to ensure the refcount at entry is maintained at |
| 120 | // all exit points; Create and Cleanup are always called in this function. |
| 121 | // |
| 122 | // The Executor requires us to use ScopedStepContainer. We wrap it in a |
| 123 | // unique_ptr so we can capture the cleanup status in the end. |
| 124 | xla_context->Ref(); |
| 125 | Status status; |
| 126 | auto step_container = absl::make_unique<ScopedStepContainer>( |
| 127 | step_id, [&status, device](const string& name) { |
| 128 | status = device->resource_manager()->Cleanup(name); |
| 129 | }); |
| 130 | TF_RETURN_IF_ERROR(step_container->Create(device->resource_manager(), |
| 131 | XlaContext::kXlaContextResourceName, |
| 132 | xla_context)); |
| 133 | GraphCompiler graph_compiler(device, graph.get(), flib, step_container.get()); |
| 134 | TF_RETURN_IF_ERROR(graph_compiler.Compile()); |
| 135 | // Explicitly clean up the step container, to capture the cleanup status. |
| 136 | step_container.reset(); |
| 137 | return Status::OK(); |
| 138 | } |
| 139 | |
| 140 | // Builds the XLA computation. |
| 141 | // - `args` is the list of input arguments |