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

r2.7 cherry-pick request: Fix missing-device unit test failures #52274

Merged
merged 1 commit into from Oct 18, 2021
Merged

r2.7 cherry-pick request: Fix missing-device unit test failures #52274

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
10 changes: 9 additions & 1 deletion tensorflow/core/common_runtime/eager/execute.cc
View file
Open in desktop
Expand Up @@ -83,6 +83,7 @@ limitations under the License.

#ifdef INTEL_MKL
#include "tensorflow/core/graph/mkl_graph_util.h"
#include "tensorflow/core/util/util.h"
#endif

namespace tensorflow {
Expand Down Expand Up @@ -773,7 +774,7 @@ Status WrapInCallOp(EagerOperation* op, EagerOperation** wrapped_op) {
ndef->set_device(op->DeviceName());

#ifdef INTEL_MKL
if (IsMklEnabled() &&
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.
Expand Down Expand Up @@ -1175,6 +1176,13 @@ Status EagerLocalExecute(EagerOperation* op, TensorHandle** retvals,
core::RefCountPtr<KernelAndDevice> kernel;
auto status = GetOrCreateKernelAndDevice(op, retvals, num_retvals, &kernel);

if (IsMKLEnabled() && kernel != nullptr && !ctx.RunEagerOpAsFunction() &&
op->Device() == kVariantDeviceNull) {
// oneDNN optimization pass relies on the op's assigned device to determine
// whether it can be rewritten.
op->SetDevice(kernel->device());
}

// Run all the registered rewrite pass after the placement, regardless whether
// the placement is successful or not. The passes can either create new ops
// (without placement) or update some fields of the input op.
Expand Down
17 changes: 7 additions & 10 deletions tensorflow/python/framework/node_file_writer_test.py
View file
Open in desktop
Expand Up @@ -148,7 +148,12 @@ def test_simple(self):
node_defs = self._get_new_node_defs()
self.assertLen(node_defs, 1)
(node_def3,) = node_defs # pylint: disable=unbalanced-tuple-unpacking
self.assertEqual(node_def3.op, 'MatMul')
if not IsMklEnabled():
self.assertEqual(node_def3.op, 'MatMul')
else:
# Under certain conditions ops can be rewritten by oneDNN optimization
# pass.
self.assertIn(node_def3.op, ['MatMul', '_MklMatMul'])
self.assertEqual(
self._get_input_dtypes(node_def3), [dtypes.float32, dtypes.float32])
self.assertEqual(self._get_input_shapes(node_def3), [(4, 3), (3, 2)])
Expand Down Expand Up @@ -184,15 +189,7 @@ def test_skipped_ops(self):
y = constant_op.constant(np.zeros((1, 1, 1, 1)).astype(np.float32))
# Duplicate ops are skipped, even if input values are different
gen_nn_ops.conv2d(x, y, [1, 1, 1, 1], 'SAME')
if not IsMklEnabled():
self.assertLen(self._get_new_node_defs(), 1)
else:
ndefs = self._get_new_node_defs()
if (len(ndefs) >= 1 and ndefs[0].op != ndefs[1].op):
# One of the ops got rewritten by oneDNN optimization pass
self.assertLen(ndefs, 2)
else:
self.assertLen(ndefs, 1)
self.assertLen(self._get_new_node_defs(), 1)

x = constant_op.constant(np.ones((1, 1, 1, 1, 1, 1)).astype(np.float32))
paddings = constant_op.constant(np.ones((6, 2)).astype(np.int32))
Expand Down