| 2199 | } |
| 2200 | |
| 2201 | bool InlineAllFunctions(GraphDef* graphdef) { |
| 2202 | if (graphdef->library().function().empty()) { |
| 2203 | VLOG(kLogLevelModelUnchanged) << "No functions to inline."; |
| 2204 | return false; |
| 2205 | } |
| 2206 | |
| 2207 | // Override "_noinline" attribute on all functions |
| 2208 | GraphDef graphdef_copy(*graphdef); |
| 2209 | for (auto& function : |
| 2210 | (*graphdef_copy.mutable_library()->mutable_function())) { |
| 2211 | auto* attributes = function.mutable_attr(); |
| 2212 | if (attributes->count(tensorflow::kNoInlineAttr) != 0) { |
| 2213 | (*attributes)[tensorflow::kNoInlineAttr].set_b(false); |
| 2214 | } |
| 2215 | } |
| 2216 | |
| 2217 | // Construct minimum resources needed to use ExpandInlineFunctions(). |
| 2218 | tensorflow::SessionOptions options; |
| 2219 | auto* device_count = options.config.mutable_device_count(); |
| 2220 | device_count->insert({"CPU", 1}); |
| 2221 | std::vector<std::unique_ptr<tensorflow::Device>> devices; |
| 2222 | TF_CHECK_OK(tensorflow::DeviceFactory::AddDevices( |
| 2223 | options, "/job:localhost/replica:0/task:0", &devices)); |
| 2224 | |
| 2225 | tensorflow::FunctionLibraryDefinition fld(tensorflow::OpRegistry::Global(), |
| 2226 | graphdef_copy.library()); |
| 2227 | tensorflow::DeviceMgr device_mgr(std::move(devices)); |
| 2228 | tensorflow::OptimizerOptions o_opts; |
| 2229 | tensorflow::ProcessFunctionLibraryRuntime pflr( |
| 2230 | &device_mgr, tensorflow::Env::Default(), TF_GRAPH_DEF_VERSION, &fld, |
| 2231 | o_opts, nullptr); |
| 2232 | tensorflow::FunctionLibraryRuntime* flr; |
| 2233 | flr = pflr.GetFLR("/job:localhost/replica:0/task:0/cpu:0"); |
| 2234 | |
| 2235 | tensorflow::Graph graph(fld); |
| 2236 | tensorflow::ImportGraphDefOptions gc_opts; |
| 2237 | gc_opts.validate_shape = false; |
| 2238 | const auto& tf_convert_status = tensorflow::ImportGraphDef( |
| 2239 | gc_opts, graphdef_copy, &graph, nullptr, nullptr); |
| 2240 | if (!tf_convert_status.ok()) { |
| 2241 | LOG(ERROR) << "tensorflow::ImportGraphDef failed with status: " |
| 2242 | << tf_convert_status.ToString(); |
| 2243 | return false; |
| 2244 | } |
| 2245 | |
| 2246 | // Iterate over the graph until there are no more nodes to be inlined. |
| 2247 | bool graph_modified = false; |
| 2248 | while (tensorflow::ExpandInlineFunctions(flr, &graph)) { |
| 2249 | graph_modified = true; |
| 2250 | } |
| 2251 | |
| 2252 | // Output inlined graph |
| 2253 | if (graph_modified) { |
| 2254 | LOG(INFO) << "Found and inlined TensorFlow functions."; |
| 2255 | graph.ToGraphDef(graphdef); |
| 2256 | } |
| 2257 | return graph_modified; |
| 2258 | } |
no test coverage detected