vulkan upload weights per layer, reduce peak ram usage (#6534)

Author: nihui

src/command.cpp

@@ -2288,6 +2288,8 @@ class VkTransferPrivate
 
     const VulkanDevice* vkdev;
 
+    uint64_t pending_upload_total;
+
     VkCommandPool compute_command_pool;
     VkCommandPool transfer_command_pool;
 
@@ -2305,6 +2307,8 @@ class VkTransferPrivate
 VkTransferPrivate::VkTransferPrivate(const VulkanDevice* _vkdev)
     : vkdev(_vkdev)
 {
+    pending_upload_total = 0;
+
     compute_command_pool = 0;
     transfer_command_pool = 0;
 
@@ -2569,6 +2573,8 @@ void VkTransfer::record_upload(const Mat& src, VkMat& dst, const Option& opt, bo
         return;
     }
 
+    d->pending_upload_total += dst.total() * dst.elemsize;
+
     if (dst.allocator->mappable)
     {
         // memcpy src_flattened to device
@@ -2843,9 +2849,65 @@ int VkTransfer::submit_and_wait()
         }
     }
 
+    d->pending_upload_total = 0;
+
     return 0;
 }
 
+int VkTransfer::reset()
+{
+    d->upload_staging_buffers.clear();
+
+    d->pending_upload_total = 0;
+
+    // reset command buffer and fence
+    {
+        VkResult ret = vkResetCommandBuffer(d->compute_command_buffer, 0);
+        if (ret != VK_SUCCESS)
+        {
+            NCNN_LOGE("vkResetCommandBuffer failed %d", ret);
+            return -1;
+        }
+    }
+    {
+        VkResult ret = vkResetFences(vkdev->vkdevice(), 1, &d->compute_command_fence);
+        if (ret != VK_SUCCESS)
+        {
+            NCNN_LOGE("vkResetFences failed %d", ret);
+            return -1;
+        }
+    }
+
+    if (!vkdev->info.unified_compute_transfer_queue())
+    {
+        {
+            VkResult ret = vkResetCommandBuffer(d->upload_command_buffer, 0);
+            if (ret != VK_SUCCESS)
+            {
+                NCNN_LOGE("vkResetCommandBuffer failed %d", ret);
+                return -1;
+            }
+        }
+        {
+            VkResult ret = vkResetFences(vkdev->vkdevice(), 1, &d->upload_command_fence);
+            if (ret != VK_SUCCESS)
+            {
+                NCNN_LOGE("vkResetFences failed %d", ret);
+                return -1;
+            }
+        }
+    }
+
+    d->begin_command_buffer();
+
+    return 0;
+}
+
+uint64_t VkTransfer::pending_upload_total() const
+{
+    return d->pending_upload_total;
+}
+
 } // namespace ncnn
 
 #endif // NCNN_VULKAN

src/command.h

@@ -97,6 +97,10 @@ class NCNN_EXPORT VkTransfer
 
     int submit_and_wait();
 
+    int reset();
+
+    uint64_t pending_upload_total() const;
+
 protected:
     const VulkanDevice* vkdev;
 

src/net.cpp

@@ -31,12 +31,6 @@ class NetPrivate
 
     Option& opt;
 
-#if NCNN_VULKAN
-
-    int upload_model();
-
-#endif // NCNN_VULKAN
-
     friend class Extractor;
     int forward_layer(int layer_index, std::vector<Mat>& blob_mats, const Option& opt) const;
 
@@ -122,52 +116,6 @@ static Option get_masked_option(const Option& opt, int featmask)
     return opt1;
 }
 
