| 70 | } // namespace |
| 71 | |
| 72 | bool XlaKernelCreator::CanCreateKernel(const FunctionLibraryRuntime& flr, |
| 73 | const NodeDef& node_def) const { |
| 74 | const FunctionDef* function_def = |
| 75 | flr.GetFunctionLibraryDefinition()->Find(node_def.name()); |
| 76 | if (function_def == nullptr) { |
| 77 | // The node def is not calling a function. Individual ops can be |
| 78 | // run directly using on-demand mode, no need to create XlaLaunch |
| 79 | // kernel for them. |
| 80 | return false; |
| 81 | } |
| 82 | |
| 83 | // If kXlaCompileAttr is set on the node_def, use its value. |
| 84 | const auto& it = node_def.attr().find(kXlaCompileAttr); |
| 85 | if (it != node_def.attr().end()) { |
| 86 | return it->second.b(); |
| 87 | } |
| 88 | |
| 89 | // kXlaCompileAttr is not set on node_def, check if it is set on |
| 90 | // FunctionDef. |
| 91 | bool xla_compile = false; |
| 92 | Status status = flr.GetFunctionLibraryDefinition()->GetAttr( |
| 93 | node_def, kXlaCompileAttr, &xla_compile); |
| 94 | if (!status.ok() || !xla_compile) { |
| 95 | if (VLOG_IS_ON(3)) { |
| 96 | if (!status.ok()) { |
| 97 | VLOG(3) << "No " << kXlaCompileAttr << " attr defined for " |
| 98 | << node_def.op() << ". status=" << status.ToString(); |
| 99 | } else { |
| 100 | VLOG(3) << node_def.op() << " is explicitly marked not to be compiled"; |
| 101 | } |
| 102 | } |
| 103 | return false; |
| 104 | } |
| 105 | return true; |
| 106 | } |
| 107 | |
| 108 | // Given a FunctionLibraryRuntime and a NodeDef calling a function in the |
| 109 | // runtime, returns this function's body in `fbody` as well as the indices |