Reland: Vulkan: Report and enable subgroup size control device extension.

Reland after a refactor of how the extension handling work in the Vulkan
backend.

The original author is David Turner <david.turner.dev@gmail.com>.

Certain Vulkan ICDs (Intel ones notably) will compile SPIR-V
shaders with an liberal, compiler-selected, subgroup size (i.e.
either 8, 16 or 32). For more context, see [1].

This can be a problem for compute, when one shader stores data
in device memory using a subgroup-size dependent layout, to be
consumed by a another shader. Problems arise when the compiler
decides to compile both shaders with different subgroup sizes.

To work-around this, the VK_EXT_subgroup_size_control device
extension was introduced recently: it allows the device to
report the min/max subgroup sizes it provides, and allows
the Vulkan program to control the subgroup size precisely
if it wants to.

This patch adds support to the Vulkan backend to report and
enable the extension if it is available. Note that:

- This changes the definition of VulkanDeviceKnobs to
  make room for the required pNext-linked chains of
  extensions.

- A helper class, PNextChainBuilder is also provided in
  UtilsVulkan.h to make it easy to build pNext-linked
  extension struct chains at runtime, as required when
  probing device propertires/features, or when
  creating a new VkDevice handle.

- This modifies VulkanDeviceInfo::GatherDeviceInfo() to
  use PNextChainBuilder to query extension features and properties
  in a single call to vkGetPhysicalDevice{Properties,Features}2.

Apart from that, there is no change in behaviour in this CL.
I.e. a later CL might force a specific subgroup size for
consistency, or introduce a new API to let Dawn clients
select a fixed subgroup size.

[1] https://bugs.freedesktop.org/show_bug.cgi?id=108875

Bug: dawn:464
Change-Id: I01e5c28e7dac66f0a57bf35532eb192912b254fa
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/23201
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn_native/vulkan/DeviceVk.cpp b/src/dawn_native/vulkan/DeviceVk.cpp
index db9ecb6..df356f6 100644
--- a/src/dawn_native/vulkan/DeviceVk.cpp
+++ b/src/dawn_native/vulkan/DeviceVk.cpp
@@ -38,6 +38,7 @@
 #include "dawn_native/vulkan/StagingBufferVk.h"
 #include "dawn_native/vulkan/SwapChainVk.h"
 #include "dawn_native/vulkan/TextureVk.h"
+#include "dawn_native/vulkan/UtilsVulkan.h"
 #include "dawn_native/vulkan/VulkanError.h"
 
 namespace dawn_native { namespace vulkan {
@@ -281,6 +282,16 @@
             }
         }
 
+        // Some device features can only be enabled using a VkPhysicalDeviceFeatures2 struct, which
+        // is supported by the VK_EXT_get_physical_properties2 instance extension, which was
+        // promoted as a core API in Vulkan 1.1.
+        //
+        // Prepare a VkPhysicalDeviceFeatures2 struct for this use case, it will only be populated
+        // if HasExt(DeviceExt::GetPhysicalDeviceProperties2) is true.
+        VkPhysicalDeviceFeatures2 features2 = {};
+        features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
+        PNextChainBuilder featuresChain(&features2);
+
         // Always require independentBlend because it is a core Dawn feature
         usedKnobs.features.independentBlend = VK_TRUE;
         // Always require imageCubeArray because it is a core Dawn feature
@@ -288,6 +299,14 @@
         // Always require fragmentStoresAndAtomics because it is required by end2end tests.
         usedKnobs.features.fragmentStoresAndAtomics = VK_TRUE;
 
