| 352 | } |
| 353 | |
| 354 | std::vector<NodeDef*> AddFuseRecv(const PartitionOptions& opts, |
| 355 | const GraphInfo& g_info, GraphDef* gdef, |
| 356 | const std::vector<const Edge*> edges, |
| 357 | NodeDef** real_recv, Status* status) { |
| 358 | int fuse_count = edges.size(); |
| 359 | |
| 360 | std::vector<DataType> dtypes(fuse_count); |
| 361 | std::vector<const Node*> srcs(fuse_count); |
| 362 | std::vector<const Node*> dsts(fuse_count); |
| 363 | std::vector<DataType> cast_dtypes(fuse_count); |
| 364 | |
| 365 | bool host_memory; |
| 366 | int host_memory_count = 0, device_memory_count = 0; |
| 367 | for (int i = 0; i < fuse_count; ++i) { |
| 368 | dtypes[i] = EdgeType(edges[i]); |
| 369 | srcs[i] = edges[i]->src(); |
| 370 | dsts[i] = edges[i]->dst(); |
| 371 | cast_dtypes[i] = dtypes[i]; |
| 372 | |
| 373 | // NOTE(yuanbyu): Only cast for cross-device send/recv. |
| 374 | if (opts.should_cast && !NeedSameDeviceSendRecv(edges[i], g_info)) { |
| 375 | cast_dtypes[i] = opts.should_cast(edges[i]); |
| 376 | } |
| 377 | |
| 378 | host_memory = false; |
| 379 | if (!edges[i]->IsControlEdge()) { |
| 380 | auto dst_it = g_info.input_types.find({dsts[i]->id(), edges[i]->dst_input()}); |
| 381 | DCHECK(dst_it != g_info.input_types.end()); |
| 382 | host_memory = (dst_it->second == HOST_MEMORY); |
| 383 | } |
| 384 | if (host_memory) { |
| 385 | ++host_memory_count; |
| 386 | } else { |
| 387 | ++device_memory_count; |
| 388 | } |
| 389 | } |
| 390 | |
| 391 | host_memory = false; |
| 392 | if (host_memory_count > device_memory_count) host_memory = true; |
| 393 | |
| 394 | // Add the fuse recv node. |
| 395 | const string fuse_recv_op = (host_memory) ? "_HostFuseRecv" : "_FuseRecv"; |
| 396 | string fuse_recv_node_name; |
| 397 | // TODO: name -> use global counter |
| 398 | for (int i = 0; i < fuse_count; ++i) { |
| 399 | if (i != 0) { |
| 400 | fuse_recv_node_name += "__"; |
| 401 | } |
| 402 | fuse_recv_node_name += srcs[i]->name(); |
| 403 | } |
| 404 | |
| 405 | NodeDefBuilder fuse_recv_builder(opts.new_name(fuse_recv_node_name), |
| 406 | fuse_recv_op); |
| 407 | SetFuseRecvAttrs(opts, edges, &fuse_recv_builder); |
| 408 | fuse_recv_builder.Device(dsts[0]->assigned_device_name()) |
| 409 | .Attr("tensor_types", cast_dtypes); |
| 410 | |
| 411 | *real_recv = gdef->add_node(); |
no test coverage detected