1 /* Copyright 2020 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_EAGER_PLACEMENT_UTILS_H_ 17 #define TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_PLACEMENT_UTILS_H_ 18 19 #include "tensorflow/c/eager/immediate_execution_operation.h" 20 #include "tensorflow/core/common_runtime/eager/context.h" 21 #include "tensorflow/core/platform/status.h" 22 #include "tensorflow/core/platform/stringpiece.h" 23 24 namespace tensorflow { 25 namespace eager { 26 27 bool IsColocationExempt(StringPiece op_name); 28 29 bool IsFunction(StringPiece op_name); 30 31 // TODO(b/154234908): Unify placement logic. 32 // TODO(b/159647422): Add C++ unit tests for placement logic. 33 34 // Pin the op to cpu if all op inputs are on the CPU, small (<64 elements) and 35 // integers (int32/int64). This can be disabled by setting the environment 36 // variable "TF_EAGER_ENABLE_SMALL_TENSOR_CPU_PINNING" to "0" or "false". 37 Status MaybePinSmallOpsToCpu( 38 bool* result, StringPiece op_name, 39 absl::Span<ImmediateExecutionTensorHandle* const> args, 40 StringPiece cpu_device_name); 41 42 // If a resource touching input is specified, all resource-touching ops run in 43 // the device the resource is, regardless of anything else that has been 44 // specified. This is identical to the graph mode behavior. 45 Status MaybePinToResourceDevice(Device** device, const EagerOperation& op); 46 } // namespace eager 47 } // namespace tensorflow 48 49 #endif // TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_PLACEMENT_UTILS_H_ 50