tensorflow/core/common_runtime/eval_const_tensor.cc
// Extracts the subgraph ending at 'target_node' that is statically computable
// and inserts into 'out_graph'. If statically computable, 'is_constant_graph'
// will be set to true.
Extracts the subgraph ending at 'target_node' that is statically computable
and inserts into 'out_graph'. If statically computable, 'is_constant_graph'
will be set to true.
Extracts the subgraph ending at ' target_node ' that is statically computable and inserts into ' out_graph ' . If statically computable ' is_constant_graph ' will be set to true .
Status ExtractConstantSubgraph(
    const Node& target_node, const ShapeRefiner& refiner,
    const std::unordered_map<string, Tensor>* cached_values, Graph* out_graph,
    bool* is_constant_graph,
    std::vector<std::pair<string, Tensor>>* const_inputs,
    InferenceContext* outer_context) {
  *is_constant_graph = false;
  std::unordered_set<string> const_inputs_added;
  if (target_node.op_def().is_stateful()) {
    return Status::OK();
  }
  if (IsMerge(&target_node)) {
    return Status::OK();
  }
  if (target_node.type_string() == "PlaceholderWithDefault") {
    return Status::OK();
  }
  if (!HasCpuKernel(target_node)) {
    return Status::OK();
  }
  struct NodeAndRecursed {
    Node* new_node = nullptr;
    bool recursed = false;
  };
  std::map<const Node*, NodeAndRecursed> old_to_new_and_recursed;
  Node* target_node_copy = out_graph->CopyNode(&target_node);
  old_to_new_and_recursed[&target_node].new_node = target_node_copy;
  old_to_new_and_recursed[&target_node].recursed = true;
  std::deque<const Edge*> edges_to_visit;
  for (const Edge* e : target_node.in_edges()) {
    if (e->IsControlEdge()) continue;
    edges_to_visit.push_back(e);
  }
  *is_constant_graph = true;
  while (!edges_to_visit.empty()) {
    const Edge* current_edge = edges_to_visit.front();
    edges_to_visit.pop_front();
    Node* current_node = current_edge->src();
    if (!current_node->IsArg() && current_node->op_def().is_stateful()) {
      *is_constant_graph = false;
      return Status::OK();
    }
    if (IsMerge(current_node)) {
      *is_constant_graph = false;
      return Status::OK();
    }
    if (IsEnter(current_node) || IsExit(current_node)) {
      *is_constant_graph = false;
      return Status::OK();
    }
    if (current_node->type_string() == "PlaceholderWithDefault") {
      *is_constant_graph = false;
      return Status::OK();
    }
    if (!HasCpuKernel(*current_node)) {
      *is_constant_graph = false;
      return Status::OK();
    }
    if (current_node->num_inputs() == 0) {
      if (outer_context && current_node->IsArg()) {
        const string& tensor_name =
            strings::StrCat(current_node->name(), ":", 0);
        if (const_inputs_added.count(tensor_name) == 0) {
          int index;
          TF_RETURN_IF_ERROR(GetArgNodeIndex(
              current_node, outer_context->num_inputs(), &index));
          const Tensor* const_tensor = outer_context->input_tensor(index);
          if (const_tensor) {
            const_inputs->emplace_back(tensor_name, *const_tensor);
            const_inputs_added.insert(tensor_name);
          } else {
            outer_context->request_input_tensor(index);
            *is_constant_graph = false;
            return Status::OK();
          }
        }
      } else if (!current_node->IsConstant()) {
        *is_constant_graph = false;
        return Status::OK();
      }
    }
     Add a copy of its node and a new edge to the new subgraph.
     Get or create the version of 'current_node' in the new graph.
    Node* current_node_copy;
     This gets or creates the NodeAndRecursed entry for current_node.
    NodeAndRecursed* node_and_recursed = &old_to_new_and_recursed[current_node];
    if (node_and_recursed->new_node == nullptr) {
       First time processing this node.
      current_node_copy = out_graph->CopyNode(current_node);
       Track the mapping from the original node to the new one.
      node_and_recursed->new_node = current_node_copy;
    } else {
      current_node_copy = node_and_recursed->new_node;
    }
     Add the edge to the destination node.
    {
      auto it = old_to_new_and_recursed.find(current_edge->dst());
      if (it == old_to_new_and_recursed.end()) {
        return errors::Internal(
            "Could not find mapping from old to new copy of destination node: ",
            current_edge->dst()->name());
      }
      Node* dst_copy = it->second.new_node;
      out_graph->AddEdge(current_node_copy, current_edge->src_output(),
                         dst_copy, current_edge->dst_input());
    }
    const string& output_tensor_name =
        strings::StrCat(current_node->name(), ":", current_edge->src_output());
     Some tensor values can be inferred. For example, a shape op
     with input shapes fully defined can have its output tensor inferred.
    Tensor tensor_inferred;
    bool successfully_inferred_tensor = false;
    TF_RETURN_IF_ERROR(TryToInferTensorOutputFromInputShapes(
        *current_edge, refiner, &tensor_inferred,
        &successfully_inferred_tensor));
    if (successfully_inferred_tensor) {
      const_inputs->emplace_back(output_tensor_name, tensor_inferred);
      const_inputs_added.insert(output_tensor_name);
      continue;
    }
     If we have a copy of the input tensor materialized already,
     then add to the list of inputs to feed and do not recurse further.
    if (cached_values != nullptr) {
      auto it = cached_values->find(output_tensor_name);
      if (it != cached_values->end() &&
          const_inputs_added.count(output_tensor_name) == 0) {
        const_inputs->emplace_back(output_tensor_name, it->second);
        const_inputs_added.insert(output_tensor_name);
        continue;
      }
    }
     If this node's inputs have not been processed already, do so now.
    if (!node_and_recursed->recursed) {
      node_and_recursed->recursed = true;
      for (const Edge* e : current_node->in_edges()) {
        if (e->IsControlEdge()) continue;
        edges_to_visit.push_back(e);
      }
    }
  }
  return Status::OK();
}
Status ExtractConstantSubgraph ( const Node & target_node , const ShapeRefiner & refiner , const std :: unordered_map < string , Tensor > * cached_values , Graph * out_graph , bool * is_constant_graph , std :: vector < std :: pair < string , Tensor > > * const_inputs , InferenceContext * outer_context ) { * is_constant_graph = false ; std :: unordered_set < string > const_inputs_added ; if ( target_node . op_def ( ) . is_stateful ( ) ) { return Status :: OK ( ) ; } if ( IsMerge ( & target_node ) ) { return Status :: OK ( ) ; } if ( target_node . type_string ( ) == " " ) { return Status :: OK ( ) ; } if ( ! HasCpuKernel ( target_node ) ) { return Status :: OK ( ) ; } struct NodeAndRecursed { Node * new_node = nullptr ; bool recursed = false ; } ; std :: map < const Node * , NodeAndRecursed > old_to_new_and_recursed ; Node * target_node_copy = out_graph -> CopyNode ( & target_node ) ; old_to_new_and_recursed [ & target_node ] . new_node = target_node_copy ; old_to_new_and_recursed [ & target_node ] . recursed = true ; std :: deque < const Edge * > edges_to_visit ; for ( const Edge * e : target_node . in_edges ( ) ) { if ( e -> IsControlEdge ( ) ) continue ; edges_to_visit . push_back ( e ) ; } * is_constant_graph = true ; while ( ! edges_to_visit . empty ( ) ) { const Edge * current_edge = edges_to_visit . front ( ) ; edges_to_visit . pop_front ( ) ; Node * current_node = current_edge -> src ( ) ; if ( ! current_node -> IsArg ( ) && current_node -> op_def ( ) . is_stateful ( ) ) { * is_constant_graph = false ; return Status :: OK ( ) ; } if ( IsMerge ( current_node ) ) { * is_constant_graph = false ; return Status :: OK ( ) ; } if ( IsEnter ( current_node ) || IsExit ( current_node ) ) { * is_constant_graph = false ; return Status :: OK ( ) ; } if ( current_node -> type_string ( ) == " " ) { * is_constant_graph = false ; return Status :: OK ( ) ; } if ( ! HasCpuKernel ( * current_node ) ) { * is_constant_graph = false ; return Status :: OK ( ) ; } if ( current_node -> num_inputs ( ) == 0 ) { if ( outer_context && current_node -> IsArg ( ) ) { const string & tensor_name = strings :: StrCat ( current_node -> name ( ) , " " , 0 ) ; if ( const_inputs_added . count ( tensor_name ) == 0 ) { int index ; TF_RETURN_IF_ERROR ( GetArgNodeIndex ( current_node , outer_context -> num_inputs ( ) , & index ) ) ; const Tensor * const_tensor = outer_context -> input_tensor ( index ) ; if ( const_tensor ) { const_inputs -> emplace_back ( tensor_name , * const_tensor ) ; const_inputs_added . insert ( tensor_name ) ; } else { outer_context -> request_input_tensor ( index ) ; * is_constant_graph = false ; return Status :: OK ( ) ; } } } else if ( ! current_node -> IsConstant ( ) ) { * is_constant_graph = false ; return Status :: OK ( ) ; } } Node * current_node_copy ; NodeAndRecursed * node_and_recursed = & old_to_new_and_recursed [ current_node ] ; if ( node_and_recursed -> new_node == nullptr ) { current_node_copy = out_graph -> CopyNode ( current_node ) ; node_and_recursed -> new_node = current_node_copy ; } else { current_node_copy = node_and_recursed -> new_node ; } { auto it = old_to_new_and_recursed . find ( current_edge -> dst ( ) ) ; if ( it == old_to_new_and_recursed . end ( ) ) { return errors :: Internal ( " " , current_edge -> dst ( ) -> name ( ) ) ; } Node * dst_copy = it -> second . new_node ; out_graph -> AddEdge ( current_node_copy , current_edge -> src_output ( ) , dst_copy , current_edge -> dst_input ( ) ) ; } const string & output_tensor_name = strings :: StrCat ( current_node -> name ( ) , " " , current_edge -> src_output ( ) ) ; Tensor tensor_inferred ; bool successfully_inferred_tensor = false ; TF_RETURN_IF_ERROR ( TryToInferTensorOutputFromInputShapes ( * current_edge , refiner , & tensor_inferred , & successfully_inferred_tensor ) ) ; if ( successfully_inferred_tensor ) { const_inputs -> emplace_back ( output_tensor_name , tensor_inferred ) ; const_inputs_added . insert ( output_tensor_name ) ; continue ; } if ( cached_values != nullptr ) { auto it = cached_values -> find ( output_tensor_name ) ; if ( it != cached_values -> end ( ) && const_inputs_added . count ( output_tensor_name ) == 0 ) { const_inputs -> emplace_back ( output_tensor_name , it -> second ) ; const_inputs_added . insert ( output_tensor_name ) ; continue ; } } if ( ! node_and_recursed -> recursed ) { node_and_recursed -> recursed = true ; for ( const Edge * e : current_node -> in_edges ( ) ) { if ( e -> IsControlEdge ( ) ) continue ; edges_to_visit . push_back ( e ) ; } } } return Status :: OK ( ) ; }
Extracts the subgraph ending at 'target_node' that is statically computable
and inserts into 'out_graph'.
Extracts the subgraph ending at ' target_node ' that is statically computable and inserts into ' out_graph ' .
// Since constant-folding runs on the CPU, do not attempt to constant-fold // operators that have no CPU kernel. // TODO(skyewm): should more of the filtering applied in input nodes below be // applied to target_node here? // Identify the possibly constant subgraph by recursively iterating backwards // through the inputs to 'target_node' until we either 1) find an already // existing input to our subgraph 'const_inputs', 2) Discover our graph is not // constant, or 3) Hit a root node. // Add the target node's inputs to seed the recursion. // TODO(skyewm): control edges will be meaningful if/when we handle control // flow (e.g. constants in cond branches are triggered via control edges). // Iterate over the set of edges to visit (backwards). // If the node is stateful, assume the graph is not constant unless it is // an Arg node which is handled later on. // During construction or import from GraphConstructor, back edges may not // be filled in. In addition, control flow constructs may depend on control // edges which aren't handled by this method. Don't constant fold through // merges at all for now. // Don't constant fold enter/exit currently either, as it's easy to end // up with a partial frame. // Placeholders should never be constant folded because their outputs are // fed by the user. Note that "Placeholder" nodes have no inputs so are // handled below. // If there is nothing more to recurse down, see if // the generator node is a constant or an Arg node whose value is available // in the `outer_context`. // If we do not already have a constant Tensor for this Arg try to // fetch it from the outer context. // Request a constant value for this Arg. If that is statically // computable, shape refiner will re-run the shape inference for // this function with this tensor's value. // Generator node is not a constant, so subgraph is not // constant. // Either the node is a constant, or the node is a potential // intermediate node on the path from a constant. // // Add a copy of its node and a new edge to the new subgraph. // Get or create the version of 'current_node' in the new graph. // This gets or creates the NodeAndRecursed entry for current_node. // First time processing this node. // Track the mapping from the original node to the new one. // Add the edge to the destination node. // Some tensor values can be inferred. For example, a shape op // with input shapes fully defined can have its output tensor inferred. // If we have a copy of the input tensor materialized already, // then add to the list of inputs to feed and do not recurse further. // If this node's inputs have not been processed already, do so now.