[GitHub] [tvm] tqchen commented on a change in pull request #7833: [Vulkan] Support uniform buffer object for passing many scalar arguments (Take 2)

[GitBox <gi...@apache.org>]

tqchen commented on a change in pull request #7833:
URL: https://github.com/apache/tvm/pull/7833#discussion_r612418998



##########
File path: src/runtime/vulkan/vulkan_common.h
##########
@@ -35,6 +35,10 @@ namespace tvm {
 namespace runtime {
 namespace vulkan {
 
+#define kMaxPushConstantsBytes 128
+
+enum ShaderMetaDataKind { kUSE_UBO = 0 };

Review comment:
       add `ShaderMetaFlagMask` document, kUseUBOMask to indicate it is a bitmask

##########
File path: src/runtime/vulkan/vulkan_common.h
##########
@@ -35,6 +35,10 @@ namespace tvm {
 namespace runtime {
 namespace vulkan {
 
+#define kMaxPushConstantsBytes 128

Review comment:
       use const int instead

##########
File path: src/runtime/vulkan/vulkan.cc
##########
@@ -1115,20 +1194,31 @@ void VulkanWrappedFunc::operator()(TVMArgs args, TVMRetValue* rv,
       write_descriptor_sets[i].dstBinding = i;
       write_descriptor_sets[i].dstArrayElement = 0;
       write_descriptor_sets[i].descriptorCount = 1;
-      write_descriptor_sets[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
       write_descriptor_sets[i].pImageInfo = 0;
       write_descriptor_sets[i].pBufferInfo = &(descriptor_buffers[i]);
       write_descriptor_sets[i].pTexelBufferView = 0;
+
+      if (pipeline->use_ubo && i == write_descriptor_sets.size() - 1) {
+        // The last binding is for UBO
+        write_descriptor_sets[i].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+      } else {
+        write_descriptor_sets[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+      }
     }
     vkUpdateDescriptorSets(vctx.device, write_descriptor_sets.size(), write_descriptor_sets.data(),
                            0, 0);
   };
-  const auto& deferred_kernel = [pipeline, wl, pack_args_storage](VulkanStreamState* state) {
+  const auto& deferred_kernel = [this, pipeline, wl, pack_args_storage, nbytes_scalars,
+                                 device_id](VulkanStreamState* state) {
+    if (pipeline->use_ubo) {
+      auto ubo = VulkanThreadEntry::ThreadLocal()->GetUniformBuffer(device_id, nbytes_scalars);

Review comment:
       move next to vkCmdPushConstants
   
   ```
   if (packed_arg_storage.size() != 0) {
      if (pipeline->use_ubo) {
        ubo logic 
     } else {
        push constant
     }
   }
   ```

##########
File path: src/runtime/vulkan/vulkan.cc
##########
@@ -1076,15 +1141,29 @@ void VulkanWrappedFunc::operator()(TVMArgs args, TVMRetValue* rv,
     binfo.range = VK_WHOLE_SIZE;
     descriptor_buffers[i] = binfo;
   }
+  const size_t nbytes_scalars = num_pack_args_ * sizeof(ArgUnion64);
+  if (pipeline->use_ubo) {
+    auto ubo = VulkanThreadEntry::ThreadLocal()->GetUniformBuffer(device_id, nbytes_scalars);
+    CHECK(ubo->host_addr) << "The UBO host buffer is not allocated";
+    VkDescriptorBufferInfo binfo;
+    binfo.buffer = ubo->vk_buf->buffer;
+    binfo.offset = 0;
+    binfo.range = VK_WHOLE_SIZE;
+    descriptor_buffers.push_back(binfo);
+  }
   if (vctx.UseImmediate()) {
     // Can safely capture by reference as this lambda is immediately executed on the calling thread.
+    if (pipeline->use_ubo) {
+      auto ubo = VulkanThreadEntry::ThreadLocal()->GetUniformBuffer(device_id, nbytes_scalars);

Review comment:
       move this next to the vkCmdPushConstants so it is clear there is a if else relation in two kinds of buffer handling(for readability)

##########
File path: src/runtime/vulkan/vulkan.cc
##########
@@ -60,35 +91,21 @@ class VulkanThreadEntry {
     pool.reset();
     streams_.clear();
     for (const auto& kv : staging_buffers_) {
-      if (!kv.second) {
-        continue;
-      }
-      auto& buf = *(kv.second);
-      if (buf.host_addr != nullptr) {
-        vkUnmapMemory(buf.device, buf.memory);
-      }
-      if (buf.memory != VK_NULL_HANDLE) {
-        vkFreeMemory(buf.device, buf.memory, nullptr);
-      }
-      if (buf.buffer != VK_NULL_HANDLE) {
-        vkDestroyBuffer(buf.device, buf.buffer, nullptr);
-      }
+      DeleteHostVisibleBuffer(kv.second.get());

Review comment:
       we need to trigger Synchronize when deleting an old Array(in the caller), stream memcpy does that explicitly https://github.com/apache/tvm/blob/e8369881ef69a46eaafd143a1c9bfa0eba2eb003/src/runtime/vulkan/vulkan.cc#L362
   
   This is to ensure that old tasks that uses the smaller staging UBO can get flushed in deferred mode.
   
   




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org