From 5f9d382856a7933a1c086bced9f9dd9121087713 Mon Sep 17 00:00:00 2001 From: lichun Date: Fri, 30 Jul 2021 11:36:30 +0800 Subject: [PATCH] get hbm and ddr size adaptively --- ge/graph/manager/graph_var_manager.cc | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/ge/graph/manager/graph_var_manager.cc b/ge/graph/manager/graph_var_manager.cc index 9bc412cb..33571578 100755 --- a/ge/graph/manager/graph_var_manager.cc +++ b/ge/graph/manager/graph_var_manager.cc @@ -826,12 +826,20 @@ Status VarManager::GetTotalMemorySize(size_t &total_mem_size) { return RT_FAILED; } size_t free_mem = 0; - rt_ret = rtMemGetInfo(&free_mem, &total_mem_size); + rt_ret = rtMemGetInfoEx(RT_MEMORYINFO_HBM, &free_mem, &total_mem_size); if (rt_ret != RT_ERROR_NONE) { REPORT_CALL_ERROR("E19999", "Call rtMemGetInfo failed, ret:0x%X", rt_ret); GELOGE(RT_FAILED, "[Call][RtMemGetInfo] failed, ret:0x%X", rt_ret); return RT_FAILED; } + if (total_mem_size == 0) { + rt_ret = rtMemGetInfoEx(RT_MEMORYINFO_DDR, &free_mem, &total_mem_size); + if (rt_ret != RT_ERROR_NONE) { + REPORT_CALL_ERROR("E19999", "Call rtMemGetInfo failed, ret:0x%X", rt_ret); + GELOGE(RT_FAILED, "[Call][RtMemGetInfo] failed, ret:0x%X", rt_ret); + return RT_FAILED; + } + } rt_ret = rtDeviceReset(GetContext().DeviceId()); if (rt_ret != RT_ERROR_NONE) { REPORT_CALL_ERROR("E19999", "Call rtDeviceReset failed, device_id:%u, ret:0x%X",