| 393 | } |
| 394 | |
| 395 | Status ShouldCompileWithXLA(const EagerOperation* op, const EagerContext* ctx, |
| 396 | bool* compile_with_xla) { |
| 397 | if (!op->is_function()) { |
| 398 | *compile_with_xla = false; |
| 399 | return Status::OK(); |
| 400 | } |
| 401 | |
| 402 | // Does node have an explicit request to compile or not? |
| 403 | Status status = op->Attrs().Get(kXlaCompileAttr, compile_with_xla); |
| 404 | if (status.ok()) { |
| 405 | VLOG(2) << "Caller explicitly requested " |
| 406 | << (*compile_with_xla ? "" : "not ") |
| 407 | << "to compile with XLA: " << op->DebugString(); |
| 408 | return Status::OK(); |
| 409 | } |
| 410 | |
| 411 | // Does FunctionDef have an explicit request to compile or not? |
| 412 | const FunctionDef* function_def = |
| 413 | ctx->pflr()->GetFunctionLibraryDefinition()->Find(op->Name()); |
| 414 | if (function_def == nullptr) { |
| 415 | return errors::NotFound("Failed to find function '", op->Name(), "'"); |
| 416 | } |
| 417 | |
| 418 | status = GetNodeAttr(AttrSlice(&function_def->attr()), kXlaCompileAttr, |
| 419 | compile_with_xla); |
| 420 | if (status.ok()) { |
| 421 | VLOG(2) << "Function definition explicitly specifies " |
| 422 | << (*compile_with_xla ? "" : "not ") << "to compile with XLA"; |
| 423 | return Status::OK(); |
| 424 | } |
| 425 | |
| 426 | // No explicit requests. Compile for XLA devices by default. |
| 427 | if (op->GetDeviceName().type == "TPU" || |
| 428 | op->GetDeviceName().type == "XLA_GPU" || |
| 429 | op->GetDeviceName().type == "XLA_CPU") { |
| 430 | VLOG(2) << "Compiling " << op->Name() |
| 431 | << " with XLA because it is running on an XLA device " |
| 432 | << op->GetDeviceName().type; |
| 433 | *compile_with_xla = true; |
| 434 | } else { |
| 435 | *compile_with_xla = false; |
| 436 | } |
| 437 | |
| 438 | return Status::OK(); |
| 439 | } |
| 440 | |
| 441 | // There are a lot of references to devices in this function and around. |
| 442 | // Here is what they mean: |
no test coverage detected