Enable zero-size tensors (#3640)

Summary:

As title.

The approach is slightly different than in PyTorch Vulkan. Instead of binding no memory, we make a small allocation. The reason for this change is to account for the possibility that some zero size tensors are used as input but the output is not zero size. In that case we still need to be able to bind the zero size tensor to a shader.

Differential Revision: D57450473

Author: SS-JIA

backends/vulkan/runtime/api/memory/Buffer.cpp

@@ -38,16 +38,18 @@ VulkanBuffer::VulkanBuffer(
       memory_{},
       owns_memory_(allocate_memory),
       handle_(VK_NULL_HANDLE) {
-  // Only allocate memory if the buffer has non-zero size
+  // If the buffer size is 0, allocate a buffer with a size of 1 byte. This is
+  // to ensure that there will be some resource that can be bound to a shader.
   if (size == 0) {
-    return;
+    buffer_properties_.size = 1u;
+    buffer_properties_.mem_range = 1u;
   }
 
   const VkBufferCreateInfo buffer_create_info{
       VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, // sType
       nullptr, // pNext
       0u, // flags
-      size, // size
+      buffer_properties_.size, // size
       buffer_properties_.buffer_usage, // usage
       VK_SHARING_MODE_EXCLUSIVE, // sharingMode
       0u, // queueFamilyIndexCount

backends/vulkan/runtime/api/memory/Image.cpp

@@ -129,11 +129,14 @@ VulkanImage::VulkanImage(
   VmaAllocatorInfo allocator_info{};
   vmaGetAllocatorInfo(allocator_, &allocator_info);
 
-  // If any dims are zero, then no memory will be allocated for the image.
+  // If any dims are zero, then allocate a 1x1x1 image texture. This is to
+  // ensure that there will be some resource that can be bound to a shader.
   if (image_props.image_extents.width == 0 ||
       image_props.image_extents.height == 0 ||
       image_props.image_extents.depth == 0) {
-    return;
+    image_properties_.image_extents.width = 1u;
+    image_properties_.image_extents.height = 1u;
+    image_properties_.image_extents.depth = 1u;
   }
 
   const VkImageCreateInfo image_create_info{

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -259,6 +259,38 @@ TEST_F(VulkanComputeAPITest, test_buffer_int8) {
   test_storage_buffer_type<int8_t, api::kQInt8>(16);
 }
 
+TEST_F(VulkanComputeAPITest, test_zero_size_tensor) {
+  // Simple test that performs a + b -> c
+
+  std::vector<int64_t> sizes = {0, 5, 7};
+  vTensor a = CREATE_FLOAT_TEXTURE(sizes, /*allocate_memory = */ true);
+  vTensor b = CREATE_FLOAT_TEXTURE(sizes, /*allocate_memory = */ true);
+  vTensor c = CREATE_FLOAT_TEXTURE(sizes, /*allocate_memory = */ true);
+
+  // Fill input tensors
+  fill_vtensor(a, 2.5f);
+  fill_vtensor(b, 1.5f);
+
+  // a + b -> c
+  record_binary_op(api::context(), "add", a, b, c);
+
+  // Extract output tensor
+  std::vector<float> data_out = extract_vtensor(c);
+
+  // Assert all tensors are empty
+  ASSERT_TRUE(a.numel() == 0);
+  ASSERT_TRUE(b.numel() == 0);
+  ASSERT_TRUE(c.numel() == 0);
+  ASSERT_TRUE(a.nbytes() == 0);
+  ASSERT_TRUE(b.nbytes() == 0);
+  ASSERT_TRUE(c.nbytes() == 0);
+
+  // Check output
+  for (size_t i = 0; i < data_out.size(); ++i) {
+    CHECK_VALUE(data_out, i, 4.0f);
+  }
+}
+
 TEST_F(VulkanComputeAPITest, texture_add_sanity_check) {
   // Simple test that performs a + b -> c