commit 3882a2a901779713ecbe6800964c7fe1c6c85e26
author Roman Dzhabarov <rdzhabarov@google.com> Thu Jul 15 18:45:03 2021 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Thu Jul 15 18:49:17 2021 -0700
tree 6e30e79f35452f0cdd826610d1bd1061f1880482
parent ffb4ca869a1b08075e9d92323c864cc489549db6

[mlir/importer] Remove ImporterBase from the stack while recursively import TF graph to a MLIR module.

The change is to reduce stack size pressure while importing the model.

PiperOrigin-RevId: 385053644
Change-Id: I86e5c2a131ecd26eb2c824cf7dc7cd99c3fcad4c

tensorflow/compiler/mlir/tensorflow/translate/import_model.cc

diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index 4266433..b218532 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -1328,13 +1328,15 @@
     }
   }
 
-  ImporterBase child_importer(graph_flib_, debug_info_, specs, module_,
-                              tf_name_to_mlir_name_, function_name_uniquifier_,
-                              func_name);
-  TF_RETURN_IF_ERROR(child_importer.PrepareConvert(*fbody->graph));
+  // std::make_unique is not possible here, using `new` to access
+  // a non-public constructor.
+  std::unique_ptr<ImporterBase> child_importer(new ImporterBase(
+      graph_flib_, debug_info_, specs, module_, tf_name_to_mlir_name_,
+      function_name_uniquifier_, func_name));
+  TF_RETURN_IF_ERROR(child_importer->PrepareConvert(*fbody->graph));
 
   TF_ASSIGN_OR_RETURN(auto func_type,
-                      child_importer.InferLibFunctionType(*fbody));
+                      child_importer->InferLibFunctionType(*fbody));
 
   absl::InlinedVector<OutputTensor, 4> arg_nodes;
   absl::InlinedVector<OutputTensor, 4> ret_nodes;
@@ -1342,7 +1344,7 @@
   GetArgsAndRetsFromFunctionBody(*fbody, &arg_nodes, &ret_nodes,
                                  &control_ret_nodes);
 
-  TF_RETURN_IF_ERROR(child_importer.Convert(
+  TF_RETURN_IF_ERROR(child_importer->Convert(
       mlir_func_name, func_type, arg_nodes, ret_nodes, control_ret_nodes,
       llvm::makeArrayRef(attributes.begin(), attributes.end())));
   return Status::OK();