Enabling Vulkan buffer device addresses support in the runtime. (#14778)

We'll need command buffer and compiler support to use it but this will
let us see if there are any issues enabling the features on our targets.
There's support for mixed SPIR-V programs using buffer device addresses
and those not and VMFBs can contain both to allow for graceful fallback
when not supported. SPIR-V executables with the `vulkan-spirv-fb-ptr`
will get the support and future changes will add descriptor set layout
flags (`IREE_HAL_DESCRIPTOR_SET_LAYOUT_FLAG_INDIRECT` or something) that
allows for routing certain descriptors into the indirect parameter
buffer.

Progress on #13945.

runtime/src/iree/hal/drivers/vulkan/api.h

@@ -90,6 +90,10 @@ enum iree_hal_vulkan_feature_bits_t {
   // to allow for queue-ordered virtual memory management.
   IREE_HAL_VULKAN_FEATURE_ENABLE_SPARSE_RESIDENCY_ALIASED =
       IREE_HAL_VULKAN_FEATURE_ENABLE_SPARSE_BINDING | (1u << 5),
+
+  // Enables buffer device addresses when supported and uses them when
+  // appropriately compiled SPIR-V executables require them.
+  IREE_HAL_VULKAN_FEATURE_ENABLE_BUFFER_DEVICE_ADDRESSES = 1u << 6,
 };
 typedef uint32_t iree_hal_vulkan_features_t;
 

runtime/src/iree/hal/drivers/vulkan/cts/CMakeLists.txt

@@ -14,7 +14,7 @@ iree_hal_cts_test_suite(
   COMPILER_TARGET_BACKEND
     "vulkan-spirv"
   EXECUTABLE_FORMAT
-    "\"SPVE\""
+    "\"vulkan-spirv-fb\""
   DEPS
     iree::hal::drivers::vulkan::registration
   EXCLUDED_TESTS

runtime/src/iree/hal/drivers/vulkan/dynamic_symbol_tables.h

@@ -293,6 +293,9 @@ namespace vulkan {
   DEV_PFN(OPTIONAL, vkSignalSemaphore)                                  \
   DEV_PFN(OPTIONAL, vkSignalSemaphoreKHR)                               \
                                                                         \
+  DEV_PFN(OPTIONAL, vkGetBufferDeviceAddress)                           \
+  DEV_PFN(OPTIONAL, vkGetBufferDeviceAddressKHR)                        \
+                                                                        \
   INS_PFN(EXCLUDED, vkCreateDebugReportCallbackEXT)                     \
   INS_PFN(OPTIONAL, vkCreateDebugUtilsMessengerEXT)                     \
   INS_PFN(EXCLUDED, vkCreateDisplayPlaneSurfaceKHR)                     \
@@ -329,7 +332,7 @@ namespace vulkan {
   INS_PFN(EXCLUDED, vkGetPhysicalDeviceExternalSemaphoreProperties)     \
   INS_PFN(EXCLUDED, vkGetPhysicalDeviceExternalSemaphorePropertiesKHR)  \
   INS_PFN(REQUIRED, vkGetPhysicalDeviceFeatures)                        \
-  INS_PFN(EXCLUDED, vkGetPhysicalDeviceFeatures2)                       \
+  INS_PFN(REQUIRED, vkGetPhysicalDeviceFeatures2)                       \
   INS_PFN(EXCLUDED, vkGetPhysicalDeviceFeatures2KHR)                    \
   INS_PFN(REQUIRED, vkGetPhysicalDeviceFormatProperties)                \
   INS_PFN(EXCLUDED, vkGetPhysicalDeviceFormatProperties2)               \

runtime/src/iree/hal/drivers/vulkan/extensibility_util.cc

@@ -213,6 +213,9 @@ iree_hal_vulkan_populate_enabled_device_extensions(
     } else if (strcmp(extension_name,
                       VK_EXT_EXTERNAL_MEMORY_HOST_EXTENSION_NAME) == 0) {
       extensions.external_memory_host = true;
+    } else if (strcmp(extension_name,
+                      VK_KHR_BUFFER_DEVICE_ADDRESS_EXTENSION_NAME) == 0) {
+      extensions.buffer_device_address = true;
     }
   }
   return extensions;
@@ -235,5 +238,12 @@ iree_hal_vulkan_infer_enabled_device_extensions(
   if (device_syms->vkGetCalibratedTimestampsEXT) {
     extensions.calibrated_timestamps = true;
   }
+  if (device_syms->vkGetMemoryHostPointerPropertiesEXT) {
+    extensions.external_memory_host = true;
+  }
+  if (device_syms->vkGetBufferDeviceAddress ||
+      device_syms->vkGetBufferDeviceAddressKHR) {
+    extensions.buffer_device_address = true;
+  }
   return extensions;
 }

runtime/src/iree/hal/drivers/vulkan/extensibility_util.h

@@ -84,6 +84,8 @@ typedef struct iree_hal_vulkan_device_extensions_t {
   bool subgroup_size_control : 1;
   // VK_EXT_external_memory_host is enabled.
   bool external_memory_host : 1;
+  // VK_KHR_buffer_device_address is enabled.
+  bool buffer_device_address : 1;
 } iree_hal_vulkan_device_extensions_t;
 
 // Returns a bitfield with all of the provided extension names.

runtime/src/iree/hal/drivers/vulkan/native_allocator.cc

@@ -293,6 +293,17 @@ static iree_status_t iree_hal_vulkan_native_allocator_commit_and_wrap(
   allocate_info.pNext = NULL;
   allocate_info.allocationSize = requirements.size;
   allocate_info.memoryTypeIndex = memory_type_index;
+  VkMemoryAllocateFlagsInfo allocate_flags_info = {};
+  allocate_flags_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO;
+  allocate_flags_info.pNext = NULL;
+  allocate_flags_info.flags = 0;
+  if (iree_all_bits_set(
+          logical_device->enabled_features(),
+          IREE_HAL_VULKAN_FEATURE_ENABLE_BUFFER_DEVICE_ADDRESSES)) {
+    allocate_flags_info.flags |= VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT;
+  }
+  allocate_flags_info.deviceMask = 0;
+  allocate_info.pNext = &allocate_flags_info;
   VkDeviceMemory device_memory = VK_NULL_HANDLE;
   VK_RETURN_IF_ERROR(logical_device->syms()->vkAllocateMemory(
                          *logical_device, &allocate_info,

runtime/src/iree/hal/drivers/vulkan/nop_executable_cache.cc

@@ -73,14 +73,34 @@ static bool iree_hal_vulkan_nop_executable_cache_can_prepare_format(
     iree_hal_executable_cache_t* base_executable_cache,
     iree_hal_executable_caching_mode_t caching_mode,
     iree_string_view_t executable_format) {
-  return iree_string_view_equal(executable_format,
-                                iree_make_cstring_view("SPVE"));
+  iree_hal_vulkan_nop_executable_cache_t* executable_cache =
+      iree_hal_vulkan_nop_executable_cache_cast(base_executable_cache);
+  if (iree_string_view_equal(executable_format,
+                             iree_make_cstring_view("vulkan-spirv-fb"))) {
+    return true;
+  } else if (iree_string_view_equal(
+                 executable_format,
+                 iree_make_cstring_view("vulkan-spirv-fb-ptr"))) {
+    return iree_all_bits_set(
+        executable_cache->logical_device->enabled_features(),
+        IREE_HAL_VULKAN_FEATURE_ENABLE_BUFFER_DEVICE_ADDRESSES);
+  }
+  return false;
 }
 
 static iree_status_t iree_hal_vulkan_nop_executable_cache_prepare_executable(
     iree_hal_executable_cache_t* base_executable_cache,
     const iree_hal_executable_params_t* executable_params,
     iree_hal_executable_t** out_executable) {
+  if (!iree_hal_vulkan_nop_executable_cache_can_prepare_format(
+          base_executable_cache, executable_params->caching_mode,
+          executable_params->executable_format)) {
+    return iree_make_status(IREE_STATUS_NOT_FOUND,
+                            "no Vulkan executable implementation registered "
+                            "for the given executable format '%.*s'",
+                            (int)executable_params->executable_format.size,
+                            executable_params->executable_format.data);
+  }
   iree_hal_vulkan_nop_executable_cache_t* executable_cache =
       iree_hal_vulkan_nop_executable_cache_cast(base_executable_cache);
   return iree_hal_vulkan_native_executable_create(

runtime/src/iree/hal/drivers/vulkan/registration/driver_module.cc

@@ -38,6 +38,9 @@ IREE_FLAG(
 IREE_FLAG(bool, vulkan_sparse_residency, true,
           "Enables the Vulkan 'sparseResidencyBuffer' feature (and others) "
           "when available.");
+IREE_FLAG(bool, vulkan_buffer_device_addresses, true,
+          "Enables the Vulkan 'bufferDeviceAddress` feature and support for "
+          "SPIR-V executables compiled to use it.");
 
 IREE_FLAG(
     bool, vulkan_dedicated_compute_queue, false,
@@ -94,6 +97,10 @@ static iree_status_t iree_hal_vulkan_create_driver_with_flags(
     driver_options.requested_features |=
         IREE_HAL_VULKAN_FEATURE_ENABLE_SPARSE_RESIDENCY_ALIASED;
   }
+  if (FLAG_vulkan_buffer_device_addresses) {
+    driver_options.requested_features |=
+        IREE_HAL_VULKAN_FEATURE_ENABLE_BUFFER_DEVICE_ADDRESSES;
+  }
 
   if (FLAG_vulkan_dedicated_compute_queue) {
     driver_options.device_options.flags |=

runtime/src/iree/hal/drivers/vulkan/vulkan_device.cc

@@ -215,6 +215,11 @@ IREE_API_EXPORT iree_status_t iree_hal_vulkan_query_extensibility_set(
   ADD_EXT(IREE_HAL_VULKAN_EXTENSIBILITY_DEVICE_EXTENSIONS_OPTIONAL,
           VK_EXT_EXTERNAL_MEMORY_HOST_EXTENSION_NAME);
 
+  // VK_KHR_buffer_device_address:
+  // Promoted to core in Vulkan 1.2 but still an extension in 1.1.
+  ADD_EXT(IREE_HAL_VULKAN_EXTENSIBILITY_DEVICE_EXTENSIONS_OPTIONAL,
+          VK_KHR_BUFFER_DEVICE_ADDRESS_EXTENSION_NAME);
+
   //===--------------------------------------------------------------------===//
   // Vulkan forward-compatibility shims
   //===--------------------------------------------------------------------===//
@@ -911,9 +916,20 @@ iree_status_t iree_hal_vulkan_device_create(
   dispatch_queue_info.pQueuePriorities = dispatch_queue_priorities.data();
 
   // Collect supported physical device features.
-  VkPhysicalDeviceFeatures physical_device_features;
-  instance_syms->vkGetPhysicalDeviceFeatures(physical_device,
-                                             &physical_device_features);
+  VkPhysicalDeviceFeatures2 available_features2;
+  memset(&available_features2, 0, sizeof(available_features2));
+  available_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
+  VkPhysicalDeviceBufferDeviceAddressFeatures
+      available_buffer_device_address_features;
+  memset(&available_buffer_device_address_features, 0,
+         sizeof(available_buffer_device_address_features));
+  available_buffer_device_address_features.sType =
+      VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES;
+  available_features2.pNext = &available_buffer_device_address_features;
+  instance_syms->vkGetPhysicalDeviceFeatures2(physical_device,
+                                              &available_features2);
+  const VkPhysicalDeviceFeatures* available_features =
+      &available_features2.features;
 
   // Create device and its queues.
   VkDeviceCreateInfo device_create_info;
@@ -927,12 +943,12 @@ iree_status_t iree_hal_vulkan_device_create(
   device_create_info.pQueueCreateInfos = queue_create_info.data();
   device_create_info.pEnabledFeatures = NULL;
 
-  VkPhysicalDeviceFeatures2 features2;
-  memset(&features2, 0, sizeof(features2));
-  features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
-  device_create_info.pNext = &features2;
-  if (physical_device_features.shaderInt64) {
-    features2.features.shaderInt64 = VK_TRUE;
+  VkPhysicalDeviceFeatures2 enabled_features2;
+  memset(&enabled_features2, 0, sizeof(enabled_features2));
+  enabled_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
+  device_create_info.pNext = &enabled_features2;
+  if (available_features->shaderInt64) {
+    enabled_features2.features.shaderInt64 = VK_TRUE;
   }
 
   iree_hal_vulkan_features_t enabled_features = 0;
@@ -946,46 +962,61 @@ iree_status_t iree_hal_vulkan_device_create(
 
   if (iree_all_bits_set(requested_features,
                         IREE_HAL_VULKAN_FEATURE_ENABLE_SPARSE_BINDING) &&
-      physical_device_features.sparseBinding) {
-    features2.features.sparseBinding = VK_TRUE;
+      available_features->sparseBinding) {
+    enabled_features2.features.sparseBinding = VK_TRUE;
     enabled_features |= IREE_HAL_VULKAN_FEATURE_ENABLE_SPARSE_BINDING;
   }
   if (iree_all_bits_set(
           requested_features,
           IREE_HAL_VULKAN_FEATURE_ENABLE_SPARSE_RESIDENCY_ALIASED) &&
-      physical_device_features.sparseResidencyBuffer &&
-      physical_device_features.sparseResidencyAliased) {
-    features2.features.sparseResidencyBuffer = VK_TRUE;
-    features2.features.sparseResidencyAliased = VK_TRUE;
+      available_features->sparseResidencyBuffer &&
+      available_features->sparseResidencyAliased) {
+    enabled_features2.features.sparseResidencyBuffer = VK_TRUE;
+    enabled_features2.features.sparseResidencyAliased = VK_TRUE;
     enabled_features |= IREE_HAL_VULKAN_FEATURE_ENABLE_SPARSE_RESIDENCY_ALIASED;
   }
 
   if (iree_all_bits_set(requested_features,
                         IREE_HAL_VULKAN_FEATURE_ENABLE_ROBUST_BUFFER_ACCESS)) {
-    if (physical_device_features.robustBufferAccess != VK_TRUE) {
+    if (available_features->robustBufferAccess != VK_TRUE) {
       return iree_make_status(
           IREE_STATUS_UNAVAILABLE,
-          "Robust buffer access not supported by physical device");
+          "robust buffer access not supported by physical device");
     }
-    features2.features.robustBufferAccess = VK_TRUE;
+    enabled_features2.features.robustBufferAccess = VK_TRUE;
     enabled_features |= IREE_HAL_VULKAN_FEATURE_ENABLE_ROBUST_BUFFER_ACCESS;
   }
 
+  VkPhysicalDeviceBufferDeviceAddressFeatures buffer_device_address_features;
+  if (iree_all_bits_set(
+          requested_features,
+          IREE_HAL_VULKAN_FEATURE_ENABLE_BUFFER_DEVICE_ADDRESSES) &&
+      available_buffer_device_address_features.bufferDeviceAddress) {
+    memset(&buffer_device_address_features, 0,
+           sizeof(buffer_device_address_features));
+    buffer_device_address_features.sType =
+        VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES;
+    buffer_device_address_features.pNext = enabled_features2.pNext;
+    enabled_features2.pNext = &buffer_device_address_features;
+    buffer_device_address_features.bufferDeviceAddress = true;
+    enabled_features |= IREE_HAL_VULKAN_FEATURE_ENABLE_BUFFER_DEVICE_ADDRESSES;
+  }
+
   VkPhysicalDeviceTimelineSemaphoreFeatures semaphore_features;
   memset(&semaphore_features, 0, sizeof(semaphore_features));
   semaphore_features.sType =
       VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_TIMELINE_SEMAPHORE_FEATURES;
-  semaphore_features.pNext = features2.pNext;
-  features2.pNext = &semaphore_features;
+  semaphore_features.pNext = enabled_features2.pNext;
+  enabled_features2.pNext = &semaphore_features;
   semaphore_features.timelineSemaphore = VK_TRUE;
 
   VkPhysicalDeviceHostQueryResetFeaturesEXT host_query_reset_features;
   if (enabled_device_extensions.host_query_reset) {
     memset(&host_query_reset_features, 0, sizeof(host_query_reset_features));
     host_query_reset_features.sType =
         VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_HOST_QUERY_RESET_FEATURES_EXT;
-    host_query_reset_features.pNext = features2.pNext;
-    features2.pNext = &host_query_reset_features;
+    host_query_reset_features.pNext = enabled_features2.pNext;
+    enabled_features2.pNext = &host_query_reset_features;
     host_query_reset_features.hostQueryReset = VK_TRUE;
   }
 
@@ -994,8 +1025,8 @@ iree_status_t iree_hal_vulkan_device_create(
     memset(&subgroup_control_features, 0, sizeof(subgroup_control_features));
     subgroup_control_features.sType =
         VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES;
-    subgroup_control_features.pNext = features2.pNext;
-    features2.pNext = &subgroup_control_features;
+    subgroup_control_features.pNext = enabled_features2.pNext;
+    enabled_features2.pNext = &subgroup_control_features;
     subgroup_control_features.subgroupSizeControl = VK_TRUE;
   }
 
@@ -1134,16 +1165,25 @@ static iree_status_t iree_hal_vulkan_device_trim(
 static iree_status_t iree_hal_vulkan_device_query_i64(
     iree_hal_device_t* base_device, iree_string_view_t category,
     iree_string_view_t key, int64_t* out_value) {
-  // iree_hal_vulkan_device_t* device =
-  //     iree_hal_vulkan_device_cast(base_device);
+  iree_hal_vulkan_device_t* device = iree_hal_vulkan_device_cast(base_device);
   *out_value = 0;
 
   if (iree_string_view_equal(category,
                              iree_make_cstring_view("hal.executable.format"))) {
-    *out_value =
-        iree_string_view_equal(key, iree_make_cstring_view("vulkan-spirv-fb"))
-            ? 1
-            : 0;
+    if (iree_string_view_equal(key,
+                               iree_make_cstring_view("vulkan-spirv-fb"))) {
+      // Base SPIR-V always supported.
+      *out_value = 1;
+    } else if (iree_string_view_equal(
+                   key, iree_make_cstring_view("vulkan-spirv-fb-ptr"))) {
+      // SPIR-V with device addresses is optionally supported based on whether
+      // we have device feature support.
+      *out_value = iree_all_bits_set(
+                       device->logical_device->enabled_features(),
+                       IREE_HAL_VULKAN_FEATURE_ENABLE_BUFFER_DEVICE_ADDRESSES)
+                       ? 1
+                       : 0;
+    }
     return iree_ok_status();
   }