From e3d02acd52ca4604b17b73d0aea3bb9097876f39 Mon Sep 17 00:00:00 2001
From: zhangmeng <775834166@qq.com>
Date: 星期日, 17 十一月 2019 12:28:53 +0800
Subject: [PATCH] update

---
 csrc/common/gpu/info.cpp |   31 +++++++++++++++++++++++++++++++
 1 files changed, 31 insertions(+), 0 deletions(-)

diff --git a/csrc/common/gpu/info.cpp b/csrc/common/gpu/info.cpp
index d38a8ae..d19de83 100644
--- a/csrc/common/gpu/info.cpp
+++ b/csrc/common/gpu/info.cpp
@@ -369,6 +369,37 @@
     return suitable_gpu;
 }
 
+int getGPUPrior(const int need, const int reserved, const int lastChoice){
+    nvGpuInfo_t gpu_info;
+
+    int ret = get_gpu_info(&gpu_info);
+    if(!ret){
+        if (gpu_info.device_count == 0) return -1;
+
+        int suitable_gpu = -1;
+        int mem_idle = need;
+        for(int i = 0; i < gpu_info.device_count; i++){
+            if (i != lastChoice){
+                int mem_free = gpu_info.devices[i].memory_free >> 20;
+                if(mem_free > mem_idle){
+                    mem_idle = mem_free;
+                    suitable_gpu = i;
+                }
+            }
+        }
+        if (suitable_gpu != -1){
+            return suitable_gpu;
+        }else{
+            if (gpu_info.device_count <= lastChoice) return -1;
+            int mem_free = (gpu_info.devices[lastChoice].memory_free >> 20) - reserved;
+            if(mem_free > need){
+                return lastChoice;
+            }
+        }
+    }
+    return -1;
+}
+
 int getGPU(const int need){
     nvGpuInfo_t gpu_buf;
 

--
Gitblit v1.8.0