Add quant_params_count_limit to VisualizeConfig

PiperOrigin-RevId: 633054360

src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc

@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "direct_flatbuffer_to_json_graph_convert.h"
 
+#include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <cstring>
 #include <memory>
@@ -649,6 +651,7 @@ absl::Status AddTensorTags(const OperatorT& op, absl::string_view op_label,
 }
 
 void AddQuantizationParameters(const std::unique_ptr<TensorT>& tensor,
+                               const size_t size_limit,
                                const EdgeType edge_type, const int rel_idx,
                                GraphNodeBuilder& builder) {
   if (tensor->quantization == nullptr) return;
@@ -662,9 +665,13 @@ void AddQuantizationParameters(const std::unique_ptr<TensorT>& tensor,
   }
   if (quant->scale.empty()) return;
 
+  const unsigned num_params = (size_limit < 0)
+                                  ? quant->scale.size()
+                                  : std::min(quant->scale.size(), size_limit);
+  if (num_params == 0) return;
   std::vector<std::string> parameters;
-  parameters.reserve(quant->scale.size());
-  for (int i = 0; i < quant->scale.size(); ++i) {
+  parameters.reserve(num_params);
+  for (int i = 0; i < num_params; ++i) {
     // Parameters will be shown as "[scale] * (q + [zero_point])"
     parameters.push_back(
         absl::StrCat(quant->scale[i], " * (q + ", quant->zero_point[i], ")"));
@@ -680,7 +687,7 @@ absl::Status AddNode(
     const Buffers& buffers, const std::vector<std::string>& func_names,
     const std::optional<const SignatureNameMap>& signature_name_map,
     const OpdefsMap& op_defs, const std::unique_ptr<FlatBufferModel>& model_ptr,
-    const int const_element_count_limit, std::vector<std::string>& node_ids,
+    const VisualizeConfig& config, std::vector<std::string>& node_ids,
     EdgeMap& edge_map, mlir::Builder mlir_builder, Subgraph& subgraph) {
   if (op.opcode_index >= op_names.size()) {
     return absl::InvalidArgumentError(
@@ -713,14 +720,16 @@ absl::Status AddNode(
     // when the input tensor is constant and not an output of a node. Thus we
     // create an auxiliary constant node to align with graph structure.
     if (EdgeInfoIncomplete(edge_map.at(tensor_index))) {
-      RETURN_IF_ERROR(AddAuxiliaryNode(
-          NodeType::kConstNode, std::vector<int>{tensor_index}, tensors,
-          buffers, signature_name_map, model_ptr, const_element_count_limit,
-          node_ids, edge_map, mlir_builder, subgraph));
+      RETURN_IF_ERROR(
+          AddAuxiliaryNode(NodeType::kConstNode, std::vector<int>{tensor_index},
+                           tensors, buffers, signature_name_map, model_ptr,
+                           config.const_element_count_limit, node_ids, edge_map,
+                           mlir_builder, subgraph));
     }
     AppendIncomingEdge(edge_map.at(tensor_index), builder);
-    AddQuantizationParameters(tensors[tensor_index], EdgeType::kInput, i,
-                              builder);
+    AddQuantizationParameters(tensors[tensor_index],
+                              config.quant_params_count_limit, EdgeType::kInput,
+                              i, builder);
   }
 
   for (int i = 0; i < op.outputs.size(); ++i) {
@@ -732,8 +741,9 @@ absl::Status AddNode(
                       .source_node_output_id = absl::StrCat(i)},
                      edge_map);
 
-    AddQuantizationParameters(tensors[tensor_index], EdgeType::kOutput, i,
-                              builder);
+    AddQuantizationParameters(tensors[tensor_index],
+                              config.quant_params_count_limit,
+                              EdgeType::kOutput, i, builder);
   }
 
   status = AddTensorTags(op, node_label, op_defs, builder);
@@ -802,8 +812,8 @@ absl::Status AddSubgraph(
     const Tensors& tensors = subgraph_t.tensors;
     RETURN_IF_ERROR(AddNode(i, *op, op_codes, op_names, tensors, buffers,
                             func_names, signature_name_map, op_defs, model_ptr,
-                            config.const_element_count_limit, node_ids,
-                            edge_map, mlir_builder, subgraph));
+                            config, node_ids, edge_map, mlir_builder,
+                            subgraph));
   }
 
   // Adds GraphOutputs node to the subgraph.

src/builtin-adapter/models_to_json_main.cc

@@ -29,6 +29,7 @@ constexpr char kInputFileFlag[] = "i";
 constexpr char kOutputFileFlag[] = "o";
 constexpr char kConstElementCountLimitFlag[] = "const_element_count_limit";
 constexpr char kDisableMlirFlag[] = "disable_mlir";
+constexpr char kQuantParamsCountLimitFlag[] = "quant_params_count_limit";
 
 namespace {
 
@@ -42,6 +43,7 @@ int main(int argc, char* argv[]) {
   // Creates and parses flags.
   std::string input_file, output_file;
   int const_element_count_limit = 16;
+  int quant_params_count_limit = 16;
   bool disable_mlir = false;
 
   std::vector<mlir::Flag> flag_list = {
@@ -61,6 +63,12 @@ int main(int argc, char* argv[]) {
           "Disable the MLIR-based conversion. If set to true, the conversion "
           "becomes from model directly to graph json",
           mlir::Flag::kOptional),
+      mlir::Flag::CreateFlag(
+          kQuantParamsCountLimitFlag, &quant_params_count_limit,
+          "The maximum number of quant parameters. If the number exceeds this "
+          "threshold, the rest of data will be elided. If the flag is not set, "
+          "the default threshold is 16 (use -1 to print all)",
+          mlir::Flag::kOptional),
   };
   mlir::Flags::Parse(&argc, const_cast<const char**>(argv), flag_list);
 
@@ -75,7 +83,7 @@ int main(int argc, char* argv[]) {
 
   // Creates visualization config.
   tooling::visualization_client::VisualizeConfig config(
-      const_element_count_limit);
+      const_element_count_limit, quant_params_count_limit);
 
   const absl::StatusOr<std::string> json_output =
       ConvertModelToJson(config, input_file, disable_mlir);

src/builtin-adapter/visualize_config.h

@@ -21,13 +21,20 @@ namespace visualization_client {
 
 struct VisualizeConfig {
   VisualizeConfig() = default;
-  explicit VisualizeConfig(const int const_element_count_limit)
-      : const_element_count_limit(const_element_count_limit) {}
+  explicit VisualizeConfig(const int const_element_count_limit = 16,
+                           const int quant_params_count_limit = 16)
+      : const_element_count_limit(const_element_count_limit),
+        quant_params_count_limit(quant_params_count_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 number of quantization parameters 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 quant_params_count_limit = 16;
 };
 
 }  // namespace visualization_client