Vulkan: Report and enable subgroup size control device extension.

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:

- The corresponding VkStructureType enum values and
  struct types are not rolled to the third-party Vulkan
  headers used by Dawn yet, so vulkan_platform.h has been
  modified to define them if necessary. This can be
  removed in the future when the Vulkan-Headers are
  updated in a different patch.

- This modifies VulkanDeviceInfo::GatherDeviceInfo() to
  use VkGetPhysicalDevice{Properties2,Features2} if the
  VK_KHR_get_device_properties2 instance extension is
  available. Otherwise, the Vulkan 1.0 APIs
  VkGetPhysicalDevice{Properties,Features} are used instead
  (and it is assumed that no subgroup size control is
  possible).

- 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.

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

Change-Id: I524af6ff3479f25b0a8bb139a062fe632c826893
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/16020
Reviewed-by: Austin Eng <enga@chromium.org>
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 e78718e..2c7e40f 100644
--- a/src/dawn_native/vulkan/DeviceVk.cpp
+++ b/src/dawn_native/vulkan/DeviceVk.cpp
@@ -40,6 +40,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 {
@@ -308,6 +309,22 @@
     ResultOrError<VulkanDeviceKnobs> Device::CreateDevice(VkPhysicalDevice physicalDevice) {
         VulkanDeviceKnobs usedKnobs = {};
 
+        // 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 |hasPhysicalDeviceFeatures2| is true.
+        //
+        bool hasPhysicalDeviceFeatures2 =
+            ToBackend(GetAdapter())->GetBackend()->GetGlobalInfo().getPhysicalDeviceProperties2;
+
+        VkPhysicalDeviceFeatures2 features2 = {
+            .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2,
+            .pNext = nullptr,
+        };
+        PNextChainBuilder featuresChain(&features2);
+
         float zero = 0.0f;
         std::vector<const char*> layersToRequest;
         std::vector<const char*> extensionsToRequest;
@@ -357,6 +374,21 @@
             extensionsToRequest.push_back(kExtensionNameKhrMaintenance1);
             usedKnobs.maintenance1 = true;
         }
+        if (mDeviceInfo.subgroupSizeControl) {
+            // This extension is part of Vulkan 1.1 which always provides support
+            // for VkPhysicalDeviceFeatures2.
+            ASSERT(hasPhysicalDeviceFeatures2);
+
+            // Always require subgroup size control if available.
+            extensionsToRequest.push_back(kExtensionNameExtSubgroupSizeControl);
+            usedKnobs.subgroupSizeControl = true;
+
+            VkPhysicalDeviceSubgroupSizeControlFeaturesEXT* dst =
+                &usedKnobs.featuresExtensions.subgroupSizeControl;
+
+            *dst = mDeviceInfo.featuresExtensions.subgroupSizeControl;
+            featuresChain.Add(dst);
+        }
 
         // Always require independentBlend because it is a core Dawn feature
         usedKnobs.features.independentBlend = VK_TRUE;
@@ -415,6 +447,15 @@
         createInfo.ppEnabledExtensionNames = extensionsToRequest.data();
         createInfo.pEnabledFeatures = &usedKnobs.features;
 
+        if (hasPhysicalDeviceFeatures2 && features2.pNext != nullptr) {
+            // IMPORTANT: To enable features that are not covered by VkPhysicalDeviceFeatures,
+            // one should include a VkPhysicalDeviceFeatures2 struct in the
+            // VkDeviceCreateInfo.pNext chain, and set VkDeviceCreateInfo.pEnabledFeatures to null.
+            features2.features = usedKnobs.features;
+            createInfo.pNext = &features2;
+            createInfo.pEnabledFeatures = nullptr;
+        }
+
         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..cdd89ba 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) everytime 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 VkStructeType 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_TYPE 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/VulkanInfo.cpp b/src/dawn_native/vulkan/VulkanInfo.cpp
index 9bc68ee..be01707 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>
@@ -79,6 +80,7 @@
     const char kExtensionNameKhrXlibSurface[] = "VK_KHR_xlib_surface";
     const char kExtensionNameFuchsiaImagePipeSurface[] = "VK_FUCHSIA_imagepipe_surface";
     const char kExtensionNameKhrMaintenance1[] = "VK_KHR_maintenance1";
