1 /* Copyright 2015 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 #include "tensorflow/core/common_runtime/device/device_id_manager.h"
17
18 #include <unordered_map>
19
20 #include "tensorflow/core/common_runtime/device/device_id.h"
21 #include "tensorflow/core/lib/core/errors.h"
22 #include "tensorflow/core/lib/core/status.h"
23 #include "tensorflow/core/platform/logging.h"
24 #include "tensorflow/core/platform/macros.h"
25 #include "tensorflow/core/platform/mutex.h"
26
27 namespace tensorflow {
28 namespace {
29 // Manages the map between TfDeviceId and platform device id.
30 class TfToPlatformDeviceIdMap {
31 public:
singleton()32 static TfToPlatformDeviceIdMap* singleton() {
33 static auto* id_map = new TfToPlatformDeviceIdMap;
34 return id_map;
35 }
36
Insert(const DeviceType & type,TfDeviceId tf_device_id,PlatformDeviceId platform_device_id)37 Status Insert(const DeviceType& type, TfDeviceId tf_device_id,
38 PlatformDeviceId platform_device_id) TF_LOCKS_EXCLUDED(mu_) {
39 std::pair<IdMapType::iterator, bool> result;
40 {
41 mutex_lock lock(mu_);
42 TypeIdMapType::iterator device_id_map_iter =
43 id_map_.insert({type.type_string(), IdMapType()}).first;
44 result = device_id_map_iter->second.insert(
45 {tf_device_id.value(), platform_device_id.value()});
46 }
47 if (!result.second && platform_device_id.value() != result.first->second) {
48 return errors::AlreadyExists(
49 "TensorFlow device (", type, ":", tf_device_id.value(),
50 ") is being mapped to multiple devices (", platform_device_id.value(),
51 " now, and ", result.first->second,
52 " previously), which is not supported. "
53 "This may be the result of providing different ",
54 type, " configurations (ConfigProto.gpu_options, for example ",
55 "different visible_device_list) when creating multiple Sessions in ",
56 "the same process. This is not currently supported, see ",
57 "https://github.com/tensorflow/tensorflow/issues/19083");
58 }
59 return Status::OK();
60 }
61
Find(const DeviceType & type,TfDeviceId tf_device_id,PlatformDeviceId * platform_device_id) const62 bool Find(const DeviceType& type, TfDeviceId tf_device_id,
63 PlatformDeviceId* platform_device_id) const TF_LOCKS_EXCLUDED(mu_) {
64 // TODO(mrry): Consider replacing this with an atomic `is_initialized` bit,
65 // to avoid writing to a shared cache line in the tf_shared_lock.
66 tf_shared_lock lock(mu_);
67 auto type_id_map_iter = id_map_.find(type.type_string());
68 if (type_id_map_iter == id_map_.end()) return false;
69 auto id_map_iter = type_id_map_iter->second.find(tf_device_id.value());
70 if (id_map_iter == type_id_map_iter->second.end()) return false;
71 *platform_device_id = id_map_iter->second;
72 return true;
73 }
74
75 private:
76 TfToPlatformDeviceIdMap() = default;
77
TestOnlyReset()78 void TestOnlyReset() TF_LOCKS_EXCLUDED(mu_) {
79 mutex_lock lock(mu_);
80 id_map_.clear();
81 }
82
83 // Map from physical device id to platform device id.
84 using IdMapType = std::unordered_map<int32, int32>;
85 // Map from DeviceType to IdMapType.
86 // We use std::string instead of DeviceType because the key should
87 // be default-initializable.
88 using TypeIdMapType = std::unordered_map<std::string, IdMapType>;
89 mutable mutex mu_;
90 TypeIdMapType id_map_ TF_GUARDED_BY(mu_);
91
92 friend class ::tensorflow::DeviceIdManager;
93 TF_DISALLOW_COPY_AND_ASSIGN(TfToPlatformDeviceIdMap);
94 };
95 } // namespace
96
InsertTfPlatformDeviceIdPair(const DeviceType & type,TfDeviceId tf_device_id,PlatformDeviceId platform_device_id)97 Status DeviceIdManager::InsertTfPlatformDeviceIdPair(
98 const DeviceType& type, TfDeviceId tf_device_id,
99 PlatformDeviceId platform_device_id) {
100 return TfToPlatformDeviceIdMap::singleton()->Insert(type, tf_device_id,
101 platform_device_id);
102 }
103
TfToPlatformDeviceId(const DeviceType & type,TfDeviceId tf_device_id,PlatformDeviceId * platform_device_id)104 Status DeviceIdManager::TfToPlatformDeviceId(
105 const DeviceType& type, TfDeviceId tf_device_id,
106 PlatformDeviceId* platform_device_id) {
107 if (TfToPlatformDeviceIdMap::singleton()->Find(type, tf_device_id,
108 platform_device_id)) {
109 return Status::OK();
110 }
111 return errors::NotFound("TensorFlow device ", type, ":", tf_device_id.value(),
112 " was not registered");
113 }
114
TestOnlyReset()115 void DeviceIdManager::TestOnlyReset() {
116 TfToPlatformDeviceIdMap::singleton()->TestOnlyReset();
117 }
118
119 } // namespace tensorflow
120