1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
30namespace tensorflow {
31
32static 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.
40class 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
84struct 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.
163Status 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.
174Status 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.
203struct 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.
217bool 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).
234inline 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