fix potential out of vram in directory processing

nihui · Mar 18, 2022 · f1271a2 · f1271a2
1 parent 3fe2504
commit f1271a2
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/main.cpp b/src/main.cpp
@@ -809,6 +809,7 @@ int main(int argc, char** argv)
     }
 
     int total_jobs_proc = 0;
+    int jobs_proc_per_gpu[16] = {0};
     for (int i=0; i<use_gpu_count; i++)
     {
         if (gpuid[i] == -1)
@@ -819,6 +820,7 @@ int main(int argc, char** argv)
         else
         {
             total_jobs_proc += jobs_proc[i];
+            jobs_proc_per_gpu[gpuid[i]] += jobs_proc[i];
         }
     }
 
@@ -836,6 +838,12 @@ int main(int argc, char** argv)
 
         uint32_t heap_budget = ncnn::get_gpu_device(gpuid[i])->get_heap_budget();
 
+        if (path_is_directory(inputpath) && path_is_directory(outputpath))
+        {
+            // multiple gpu jobs share the same heap
+            heap_budget /= jobs_proc_per_gpu[gpuid[i]];
+        }
+
         // more fine-grained tilesize policy here
         if (model.find(PATHSTR("models-nose")) != path_t::npos || model.find(PATHSTR("models-se")) != path_t::npos)
         {