From c3865e746b3eb4d0a789874bc9efa55f8e598b7f Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Mon, 25 Mar 2019 12:12:40 -0700
Subject: [PATCH] Disable multi-device function inlining after partitioning
 (GraphOptimizer)

PiperOrigin-RevId: 240191756
---
 tensorflow/core/common_runtime/function.cc        | 11 +++++++++--
 tensorflow/core/common_runtime/function.h         |  4 ++++
 tensorflow/core/common_runtime/graph_optimizer.cc |  6 ++++++
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/function.cc b/tensorflow/core/common_runtime/function.cc
index dabd977d18e..ccfbf4473a8 100644
--- a/tensorflow/core/common_runtime/function.cc
+++ b/tensorflow/core/common_runtime/function.cc
@@ -1471,8 +1471,10 @@ using OutputControlSrc = InlineFunctionBodyOptions::OutputControlSource;
 }  // namespace
 
 string InlineFunctionBodyOptions::DebugString() const {
-  return absl::StrCat("ignore_noinline=", ignore_noinline ? "true" : "false",
-                      ", override_device=", override_device ? "true" : "false",
+  const auto true_false = [](bool b) { return b ? "true" : "false"; };
+  return absl::StrCat("disable_inlining=", true_false(disable_inlining),
+                      ", ignore_noinline=", true_false(ignore_noinline),
+                      ", override_device=", true_false(ignore_noinline),
                       ", output_control_src=",
                       output_control_src == OutputControlSrc::kDataOutputs
                           ? "DataOutputs"
@@ -1520,6 +1522,11 @@ Status ValidateInlining(const Node* node, const FunctionBody* fbody,
     }
   }
 
+  if (options.disable_inlining) {
+    return errors::InvalidArgument(
+        "Function inlining explicitly disabled by 'options.disable_inlining'");
+  }
+
   if (!options.ignore_noinline) {
     TF_RETURN_IF_ERROR(ValidateNoInline(fbody));
   }
diff --git a/tensorflow/core/common_runtime/function.h b/tensorflow/core/common_runtime/function.h
index 7b2ce35cb4d..cbaf74de002 100644
--- a/tensorflow/core/common_runtime/function.h
+++ b/tensorflow/core/common_runtime/function.h
@@ -168,6 +168,10 @@ struct InlineFunctionBodyOptions {
   //   b) data returns               (`ret` field in FunctionDef)
   enum class OutputControlSource { kDataOutputs, kControlOutputs };
 
+  // If 'true' function inlining is completely disabled. This allows to control
+  // function inlining for different types of function calls (see
+  // 'ExpandInlineFunctionsOptions' below).
+  bool disable_inlining = false;
   // Ignore '_noinline' function attribute.
   bool ignore_noinline = false;
   // If 'true' function inlining will override explicitly specified devices
diff --git a/tensorflow/core/common_runtime/graph_optimizer.cc b/tensorflow/core/common_runtime/graph_optimizer.cc
index 465cddfe1ab..2c66beca5e5 100644
--- a/tensorflow/core/common_runtime/graph_optimizer.cc
+++ b/tensorflow/core/common_runtime/graph_optimizer.cc
@@ -89,6 +89,12 @@ void GraphOptimizer::Optimize(
     if (opts_.do_function_inlining()) {
       ExpandInlineFunctionsOptions expand_inline_opts;
       expand_inline_opts.native_options.override_device = true;
+      // GraphOptimizer is running:
+      //   (1) After partitioning when executing with a Session API.
+      //   (2) For a single device function body after instantiation.
+      // We can't inline multi-device functions in these cases, because it might
+      // lead to multiple device assignments.
+      expand_inline_opts.multi_device_options.disable_inlining = true;
 
       bool was_mutated = ExpandInlineFunctions(runtime, g, expand_inline_opts);
       if (was_mutated) {