1 | /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_COMPOSITE_DEVICE_H_ |
17 | #define TENSORFLOW_CORE_COMMON_RUNTIME_COMPOSITE_DEVICE_H_ |
18 | |
19 | #include "absl/strings/string_view.h" |
20 | #include "tensorflow/core/common_runtime/device.h" |
21 | #include "tensorflow/core/framework/allocator.h" |
22 | #include "tensorflow/core/framework/device_attributes.pb.h" |
23 | #include "tensorflow/core/lib/core/errors.h" |
24 | #include "tensorflow/core/lib/core/status.h" |
25 | |
26 | namespace tensorflow { |
27 | |
28 | extern const char* const kCompositeDeviceType; |
29 | |
30 | // A virtual device which represents a set of devices. We don't execute any |
31 | // op on this virtial device. |
32 | class CompositeDevice : public Device { |
33 | public: |
34 | Status Sync() override { |
35 | return errors::Internal( |
36 | "Sync() should never been invoked on CompositeDevice." ); |
37 | } |
38 | |
39 | Allocator* GetAllocator(AllocatorAttributes) override { return nullptr; } |
40 | |
41 | const std::vector<string>* underlying_devices() const { |
42 | return &underlying_devices_; |
43 | } |
44 | |
45 | // Helper for creating a CompositeDevice on the same task as the given host |
46 | // CPU. |
47 | static std::unique_ptr<CompositeDevice> MakeDevice( |
48 | const std::vector<string>& underlying_devices, const int unique_device_id, |
49 | const DeviceNameUtils::ParsedName& host_name, Status* status); |
50 | |
51 | // Helper for creating a CompositeDevice with the given device name. |
52 | static std::unique_ptr<CompositeDevice> MakeDevice( |
53 | const std::vector<string>& underlying_devices, const string& device_name, |
54 | Status* status); |
55 | |
56 | bool IsRemoteCallAllowed() const override { return false; } |
57 | |
58 | private: |
59 | CompositeDevice(const DeviceAttributes& device_attributes, |
60 | const std::vector<string>& underlying_devices) |
61 | : Device(/*env=*/nullptr, device_attributes), |
62 | underlying_devices_(underlying_devices) {} |
63 | |
64 | const std::vector<string> underlying_devices_; |
65 | |
66 | TF_DISALLOW_COPY_AND_ASSIGN(CompositeDevice); |
67 | }; |
68 | |
69 | } // namespace tensorflow |
70 | |
71 | #endif // TENSORFLOW_CORE_COMMON_RUNTIME_COMPOSITE_DEVICE_H_ |
72 | |