Add a config field to GraphNode to allow for custom configuration of …

…nodes

Users can set `pinToGroupTop` for a node to be pinned to the top of a layer.

PiperOrigin-RevId: 677915322

src/builtin-adapter/formats/schema_structs.cc

@@ -64,6 +64,14 @@ llvm::json::Object GraphEdge::Json() {
   return json_edge;
 }
 
+const char GraphNodeConfig::kPinToGroupTop[] = "pinToGroupTop";
+
+llvm::json::Object GraphNodeConfig::Json() {
+  llvm::json::Object json_config;
+  json_config[kPinToGroupTop] = pin_to_group_top;
+  return json_config;
+}
+
 const char GraphNode::kNodeId[] = "id";
 const char GraphNode::kNodeLabel[] = "label";
 const char GraphNode::kNodeName[] = "namespace";
@@ -72,6 +80,7 @@ const char GraphNode::kNodeAttrs[] = "attrs";
 const char GraphNode::kIncomingEdges[] = "incomingEdges";
 const char GraphNode::kInputsMetadata[] = "inputsMetadata";
 const char GraphNode::kOutputsMetadata[] = "outputsMetadata";
+const char GraphNode::kConfig[] = "config";
 
 llvm::json::Object GraphNode::Json() {
   llvm::json::Object json_node;
@@ -105,6 +114,11 @@ llvm::json::Object GraphNode::Json() {
   for (Metadata& metadata : outputs_metadata) {
     json_outputs_metadata->push_back(metadata.Json());
   }
+
+  if (config.has_value()) {  // Only add config if it exists
+    json_node[kConfig] = config->Json();
+  }
+
   return json_node;
 }
 

src/builtin-adapter/formats/schema_structs.h

@@ -16,6 +16,7 @@
 #ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_GOOGLE_TOOLING_FORMATS_SCHEMA_STRUCTS_H_
 #define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_GOOGLE_TOOLING_FORMATS_SCHEMA_STRUCTS_H_
 
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -64,6 +65,17 @@ struct GraphEdge {
   static const char kEdgeMetadata[];
 };
 
+// Configuration for a graph node.
+struct GraphNodeConfig {
+  // Whether to pin the node to the top of the group it belongs to.
+  bool pin_to_group_top = false;
+
+  llvm::json::Object Json();
+
+ private:
+  static const char kPinToGroupTop[];
+};
+
 struct GraphNode {
   std::string node_id;
   std::string node_label;
@@ -73,6 +85,7 @@ struct GraphNode {
   std::vector<GraphEdge> incoming_edges;
   std::vector<Metadata> inputs_metadata;
   std::vector<Metadata> outputs_metadata;
+  std::optional<GraphNodeConfig> config;
 
   llvm::json::Object Json();
 
@@ -85,6 +98,7 @@ struct GraphNode {
   static const char kIncomingEdges[];
   static const char kInputsMetadata[];
   static const char kOutputsMetadata[];
+  static const char kConfig[];
 };
 
 struct Subgraph {

src/builtin-adapter/graphnode_builder.cc

@@ -110,5 +110,13 @@ void GraphNodeBuilder::AppendAttrToMetadata(const EdgeType edge_type,
   }
 }
 
+void GraphNodeBuilder::SetPinToGroupTop(bool pin_to_group_top) {
+  if (node_.config.has_value()) {
+    node_.config->pin_to_group_top = pin_to_group_top;
+  } else {
+    node_.config = GraphNodeConfig{.pin_to_group_top = pin_to_group_top};
+  }
+}
+
 }  // namespace visualization_client
 }  // namespace tooling

src/builtin-adapter/graphnode_builder.h

@@ -77,6 +77,11 @@ class GraphNodeBuilder {
                             absl::string_view attr_key,
                             absl::string_view attr_value);
 
+  // Sets the node to be pinned to the top of the group it belongs to.
+  // User needs to ensure the node's namespace is indeed a layer for the pinning
+  // to work.
+  void SetPinToGroupTop(bool pin_to_group_top);
+
   // Returns the node that has been created by this class.
   GraphNode Build() && { return std::move(node_); }