[go: nahoru, domu]
Skip to content

Commit

Permalink
Clone the module into a new context to erase the unused tensors from …
Browse files Browse the repository at this point in the history 
…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: 645520361
  • Loading branch information
@vamsimanchala @tensorflower-gardener
vamsimanchala authored and tensorflower-gardener committed Jun 28, 2024
1 parent f346a1f commit 865aa92
Show file tree
Hide file tree
Showing 18 changed files with 369 additions and 139 deletions.
12 changes: 12 additions & 0 deletions tensorflow/compiler/mlir/lite/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1411,6 +1411,7 @@ cc_library(
],
deps = [
":common",
":const_tensor_utils",
":flatbuffer_translate_lib",
":tensorflow_lite",
":tf_tfl_passes",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
],
)

Expand Down Expand Up @@ -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",
Expand Down
4 changes: 3 additions & 1 deletion tensorflow/compiler/mlir/lite/debug/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
],
)
Expand Down
10 changes: 9 additions & 1 deletion tensorflow/compiler/mlir/lite/debug/debug.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
10 changes: 6 additions & 4 deletions tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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();
Expand All @@ -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);
}
Expand Down
9 changes: 5 additions & 4 deletions tensorflow/compiler/mlir/lite/python/jax_to_tfl_flatbuffer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.");
}
Expand All @@ -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;
Expand Down
17 changes: 10 additions & 7 deletions tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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()) {
Expand Down Expand Up @@ -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;
}

Expand Down
10 changes: 6 additions & 4 deletions tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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;

Expand All @@ -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;
Expand Down
3 changes: 2 additions & 1 deletion tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
25 changes: 14 additions & 11 deletions tensorflow/compiler/mlir/lite/tf_tfl_translate.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ limitations under the License.
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>

#include "absl/status/statusor.h"
Expand Down Expand Up @@ -88,9 +89,7 @@ int main(int argc, char **argv) {

mlir::DialectRegistry registry;
mlir::func::registerAllExtensions(registry);
MLIRContext context(registry);
llvm::SourceMgr source_mgr;
mlir::SourceMgrDiagnosticHandler sourceMgrHandler(source_mgr, &context);
auto context = std::make_unique<mlir::MLIRContext>(registry);

if (input_mlir) {
// TODO(@zichuanwei): hack to enable mlir conversion via this tool, will get
Expand All @@ -99,7 +98,7 @@ int main(int argc, char **argv) {
RegisterAllTensorFlowDialects(registry);
registry
.insert<mlir::func::FuncDialect, mlir::stablehlo::StablehloDialect>();
context.appendDialectRegistry(registry);
context->appendDialectRegistry(registry);
}

absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> module;
Expand Down Expand Up @@ -146,7 +145,7 @@ int main(int argc, char **argv) {
custom_opdefs.end());
module = tensorflow::ImportSavedModel(
input_file_name, saved_model_version, tags, extra_opdefs,
exported_names, specs, /*enable_variable_lifting=*/true, &context,
exported_names, specs, /*enable_variable_lifting=*/true, context.get(),
&bundle);
} else if (import_hlo) {
// HLO import path.
Expand All @@ -161,19 +160,22 @@ int main(int argc, char **argv) {

auto content = buffer->getBuffer();
if (hlo_import_type == HloImportType::hlotxt) {
module = xla::HloTextToMlirHloTranslateFunction(content, &context, false);
module =
xla::HloTextToMlirHloTranslateFunction(content, context.get(), false);
} else if (hlo_import_type == HloImportType::proto) {
module = xla::HloToMlirHloTranslateFunction(content, &context, false);
module =
xla::HloToMlirHloTranslateFunction(content, context.get(), false);
} else {
module = mlir::OwningOpRef<mlir::ModuleOp>(
mlir::parseSourceString<mlir::ModuleOp>(content, &context));
mlir::parseSourceString<mlir::ModuleOp>(content, context.get()));
}
} else {
// Graphdef import path.
llvm::SourceMgr source_mgr;
module = tensorflow::LoadFromGraphdefOrMlirSource(
input_file_name, input_mlir, use_splatted_constant, custom_opdefs,
specs, debug_info_file, input_arrays, input_dtypes, input_shapes,
output_arrays, control_output_arrays, &source_mgr, &context);
output_arrays, control_output_arrays, &source_mgr, context.get());
}

// If errors occur, the library call in the above already logged the error
Expand Down Expand Up @@ -251,9 +253,10 @@ int main(int argc, char **argv) {

std::string result;
auto status = tensorflow::ConvertTFExecutorToTFLOrFlatbuffer(
module.value().get(), output_mlir, toco_flags, pass_config, tags,
std::move(context), std::move(module.value()), toco_flags, pass_config,
tags,
/*saved_model_dir=*/"", std::move(bundle), &result,
serialize_stablehlo_ops);
serialize_stablehlo_ops, /*output_mlir*/ output_mlir);
if (!status.ok()) {
llvm::errs() << status.message() << '\n';
return kTrFailure;
Expand Down
Loading

0 comments on commit 865aa92

Please sign in to comment.