| 335 | |
| 336 | namespace { |
| 337 | Status InferAllocAttr(const Node* n, const Node* dst, |
| 338 | const DeviceNameUtils::ParsedName& local_dev_name, |
| 339 | AllocatorAttributes* attr) { |
| 340 | Status s; |
| 341 | // Note that it's possible for *n to be a Recv and *dst to be a Send, |
| 342 | // so these two cases are not mutually exclusive. |
| 343 | if (IsRecv(n)) { |
| 344 | string src_name; |
| 345 | s = GetNodeAttr(n->attrs(), "send_device", &src_name); |
| 346 | if (!s.ok()) return s; |
| 347 | DeviceNameUtils::ParsedName parsed_src_name; |
| 348 | if (!DeviceNameUtils::ParseFullName(src_name, &parsed_src_name)) { |
| 349 | s = errors::Internal("Bad send_device attr '", src_name, "' in node ", |
| 350 | n->name()); |
| 351 | return s; |
| 352 | } |
| 353 | if (!DeviceNameUtils::IsSameAddressSpace(parsed_src_name, local_dev_name)) { |
| 354 | // Value is going to be the sink of an RPC. |
| 355 | attr->set_nic_compatible(true); |
| 356 | VLOG(2) << "node " << n->name() << " is the sink of an RPC in"; |
| 357 | } else if ((local_dev_name.type == "CPU" || n->IsHostRecv()) && |
| 358 | parsed_src_name.type != "CPU") { |
| 359 | // Value is going to be the sink of a local DMA from GPU to CPU (or |
| 360 | // other types of accelerators). |
| 361 | attr->set_gpu_compatible(true); |
| 362 | VLOG(2) << "node " << n->name() << " is the sink of a gpu->cpu copy"; |
| 363 | } else { |
| 364 | VLOG(2) << "default alloc case local type " << local_dev_name.type |
| 365 | << " remote type " << parsed_src_name.type; |
| 366 | } |
| 367 | } |
| 368 | if (IsSend(dst)) { |
| 369 | string dst_name; |
| 370 | s = GetNodeAttr(dst->attrs(), "recv_device", &dst_name); |
| 371 | if (!s.ok()) return s; |
| 372 | DeviceNameUtils::ParsedName parsed_dst_name; |
| 373 | if (!DeviceNameUtils::ParseFullName(dst_name, &parsed_dst_name)) { |
| 374 | s = errors::Internal("Bad recv_device attr '", dst_name, "' in node ", |
| 375 | n->name()); |
| 376 | return s; |
| 377 | } |
| 378 | if (!DeviceNameUtils::IsSameAddressSpace(parsed_dst_name, local_dev_name)) { |
| 379 | // Value is going to be the source of an RPC. |
| 380 | attr->set_nic_compatible(true); |
| 381 | VLOG(2) << "node " << n->name() << " is the source of an RPC out"; |
| 382 | } else if ((local_dev_name.type == "CPU" || dst->IsHostSend()) && |
| 383 | parsed_dst_name.type != "CPU") { |
| 384 | // Value is going to be the source of a local DMA from CPU to GPU (or |
| 385 | // other types of accelerators). |
| 386 | // Note that this does not cover the case where the allocation of the |
| 387 | // output tensor is not generated by the src: n. |
| 388 | attr->set_gpu_compatible(true); |
| 389 | VLOG(2) << "node " << n->name() << " is the source of a cpu->gpu copy"; |
| 390 | } else { |
| 391 | VLOG(2) << "default alloc case local type " << local_dev_name.type |
| 392 | << " remote type " << parsed_dst_name.type; |
| 393 | } |
| 394 | } |
no test coverage detected