1 /**
2 * Copyright 2024 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 "pybind_api/hal/memory_py.h"
18 #include "runtime/hardware/device_context.h"
19 #include "runtime/hardware/device_context_manager.h"
20
21 namespace mindspore {
22 namespace hal {
MemoryStats(const std::string & device_target)23 py::dict MemoryStats(const std::string &device_target) {
24 auto device_ctx = device::DeviceContextManager::GetInstance().GetDeviceContext(device_target);
25 if (device_ctx == nullptr) {
26 MS_LOG(EXCEPTION) << "Device context of device " << device_target << " is not created yet.";
27 }
28
29 // Memory statistics result to be returned.
30 py::dict memory_stats;
31 py::dict commom_mem_pool_stats;
32 py::dict persistent_mem_pool_stats;
33 // Peak memory statistics.
34 // py::dict peak_mem_stats;
35
36 size_t total_mem_size = device_ctx->device_res_manager_->GetTotalMemStatistics();
37 size_t total_used_mem_size = device_ctx->device_res_manager_->GetTotalUsedMemStatistics();
38 size_t total_idle_mem_size = device_ctx->device_res_manager_->GetTotalIdleMemStatistics();
39 size_t total_eager_free_mem_size = device_ctx->device_res_manager_->GetTotalEagerFreeMemStatistics();
40 size_t used_mem_peak_size = device_ctx->device_res_manager_->GetUsedMemPeakStatistics();
41 size_t reserved_mem_peak_size = device_ctx->device_res_manager_->GetReservedMemPeakStatistics();
42 std::unordered_map<std::string, std::size_t> block_counts_stats =
43 device_ctx->device_res_manager_->GetBlockCountsStatistics();
44 std::unordered_map<std::string, std::size_t> block_unit_size_stats =
45 device_ctx->device_res_manager_->GetBlockUnitSizeStatistics();
46 std::unordered_map<device::DeviceMemPtr, std::unordered_map<std::string, size_t>> common_mem_blocks_info =
47 device_ctx->device_res_manager_->GetCommonMemBlocksInfoStatistics();
48 std::unordered_map<device::DeviceMemPtr, std::unordered_map<std::string, size_t>> persistent_mem_blocks_info =
49 device_ctx->device_res_manager_->GetPersistentMemBlocksInfoStatistics();
50
51 memory_stats["total_reserved_memory"] = total_mem_size;
52 memory_stats["total_allocatd_memory"] = total_used_mem_size;
53 memory_stats["total_idle_memory"] = total_idle_mem_size;
54 memory_stats["total_eager_free_memory"] = total_eager_free_mem_size;
55 memory_stats["max_reserved_memory"] = reserved_mem_peak_size;
56 memory_stats["max_allocated_memory"] = used_mem_peak_size;
57 commom_mem_pool_stats["block_unit_size"] = block_unit_size_stats["common_mem_pool"];
58 commom_mem_pool_stats["block_counts"] = block_counts_stats["common_mem_pool"];
59 commom_mem_pool_stats["blocks_info"] = common_mem_blocks_info;
60 persistent_mem_pool_stats["block_counts"] = block_counts_stats["persistent_mem_pool"];
61 persistent_mem_pool_stats["block_unit_size"] = block_unit_size_stats["persistent_mem_pool"];
62 persistent_mem_pool_stats["blocks_info"] = persistent_mem_blocks_info;
63 memory_stats["commom_mem_pool_stats"] = commom_mem_pool_stats;
64 memory_stats["persistent_mem_pool_stats"] = persistent_mem_pool_stats;
65 return memory_stats;
66 }
67
ResetMaxMemoryReserved(const std::string & device_target)68 void ResetMaxMemoryReserved(const std::string &device_target) {
69 auto device_ctx = device::DeviceContextManager::GetInstance().GetDeviceContext(device_target);
70 if (device_ctx == nullptr) {
71 MS_LOG(EXCEPTION) << "Device context of device " << device_target << " is not created yet.";
72 }
73
74 device_ctx->device_res_manager_->ResetMaxMemoryReserved();
75 }
76
ResetMaxMemoryAllocated(const std::string & device_target)77 void ResetMaxMemoryAllocated(const std::string &device_target) {
78 auto device_ctx = device::DeviceContextManager::GetInstance().GetDeviceContext(device_target);
79 if (device_ctx == nullptr) {
80 MS_LOG(EXCEPTION) << "Device context of device " << device_target << " is not created yet.";
81 }
82
83 device_ctx->device_res_manager_->ResetMaxMemoryAllocated();
84 }
85
RegMemory(py::module * m)86 void RegMemory(py::module *m) {
87 (void)m->def("_memory_stats", &mindspore::hal::MemoryStats, "Get memory pool's statistics.");
88 (void)m->def("_reset_max_mem_reserved", &mindspore::hal::ResetMaxMemoryReserved,
89 "Reset the maximum recorded memory reserved.");
90 (void)m->def("_reset_max_mem_allocated", &mindspore::hal::ResetMaxMemoryAllocated,
91 "Reset the maximum recorded memory allocated.");
92 }
93 } // namespace hal
94 } // namespace mindspore
95