Disable multi-device function inlining after partitioning (GraphOptimizer)

PiperOrigin-RevId: 240191756

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));
   }

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

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) {