Add virtual_device to TPU library

PiperOrigin-RevId: 318557230
Change-Id: I5898a269e816011715c646679b3dafe0ccbcb8d5

tensorflow/core/tpu/BUILD

@@ -154,3 +154,14 @@ cc_library(
     hdrs = ["tpu_library_init_fns.inc"],
     visibility = ["//visibility:public"],
 )
+
+cc_library(
+    name = "virtual_device",
+    srcs = ["virtual_device.cc"],
+    hdrs = ["virtual_device.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:protos_all_cc",
+    ],
+)

tensorflow/core/tpu/virtual_device.cc

@@ -0,0 +1,97 @@
+/* Copyright 2020 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/tpu/virtual_device.h"
+
+#include "tensorflow/core/framework/tensor.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class VirtualDeviceContext : public DeviceContext {
+ public:
+  VirtualDeviceContext() = default;
+
+  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
+                             Tensor* device_tensor, StatusCallback done,
+                             bool sync_dst_compute) const override;
+  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
+                             StringPiece tensor_name, Device* device,
+                             Tensor* cpu_tensor, StatusCallback done) override;
+  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
+                              Tensor* output_tensor,
+                              StatusCallback done) const override;
+};
+
+void VirtualDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
+                                                 Device* device,
+                                                 Tensor* device_tensor,
+                                                 StatusCallback done,
+                                                 bool sync_dst_compute) const {
+  *device_tensor = *cpu_tensor;
+  done(Status::OK());
+}
+
+void VirtualDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
+                                                 StringPiece tensor_name,
+                                                 Device* device,
+                                                 Tensor* cpu_tensor,
+                                                 StatusCallback done) {
+  *cpu_tensor = *device_tensor;
+  done(Status::OK());
+}
+
+void VirtualDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor,
+                                                  Device* device,
+                                                  Tensor* output_tensor,
+                                                  StatusCallback done) const {
+  *output_tensor = *input_tensor;
+  done(Status::OK());
+}
+
+}  // namespace
+
+// VirtualDevice
+
+VirtualDevice::VirtualDevice(Env* env,
+                             const DeviceAttributes& device_attributes)
+    : Device(env, device_attributes) {}
+
+Status VirtualDevice::Sync() { return Status::OK(); }
+
+Allocator* VirtualDevice::GetAllocator(AllocatorAttributes attr) {
+  // Tensors always live on the host.
+  return cpu_allocator();
+}
+
+Status VirtualDevice::MakeTensorFromProto(const TensorProto& tensor_proto,
+                                          const AllocatorAttributes alloc_attrs,
+                                          Tensor* tensor) {
+  Tensor parsed(tensor_proto.dtype());
+  Allocator* allocator = cpu_allocator();
+  if (!parsed.FromProto(allocator, tensor_proto)) {
+    return errors::InvalidArgument("Cannot parse tensor from proto: ",
+                                   tensor_proto.DebugString());
+  }
+  *tensor = parsed;
+  return Status::OK();
+}
+
+Status VirtualDevice::TryGetDeviceContext(DeviceContext** out_context) {
+  *out_context = new VirtualDeviceContext;
+  return Status::OK();
+}
+
+}  // namespace tensorflow

tensorflow/core/tpu/virtual_device.h

@@ -0,0 +1,39 @@
+/* Copyright 2020 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_TPU_VIRTUAL_DEVICE_H_
+#define TENSORFLOW_CORE_TPU_VIRTUAL_DEVICE_H_
+
+#include "tensorflow/core/common_runtime/device.h"
+
+namespace tensorflow {
+
+// A dummy device that exists primarily for operator placement, without
+// corresponding directly to a piece of hardware.
+class VirtualDevice : public Device {
+ public:
+  VirtualDevice(Env* env, const DeviceAttributes& device_attributes);
+
+  Status Sync() override;
+  Allocator* GetAllocator(AllocatorAttributes attr) override;
+  Status MakeTensorFromProto(const TensorProto& tensor_proto,
+                             const AllocatorAttributes alloc_attrs,
+                             Tensor* tensor) override;
+  Status TryGetDeviceContext(DeviceContext** out_context) override;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TPU_VIRTUAL_DEVICE_H_