diff --git a/backends/vulkan/runtime/VulkanBackend.cpp b/backends/vulkan/runtime/VulkanBackend.cpp
index 03579b888dd..6d0672377e8 100644
--- 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
index 3f0f90119dd..731a79b1f7f 100644
--- 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
index 6677ecee30b..9a04b52b80a 100644
--- 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
index 89424abb2a0..019c3ab736f 100644
--- 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
index 3ce34b9949a..a6bd8b0531c 100644
--- 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
index 03cef80e673..3c1f9e2056b 100644
--- 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
index 6b085277579..5f1e579a14f 100644
--- 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(
diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp
index 789a0945366..1834b602131 100644
--- a/backends/vulkan/runtime/graph/ComputeGraph.cpp
+++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp
@@ -86,7 +86,7 @@ api::StorageType ComputeGraph::suggested_storage_type() {
   if (config_.enableStorageTypeOverride) {
     return config_.storageTypeOverride;
   }
-  return api::StorageType::TEXTURE_3D;
+  return api::kTexture3D;
 }
 
 api::GPUMemoryLayout ComputeGraph::suggested_memory_layout(
@@ -95,14 +95,14 @@ api::GPUMemoryLayout ComputeGraph::suggested_memory_layout(
     return config_.memoryLayoutOverride;
   }
   if (sizes.size() < 3) {
-    return api::GPUMemoryLayout::TENSOR_WIDTH_PACKED;
+    return api::kWidthPacked;
   }
   // For 3 dimensional tensors that only have a channels dimension of 1, still
   // prefer width packed.
   if (api::utils::val_at(-3, sizes) == 1) {
-    return api::GPUMemoryLayout::TENSOR_WIDTH_PACKED;
+    return api::kWidthPacked;
   }
-  return api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED;
+  return api::kChannelsPacked;
 }
 
 ValueRef ComputeGraph::add_tensor(
diff --git a/backends/vulkan/runtime/graph/GraphConfig.cpp b/backends/vulkan/runtime/graph/GraphConfig.cpp
index 7e3d1dfbf98..98b2d9a4263 100644
--- a/backends/vulkan/runtime/graph/GraphConfig.cpp
+++ b/backends/vulkan/runtime/graph/GraphConfig.cpp
@@ -48,16 +48,16 @@ GraphConfig::GraphConfig() {
   // of memory, increase this safety factor.
   descriptorPoolSafetyFactor = 1.25;
 
-  // For now, force TEXTURE_3D storage as we are still developing shader
-  // support for buffer storage type.
+  // For now, force kTexture3D storage as we are still developing shader support
+  // for buffer storage type.
   enableStorageTypeOverride = true;
-  storageTypeOverride = api::StorageType::TEXTURE_3D;
+  storageTypeOverride = api::kTexture3D;
 
-  // For now, force TENSOR_WIDTH_PACKED memory layout by default as we are still
+  // For now, force kWidthPacked memory layout by default as we are still
   // developing support for other memory layouts. In the future memory layout
   // settings will be serialized as part of the graph.
   enableMemoryLayoutOverride = true;
-  memoryLayoutOverride = api::GPUMemoryLayout::TENSOR_WIDTH_PACKED;
+  memoryLayoutOverride = api::kWidthPacked;
 }
 
 void GraphConfig::setStorageTypeOverride(api::StorageType storage_type) {
diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/Conv2d.cpp
index 03f9b40c2f3..0b5353229c4 100644
--- a/backends/vulkan/runtime/graph/ops/impl/Conv2d.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/Conv2d.cpp
@@ -56,10 +56,7 @@ ValueRef prepack_biases(ComputeGraph& graph, const ValueRef vref) {
     VK_THROW("aten.convolution.default: Null bias is not supported yet!");
   }
 
-  ValueRef v = graph.add_tensor_like(
-      vref,
-      api::StorageType::TEXTURE_2D,
-      api::GPUMemoryLayout::TENSOR_WIDTH_PACKED);
+  ValueRef v = graph.add_tensor_like(vref, api::kTexture2D, api::kWidthPacked);
   vTensor& t = graph.get_val(v).toTensor();
 
   api::ShaderInfo shader = get_nchw_to_image_shader(t);
@@ -106,8 +103,8 @@ ValueRef prepack_weights(ComputeGraph& graph, const ValueRef vref) {
   ValueRef v = graph.add_tensor(
       final_sizes,
       graph.get_val(vref).toTensorRef().dtype,
-      api::StorageType::TEXTURE_2D,
-      api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED);
+      api::kTexture2D,
+      api::kChannelsPacked);
   vTensor& t = graph.get_val(v).toTensor();
 
   api::utils::uvec3 global_size = t.extents();
@@ -138,10 +135,8 @@ void check_conv2d_args(const vTensor& in, const vTensor& out) {
     VK_THROW(
         "aten.convolution.default: input batch size > 1 is not supported yet!");
   }
-  VK_CHECK_COND(
-      check_memory_layout_is(in, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED));
-  VK_CHECK_COND(check_memory_layout_is(
-      out, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED));
+  VK_CHECK_COND(check_memory_layout_is(in, api::kChannelsPacked));
+  VK_CHECK_COND(check_memory_layout_is(out, api::kChannelsPacked));
 }
 
 struct Conv2dParams final {
diff --git a/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp b/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp
index a20cce2ce87..ed49926baef 100644
--- a/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp
@@ -25,9 +25,8 @@ void check_matmul_args(
   VK_CHECK_COND(check_same_ndim(mat1, mat2));
 
   VK_CHECK_COND(
-      check_memory_layout_is(
-          mat1, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED) ||
-      check_memory_layout_is(mat1, api::GPUMemoryLayout::TENSOR_WIDTH_PACKED));
+      check_memory_layout_is(mat1, api::kChannelsPacked) ||
+      check_memory_layout_is(mat1, api::kWidthPacked));
   VK_CHECK_COND(check_same_memory_layout(mat1, out));
 
   VK_CHECK_COND(check_same_sizes_at(mat1, -1, mat2, -2));
@@ -61,13 +60,12 @@ void add_matmul_node(
     const ValueRef mat1,
     const ValueRef mat2,
     const ValueRef out) {
-  ValueRef arg1 = prepack_if_tensor_ref(
-      graph, mat1, api::GPUMemoryLayout::TENSOR_WIDTH_PACKED);
+  ValueRef arg1 = prepack_if_tensor_ref(graph, mat1, api::kWidthPacked);
 
-  api::GPUMemoryLayout mat2_layout = graph.memory_layout_of(arg1) ==
-          api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED
-      ? api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED
-      : api::GPUMemoryLayout::TENSOR_HEIGHT_PACKED;
+  api::GPUMemoryLayout mat2_layout =
+      graph.memory_layout_of(arg1) == api::kChannelsPacked
+      ? api::kChannelsPacked
+      : api::kHeightPacked;
 
   ValueRef arg2 = prepack_if_tensor_ref(graph, mat2, mat2_layout);
 
diff --git a/backends/vulkan/runtime/graph/ops/impl/Pool.cpp b/backends/vulkan/runtime/graph/ops/impl/Pool.cpp
index d5f16cd98a8..2ae6a56ce5d 100644
--- a/backends/vulkan/runtime/graph/ops/impl/Pool.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/Pool.cpp
@@ -52,10 +52,8 @@ void resize_max_pool2d_node(
 }
 
 void check_max_pool2d_args(const vTensor& in, const vTensor& out) {
-  VK_CHECK_COND(
-      check_memory_layout_is(in, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED));
-  VK_CHECK_COND(check_memory_layout_is(
-      out, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED));
+  VK_CHECK_COND(check_memory_layout_is(in, api::kChannelsPacked));
+  VK_CHECK_COND(check_memory_layout_is(out, api::kChannelsPacked));
 }
 
 void add_max_pool2d_node(
diff --git a/backends/vulkan/runtime/graph/ops/impl/Sum.cpp b/backends/vulkan/runtime/graph/ops/impl/Sum.cpp
index a8cd31c5d61..ff235a2357f 100644
--- a/backends/vulkan/runtime/graph/ops/impl/Sum.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/Sum.cpp
@@ -46,10 +46,8 @@ void resize_sum_node(
 }
 
 void check_sum_args(const vTensor& in, const vTensor& out) {
-  VK_CHECK_COND(
-      check_memory_layout_is(in, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED));
-  VK_CHECK_COND(check_memory_layout_is(
-      out, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED));
+  VK_CHECK_COND(check_memory_layout_is(in, api::kChannelsPacked));
+  VK_CHECK_COND(check_memory_layout_is(out, api::kChannelsPacked));
 }
 
 void add_sum_dim_node(
@@ -106,8 +104,7 @@ ValueRef add_node(
     const api::ScalarType dtype = api::kFloat) {
   vTensor& v_input = graph.get_val(input).toTensor();
   std::vector<int64_t> output_size = calc_out_sizes(v_input, dim, keepdim);
-  return graph.add_tensor(
-      output_size, dtype, api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED);
+  return graph.add_tensor(output_size, dtype, api::kChannelsPacked);
 }
 
 void add_sum_dim_IntList(
diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp
index 307d676d3b2..842cfa2f4fc 100644
--- a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp
@@ -78,13 +78,13 @@ bool is_packed_dim_broadcasted(const vTensor& sndr, const vTensor& rcvr) {
   // We assume that the tensors are broadcastable. If values aren't equal at
   // some index, then the value of rcvr is 1 and hence should be broadcasted.
   switch (sndr.gpu_memory_layout()) {
-    case api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED:
+    case api::kChannelsPacked:
       return api::utils::val_at(-3, sndr.sizes()) >
           api::utils::val_at(-3, rcvr.sizes());
-    case api::GPUMemoryLayout::TENSOR_HEIGHT_PACKED:
+    case api::kHeightPacked:
       return api::utils::val_at(-2, sndr.sizes()) >
           api::utils::val_at(-2, rcvr.sizes());
-    case api::GPUMemoryLayout::TENSOR_WIDTH_PACKED:
+    case api::kWidthPacked:
       return api::utils::val_at(-1, sndr.sizes()) >
           api::utils::val_at(-1, rcvr.sizes());
   }
diff --git a/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp b/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp
index 4f29e93e161..e472690327b 100644
--- a/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp
+++ b/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp
@@ -28,10 +28,10 @@ void apply_dtype_suffix(std::stringstream& kernel_name, const vTensor& tensor) {
 
 void apply_ndim_suffix(std::stringstream& kernel_name, const vTensor& tensor) {
   switch (tensor.storage_type()) {
-    case api::StorageType::TEXTURE_3D:
+    case api::kTexture3D:
       kernel_name << "_3d";
       break;
-    case api::StorageType::TEXTURE_2D:
+    case api::kTexture2D:
       kernel_name << "_2d";
       break;
     default:
@@ -43,13 +43,13 @@ void apply_memory_layout_suffix(
     std::stringstream& kernel_name,
     const vTensor& tensor) {
   switch (tensor.gpu_memory_layout()) {
-    case api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED:
+    case api::kChannelsPacked:
       kernel_name << "_C_packed";
       break;
-    case api::GPUMemoryLayout::TENSOR_HEIGHT_PACKED:
+    case api::kHeightPacked:
       kernel_name << "_H_packed";
       break;
-    case api::GPUMemoryLayout::TENSOR_WIDTH_PACKED:
+    case api::kWidthPacked:
       kernel_name << "_W_packed";
       break;
     default:
diff --git a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp
index 7f5ae409d44..94228321f79 100644
--- a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp
+++ b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp
@@ -97,10 +97,10 @@ api::ShaderInfo get_nchw_to_image_shader(const vTensor& v_dst) {
   std::stringstream kernel_name;
 
   switch (v_dst.storage_type()) {
-    case api::StorageType::TEXTURE_3D:
+    case api::kTexture3D:
       kernel_name << "nchw_to_image3d";
       break;
-    case api::StorageType::TEXTURE_2D:
+    case api::kTexture2D:
       kernel_name << "nchw_to_image2d";
       break;
     default:
@@ -121,10 +121,10 @@ api::ShaderInfo get_image_to_nchw_shader(const vTensor& v_src) {
   std::stringstream kernel_name;
 
   switch (v_src.storage_type()) {
-    case api::StorageType::TEXTURE_3D:
+    case api::kTexture3D:
       kernel_name << "image3d_to_nchw";
       break;
-    case api::StorageType::TEXTURE_2D:
+    case api::kTexture2D:
       kernel_name << "image2d_to_nchw";
       break;
     default: