From 56f28ac3efb5b14c6c0c7c906c4c74a3e14de428 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Fri, 21 May 2021 09:23:06 -0700
Subject: [PATCH 1/9] [Vulkan] Enable instance/device extensions

- Vulkan requires that extensions be explicitly enabled if used.
  Explicitly list out which extensions are required (currently none)
  and which are optional.
---
 src/runtime/vulkan/vulkan.cc       | 152 +++++++++++++++--------------
 src/runtime/vulkan/vulkan_common.h |  25 ++++-
 2 files changed, 101 insertions(+), 76 deletions(-)

diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index b7fe2b1ceb21..111b09228131 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -427,6 +427,33 @@ class VulkanDeviceAPI final : public DeviceAPI {
   }
 
  private:
+  std::vector<const char*> find_enabled_extensions(
+      const std::vector<VkExtensionProperties>& ext_prop,
+      const std::vector<const char*>& required_extensions,
+      const std::vector<const char*>& optional_extensions) {
+    std::set<std::string> available_extensions;
+    for (const auto& prop : ext_prop) {
+      if (prop.specVersion > 0) {
+        available_extensions.insert(prop.extensionName);
+      }
+    }
+
+    std::vector<const char*> enabled_extensions;
+    for (const auto& ext : required_extensions) {
+      ICHECK(available_extensions.count(ext))
+          << "Required vulkan extension \"" << ext << "\" not supported by driver";
+      enabled_extensions.push_back(ext);
+    }
+
+    for (const auto& ext : optional_extensions) {
+      if (available_extensions.count(ext)) {
+        enabled_extensions.push_back(ext);
+      }
+    }
+
+    return enabled_extensions;
+  }
+
   VkInstance instance_{nullptr};
   // The physical devices, have 1 to 1 mapping to devices
   std::vector<VulkanContext> context_;
@@ -516,20 +543,6 @@ void VulkanDeviceAPI::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv)
 }
 
 VulkanDeviceAPI::VulkanDeviceAPI() {
-  VkApplicationInfo app_info;
-  app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
-  app_info.pNext = nullptr;
-  app_info.pApplicationName = "TVM";
-  app_info.applicationVersion = 0;
-  app_info.pEngineName = "";
-  app_info.engineVersion = 0;
-  app_info.apiVersion = VK_MAKE_VERSION(1, 0, 0);
-
-  VkInstanceCreateInfo inst_info;
-  inst_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
-  inst_info.pNext = nullptr;
-  inst_info.flags = 0;
-
   const auto layers = []() -> std::vector<const char*> {
     uint32_t inst_layer_prop_count;
     VULKAN_CALL(vkEnumerateInstanceLayerProperties(&inst_layer_prop_count, nullptr));
@@ -554,29 +567,42 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     return l;
   }();
 
-  const auto instance_extensions = []() -> std::vector<const char*> {
+  const auto instance_extensions = [this]() {
+    std::vector<const char*> required_extensions{"VK_KHR_get_physical_device_properties2"};
+    std::vector<const char*> optional_extensions{};
+
     uint32_t inst_extension_prop_count;
     VULKAN_CALL(
         vkEnumerateInstanceExtensionProperties(nullptr, &inst_extension_prop_count, nullptr));
     std::vector<VkExtensionProperties> inst_extension_prop(inst_extension_prop_count);
     VULKAN_CALL(vkEnumerateInstanceExtensionProperties(nullptr, &inst_extension_prop_count,
                                                        inst_extension_prop.data()));
-    std::vector<const char*> extensions;
-    for (const auto& ip : inst_extension_prop) {
-      if (std::strcmp(ip.extensionName, "VK_KHR_get_physical_device_properties2") == 0) {
-        extensions.push_back("VK_KHR_get_physical_device_properties2");
-      }
-    }
-    return extensions;
-  }();
 
-  inst_info.pApplicationInfo = &app_info;
-  inst_info.enabledLayerCount = layers.size();
-  inst_info.ppEnabledLayerNames = layers.data();
-  inst_info.enabledExtensionCount = instance_extensions.size();
-  inst_info.ppEnabledExtensionNames = instance_extensions.data();
+    return find_enabled_extensions(inst_extension_prop, required_extensions, optional_extensions);
+  }();
 
-  VULKAN_CALL(vkCreateInstance(&inst_info, nullptr, &instance_));
+  {
+    VkApplicationInfo app_info;
+    app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
+    app_info.pNext = nullptr;
+    app_info.pApplicationName = "TVM";
+    app_info.applicationVersion = 0;
+    app_info.pEngineName = "";
+    app_info.engineVersion = 0;
+    app_info.apiVersion = VK_MAKE_VERSION(1, 0, 0);
+
+    VkInstanceCreateInfo inst_info;
+    inst_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
+    inst_info.pNext = nullptr;
+    inst_info.flags = 0;
+    inst_info.pApplicationInfo = &app_info;
+    inst_info.enabledLayerCount = layers.size();
+    inst_info.ppEnabledLayerNames = layers.data();
+    inst_info.enabledExtensionCount = instance_extensions.size();
+    inst_info.ppEnabledExtensionNames = instance_extensions.data();
+
+    VULKAN_CALL(vkCreateInstance(&inst_info, nullptr, &instance_));
+  }
 
   uint32_t phy_dev_count = 0;
   VULKAN_CALL(vkEnumeratePhysicalDevices(instance_, &phy_dev_count, nullptr));
@@ -603,32 +629,30 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     ctx.phy_device = phy_dev;
     vkGetPhysicalDeviceProperties(ctx.phy_device, &(ctx.phy_device_prop));
 
-    const auto extensions = [&]() {
+    const auto device_extensions = [&]() {
+      std::vector<const char*> required_extensions{};
+      std::vector<const char*> optional_extensions{
+          "VK_KHR_driver_properties",
+          "VK_KHR_storage_buffer_storage_class",
+          "VK_KHR_8bit_storage",
+          "VK_KHR_16bit_storage",
+          "VK_KHR_shader_float16_int8",
+          "VK_KHR_push_descriptor",
+          "VK_KHR_descriptor_update_template",
+          "VK_KHR_get_memory_requirements2",
+          "VK_KHR_dedicated_allocation",
+          "VK_KHR_spirv_1_4",
+      };
+
       uint32_t device_extension_prop_count;
       VULKAN_CALL(vkEnumerateDeviceExtensionProperties(ctx.phy_device, nullptr,
                                                        &device_extension_prop_count, nullptr));
       std::vector<VkExtensionProperties> device_extension_prop(device_extension_prop_count);
       VULKAN_CALL(vkEnumerateDeviceExtensionProperties(
           ctx.phy_device, nullptr, &device_extension_prop_count, device_extension_prop.data()));
-      std::vector<const char*> extensions;
-      for (const auto& dp : device_extension_prop) {
-        if ((std::strcmp(dp.extensionName, "VK_KHR_push_descriptor") == 0) && dp.specVersion > 0) {
-          extensions.push_back("VK_KHR_push_descriptor");
-        }
-        if ((std::strcmp(dp.extensionName, "VK_KHR_descriptor_update_template") == 0) &&
-            dp.specVersion > 0) {
-          extensions.push_back("VK_KHR_descriptor_update_template");
-        }
-        if ((std::strcmp(dp.extensionName, "VK_KHR_get_memory_requirements2") == 0) &&
-            dp.specVersion > 0) {
-          extensions.push_back("VK_KHR_get_memory_requirements2");
-        }
-        if ((std::strcmp(dp.extensionName, "VK_KHR_dedicated_allocation") == 0) &&
-            dp.specVersion > 0) {
-          extensions.push_back("VK_KHR_dedicated_allocation");
-        }
-      }
-      return extensions;
+
+      return find_enabled_extensions(device_extension_prop, required_extensions,
+                                     optional_extensions);
     }();
 
     // All TVM-generated spirv shaders are marked as requiring int64
@@ -644,8 +668,8 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     device_create_info.pQueueCreateInfos = &queue_create_info;
     device_create_info.enabledLayerCount = 0;
     device_create_info.ppEnabledLayerNames = nullptr;
-    device_create_info.enabledExtensionCount = extensions.size();
-    device_create_info.ppEnabledExtensionNames = extensions.data();
+    device_create_info.enabledExtensionCount = device_extensions.size();
+    device_create_info.ppEnabledExtensionNames = device_extensions.data();
     device_create_info.pEnabledFeatures = &enabled_features;
     VULKAN_CALL(vkCreateDevice(phy_dev, &device_create_info, nullptr, &(ctx.device)));
     ctx.queue_mutex.reset(new std::mutex());
@@ -718,28 +742,16 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
       }
     }
     ICHECK_GE(win_rank, 0) << "Cannot find suitable local memory on device.";
