Prefix the signature name to private functions' names when importing a signature.
Previously, it is assumed that the private functions' names are unique. However, it can be broken by graph rewrites such functionalization and grappler optimization. Prefixing the signature name to these private functions will make sure they are unique in the generated MLIR module.
PiperOrigin-RevId: 357825986
Change-Id: Ie06befa62ccfe400382ffd88e5ae8bb3cadd0ba7
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/control_flow_duplicate_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/control_flow_duplicate_v1.py
index 67563fa..ab786ac 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/control_flow_duplicate_v1.py
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/control_flow_duplicate_v1.py
@@ -27,19 +27,16 @@
# CHECK: func {{.*}} tf_saved_model.exported_names = ["key_1"]
# CHECK: "tf.If"
-# CHECK-SAME: else_branch = @[[else_1:"key_1/[a-zA-Z_0-9]+"]]
-# CHECK-SAME: then_branch = @[[then_1:"key_1/[a-zA-Z_0-9]+"]]
+# CHECK-SAME: else_branch = @[[else:[a-zA-Z_0-9]+]]
+# CHECK-SAME: then_branch = @[[then:[a-zA-Z_0-9]+]]
# CHECK: func {{.*}} tf_saved_model.exported_names = ["key_2"]
# CHECK: "tf.If"
-# CHECK-SAME: else_branch = @[[else_2:"key_2/[a-zA-Z_0-9]+"]]
-# CHECK-SAME: then_branch = @[[then_2:"key_2/[a-zA-Z_0-9]+"]]
+# CHECK-SAME: else_branch = @[[else]]
+# CHECK-SAME: then_branch = @[[then]]
-# CHECK: func private @[[else_1]](
-# CHECK: func private @[[then_1]](
-
-# CHECK: func private @[[else_2]](
-# CHECK: func private @[[then_2]](
+# CHECK: func private @[[else]](
+# CHECK: func private @[[then]](
def Test():
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index c5f475b..26bb079 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -3487,8 +3487,7 @@
// Moves the functions in `sub_module` to `module_` and skips the duplicate
// functions.
- Status MoveConvertedFunctionsToModule(absl::string_view name,
- mlir::ModuleOp sub_module);
+ void MoveConvertedFunctionsToModule(mlir::ModuleOp sub_module);
GraphImportConfig::InputArrays ParseInputArrays(
llvm::ArrayRef<std::pair<std::string, TensorInfo>> inputs);
@@ -3528,34 +3527,14 @@
return results;
}
-Status SavedModelSignatureDefImporterLite::MoveConvertedFunctionsToModule(
- absl::string_view name, mlir::ModuleOp sub_module) {
- mlir::Builder builder(sub_module.getContext());
- mlir::SymbolTable sub_module_symbol_table(sub_module);
-
- // Prefix private functions with the unique signature name, so that it cannot
- // collide with private functions used in the other signatures.
+void SavedModelSignatureDefImporterLite::MoveConvertedFunctionsToModule(
+ mlir::ModuleOp sub_module) {
+ // Iterate through all functions and insert the ones that do not already exist
+ // in `module_`.
for (auto func : sub_module.getOps<mlir::FuncOp>()) {
- if (mlir::tf_saved_model::IsExported(func)) continue;
-
- std::string new_sym_name = absl::StrCat(name, "/", func.sym_name().str());
- if (mlir::failed(sub_module_symbol_table.replaceAllSymbolUses(
- func, new_sym_name, sub_module)))
- return tensorflow::errors::InvalidArgument(absl::StrCat(
- "SavedModelSignatureDefImporterLite: failed to assign a unique "
- "name to the private function used in a signature: ",
- func.sym_name().str()));
-
- mlir::SymbolTable::setSymbolName(func, new_sym_name);
- }
-
- // Copy all functions used by this signature to the final MLIR module.
- for (auto func : sub_module.getOps<mlir::FuncOp>()) {
- DCHECK(symbol_table_.lookup(func.sym_name()) == nullptr);
+ if (symbol_table_.lookup(func.getName())) continue;
symbol_table_.insert(func.clone());
}
-
- return Status::OK();
}
Status SavedModelSignatureDefImporterLite::ConvertInitializer(
@@ -3596,7 +3575,9 @@
"__tf_saved_model_session_initializer_", target_node_name)}));
// Move the converted functions to top level MLIR module.
- return MoveConvertedFunctionsToModule(target_node_name, *sub_module);
+ MoveConvertedFunctionsToModule(*sub_module);
+
+ return Status::OK();
}
StatusOr<mlir::OwningModuleRef>
@@ -3666,7 +3647,9 @@
}
// Move the converted functions to top level MLIR module.
- return MoveConvertedFunctionsToModule(sig_def_key, *sub_module);
+ MoveConvertedFunctionsToModule(*sub_module);
+
+ return Status::OK();
}
GraphImportConfig::InputArrays
