1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "runtime/hardware/device_context_manager.h"
18
19 namespace mindspore {
20 namespace device {
Register(const std::string & device_name,DeviceContextCreator && device_context_creator)21 void DeviceContextManager::Register(const std::string &device_name, DeviceContextCreator &&device_context_creator) {
22 if (device_context_creators_.find(device_name) == device_context_creators_.end()) {
23 (void)device_context_creators_.emplace(device_name, device_context_creator);
24 }
25 }
26
ClearDeviceContexts()27 void DeviceContextManager::ClearDeviceContexts() {
28 for (auto &iter : device_contexts_) {
29 MS_LOG(INFO) << "Release device " << iter.first;
30 MS_EXCEPTION_IF_NULL(iter.second);
31 iter.second->Destroy();
32 }
33 device_contexts_.clear();
34 }
35
GetOrCreateDeviceContext(const DeviceContextKey & device_context_key)36 DeviceContext *DeviceContextManager::GetOrCreateDeviceContext(const DeviceContextKey &device_context_key) {
37 std::string device_context_key_str = device_context_key.ToString();
38
39 auto device_context_iter = device_contexts_.find(device_context_key_str);
40 if (device_context_iter != device_contexts_.end()) {
41 return device_context_iter->second.get();
42 }
43
44 std::shared_ptr<DeviceContext> device_context;
45 auto creator_iter = device_context_creators_.find(device_context_key.device_name_);
46 if (creator_iter != device_context_creators_.end()) {
47 device_context = (creator_iter->second)(device_context_key);
48 MS_EXCEPTION_IF_NULL(device_context);
49 device_contexts_[device_context_key_str] = device_context;
50 } else {
51 MS_LOG(EXCEPTION) << "There is no device context creator for " << device_context_key.device_name_
52 << " with device id " << device_context_key.device_id_;
53 }
54 return device_context.get();
55 }
56
UpdateDeviceContextKey(const DeviceContextKey & old_key,const DeviceContextKey & new_key)57 void DeviceContextManager::UpdateDeviceContextKey(const DeviceContextKey &old_key, const DeviceContextKey &new_key) {
58 std::string old_key_str = old_key.ToString();
59 std::string new_key_str = new_key.ToString();
60
61 auto handle = device_contexts_.extract(old_key_str);
62 if (handle.empty()) {
63 MS_LOG(EXCEPTION) << "Can not find device context for: " << old_key_str;
64 }
65
66 handle.key() = new_key_str;
67 (void)device_contexts_.insert(std::move(handle));
68 }
69 } // namespace device
70 } // namespace mindspore
71