aboutsummaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
Diffstat
-rw-r--r--tensorflow/core/graph/testlib.cc13
-rw-r--r--tensorflow/core/kernels/BUILD7
-rw-r--r--tensorflow/core/kernels/constant_op.cc38
-rw-r--r--tensorflow/core/kernels/constant_op.h20
-rw-r--r--tensorflow/core/kernels/host_constant_op.cc78
-rw-r--r--tensorflow/core/kernels/host_constant_op.h42
6 files changed, 68 insertions, 130 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.
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 46c234d057..7246de147a 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -630,7 +630,6 @@ cc_library(
":gather_nd_op",
":gather_op",
":guarantee_const_op",
- ":host_constant_op",
":identity_n_op",
":identity_op",
":inplace_ops",
@@ -705,12 +704,6 @@ tf_kernel_library(
)
tf_kernel_library(
- name = "host_constant_op",
- prefix = "host_constant_op",
- deps = ARRAY_DEPS,
-)
-
-tf_kernel_library(
name = "diag_op",
prefix = "diag_op",
deps = ARRAY_DEPS,
diff --git a/tensorflow/core/kernels/constant_op.cc b/tensorflow/core/kernels/constant_op.cc
index 375819a8a2..a888422d49 100644
--- a/tensorflow/core/kernels/constant_op.cc
+++ b/tensorflow/core/kernels/constant_op.cc
@@ -140,6 +140,44 @@ REGISTER_SYCL_KERNEL(SYCL, bool);
#undef REGISTER_SYCL_KERNEL
#endif
+HostConstantOp::HostConstantOp(OpKernelConstruction* ctx)
+ : OpKernel(ctx), tensor_(ctx->output_type(0)) {
+ const TensorProto* proto = nullptr;
+ AllocatorAttributes alloc_attr;
+ alloc_attr.set_on_host(true);
+ OP_REQUIRES_OK(ctx, ctx->GetAttr("value", &proto));
+ OP_REQUIRES_OK(
+ ctx, ctx->device()->MakeTensorFromProto(*proto, alloc_attr, &tensor_));
+ OP_REQUIRES(
+ ctx, ctx->output_type(0) == tensor_.dtype(),
+ errors::InvalidArgument("Type mismatch between value (",
+ DataTypeString(tensor_.dtype()), ") and dtype (",
+ DataTypeString(ctx->output_type(0)), ")"));
+}
+
+void HostConstantOp::Compute(OpKernelContext* ctx) {
+ ctx->set_output(0, tensor_);
+}
+
+#if GOOGLE_CUDA
+// A special GPU kernel for int32.
+// TODO(b/25387198): Also enable int32 in device memory. This kernel
+// registration requires all int32 inputs and outputs to be in host memory.
+REGISTER_KERNEL_BUILDER(Name("Const")
+ .Device(DEVICE_GPU)
+ .HostMemory("output")
+ .TypeConstraint<int32>("dtype"),
+ HostConstantOp);
+#endif
+
+#ifdef TENSORFLOW_USE_SYCL
+REGISTER_KERNEL_BUILDER(Name("Const")
+ .Device(DEVICE_SYCL)
+ .HostMemory("output")
+ .TypeConstraint<int32>("dtype"),
+ HostConstantOp);
+#endif // TENSORFLOW_USE_SYCL
+
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
#ifdef TENSORFLOW_USE_SYCL
diff --git a/tensorflow/core/kernels/constant_op.h b/tensorflow/core/kernels/constant_op.h
index 77ba441863..b98153e347 100644
--- a/tensorflow/core/kernels/constant_op.h
+++ b/tensorflow/core/kernels/constant_op.h
@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CORE_KERNELS_CONSTANT_OP_H_
-#define TENSORFLOW_CORE_KERNELS_CONSTANT_OP_H_
+#ifndef TENSORFLOW_KERNELS_CONSTANT_OP_H_
+#define TENSORFLOW_KERNELS_CONSTANT_OP_H_
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/op_kernel.h"
@@ -36,6 +36,20 @@ class ConstantOp : public OpKernel {
TF_DISALLOW_COPY_AND_ASSIGN(ConstantOp);
};
+// HostConstantOp differs from ConstantOp in that its output is always
+// in host memory.
+class HostConstantOp : public OpKernel {
+ public:
+ explicit HostConstantOp(OpKernelConstruction* ctx);
+ void Compute(OpKernelContext* ctx) override;
+ bool IsExpensive() override { return false; }
+ ~HostConstantOp() override {}
+
+ private:
+ Tensor tensor_;
+ TF_DISALLOW_COPY_AND_ASSIGN(HostConstantOp);
+};
+
class PlaceholderOp : public OpKernel {
public:
explicit PlaceholderOp(OpKernelConstruction* ctx);
@@ -47,4 +61,4 @@ class PlaceholderOp : public OpKernel {
} // namespace tensorflow
-#endif // TENSORFLOW_CORE_KERNELS_CONSTANT_OP_H_
+#endif // TENSORFLOW_KERNELS_CONSTANT_OP_H_
diff --git a/tensorflow/core/kernels/host_constant_op.cc b/tensorflow/core/kernels/host_constant_op.cc
deleted file mode 100644
index d08a7c9bd2..0000000000
--- a/tensorflow/core/kernels/host_constant_op.cc
+++ /dev/null
@@ -1,78 +0,0 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#include "tensorflow/core/kernels/host_constant_op.h"
-
-#include "tensorflow/core/framework/allocator.h"
-#include "tensorflow/core/framework/op_kernel.h"
-#include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/platform/macros.h"
-
-namespace tensorflow {
-
-_HostConstantOp::_HostConstantOp(OpKernelConstruction* ctx)
- : OpKernel(ctx), tensor_(ctx->output_type(0)) {
- const TensorProto* proto = nullptr;
- AllocatorAttributes alloc_attr;
- alloc_attr.set_on_host(true);
- OP_REQUIRES_OK(ctx, ctx->GetAttr("value", &proto));
- OP_REQUIRES_OK(
- ctx, ctx->device()->MakeTensorFromProto(*proto, alloc_attr, &tensor_));
- OP_REQUIRES(
- ctx, ctx->output_type(0) == tensor_.dtype(),
- errors::InvalidArgument("Type mismatch between value (",
- DataTypeString(tensor_.dtype()), ") and dtype (",
- DataTypeString(ctx->output_type(0)), ")"));
-}
-
-void _HostConstantOp::Compute(OpKernelContext* ctx) {
- ctx->set_output(0, tensor_);
-}
-
-#if GOOGLE_CUDA
-// A special GPU kernel for int32.
-// TODO(b/25387198): Also enable int32 in device memory. This kernel
-// registration requires all int32 inputs and outputs to be in host memory.
-REGISTER_KERNEL_BUILDER(Name("Const")
- .Device(DEVICE_GPU)
- .HostMemory("output")
- .TypeConstraint<int32>("dtype"),
- _HostConstantOp);
-#endif
-
-#ifdef TENSORFLOW_USE_SYCL
-REGISTER_KERNEL_BUILDER(Name("Const")
- .Device(DEVICE_SYCL)
- .HostMemory("output")
- .TypeConstraint<int32>("dtype"),
- _HostConstantOp);
-#endif // TENSORFLOW_USE_SYCL
-
-// HostConst: forced to generate output on the host.
-// Only used in tests; 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
-
-} // end namespace tensorflow
-
diff --git a/tensorflow/core/kernels/host_constant_op.h b/tensorflow/core/kernels/host_constant_op.h
deleted file mode 100644
index 1b887ea1aa..0000000000
--- a/tensorflow/core/kernels/host_constant_op.h
+++ /dev/null
@@ -1,42 +0,0 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#ifndef TENSORFLOW_CORE_KERNELS_HOST_CONSTANT_OP_H_
-#define TENSORFLOW_CORE_KERNELS_HOST_CONSTANT_OP_H_
-
-#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
-#include "tensorflow/core/framework/op_kernel.h"
-#include "tensorflow/core/framework/tensor_types.h"
-#include "tensorflow/core/platform/macros.h"
-
-namespace tensorflow {
-
-// HostConstantOp differs from ConstantOp in that its output is always
-// in host memory.
-class _HostConstantOp : public OpKernel {
- public:
- explicit _HostConstantOp(OpKernelConstruction* ctx);
- void Compute(OpKernelContext* ctx) override;
- bool IsExpensive() override { return false; }
- ~_HostConstantOp() override {}
-
- private:
- Tensor tensor_;
- TF_DISALLOW_COPY_AND_ASSIGN(_HostConstantOp);
-};
-
-} // namespace tensorflow
-
-#endif // TENSORFLOW_CORE_KERNELS_HOST_CONSTANT_OP_H_
generated by cgit v1.2.3 (git 2.39.1) at 2024-06-26 15:54:05 +0000