Add a new converter flag `_experimental_guarantee_all_funcs_one_use` to allow cloning function calls so that each function have exactly one use.
By default the flag is False, since duplicating FuncOps will significantly bloat the model size in certain scenarios. With this flag set to True, user can enable the `mlir::TF::CreateGuaranteeAllFuncsOneUsePass()` pass so that function inlining can succeed when multiple `StatefulPartitionedCall` ops call into the same function.
PiperOrigin-RevId: 433335089
diff --git a/tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc b/tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc
index c52f31c..113e625 100644
--- a/tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc
+++ b/tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc
@@ -1,4 +1,4 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -110,6 +110,8 @@
pass_config.enable_dynamic_update_slice =
toco_flags.enable_dynamic_update_slice();
pass_config.preserve_assert_op = toco_flags.preserve_assert_op();
+ pass_config.guarantee_all_funcs_one_use =
+ toco_flags.guarantee_all_funcs_one_use();
return internal::ConvertMLIRToTFLiteFlatBuffer(
model_flags, toco_flags, std::move(module), pass_config,
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
index 1d48124..7d22e6a 100644
--- a/tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.cc
+++ b/tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.cc
@@ -1,4 +1,4 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -194,6 +194,8 @@
pass_config.enable_dynamic_update_slice =
toco_flags.enable_dynamic_update_slice();
pass_config.preserve_assert_op = toco_flags.preserve_assert_op();
+ pass_config.guarantee_all_funcs_one_use =
+ toco_flags.guarantee_all_funcs_one_use();
// TODO(b/153507667): Pass the session object when importing logic is removed.
auto status = internal::ConvertMLIRToTFLiteFlatBuffer(
diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py
index dd71322..2a4bdd6 100644
--- a/tensorflow/lite/python/convert.py
+++ b/tensorflow/lite/python/convert.py
@@ -1,5 +1,5 @@
# Lint as: python2, python3
-# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
@@ -502,6 +502,7 @@
use_fake_quant_num_bits=False,
enable_dynamic_update_slice=False,
preserve_assert_op=False,
+ guarantee_all_funcs_one_use=False,
**_):
"""Builds protocol buffer describing a conversion of a model.
@@ -582,6 +583,10 @@
enable_dynamic_update_slice: Enable to convert to DynamicUpdateSlice op.
(default: False).
preserve_assert_op: Whether to preserve `TF::AssertOp` (default: False).
+ guarantee_all_funcs_one_use: Whether to clone functions so that each
+ function only has a single use. This option will be helpful if the
+ conversion fails when the `PartitionedCall` or `StatefulPartitionedCall`
+ can't be properly inlined (default: False).
Returns:
conversion_flags: protocol buffer describing the conversion process.
@@ -639,6 +644,7 @@
enable_mlir_dynamic_range_quantizer)
conversion_flags.enable_dynamic_update_slice = enable_dynamic_update_slice
conversion_flags.preserve_assert_op = preserve_assert_op
+ conversion_flags.guarantee_all_funcs_one_use = guarantee_all_funcs_one_use
if tf_quantization_mode:
conversion_flags.tf_quantization_mode = tf_quantization_mode
conversion_flags.disable_infer_tensor_range = disable_infer_tensor_range
diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py
index 1c230c7..8b90262 100644
--- a/tensorflow/lite/python/lite.py
+++ b/tensorflow/lite/python/lite.py
@@ -1,5 +1,5 @@
# Lint as: python2, python3
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
@@ -553,6 +553,7 @@
self._metadata.environment.modelType = self._get_original_model_type()
self._experimental_enable_dynamic_update_slice = False
self._experimental_preserve_assert_op = False
+ self._experimental_guarantee_all_funcs_one_use = False
# When the value is true, the MLIR quantantizer triggers dynamic range
# quantization in MLIR instead of the old quantizer. Used only if
@@ -673,6 +674,8 @@
self._experimental_enable_dynamic_update_slice,
"preserve_assert_op":
self._experimental_preserve_assert_op,
+ "guarantee_all_funcs_one_use":
+ self._experimental_guarantee_all_funcs_one_use,
}
if self.saved_model_dir:
diff --git a/tensorflow/lite/toco/toco_flags.proto b/tensorflow/lite/toco/toco_flags.proto
index 9bc5416..0247d67 100644
--- a/tensorflow/lite/toco/toco_flags.proto
+++ b/tensorflow/lite/toco/toco_flags.proto
@@ -1,4 +1,4 @@
-// Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+// Copyright 2022 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
@@ -38,7 +38,7 @@
// of as properties of models, instead describing how models are to be
// processed in the context of the present tooling job.
//
-// Next ID to use: 50.
+// Next ID to use: 51.
message TocoFlags {
// Input file format
optional FileFormat input_format = 1;
@@ -297,4 +297,7 @@
// Whether to preserve `TF::AssertOp`.
optional bool preserve_assert_op = 49 [default = false];
+
+ // Whether to ensure each function has a single use.
+ optional bool guarantee_all_funcs_one_use = 50 [default = false];
}
