Add blacklist ops to PinToHostOptimizer. Fix test.

PiperOrigin-RevId: 214195020

1 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
index 8a65cd3ec3..c8f9311b2e 100644
--- a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
@@ -156,6 +156,9 @@ string TryFindHostDevice(const gtl::FlatSet<string>& devices,
   // We couldn't find an appropriate Host device, return original device.
   return device;
 }
+
+// All the nodes that should be blacklisted and not swapped.
+bool IsBlacklisted(const NodeDef& node) { return IsCollective(node); }
 }  // end namespace internal
 
 Status PinToHostOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
@@ -185,6 +188,11 @@ Status PinToHostOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
       continue;
     }
 
+    // Skip these node types.
+    if (internal::IsBlacklisted(node)) {
+      continue;
+    }
+
     // Check the node can be run on CPU.
     Status s = FindKernelDef(DEVICE_CPU, node, nullptr, nullptr);
     if (!s.ok()) {