Disable multi-device function inlining after partitioning (GraphOptimizer)
PiperOrigin-RevId: 240191756
This commit is contained in:
parent
e38c98c650
commit
c3865e746b
@ -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));
|
||||||
}
|
}
|
||||||
|
@ -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
|
||||||
|
@ -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) {
|
||||||
|
Loading…
Reference in New Issue
Block a user