From a325b0ae208b3295dca574f93c214f7c924cea11 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 15 Mar 2018 14:40:12 -0700 Subject: [PATCH] Adds a GPU kernel registration for PlaceholderWithDefault, so we can avoid the issue of using it with a registered GPU device without soft placement. PiperOrigin-RevId: 189248024 --- tensorflow/core/kernels/identity_op.cc | 34 ++++++++++++++-------- tensorflow/python/kernel_tests/constant_op_test.py | 9 +++--- 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/tensorflow/core/kernels/identity_op.cc b/tensorflow/core/kernels/identity_op.cc index a18a72c..b5603fe 100644 --- a/tensorflow/core/kernels/identity_op.cc +++ b/tensorflow/core/kernels/identity_op.cc @@ -101,6 +101,10 @@ REGISTER_SYCL_HOST_KERNEL(bool); REGISTER_KERNEL_BUILDER(Name("DebugGradientIdentity") \ .Device(DEVICE_GPU) \ .TypeConstraint("T"), \ + IdentityOp); \ + REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault") \ + .Device(DEVICE_GPU) \ + .TypeConstraint("dtype"), \ IdentityOp) TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); @@ -112,18 +116,24 @@ REGISTER_GPU_KERNEL(Variant); // A special GPU kernel for int32 and bool. // TODO(b/25387198): Also enable int32 in device memory. This kernel // registration requires all int32 inputs and outputs to be in host memory. -#define REGISTER_GPU_HOST_KERNEL(type) \ - REGISTER_KERNEL_BUILDER(Name("Identity") \ - .Device(DEVICE_GPU) \ - .HostMemory("input") \ - .HostMemory("output") \ - .TypeConstraint("T"), \ - IdentityOp); \ - REGISTER_KERNEL_BUILDER(Name("RefIdentity") \ - .Device(DEVICE_GPU) \ - .HostMemory("input") \ - .HostMemory("output") \ - .TypeConstraint("T"), \ +#define REGISTER_GPU_HOST_KERNEL(type) \ + REGISTER_KERNEL_BUILDER(Name("Identity") \ + .Device(DEVICE_GPU) \ + .HostMemory("input") \ + .HostMemory("output") \ + .TypeConstraint("T"), \ + IdentityOp); \ + REGISTER_KERNEL_BUILDER(Name("RefIdentity") \ + .Device(DEVICE_GPU) \ + .HostMemory("input") \ + .HostMemory("output") \ + .TypeConstraint("T"), \ + IdentityOp); \ + REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault") \ + .Device(DEVICE_GPU) \ + .HostMemory("input") \ + .HostMemory("output") \ + .TypeConstraint("dtype"), \ IdentityOp) REGISTER_GPU_HOST_KERNEL(int32); diff --git a/tensorflow/python/kernel_tests/constant_op_test.py b/tensorflow/python/kernel_tests/constant_op_test.py index ffbdb0e..18796f7 100644 --- a/tensorflow/python/kernel_tests/constant_op_test.py +++ b/tensorflow/python/kernel_tests/constant_op_test.py @@ -881,7 +881,7 @@ versions { class PlaceholderWithDefaultTest(test.TestCase): def testFullShape(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): p = array_ops.placeholder_with_default([[2, 2], [2, 2]], shape=[2, 2]) a = array_ops.identity(p) self.assertAllEqual([[2, 2], [2, 2]], a.eval()) @@ -892,7 +892,7 @@ class PlaceholderWithDefaultTest(test.TestCase): a.eval(feed_dict={p: [[6, 6, 6], [6, 6, 6]]}) def testPartialShape(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): p = array_ops.placeholder_with_default([1, 2, 3], shape=[None]) a = array_ops.identity(p) self.assertAllEqual([1, 2, 3], a.eval()) @@ -902,7 +902,7 @@ class PlaceholderWithDefaultTest(test.TestCase): a.eval(feed_dict={p: [[2, 2], [2, 2]]}) def testNoShape(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): p = array_ops.placeholder_with_default([17], shape=None) a = array_ops.identity(p) self.assertAllEqual([17], a.eval()) @@ -911,11 +911,12 @@ class PlaceholderWithDefaultTest(test.TestCase): [[3, 3], [3, 3]], a.eval(feed_dict={p: [[3, 3], [3, 3]]})) def testGradient(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): x = array_ops.placeholder(dtypes_lib.float32, [5, 7]) y = array_ops.placeholder_with_default(x, None) err = gradient_checker.compute_gradient_error(x, [5, 7], y, [5, 7]) self.assertLess(err, 1e-3) + if __name__ == "__main__": test.main() -- 2.7.4