| 149 | } |
| 150 | |
| 151 | Status XlaKernelCreator::CreateKernel(FunctionLibraryRuntime* flr, |
| 152 | const NodeDef& node_def, |
| 153 | std::unique_ptr<OpKernel>* kernel) const { |
| 154 | if (!CanCreateKernel(*flr, node_def)) { |
| 155 | return errors::Internal("Invalid node: ", node_def.ShortDebugString()); |
| 156 | } |
| 157 | |
| 158 | VLOG(3) << "Attempting to create XlaLaunchOp for " << node_def.DebugString(); |
| 159 | |
| 160 | // Make sure that kernels have been registered on the JIT device. |
| 161 | XlaOpRegistry::RegisterCompilationKernels(); |
| 162 | std::vector<RecursiveCompilabilityChecker::UncompilableNodeInfo> |
| 163 | uncompilable_node_info; |
| 164 | if (!IsCompilable(flr, node_def, &uncompilable_node_info)) { |
| 165 | string message = absl::StrCat( |
| 166 | "Function invoked by the following node is not compilable: ", |
| 167 | node_def.ShortDebugString(), ".\n"); |
| 168 | absl::StrAppend(&message, "Uncompilable nodes:\n"); |
| 169 | for (const auto& node_info : uncompilable_node_info) { |
| 170 | string node_message = |
| 171 | absl::StrCat("\t", node_info.name, ": ", |
| 172 | node_info.uncompilable_reason, "\n", "\tStacktrace:\n"); |
| 173 | for (const auto& stack_frame : node_info.stack_trace) { |
| 174 | absl::StrAppendFormat(&node_message, "\t\tNode: %s, function: %s\n", |
| 175 | stack_frame.name, stack_frame.function_name); |
| 176 | } |
| 177 | absl::StrAppend(&message, node_message); |
| 178 | } |
| 179 | VLOG(1) << message; |
| 180 | // node_def is calling a function that XLA can't compile. |
| 181 | return errors::InvalidArgument(message); |
| 182 | } |
| 183 | |
| 184 | // Get function body, constant args, and resource args. |
| 185 | const FunctionBody* fbody = nullptr; |
| 186 | std::vector<int> constant_arg_indices; |
| 187 | std::vector<int> resource_arg_indices; |
| 188 | TF_RETURN_IF_ERROR(GetBodyAndConstantsAndResources( |
| 189 | flr, node_def, &fbody, &constant_arg_indices, &resource_arg_indices)); |
| 190 | |
| 191 | // Set input and output memory types. |
| 192 | MemoryTypeVector input_memory_types(fbody->arg_types.size(), DEVICE_MEMORY); |
| 193 | // These indices are used only for optimization purposes. They allow us |
| 194 | // to loop over constant_arg_indices and resource_arg_indices only once |
| 195 | // while iterating over all the function arguments checking if it is a |
| 196 | // resource or a constant. |
| 197 | // The reason we optimized this code is because functions can have a lot of |
| 198 | // captured arguments. For example, the backward pass of ResNet50 takes in all |
| 199 | // 214 variables and a similar number of activations. |
| 200 | SinglePassSearch constants_search(&constant_arg_indices); |
| 201 | SinglePassSearch resources_search(&resource_arg_indices); |
| 202 | for (int i = 0; i < fbody->arg_types.size(); ++i) { |
| 203 | if (resources_search.ScanForValue(i) || constants_search.ScanForValue(i)) { |
| 204 | // Compile-time constants and resource handles are expected to be in |
| 205 | // host memory. |
| 206 | input_memory_types[i] = HOST_MEMORY; |
| 207 | } |
| 208 | } |