Factor out tf_saved_model passes from tensorflow_passes
Since tf_saved_model passes now requires CPU op kernels to capture resource
variables, it would be better to be split from the tensorflow passes to keep
use cases, that are not required tf_saved_model dialect, small.
PiperOrigin-RevId: 318143061
Change-Id: I0b30efd1dcb8945cdfe6184664f66cd757e0ea67
diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index 8d4efeb..f3ad2e1 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -839,6 +839,7 @@
"//tensorflow/compiler/mlir/tensorflow:tensorflow_dialect_registration",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow:tf_graph_optimization_pass",
+ "//tensorflow/compiler/mlir/tensorflow:tf_saved_model_passes",
"//tensorflow/compiler/mlir/tensorflow:translate_lib",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:Analysis",
diff --git a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
index 3fa2eae..de07ce7 100644
--- a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
+++ b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
@@ -26,6 +26,7 @@
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/decode_constant.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.h"
namespace mlir {
diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index db31d4f..5de8db4 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -470,6 +470,34 @@
)
cc_library(
+ name = "tf_saved_model_passes",
+ srcs = [
+ "transforms/freeze_global_tensors.cc",
+ "transforms/lift_variables_pass_registration.cc",
+ "transforms/optimize_global_tensors.cc",
+ ],
+ hdrs = [
+ "transforms/tf_saved_model_passes.h",
+ ],
+ deps = [
+ ":lift_variables_pass",
+ ":tensorflow",
+ ":tensorflow_passes",
+ ":tensorflow_types",
+ "//tensorflow/core:core_cpu",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "@llvm-project//llvm:Support",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:Pass",
+ "@llvm-project//mlir:StandardOps",
+ "@llvm-project//mlir:Support",
+ "@llvm-project//mlir:Transforms",
+ ],
+ alwayslink = 1,
+)
+
+cc_library(
name = "tensorflow_passes",
srcs = [
"transforms/annotate_parameter_replication.cc",
@@ -487,7 +515,6 @@
"transforms/executor_tpuv1_island_coarsening.cc",
"transforms/executor_tpuv1_outline_tpu_island.cc",
"transforms/fold_switch.cc",
- "transforms/freeze_global_tensors.cc",
"transforms/functional_control_flow_to_cfg.cc",
"transforms/functional_control_flow_to_regions.cc",
"transforms/fused_kernel_matcher.cc",
@@ -499,7 +526,6 @@
"transforms/layout_optimization.cc",
"transforms/materialize_mlir_passthrough_op.cc",
"transforms/optimize.cc",
- "transforms/optimize_global_tensors.cc",
"transforms/parallel_execute_to_islands.cc",
"transforms/promote_resources_to_args.cc",
"transforms/readonly_references_to_resources.cc",
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h
index 0eaee95..cff4936 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h
@@ -45,12 +45,6 @@
::tensorflow::Session* session_;
};
-// Creates as pass that creates GlobalTensorOp for each variable from function
-// arguments and converts the function arguments to the corresponding saved
-// model arguments.
-std::unique_ptr<OperationPass<ModuleOp>> CreateLiftVariablesPass(
- ::tensorflow::Session* session);
-
} // namespace tf_saved_model
} // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass_registration.cc b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass_registration.cc
new file mode 100644
index 0000000..672c169
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass_registration.cc
@@ -0,0 +1,27 @@
+/* Copyright 2020 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.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h"
+
+namespace mlir {
+namespace tf_saved_model {
+
+std::unique_ptr<OperationPass<ModuleOp>> CreateLiftVariablesPass(
+ ::tensorflow::Session* session) {
+ return std::make_unique<LiftVariablesPass>(session);
+}
+
+} // namespace tf_saved_model
+} // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
index 168b317..795f487 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
@@ -296,16 +296,6 @@
} // namespace TFTPU
-namespace tf_saved_model {
-
-// Creates a pass that optimizes tf_saved_model.global_tensor ops.
-std::unique_ptr<OperationPass<ModuleOp>> CreateOptimizeGlobalTensorsPass();
-
-// Creates a pass that freezes tf_saved_model.global_tensor ops.
-std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeGlobalTensorsPass();
-
-} // namespace tf_saved_model
-
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_PASSES_H_
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h b/tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h
new file mode 100644
index 0000000..242c4c0
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h
@@ -0,0 +1,42 @@
+/* Copyright 2020 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.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_TF_SAVED_MODEL_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_TF_SAVED_MODEL_PASSES_H_
+
+#include <memory>
+
+#include "mlir/Pass/Pass.h" // from @llvm-project
+#include "tensorflow/core/public/session.h"
+
+namespace mlir {
+namespace tf_saved_model {
+
+// Creates a pass that optimizes tf_saved_model.global_tensor ops.
+std::unique_ptr<OperationPass<ModuleOp>> CreateOptimizeGlobalTensorsPass();
+
+// Creates a pass that freezes tf_saved_model.global_tensor ops.
+std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeGlobalTensorsPass();
+
+// Creates as pass that creates GlobalTensorOp for each variable from function
+// arguments and converts the function arguments to the corresponding saved
+// model arguments.
+std::unique_ptr<OperationPass<ModuleOp>> CreateLiftVariablesPass(
+ ::tensorflow::Session* session);
+
+} // namespace tf_saved_model
+} // namespace mlir
+
+#endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_TF_SAVED_MODEL_PASSES_H_
diff --git a/tensorflow/compiler/mlir/tfjs/BUILD b/tensorflow/compiler/mlir/tfjs/BUILD
index b4b72ae..7d3091f 100644
--- a/tensorflow/compiler/mlir/tfjs/BUILD
+++ b/tensorflow/compiler/mlir/tfjs/BUILD
@@ -132,6 +132,7 @@
"//tensorflow/compiler/mlir/tensorflow:tensorflow_dialect_registration",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow:tf_graph_optimization_pass",
+ "//tensorflow/compiler/mlir/tensorflow:tf_saved_model_passes",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
diff --git a/tensorflow/compiler/mlir/tfjs/tf_tfjs_passes.cc b/tensorflow/compiler/mlir/tfjs/tf_tfjs_passes.cc
index d48d909..b7e9562 100644
--- a/tensorflow/compiler/mlir/tfjs/tf_tfjs_passes.cc
+++ b/tensorflow/compiler/mlir/tfjs/tf_tfjs_passes.cc
@@ -21,6 +21,7 @@
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h"
#include "tensorflow/compiler/mlir/tfjs/transforms/passes.h"
namespace tensorflow {
