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 #include "tensorflow/core/tfrt/eager/core_runtime/op_handler_selector.h"
17
18 #include "tensorflow/core/common_runtime/eager/attr_builder.h"
19 #include "tensorflow/core/common_runtime/eager/context.h"
20 #include "tensorflow/core/common_runtime/eager/placement_utils.h"
21 #include "tensorflow/core/platform/errors.h"
22 #include "tensorflow/core/tfrt/eager/core_runtime/op_handler_registry.h"
23 #include "tfrt/core_runtime/core_runtime.h" // from @tf_runtime
24 #include "tfrt/host_context/device.h" // from @tf_runtime
25 #include "tfrt/host_context/host_context.h" // from @tf_runtime
26 #include "tfrt/support/error_util.h" // from @tf_runtime
27 #include "tfrt/support/ref_count.h" // from @tf_runtime
28
29 namespace tfrt {
30 namespace tf {
31
EagerOpHandlerSelector(CoreRuntime * core_runtime,EagerContext * eager_context,OpHandler * fallback_op_handler,bool pin_small_ops_to_cpu)32 EagerOpHandlerSelector::EagerOpHandlerSelector(CoreRuntime* core_runtime,
33 EagerContext* eager_context,
34 OpHandler* fallback_op_handler,
35 bool pin_small_ops_to_cpu)
36 : core_runtime_(core_runtime),
37 eager_context_(eager_context),
38 cpu_device_(core_runtime->GetHostContext()->GetHostDevice()),
39 cpu_op_handler_(core_runtime_->GetOpHandler(cpu_device_.name())),
40 fallback_op_handler_(fallback_op_handler),
41 pin_small_ops_to_cpu_(pin_small_ops_to_cpu) {
42 assert(cpu_op_handler_);
43 assert(fallback_op_handler_);
44 }
45
~EagerOpHandlerSelector()46 EagerOpHandlerSelector::~EagerOpHandlerSelector() {}
47
SelectFromArguments(const ImmediateExecutionOperation & op,OpHandler ** op_handler)48 Status EagerOpHandlerSelector::SelectFromArguments(
49 const ImmediateExecutionOperation& op, OpHandler** op_handler) {
50 // If the op contains resource handle, place the op on the device of the
51 // resource.
52 // TODO(tfrt-devs): Unify this logic with MaybePinToResourceDevice in Eager
53 // runtime.
54 for (int i = 0; i < op.GetInputs().size(); i++) {
55 auto& handle = op.GetInputs()[i];
56 Status s;
57 if (handle->DataType() == tensorflow::DT_RESOURCE) {
58 auto device_name = handle->DeviceName(&s);
59 TF_RETURN_IF_ERROR(s);
60 *op_handler = core_runtime_->GetOpHandler(device_name);
61 if (*op_handler != nullptr) {
62 DVLOG(1) << "Setting device of operation " << op.Name() << " to "
63 << device_name << " because input #" << i
64 << " is a resource in this device.";
65 return ::tensorflow::OkStatus();
66 }
67 }
68 }
69
70 // Pin the op to cpu op handler if it is a small ops and all its inputs
71 // are on cpu already.
72 if (pin_small_ops_to_cpu_) {
73 bool pin_to_cpu;
74 TF_RETURN_IF_ERROR(tensorflow::eager::MaybePinSmallOpsToCpu(
75 &pin_to_cpu, op.Name(), op.GetInputs(),
76 {cpu_device_.name().data(), cpu_device_.name().size()}));
77 if (pin_to_cpu) {
78 *op_handler = cpu_op_handler_;
79 return ::tensorflow::OkStatus();
80 }
81 }
82
83 // Note: The output op_handler is nullptr.
84 return ::tensorflow::OkStatus();
85 }
86
SelectFromNodeDef(const ImmediateExecutionOperation & op,const NodeDef * ndef,OpHandler ** op_handler)87 Status EagerOpHandlerSelector::SelectFromNodeDef(
88 const ImmediateExecutionOperation& op, const NodeDef* ndef,
89 OpHandler** op_handler) {
90 const auto& requested_device = op.DeviceName();
91
92 // TODO(fishx): Use TFRT native op registry to select op handler.
93
94 // TODO(fishx): Add a cache for following device placement using current TF.
95 // Use EagerContext from current tf to select op handler for this op.
96 tensorflow::DeviceNameUtils::ParsedName device_parsed_name;
97 if (!tensorflow::DeviceNameUtils::ParseFullName(requested_device,
98 &device_parsed_name)) {
99 return tensorflow::errors::InvalidArgument("Failed to parse device name: ",
100 requested_device);
101 }
102
103 tensorflow::Device* device;
104 TF_RETURN_IF_ERROR(
105 eager_context_->SelectDevice(device_parsed_name, *ndef, &device));
106
107 *op_handler = core_runtime_->GetOpHandler(device->name());
108
109 if (!(*op_handler)) *op_handler = fallback_op_handler_;
110
111 return ::tensorflow::OkStatus();
112 }
113
114 } // namespace tf
115 } // namespace tfrt
116