| 301 | } |
| 302 | |
| 303 | Status XlaCompilationCache::CompileAsynchronous( |
| 304 | Entry* entry, const XlaCompiler::Options& options, |
| 305 | const std::vector<XlaCompiler::Argument>& args, |
| 306 | const string &function_name, |
| 307 | const std::function<Status(XlaCompiler* compiler, |
| 308 | const std::vector<XlaCompiler::Argument>& args, |
| 309 | XlaCompiler::CompilationResult*)>& compile_fn) { |
| 310 | entry->compile_state = CompileState::kCompiling; // still under caller's lock. |
| 311 | { |
| 312 | mutex_lock lock(async_compilation_.async_compilation_mu_); |
| 313 | async_compilation_.nrof_ongoing_compilations++; |
| 314 | } |
| 315 | // don't move the above code into the thread function!!! |
| 316 | |
| 317 | // passing options by value into the lamba increases the refcount on |
| 318 | // options.device_allocator, keeping it alive for the duration of the |
| 319 | // compilation |
| 320 | // passing args by value as well. Doing this here only when an asynchronous |
| 321 | // compilation is performed, as copying many args incurs overhead, |
| 322 | async_compilation_.compiler_threads.Schedule([=] { |
| 323 | Entry tmp; |
| 324 | VLOG(2) << "Starting asynchronous compilation of cluster " |
| 325 | << function_name << '.'; |
| 326 | (void)CompileStrict(&tmp, options, args, function_name, compile_fn); |
| 327 | VLOG(2) << "Finished asynchronous compililation of cluster " |
| 328 | << function_name << '.'; |
| 329 | { |
| 330 | mutex_lock lock(async_compilation_.async_compilation_mu_); |
| 331 | async_compilation_.nrof_ongoing_compilations--; |
| 332 | } |
| 333 | { // populate original entry with compilation result |
| 334 | mutex_lock entry_lock(entry->mu); |
| 335 | entry->compilation_result = tmp.compilation_result; |
| 336 | entry->compile_state = tmp.compile_state; |
| 337 | entry->compilation_status = tmp.compilation_status; |
| 338 | entry->executable = std::move(tmp.executable); |
| 339 | } |
| 340 | } |
| 341 | ); |
| 342 | return Status::OK(); |
| 343 | } |
| 344 | |
| 345 | Status XlaCompilationCache::CompileImpl( |
| 346 | const XlaCompiler::Options& options, const NameAttrList& function, |