No public description

PiperOrigin-RevId: 691662609
KevinKaChunLee · Oct 31, 2024 · 01f1eb4 · 01f1eb4
1 parent c90bc64
commit 01f1eb4
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 15 deletions.
diff --git a/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc b/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc
@@ -434,21 +434,6 @@ absl::Status AddAuxiliaryNode(
   return absl::OkStatus();
 }
 
-// Logic referred from `CustomOptionsToAttributes` in
-// tensorflow/compiler/mlir/lite/flatbuffer_operator.cc.
-void CustomOptionsToAttributes(
-    const std::vector<uint8_t>& custom_options, mlir::Builder mlir_builder,
-    llvm::SmallVectorImpl<mlir::NamedAttribute>& attributes) {
-  const flexbuffers::Map& map = flexbuffers::GetRoot(custom_options).AsMap();
-  const flexbuffers::TypedVector& keys = map.Keys();
-  for (size_t i = 0; i < keys.size(); ++i) {
-    const char* key = keys[i].AsKey();
-    const flexbuffers::Reference& value = map[key];
-    attributes.emplace_back(mlir_builder.getNamedAttr(
-        key, mlir_builder.getStringAttr(value.ToString())));
-  }
-}
-
 absl::Status SubgraphIdxToAttributes(
     const tflite::OperatorT& op, const std::vector<std::string>& func_names,
     mlir::Builder mlir_builder,
@@ -898,5 +883,28 @@ absl::StatusOr<std::string> ConvertFlatbufferDirectlyToJson(
   return llvm::formatv("{0:2}", json_result);
 }
 
+// Logic referred from `CustomOptionsToAttributes` in
+// tensorflow/compiler/mlir/lite/flatbuffer_operator.cc.
+void CustomOptionsToAttributes(
+    const std::vector<uint8_t>& custom_options, mlir::Builder mlir_builder,
+    llvm::SmallVectorImpl<mlir::NamedAttribute>& attributes) {
+  if (custom_options.empty()) {
+    // Avoid calling flexbuffers::GetRoot() with empty data. Otherwise it will
+    // crash.
+    //
+    // TODO(yijieyang): We should use a default value for input custom_options
+    // that is not empty to avoid this check.
+    return;
+  }
+  const flexbuffers::Map& map = flexbuffers::GetRoot(custom_options).AsMap();
+  const flexbuffers::TypedVector& keys = map.Keys();
+  for (size_t i = 0; i < keys.size(); ++i) {
+    const char* key = keys[i].AsKey();
+    const flexbuffers::Reference& value = map[key];
+    attributes.emplace_back(mlir_builder.getNamedAttr(
+        key, mlir_builder.getStringAttr(value.ToString())));
+  }
+}
+
 }  // namespace visualization_client
 }  // namespace tooling
diff --git a/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.h b/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.h
@@ -17,9 +17,13 @@
 #define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_GOOGLE_TOOLING_DIRECT_FLATBUFFER_TO_JSON_GRAPH_CONVERT_H_
 
 #include <string>
+#include <vector>
 
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
 #include "visualize_config.h"
 
 namespace tooling {
@@ -30,6 +34,13 @@ namespace visualization_client {
 absl::StatusOr<std::string> ConvertFlatbufferDirectlyToJson(
     const VisualizeConfig& config, absl::string_view model_path);
 
+// Converts custom options to attributes.
+// Logic referred from `CustomOptionsToAttributes` in
+// tensorflow/compiler/mlir/lite/flatbuffer_operator.cc.
+void CustomOptionsToAttributes(
+    const std::vector<uint8_t>& custom_options, mlir::Builder mlir_builder,
+    llvm::SmallVectorImpl<mlir::NamedAttribute>& attributes);
+
 }  // namespace visualization_client
 }  // namespace tooling