+        if (mDeviceInfo.HasExt(DeviceExt::SubgroupSizeControl)) {
+            ASSERT(usedKnobs.HasExt(DeviceExt::SubgroupSizeControl));
+
+            // Always request all the features from VK_EXT_subgroup_size_control when available.
+            usedKnobs.subgroupSizeControlFeatures = mDeviceInfo.subgroupSizeControlFeatures;
+            featuresChain.Add(&usedKnobs.subgroupSizeControlFeatures);
+        }
+
         if (IsExtensionEnabled(Extension::TextureCompressionBC)) {
             ASSERT(ToBackend(GetAdapter())->GetDeviceInfo().features.textureCompressionBC ==
                    VK_TRUE);
@@ -349,7 +368,17 @@
         createInfo.ppEnabledLayerNames = nullptr;
         createInfo.enabledExtensionCount = static_cast<uint32_t>(extensionNames.size());
         createInfo.ppEnabledExtensionNames = extensionNames.data();
-        createInfo.pEnabledFeatures = &usedKnobs.features;
+
+        // When we have DeviceExt::GetPhysicalDeviceProperties2, use features2 so that features not
+        // covered by VkPhysicalDeviceFeatures can be enabled.
+        if (mDeviceInfo.HasExt(DeviceExt::GetPhysicalDeviceProperties2)) {
+            features2.features = usedKnobs.features;
+            createInfo.pNext = &features2;
+            createInfo.pEnabledFeatures = nullptr;
+        } else {
+            ASSERT(features2.pNext == nullptr);
+            createInfo.pEnabledFeatures = &usedKnobs.features;
+        }
 
         DAWN_TRY(CheckVkSuccess(fn.CreateDevice(physicalDevice, &createInfo, nullptr, &mVkDevice),
                                 "vkCreateDevice"));
diff --git a/src/dawn_native/vulkan/UtilsVulkan.h b/src/dawn_native/vulkan/UtilsVulkan.h
index 02ef6d3..462c08c 100644
--- a/src/dawn_native/vulkan/UtilsVulkan.h
+++ b/src/dawn_native/vulkan/UtilsVulkan.h
@@ -21,6 +21,70 @@
 
 namespace dawn_native { namespace vulkan {
 
+    // A Helper type used to build a pNext chain of extension structs.
+    // Usage is:
+    //   1) Create instance, passing the address of the first struct in the
+    //      chain. This will parse the existing |pNext| chain in it to find
+    //      its tail.
+    //
+    //   2) Call Add(&vk_struct) every time a new struct needs to be appended
+    //      to the chain.
+    //
+    //   3) Alternatively, call Add(&vk_struct, VK_STRUCTURE_TYPE_XXX) to
+    //      initialize the struct with a given VkStructureType value while
+    //      appending it to the chain.
+    //
+    // Examples:
+    //     VkPhysicalFeatures2 features2 = {
+    //       .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2,
+    //       .pNext = nullptr,
+    //     };
+    //
+    //     PNextChainBuilder featuresChain(&features2);
+    //
+    //     featuresChain.Add(&featuresExtensions.subgroupSizeControl,
+    //                       VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT);
+    //
+    struct PNextChainBuilder {
+        // Constructor takes the address of a Vulkan structure instance, and
+        // walks its pNext chain to record the current location of its tail.
+        //
+        // NOTE: Some VK_STRUCT_TYPEs define their pNext field as a const void*
+        // which is why the VkBaseOutStructure* casts below are necessary.
+        template <typename VK_STRUCT_TYPE>
+        explicit PNextChainBuilder(VK_STRUCT_TYPE* head)
+            : mTailPtr(&reinterpret_cast<VkBaseOutStructure*>(head)->pNext) {
+            // Find the end of the current chain.
+            while (*mTailPtr) {
+                mTailPtr = &(*mTailPtr)->pNext;
+            }
+        }
+
+        // Add one item to the chain. |vk_struct| must be a Vulkan structure
+        // that is already initialized.
+        template <typename VK_STRUCT_TYPE>
+        void Add(VK_STRUCT_TYPE* vkStruct) {
+            // Sanity checks to ensure proper type safety.
+            static_assert(
+                offsetof(VK_STRUCT_TYPE, sType) == offsetof(VkBaseOutStructure, sType) &&
+                    offsetof(VK_STRUCT_TYPE, pNext) == offsetof(VkBaseOutStructure, pNext),
+                "Argument type is not a proper Vulkan structure type");
+            vkStruct->pNext = nullptr;
+            *mTailPtr = reinterpret_cast<VkBaseOutStructure*>(vkStruct);
+            mTailPtr = &(*mTailPtr)->pNext;
+        }
+
+        // A variant of Add() above that also initializes the |sType| field in |vk_struct|.
+        template <typename VK_STRUCT_TYPE>
+        void Add(VK_STRUCT_TYPE* vkStruct, VkStructureType sType) {
+            vkStruct->sType = sType;
+            Add(vkStruct);
+        }
+
+      private:
+        VkBaseOutStructure** mTailPtr;
+    };
+
     VkCompareOp ToVulkanCompareOp(wgpu::CompareFunction op);
 
     Extent3D ComputeTextureCopyExtent(const TextureCopy& textureCopy, const Extent3D& copySize);
diff --git a/src/dawn_native/vulkan/VulkanExtensions.cpp b/src/dawn_native/vulkan/VulkanExtensions.cpp
index 9015a6d..1e7f23d 100644
--- a/src/dawn_native/vulkan/VulkanExtensions.cpp
+++ b/src/dawn_native/vulkan/VulkanExtensions.cpp
@@ -164,6 +164,7 @@
         {DeviceExt::DebugMarker, "VK_EXT_debug_marker", NeverPromoted},
         {DeviceExt::ImageDrmFormatModifier, "VK_EXT_image_drm_format_modifier", NeverPromoted},
         {DeviceExt::Swapchain, "VK_KHR_swapchain", NeverPromoted},
+        {DeviceExt::SubgroupSizeControl, "VK_EXT_subgroup_size_control", NeverPromoted},
         //
     }};
 
@@ -191,7 +192,8 @@
     }
 
     DeviceExtSet EnsureDependencies(const DeviceExtSet& advertisedExts,
-                                    const InstanceExtSet& instanceExts) {
+                                    const InstanceExtSet& instanceExts,
+                                    uint32_t icdVersion) {
         // This is very similar to EnsureDependencies for instanceExtSet. See comment there for
         // an explanation of what happens.
         DeviceExtSet visitedSet;
@@ -288,6 +290,13 @@
                                       HasDep(DeviceExt::StorageBufferStorageClass);
                     break;
 
+                case DeviceExt::SubgroupSizeControl:
+                    // Using the extension requires DeviceExt::GetPhysicalDeviceProperties2, but we
+                    // don't need to check for it as it also requires Vulkan 1.1 in which
+                    // VK_KHR_get_physical_device_properties2 was promoted.
+                    hasDependencies = icdVersion >= VulkanVersion_1_1;
+                    break;
+
                 default:
                     UNREACHABLE();
                     break;
diff --git a/src/dawn_native/vulkan/VulkanExtensions.h b/src/dawn_native/vulkan/VulkanExtensions.h
index 5c59bc8..ba6abc2 100644
--- a/src/dawn_native/vulkan/VulkanExtensions.h
+++ b/src/dawn_native/vulkan/VulkanExtensions.h
@@ -101,6 +101,7 @@
         DebugMarker,
         ImageDrmFormatModifier,
         Swapchain,
+        SubgroupSizeControl,
 
         EnumCount,
     };
