aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/common_runtime/threadpool_device.h
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/core/common_runtime/threadpool_device.h')
-rw-r--r--tensorflow/core/common_runtime/threadpool_device.h31
1 files changed, 31 insertions, 0 deletions
diff --git a/tensorflow/core/common_runtime/threadpool_device.h b/tensorflow/core/common_runtime/threadpool_device.h
new file mode 100644
index 0000000000..5b0347231f
--- /dev/null
+++ b/tensorflow/core/common_runtime/threadpool_device.h
@@ -0,0 +1,31 @@
+#ifndef TENSORFLOW_COMMON_RUNTIME_THREADPOOL_DEVICE_H_
+#define TENSORFLOW_COMMON_RUNTIME_THREADPOOL_DEVICE_H_
+
+#include "tensorflow/core/common_runtime/device_factory.h"
+#include "tensorflow/core/common_runtime/local_device.h"
+
+namespace tensorflow {
+
+// CPU device implementation.
+class ThreadPoolDevice : public LocalDevice {
+ public:
+ ThreadPoolDevice(const SessionOptions& options, const string& name,
+ Bytes memory_limit, BusAdjacency bus_adjacency,
+ Allocator* allocator);
+ ~ThreadPoolDevice() override;
+
+ void Compute(OpKernel* op_kernel, OpKernelContext* context) override;
+ Allocator* GetAllocator(AllocatorAttributes attr) override;
+ Status MakeTensorFromProto(const TensorProto& tensor_proto,
+ const AllocatorAttributes alloc_attrs,
+ Tensor* tensor) override;
+
+ Status Sync() override { return Status::OK(); }
+
+ private:
+ Allocator* allocator_; // Not owned
+};
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_COMMON_RUNTIME_THREADPOOL_DEVICE_H_