| 1296 | } |
| 1297 | |
| 1298 | Status ColocationGraph::InitializeMember(const Node& node, Member* member) { |
| 1299 | TF_RETURN_IF_ERROR(member->SetParentAndSupportedDevices( |
| 1300 | node, device_types_, &local_address_spec_)); |
| 1301 | |
| 1302 | if (node.has_assigned_device_name()) { |
| 1303 | TF_RETURN_IF_ERROR(InitializeMemberWithAssignedDevice( |
| 1304 | node.assigned_device_name(), node.type_string(), member)); |
| 1305 | } else { |
| 1306 | // This node has not yet been assigned to a device, so we |
| 1307 | // calculate any constraints due to the set of registered |
| 1308 | // kernels and any (partial) user-provided device specification |
| 1309 | // in the NodeDef. |
| 1310 | |
| 1311 | // If no kernels are registered for this op type, fail with an error. |
| 1312 | if (member->supported_device_types().empty()) { |
| 1313 | std::set<string> registered_device_types; |
| 1314 | for (Device* d : device_set_.devices()) { |
| 1315 | registered_device_types.insert(d->device_type()); |
| 1316 | } |
| 1317 | return errors::InvalidArgument( |
| 1318 | "No OpKernel was registered to support Op '", node.type_string(), |
| 1319 | "' used by ", errors::FormatNodeNameForError(node.name()), |
| 1320 | "with these attrs: [", node.attrs().DebugString(), |
| 1321 | "]\n" |
| 1322 | "Registered devices: [", |
| 1323 | absl::StrJoin(registered_device_types, ", "), "]\n", |
| 1324 | "Registered kernels:\n", KernelsRegisteredForOp(node.type_string())); |
| 1325 | } |
| 1326 | |
| 1327 | // If the NodeDef contains a device, then we interpret it as a |
| 1328 | // (partial) device specification. |
| 1329 | if (!node.requested_device().empty()) { |
| 1330 | if (IsRefOrResourceGeneratorNode(node)) { |
| 1331 | // Treat requested device on resource generating nodes as assigned |
| 1332 | // device so that we don't override it. |
| 1333 | TF_RETURN_IF_ERROR(member->SetResourceDeviceName(node)); |
| 1334 | } else { |
| 1335 | // The user has specified a device in the NodeDef, try to find a |
| 1336 | // valid device matching their specification in the set of |
| 1337 | // devices. |
| 1338 | // NOTE: The full name may specify a device that is not in |
| 1339 | // n.supported_device_types(), but we check that in AssignDevice(). |
| 1340 | TF_RETURN_IF_ERROR(member->SetRequestedDeviceName(node)); |
| 1341 | } |
| 1342 | } |
| 1343 | } |
| 1344 | |
| 1345 | // set gpu stream idx, default value is -1, which means multi-stream is disable. |
| 1346 | const AttrValue *gpu_stream_idx_attr = node.attrs().Find("_stream_id"); |
| 1347 | if (gpu_stream_idx_attr != nullptr) { |
| 1348 | int gpu_stream_idx = gpu_stream_idx_attr->i(); |
| 1349 | member->SetGpuStreamIdx(gpu_stream_idx); |
| 1350 | } |
| 1351 | |
| 1352 | return Status::OK(); |
| 1353 | } |
| 1354 | |
| 1355 | // Returns a list of devices having type in supported_device_types. The |
nothing calls this directly
no test coverage detected