| 268 | } |
| 269 | |
| 270 | NodeDef* AddRecv(const PartitionOptions& opts, const GraphInfo& g_info, |
| 271 | GraphDef* gdef, const Edge* edge, NodeDef** real_recv, |
| 272 | Status* status) { |
| 273 | const DataType dtype = EdgeType(edge); |
| 274 | const Node* src = edge->src(); |
| 275 | const Node* dst = edge->dst(); |
| 276 | const int dst_port = edge->dst_input(); |
| 277 | DataType cast_dtype = dtype; |
| 278 | |
| 279 | // NOTE(yuanbyu): Only cast for cross-device send/recv. |
| 280 | if (opts.should_cast && !NeedSameDeviceSendRecv(edge, g_info)) { |
| 281 | cast_dtype = opts.should_cast(edge); |
| 282 | } |
| 283 | |
| 284 | // host_memory = true iff we need to use HostRecv/HostCast. |
| 285 | // Also log the introduction of the send-recv pair, for performance debugging. |
| 286 | bool host_memory = false; |
| 287 | if (!edge->IsControlEdge()) { |
| 288 | auto dst_it = g_info.input_types.find({dst->id(), dst_port}); |
| 289 | DCHECK(dst_it != g_info.input_types.end()); |
| 290 | host_memory = (dst_it->second == HOST_MEMORY); |
| 291 | bool src_host_memory = false; |
| 292 | if (VLOG_IS_ON(1)) { |
| 293 | const int src_port = edge->src_output(); |
| 294 | auto src_it = g_info.output_types.find({src->id(), src_port}); |
| 295 | DCHECK(src_it != g_info.output_types.end()); |
| 296 | src_host_memory = (src_it->second == HOST_MEMORY); |
| 297 | } |
| 298 | VLOG(1) << "Receiving data" |
| 299 | << " from " << src->name() << " (" << src->type_string() << ")" |
| 300 | << " on " << src->assigned_device_name() << " in " |
| 301 | << (src_host_memory ? "host memory" : "device memory") << " for " |
| 302 | << dst->name() << " (" << dst->type_string() << ")" |
| 303 | << " on " << dst->assigned_device_name() << " in " |
| 304 | << (host_memory ? "host memory" : "device memory"); |
| 305 | } else { |
| 306 | // Log control-edge transfers too, but don't mention memory space since it's |
| 307 | // irrelevant. |
| 308 | VLOG(1) << "Receiving control" |
| 309 | << " from " << src->name() << " (" << src->type_string() << ")" |
| 310 | << " on " << src->assigned_device_name() << " for " << dst->name() |
| 311 | << " (" << dst->type_string() << ")" |
| 312 | << " on " << dst->assigned_device_name(); |
| 313 | } |
| 314 | |
| 315 | // Add the recv node. |
| 316 | const string recv_op = (host_memory) ? "_HostRecv" : "_Recv"; |
| 317 | NodeDefBuilder recv_builder(opts.new_name(src->name()), recv_op, |
| 318 | NodeDebugInfo(*src)); |
| 319 | SetSendRecvAttrs(opts, edge, &recv_builder); |
| 320 | recv_builder.Device(dst->assigned_device_name()) |
| 321 | .Attr("tensor_type", cast_dtype); |
| 322 | NodeDef* recv = gdef->add_node(); |
| 323 | *status = recv_builder.Finalize(recv, /*consume=*/true); |
| 324 | if (!status->ok()) return nullptr; |
| 325 | *real_recv = recv; |
| 326 | |
| 327 | // Add the cast node (from cast_dtype to dtype) or an Identity node. |
no test coverage detected