1 | /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_INLINE_FUNCTION_UTILS_H_ |
17 | #define TENSORFLOW_CORE_COMMON_RUNTIME_INLINE_FUNCTION_UTILS_H_ |
18 | |
19 | #include <functional> |
20 | #include <memory> |
21 | |
22 | #include "absl/types/optional.h" |
23 | #include "tensorflow/core/common_runtime/device.h" |
24 | #include "tensorflow/core/common_runtime/function_body.h" |
25 | #include "tensorflow/core/common_runtime/lower_function_call_inline_policy.h" |
26 | #include "tensorflow/core/framework/function.h" |
27 | #include "tensorflow/core/graph/graph.h" |
28 | #include "tensorflow/core/protobuf/config.pb.h" |
29 | |
30 | namespace tensorflow { |
31 | |
32 | static constexpr const char* const kNoInlineAttr = "_noinline" ; |
33 | |
34 | // Optionally override device assignment for nodes added to the graph for |
35 | // inlined functions: |
36 | // (1) Identity nodes added in place of function input arguments. |
37 | // (2) Identity nodes added in place of function return values. |
38 | // (3) Special NoOp nodes that enforce side-effects execution order. |
39 | // (4) All nodes inside function body specified in FunctionDef. |
40 | class InlinedFunctionBodyPlacer { |
41 | public: |
42 | virtual ~InlinedFunctionBodyPlacer() = default; |
43 | |
44 | virtual absl::optional<string> InputNodeDevice(int input_index) const = 0; |
45 | virtual absl::optional<string> OutputNodeDevice(int output_index) const = 0; |
46 | // Returns true if the added input/output identity nodes should be colocated |
47 | // with the corresponding input/output from the function body. |
48 | virtual bool ColocateInputOutputIdentities() const = 0; |
49 | virtual absl::optional<string> ControlNodeDevice() const = 0; |
50 | virtual absl::optional<string> BodyNodeDevice(const NodeDef& ndef) const = 0; |
51 | |
52 | // LINT.IfChange |
53 | // Place input nodes on the same device as the corresponding caller input |
54 | // node. Do not specify any placement for all other nodes. |
55 | static std::unique_ptr<InlinedFunctionBodyPlacer> DefaultPlacer( |
56 | const Graph& graph, const Node& caller); |
57 | |
58 | // Place all nodes on the same device as caller node. |
59 | static std::unique_ptr<InlinedFunctionBodyPlacer> SingleDevicePlacer( |
60 | const Graph& graph, const Node& caller); |
61 | |
62 | // Place input nodes on the same device as the corresponding caller input |
63 | // node. Do not place output node. Place control nodes on the same device as |
64 | // caller node. For all function body nodes set job, replica and task |
65 | // parts of the device assignment to match function caller node where those |
66 | // are unspecified. |
67 | static std::unique_ptr<InlinedFunctionBodyPlacer> MultiDevicePlacer( |
68 | const Graph& graph, const Node& caller); |
69 | // LINT.ThenChange(lower_function_call_inline_policy.h) |
70 | |
71 | using Factory = std::function<std::unique_ptr<InlinedFunctionBodyPlacer>( |
72 | const Graph&, const Node&)>; |
73 | |
74 | struct Config { |
75 | string name; |
76 | Factory get; |
77 | }; |
78 | |
79 | static Config Default() { return {"default" , DefaultPlacer}; } |
80 | static Config SingleDevice() { return {"single_device" , SingleDevicePlacer}; } |
81 | static Config MultiDevice() { return {"multi_device" , MultiDevicePlacer}; } |
82 | }; |
83 | |
84 | struct InlineFunctionBodyOptions { |
85 | // All nodes that have incoming control edge *from* the function call node, |
86 | // will be forwarded to the "output control node". There are two options for |
87 | // choosing which nodes will have a control edge *to* the "output control |
88 | // node": |
89 | // a) control returns (`control_ret` field in FunctionDef) |
90 | // b) data returns (`ret` field in FunctionDef) |
91 | enum class OutputControlSource { kDataOutputs, kControlOutputs }; |
92 | |
93 | // Keep a node in a graph with the same name as the function call node: |
94 | // |
95 | // a) DoNotKeep: Function call node is fully inlined, and there is no node in |
96 | // a graph with the same name. |
97 | // |
98 | // b) Fetchable: Add an IdentityN node to the graph in place of the inlined |
99 | // function call node. It will have a control edge from inlined |
100 | // 'output_control_node' and data edges from function output nodes. |
101 | // The IdentityN node will be placed on the same device as the caller node. |
102 | // |
103 | // This is mostly for compatibility with Tensorflow v1 and sessions. |
104 | // When we prepare a graph for execution in |
105 | // GraphExecutionState::MakeForBaseGraph we don't know what nodes will be |
106 | // fetched, so we can't safely remove any of them. When graph executed as a |
107 | // function it has 'Retval' nodes for all fetched tensors, and we can |
108 | // safely inline function calls. |
109 | // |
110 | // c) Targetable: Add a NoOp node to the graph in place of the inlined |
111 | // function call node. It will have a control edge from inline |
112 | // 'output_control_node' and no data edges. NoOp node will be placed on the |
113 | // same device as the caller node. This will keep the inlined function call |
114 | // node a valid 'session.run' target, and also will keep it a valid control |
115 | // output node. |
116 | enum class KeepCallerNode { kDoNotKeep, kFetchable, kTargetable }; |
117 | |
118 | // If 'true' function inlining is completely disabled. This allows to control |
119 | // function inlining for different types of function calls (see |
120 | // 'ExpandInlineFunctionsOptions' below). |
121 | bool disable_inlining = false; |
122 | // Ignore '_noinline' function attribute. |
123 | bool ignore_noinline = false; |
124 | // If 'true' function inlining will inline functions in implementation |
125 | // selection group. Normally those functions should not be inlined; they will |
126 | // be handled by Grappler. |
127 | bool inline_impl_selection_group_functions = false; |
128 | // Controls if we want to keep a node with the name as the function call node |
129 | // in a graph after function inlining. |
130 | KeepCallerNode keep_caller_node = KeepCallerNode::kDoNotKeep; |
131 | // For compatibility with Tensorflow v1 by default we will use data outputs. |
132 | // Control returns were added to Tensorflow v2 with automatic control |
133 | // dependencies tracking in Eager mode. |
134 | OutputControlSource output_control_src = OutputControlSource::kDataOutputs; |
135 | // Inlined function body placer decides what requested device assignments |
136 | // should be added to the nodes added to the graph. See documentation above |
137 | // for available strategies. |
138 | InlinedFunctionBodyPlacer::Config inlined_function_body_placer = |
139 | InlinedFunctionBodyPlacer::Default(); |
140 | // If true, frame names in the function body will be |
141 | // made unique in the resulting graph (e.g. by prepending a unique prefix). |
142 | // NOTE(mrry): Only set this option to false when there is a single function |
143 | // call in the graph (e.g. when making a remote function call via |
144 | // ClusterFunctionLibraryRuntime). This option is provided because the graph |
145 | // partitioner generates frame names that must remain unmodified across all |
146 | // partitions of a multi-device function. |
147 | bool uniquify_frame_names = true; |
148 | |
149 | // A human-readable debug string for this options. |
150 | string DebugString() const; |
151 | }; |
152 | |
153 | // Returns 'OkStatus()' iff the function '*fbody' can be inlined at 'node' |
154 | // based on the type signature of 'node' and 'fbody': |
155 | // |
156 | // (1) Caller node has the same number of inputs and outputs as the function. |
157 | // (2) Caller node inputs and outputs have the same data types as function |
158 | // inputs and returns. |
159 | // (3) Validation rules defined in InlineFunctionBodyOptions. |
160 | // |
161 | // If function can't be safely inlined, returns error message with details why |
162 | // inlining is not possible or safe. |
163 | Status ValidateInlining(const Node* node, const FunctionBody* fbody, |
164 | const InlineFunctionBodyOptions& options); |
165 | |
166 | // Given a "caller" in graph "g", which is a function call of a function |
167 | // to "fbody". Replaces the "caller" with fbody->graph and connects |
168 | // edges properly. "override_device" specifies whether inlining should replace |
169 | // explicitly specified devices inside fbody with the callee's device. |
170 | // |
171 | // Returns 'OkStatus()' if function was successfully inlined into the graph. |
172 | // If function inlining is not possible returns an error with a reason, and |
173 | // leaves the graph in unmodified state. |
174 | Status InlineFunctionBody(const FunctionLibraryDefinition& flib_def, Graph* g, |
175 | Node* caller, const FunctionBody* fbody, |
176 | const InlineFunctionBodyOptions& options); |
177 | |
178 | // There are three types of function calls that could be invoked during |
179 | // *Tensorflow graph execution*: |
180 | // |
181 | // 1) Native function call (node.type_string() is the function name). These |
182 | // functions are always executed on a single-device, which is the device of |
183 | // the function call node. |
184 | // |
185 | // 2) Multi-device function calls (PartitionedCall or StatefulPartitionedCall |
186 | // ops) can execute on multiple devices and accept DT_RESOURCE inputs that |
187 | // belong to different devices. This type of functions was added in |
188 | // Tensorflow 2.0 Eager mode, and it has control outputs to represent |
189 | // side-effects that must always execute (see `control_ret` in FunctionDef). |
190 | // |
191 | // 3) SymbolicGradient has been deprecated for a while, but we still keep it and |
192 | // use `native` options for inlining for compatibility. |
193 | // |
194 | // We need to have distinct inlining rules for compatibility with Tensorflow v1. |
195 | // |
196 | // There are few other places in Tensorflow that could execute functions: |
197 | // |
198 | // 1) common_runtime/eager/kernel_and_device.{h,cc} - executes "top level" |
199 | // functions directly via function library runtime, without going through |
200 | // the graph. |
201 | // 2) tf.data pipelines - also execute functions directly via function library |
202 | // runtime with custom executors. |
203 | struct ExpandInlineFunctionsOptions { |
204 | ExpandInlineFunctionsOptions() : native_options(), multi_device_options() { |
205 | using OutputControlSrc = InlineFunctionBodyOptions::OutputControlSource; |
206 | multi_device_options.output_control_src = OutputControlSrc::kControlOutputs; |
207 | } |
208 | |
209 | InlineFunctionBodyOptions native_options; |
210 | InlineFunctionBodyOptions multi_device_options; |
211 | }; |
212 | |
213 | // WARNING(ezhulenev): PLEASE DO NOT USE THIS FUNCTION. This is a temporary |
214 | // workaround that will be enabled only during the function inlining unification |
215 | // (b/126811947). Contact ezhulenev@ if you think you need it. |
216 | // TODO(ezhulenev): Delete this function. |
217 | bool ExpandInlineFunctions(FunctionLibraryRuntime* lib, Graph* graph, |
218 | const ExpandInlineFunctionsOptions& options); |
219 | |
220 | // For each node in "graph", if "lib" indicates that the node is a |
221 | // function call, inline the function body. Returns true if at least |
222 | // one node is inlined. |
223 | // |
224 | // This routine goes through "graph" nodes once and applies the |
225 | // inlining. The caller may decide to apply the inlining on "graph" |
226 | // multiple times by calling ExpandInlineFunctions a few times. |
227 | // |
228 | // Function calls that can't be safely inlined into the graph (ValidateInlining |
229 | // returns error), are ignored. |
230 | // |
231 | // TODO(ezhulenev): We do not FunctionLibraryRuntime for this. We need just the |
232 | // FunctionLibraryDefinition and FunctionDefToBodyHelper to implement this (see |
233 | // lower_function_call.cc). |
234 | inline bool ExpandInlineFunctions(FunctionLibraryRuntime* lib, Graph* graph) { |
235 | return ExpandInlineFunctions(lib, graph, ExpandInlineFunctionsOptions()); |
236 | } |
237 | |
238 | } // end namespace tensorflow |
239 | |
240 | #endif // TENSORFLOW_CORE_COMMON_RUNTIME_INLINE_FUNCTION_UTILS_H_ |
241 | |