[Graph] Avoid calling Node::GetNodeClassForOp() when a node is copied.

Instead, copy the class from the original node. This change also modifies the `kNodeClassTable` to use `absl::flat_hash_map<>`.

PiperOrigin-RevId: 306945263
Change-Id: I8eb1c80b57fdf204fbc7072a55615dd688025e87

tensorflow/core/graph/graph.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/framework/node_properties.h"
@@ -38,11 +39,11 @@ namespace tensorflow {
 const int Graph::kControlSlot = -1;
 
 // Node
+Node::NodeClass Node::GetNodeClassForOp(const string& ts) {
+  static const absl::flat_hash_map<string, Node::NodeClass>* kNodeClassTable =
 #define REF_CLASS(key, value) \
   {key, value}, { "Ref" key, value }
-
+      new absl::flat_hash_map<string, Node::NodeClass>({
-const std::unordered_map<string, Node::NodeClass>& Node::kNodeClassTable =
-    *new std::unordered_map<string, Node::NodeClass>({
           // Keep in same order as NodeClass values
           REF_CLASS("Switch", NC_SWITCH),
           REF_CLASS("_SwitchN", NC_SWITCH),
@@ -91,12 +92,10 @@ const std::unordered_map<string, Node::NodeClass>& Node::kNodeClassTable =
           {"_DeviceRetval", NC_RETVAL},
           {"_XlaMerge", NC_MERGE},
       });
-
 #undef REF_CLASS
 
-Node::NodeClass Node::GetNodeClassForOp(const string& ts) {
+  auto it = kNodeClassTable->find(ts);
-  auto it = kNodeClassTable.find(ts);
+  if (it != kNodeClassTable->end()) {
-  if (it != kNodeClassTable.end()) {
     return it->second;
   } else {
     return NC_OTHER;
@@ -127,7 +126,7 @@ Node::Node()
 
 void Node::Initialize(int id, int cost_id,
                       std::shared_ptr<NodeProperties> props,
-                      bool is_function_op) {
+                      Node::NodeClass node_class) {
   DCHECK_EQ(id_, -1);
   DCHECK(in_edges_.empty());
   DCHECK(out_edges_.empty());
@@ -135,12 +134,7 @@ void Node::Initialize(int id, int cost_id,
   cost_id_ = cost_id;
 
   props_ = std::move(props);
-  // Initialize the class_ based on the type string
+  class_ = node_class;
-  if (is_function_op) {
-    class_ = NC_FUNCTION_OP;
-  } else {
-    class_ = GetNodeClassForOp(props_->node_def.op());
-  }
 }
 
 void Node::Clear() {
@@ -423,18 +417,21 @@ Node* Graph::AddNode(NodeDef node_def, Status* status) {
     return nullptr;
   }
 
+  Node::NodeClass node_class = op_reg_data->is_function_op
+                                   ? Node::NC_FUNCTION_OP
+                                   : Node::GetNodeClassForOp(node_def.op());
+
   Node* node = AllocateNode(
       std::make_shared<NodeProperties>(&op_reg_data->op_def,
                                        std::move(node_def), inputs, outputs),
-      nullptr, op_reg_data->is_function_op);
+      nullptr, node_class);
   return node;
 }
 
 Node* Graph::CopyNode(const Node* node) {
   DCHECK(!node->IsSource());
   DCHECK(!node->IsSink());
-  Node* copy =
+  Node* copy = AllocateNode(node->props_, node, node->class_);
-      AllocateNode(node->props_, node, node->class_ == Node::NC_FUNCTION_OP);
   copy->set_assigned_device_name(node->assigned_device_name());
 
   // Since the OpDef of a function may be owned by the Graph that owns 'node',
@@ -759,7 +756,7 @@ Status Graph::IsValidInputTensor(const Node* node, int idx) const {
 }
 
 Node* Graph::AllocateNode(std::shared_ptr<NodeProperties> props,
-                          const Node* cost_node, bool is_function_op) {
+                          const Node* cost_node, Node::NodeClass node_class) {
   Node* node = nullptr;
   if (free_nodes_.empty()) {
     node = new (arena_.Alloc(sizeof(Node))) Node;  // placement new
@@ -770,7 +767,7 @@ Node* Graph::AllocateNode(std::shared_ptr<NodeProperties> props,
   node->graph_ = this;
   const int id = nodes_.size();
   int cost_id = cost_node ? cost_node->cost_id() : id;
-  node->Initialize(id, cost_id, std::move(props), is_function_op);
+  node->Initialize(id, cost_id, std::move(props), node_class);
   nodes_.push_back(node);
   ++num_nodes_;
   return node;

tensorflow/core/graph/graph.h

@@ -236,10 +236,6 @@ class Node {
   friend class Graph;
   Node();
 
-
-  void Initialize(int id, int cost_id, std::shared_ptr<NodeProperties> props,
-                  bool is_function_op);
-
   // Releases memory from props_, in addition to restoring *this to its
   // uninitialized state.
   void Clear();
@@ -291,7 +287,8 @@ class Node {
     NC_OTHER  // Not a special kind of node
   };
 
-  static const std::unordered_map<string, NodeClass>& kNodeClassTable;
+  void Initialize(int id, int cost_id, std::shared_ptr<NodeProperties> props,
+                  NodeClass node_class);
 
   static NodeClass GetNodeClassForOp(const string& ts);
 
@@ -692,7 +689,7 @@ class Graph {
   //
   // Ownership of the returned Node is not transferred to caller.
   Node* AllocateNode(std::shared_ptr<NodeProperties> props,
-                     const Node* cost_node, bool is_function_op);
+                     const Node* cost_node, Node::NodeClass node_class);
   void ReleaseNode(Node* node);
   // Insert edge in free_edges_ for possible reuse.
   void RecycleEdge(const Edge* edge);