Turn on function optimization by default

PiperOrigin-RevId: 188722505

diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index e9ed5c4..98a18e4 100644 (file)
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -3205,6 +3205,7 @@ tf_cc_test(
         "//tensorflow/core/kernels:dense_update_ops",
         "//tensorflow/core/kernels:fifo_queue_op",
         "//tensorflow/core/kernels:function_ops",
+        "//tensorflow/core/kernels:identity_n_op",
         "//tensorflow/core/kernels:identity_op",
         "//tensorflow/core/kernels:matmul_op",
         "//tensorflow/core/kernels:ops_util",
@@ -3247,6 +3248,7 @@ tf_cc_test(
         "//tensorflow/core/kernels:fifo_queue_op",
         "//tensorflow/core/kernels:function_ops",
         "//tensorflow/core/kernels:identity_op",
+        "//tensorflow/core/kernels:identity_n_op",
         "//tensorflow/core/kernels:matmul_op",
         "//tensorflow/core/kernels:ops_util",
         "//tensorflow/core/kernels:queue_ops",

diff --git a/tensorflow/core/grappler/optimizers/function_optimizer.cc b/tensorflow/core/grappler/optimizers/function_optimizer.cc
index d8a237c..87160f6 100644 (file)
--- a/tensorflow/core/grappler/optimizers/function_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer.cc
@@ -136,6 +136,11 @@ Status FunctionOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
     if (func.attr().count("_noinline") != 0) {
       continue;
     }
+    // Don't touch anything marked XLA to prevent XLA failures further down the
+    // road.
+    if (func.attr().count("_XlaCompile") != 0) {
+      continue;
+    }
     // Can't create IdentityN nodes with no input or output: skip these
     // functions for now.
     if (func.signature().input_arg_size() == 0 ||

diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index 6fa8c03..3a76493 100644 (file)
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -94,7 +94,7 @@ Status MetaOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
     if (!cfg_.disable_model_pruning()) {
       optimizers.push_back(std::unique_ptr<GraphOptimizer>(new ModelPruner()));
     }
-    if (cfg_.function_optimization() == RewriterConfig::ON) {
+    if (cfg_.function_optimization() != RewriterConfig::OFF) {
       optimizers.push_back(
           std::unique_ptr<GraphOptimizer>(new FunctionOptimizer()));
     }
@@ -231,7 +231,7 @@ void MetaOptimizer::Feedback(Cluster* cluster, const GrapplerItem& item,
 bool MetaOptimizerEnabled(const RewriterConfig& cfg) {
   return !cfg.disable_model_pruning() ||
          cfg.layout_optimizer() != RewriterConfig::OFF ||
-         cfg.function_optimization() == RewriterConfig::ON ||
+         cfg.function_optimization() != RewriterConfig::OFF ||
          cfg.constant_folding() != RewriterConfig::OFF ||
          cfg.arithmetic_optimization() != RewriterConfig::OFF ||
          cfg.loop_optimization() == RewriterConfig::ON ||

diff --git a/tensorflow/core/protobuf/rewriter_config.proto b/tensorflow/core/protobuf/rewriter_config.proto
index 0ccf214..b1fceaa 100644 (file)
--- a/tensorflow/core/protobuf/rewriter_config.proto
+++ b/tensorflow/core/protobuf/rewriter_config.proto
@@ -44,7 +44,7 @@ message RewriterConfig {
   Toggle dependency_optimization = 8;
   // Loop optimizations (default is OFF).
   Toggle loop_optimization = 9;
-  // Function optimizations (default is OFF).
+  // Function optimizations (default is ON).
   Toggle function_optimization = 10;
   // If true, don't remove unnecessary ops from the graph
   bool disable_model_pruning = 2;