@@ -132,7 +133,8 @@
     // extensions that don't have all their transitive dependencies in advertisedExts or in
     // instanceExts.
     DeviceExtSet EnsureDependencies(const DeviceExtSet& advertisedExts,
-                                    const InstanceExtSet& instanceExts);
+                                    const InstanceExtSet& instanceExts,
+                                    uint32_t icdVersion);
 
 }}  // namespace dawn_native::vulkan
 
diff --git a/src/dawn_native/vulkan/VulkanInfo.cpp b/src/dawn_native/vulkan/VulkanInfo.cpp
index 6b2346e..5a66c67 100644
--- a/src/dawn_native/vulkan/VulkanInfo.cpp
+++ b/src/dawn_native/vulkan/VulkanInfo.cpp
@@ -17,6 +17,7 @@
 #include "common/Log.h"
 #include "dawn_native/vulkan/AdapterVk.h"
 #include "dawn_native/vulkan/BackendVk.h"
+#include "dawn_native/vulkan/UtilsVulkan.h"
 #include "dawn_native/vulkan/VulkanError.h"
 
 #include <cstring>
@@ -179,9 +180,8 @@
         const VulkanGlobalInfo& globalInfo = adapter.GetBackend()->GetGlobalInfo();
         const VulkanFunctions& vkFunctions = adapter.GetBackend()->GetFunctions();
 
-        // Gather general info about the device
+        // Query the device properties first to get the ICD's `apiVersion`
         vkFunctions.GetPhysicalDeviceProperties(physicalDevice, &info.properties);
-        vkFunctions.GetPhysicalDeviceFeatures(physicalDevice, &info.features);
 
         // Gather info about device memory.
         {
@@ -245,30 +245,54 @@
             }
 
             MarkPromotedExtensions(&info.extensions, info.properties.apiVersion);
-            info.extensions = EnsureDependencies(info.extensions, globalInfo.extensions);
+            info.extensions = EnsureDependencies(info.extensions, globalInfo.extensions,
+                                                 info.properties.apiVersion);
         }
 
