Add display limit to the attribute string

The custom op attributes may contain strings that are too long to display, and that may crash the browser. Truncate the string to a reasonable length.

PiperOrigin-RevId: 693280782

src/builtin-adapter/translate_helpers.cc

@@ -97,7 +97,8 @@ inline std::string GetTypeString(const mlir::Type& t) {
   return result;
 }
 
-void AppendNodeAttrs(const int const_element_count_limit, Operation& operation,
+void AppendNodeAttrs(const int const_element_count_limit,
+                     const int attribute_string_limit, Operation& operation,
                      GraphNodeBuilder& builder) {
   std::string value;
   llvm::raw_string_ostream sstream(value);
@@ -122,6 +123,14 @@ void AppendNodeAttrs(const int const_element_count_limit, Operation& operation,
       // Special handles `value` attribute to represent the tensor data.
       builder.AppendNodeAttribute(kValue, value);
     } else {
+      // The custom op attributes may contain strings that are too long to
+      // display, and that may crash the browser. Truncate the string to a
+      // reasonable length.
+      if (attribute_string_limit >= 0 &&
+          value.length() > attribute_string_limit) {
+        value.resize(attribute_string_limit);
+        absl::StrAppend(&value, "... [truncated]");
+      }
       builder.AppendNodeAttribute(name, value);
     }
     value.clear();
@@ -431,7 +440,8 @@ absl::StatusOr<Subgraph> TfFunctionToSubgraph(const VisualizeConfig& config,
     llvm::StringRef node_name = GetTfNodeName(operation);
     GraphNodeBuilder builder;
     builder.SetNodeInfo(node_id, node_label, node_name);
-    AppendNodeAttrs(config.const_element_count_limit, operation, builder);
+    AppendNodeAttrs(config.const_element_count_limit,
+                    config.attribute_string_limit, operation, builder);
     for (int input_index = 0; input_index < operation.getNumOperands();
          ++input_index) {
       mlir::Value val = operation.getOperand(input_index);
@@ -483,7 +493,8 @@ absl::StatusOr<Subgraph> TfliteFunctionToSubgraph(const VisualizeConfig& config,
     std::string node_name = GenerateTfliteNodeName(node_label, operation);
     GraphNodeBuilder builder;
     builder.SetNodeInfo(node_id, node_label, node_name);
-    AppendNodeAttrs(config.const_element_count_limit, operation, builder);
+    AppendNodeAttrs(config.const_element_count_limit,
+                    config.attribute_string_limit, operation, builder);
     absl::Status append_subgraph_status =
         TfliteMaybeAppendSubgraphs(operation, builder);
     if (!append_subgraph_status.ok()) {
@@ -565,7 +576,8 @@ absl::StatusOr<Subgraph> StablehloFunctionToSubgraph(
     builder.SetNodeId(node_id);
     builder.SetNodeLabel(node_label);
     AddJaxNodeNameAndAttribute(operation, builder);
-    AppendNodeAttrs(config.const_element_count_limit, operation, builder);
+    AppendNodeAttrs(config.const_element_count_limit,
+                    config.attribute_string_limit, operation, builder);
     absl::Status append_subgraph_status =
         StablehloMaybeAppendSubgraphs(operation);
     if (!append_subgraph_status.ok()) {

src/builtin-adapter/visualize_config.h

@@ -22,12 +22,27 @@ namespace visualization_client {
 struct VisualizeConfig {
   VisualizeConfig() = default;
   explicit VisualizeConfig(const int const_element_count_limit)
-      : const_element_count_limit(const_element_count_limit) {}
+      : const_element_count_limit(const_element_count_limit) {
+    // If const_element_count_limit is set to -1, then we should also set
+    // attribute_string_limit to -1 to print all attributes.
+    if (const_element_count_limit == -1) {
+      attribute_string_limit = -1;
+    }
+  }
+  VisualizeConfig(const int const_element_count_limit,
+                  const int attribute_string_limit)
+      : const_element_count_limit(const_element_count_limit),
+        attribute_string_limit(attribute_string_limit) {}
 
   // The maximum number of constant elements to be displayed. If the number
   // exceeds this threshold, the rest of data will be elided. The default
   // threshold is set to 16 (use -1 to print all).
   int const_element_count_limit = 16;
+
+  // The maximum length of attribute string to be displayed. If the length
+  // exceeds this threshold, the rest of data will be elided. The default
+  // threshold is set to 1024 (use -1 to print all).
+  int attribute_string_limit = 1024;
 };
 
 }  // namespace visualization_client