+    const char kExtensionNameExtSubgroupSizeControl[] = "VK_EXT_subgroup_size_control";
 
     ResultOrError<VulkanGlobalInfo> GatherGlobalInfo(const Backend& backend) {
         VulkanGlobalInfo info = {};
@@ -224,8 +226,43 @@
         const VulkanFunctions& vkFunctions = adapter.GetBackend()->GetFunctions();
 
         // Gather general info about the device
-        vkFunctions.GetPhysicalDeviceProperties(physicalDevice, &info.properties);
-        vkFunctions.GetPhysicalDeviceFeatures(physicalDevice, &info.features);
+
+        // Use VkGetPhysicalDeviceFeatures2() when available, otherwise fallback
+        // to VkGetPhysicalDeviceFeatures().
+        if (adapter.GetBackend()->GetGlobalInfo().getPhysicalDeviceProperties2) {
+            // NOTE: VkGetPhysicalDevice{Features2,Properties2} will populate
+            // the |features2.features| and |properties2.properties| struct, as
+            // well as the extension structs in their respective pNext chains.
+            VkPhysicalDeviceFeatures2 features2 = {
+                .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2,
+            };
+
+            VkPhysicalDeviceProperties2 properties2 = {
+                .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2,
+            };
+
+            PNextChainBuilder propsChain(&properties2);
+            PNextChainBuilder featuresChain(&features2);
+
+            // Add all known properties extension structs to the chain.
+            propsChain.Add(&info.propertiesExtensions.subgroupSizeControl,
+                           VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_PROPERTIES_EXT);
+
+            vkFunctions.GetPhysicalDeviceProperties2(physicalDevice, &properties2);
+
+            // Add all known features extension structs to the chain.
+            featuresChain.Add(&info.featuresExtensions.subgroupSizeControl,
+                              VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT);
+
+            vkFunctions.GetPhysicalDeviceFeatures2(physicalDevice, &features2);
+
+            info.features = features2.features;
+            info.properties = properties2.properties;
+        } else {
+            // There is no way to query extensions, do read |properties| and |features| directly.
+            vkFunctions.GetPhysicalDeviceProperties(physicalDevice, &info.properties);
+            vkFunctions.GetPhysicalDeviceFeatures(physicalDevice, &info.features);
+        }
 
         // Gather info about device memory.
         {
@@ -311,6 +348,9 @@
                 if (IsExtensionName(extension, kExtensionNameKhrMaintenance1)) {
                     info.maintenance1 = true;
                 }
+                if (IsExtensionName(extension, kExtensionNameExtSubgroupSizeControl)) {
+                    info.subgroupSizeControl = true;
+                }
             }
         }
 
diff --git a/src/dawn_native/vulkan/VulkanInfo.h b/src/dawn_native/vulkan/VulkanInfo.h
index 13a0fdd..82ecea3 100644
--- a/src/dawn_native/vulkan/VulkanInfo.h
+++ b/src/dawn_native/vulkan/VulkanInfo.h
@@ -52,6 +52,7 @@
     extern const char kExtensionNameKhrXlibSurface[];
     extern const char kExtensionNameFuchsiaImagePipeSurface[];
     extern const char kExtensionNameKhrMaintenance1[];
+    extern const char kExtensionNameExtSubgroupSizeControl[];
 
     // Global information - gathered before the instance is created
     struct VulkanGlobalKnobs {
@@ -86,6 +87,12 @@
     struct VulkanDeviceKnobs {
         VkPhysicalDeviceFeatures features;
 
+        // The physical device features that are only accessible when
+        // getPhysicalDeviceProperties2 is true.
+        struct {
+            VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroupSizeControl;
+        } featuresExtensions;
+
         // Extensions, promoted extensions are set to true if their core version is supported.
         bool debugMarker = false;
         bool externalMemory = false;
@@ -98,10 +105,18 @@
         bool externalSemaphoreZirconHandle = false;
         bool swapchain = false;
         bool maintenance1 = false;
+        bool subgroupSizeControl = false;
     };
 
     struct VulkanDeviceInfo : VulkanDeviceKnobs {
         VkPhysicalDeviceProperties properties;
+
+        // The physical device properties that are only accessible when
+        // getPhysicalDeviceProperties2 is true.
+        struct {
+            VkPhysicalDeviceSubgroupSizeControlPropertiesEXT subgroupSizeControl;
+        } propertiesExtensions;
+
         std::vector<VkQueueFamilyProperties> queueFamilies;
 
         std::vector<VkMemoryType> memoryTypes;