| 210 | } |
| 211 | |
| 212 | NodeDef* AddSend(const PartitionOptions& opts, const GraphInfo& g_info, |
| 213 | GraphDef* gdef, const Edge* edge, |
| 214 | NodeDefBuilder::NodeOut send_from, int64 start_time, |
| 215 | Status* status) { |
| 216 | const DataType dtype = send_from.data_type; |
| 217 | const DataType cast_dtype = opts.should_cast ? opts.should_cast(edge) : dtype; |
| 218 | const Node* src = edge->src(); |
| 219 | const int src_port = edge->src_output(); |
| 220 | |
| 221 | // host_memory = true iff we need to use HostSend/HostCast. |
| 222 | bool host_memory = false; |
| 223 | if (!edge->IsControlEdge()) { |
| 224 | auto src_it = g_info.output_types.find({src->id(), src_port}); |
| 225 | DCHECK(src_it != g_info.output_types.end()); |
| 226 | host_memory = (src_it->second == HOST_MEMORY); |
| 227 | } |
| 228 | |
| 229 | // Add a cast node that casts dtype to cast_dtype. |
| 230 | // NOTE(yuanbyu): Only cast for cross-device send/recv. |
| 231 | if (dtype != cast_dtype && !NeedSameDeviceSendRecv(edge, g_info)) { |
| 232 | const string cast_op = (host_memory) ? "_HostCast" : "Cast"; |
| 233 | NodeDefBuilder cast_builder(opts.new_name(src->name()), cast_op, |
| 234 | NodeDebugInfo(*src)); |
| 235 | cast_builder.Device(src->assigned_device_name()).Input(send_from); |
| 236 | if (opts.scheduling_for_recvs) { |
| 237 | cast_builder.Attr("_start_time", start_time); |
| 238 | } |
| 239 | cast_builder.Attr("DstT", cast_dtype); |
| 240 | |
| 241 | if (cast_dtype == DT_BFLOAT16) { |
| 242 | // the below attribute specifies that the cast to bfloat16 should use |
| 243 | // truncation. This is needed to retain legacy behavior when we change |
| 244 | // the default bfloat16 casts to use rounding instead of truncation |
| 245 | cast_builder.Attr("Truncate", true); |
| 246 | } |
| 247 | |
| 248 | NodeDef* cast = gdef->add_node(); |
| 249 | *status = cast_builder.Finalize(cast, /*consume=*/true); |
| 250 | if (!status->ok()) return nullptr; |
| 251 | |
| 252 | // Connect the Send op to the cast. |
| 253 | send_from.Reset(cast->name(), 0, cast_dtype); |
| 254 | } |
| 255 | |
| 256 | // Add the send node. |
| 257 | const string send_op = (host_memory) ? "_HostSend" : "_Send"; |
| 258 | NodeDefBuilder send_builder(opts.new_name(src->name()), send_op, |
| 259 | NodeDebugInfo(*src)); |
| 260 | SetSendRecvAttrs(opts, edge, &send_builder); |
| 261 | send_builder.Device(src->assigned_device_name()).Input(send_from); |
| 262 | if (opts.scheduling_for_recvs) { |
| 263 | send_builder.Attr("_start_time", start_time); |
| 264 | } |
| 265 | NodeDef* send = gdef->add_node(); |
| 266 | *status = send_builder.Finalize(send, /*consume=*/true); |
| 267 | return send; |
| 268 | } |
| 269 |
no test coverage detected