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/gpu/gpu_id_manager.h"
17
18 #include <unordered_map>
19
20 #include "tensorflow/core/common_runtime/gpu/gpu_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 TfGpuId and platform GPU id.
30 class TfToPlatformGpuIdMap {
31 public:
singleton()32 static TfToPlatformGpuIdMap* singleton() {
33 static auto* id_map = new TfToPlatformGpuIdMap;
34 return id_map;
35 }
36
Insert(TfGpuId tf_gpu_id,PlatformGpuId platform_gpu_id)37 Status Insert(TfGpuId tf_gpu_id, PlatformGpuId platform_gpu_id)
38 LOCKS_EXCLUDED(mu_) {
39 std::pair<IdMapType::iterator, bool> result;
40 {
41 mutex_lock lock(mu_);
42 result = id_map_.insert({tf_gpu_id.value(), platform_gpu_id.value()});
43 }
44 if (!result.second && platform_gpu_id.value() != result.first->second) {
45 return errors::AlreadyExists(
46 "TensorFlow device (GPU:", tf_gpu_id.value(),
47 ") is being mapped to "
48 "multiple CUDA devices (",
49 platform_gpu_id.value(), " now, and ", result.first->second,
50 " previously), which is not supported. "
51 "This may be the result of providing different GPU configurations "
52 "(ConfigProto.gpu_options, for example different visible_device_list)"
53 " when creating multiple Sessions in the same process. This is not "
54 " currently supported, see "
55 "https://github.com/tensorflow/tensorflow/issues/19083");
56 }
57 return Status::OK();
58 }
59
Find(TfGpuId tf_gpu_id,PlatformGpuId * platform_gpu_id) const60 bool Find(TfGpuId tf_gpu_id, PlatformGpuId* platform_gpu_id) const
61 LOCKS_EXCLUDED(mu_) {
62 mutex_lock lock(mu_);
63 auto result = id_map_.find(tf_gpu_id.value());
64 if (result == id_map_.end()) return false;
65 *platform_gpu_id = result->second;
66 return true;
67 }
68
69 private:
70 TfToPlatformGpuIdMap() = default;
71
TestOnlyReset()72 void TestOnlyReset() LOCKS_EXCLUDED(mu_) {
73 mutex_lock lock(mu_);
74 id_map_.clear();
75 }
76
77 using IdMapType = std::unordered_map<int32, int32>;
78 mutable mutex mu_;
79 IdMapType id_map_ GUARDED_BY(mu_);
80
81 friend class ::tensorflow::GpuIdManager;
82 TF_DISALLOW_COPY_AND_ASSIGN(TfToPlatformGpuIdMap);
83 };
84 } // namespace
85
InsertTfPlatformGpuIdPair(TfGpuId tf_gpu_id,PlatformGpuId platform_gpu_id)86 Status GpuIdManager::InsertTfPlatformGpuIdPair(TfGpuId tf_gpu_id,
87 PlatformGpuId platform_gpu_id) {
88 return TfToPlatformGpuIdMap::singleton()->Insert(tf_gpu_id, platform_gpu_id);
89 }
90
TfToPlatformGpuId(TfGpuId tf_gpu_id,PlatformGpuId * platform_gpu_id)91 Status GpuIdManager::TfToPlatformGpuId(TfGpuId tf_gpu_id,
92 PlatformGpuId* platform_gpu_id) {
93 if (TfToPlatformGpuIdMap::singleton()->Find(tf_gpu_id, platform_gpu_id)) {
94 return Status::OK();
95 }
96 return errors::NotFound("TensorFlow device GPU:", tf_gpu_id.value(),
97 " was not registered");
98 }
99
TestOnlyReset()100 void GpuIdManager::TestOnlyReset() {
101 TfToPlatformGpuIdMap::singleton()->TestOnlyReset();
102 }
103
104 } // namespace tensorflow
105