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