| 267 | } |
| 268 | |
| 269 | static Status CompileToLocalExecutable( |
| 270 | OpKernelContext* ctx, const NameAttrList& function, |
| 271 | const XlaPlatformInfo& platform_info, absl::Span<const int> resources, |
| 272 | absl::Span<const int> constants, bool async, bool lazy, |
| 273 | xla::LocalClient** client, std::map<int, OptionalTensor>* variables, |
| 274 | const XlaCompiler::CompilationResult** kernel, |
| 275 | xla::LocalExecutable** executable) { |
| 276 | // We store information about the JIT-compiled XLA computation |
| 277 | // in the ResourceMgr. |
| 278 | ResourceMgr* rm = ctx->resource_manager(); |
| 279 | if (!rm) { |
| 280 | return errors::Internal("No resource manager."); |
| 281 | } |
| 282 | |
| 283 | XlaCompilationCache* cache; |
| 284 | TF_RETURN_IF_ERROR(rm->LookupOrCreate<XlaCompilationCache>( |
| 285 | rm->default_container(), "xla_cache", &cache, |
| 286 | [&](XlaCompilationCache** cache) { |
| 287 | return BuildCompilationCache(ctx, platform_info, cache); |
| 288 | })); |
| 289 | // Hold the reference to the JIT during evaluation. (We could probably |
| 290 | // free it sooner because the ResourceMgr will retain a reference, but |
| 291 | // this is more obviously correct.) |
| 292 | core::ScopedUnref cache_ref(cache); |
| 293 | |
| 294 | TF_RETURN_IF_ERROR(SnapshotResourceVariables(ctx, resources, variables)); |
| 295 | *client = static_cast<xla::LocalClient*>(cache->client()); |
| 296 | |
| 297 | XlaCompiler::Options options; |
| 298 | options.client = *client; |
| 299 | if (ctx->op_device_context() != nullptr) { |
| 300 | options.device_ordinal = |
| 301 | ctx->op_device_context()->stream()->parent()->device_ordinal(); |
| 302 | } |
| 303 | options.device_type = cache->device_type(); |
| 304 | options.flib_def = ctx->function_library()->GetFunctionLibraryDefinition(); |
| 305 | options.graph_def_version = ctx->function_library()->graph_def_version(); |
| 306 | options.allow_cpu_custom_calls = |
| 307 | (platform_info.platform_id() == se::host::kHostPlatformId); |
| 308 | options.device_allocator = GetAllocator(ctx, platform_info); |
| 309 | if (platform_info.xla_device_metadata()) { |
| 310 | options.shape_representation_fn = |
| 311 | platform_info.xla_device_metadata()->shape_representation_fn(); |
| 312 | } |
| 313 | |
| 314 | std::map<int, Tensor> constant_args; |
| 315 | for (int i : constants) { |
| 316 | constant_args.insert({i, ctx->input(i)}); |
| 317 | } |
| 318 | XlaCompiler::CompileOptions compile_options; |
| 319 | compile_options.is_entry_computation = true; |
| 320 | // If we resolve constants we never emit them on the device, meaning that if |
| 321 | // they are needed by a following computation the host has to transfer |
| 322 | // them. Not resolving constants is expected to be faster than resolving |
| 323 | // constants. |
| 324 | compile_options.resolve_compile_time_constants = true; |
| 325 | // Optimization: where possible, have the computation return a naked array |
| 326 | // rather than a one-element tuple. |
no test coverage detected