| 1160 | } |
| 1161 | |
| 1162 | Status EncapsulateSubgraphsPass::Run( |
| 1163 | const GraphOptimizationPassOptions& options) { |
| 1164 | VLOG(1) << "EncapsulateSubgraphsPass::Run"; |
| 1165 | if (VLOG_IS_ON(1)) { |
| 1166 | DumpGraphToFile("encapsulate_subgraphs_before", **options.graph, |
| 1167 | options.flib_def); |
| 1168 | } |
| 1169 | |
| 1170 | std::unique_ptr<Graph> graph_out; |
| 1171 | FunctionLibraryDefinition* const library = options.flib_def; |
| 1172 | |
| 1173 | // Constant folding below might need to run part of the function to compute |
| 1174 | // constants. Create an FunctionLibraryRuntime with a single CPU device |
| 1175 | // that can run the part of the function. |
| 1176 | // NOTE: If this turns out to be slow, we can cache the FLRs keyed by |
| 1177 | // `options`. |
| 1178 | SessionOptions session_options; |
| 1179 | auto* device_count = session_options.config.mutable_device_count(); |
| 1180 | device_count->insert({"CPU", 1}); |
| 1181 | std::vector<std::unique_ptr<Device>> devices; |
| 1182 | |
| 1183 | DeviceFactory* cpu_factory = DeviceFactory::GetFactory("CPU"); |
| 1184 | if (!cpu_factory) { |
| 1185 | return errors::NotFound( |
| 1186 | "CPU Factory not registered. Can't run EncapsulateSubgraphsPass"); |
| 1187 | } |
| 1188 | TF_RETURN_IF_ERROR(cpu_factory->CreateDevices( |
| 1189 | session_options, "/job:localhost/replica:0/task:0", &devices)); |
| 1190 | if (devices.empty()) { |
| 1191 | return errors::NotFound( |
| 1192 | "Failed to create a CPU device for EncapsulateSubgraphsPass"); |
| 1193 | } |
| 1194 | |
| 1195 | std::unique_ptr<DeviceMgr> device_mgr = |
| 1196 | absl::make_unique<DeviceMgr>(std::move(devices)); |
| 1197 | OptimizerOptions opts; |
| 1198 | std::unique_ptr<ProcessFunctionLibraryRuntime> pflr( |
| 1199 | new ProcessFunctionLibraryRuntime(device_mgr.get(), |
| 1200 | options.session_options->env, |
| 1201 | TF_GRAPH_DEF_VERSION, library, opts)); |
| 1202 | FunctionLibraryRuntime* flr = |
| 1203 | pflr->GetFLR("/job:localhost/replica:0/task:0/device:CPU:0"); |
| 1204 | if (flr == nullptr) { |
| 1205 | return errors::Internal( |
| 1206 | "Failed to create and retrieve function library runtime to run " |
| 1207 | "constant folding"); |
| 1208 | } |
| 1209 | |
| 1210 | auto rewrite_subgraph = |
| 1211 | [flr](const std::vector<OutputTensor>& arg_source_tensors, |
| 1212 | std::unique_ptr<Graph>* subgraph, |
| 1213 | std::vector<int>* input_permutation, |
| 1214 | std::vector<int>* output_permutation, NodeDef* node) { |
| 1215 | // Optimize the subgraph. |
| 1216 | // Do not constant fold nodes that output DT_VARIANT type tensors. |
| 1217 | // XLA does not support Const nodes of Variant type since it needs |
| 1218 | // to know the original ops to be able to compile them to the relevant |
| 1219 | // XLA form. |
nothing calls this directly
no test coverage detected