Introduce convenience constexpr for StorageTypes and `GPUMemoryLayo…

…ut`s (#2948) Summary: Pull Request resolved: #2948 ## Context Introduce the following convenience `constexpr`: * `api::kBuffer`, `api::kTexture3D`, and `api::kTexture2D` * `api::kWidthPacked`, `api::kHeightPacked`, and `api::kChannelsPacked` Also remove the `api::StorageType::UNKNOWN` enum entry as it doesn't really serve any purpose. ghstack-source-id: 221871428 bypass-github-pytorch-ci-checks Reviewed By: copyrightly, jorgep31415 Differential Revision: D55811278 fbshipit-source-id: 26dc1706ac2605c13f247d08a21863ff3ef94488
pytorch · Apr 9, 2024 · b26eee8 · b26eee8
1 parent 4599650
commit b26eee8
Show file tree

Hide file tree

Showing 16 changed files with 95 additions and 144 deletions.
diff --git a/backends/vulkan/runtime/VulkanBackend.cpp b/backends/vulkan/runtime/VulkanBackend.cpp
@@ -77,26 +77,26 @@ api::StorageType get_storage_type(
     const vkgraph::VkStorageType& vk_storage_type) {
   switch (vk_storage_type) {
     case vkgraph::VkStorageType::BUFFER:
-      return api::StorageType::BUFFER;
+      return api::kBuffer;
     case vkgraph::VkStorageType::TEXTURE_3D:
-      return api::StorageType::TEXTURE_3D;
+      return api::kTexture3D;
     case vkgraph::VkStorageType::TEXTURE_2D:
-      return api::StorageType::TEXTURE_2D;
+      return api::kTexture2D;
     default:
       break;
   }
-  return api::StorageType::UNKNOWN;
+  VK_THROW("Invalid storage type encountered!");
 }
 
 api::GPUMemoryLayout get_memory_layout(
     const vkgraph::VkMemoryLayout& vk_memory_layout) {
   switch (vk_memory_layout) {
     case vkgraph::VkMemoryLayout::TENSOR_WIDTH_PACKED:
-      return api::GPUMemoryLayout::TENSOR_WIDTH_PACKED;
+      return api::kWidthPacked;
     case vkgraph::VkMemoryLayout::TENSOR_HEIGHT_PACKED:
-      return api::GPUMemoryLayout::TENSOR_HEIGHT_PACKED;
+      return api::kHeightPacked;
     case vkgraph::VkMemoryLayout::TENSOR_CHANNELS_PACKED:
-      return api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED;
+      return api::kChannelsPacked;
     default:
       break;
   }

diff --git a/backends/vulkan/runtime/api/Shader.cpp b/backends/vulkan/runtime/api/Shader.cpp
@@ -23,38 +23,19 @@ ShaderInfo::ShaderInfo()
           0u,
       } {}
 
-ShaderInfo::ShaderInfo(
-    std::string name,
-    const uint32_t* const spirv_bin,
-    const uint32_t size,
-    std::vector<VkDescriptorType>  layout)
-    : src_code{
-          spirv_bin,
-          size,
-      },
-      kernel_name{std::move(name)},
-      kernel_layout{std::move(layout)} {}
-
 ShaderInfo::ShaderInfo(
     std::string name,
     const uint32_t* const spirv_bin,
     const uint32_t size,
     std::vector<VkDescriptorType>  layout,
-    const std::vector<uint32_t>& tile_size,
-    const StorageType bias_storage_type,
-    const StorageType weight_storage_type)
+    const utils::uvec3 tile_size)
     : src_code{
           spirv_bin,
           size,
       },
       kernel_name{std::move(name)},
       kernel_layout{std::move(layout)},
