Clone the module into a new context to erase the unused tensors from …

…memory. Ideally, unused tensors cleaned-up during DCE must not consume any memory but they currently do. This is due to how MLIRContext is designed and this has been addressed with the introduction of DenseResourceElementsAttr in MLIR. This is a temporary workaround until we are able to use the DenseResourceElementsAttr to store the model checkpoint weights. PiperOrigin-RevId: 647742961
tensorflow · copybara-service · Jun 28, 2024 · Jun 28, 2024 · Jun 28, 2024 · df25ab1dafc77af272c39f513329c5a21456245a
commit df25ab1dafc77af272c39f513329c5a21456245a
diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
@@ -1411,6 +1411,7 @@ cc_library(
     ],
     deps = [
         ":common",
+        ":const_tensor_utils",
         ":flatbuffer_translate_lib",
         ":tensorflow_lite",
         ":tf_tfl_passes",
@@ -1445,6 +1446,7 @@ cc_library(
         "//tensorflow/core:core_cpu_base",
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/ir/types:Dialect",
         "//tensorflow/lite/python/metrics:converter_error_data_proto_cc",
         "//tensorflow/lite/toco:toco_flags_proto_cc",
         "//tensorflow/lite/tools/optimize:quantize_weights",
@@ -1456,14 +1458,22 @@ cc_library(
         "@com_google_absl//absl/types:span",
         "@flatbuffers//:runtime_cc",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BytecodeWriter",
+        "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:FuncExtensions",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
         "@llvm-project//mlir:ReconcileUnrealizedCasts",
         "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:protobuf",
         "@local_tsl//tsl/platform:statusor",
+        "@stablehlo//:stablehlo_ops",
+        "@stablehlo//:vhlo_ops",
     ],
 )
 
@@ -1501,6 +1511,8 @@ cc_library(
         "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
         "//tensorflow/core:protos_all_cc",
         "@com_google_absl//absl/base",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/meta:type_traits",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

diff --git a/tensorflow/compiler/mlir/lite/debug/BUILD b/tensorflow/compiler/mlir/lite/debug/BUILD
@@ -22,6 +22,7 @@ cc_library(
     visibility = ["//tensorflow/compiler/mlir/lite:__subpackages__"],
     deps = [
         ":debug_options_proto_cc",
+        "//tensorflow/compiler/mlir/lite/metrics:error_collector_inst",
         "//tensorflow/core:portable_gif_internal",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status",
@@ -32,10 +33,11 @@ cc_library(
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
         "@local_tsl//tsl/lib/io:buffered_file",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:path",
-        "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:stringpiece",
     ],
 )

diff --git a/tensorflow/compiler/mlir/lite/debug/debug.cc b/tensorflow/compiler/mlir/lite/debug/debug.cc
@@ -34,20 +34,24 @@ limitations under the License.
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/raw_ostream.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
 #include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Transforms/ViewOpGraph.h"  // from @llvm-project
 #include "re2/re2.h"  // IWYU pragma: keep
 #include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
+#include "tensorflow/compiler/mlir/lite/metrics/error_collector_inst.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tsl/lib/io/buffered_file.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/file_system.h"
 #include "tsl/platform/path.h"
-#include "tsl/platform/status.h"
 #include "tsl/platform/stringpiece.h"
+
 // IWYU pragma: no_include "util/regexp/re2/re2.h"
 
 namespace tensorflow {
@@ -342,6 +346,10 @@ void InitPassManager(mlir::PassManager& pm,
   if (options.enable_timing()) {
     pm.enableTiming();
   }
+
+  pm.addInstrumentation(
+      std::make_unique<mlir::TFL::ErrorCollectorInstrumentation>(
+          pm.getContext()));
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc b/tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.h"
 
+#include <memory>
 #include <optional>
 #include <string>
 #include <utility>
@@ -44,7 +45,7 @@ absl::Status ConvertGraphDefToTFLiteFlatBuffer(
     const GraphDebugInfo& debug_info, const GraphDef& input,
     std::string* result) {
   using ::tflite::optimize::ReducedPrecisionSupport;
-  mlir::MLIRContext context;
+  auto context = std::make_unique<mlir::MLIRContext>();
   GraphImportConfig specs;
   mlir::quant::QuantizationSpecs quant_specs;
 
@@ -85,8 +86,8 @@ absl::Status ConvertGraphDefToTFLiteFlatBuffer(
   // Register all custom ops, including user-specified custom ops.
   TF_RETURN_IF_ERROR(internal::RegisterAllCustomOps(toco_flags));
 
-  TF_ASSIGN_OR_RETURN(
-      auto module, ConvertGraphdefToMlir(input, debug_info, specs, &context));
+  TF_ASSIGN_OR_RETURN(auto module, ConvertGraphdefToMlir(input, debug_info,
+                                                         specs, context.get()));
 
   mlir::TFL::PassConfig pass_config(quant_specs);
   bool emit_builtin_tflite_ops = !toco_flags.force_select_tf_ops();
@@ -112,7 +113,8 @@ absl::Status ConvertGraphDefToTFLiteFlatBuffer(
   // StableHLO Quantizer is not supported for GraphDef inputs, so
   // quantization_py_function_lib is set to nullptr.
   return internal::ConvertMLIRToTFLiteFlatBuffer(
-      model_flags, toco_flags, std::move(module), pass_config,
+      model_flags, toco_flags, std::move(context), std::move(module),
+      pass_config,
       /*saved_model_tags=*/{}, result, /*saved_model_bundle=*/nullptr,
       /*quantization_py_function_lib=*/nullptr);
 }

diff --git a/tensorflow/compiler/mlir/lite/python/jax_to_tfl_flatbuffer.cc b/tensorflow/compiler/mlir/lite/python/jax_to_tfl_flatbuffer.cc
@@ -123,7 +123,7 @@ absl::Status ConvertJaxToTFLiteFlatBuffer(const std::string& input,
                                           const toco::ModelFlags& model_flags,
                                           toco::TocoFlags& toco_flags,
                                           std::string* result) {
-  mlir::MLIRContext context;
+  auto context = std::make_unique<mlir::MLIRContext>();
   mlir::quant::QuantizationSpecs quant_specs;
 
   // Parse input arrays.
@@ -162,9 +162,9 @@ absl::Status ConvertJaxToTFLiteFlatBuffer(const std::string& input,
 
   mlir::OwningOpRef<mlir::ModuleOp> module;
   if (model_flags.hlo_file_type() == toco::ModelFlags::HLO_TEXT) {
-    module = HloTextToMlirHloTranslateFunction(input, &context, false);
+    module = HloTextToMlirHloTranslateFunction(input, context.get(), false);
   } else if (model_flags.hlo_file_type() == toco::ModelFlags::HLO_PROTO) {
-    module = HloToMlirHloTranslateFunction(input, &context, false);
+    module = HloToMlirHloTranslateFunction(input, context.get(), false);
   } else {
     return errors::InvalidArgument("unknown hlo format type.");
   }
@@ -191,7 +191,8 @@ absl::Status ConvertJaxToTFLiteFlatBuffer(const std::string& input,
   // StableHLO Quantizer is not supported for JAX input models, so
   // quantization_py_function_lib is set to nullptr.
   auto status = internal::ConvertMLIRToTFLiteFlatBuffer(
-      model_flags, toco_flags, std::move(module), pass_config,
+      model_flags, toco_flags, std::move(context), std::move(module),
+      pass_config,
       /*saved_model_tags=*/{}, result, /*saved_model_bundle=*/nullptr,
       /*quantization_py_function_lib=*/nullptr);
   return status;

diff --git a/tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.cc b/tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.cc
@@ -137,7 +137,7 @@ Status ConvertSavedModelToTFLiteFlatBuffer(
     const toco::ModelFlags& model_flags, toco::TocoFlags& toco_flags,
     std::string* result,
     const PyFunctionLibrary* quantization_py_function_lib) {
-  mlir::MLIRContext context;
+  auto context = std::make_unique<mlir::MLIRContext>();
   mlir::quant::QuantizationSpecs quant_specs;
 
   // Parse input arrays.
@@ -177,10 +177,11 @@ Status ConvertSavedModelToTFLiteFlatBuffer(
   auto bundle = std::make_unique<tensorflow::SavedModelBundle>();
   TF_ASSIGN_OR_RETURN(
       auto module,
-      ImportSavedModel(
-          model_flags.saved_model_dir(), model_flags.saved_model_version(),
-          tags, absl::MakeSpan(custom_opdefs), exported_names, specs,
-          !toco_flags.enable_tflite_resource_variables(), &context, &bundle));
+      ImportSavedModel(model_flags.saved_model_dir(),
+                       model_flags.saved_model_version(), tags,
+                       absl::MakeSpan(custom_opdefs), exported_names, specs,
+                       !toco_flags.enable_tflite_resource_variables(),
+                       context.get(), &bundle));
 
   if (!model_flags.input_arrays().empty() ||
       !model_flags.output_arrays().empty()) {
@@ -240,8 +241,10 @@ Status ConvertSavedModelToTFLiteFlatBuffer(
 
   // TODO(b/153507667): Pass the session object when importing logic is removed.
   auto status = internal::ConvertMLIRToTFLiteFlatBuffer(
-      model_flags, toco_flags, std::move(module), pass_config, tags, result,
-      std::move(bundle), quantization_py_function_lib);
+      model_flags, toco_flags, std::move(context), std::move(module),
+      pass_config, tags, result, std::move(bundle),
+      quantization_py_function_lib);
+
   return status;
 }
 

diff --git a/tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc b/tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/OwningOpRef.h"  // from @llvm-project
 #include "mlir/IR/Visitors.h"  // from @llvm-project
 #include "mlir/Support/FileUtilities.h"  // from @llvm-project
@@ -324,12 +325,12 @@ absl::Status PopulateQuantizationSpecs(
 
 absl::Status ConvertMLIRToTFLiteFlatBuffer(
     const toco::ModelFlags& model_flags, toco::TocoFlags& toco_flags,
+    std::unique_ptr<mlir::MLIRContext>&& context,
     mlir::OwningOpRef<mlir::ModuleOp> module,
     const mlir::TFL::PassConfig& pass_config,
     const std::unordered_set<std::string>& saved_model_tags,
-    std::string* result, std::unique_ptr<SavedModelBundle> saved_model_bundle,
+    std::string* result, std::unique_ptr<SavedModelBundle>&& saved_model_bundle,
     const PyFunctionLibrary* quantization_py_function_lib) {
-
   mlir::TFL::PassConfig pass_config_copy = pass_config;
   pass_config_copy.outline_tf_while = true;
 
@@ -345,9 +346,10 @@ absl::Status ConvertMLIRToTFLiteFlatBuffer(
   });
 
   auto status = ConvertTFExecutorToTFLOrFlatbuffer(
-      module.get(), /*export_to_mlir=*/false, toco_flags, pass_config_copy,
+      std::move(context), std::move(module), toco_flags, pass_config_copy,
       saved_model_tags, model_flags.saved_model_dir(),
-      std::move(saved_model_bundle), result, /*serialize_stablehlo_ops=*/false,
+      std::move(saved_model_bundle), result,
+      /*serialize_stablehlo_ops=*/false, /*export_to_mlir=*/false,
       quantization_py_function_lib);
 
   return status;

diff --git a/tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h b/tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h
@@ -58,10 +58,11 @@ Status PopulateQuantizationSpecs(
 // This will also run relevant passes as well.
 Status ConvertMLIRToTFLiteFlatBuffer(
     const toco::ModelFlags& model_flags, toco::TocoFlags& toco_flags,
+    std::unique_ptr<mlir::MLIRContext>&& context,
     mlir::OwningOpRef<mlir::ModuleOp> module,
     const mlir::TFL::PassConfig& pass_config,
     const std::unordered_set<std::string>& saved_model_tags, string* result,
-    std::unique_ptr<SavedModelBundle> saved_model_bundle,
+    std::unique_ptr<SavedModelBundle>&& saved_model_bundle,
     const quantization::PyFunctionLibrary* quantization_py_function_lib);
 
 // Give a warning for any unused flags that have been specified.