tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
// Checks if a node's input port is Host friendly.
// Roughly this means checking if the input port is on Host memory.
Checks if a node's input port is Host friendly.
Roughly this means checking if the input port is on Host memory.
Checks if a node ' s input port is Host friendly . Roughly this means checking if the input port is on Host memory .
bool IsNodeInputPortHostFriendly(const NodeDef& node, int port_id) {
  if (absl::StrContains(node.device(), DEVICE_CPU)) {
    return true;
  }
  const OpDef* op = nullptr;
  Status s = OpRegistry::Global()->LookUpOpDef(node.op(), &op);
  if (!s.ok()) {
    LOG(WARNING) << "Could not find OpDef for : " << node.op();
    return false;
  }
  const int input_arg_id = OpInputPortIdToArgId(node, *op, port_id);
  const KernelDef* kernel = nullptr;
  s = internal::TryFindKernelDef(
      {node.device().c_str(), DEVICE_GPU, DEVICE_CPU}, node, &kernel);
  if (!s.ok()) {
    LOG(INFO) << "Could not find KernelDef for: " << node.op();
    return false;
  }
  for (const string& host_memory_arg : kernel->host_memory_arg()) {
    if (op->input_arg(input_arg_id).name() == host_memory_arg) {
      return true;
    }
  }
  return false;
}
bool IsNodeInputPortHostFriendly ( const NodeDef & node , int port_id ) { if ( absl :: StrContains ( node . device ( ) , DEVICE_CPU ) ) { return true ; } const OpDef * op = nullptr ; Status s = OpRegistry :: Global ( ) -> LookUpOpDef ( node . op ( ) , & op ) ; if ( ! s . ok ( ) ) { LOG ( WARNING ) << " " << node . op ( ) ; return false ; } const int input_arg_id = OpInputPortIdToArgId ( node , * op , port_id ) ; const KernelDef * kernel = nullptr ; s = internal :: TryFindKernelDef ( { node . device ( ) . c_str ( ) , DEVICE_GPU , DEVICE_CPU } , node , & kernel ) ; if ( ! s . ok ( ) ) { LOG ( INFO ) << " " << node . op ( ) ; return false ; } for ( const string & host_memory_arg : kernel -> host_memory_arg ( ) ) { if ( op -> input_arg ( input_arg_id ) . name ( ) == host_memory_arg ) { return true ; } } return false ; }
Checks if a node's input port is Host friendly.
Checks if a node ' s input port is Host friendly .
// If node is on Host, assume its inputs are Host friendly. // Check if op's input port is pinned to HostMemory. // Find the kernel. // Check if the input_arg is pinned to Host.