Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[oneDNN] Fixing eager-related nightly test failures #52204

Merged
merged 2 commits into from Oct 11, 2021
Merged

[oneDNN] Fixing eager-related nightly test failures #52204

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dc41c70
Fixing eager-related testes
mahmoud-abuzaina Sep 30, 2021
36e3320
making GetMKLNodeDef function an inline function
sachinmuradi Oct 6, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
25 changes: 19 additions & 6 deletions tensorflow/core/common_runtime/eager/execute.cc
View file
Open in desktop
Expand Up @@ -727,6 +727,16 @@ Status PopulateRetMap(FunctionDef* fdef, const AbstractOpAttrs* op_attrs,
return Status::OK();
}

#ifdef INTEL_MKL
inline void GetMKLNodeDef(NodeDef* ndef) {
// All MKL eager ops have `_kernel` private attribute that needs to be set
// to a fixed label.
AttrValue attr_kernel;
attr_kernel.set_s(mkl_op_registry::kMklNameChangeOpLabel);
(*ndef->mutable_attr()).insert({"_kernel", attr_kernel});
}
#endif // INTEL_MKL

Status WrapInCallOp(EagerOperation* op, EagerOperation** wrapped_op) {
DCHECK(!op->is_function());
const OpDef& opdef = OpRegistry::Global()->LookUp(op->Name())->op_def;
Expand Down Expand Up @@ -776,11 +786,7 @@ Status WrapInCallOp(EagerOperation* op, EagerOperation** wrapped_op) {
#ifdef INTEL_MKL
if (IsMKLEnabled() &&
absl::StartsWith(op->Name(), mkl_op_registry::kMklOpPrefix)) {
// All MKL eager ops have `_kernel` private attribute that needs to be set
// to a fixed label.
AttrValue attr_kernel;
attr_kernel.set_s(mkl_op_registry::kMklNameChangeOpLabel);
(*ndef->mutable_attr()).insert({"_kernel", attr_kernel});
GetMKLNodeDef(ndef);
}
#endif // INTEL_MKL

Expand Down Expand Up @@ -933,7 +939,14 @@ Status GetOrCreateKernelAndDevice(
// place the wrapped op on a GPU (if one is available) which leads to
// errors because placer pins the function output nodes to GPU thereby
// forcing a H2D copy of the dataset variant which is not supported.
const NodeDef& ndef = op->MutableAttrs()->BuildNodeDef();
auto ndef = op->MutableAttrs()->BuildNodeDef();
#ifdef INTEL_MKL
if (IsMKLEnabled() &&
absl::StartsWith(op->Name(), mkl_op_registry::kMklOpPrefix)) {
GetMKLNodeDef(&ndef);
}
#endif // INTEL_MKL

TF_RETURN_IF_ERROR(ctx.SelectDevice(preferred_device, ndef, &device));

VLOG(1) << "PreferredDevice " << op->Name() << ": " << preferred_device;
Expand Down