[ET-VK][runtime] Derive SPIR-V requirements and enable Vulkan shader features

backends/vulkan/runtime/VulkanBackend.cpp

@@ -717,8 +717,14 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     new (compute_graph) ComputeGraph(graph_config);
 
     const NamedDataMap* named_data_map = context.get_named_data_map();
-    Error err = compileModel(
-        processed->data(), processed->size(), compute_graph, named_data_map);
+    Error err = Error::Ok;
+    try {
+      err = compileModel(
+          processed->data(), processed->size(), compute_graph, named_data_map);
+    } catch (const vkapi::ShaderNotSupportedError& error) {
+      ET_LOG(Error, "Unsupported Vulkan shader: %s", error.what());
+      err = Error::NotSupported;
+    }
 
     // This backend does not need its processed data after compiling the
     // model.

backends/vulkan/runtime/api/Context.cpp

@@ -88,6 +88,16 @@ void Context::report_shader_dispatch_end() {
 }
 
 void Context::check_device_capabilities(const vkapi::ShaderInfo& shader) {
+  if (!shader.required_capabilities.empty()) {
+    for (const uint32_t capability : shader.required_capabilities) {
+      if (!adapter_p_->supports_spirv_capability(capability)) {
+        throw vkapi::ShaderNotSupportedError(
+            shader.kernel_name,
+            utils::concat_str("SPIR-V capability ", capability));
+      }
+    }
+    return;
+  }
   if (shader.requires_shader_int16) {
     if (!adapter_p_->supports_int16_shader_types()) {
       throw vkapi::ShaderNotSupportedError(

backends/vulkan/runtime/gen_vulkan_spv.py

@@ -17,14 +17,16 @@
 import os
 import re
 import shutil
+import struct
 import sys
+from enum import IntEnum
 from itertools import product
 from multiprocessing.pool import ThreadPool
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 import subprocess
 import textwrap
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import yaml
@@ -1349,6 +1351,23 @@ def compile_spirv(shader_paths_pair) -> Tuple[str, str]:
 ##############################################
 
 
+class SpvCapability(IntEnum):
+    FLOAT64 = 10
+    INT64 = 11
+    INT16 = 22
+    STORAGE_BUFFER_16_BIT_ACCESS = 4433
+    UNIFORM_AND_STORAGE_BUFFER_16_BIT_ACCESS = 4434
+    STORAGE_PUSH_CONSTANT_16 = 4435
+    STORAGE_INPUT_OUTPUT_16 = 4436
+    STORAGE_BUFFER_8_BIT_ACCESS = 4448
+    UNIFORM_AND_STORAGE_BUFFER_8_BIT_ACCESS = 4449
+    STORAGE_PUSH_CONSTANT_8 = 4450
+    DOT_PRODUCT_INPUT_ALL = 6016
+    DOT_PRODUCT_INPUT_4X8_BIT = 6017
+    DOT_PRODUCT_INPUT_4X8_BIT_PACKED = 6018
+    DOT_PRODUCT = 6019
+
+
 @dataclass
 class ShaderInfo:
     tile_size: List[int]
@@ -1368,6 +1387,7 @@ class ShaderInfo:
     #        `SUBGROUP_SIZE` template parameter (single source of truth for
     #        both GLSL substitution and the Vulkan pipeline pin).
     required_subgroup_size: int = 0
+    required_capabilities: List[int] = field(default_factory=list)
 
 
 def getName(filePath: str) -> str:
@@ -1446,11 +1466,6 @@ def findRegisterFor(lineStr: str) -> Tuple[str, List[str]]:
     return (matches_list[0], matches_list[1:])
 
 
-def isExtensionRequireLine(lineStr: str) -> bool:
-    extension_require_id = r"^#extension ([A-Za-z0-9_]+)\s*:\s*require"
-    return re.search(extension_require_id, lineStr) is not None
-
-
 typeIdMapping = {
     r"image[123]D\b": "VK_DESCRIPTOR_TYPE_STORAGE_IMAGE",
     r"sampler[123]D\b": "VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER",
@@ -1468,7 +1483,33 @@ def determineDescriptorType(lineStr: str) -> str:
     )
 
 
-def getShaderInfo(srcFilePath: str) -> ShaderInfo:  # noqa: C901
+def getSpvCapabilities(spvPath: str) -> Set[int]:
+    with open(spvPath, "rb") as spvFile:
+        data = spvFile.read()
+    if len(data) < 20 or len(data) % 4 != 0:
+        raise ValueError(f"Invalid SPIR-V header or byte length: {spvPath}")
+    words = struct.unpack(f"<{len(data) // 4}I", data)
+    if words[0] != 0x07230203:
+        raise ValueError(f"Invalid SPIR-V magic: {spvPath}")
+
+    capabilities = set()
+    offset = 5
+    while offset < len(words):
+        word_count = words[offset] >> 16
+        opcode = words[offset] & 0xFFFF
+        if word_count == 0 or offset + word_count > len(words):
+            raise ValueError(f"Invalid SPIR-V instruction at word {offset}: {spvPath}")
+        if opcode == 17:  # OpCapability
+            if word_count != 2:
+                raise ValueError(
+                    f"Invalid SPIR-V OpCapability at word {offset}: {spvPath}"
+                )
+            capabilities.add(words[offset + 1])
+        offset += word_count
+    return capabilities
+
+
+def getShaderInfo(srcFilePath: str, spvFilePath: str) -> ShaderInfo:  # noqa: C901
     shader_info = ShaderInfo([], [], "")
     with open(srcFilePath) as srcFile:
         for line in srcFile:
@@ -1484,20 +1525,37 @@ def getShaderInfo(srcFilePath: str) -> ShaderInfo:  # noqa: C901
                 shader_info.bias_storage_type = getBiasStorageType(line)
             if isRegisterForLine(line):
                 shader_info.register_for = findRegisterFor(line)
-            if isExtensionRequireLine(line):
-                if "GL_EXT_shader_explicit_arithmetic_types_int16" in line:
-                    shader_info.requires_shader_int16_ext = True
-                if "GL_EXT_shader_16bit_storage" in line:
-                    shader_info.requires_16bit_storage_ext = True
-                if "GL_EXT_shader_8bit_storage" in line:
-                    shader_info.requires_8bit_storage_ext = True
-                if "GL_EXT_integer_dot_product" in line:
-                    shader_info.requires_integer_dot_product_ext = True
-                if "GL_EXT_shader_explicit_arithmetic_types_int64" in line:
-                    shader_info.requires_shader_int64_ext = True
-                if "GL_EXT_shader_explicit_arithmetic_types_float64" in line:
-                    shader_info.requires_shader_float64_ext = True
-
+    capabilities = getSpvCapabilities(spvFilePath)
+    shader_info.required_capabilities = sorted(capabilities)
+    shader_info.requires_shader_int16_ext = SpvCapability.INT16 in capabilities
+    shader_info.requires_shader_int64_ext = SpvCapability.INT64 in capabilities
+    shader_info.requires_shader_float64_ext = SpvCapability.FLOAT64 in capabilities
+    shader_info.requires_16bit_storage_ext = bool(
+        capabilities
+        & {
+            SpvCapability.STORAGE_BUFFER_16_BIT_ACCESS,
+            SpvCapability.UNIFORM_AND_STORAGE_BUFFER_16_BIT_ACCESS,
+            SpvCapability.STORAGE_PUSH_CONSTANT_16,
+            SpvCapability.STORAGE_INPUT_OUTPUT_16,
+        }
+    )
+    shader_info.requires_8bit_storage_ext = bool(
+        capabilities
+        & {
+            SpvCapability.STORAGE_BUFFER_8_BIT_ACCESS,
+            SpvCapability.UNIFORM_AND_STORAGE_BUFFER_8_BIT_ACCESS,
+            SpvCapability.STORAGE_PUSH_CONSTANT_8,
+        }
+    )
+    shader_info.requires_integer_dot_product_ext = bool(
+        capabilities
+        & {
+            SpvCapability.DOT_PRODUCT_INPUT_ALL,
+            SpvCapability.DOT_PRODUCT_INPUT_4X8_BIT,
+            SpvCapability.DOT_PRODUCT_INPUT_4X8_BIT_PACKED,
+            SpvCapability.DOT_PRODUCT,
+        }
+    )
     return shader_info
 
 
@@ -1576,6 +1634,7 @@ def to_cpp_str(val: bool):
         to_cpp_str(shader_info.requires_shader_int64_ext),
         to_cpp_str(shader_info.requires_shader_float64_ext),
         str(shader_info.required_subgroup_size),
+        "{" + ", ".join(str(x) for x in shader_info.required_capabilities) + "}",
     ]
 
     shader_info_str = textwrap.indent(
@@ -1621,7 +1680,7 @@ def genCppFiles(
         sizeBytes, spv_bin_str = generateSpvBinStr(spvPath, name)
         spv_bin_strs.append(spv_bin_str)
 
-        shader_info = getShaderInfo(srcPath)
+        shader_info = getShaderInfo(srcPath, spvPath)
 
         register_shader_info_strs.append(
             generateShaderInfoStr(shader_info, name, sizeBytes)

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -880,6 +880,7 @@ void ComputeGraph::register_pipeline_to_create(
     const LocalWorkGroup& lwg,
     const vkapi::SpecVarList& spec_vars,
     const std::vector<PushConstantDataInfo>& push_constants) {
+  context()->check_device_capabilities(shader_info);
   VkDescriptorSetLayout shader_layout =
       context()->shader_layout_cache().retrieve(shader_info.kernel_layout);
 

backends/vulkan/runtime/vk_api/Adapter.cpp

@@ -154,6 +154,29 @@ VkDevice create_logical_device(
   enabled_features.shaderFloat64 =
       physical_device.supports_float64_shader_types ? VK_TRUE : VK_FALSE;
 
+  enabled_features.shaderImageGatherExtended =
+      physical_device.shader_features.shaderImageGatherExtended;
+  enabled_features.shaderStorageImageMultisample =
+      physical_device.shader_features.shaderStorageImageMultisample;
+  enabled_features.shaderUniformBufferArrayDynamicIndexing =
+      physical_device.shader_features.shaderUniformBufferArrayDynamicIndexing;
+  enabled_features.shaderSampledImageArrayDynamicIndexing =
+      physical_device.shader_features.shaderSampledImageArrayDynamicIndexing;
+  enabled_features.shaderStorageBufferArrayDynamicIndexing =
+      physical_device.shader_features.shaderStorageBufferArrayDynamicIndexing;
+  enabled_features.shaderStorageImageArrayDynamicIndexing =
+      physical_device.shader_features.shaderStorageImageArrayDynamicIndexing;
+  enabled_features.imageCubeArray =
+      physical_device.shader_features.imageCubeArray;
+  enabled_features.shaderResourceMinLod =
+      physical_device.shader_features.shaderResourceMinLod;
+  enabled_features.shaderStorageImageExtendedFormats =
+      physical_device.shader_features.shaderStorageImageExtendedFormats;
+  enabled_features.shaderStorageImageReadWithoutFormat =
+      physical_device.shader_features.shaderStorageImageReadWithoutFormat;
+  enabled_features.shaderStorageImageWriteWithoutFormat =
+      physical_device.shader_features.shaderStorageImageWriteWithoutFormat;
+
   VkDeviceCreateInfo device_create_info{
       VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, // sType
       nullptr, // pNext
@@ -685,5 +708,123 @@ uint32_t resolve_required_subgroup_size(
   return resolved;
 }
 
+bool Adapter::supports_spirv_capability(const uint32_t value) {
+  const auto capability = static_cast<SpirvCapability>(value);
+  switch (capability) {
+    case SpirvCapability::MATRIX:
+    case SpirvCapability::SHADER:
+    case SpirvCapability::SAMPLED_1D:
+    case SpirvCapability::IMAGE_1D:
+    case SpirvCapability::SAMPLED_BUFFER:
+    case SpirvCapability::IMAGE_BUFFER:
+    case SpirvCapability::IMAGE_QUERY:
+      return true;
+    default:
+      break;
+  }
+  const auto& features = physical_device_.shader_features;
+  switch (capability) {
+    case SpirvCapability::FLOAT16:
+      return supports_float16_shader_types();
+    case SpirvCapability::FLOAT64:
+      return supports_float64_shader_types();
+    case SpirvCapability::INT64:
+      return supports_int64_shader_types();
+    case SpirvCapability::INT16:
+      return supports_int16_shader_types();
+    case SpirvCapability::INT8:
+      return supports_int8_shader_types();
+    case SpirvCapability::IMAGE_GATHER_EXTENDED:
+      return features.shaderImageGatherExtended == VK_TRUE;
+    case SpirvCapability::STORAGE_IMAGE_MULTISAMPLE:
+      return features.shaderStorageImageMultisample == VK_TRUE;
+    case SpirvCapability::UNIFORM_BUFFER_ARRAY_DYNAMIC_INDEXING:
+      return features.shaderUniformBufferArrayDynamicIndexing == VK_TRUE;
+    case SpirvCapability::SAMPLED_IMAGE_ARRAY_DYNAMIC_INDEXING:
+      return features.shaderSampledImageArrayDynamicIndexing == VK_TRUE;
+    case SpirvCapability::STORAGE_BUFFER_ARRAY_DYNAMIC_INDEXING:
+      return features.shaderStorageBufferArrayDynamicIndexing == VK_TRUE;
+    case SpirvCapability::STORAGE_IMAGE_ARRAY_DYNAMIC_INDEXING:
+      return features.shaderStorageImageArrayDynamicIndexing == VK_TRUE;
+    case SpirvCapability::IMAGE_CUBE_ARRAY:
+    case SpirvCapability::SAMPLED_CUBE_ARRAY:
+      return features.imageCubeArray == VK_TRUE;
+    case SpirvCapability::MIN_LOD:
+      return features.shaderResourceMinLod == VK_TRUE;
+    case SpirvCapability::STORAGE_IMAGE_EXTENDED_FORMATS:
+      return features.shaderStorageImageExtendedFormats == VK_TRUE;
+    case SpirvCapability::STORAGE_IMAGE_READ_WITHOUT_FORMAT:
+      return features.shaderStorageImageReadWithoutFormat == VK_TRUE;
+    case SpirvCapability::STORAGE_IMAGE_WRITE_WITHOUT_FORMAT:
+      return features.shaderStorageImageWriteWithoutFormat == VK_TRUE;
+    case SpirvCapability::GROUP_NON_UNIFORM:
+      return supports_subgroup_compute_basic();
+    case SpirvCapability::GROUP_NON_UNIFORM_VOTE:
+      return supports_subgroup_compute_vote();
+    case SpirvCapability::GROUP_NON_UNIFORM_ARITHMETIC:
+      return supports_subgroup_compute_arithmetic();
+    case SpirvCapability::GROUP_NON_UNIFORM_BALLOT:
+      return supports_subgroup_compute_ballot();
+    case SpirvCapability::GROUP_NON_UNIFORM_SHUFFLE:
+      return supports_subgroup_compute_shuffle();
+    case SpirvCapability::GROUP_NON_UNIFORM_SHUFFLE_RELATIVE:
+      return supports_subgroup_compute_shuffle_relative();
+    case SpirvCapability::GROUP_NON_UNIFORM_CLUSTERED:
+      return supports_subgroup_compute_clustered();
+    case SpirvCapability::GROUP_NON_UNIFORM_QUAD:
+      return supports_subgroup_compute_quad();
+#if defined(VK_KHR_16bit_storage) && !defined(ETVK_FORCE_NO_EXTENSIONS)
+    case SpirvCapability::STORAGE_BUFFER_16_BIT_ACCESS:
+      return supports_16bit_storage_buffers();
+    case SpirvCapability::UNIFORM_AND_STORAGE_BUFFER_16_BIT_ACCESS:
+      return physical_device_.shader_16bit_storage
+                 .uniformAndStorageBuffer16BitAccess == VK_TRUE;
+    case SpirvCapability::STORAGE_PUSH_CONSTANT_16:
+      return physical_device_.shader_16bit_storage.storagePushConstant16 ==
+          VK_TRUE;
+    case SpirvCapability::STORAGE_INPUT_OUTPUT_16:
+      return physical_device_.shader_16bit_storage.storageInputOutput16 ==
+          VK_TRUE;
+#endif
+#if defined(VK_KHR_8bit_storage) && !defined(ETVK_FORCE_NO_EXTENSIONS)
+    case SpirvCapability::STORAGE_BUFFER_8_BIT_ACCESS:
+      return supports_8bit_storage_buffers();
+    case SpirvCapability::UNIFORM_AND_STORAGE_BUFFER_8_BIT_ACCESS:
+      return physical_device_.shader_8bit_storage
+                 .uniformAndStorageBuffer8BitAccess == VK_TRUE;
+    case SpirvCapability::STORAGE_PUSH_CONSTANT_8:
+      return physical_device_.shader_8bit_storage.storagePushConstant8 ==
+          VK_TRUE;
+#endif
+    case SpirvCapability::DOT_PRODUCT_INPUT_ALL:
+    case SpirvCapability::DOT_PRODUCT_INPUT_4X8_BIT:
+    case SpirvCapability::DOT_PRODUCT_INPUT_4X8_BIT_PACKED:
+    case SpirvCapability::DOT_PRODUCT:
+      return supports_int8_dot_product();
+    case SpirvCapability::COOPERATIVE_MATRIX_KHR:
+      return supports_cooperative_matrix();
+#if defined(VK_NV_cooperative_matrix2) && !defined(ETVK_FORCE_NO_EXTENSIONS)
+    case SpirvCapability::COOPERATIVE_MATRIX_REDUCTIONS_NV:
+      return physical_device_.cooperative_matrix2_features
+                 .cooperativeMatrixReductions == VK_TRUE;
+    case SpirvCapability::COOPERATIVE_MATRIX_CONVERSIONS_NV:
+      return physical_device_.cooperative_matrix2_features
+                 .cooperativeMatrixConversions == VK_TRUE;
+    case SpirvCapability::COOPERATIVE_MATRIX_PER_ELEMENT_OPERATIONS_NV:
+      return physical_device_.cooperative_matrix2_features
+                 .cooperativeMatrixPerElementOperations == VK_TRUE;
+    case SpirvCapability::COOPERATIVE_MATRIX_TENSOR_ADDRESSING_NV:
+      return physical_device_.cooperative_matrix2_features
+                 .cooperativeMatrixTensorAddressing == VK_TRUE;
+    case SpirvCapability::COOPERATIVE_MATRIX_BLOCK_LOADS_NV:
+      return physical_device_.cooperative_matrix2_features
+                 .cooperativeMatrixBlockLoads == VK_TRUE;
+#endif
+    default:
+      // Unhandled requirements have no logical-device feature enablement path.
+      return false;
+  }
+}
+
 } // namespace vkapi
 } // namespace vkcompute

backends/vulkan/runtime/vk_api/Adapter.h

@@ -192,6 +192,8 @@ class Adapter final {
 
   // Physical Device Features
 
+  bool supports_spirv_capability(uint32_t capability);
+
   inline bool supports_16bit_storage_buffers() {
 #ifdef ETVK_FORCE_NO_EXTENSIONS
     return false;