TODO(irving): Replace with const Node& version below.
| 1223 | |
| 1224 | // TODO(irving): Replace with const Node& version below. |
| 1225 | Status FindKernelRegistration( |
| 1226 | const DeviceType& device_type, StringPiece node_name, |
| 1227 | bool has_experimental_debug_info, |
| 1228 | const NodeDef_ExperimentalDebugInfo& experimental_debug_info, |
| 1229 | StringPiece node_op, AttrSlice node_attrs, const KernelRegistration** reg, |
| 1230 | bool* was_attr_mismatch) { |
| 1231 | *reg = nullptr; |
| 1232 | *was_attr_mismatch = false; |
| 1233 | // Label defaults to empty if not found in NodeDef. |
| 1234 | const string& label = GetNodeAttrString(node_attrs, kKernelAttr); |
| 1235 | |
| 1236 | const string key = Key(node_op, device_type, label); |
| 1237 | auto typed_registry = GlobalKernelRegistryTyped(); |
| 1238 | tf_shared_lock lock(typed_registry->mu); |
| 1239 | auto regs = typed_registry->registry.equal_range(key); |
| 1240 | for (auto iter = regs.first; iter != regs.second; ++iter) { |
| 1241 | // If there is a kernel registered for the op and device_type, |
| 1242 | // check that the attrs match. |
| 1243 | bool match; |
| 1244 | TF_RETURN_IF_ERROR(KernelAttrsMatch(iter->second.def, node_attrs, &match)); |
| 1245 | if (match) { |
| 1246 | if (*reg != nullptr) { |
| 1247 | return errors::InvalidArgument( |
| 1248 | "Multiple OpKernel registrations match NodeDef '", |
| 1249 | FormatNodeDefForError(node_name, has_experimental_debug_info, |
| 1250 | experimental_debug_info), |
| 1251 | "': '", ProtoShortDebugString((*reg)->def), "' and '", |
| 1252 | ProtoShortDebugString(iter->second.def), "'"); |
| 1253 | } |
| 1254 | *reg = &iter->second; |
| 1255 | } else { |
| 1256 | *was_attr_mismatch = true; |
| 1257 | } |
| 1258 | } |
| 1259 | // Check if no device specific registrations found. If not, try finding a |
| 1260 | // default kernel. |
| 1261 | if (*reg == nullptr && |
| 1262 | !IsSymbolicExecutionDevice(device_type.type_string())) { |
| 1263 | const string default_key = Key(node_op, DEVICE_DEFAULT, label); |
| 1264 | auto regs = typed_registry->registry.equal_range(default_key); |
| 1265 | for (auto iter = regs.first; iter != regs.second; ++iter) { |
| 1266 | // If there is a kernel registered for the op and device_type, |
| 1267 | // check that the attrs match. |
| 1268 | bool match; |
| 1269 | TF_RETURN_IF_ERROR( |
| 1270 | KernelAttrsMatch(iter->second.def, node_attrs, &match)); |
| 1271 | if (match) { |
| 1272 | if (*reg != nullptr) { |
| 1273 | return errors::InvalidArgument( |
| 1274 | "Multiple Default OpKernel registrations match NodeDef '", |
| 1275 | FormatNodeDefForError(node_name, has_experimental_debug_info, |
| 1276 | experimental_debug_info), |
| 1277 | "': '", ProtoShortDebugString((*reg)->def), "' and '", |
| 1278 | ProtoShortDebugString(iter->second.def), "'"); |
| 1279 | } |
| 1280 | *reg = &iter->second; |
| 1281 | } else { |
| 1282 | *was_attr_mismatch = true; |
no test coverage detected