-        // Gather additional information for some of the extensions
-        {
-            if (info.extensions.Has(DeviceExt::ShaderFloat16Int8)) {
-                info.shaderFloat16Int8Features.sType =
-                    VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES_KHR;
+        // Gather general and extension features and properties
+        //
+        // Use vkGetPhysicalDevice{Features,Properties}2 if required to gather information about
+        // the extensions. DeviceExt::GetPhysicalDeviceProperties2 is guaranteed to be available
+        // because these extensions (transitively) depend on it in `EnsureDependencies`
+        VkPhysicalDeviceFeatures2 features2 = {};
+        features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
+        PNextChainBuilder featuresChain(&features2);
 
-                VkPhysicalDeviceFeatures2KHR physicalDeviceFeatures2 = {};
-                physicalDeviceFeatures2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2_KHR;
-                physicalDeviceFeatures2.pNext = &info.shaderFloat16Int8Features;
-                vkFunctions.GetPhysicalDeviceFeatures2(physicalDevice, &physicalDeviceFeatures2);
-            }
+        VkPhysicalDeviceProperties2 properties2 = {};
+        properties2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;
+        PNextChainBuilder propertiesChain(&properties2);
 
-            if (info.extensions.Has(DeviceExt::_16BitStorage)) {
-                info._16BitStorageFeatures.sType =
-                    VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES;
+        if (info.extensions.Has(DeviceExt::ShaderFloat16Int8)) {
+            featuresChain.Add(&info.shaderFloat16Int8Features,
+                              VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES_KHR);
+        }
 
-                VkPhysicalDeviceFeatures2 physicalDeviceFeatures2 = {};
-                physicalDeviceFeatures2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
-                physicalDeviceFeatures2.pNext = &info._16BitStorageFeatures;
-                vkFunctions.GetPhysicalDeviceFeatures2(physicalDevice, &physicalDeviceFeatures2);
-            }
+        if (info.extensions.Has(DeviceExt::_16BitStorage)) {
+            featuresChain.Add(&info._16BitStorageFeatures,
+                              VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES);
+        }
+
+        if (info.extensions.Has(DeviceExt::SubgroupSizeControl)) {
+            featuresChain.Add(&info.subgroupSizeControlFeatures,
+                              VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT);
+            propertiesChain.Add(
+                &info.subgroupSizeControlProperties,
+                VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_PROPERTIES_EXT);
+        }
+
+        // If we have DeviceExt::GetPhysicalDeviceProperties2, use features2 and properties2 so
+        // that features no covered by VkPhysicalDevice{Features,Properties} can be queried.
+        //
+        // Note that info.properties has already been filled at the start of this function to get
+        // `apiVersion`.
+        ASSERT(info.properties.apiVersion != 0);
+        if (info.extensions.Has(DeviceExt::GetPhysicalDeviceProperties2)) {
+            vkFunctions.GetPhysicalDeviceProperties2(physicalDevice, &properties2);
+            vkFunctions.GetPhysicalDeviceFeatures2(physicalDevice, &features2);
+            info.features = features2.features;
+        } else {
+            ASSERT(features2.pNext == nullptr && properties2.pNext == nullptr);
+            vkFunctions.GetPhysicalDeviceFeatures(physicalDevice, &info.features);
         }
 
         // TODO(cwallez@chromium.org): gather info about formats
diff --git a/src/dawn_native/vulkan/VulkanInfo.h b/src/dawn_native/vulkan/VulkanInfo.h
index 6d0a756f..e051247 100644
--- a/src/dawn_native/vulkan/VulkanInfo.h
+++ b/src/dawn_native/vulkan/VulkanInfo.h
@@ -54,6 +54,7 @@
         VkPhysicalDeviceFeatures features;
         VkPhysicalDeviceShaderFloat16Int8FeaturesKHR shaderFloat16Int8Features;
         VkPhysicalDevice16BitStorageFeaturesKHR _16BitStorageFeatures;
+        VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroupSizeControlFeatures;
 
         bool HasExt(DeviceExt ext) const;
         DeviceExtSet extensions;
@@ -61,6 +62,8 @@
 
     struct VulkanDeviceInfo : VulkanDeviceKnobs {
         VkPhysicalDeviceProperties properties;
+        VkPhysicalDeviceSubgroupSizeControlPropertiesEXT subgroupSizeControlProperties;
+
         std::vector<VkQueueFamilyProperties> queueFamilies;
 
         std::vector<VkMemoryType> memoryTypes;