1 | /* Copyright 2021 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_LOWER_FUNCTION_CALL_INLINE_POLICY_H_ |
17 | #define TENSORFLOW_CORE_COMMON_RUNTIME_LOWER_FUNCTION_CALL_INLINE_POLICY_H_ |
18 | |
19 | #include "tensorflow/core/graph/graph.h" |
20 | |
21 | namespace tensorflow { |
22 | |
23 | // LINT.IfChange |
24 | enum class FunctionCallInlinePolicy { |
25 | // Place input nodes on the same device as the corresponding caller input |
26 | // node. Do not specify any placement for all other nodes. |
27 | kDefaultPlacer, |
28 | |
29 | // Place all nodes on the same device as caller node. |
30 | kSingleDevicePlacer, |
31 | |
32 | // Place input nodes on the same device as the corresponding caller input |
33 | // node. Do not place output node. Place control nodes on the same device as |
34 | // caller node. For all function body nodes overrides job, replica and task |
35 | // parts of the device assignment to match function caller node. |
36 | kMultiDevicePlacer |
37 | }; |
38 | // LINT.ThenChange(inline_function_utils.h,\ |
39 | // ../../compiler/mlir/tensorflow/ir/tf_ops.cc) |
40 | |
41 | struct LowerFunctionalOpsConstants { |
42 | static constexpr const char* const kLowerUsingSwitchMergeAttr = |
43 | "_lower_using_switch_merge" ; |
44 | static constexpr const char* const kLowerAsMultiDeviceFunctionAttr = |
45 | "_lower_as_multi_device_function" ; |
46 | }; |
47 | |
48 | // Inliner policy used in common runtime's lower function call op. |
49 | |
50 | // Returns the function call inline policy to use for a given call. |
51 | FunctionCallInlinePolicy GetFunctionCallInlinePolicy(const Node* n); |
52 | |
53 | // Overload of GetFunctionCallInlinePolicy that doesn't require an op but only |
54 | // the features required. |
55 | FunctionCallInlinePolicy GetFunctionCallInlinePolicy( |
56 | bool is_partioned_call, bool has_lower_as_multi_device_function_attr); |
57 | |
58 | } // namespace tensorflow |
59 | |
60 | #endif // TENSORFLOW_CORE_COMMON_RUNTIME_LOWER_FUNCTION_CALL_INLINE_POLICY_H_ |
61 | |