Fix erroneous extension_features initialization

Summary:
## Context

Some devices were hanging when trying to run Vulkan due to a bug in pytorch#6910.

The issue is that `extension_features` was being initialized with `&shader_float16_int8_types` before `shader_float16_int8_types` was initialized.

## Changes

This diff fixes the bug and changes extension_features to set to an appropriate value based on which extensions are available according to the vulkan headers being used.

Differential Revision: D66118505

Author: SS-JIA

backends/vulkan/runtime/vk_api/Device.cpp

@@ -23,11 +23,8 @@ PhysicalDevice::PhysicalDevice(VkPhysicalDevice physical_device_handle)
       properties{},
       memory_properties{},
 #ifdef VK_KHR_16bit_storage
-      extension_features(&shader_16bit_storage),
       shader_16bit_storage{
           VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES},
-#else
-      extension_features{nullptr},
 #endif /* VK_KHR_16bit_storage */
 #ifdef VK_KHR_8bit_storage
       shader_8bit_storage{
@@ -37,6 +34,7 @@ PhysicalDevice::PhysicalDevice(VkPhysicalDevice physical_device_handle)
       shader_float16_int8_types{
           VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES_KHR},
 #endif /* VK_KHR_shader_float16_int8 */
+      extension_features{nullptr},
       queue_families{},
       num_compute_queues(0),
       supports_int16_shader_types(false),
@@ -53,10 +51,13 @@ PhysicalDevice::PhysicalDevice(VkPhysicalDevice physical_device_handle)
   // Create linked list to query availability of extensions
 
 #ifdef VK_KHR_16bit_storage
+  extension_features = &shader_16bit_storage;
   features2.pNext = &shader_16bit_storage;
 #elif defined(VK_KHR_8bit_storage)
+  extension_features = &shader_8bit_storage;
   features2.pNext = &shader_8bit_storage;
 #elif defined(VK_KHR_shader_float16_int8)
+  extension_features = &shader_float16_int8_types;
   features2.pNext = &shader_float16_int8_types;
 #endif /* VK_KHR_16bit_storage */
 

backends/vulkan/runtime/vk_api/Device.h

@@ -26,9 +26,6 @@ struct PhysicalDevice final {
   VkPhysicalDeviceProperties properties;
   VkPhysicalDeviceMemoryProperties memory_properties;
 
-  // Head of the linked list of extensions to be requested
-  void* extension_features;
-
   // Additional features available from extensions
 #ifdef VK_KHR_16bit_storage
   VkPhysicalDevice16BitStorageFeatures shader_16bit_storage;
@@ -40,6 +37,9 @@ struct PhysicalDevice final {
   VkPhysicalDeviceShaderFloat16Int8Features shader_float16_int8_types;
 #endif /* VK_KHR_shader_float16_int8 */
 
+  // Head of the linked list of extensions to be requested
+  void* extension_features;
+
   // Available GPU queues
   std::vector<VkQueueFamilyProperties> queue_families;