xref: /aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/inline_function_utils.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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 
Default()79   static Config Default() { return {"default", DefaultPlacer}; }
SingleDevice()80   static Config SingleDevice() { return {"single_device", SingleDevicePlacer}; }
MultiDevice()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 'Status::OK()' 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 'Status::OK()' 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 {
ExpandInlineFunctionsOptionsExpandInlineFunctionsOptions204   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).
ExpandInlineFunctions(FunctionLibraryRuntime * lib,Graph * graph)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