-    auto has_extension = [&extensions](const char* query) {
-      return std::any_of(extensions.begin(), extensions.end(),
+    auto has_extension = [&device_extensions](const char* query) {
+      return std::any_of(device_extensions.begin(), device_extensions.end(),
                          [&](const char* extension) { return std::strcmp(query, extension) == 0; });
     };
 
 #ifdef USE_VULKAN_IMMEDIATE_MODE
     if (has_extension("VK_KHR_push_descriptor") &&
         has_extension("VK_KHR_descriptor_update_template")) {
-      ctx.descriptor_template_khr_functions = std::unique_ptr<VulkanDescriptorTemplateKHRFunctions>(
-          new VulkanDescriptorTemplateKHRFunctions());
-      ctx.descriptor_template_khr_functions->vkCreateDescriptorUpdateTemplateKHR =
-          CHECK_NOTNULL((PFN_vkCreateDescriptorUpdateTemplateKHR)vkGetDeviceProcAddr(
-              ctx.device, "vkCreateDescriptorUpdateTemplateKHR"));
-      ctx.descriptor_template_khr_functions->vkDestroyDescriptorUpdateTemplateKHR =
-          CHECK_NOTNULL((PFN_vkDestroyDescriptorUpdateTemplateKHR)vkGetDeviceProcAddr(
-              ctx.device, "vkDestroyDescriptorUpdateTemplateKHR"));
-      ctx.descriptor_template_khr_functions->vkUpdateDescriptorSetWithTemplateKHR =
-          CHECK_NOTNULL((PFN_vkUpdateDescriptorSetWithTemplateKHR)vkGetDeviceProcAddr(
-              ctx.device, "vkUpdateDescriptorSetWithTemplateKHR"));
-      ctx.descriptor_template_khr_functions->vkCmdPushDescriptorSetWithTemplateKHR =
-          CHECK_NOTNULL((PFN_vkCmdPushDescriptorSetWithTemplateKHR)vkGetDeviceProcAddr(
-              ctx.device, "vkCmdPushDescriptorSetWithTemplateKHR"));
+      ctx.descriptor_template_khr_functions =
+          std::make_unique<VulkanDescriptorTemplateKHRFunctions>(ctx.device);
     }
 #endif
 
@@ -747,11 +759,7 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     if (has_extension("VK_KHR_get_memory_requirements2") &&
         has_extension("VK_KHR_dedicated_allocation")) {
       ctx.get_buffer_memory_requirements_2_functions =
-          std::unique_ptr<VulkanGetBufferMemoryRequirements2Functions>(
-              new VulkanGetBufferMemoryRequirements2Functions());
-      ctx.get_buffer_memory_requirements_2_functions->vkGetBufferMemoryRequirements2KHR =
-          CHECK_NOTNULL((PFN_vkGetBufferMemoryRequirements2KHR)vkGetDeviceProcAddr(
-              ctx.device, "vkGetBufferMemoryRequirements2KHR"));
+          std::make_unique<VulkanGetBufferMemoryRequirements2Functions>(ctx.device);
     }
 #endif
     context_.push_back(std::move(ctx));
diff --git a/src/runtime/vulkan/vulkan_common.h b/src/runtime/vulkan/vulkan_common.h
index 2ef879a487a6..8cdbaacbbb19 100644
--- a/src/runtime/vulkan/vulkan_common.h
+++ b/src/runtime/vulkan/vulkan_common.h
@@ -87,10 +87,10 @@ inline const char* VKGetErrorString(VkResult error) {
  * \brief Protected Vulkan call
  * \param func Expression to call.
  */
-#define VULKAN_CHECK_ERROR(__e)                                      \
-  {                                                                  \
-    ICHECK(__e == VK_SUCCESS) << "Vulan Error, code=" << __e << ": " \
-                              << vulkan::VKGetErrorString(__e);      \
+#define VULKAN_CHECK_ERROR(__e)                                       \
+  {                                                                   \
+    ICHECK(__e == VK_SUCCESS) << "Vulkan Error, code=" << __e << ": " \
+                              << vulkan::VKGetErrorString(__e);       \
   }
 
 #define VULKAN_CALL(func)    \
@@ -100,6 +100,18 @@ inline const char* VKGetErrorString(VkResult error) {
   }
 
 struct VulkanDescriptorTemplateKHRFunctions {
+  explicit VulkanDescriptorTemplateKHRFunctions(VkDevice device) {
+    vkCreateDescriptorUpdateTemplateKHR = (PFN_vkCreateDescriptorUpdateTemplateKHR)ICHECK_NOTNULL(
+        vkGetDeviceProcAddr(device, "vkCreateDescriptorUpdateTemplateKHR"));
+    vkDestroyDescriptorUpdateTemplateKHR = (PFN_vkDestroyDescriptorUpdateTemplateKHR)ICHECK_NOTNULL(
+        vkGetDeviceProcAddr(device, "vkDestroyDescriptorUpdateTemplateKHR"));
+    vkUpdateDescriptorSetWithTemplateKHR = (PFN_vkUpdateDescriptorSetWithTemplateKHR)ICHECK_NOTNULL(
+        vkGetDeviceProcAddr(device, "vkUpdateDescriptorSetWithTemplateKHR"));
+    vkCmdPushDescriptorSetWithTemplateKHR =
+        (PFN_vkCmdPushDescriptorSetWithTemplateKHR)ICHECK_NOTNULL(
+            vkGetDeviceProcAddr(device, "vkCmdPushDescriptorSetWithTemplateKHR"));
+  }
+
   PFN_vkCreateDescriptorUpdateTemplateKHR vkCreateDescriptorUpdateTemplateKHR{nullptr};
   PFN_vkDestroyDescriptorUpdateTemplateKHR vkDestroyDescriptorUpdateTemplateKHR{nullptr};
   PFN_vkUpdateDescriptorSetWithTemplateKHR vkUpdateDescriptorSetWithTemplateKHR{nullptr};
@@ -107,6 +119,11 @@ struct VulkanDescriptorTemplateKHRFunctions {
 };
 
 struct VulkanGetBufferMemoryRequirements2Functions {
+  explicit VulkanGetBufferMemoryRequirements2Functions(VkDevice device) {
+    vkGetBufferMemoryRequirements2KHR = (PFN_vkGetBufferMemoryRequirements2KHR)ICHECK_NOTNULL(
+        vkGetDeviceProcAddr(device, "vkGetBufferMemoryRequirements2KHR"));
+  }
+
   PFN_vkGetBufferMemoryRequirements2KHR vkGetBufferMemoryRequirements2KHR{nullptr};
 };
 

From 1882c270c4cdd4ca8da9e34e678c1b6236b342c9 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Mon, 24 May 2021 11:16:37 -0700
Subject: [PATCH 2/9] [Vulkan] Extract device information from vulkan API.

- Based on vkGetPhysicalDeviceProperties and
  vkGetPhysicalDeviceFeatures, determine which Vulkan capabilities are
  supported, pack into a Target.
---
 src/runtime/vulkan/vulkan.cc       | 166 ++++++++++++++++++++++++++++-
 src/runtime/vulkan/vulkan_common.h |   4 +-
 src/target/target_kind.cc          |  27 +++++
 3 files changed, 193 insertions(+), 4 deletions(-)

diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index 111b09228131..a1114087528d 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -21,6 +21,7 @@
 #include <dmlc/thread_local.h>
 #include <tvm/runtime/device_api.h>
 #include <tvm/runtime/registry.h>
+#include <tvm/target/target.h>
 #include <vulkan/vulkan.h>
 #include <vulkan/vulkan_core.h>
 
@@ -262,6 +263,10 @@ class VulkanDeviceAPI final : public DeviceAPI {
     delete pbuf;
   }
 
+  Target GetDeviceDescription(VkInstance instance, VkPhysicalDevice dev,
+                              const std::vector<const char*>& instance_extensions,
+                              const std::vector<const char*>& device_extensions);
+
  protected:
   void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset, size_t size,
                       Device dev_from, Device dev_to, DLDataType type_hint,
@@ -459,6 +464,157 @@ class VulkanDeviceAPI final : public DeviceAPI {
   std::vector<VulkanContext> context_;
 };
 
+Target VulkanDeviceAPI::GetDeviceDescription(VkInstance instance, VkPhysicalDevice dev,
+                                             const std::vector<const char*>& instance_extensions,
+                                             const std::vector<const char*>& device_extensions) {
+  auto has_extension = [&](const char* query) {
+    return std::any_of(device_extensions.begin(), device_extensions.end(),
+                       [&](const char* extension) { return std::strcmp(query, extension) == 0; }) ||
+           std::any_of(instance_extensions.begin(), instance_extensions.end(),
+                       [&](const char* extension) { return std::strcmp(query, extension) == 0; });
+  };
+
+  // Declare output locations for properties
+  VkPhysicalDeviceProperties2 properties = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2};
+  VkPhysicalDeviceDriverProperties driver = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DRIVER_PROPERTIES};
+  VkPhysicalDeviceSubgroupProperties subgroup = {
+      VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_PROPERTIES};
+
+  // Need to do initial query in order to check the apiVersion.
+  vkGetPhysicalDeviceProperties(dev, &properties.properties);
+
+  // Set up linked list for property query
+  {
+    void** pp_next = &properties.pNext;
+    if (has_extension("VK_KHR_driver_properties")) {
+      *pp_next = &driver;
+      pp_next = &driver.pNext;
+    }
+    if (properties.properties.apiVersion >= VK_API_VERSION_1_1) {
+      *pp_next = &subgroup;
+      pp_next = &subgroup.pNext;
+    }
+  }
+
+  // Declare output locations for features
+  VkPhysicalDeviceFeatures2 features = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2};
+  VkPhysicalDevice8BitStorageFeatures storage_8bit = {
+      VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_8BIT_STORAGE_FEATURES};
+  VkPhysicalDevice16BitStorageFeatures storage_16bit = {
+      VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES};
+  VkPhysicalDeviceShaderFloat16Int8Features float16_int8 = {
+      VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES};
+
+  // Set up linked list for feature query
+  {
+    void** pp_next = &features.pNext;
+    if (has_extension("VK_KHR_8bit_storage")) {
+      *pp_next = &storage_8bit;
+      pp_next = &storage_8bit.pNext;
+    }
+    if (has_extension("VK_KHR_16bit_storage")) {
+      *pp_next = &storage_16bit;
+      pp_next = &storage_16bit.pNext;
+    }
+    if (has_extension("VK_KHR_shader_float16_int8")) {
+      *pp_next = &float16_int8;
+      pp_next = &float16_int8.pNext;
+    }
+  }
+
+  if (has_extension("VK_KHR_get_physical_device_properties2")) {
+    // Preferred method, call to get all properties that can be queried.
+    auto vkGetPhysicalDeviceProperties2KHR = (PFN_vkGetPhysicalDeviceProperties2KHR)ICHECK_NOTNULL(
+        vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceProperties2KHR"));
+    vkGetPhysicalDeviceProperties2KHR(dev, &properties);
+
+    auto vkGetPhysicalDeviceFeatures2KHR = (PFN_vkGetPhysicalDeviceFeatures2KHR)ICHECK_NOTNULL(
+        vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceFeatures2KHR"));
+    vkGetPhysicalDeviceFeatures2KHR(dev, &features);
+  } else {
+    // Fallback, get as many features as we can from the Vulkan1.0
+    // API.  Corresponding vkGetPhysicalDeviceProperties was already done earlier.
+    vkGetPhysicalDeviceFeatures(dev, &features.features);
+  }
+
+  //// Now, extracting all the information from the vulkan query.
+
+  // Not technically needed, because VK_SHADER_STAGE_COMPUTE_BIT will
+  // be set so long at least one queue has VK_QUEUE_COMPUTE_BIT, but
+  // preferring the explicit check.
+  uint32_t supported_subgroup_operations =
+      (subgroup.supportedStages & VK_SHADER_STAGE_COMPUTE_BIT) ? subgroup.supportedOperations : 0;
+
+  // Even if we can't query it, warp size must be at least 1.  Must
+  // also be defined, as `transpose` operation requires it.
+  uint32_t thread_warp_size = std::max(subgroup.subgroupSize, 1U);
+
+  // By default, use the maximum API version that the driver allows,
+  // so that any supported features can be used by TVM shaders.
+  // However, if we can query the conformance version, then limit to
+  // only using the api version that passes the vulkan conformance
+  // tests.
+  uint32_t vulkan_api_version = properties.properties.apiVersion;
+  if (has_extension("VK_KHR_driver_properties")) {
+    auto api_major = VK_VERSION_MAJOR(vulkan_api_version);
+    auto api_minor = VK_VERSION_MINOR(vulkan_api_version);
+    if ((api_major > driver.conformanceVersion.major) ||
+        ((api_major == driver.conformanceVersion.major) &&
+         (api_minor > driver.conformanceVersion.minor))) {
+      vulkan_api_version =
+          VK_MAKE_VERSION(driver.conformanceVersion.major, driver.conformanceVersion.minor, 0);
+    }
+  }
+
+  // From "Versions and Formats" section of Vulkan spec.
+  uint32_t max_spirv_version = 0x10000;
+  if (vulkan_api_version >= VK_API_VERSION_1_2) {
+    max_spirv_version = 0x10500;
+  } else if (has_extension("VK_KHR_spirv_1_4")) {
+    max_spirv_version = 0x10400;
+  } else if (vulkan_api_version >= VK_API_VERSION_1_1) {
+    max_spirv_version = 0x10300;
+  }
+
+  Map<String, ObjectRef> config = {
+      {"kind", String("vulkan")},
+      // Feature support
+      {"supports_float16", Bool(float16_int8.shaderFloat16)},
+      {"supports_float32", Bool(true)},
+      {"supports_float64", Bool(features.features.shaderFloat64)},
+      {"supports_int8", Bool(float16_int8.shaderInt8)},
+      {"supports_int16", Bool(features.features.shaderInt16)},
+      {"supports_int32", Bool(true)},
+      {"supports_int64", Bool(features.features.shaderInt64)},
+      {"supports_8bit_buffer", Bool(storage_8bit.storageBuffer8BitAccess)},
+      {"supports_16bit_buffer", Bool(storage_16bit.storageBuffer16BitAccess)},
+      {"supports_storage_buffer_storage_class",
+       Bool(has_extension("VK_KHR_storage_buffer_storage_class"))},
+      {"supported_subgroup_operations", Integer(supported_subgroup_operations)},
+      // Physical device limits
+      {"max_num_threads", Integer(properties.properties.limits.maxComputeWorkGroupInvocations)},
+      {"thread_warp_size", Integer(thread_warp_size)},
+      {"max_block_size_x", Integer(properties.properties.limits.maxComputeWorkGroupSize[0])},
+      {"max_block_size_y", Integer(properties.properties.limits.maxComputeWorkGroupSize[1])},
+      {"max_block_size_z", Integer(properties.properties.limits.maxComputeWorkGroupSize[2])},
+      {"max_push_constants_size", Integer(properties.properties.limits.maxPushConstantsSize)},
+      {"max_uniform_buffer_range", Integer(properties.properties.limits.maxUniformBufferRange)},
+      {"max_storage_buffer_range",
+       Integer(IntImm(DataType::UInt(32), properties.properties.limits.maxStorageBufferRange))},
+      {"max_per_stage_descriptor_storage_buffer",
+       Integer(properties.properties.limits.maxPerStageDescriptorStorageBuffers)},
+      {"max_shared_memory_per_block",
+       Integer(properties.properties.limits.maxComputeSharedMemorySize)},
+      // Other device properties
+      {"device_name", String(properties.properties.deviceName)},
+      {"driver_version", Integer(properties.properties.driverVersion)},
+      {"vulkan_api_version", Integer(vulkan_api_version)},
+      {"max_spirv_version", Integer(max_spirv_version)},
+  };
+
+  return Target(config);
+}
+
 void VulkanDeviceAPI::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) {
   size_t index = static_cast<size_t>(dev.device_id);
   if (kind == kExist) {
@@ -470,6 +626,8 @@ void VulkanDeviceAPI::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv)
   VkPhysicalDeviceProperties phy_prop;
   vkGetPhysicalDeviceProperties(vctx.phy_device, &phy_prop);
 
+  // TODO: Replace with lookups to previously queried values.
+
   switch (kind) {
     case kMaxThreadsPerBlock: {
       int64_t value = phy_prop.limits.maxComputeWorkGroupInvocations;
@@ -568,8 +726,8 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
   }();
 
   const auto instance_extensions = [this]() {
-    std::vector<const char*> required_extensions{"VK_KHR_get_physical_device_properties2"};
-    std::vector<const char*> optional_extensions{};
+    std::vector<const char*> required_extensions{};
+    std::vector<const char*> optional_extensions{"VK_KHR_get_physical_device_properties2"};
 
     uint32_t inst_extension_prop_count;
     VULKAN_CALL(
@@ -660,6 +818,10 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     VkPhysicalDeviceFeatures enabled_features = {};
     enabled_features.shaderInt64 = VK_TRUE;
 
+    // TODO: Pass the VkPhysicalDeviceFeatures2 as the pNext chain if
+    // supported, otherwise pass VkPhysicalDeviceFeatures as
+    // pEnabledFeatures.
+
     VkDeviceCreateInfo device_create_info;
     device_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
     device_create_info.pNext = nullptr;
diff --git a/src/runtime/vulkan/vulkan_common.h b/src/runtime/vulkan/vulkan_common.h
index 8cdbaacbbb19..9d4ef6e009c1 100644
--- a/src/runtime/vulkan/vulkan_common.h
+++ b/src/runtime/vulkan/vulkan_common.h
@@ -128,7 +128,7 @@ struct VulkanGetBufferMemoryRequirements2Functions {
 };
 
 struct VulkanContext {
-  // phyiscal device
+  // physical device
   VkPhysicalDevice phy_device{nullptr};
   // Phyiscal device property
   VkPhysicalDeviceProperties phy_device_prop;
@@ -153,7 +153,7 @@ struct VulkanContext {
   // Queue family index.
   VkQueueFamilyProperties queue_prop;
 
-  bool UseImmediate() const { return descriptor_template_khr_functions.get() != nullptr; }
+  bool UseImmediate() const { return descriptor_template_khr_functions != nullptr; }
 };
 
 }  // namespace vulkan
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index e06b2c05d3bf..14ce81670c53 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -273,8 +273,35 @@ TVM_REGISTER_TARGET_KIND("metal", kDLMetal)
 
 TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
     .add_attr_option<Bool>("system-lib")
+    // Feature support
+    .add_attr_option<Bool>("supports_float16")
+    .add_attr_option<Bool>("supports_float32", Bool(true))
+    .add_attr_option<Bool>("supports_float64")
+    .add_attr_option<Bool>("supports_int8")
+    .add_attr_option<Bool>("supports_int16")
+    .add_attr_option<Bool>("supports_int32", Bool(true))
+    .add_attr_option<Bool>("supports_int64")
+    .add_attr_option<Bool>("supports_8bit_buffer")
+    .add_attr_option<Bool>("supports_16bit_buffer")
+    .add_attr_option<Bool>("supports_storage_buffer_storage_class")
+    .add_attr_option<Integer>("supported_subgroup_operations")
+    // Physical device limits
     .add_attr_option<Integer>("max_num_threads", Integer(256))
     .add_attr_option<Integer>("thread_warp_size", Integer(1))
+    .add_attr_option<Integer>("max_block_size_x")
+    .add_attr_option<Integer>("max_block_size_y")
+    .add_attr_option<Integer>("max_block_size_z")
+    .add_attr_option<Integer>("max_push_constants_size")
+    .add_attr_option<Integer>("max_uniform_buffer_range")
+    .add_attr_option<Integer>("max_storage_buffer_range")
+    .add_attr_option<Integer>("max_per_stage_descriptor_storage_buffer")
+    .add_attr_option<Integer>("max_shared_memory_per_block")
+    // Other device properties
+    .add_attr_option<String>("device_name")
+    .add_attr_option<Integer>("driver_version")
+    .add_attr_option<Integer>("vulkan_api_version")
+    .add_attr_option<Integer>("max_spirv_version")
+    // Tags
     .set_default_keys({"vulkan", "gpu"});
 
 TVM_REGISTER_TARGET_KIND("webgpu", kDLWebGPU)

From ef731c2e05863b7cebd4cecc794434902bdbd70f Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Mon, 24 May 2021 11:17:16 -0700
Subject: [PATCH 3/9] [Vulkan] Query instance-supported apiVersion before
 creating instance

- Previously, vkCreateInstance was called to initialize Vulkan 1.0.
---
 src/runtime/vulkan/vulkan.cc | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index a1114087528d..74c834ed745e 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -626,8 +626,6 @@ void VulkanDeviceAPI::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv)
   VkPhysicalDeviceProperties phy_prop;
   vkGetPhysicalDeviceProperties(vctx.phy_device, &phy_prop);
 
-  // TODO: Replace with lookups to previously queried values.
-
   switch (kind) {
     case kMaxThreadsPerBlock: {
       int64_t value = phy_prop.limits.maxComputeWorkGroupInvocations;
@@ -739,6 +737,20 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     return find_enabled_extensions(inst_extension_prop, required_extensions, optional_extensions);
   }();
 
+  const auto instance_api_version = []() {
+    uint32_t api_version = VK_MAKE_VERSION(1, 0, 0);
+
+    // Result from vkGetInstanceProcAddr is NULL if driver only
+    // supports vulkan 1.0.
+    auto vkEnumerateInstanceVersion =
+        (PFN_vkEnumerateInstanceVersion)vkGetInstanceProcAddr(NULL, "vkEnumerateInstanceVersion");
+    if (vkEnumerateInstanceVersion) {
+      vkEnumerateInstanceVersion(&api_version);
+    }
+
+    return api_version;
+  }();
+
   {
     VkApplicationInfo app_info;
     app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
@@ -747,7 +759,7 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     app_info.applicationVersion = 0;
     app_info.pEngineName = "";
     app_info.engineVersion = 0;
-    app_info.apiVersion = VK_MAKE_VERSION(1, 0, 0);
+    app_info.apiVersion = instance_api_version;
 
     VkInstanceCreateInfo inst_info;
     inst_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
@@ -818,10 +830,6 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     VkPhysicalDeviceFeatures enabled_features = {};
     enabled_features.shaderInt64 = VK_TRUE;
 
-    // TODO: Pass the VkPhysicalDeviceFeatures2 as the pNext chain if
-    // supported, otherwise pass VkPhysicalDeviceFeatures as
-    // pEnabledFeatures.
-
     VkDeviceCreateInfo device_create_info;
     device_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
     device_create_info.pNext = nullptr;

From fdf16c3a5a7315cbf2fdddd770c0140976d511a1 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Mon, 24 May 2021 12:00:20 -0700
Subject: [PATCH 4/9] [Vulkan] Moved options for dedicated allocation and push
 descriptors to environment variables

- Query support for dedicated allocation and push descriptors along
  with the rest of the device support.  Move the options to disable
  their use from compile-time variables to environment variables
  `TVM_VULKAN_DISABLE_PUSH_DESCRIPTOR` and
  `TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION`.
---
 cmake/modules/Vulkan.cmake         | 12 ---------
 src/runtime/vulkan/vulkan.cc       | 42 +++++++++++++++++++++---------
 src/runtime/vulkan/vulkan_common.h |  4 +++
 src/target/target_kind.cc          |  2 ++
 4 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/cmake/modules/Vulkan.cmake b/cmake/modules/Vulkan.cmake
index 095790f08547..6c5b4d72be2f 100644
--- a/cmake/modules/Vulkan.cmake
+++ b/cmake/modules/Vulkan.cmake
@@ -19,10 +19,6 @@
 find_vulkan(${USE_VULKAN})
 
 # Extra Vulkan runtime options, exposed for advanced users.
-tvm_option(USE_VULKAN_IMMEDIATE_MODE "Use Vulkan Immediate mode
-(KHR_push_descriptor extension)" ON IF USE_VULKAN)
-tvm_option(USE_VULKAN_DEDICATED_ALLOCATION "Use Vulkan dedicated allocations" ON
-IF USE_VULKAN)
 tvm_option(USE_VULKAN_VALIDATION "Enable Vulkan API validation layers" OFF
   IF USE_VULKAN)
 
@@ -39,14 +35,6 @@ if(USE_VULKAN)
   list(APPEND TVM_LINKER_LIBS ${Vulkan_SPIRV_TOOLS_LIBRARY})
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${Vulkan_LIBRARY})
 
-  if(USE_VULKAN_IMMEDIATE_MODE)
-    message(STATUS "Build with Vulkan immediate mode")
-    add_definitions(-DUSE_VULKAN_IMMEDIATE_MODE=1)
-  endif()
-  if(USE_VULKAN_DEDICATED_ALLOCATION)
-    message(STATUS "Build with Vulkan dedicated allocation")
-    add_definitions(-DUSE_VULKAN_DEDICATED_ALLOCATION=1)
-  endif()
   if(USE_VULKAN_VALIDATION)
     message(STATUS "Build with Vulkan API validation")
     add_definitions(-DUSE_VULKAN_VALIDATION=1)
diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index 74c834ed745e..975c761434a8 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -27,6 +27,7 @@
 
 #include <algorithm>
 #include <array>
+#include <cstdlib>
 #include <cstring>
 
 #include "../file_utils.h"
@@ -576,6 +577,28 @@ Target VulkanDeviceAPI::GetDeviceDescription(VkInstance instance, VkPhysicalDevi
     max_spirv_version = 0x10300;
   }
 
+  // Support is available based on these extensions, but allow it to
+  // be disabled based on an environment variable.
+  bool supports_push_descriptor =
+      has_extension("VK_KHR_push_descriptor") && has_extension("VK_KHR_descriptor_update_template");
+  {
+    const char* disable = std::getenv("TVM_VULKAN_DISABLE_PUSH_DESCRIPTOR");
+    if (disable && *disable) {
+      supports_push_descriptor = false;
+    }
+  }
+
+  // Support is available based on these extensions, but allow it to
+  // be disabled based on an environment variable.
+  bool supports_dedicated_allocation = has_extension("VK_KHR_get_memory_requirements2") &&
+                                       has_extension("VK_KHR_dedicated_allocation");
+  {
+    const char* disable = std::getenv("TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION");
+    if (disable && *disable) {
+      supports_dedicated_allocation = false;
+    }
+  }
+
   Map<String, ObjectRef> config = {
       {"kind", String("vulkan")},
       // Feature support
@@ -590,6 +613,8 @@ Target VulkanDeviceAPI::GetDeviceDescription(VkInstance instance, VkPhysicalDevi
       {"supports_16bit_buffer", Bool(storage_16bit.storageBuffer16BitAccess)},
       {"supports_storage_buffer_storage_class",
        Bool(has_extension("VK_KHR_storage_buffer_storage_class"))},
+      {"supports_push_descriptor", Bool(supports_push_descriptor)},
+      {"supports_dedicated_allocation", Bool(supports_dedicated_allocation)},
       {"supported_subgroup_operations", Integer(supported_subgroup_operations)},
       // Physical device limits
       {"max_num_threads", Integer(properties.properties.limits.maxComputeWorkGroupInvocations)},
@@ -825,6 +850,8 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
                                      optional_extensions);
     }();
 
+    ctx.target = GetDeviceDescription(instance_, phy_dev, instance_extensions, device_extensions);
+
     // All TVM-generated spirv shaders are marked as requiring int64
     // support, so we need to request it from the device, too.
     VkPhysicalDeviceFeatures enabled_features = {};
@@ -912,26 +939,17 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
       }
     }
     ICHECK_GE(win_rank, 0) << "Cannot find suitable local memory on device.";
-    auto has_extension = [&device_extensions](const char* query) {
-      return std::any_of(device_extensions.begin(), device_extensions.end(),
-                         [&](const char* extension) { return std::strcmp(query, extension) == 0; });
-    };
 
-#ifdef USE_VULKAN_IMMEDIATE_MODE
-    if (has_extension("VK_KHR_push_descriptor") &&
-        has_extension("VK_KHR_descriptor_update_template")) {
+    if (ctx.target->GetAttr<Bool>("supports_push_descriptor").value()) {
       ctx.descriptor_template_khr_functions =
           std::make_unique<VulkanDescriptorTemplateKHRFunctions>(ctx.device);
     }
-#endif
 
-#ifdef USE_VULKAN_DEDICATED_ALLOCATION
-    if (has_extension("VK_KHR_get_memory_requirements2") &&
-        has_extension("VK_KHR_dedicated_allocation")) {
+    if (ctx.target->GetAttr<Bool>("supports_dedicated_allocation").value()) {
       ctx.get_buffer_memory_requirements_2_functions =
           std::make_unique<VulkanGetBufferMemoryRequirements2Functions>(ctx.device);
     }
-#endif
+
     context_.push_back(std::move(ctx));
   }
 
diff --git a/src/runtime/vulkan/vulkan_common.h b/src/runtime/vulkan/vulkan_common.h
index 9d4ef6e009c1..14ecdba6ca40 100644
--- a/src/runtime/vulkan/vulkan_common.h
+++ b/src/runtime/vulkan/vulkan_common.h
@@ -24,6 +24,7 @@
 #include <tvm/runtime/device_api.h>
 #include <tvm/runtime/logging.h>
 #include <tvm/runtime/packed_func.h>
+#include <tvm/target/target.h>
 #include <vulkan/vulkan.h>
 
 #include <memory>
@@ -130,8 +131,11 @@ struct VulkanGetBufferMemoryRequirements2Functions {
 struct VulkanContext {
   // physical device
   VkPhysicalDevice phy_device{nullptr};
+
   // Phyiscal device property
   VkPhysicalDeviceProperties phy_device_prop;
+  // Target that best represents this physical device
+  Target target;
   // Memory type index for staging.
   uint32_t staging_mtype_index{0};
   // whether staging is coherent
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index 14ce81670c53..862efa88fbb5 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -284,6 +284,8 @@ TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
     .add_attr_option<Bool>("supports_8bit_buffer")
     .add_attr_option<Bool>("supports_16bit_buffer")
     .add_attr_option<Bool>("supports_storage_buffer_storage_class")
+    .add_attr_option<Bool>("supports_push_descriptor")
+    .add_attr_option<Bool>("supports_dedicated_allocation")
     .add_attr_option<Integer>("supported_subgroup_operations")
     // Physical device limits
     .add_attr_option<Integer>("max_num_threads", Integer(256))

From cada3fdb91d2445e0a56a146f64100f62f2e214f Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Mon, 24 May 2021 13:55:23 -0700
Subject: [PATCH 5/9] [Vulkan] Move option for vulkan validation layers to
 environment variable

- Moved to enable faster use as a debug tool.  If
  `TVM_VULKAN_ENABLE_VALIDATION_LAYERS` is a non-empty string,
  validation layers will be enabled.
---
 cmake/modules/Vulkan.cmake   |  9 ---------
 src/runtime/vulkan/vulkan.cc | 27 ++++++++++++++-------------
 2 files changed, 14 insertions(+), 22 deletions(-)

diff --git a/cmake/modules/Vulkan.cmake b/cmake/modules/Vulkan.cmake
index 6c5b4d72be2f..4dc9bd664d8a 100644
--- a/cmake/modules/Vulkan.cmake
+++ b/cmake/modules/Vulkan.cmake
@@ -18,10 +18,6 @@
 # Be compatible with older version of CMake
 find_vulkan(${USE_VULKAN})
 
-# Extra Vulkan runtime options, exposed for advanced users.
-tvm_option(USE_VULKAN_VALIDATION "Enable Vulkan API validation layers" OFF
-  IF USE_VULKAN)
-
 if(USE_VULKAN)
   if(NOT Vulkan_FOUND)
     message(FATAL_ERROR "Cannot find Vulkan, USE_VULKAN=" ${USE_VULKAN})
@@ -34,9 +30,4 @@ if(USE_VULKAN)
   list(APPEND COMPILER_SRCS ${COMPILER_VULKAN_SRCS})
   list(APPEND TVM_LINKER_LIBS ${Vulkan_SPIRV_TOOLS_LIBRARY})
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${Vulkan_LIBRARY})
-
-  if(USE_VULKAN_VALIDATION)
-    message(STATUS "Build with Vulkan API validation")
-    add_definitions(-DUSE_VULKAN_VALIDATION=1)
-  endif()
 endif(USE_VULKAN)
diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index 975c761434a8..32393e704cca 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -730,20 +730,21 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     std::vector<VkLayerProperties> inst_layer_prop(inst_layer_prop_count);
     VULKAN_CALL(vkEnumerateInstanceLayerProperties(&inst_layer_prop_count, inst_layer_prop.data()));
     std::vector<const char*> l;
-    for (const auto& lp : inst_layer_prop) {
-      // TODO(tulloch): add CMAKE options.
-      (void)lp;  // suppress unused variable warning.
-#ifdef USE_VULKAN_VALIDATION
-      if (std::strcmp(lp.layerName, "VK_LAYER_LUNARG_standard_validation") == 0) {
-        l.push_back("VK_LAYER_LUNARG_standard_validation");
-      }
-      if (std::strcmp(lp.layerName, "VK_LAYER_LUNARG_parameter_validation") == 0) {
-        l.push_back("VK_LAYER_LUNARG_parameter_validation");
-      }
-      if (std::strcmp(lp.layerName, "VK_LAYER_KHRONOS_validation") == 0) {
-        l.push_back("VK_LAYER_KHRONOS_validation");
+
+    const char* enable = std::getenv("TVM_VULKAN_ENABLE_VALIDATION_LAYERS");
+    bool validation_enabled = enable && *enable;
+    if (validation_enabled) {
+      for (const auto& lp : inst_layer_prop) {
+        if (std::strcmp(lp.layerName, "VK_LAYER_LUNARG_standard_validation") == 0) {
+          l.push_back("VK_LAYER_LUNARG_standard_validation");
+        }
+        if (std::strcmp(lp.layerName, "VK_LAYER_LUNARG_parameter_validation") == 0) {
+          l.push_back("VK_LAYER_LUNARG_parameter_validation");
+        }
+        if (std::strcmp(lp.layerName, "VK_LAYER_KHRONOS_validation") == 0) {
+          l.push_back("VK_LAYER_KHRONOS_validation");
+        }
       }
-#endif
     }
     return l;
   }();

From 457119bf4d55b038076763ab8adb31abb6219cb1 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Mon, 24 May 2021 14:16:41 -0700
Subject: [PATCH 6/9] [Vulkan] Explicitly enable vulkan features in device
 creation

- Vulkan requires that features be explicitly enabled before use.  For
  each feature that the device supports and a shader might use,
  declare it in the call to `vkCreateDevice`.
---
 src/runtime/vulkan/vulkan.cc | 91 +++++++++++++++++++++++++++++-------
 1 file changed, 74 insertions(+), 17 deletions(-)

diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index 32393e704cca..b8805f970def 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -763,6 +763,11 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     return find_enabled_extensions(inst_extension_prop, required_extensions, optional_extensions);
   }();
 
+  auto has_instance_extension = [&instance_extensions](const char* query) {
+    return std::any_of(instance_extensions.begin(), instance_extensions.end(),
+                       [&](const char* extension) { return std::strcmp(query, extension) == 0; });
+  };
+
   const auto instance_api_version = []() {
     uint32_t api_version = VK_MAKE_VERSION(1, 0, 0);
 
@@ -852,24 +857,76 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     }();
 
     ctx.target = GetDeviceDescription(instance_, phy_dev, instance_extensions, device_extensions);
+    std::cout << "Target: " << ctx.target << std::endl;
+
+    {
+      // Enable all features we may use that a device supports.
+      VkPhysicalDeviceFeatures2 enabled_features = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2};
+      VkPhysicalDevice8BitStorageFeatures storage_8bit = {
+          VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_8BIT_STORAGE_FEATURES};
+      VkPhysicalDevice16BitStorageFeatures storage_16bit = {
+          VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES};
+      VkPhysicalDeviceShaderFloat16Int8Features float16_int8 = {
+          VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES};
+
+      void** pp_next = &enabled_features.pNext;
+      bool needs_float16_int8 = false;
+
+      auto has_support = [&](const char* name) { return ctx.target->GetAttr<Bool>(name).value(); };
+      if (has_support("supports_float16")) {
+        float16_int8.shaderFloat16 = true;
+        needs_float16_int8 = true;
+      }
+      if (has_support("supports_float64")) {
+        enabled_features.features.shaderFloat64 = true;
+      }
+      if (has_support("supports_int8")) {
+        float16_int8.shaderInt8 = true;
+        needs_float16_int8 = true;
+      }
+      if (has_support("supports_int16")) {
+        enabled_features.features.shaderInt16 = true;
+      }
+      if (has_support("supports_int64")) {
+        enabled_features.features.shaderInt64 = true;
+      }
+      if (has_support("supports_8bit_buffer")) {
+        storage_8bit.storageBuffer8BitAccess = true;
+        *pp_next = &storage_8bit;
+        pp_next = &storage_8bit.pNext;
+      }
+      if (has_support("supports_16bit_buffer")) {
+        storage_16bit.storageBuffer16BitAccess = true;
+        *pp_next = &storage_16bit;
+        pp_next = &storage_16bit.pNext;
+      }
+
+      if (needs_float16_int8) {
+        *pp_next = &float16_int8;
+        pp_next = &float16_int8.pNext;
+      }
+
+      VkDeviceCreateInfo device_create_info;
+      device_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
+      device_create_info.pNext = nullptr;
+      device_create_info.flags = 0;
+      device_create_info.queueCreateInfoCount = 1;
+      device_create_info.pQueueCreateInfos = &queue_create_info;
+      device_create_info.enabledLayerCount = 0;
+      device_create_info.ppEnabledLayerNames = nullptr;
+      device_create_info.enabledExtensionCount = device_extensions.size();
+      device_create_info.ppEnabledExtensionNames = device_extensions.data();
+
+      if (has_instance_extension("VK_KHR_get_physical_device_properties2")) {
+        device_create_info.pEnabledFeatures = nullptr;
+        device_create_info.pNext = &enabled_features;
+      } else {
+        device_create_info.pNext = nullptr;
+        device_create_info.pEnabledFeatures = &enabled_features.features;
+      }
+      VULKAN_CALL(vkCreateDevice(phy_dev, &device_create_info, nullptr, &(ctx.device)));
+    }
 
-    // All TVM-generated spirv shaders are marked as requiring int64
-    // support, so we need to request it from the device, too.
-    VkPhysicalDeviceFeatures enabled_features = {};
-    enabled_features.shaderInt64 = VK_TRUE;
-
-    VkDeviceCreateInfo device_create_info;
-    device_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
-    device_create_info.pNext = nullptr;
-    device_create_info.flags = 0;
-    device_create_info.queueCreateInfoCount = 1;
-    device_create_info.pQueueCreateInfos = &queue_create_info;
-    device_create_info.enabledLayerCount = 0;
-    device_create_info.ppEnabledLayerNames = nullptr;
-    device_create_info.enabledExtensionCount = device_extensions.size();
-    device_create_info.ppEnabledExtensionNames = device_extensions.data();
-    device_create_info.pEnabledFeatures = &enabled_features;
-    VULKAN_CALL(vkCreateDevice(phy_dev, &device_create_info, nullptr, &(ctx.device)));
     ctx.queue_mutex.reset(new std::mutex());
     vkGetDeviceQueue(ctx.device, queue_family_index, 0, &(ctx.queue));
     ctx.queue_family_index = queue_family_index;

From 32741fe3890d74f3dff4789033b9208371067d7c Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Mon, 24 May 2021 14:34:51 -0700
Subject: [PATCH 7/9] [Vulkan] Avoid repeated queries for device attributes.

- Implement `VulkanDeviceAPI::GetAttr` based on the per-device values
  stored in the Target.  This pulls all logic for querying device
  parameters is in a single location.
---
 python/tvm/_ffi/runtime_ctypes.py |  2 +-
 src/runtime/vulkan/vulkan.cc      | 47 +++++++++++++------------------
 2 files changed, 20 insertions(+), 29 deletions(-)

diff --git a/python/tvm/_ffi/runtime_ctypes.py b/python/tvm/_ffi/runtime_ctypes.py
index 4eda5e8cc332..efea47752f6d 100644
--- a/python/tvm/_ffi/runtime_ctypes.py
+++ b/python/tvm/_ffi/runtime_ctypes.py
@@ -376,7 +376,7 @@ def api_version(self):
             The version of the SDK
 
         """
-        return self._GetDeviceAttr(self.device_type, self.device_id, 12)
+        return self._GetDeviceAttr(self.device_type, self.device_id, 11)
 
     @property
     def driver_version(self):
diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index b8805f970def..1c1809aebd78 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -647,39 +647,24 @@ void VulkanDeviceAPI::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv)
     return;
   }
   ICHECK_LT(index, context_.size()) << "Invalid device id " << index;
-  const auto& vctx = context(index);
-  VkPhysicalDeviceProperties phy_prop;
-  vkGetPhysicalDeviceProperties(vctx.phy_device, &phy_prop);
+
+  const auto& target = context(index).target;
 
   switch (kind) {
     case kMaxThreadsPerBlock: {
-      int64_t value = phy_prop.limits.maxComputeWorkGroupInvocations;
-      *rv = value;
+      *rv = target->GetAttr<Integer>("max_num_threads").value();
       break;
     }
     case kMaxSharedMemoryPerBlock: {
-      int64_t value = phy_prop.limits.maxComputeSharedMemorySize;
-      *rv = value;
+      *rv = target->GetAttr<Integer>("max_shared_memory_per_block");
       break;
     }
     case kWarpSize: {
-      VkPhysicalDeviceSubgroupProperties subgroup_prop;
-      subgroup_prop.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_PROPERTIES;
-      subgroup_prop.pNext = NULL;
-
-      VkPhysicalDeviceProperties2 phy_prop2;
-      phy_prop2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;
-      phy_prop2.pNext = &subgroup_prop;
-
-      vkGetPhysicalDeviceProperties2(vctx.phy_device, &phy_prop2);
-      int64_t subgroup_size = subgroup_prop.subgroupSize;
-      ICHECK(subgroup_size >= 1);
-
-      *rv = subgroup_size;
+      *rv = target->GetAttr<Integer>("thread_warp_size").value();
       break;
     }
     case kComputeVersion: {
-      int64_t value = phy_prop.apiVersion;
+      int64_t value = target->GetAttr<Integer>("vulkan_api_version").value();
       std::ostringstream os;
       os << VK_VERSION_MAJOR(value) << "." << VK_VERSION_MINOR(value) << "."
          << VK_VERSION_PATCH(value);
@@ -687,33 +672,39 @@ void VulkanDeviceAPI::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv)
       break;
     }
     case kDeviceName:
-      *rv = std::string(phy_prop.deviceName);
+      *rv = target->GetAttr<String>("device_name").value();
       break;
+
     case kMaxClockRate:
       break;
+
     case kMultiProcessorCount:
       break;
+
     case kExist:
       break;
+
     case kMaxThreadDimensions: {
-      int64_t dims[3];
-      dims[0] = phy_prop.limits.maxComputeWorkGroupSize[0];
-      dims[1] = phy_prop.limits.maxComputeWorkGroupSize[1];
-      dims[2] = phy_prop.limits.maxComputeWorkGroupSize[2];
       std::stringstream ss;  // use json string to return multiple int values;
-      ss << "[" << dims[0] << ", " << dims[1] << ", " << dims[2] << "]";
+      ss << "[" << target->GetAttr<Integer>("max_block_size_x").value() << ", "
+         << target->GetAttr<Integer>("max_block_size_y").value() << ", "
+         << target->GetAttr<Integer>("max_block_size_z").value() << "]";
       *rv = ss.str();
       break;
     }
+
     case kMaxRegistersPerBlock:
       break;
+
     case kGcnArch:
       break;
+
     case kApiVersion:
       *rv = VK_HEADER_VERSION;
       break;
+
     case kDriverVersion: {
-      int64_t value = phy_prop.driverVersion;
+      int64_t value = target->GetAttr<Integer>("driver_version").value();
       std::ostringstream os;
       os << VK_VERSION_MAJOR(value) << "." << VK_VERSION_MINOR(value) << "."
          << VK_VERSION_PATCH(value);

From d95e77e757168334b5c8f50bf1a0f9b0a83fd9ec Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Tue, 25 May 2021 08:53:28 -0700
Subject: [PATCH 8/9] [Vulkan] Implement "from_device" flag for the vulkan
 target.

- With the number of device capabilities that may or may not be
  supported by a vulkan driver, it can be tedious to input them.
  Specifying "-from_device=0" now indicate that any unspecified values
  should be read from the device.
---
 src/runtime/vulkan/vulkan.cc |  7 ++++-
 src/target/target_kind.cc    | 51 ++++++++++++++++++++++++++++++++----
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/src/runtime/vulkan/vulkan.cc b/src/runtime/vulkan/vulkan.cc
index 1c1809aebd78..8982ea32648b 100644
--- a/src/runtime/vulkan/vulkan.cc
+++ b/src/runtime/vulkan/vulkan.cc
@@ -432,6 +432,8 @@ class VulkanDeviceAPI final : public DeviceAPI {
     return context_[device_id];
   }
 
+  Target GenerateTarget(size_t device_id) const { return context(device_id).target; }
+
  private:
   std::vector<const char*> find_enabled_extensions(
       const std::vector<VkExtensionProperties>& ext_prop,
@@ -848,7 +850,6 @@ VulkanDeviceAPI::VulkanDeviceAPI() {
     }();
 
     ctx.target = GetDeviceDescription(instance_, phy_dev, instance_extensions, device_extensions);
-    std::cout << "Target: " << ctx.target << std::endl;
 
     {
       // Enable all features we may use that a device supports.
@@ -1580,6 +1581,10 @@ TVM_REGISTER_GLOBAL("device_api.vulkan").set_body([](TVMArgs args, TVMRetValue*
   *rv = static_cast<void*>(ptr);
 });
 
+TVM_REGISTER_GLOBAL("device_api.vulkan.generate_target").set_body_typed([](int device_id) {
+  return VulkanDeviceAPI::Global()->GenerateTarget(device_id);
+});
+
 }  // namespace vulkan
 }  // namespace runtime
 }  // namespace tvm
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index 862efa88fbb5..08e998e0f035 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -209,6 +209,45 @@ Map<String, ObjectRef> UpdateROCmAttrs(Map<String, ObjectRef> attrs) {
   return attrs;
 }
 
+/*!
+ * \brief Update the attributes in the Vulkan target.
+ * \param attrs The original attributes
+ * \return The updated attributes
+ */
+Map<String, ObjectRef> UpdateVulkanAttrs(Map<String, ObjectRef> attrs) {
+  if (attrs.count("from_device")) {
+    int device_id = Downcast<Integer>(attrs.at("from_device"));
+    const PackedFunc* generate_target = runtime::Registry::Get("device_api.vulkan.generate_target");
+    ICHECK(generate_target)
+        << "Requested to read Vulkan parameters from device, but no Vulkan runtime available";
+    Target target = (*generate_target)(device_id).AsObjectRef<Target>();
+    for (auto& kv : target->Export()) {
+      if (!attrs.count(kv.first)) {
+        attrs.Set(kv.first, kv.second);
+      }
+    }
+
+    attrs.erase("from_device");
+  }
+
+  // Set defaults here, rather than in the .add_attr_option() calls.
+  // The priority should be user-specified > device-query > default,
+  // but defaults defined in .add_attr_option() are already applied by
+  // this point.  Longer-term, would be good to add a
+  // `DeviceAPI::GenerateTarget` function and extend "from_device" to
+  // work for all runtimes.
+  std::unordered_map<String, ObjectRef> defaults = {{"supports_float32", Bool(true)},
+                                                    {"supports_int32", Bool(true)},
+                                                    {"max_num_threads", Integer(256)},
+                                                    {"thread_warp_size", Integer(1)}};
+  for (const auto& kv : defaults) {
+    if (!attrs.count(kv.first)) {
+      attrs.Set(kv.first, kv.second);
+    }
+  }
+  return attrs;
+}
+
 /**********  Register Target kinds and attributes  **********/
 
 TVM_REGISTER_TARGET_KIND("llvm", kDLCPU)
@@ -273,13 +312,14 @@ TVM_REGISTER_TARGET_KIND("metal", kDLMetal)
 
 TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
     .add_attr_option<Bool>("system-lib")
+    .add_attr_option<Bool>("from_device")
     // Feature support
     .add_attr_option<Bool>("supports_float16")
-    .add_attr_option<Bool>("supports_float32", Bool(true))
+    .add_attr_option<Bool>("supports_float32")
     .add_attr_option<Bool>("supports_float64")
     .add_attr_option<Bool>("supports_int8")
     .add_attr_option<Bool>("supports_int16")
-    .add_attr_option<Bool>("supports_int32", Bool(true))
+    .add_attr_option<Bool>("supports_int32")
     .add_attr_option<Bool>("supports_int64")
     .add_attr_option<Bool>("supports_8bit_buffer")
     .add_attr_option<Bool>("supports_16bit_buffer")
@@ -288,8 +328,8 @@ TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
     .add_attr_option<Bool>("supports_dedicated_allocation")
     .add_attr_option<Integer>("supported_subgroup_operations")
     // Physical device limits
-    .add_attr_option<Integer>("max_num_threads", Integer(256))
-    .add_attr_option<Integer>("thread_warp_size", Integer(1))
+    .add_attr_option<Integer>("max_num_threads")
+    .add_attr_option<Integer>("thread_warp_size")
     .add_attr_option<Integer>("max_block_size_x")
     .add_attr_option<Integer>("max_block_size_y")
     .add_attr_option<Integer>("max_block_size_z")
@@ -304,7 +344,8 @@ TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
     .add_attr_option<Integer>("vulkan_api_version")
     .add_attr_option<Integer>("max_spirv_version")
     // Tags
-    .set_default_keys({"vulkan", "gpu"});
+    .set_default_keys({"vulkan", "gpu"})
+    .set_attrs_preprocessor(UpdateVulkanAttrs);
 
 TVM_REGISTER_TARGET_KIND("webgpu", kDLWebGPU)
     .add_attr_option<Bool>("system-lib")

From 1d262d7545c76f05dd6a29337919cb28d83f26dc Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Tue, 25 May 2021 09:58:10 -0700
Subject: [PATCH 9/9] [Vulkan][Codegen] Read vulkan device capabilities/limits
 from Target

- Previously, the codegen assumed that all device features were
  present.  Now, the codegen reads device capabilities from the
  Target, and throws an error if codegen would require use of an
  unsupported feature.
---
 src/target/spirv/spirv_support.cc | 50 ++++++++++++++++++++++++-------
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/src/target/spirv/spirv_support.cc b/src/target/spirv/spirv_support.cc
index ff9aee406574..e06bde08895d 100644
--- a/src/target/spirv/spirv_support.cc
+++ b/src/target/spirv/spirv_support.cc
@@ -35,17 +35,45 @@ SPIRVSupport::SPIRVSupport(tvm::Target target) {
   ICHECK_EQ(target->kind->device_type, kDLVulkan)
       << "SPIRVSupport can only be checked for vulkan device type";
 
-  // Currently, this codifies the assumptions that were present and
-  // implicit in previous implementations.  In the future, this will
-  // pull information from the specified `Target`.
-
-  supports_storage_buffer_storage_class = (SPV_VERSION >= 0x10300);
-  supports_storage_buffer_8bit_access = true;
-  supports_storage_buffer_16bit_access = true;
-  supports_float16 = true;
-  supports_int8 = true;
-  supports_int16 = true;
-  supports_int64 = true;
+  if (target->GetAttr<Integer>("supported_subgroup_operations")) {
+    supported_subgroup_operations =
+        target->GetAttr<Integer>("supported_subgroup_operations").value();
+  }
+  if (target->GetAttr<Integer>("max_push_constants_size")) {
+    max_push_constants_size = target->GetAttr<Integer>("max_push_constants_size").value();
+  }
+  if (target->GetAttr<Integer>("max_uniform_buffer_range")) {
+    max_uniform_buffer_range = target->GetAttr<Integer>("max_uniform_buffer_range").value();
+  }
+  if (target->GetAttr<Integer>("max_storage_buffer_range")) {
+    max_storage_buffer_range = target->GetAttr<Integer>("max_storage_buffer_range").value();
+  }
+  if (target->GetAttr<Integer>("max_per_stage_descriptor_storage_buffer")) {
+    max_per_stage_descriptor_storage_buffers =
+        target->GetAttr<Integer>("max_per_stage_descriptor_storage_buffer").value();
+  }
+  if (target->GetAttr<Bool>("supports_storage_buffer_storage_class")) {
+    supports_storage_buffer_storage_class =
+        target->GetAttr<Bool>("supports_storage_buffer_storage_class").value();
+  }
+  if (target->GetAttr<Bool>("supports_8bit_buffer")) {
+    supports_storage_buffer_8bit_access = target->GetAttr<Bool>("supports_8bit_buffer").value();
+  }
+  if (target->GetAttr<Bool>("supports_16bit_buffer")) {
+    supports_storage_buffer_16bit_access = target->GetAttr<Bool>("supports_16bit_buffer").value();
+  }
+  if (target->GetAttr<Bool>("supports_float16")) {
+    supports_float16 = target->GetAttr<Bool>("supports_float16").value();
+  }
+  if (target->GetAttr<Bool>("supports_int8")) {
+    supports_int8 = target->GetAttr<Bool>("supports_int8").value();
+  }
+  if (target->GetAttr<Bool>("supports_int16")) {
+    supports_int16 = target->GetAttr<Bool>("supports_int16").value();
+  }
+  if (target->GetAttr<Bool>("supports_int64")) {
+    supports_int64 = target->GetAttr<Bool>("supports_int64").value();
+  }
 }
 
 }  // namespace codegen