Add _HostSend and _HostRecv to grappler

PiperOrigin-RevId: 181774069

diff --git a/tensorflow/core/grappler/costs/virtual_scheduler.cc b/tensorflow/core/grappler/costs/virtual_scheduler.cc
index 29b401304df13f28659fcd6a9c0f267c7aebdf10..0af889f886220b1f2114aa6c591d8e19937d4eb3 100644 (file)
--- a/tensorflow/core/grappler/costs/virtual_scheduler.cc
+++ b/tensorflow/core/grappler/costs/virtual_scheduler.cc
@@ -341,7 +341,7 @@ Status VirtualScheduler::Init() {
   // to _Recv as control dependency when creating GrapplerItem.
   std::unordered_map<string, const NodeDef*> name_to_send;
   for (const auto& node : graph.node()) {
-    if (node.op() == "_Send") {
+    if (IsSend(node)) {
       const auto& attr = node.attr();
       name_to_send[attr.at("tensor_name").s()] = &node;
     }

diff --git a/tensorflow/core/grappler/op_types.cc b/tensorflow/core/grappler/op_types.cc
index 29e8e29ace0ad6b4d23be598fbbcf2fa0a477601..19d356cb1771ccd8a57dccaf2eb7bdc4823f7fcf 100644 (file)
--- a/tensorflow/core/grappler/op_types.cc
+++ b/tensorflow/core/grappler/op_types.cc
@@ -235,7 +235,9 @@ bool IsReciprocalGrad(const NodeDef& node) {
   return node.op() == "ReciprocalGrad";
 }
 
-bool IsRecv(const NodeDef& node) { return node.op() == "_Recv"; }
+bool IsRecv(const NodeDef& node) {
+  return node.op() == "_Recv" || node.op() == "_HostRecv";
+}
 
 bool IsReduction(const NodeDef& node) {
   const auto& op = node.op();
@@ -262,7 +264,9 @@ bool IsSelect(const NodeDef& node) { return node.op() == "Select"; }
 
 bool IsSeluGrad(const NodeDef& node) { return node.op() == "SeluGrad"; }
 
-bool IsSend(const NodeDef& node) { return node.op() == "_Send"; }
+bool IsSend(const NodeDef& node) {
+  return node.op() == "_Send" || node.op() == "_HostSend";
+}
 
 bool IsShape(const NodeDef& node) { return node.op() == "Shape"; }