tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
// Checks if a node's output port is host friendly.
// Roughly this means checking if the output port is on Host memory.
Checks if a node's output port is host friendly.
Roughly this means checking if the output port is on Host memory.
Checks if a node ' s output port is host friendly . Roughly this means checking if the output port is on Host memory .
Status IsNodeOutputPortHostFriendly(const GraphView& graph,
                                    GraphProperties* properties,
                                    const NodeDef& node, int port_id,
                                    bool* is_candidate) {
  *is_candidate = false;
  if (IsDenylisted(node)) {
    return Status::OK();
  }
  if (!properties->has_properties()) {
    TF_RETURN_IF_ERROR(properties->InferStatically(
        false, false,
        false));
  }
  const auto& output_properties = properties->GetOutputProperties(node.name());
  int output_properties_size = output_properties.size();
  if (port_id >= output_properties_size) {
    LOG(WARNING) << "port_id=" << port_id
                 << " but output_properties.size()=" << output_properties.size()
                 << "\n"
                 << node.DebugString();
    return Status::OK();
  }
  if (!IsTensorSmall(output_properties[port_id])) {
    return Status::OK();
  }
  if (IsIdentity(node) || IsIdentityNSingleInput(node)) {
    for (const auto& fanin : graph.GetFanins(node, false)) {
      bool fanin_candidate = false;
      TF_RETURN_IF_ERROR(IsNodeOutputPortHostFriendly(
          graph, properties, *fanin.node, fanin.port_id, &fanin_candidate));
      if (!fanin_candidate) {
        return Status::OK();
      }
    }
    *is_candidate = true;
    return Status::OK();
  }
  if (absl::StrContains(node.device(), DEVICE_CPU)) {
    *is_candidate = true;
    return Status::OK();
  }
  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 Status::OK();
  }
  const int output_arg_id = OpOutputPortIdToArgId(node, *op, port_id);
  if (output_arg_id < 0) {
    LOG(WARNING) << "Invalid port: " << port_id << "!\n"
                 << node.DebugString() << "\n"
                 << op->DebugString();
    return Status::OK();
  }
  const KernelDef* kernel = nullptr;
  s = TryFindKernelDef({node.device().c_str(), DEVICE_GPU, DEVICE_CPU}, node,
                       &kernel);
  if (!s.ok()) {
    LOG(INFO) << "Could not find KernelDef for: " << node.op();
    return Status::OK();
  }
  for (const string& host_memory_arg : kernel->host_memory_arg()) {
    if (op->output_arg(output_arg_id).name() == host_memory_arg) {
      *is_candidate = true;
      break;
    }
  }
  return Status::OK();
}
Status IsNodeOutputPortHostFriendly ( const GraphView & graph , GraphProperties * properties , const NodeDef & node , int port_id , bool * is_candidate ) { * is_candidate = false ; if ( IsDenylisted ( node ) ) { return Status :: OK ( ) ; } if ( ! properties -> has_properties ( ) ) { TF_RETURN_IF_ERROR ( properties -> InferStatically ( false , false , false ) ) ; } const auto & output_properties = properties -> GetOutputProperties ( node . name ( ) ) ; int output_properties_size = output_properties . size ( ) ; if ( port_id >= output_properties_size ) { LOG ( WARNING ) << " " << port_id << " " << output_properties . size ( ) << " \n " << node . DebugString ( ) ; return Status :: OK ( ) ; } if ( ! IsTensorSmall ( output_properties [ port_id ] ) ) { return Status :: OK ( ) ; } if ( IsIdentity ( node ) || IsIdentityNSingleInput ( node ) ) { for ( const auto & fanin : graph . GetFanins ( node , false ) ) { bool fanin_candidate = false ; TF_RETURN_IF_ERROR ( IsNodeOutputPortHostFriendly ( graph , properties , * fanin . node , fanin . port_id , & fanin_candidate ) ) ; if ( ! fanin_candidate ) { return Status :: OK ( ) ; } } * is_candidate = true ; return Status :: OK ( ) ; } if ( absl :: StrContains ( node . device ( ) , DEVICE_CPU ) ) { * is_candidate = true ; return Status :: OK ( ) ; } const OpDef * op = nullptr ; Status s = OpRegistry :: Global ( ) -> LookUpOpDef ( node . op ( ) , & op ) ; if ( ! s . ok ( ) ) { LOG ( WARNING ) << " " << node . op ( ) ; return Status :: OK ( ) ; } const int output_arg_id = OpOutputPortIdToArgId ( node , * op , port_id ) ; if ( output_arg_id < 0 ) { LOG ( WARNING ) << " " << port_id << " \n " << node . DebugString ( ) << " \n " << op -> DebugString ( ) ; return Status :: OK ( ) ; } const KernelDef * kernel = nullptr ; s = TryFindKernelDef ( { node . device ( ) . c_str ( ) , DEVICE_GPU , DEVICE_CPU } , node , & kernel ) ; if ( ! s . ok ( ) ) { LOG ( INFO ) << " " << node . op ( ) ; return Status :: OK ( ) ; } for ( const string & host_memory_arg : kernel -> host_memory_arg ( ) ) { if ( op -> output_arg ( output_arg_id ) . name ( ) == host_memory_arg ) { * is_candidate = true ; break ; } } return Status :: OK ( ) ; }
Checks if a node's output port is host friendly.
Checks if a node ' s output port is host friendly .
// Make sure we are not a denylisted op. // Check to make sure we have the right properties (i.e., statically shaped). // This is an expensive call, call it lazily. /*assume_valid_feeds=*/ /*aggressive_shape_inference=*/ /*include_tensor_values=*/ // These nodes may be optimized away downstream (even if pinned to Host), we // should (recursively) check their source. // Check if op's device is on CPU. // Check if op's output port is pinned to HostMemory. // Map the port_id to output_arg_id. // Find the kernel. // Check if the output_arg is pinned to Host.