| 162 | } |
| 163 | |
| 164 | Status PartitionedCallOp::Instantiate(FunctionLibraryRuntime* lib, |
| 165 | OpKernelContext* ctx, |
| 166 | std::vector<Tensor>* inputs, |
| 167 | FunctionLibraryRuntime::Handle* handle) { |
| 168 | FunctionLibraryRuntime::InstantiateOptions opts; |
| 169 | // TODO(b/141576771): Propagate ctxt's ConfigProto into opts.config_proto. |
| 170 | |
| 171 | #ifndef __ANDROID__ |
| 172 | // Android tf library does not include grappler. |
| 173 | grappler::GrapplerItem::OptimizationOptions optimization_options; |
| 174 | // Tensorflow 2.0 in eager mode with automatic control dependencies will |
| 175 | // prune all nodes that are not in the transitive fanin of the fetch nodes. |
| 176 | // However because the function will be executed via FunctionLibraryRuntime, |
| 177 | // and current function implementation does not prune stateful and dataset |
| 178 | // ops, we rely on Grappler to do the correct graph pruning. |
| 179 | optimization_options.allow_pruning_stateful_and_dataset_ops = true; |
| 180 | |
| 181 | // All the nested function calls will be executed and optimized via |
| 182 | // PartitionedCallOp, there is no need to optimize functions now. |
| 183 | optimization_options.optimize_function_library = false; |
| 184 | |
| 185 | opts.optimize_graph_fn = |
| 186 | std::bind(grappler::OptimizeGraph, std::placeholders::_1, |
| 187 | std::placeholders::_2, std::placeholders::_3, |
| 188 | std::placeholders::_4, std::placeholders::_5, *config_proto_, |
| 189 | func_->name(), optimization_options, std::placeholders::_6); |
| 190 | #endif |
| 191 | |
| 192 | // In some contexts like running the graph to evaluate constants, |
| 193 | // the FLR won't have any device. |
| 194 | opts.target = lib->device() == nullptr ? "" : lib->device()->name(); |
| 195 | opts.is_multi_device_function = true; |
| 196 | opts.graph_collector = ctx->graph_collector(); |
| 197 | opts.executor_type = executor_type_; |
| 198 | |
| 199 | OpInputList args; |
| 200 | TF_RETURN_IF_ERROR(ctx->input_list("args", &args)); |
| 201 | Device* cpu_device; |
| 202 | TF_RETURN_IF_ERROR(lib->device_mgr()->LookupDevice("CPU:0", &cpu_device)); |
| 203 | |
| 204 | inputs->reserve(args.size()); |
| 205 | for (const Tensor& tensor : args) { |
| 206 | inputs->push_back(tensor); |
| 207 | DataType dtype = tensor.dtype(); |
| 208 | if (dtype == DT_RESOURCE) { |
| 209 | const ResourceHandle& handle = tensor.flat<ResourceHandle>()(0); |
| 210 | opts.input_devices.push_back(handle.device()); |
| 211 | } else if (MTypeFromDType(dtype) == HOST_MEMORY) { |
| 212 | opts.input_devices.push_back(cpu_device->name()); |
| 213 | } else { |
| 214 | opts.input_devices.push_back(opts.target); |
| 215 | } |
| 216 | } |
| 217 | |
| 218 | TF_RETURN_IF_ERROR( |
| 219 | FillOutputDevices(*lib, *cpu_device, AttrSlice(&func_->attr()), &opts)); |
| 220 | |
| 221 | TF_RETURN_IF_ERROR( |
nothing calls this directly
no test coverage detected