• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2015 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/framework/memory_types.h"
17 
18 #include <utility>
19 
20 #include "tensorflow/core/framework/kernel_def.pb.h"
21 #include "tensorflow/core/framework/node_def.pb.h"
22 #include "tensorflow/core/framework/node_def_util.h"
23 #include "tensorflow/core/framework/op_kernel.h"
24 #include "tensorflow/core/framework/types.h"
25 #include "tensorflow/core/lib/core/errors.h"
26 #include "tensorflow/core/platform/types.h"
27 
28 namespace tensorflow {
29 
30 namespace {
31 // Returns the largest endpoint of anything in the name_map.
GetTotal(const NameRangeMap & name_map)32 int GetTotal(const NameRangeMap& name_map) {
33   int total = 0;
34   for (const auto& item : name_map) {
35     total = std::max(total, item.second.second);
36   }
37   return total;
38 }
39 
40 // Fills memory_types for either input or output, setting everything
41 // to DEVICE_MEMORY except those args in host_memory_args.  Removes
42 // elements of host_memory_args that were used.
MemoryTypesHelper(const NameRangeMap & name_map,std::vector<string> * host_memory_args,MemoryTypeVector * memory_types)43 void MemoryTypesHelper(const NameRangeMap& name_map,
44                        std::vector<string>* host_memory_args,
45                        MemoryTypeVector* memory_types) {
46   // Update args that have been marked as in "HOST_MEMORY".
47   size_t keep = 0;
48   for (size_t i = 0; i < host_memory_args->size(); ++i) {
49     auto iter = name_map.find((*host_memory_args)[i]);
50     if (iter != name_map.end()) {
51       for (int j = iter->second.first; j < iter->second.second; ++j) {
52         (*memory_types)[j] = HOST_MEMORY;
53       }
54     } else {
55       // (*host_memory_args)[i] not found, save it for the next pass.
56       if (i > keep) (*host_memory_args)[keep] = (*host_memory_args)[i];
57       ++keep;
58     }
59   }
60   host_memory_args->resize(keep);
61 }
62 
IsFunctionCallOp(const string & op_type)63 bool IsFunctionCallOp(const string& op_type) {
64   return op_type == "SymbolicGradient" || op_type == "PartitionedCall" ||
65          op_type == "StatefulPartitionedCall" || op_type == "While";
66 }
67 
68 }  // namespace
69 
MTypeFromDType(const DataType dtype)70 MemoryType MTypeFromDType(const DataType dtype) {
71   return (dtype == DT_INT32 || DataTypeAlwaysOnHost(dtype)) ? HOST_MEMORY
72                                                             : DEVICE_MEMORY;
73 }
74 
MemoryTypesForNode(const OpRegistryInterface * op_registry,const DeviceType & device_type,const NodeDef & ndef,MemoryTypeVector * inp_mtypes,MemoryTypeVector * out_mtypes)75 Status MemoryTypesForNode(const OpRegistryInterface* op_registry,
76                           const DeviceType& device_type, const NodeDef& ndef,
77                           MemoryTypeVector* inp_mtypes,
78                           MemoryTypeVector* out_mtypes) {
79   // Look up the Op registered for this op name.
80   const OpDef* op_def;
81   TF_RETURN_IF_ERROR(op_registry->LookUpOpDef(ndef.op(), &op_def));
82 
83   // Look up the Kernel registered for this node def.
84   const KernelDef* kdef = nullptr;
85   Status status =
86       FindKernelDef(device_type, ndef, &kdef, nullptr /* kernel_class_name */);
87 
88   DataTypeVector inp_dtypes;
89   DataTypeVector out_dtypes;
90   TF_RETURN_IF_ERROR(
91       InOutTypesForNode(ndef, *op_def, &inp_dtypes, &out_dtypes));
92 
93   inp_mtypes->clear();
94   out_mtypes->clear();
95 
96   // For functions (which have no KernelDef) and their gradients, we can only
97   // best-effort derive the memory type from the data type. For now, we assume
98   // int32 is always on host memory and other types are always on device memory.
99   // TODO(zhifengc,phawkins): We should do type inference over function bodies
100   // to derive the correct input/output memory types. We should also split
101   // host-memory and non host-memory arguments into separate type lists.
102   if (!status.ok() || IsFunctionCallOp(ndef.op())) {
103     for (const auto& t : inp_dtypes) inp_mtypes->push_back(MTypeFromDType(t));
104     for (const auto& t : out_dtypes) out_mtypes->push_back(MTypeFromDType(t));
105     return Status::OK();
106   }
107 
108   // Gets the input/output names and their corresponding endpoint ranges.
109   NameRangeMap inp_names;
110   NameRangeMap out_names;
111   TF_RETURN_IF_ERROR(NameRangesForNode(ndef, *op_def, &inp_names, &out_names));
112 
113   // Now that we know the size, fill with the default 'DEVICE_MEMORY'.
114   inp_mtypes->resize(GetTotal(inp_names), DEVICE_MEMORY);
115   out_mtypes->resize(GetTotal(out_names), DEVICE_MEMORY);
116 
117   // Fills in host memory types based on the kernel def.
118   const auto& from_proto = kdef->host_memory_arg();
119   std::vector<string> host_memory_args(from_proto.begin(), from_proto.end());
120   MemoryTypesHelper(inp_names, &host_memory_args, inp_mtypes);
121   MemoryTypesHelper(out_names, &host_memory_args, out_mtypes);
122   if (!host_memory_args.empty()) {
123     return errors::InvalidArgument(
124         "HostMemory args '", str_util::Join(host_memory_args, "', '"),
125         "' not found in OpDef: ", SummarizeOpDef(*op_def));
126   }
127   CHECK_LE(inp_mtypes->size(), inp_dtypes.size());
128   CHECK_LE(out_mtypes->size(), out_dtypes.size());
129 
130   // Mark e.g. all resource and string types as host memory.
131   for (int i = 0; i < inp_mtypes->size(); ++i) {
132     if (DataTypeAlwaysOnHost(inp_dtypes[i])) {
133       (*inp_mtypes)[i] = HOST_MEMORY;
134     }
135   }
136   for (int i = 0; i < out_mtypes->size(); ++i) {
137     if (DataTypeAlwaysOnHost(out_dtypes[i])) {
138       (*out_mtypes)[i] = HOST_MEMORY;
139     }
140   }
141 
142   std::vector<int32> hostmem_attr;
143   if (GetNodeAttr(ndef, "_input_hostmem", &hostmem_attr).ok()) {
144     for (int32 i : hostmem_attr) {
145       if (0 <= i && i < inp_mtypes->size()) {
146         (*inp_mtypes)[i] = HOST_MEMORY;
147       }
148     }
149   }
150   if (GetNodeAttr(ndef, "_output_hostmem", &hostmem_attr).ok()) {
151     for (int32 i : hostmem_attr) {
152       if (0 <= i && i < out_mtypes->size()) {
153         (*out_mtypes)[i] = HOST_MEMORY;
154       }
155     }
156   }
157 
158   return Status::OK();
159 }
160 
161 }  // namespace tensorflow
162