Disable multi-device function inlining after partitioning (GraphOptimizer)

PiperOrigin-RevId: 240191756
This commit is contained in:
Eugene Zhulenev 2019-03-25 12:12:40 -07:00 committed by TensorFlower Gardener
parent e38c98c650
commit c3865e746b
3 changed files with 19 additions and 2 deletions

View File

@ -1471,8 +1471,10 @@ using OutputControlSrc = InlineFunctionBodyOptions::OutputControlSource;
} // namespace } // namespace
string InlineFunctionBodyOptions::DebugString() const { string InlineFunctionBodyOptions::DebugString() const {
return absl::StrCat("ignore_noinline=", ignore_noinline ? "true" : "false", const auto true_false = [](bool b) { return b ? "true" : "false"; };
", override_device=", override_device ? "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=",
output_control_src == OutputControlSrc::kDataOutputs output_control_src == OutputControlSrc::kDataOutputs
? "DataOutputs" ? "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) { if (!options.ignore_noinline) {
TF_RETURN_IF_ERROR(ValidateNoInline(fbody)); TF_RETURN_IF_ERROR(ValidateNoInline(fbody));
} }

View File

@ -168,6 +168,10 @@ struct InlineFunctionBodyOptions {
// b) data returns (`ret` field in FunctionDef) // b) data returns (`ret` field in FunctionDef)
enum class OutputControlSource { kDataOutputs, kControlOutputs }; 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. // Ignore '_noinline' function attribute.
bool ignore_noinline = false; bool ignore_noinline = false;
// If 'true' function inlining will override explicitly specified devices // If 'true' function inlining will override explicitly specified devices

View File

@ -89,6 +89,12 @@ void GraphOptimizer::Optimize(
if (opts_.do_function_inlining()) { if (opts_.do_function_inlining()) {
ExpandInlineFunctionsOptions expand_inline_opts; ExpandInlineFunctionsOptions expand_inline_opts;
expand_inline_opts.native_options.override_device = true; 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); bool was_mutated = ExpandInlineFunctions(runtime, g, expand_inline_opts);
if (was_mutated) { if (was_mutated) {