-      tile_size(tile_size),
-      bias_storage_type(bias_storage_type),
-      weight_storage_type(weight_storage_type) {
-  for (uint64_t i = 0; i < tile_size.size(); ++i) {
-    out_tile_size.data[i] = tile_size[i];
-  }
+      out_tile_size(tile_size) {
 }
 
 bool operator==(const ShaderInfo& _1, const ShaderInfo& _2) {

diff --git a/backends/vulkan/runtime/api/Shader.h b/backends/vulkan/runtime/api/Shader.h
@@ -62,25 +62,14 @@ struct ShaderInfo final {
   // Shader Metadata
   utils::uvec3 out_tile_size{1u, 1u, 1u};
 
-  std::vector<uint32_t> tile_size;
-  StorageType bias_storage_type{StorageType::UNKNOWN};
-  StorageType weight_storage_type{StorageType::UNKNOWN};
-
   explicit ShaderInfo();
-  explicit ShaderInfo(std::string, const char*);
-  explicit ShaderInfo(
-      std::string,
-      const uint32_t*,
-      const uint32_t,
-      std::vector<VkDescriptorType>);
+
   explicit ShaderInfo(
       std::string,
       const uint32_t*,
       const uint32_t,
       std::vector<VkDescriptorType>,
-      const std::vector<uint32_t>& tile_size,
-      const StorageType bias_storage_type,
-      const StorageType weight_storage_type);
+      const utils::uvec3 tile_size);
 };
 
 bool operator==(const ShaderInfo& _1, const ShaderInfo& _2);

diff --git a/backends/vulkan/runtime/api/Tensor.cpp b/backends/vulkan/runtime/api/Tensor.cpp
@@ -67,20 +67,20 @@ std::vector<int64_t> calc_strides(
     const api::GPUMemoryLayout memory_layout,
     const api::StorageType storage_type) {
   switch (storage_type) {
-    case api::StorageType::BUFFER:
+    case api::kBuffer:
       switch (memory_layout) {
-        case api::GPUMemoryLayout::TENSOR_WIDTH_PACKED:
+        case api::kWidthPacked:
           return calc_contiguous_strides(sizes);
           break;
-        case api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED:
+        case api::kChannelsPacked:
           return calc_channels_last_strides(sizes);
           break;
         default:
           VK_THROW("Invalid memory format used to create vTensor!");
       }
       break;
-    case api::StorageType::TEXTURE_3D:
-    case api::StorageType::TEXTURE_2D:
+    case api::kTexture3D:
+    case api::kTexture2D:
       return std::vector<int64_t>(sizes.size());
     default:
       VK_THROW("Invalid storage type used to create vTensor!");
@@ -99,10 +99,8 @@ std::vector<int64_t> calc_gpu_sizes(
     const std::vector<int64_t>& sizes,
     const api::GPUMemoryLayout memory_layout,
     const api::StorageType storage_type) {
-  VK_CHECK_COND(storage_type != api::StorageType::UNKNOWN);
-
   std::vector<int64_t> gpu_sizes;
-  if (storage_type == api::StorageType::BUFFER) {
+  if (storage_type == api::kBuffer) {
     gpu_sizes.resize(sizes.size());
     for (size_t i = 0; i < sizes.size(); i++) {
       gpu_sizes.at(i) = sizes.at(i);
@@ -127,21 +125,21 @@ std::vector<int64_t> calc_gpu_sizes(
 
   size_t ndim = gpu_sizes.size();
   switch (memory_layout) {
-    case api::GPUMemoryLayout::TENSOR_WIDTH_PACKED:
+    case api::kWidthPacked:
       if (ndim >= 1) {
         gpu_sizes.at(ndim - 1) =
             api::utils::align_up(api::utils::val_at(-1, sizes), INT64_C(4));
       }
       break;
 
-    case api::GPUMemoryLayout::TENSOR_HEIGHT_PACKED:
+    case api::kHeightPacked:
       if (ndim >= 2) {
         gpu_sizes.at(ndim - 2) =
             api::utils::align_up(api::utils::val_at(-2, sizes), INT64_C(4));
       }
       break;
 
-    case api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED:
+    case api::kChannelsPacked:
       if (ndim >= 3) {
         gpu_sizes.at(ndim - 3) =
             api::utils::align_up(api::utils::val_at(-3, sizes), INT64_C(4));
@@ -162,7 +160,7 @@ api::utils::uvec3 create_image_extents(
     const api::GPUMemoryLayout memory_layout) {
   size_t ndim = gpu_sizes.size();
 
-  if (storage_type == api::StorageType::BUFFER) {
+  if (storage_type == api::kBuffer) {
     // image extents do not apply to buffer storage
     return {0u, 0u, 0u};
   } else {
@@ -177,15 +175,15 @@ api::utils::uvec3 create_image_extents(
     uint32_t batch = safe_downcast<uint32_t>(val_at(-4, gpu_sizes));
 
     switch (memory_layout) {
-      case api::GPUMemoryLayout::TENSOR_WIDTH_PACKED:
+      case api::kWidthPacked:
         VK_CHECK_COND(width % 4 == 0, "Channels must be divisible by 4!");
         width /= 4;
         break;
-      case api::GPUMemoryLayout::TENSOR_HEIGHT_PACKED:
+      case api::kHeightPacked:
         VK_CHECK_COND(height % 4 == 0, "Channels must be divisible by 4!");
         height /= 4;
         break;
-      case api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED:
+      case api::kChannelsPacked:
         VK_CHECK_COND(channels % 4 == 0, "Channels must be divisible by 4!");
         channels /= 4;
         break;
@@ -326,41 +324,35 @@ std::shared_ptr<api::UniformParamsBuffer> vTensor::extents_ubo() {
 
 VmaAllocationCreateInfo vTensor::get_allocation_create_info() const {
   switch (storage_type()) {
-    case api::StorageType::BUFFER:
+    case api::kBuffer:
       return view_->buffer_.allocation_create_info();
-    case api::StorageType::TEXTURE_2D:
-    case api::StorageType::TEXTURE_3D:
+    case api::kTexture2D:
+    case api::kTexture3D:
       return view_->image_.allocation_create_info();
-    case api::StorageType::UNKNOWN:
-      break;
   }
   return {};
 }
 
 VkMemoryRequirements vTensor::get_memory_requirements() const {
   switch (storage_type()) {
-    case api::StorageType::BUFFER:
+    case api::kBuffer:
       return view_->buffer_.get_memory_requirements();
-    case api::StorageType::TEXTURE_2D:
-    case api::StorageType::TEXTURE_3D:
+    case api::kTexture2D:
+    case api::kTexture3D:
       return view_->image_.get_memory_requirements();
-    case api::StorageType::UNKNOWN:
-      break;
   }
   return {};
 }
 
 void vTensor::bind_allocation(const api::MemoryAllocation& allocation) {
   switch (storage_type()) {
-    case api::StorageType::BUFFER:
+    case api::kBuffer:
       view_->buffer_.bind_allocation(allocation);
       break;
-    case api::StorageType::TEXTURE_2D:
-    case api::StorageType::TEXTURE_3D:
+    case api::kTexture2D:
+    case api::kTexture3D:
       view_->image_.bind_allocation(allocation);
       break;
-    case api::StorageType::UNKNOWN:
-      break;
   }
 }
 
@@ -397,7 +389,7 @@ void vTensor::reallocate(const std::vector<int64_t>& new_sizes) {
 
 void vTensor::virtual_resize(const std::vector<int64_t>& new_sizes) {
   update_size_metadata(new_sizes);
-  if (storage_type() == api::StorageType::BUFFER) {
+  if (storage_type() == api::kBuffer) {
     if (gpu_nbytes() > view_->buffer_.mem_size()) {
       VK_THROW(
           "Cannot virtual_resize a vTensor with sizes that require a larger "
@@ -446,11 +438,11 @@ api::VulkanImage allocate_image(
   VkImageViewType image_view_type = VK_IMAGE_VIEW_TYPE_3D;
 
   switch (storage_type) {
-    case api::StorageType::TEXTURE_3D:
+    case api::kTexture3D:
       image_type = VK_IMAGE_TYPE_3D;
       image_view_type = VK_IMAGE_VIEW_TYPE_3D;
       break;
-    case api::StorageType::TEXTURE_2D:
+    case api::kTexture2D:
       image_type = VK_IMAGE_TYPE_2D;
       image_view_type = VK_IMAGE_VIEW_TYPE_2D;
       break;
@@ -481,7 +473,7 @@ api::VulkanBuffer allocate_buffer(
   api::Adapter* adapter_ptr = context_ptr->adapter_ptr();
 
   switch (storage_type) {
-    case api::StorageType::BUFFER:
+    case api::kBuffer:
       break;
     default:
       // Return an empty VulkanBuffer if Buffer storage is not used

diff --git a/backends/vulkan/runtime/api/Tensor.h b/backends/vulkan/runtime/api/Tensor.h
@@ -103,9 +103,8 @@ class vTensor final {
       api::Context* context,
       const std::vector<int64_t>& sizes,
       const api::ScalarType dtype,
-      const api::StorageType storage_type = api::StorageType::TEXTURE_3D,
-      const api::GPUMemoryLayout memory_layout =
-          api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED,
+      const api::StorageType storage_type = api::kTexture3D,
+      const api::GPUMemoryLayout memory_layout = api::kChannelsPacked,
       const bool allocate_memory = true);
 
   // Default constructor for quantized vTensor
@@ -115,9 +114,8 @@ class vTensor final {
       double q_scale,
       int64_t q_zero_point,
       const api::ScalarType dtype,
-      const api::StorageType storage_type = api::StorageType::TEXTURE_3D,
-      const api::GPUMemoryLayout memory_layout =
-          api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED);
+      const api::StorageType storage_type = api::kTexture3D,
+      const api::GPUMemoryLayout memory_layout = api::kChannelsPacked);
 
   // Copy Constructor and Assignment; Ideally copying  would be disabled
   // (see the reasoning for move assignment below) but it is required for

diff --git a/backends/vulkan/runtime/api/Types.h b/backends/vulkan/runtime/api/Types.h
@@ -162,13 +162,16 @@ VK_FORALL_SCALAR_TYPES(SPECIALIZE_ScalarTypeToCType)
  *
  * UNKNOWN is not expected to be used.
  */
-enum class StorageType {
+enum class StorageType : uint8_t {
   BUFFER,
   TEXTURE_3D,
   TEXTURE_2D,
-  UNKNOWN,
 };
 
+static constexpr StorageType kBuffer = StorageType::BUFFER;
+static constexpr StorageType kTexture3D = StorageType::TEXTURE_3D;
+static constexpr StorageType kTexture2D = StorageType::TEXTURE_2D;
+
 /**
  * The enum below is used to describe how tensor data is laid out when stored in
  * GPU memory. The name of the enum describes which dimension is tightly packed;
@@ -182,11 +185,20 @@ enum class StorageType {
  * strides of the tensor will be used instead to convert between logical tensor
  * coordinates and linear access indices.
  */
-enum class GPUMemoryLayout : uint32_t {
+enum class GPUMemoryLayout : uint8_t {
   TENSOR_WIDTH_PACKED = 0u,
   TENSOR_HEIGHT_PACKED = 1u,
   TENSOR_CHANNELS_PACKED = 2u,
 };
 
+static constexpr GPUMemoryLayout kWidthPacked =
+    GPUMemoryLayout::TENSOR_WIDTH_PACKED;
+
+static constexpr GPUMemoryLayout kHeightPacked =
+    GPUMemoryLayout::TENSOR_HEIGHT_PACKED;
+
+static constexpr GPUMemoryLayout kChannelsPacked =
+    GPUMemoryLayout::TENSOR_CHANNELS_PACKED;
+
 } // namespace api
 } // namespace vkcompute
diff --git a/backends/vulkan/runtime/api/gen_vulkan_spv.py b/backends/vulkan/runtime/api/gen_vulkan_spv.py
@@ -543,13 +543,6 @@ def findRegisterFor(lineStr: str) -> Tuple[str, List[str]]:
     r"\buniform\b": "VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER",
 }
 
-storageTypeToEnum = {
-    "TEXTURE_2D": "api::StorageType::TEXTURE_2D",
-    "TEXTURE_3D": "api::StorageType::TEXTURE_3D",
-    "BUFFER": "api::StorageType::BUFFER",
-    "": "api::StorageType::UNKNOWN",
-}
-
 
 def determineDescriptorType(lineStr: str) -> str:
     for identifier, typeNum in typeIdMapping.items():
@@ -632,7 +625,7 @@ def generateShaderInfoStr(shader_info: ShaderInfo, name: str, sizeBytes: int) ->
     tile_size = (
         f"{{{', '.join(str(x) for x in shader_info.tile_size)}}}"
         if (len(shader_info.tile_size) > 0)
-        else "std::vector<uint32_t>()"
+        else "{1, 1, 1}"
     )
 
     shader_info_layouts = "{{{}}}".format(",\n ".join(shader_info.layouts))
@@ -643,8 +636,6 @@ def generateShaderInfoStr(shader_info: ShaderInfo, name: str, sizeBytes: int) ->
         str(sizeBytes),
         shader_info_layouts,
         tile_size,
-        storageTypeToEnum[shader_info.weight_storage_type],
-        storageTypeToEnum[shader_info.bias_storage_type],
     ]
 
     shader_info_str = textwrap.indent(