-#if NCNN_VULKAN
-int NetPrivate::upload_model()
-{
-    ncnn::VkTransfer cmd(vkdev);
-
-    // create gpu device allocator if null
-    if (!weight_vkallocator)
-    {
-        if (opt.use_weights_in_host_memory)
-        {
-            weight_vkallocator = new VkHostAllocator(vkdev);
-        }
-        else
-        {
-            weight_vkallocator = new VkWeightAllocator(vkdev);
-        }
-    }
-    if (!weight_staging_vkallocator)
-    {
-        weight_staging_vkallocator = new VkWeightStagingAllocator(vkdev);
-    }
-
-    Option opt_upload = opt;
-    opt_upload.blob_allocator = 0;
-    opt_upload.workspace_allocator = 0;
-    opt_upload.blob_vkallocator = weight_vkallocator;
-    opt_upload.workspace_vkallocator = weight_vkallocator;
-    opt_upload.staging_vkallocator = weight_staging_vkallocator;
-
-    for (size_t i = 0; i < layers.size(); i++)
-    {
-        if (layers[i]->support_vulkan)
-        {
-            int uret = layers[i]->upload_model(cmd, get_masked_option(opt_upload, layers[i]->featmask));
-            if (uret != 0)
-            {
-                NCNN_LOGE("layer upload_model %d failed", (int)i);
-                return -1;
-            }
-        }
-    }
-
-    return cmd.submit_and_wait();
-}
-#endif // NCNN_VULKAN
-
 int NetPrivate::forward_layer(int layer_index, std::vector<Mat>& blob_mats, const Option& opt) const
 {
     const Layer* layer = layers[layer_index];
@@ -1633,6 +1581,9 @@ int Net::load_model(const DataReader& dr)
     int ret = 0;
 
 #if NCNN_VULKAN
+    ncnn::VkTransfer* cmd_upload = 0;
+    Option opt_upload = opt;
+
     if (opt.use_vulkan_compute)
     {
         if (!opt.pipeline_cache)
@@ -1641,6 +1592,31 @@ int Net::load_model(const DataReader& dr)
                 d->pipeline_cache = new PipelineCache(d->vkdev);
             opt.pipeline_cache = d->pipeline_cache;
         }
+
+        cmd_upload = new ncnn::VkTransfer(d->vkdev);
+
+        // create gpu device allocator if null
+        if (!d->weight_vkallocator)
+        {
+            if (opt.use_weights_in_host_memory)
+            {
+                d->weight_vkallocator = new VkHostAllocator(d->vkdev);
+            }
+            else
+            {
+                d->weight_vkallocator = new VkWeightAllocator(d->vkdev);
+            }
+        }
+        if (!d->weight_staging_vkallocator)
+        {
+            d->weight_staging_vkallocator = new VkWeightStagingAllocator(d->vkdev);
+        }
+
+        opt_upload.blob_allocator = 0;
+        opt_upload.workspace_allocator = 0;
+        opt_upload.blob_vkallocator = d->weight_vkallocator;
+        opt_upload.workspace_vkallocator = d->weight_vkallocator;
+        opt_upload.staging_vkallocator = d->weight_staging_vkallocator;
     }
 #endif // NCNN_VULKAN
 
@@ -1682,6 +1658,36 @@ int Net::load_model(const DataReader& dr)
             ret = -1;
             break;
         }
+
+#if NCNN_VULKAN
+        if (layer->support_vulkan && opt.use_vulkan_compute && cmd_upload)
+        {
+            int uret = layer->upload_model(*cmd_upload, get_masked_option(opt_upload, layer->featmask));
+            if (uret != 0)
+            {
+#if NCNN_STRING
+                NCNN_LOGE("layer upload_model %d %s failed", i, layer->name.c_str());
+#else
+                NCNN_LOGE("layer upload_model %d failed", i);
+#endif
+                ret = -1;
+                break;
+            }
+
+            // commit as soon as we collect 256M pending
+            if (cmd_upload->pending_upload_total() > 256 * 1024 * 1024)
+            {
+                uret = cmd_upload->submit_and_wait();
+                if (uret != 0)
+                {
+                    ret = -1;
+                    break;
+                }
+
+                cmd_upload->reset();
+            }
+        }
+#endif // NCNN_VULKAN
     }
 
     if (opt.use_local_pool_allocator)
@@ -1705,9 +1711,14 @@ int Net::load_model(const DataReader& dr)
     }
 
 #if NCNN_VULKAN
-    if (ret == 0 && opt.use_vulkan_compute)
+    if (ret == 0 && opt.use_vulkan_compute && cmd_upload)
+    {
+        ret = cmd_upload->submit_and_wait();
+    }
+
+    if (cmd_upload)
     {
-        ret = d->upload_model();
+        delete cmd_upload;
     }
 #endif // NCNN_VULKAN