diff options
| author |  Gunhan Gulsoy <gunan@google.com> | 2018-08-14 14:41:38 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-14 14:45:13 -0700 |
| commit | a00fba38681fa28d6d20cd730fe362864b819a0d (patch) | |
| tree | 4b271a00efedd7ec399b9f4ac24baed657d23462 /tensorflow/core/graph | |
| parent | 3f40e007e3406f7145b2398ff14371adca3ebd1c (diff) | |
Automated rollback of commit cea262e16a004d73295259c42f21e2655da3df13
PiperOrigin-RevId: 208716358
Diffstat (limited to 'tensorflow/core/graph')
| -rw-r--r-- | tensorflow/core/graph/testlib.cc | 13 |
1 files changed, 13 insertions, 0 deletions
diff --git a/tensorflow/core/graph/testlib.cc b/tensorflow/core/graph/testlib.cc index 934ed70e89..67b252cb6c 100644 --- a/tensorflow/core/graph/testlib.cc +++ b/tensorflow/core/graph/testlib.cc @@ -21,15 +21,28 @@ limitations under the License. #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/kernels/constant_op.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/logging.h" namespace tensorflow { +// HostConst: forced to generate output on the host. +// Only used by testlib; no op is registered for this kernel +// externally (i.e., in array_ops.cc) +REGISTER_KERNEL_BUILDER(Name("HostConst").Device(DEVICE_CPU), HostConstantOp); +REGISTER_KERNEL_BUILDER( + Name("HostConst").Device(DEVICE_GPU).HostMemory("output"), HostConstantOp); +#ifdef TENSORFLOW_USE_SYCL +REGISTER_KERNEL_BUILDER( + Name("HostConst").Device(DEVICE_SYCL).HostMemory("output"), HostConstantOp); +#endif // TENSORFLOW_USE_SYCL + // Register the HostConst Op // Returns a constant tensor on the host. Useful for writing C++ tests // and benchmarks which run on GPU but require arguments pinned to the host. |