Implementing Vulkan dispatch tracing.

iree/hal/vulkan/BUILD

@@ -72,6 +72,8 @@ cc_library(
         "status_util.h",
         "timepoint_util.cc",
         "timepoint_util.h",
+        "tracing.cc",
+        "tracing.h",
         "vma_allocator.cc",
         "vma_allocator.h",
         "vma_buffer.cc",

iree/hal/vulkan/CMakeLists.txt

@@ -61,6 +61,8 @@ iree_cc_library(
     "status_util.h"
     "timepoint_util.cc"
     "timepoint_util.h"
+    "tracing.cc"
+    "tracing.h"
     "vma_allocator.cc"
     "vma_allocator.h"
     "vma_buffer.cc"

iree/hal/vulkan/api.h

@@ -37,10 +37,20 @@ extern "C" {
 enum iree_hal_vulkan_feature_e {
   // Use VK_LAYER_KHRONOS_standard_validation to validate Vulkan API usage.
   // Has a significant performance penalty and is *not* a security mechanism.
-  IREE_HAL_VULKAN_FEATURE_ENABLE_VALIDATION_LAYERS = 1 << 0,
+  IREE_HAL_VULKAN_FEATURE_ENABLE_VALIDATION_LAYERS = 1u << 0,
 
   // Use VK_EXT_debug_utils, record markers, and log errors.
-  IREE_HAL_VULKAN_FEATURE_ENABLE_DEBUG_UTILS = 1 << 1,
+  IREE_HAL_VULKAN_FEATURE_ENABLE_DEBUG_UTILS = 1u << 1,
+
+  // Enables tracing of command buffers when IREE tracing is enabled.
+  // May take advantage of additional extensions for more accurate timing or
+  // hardware-specific performance counters.
+  //
+  // NOTE: tracing has a non-trivial overhead and will skew the timing of
+  // submissions and introduce false barriers between dispatches. Use this to
+  // identify slow dispatches and refine from there; be wary of whole-program
+  // tracing with this enabled.
+  IREE_HAL_VULKAN_FEATURE_ENABLE_TRACING = 1u << 2,
 };
 typedef uint64_t iree_hal_vulkan_features_t;
 

iree/hal/vulkan/command_queue.h

@@ -22,6 +22,7 @@
 #include "iree/hal/api.h"
 #include "iree/hal/vulkan/dynamic_symbols.h"
 #include "iree/hal/vulkan/handle_util.h"
+#include "iree/hal/vulkan/tracing.h"
 #include "iree/hal/vulkan/util/arena.h"
 
 namespace iree {
@@ -42,6 +43,15 @@ class CommandQueue {
     return logical_device_->syms();
   }
 
+  VkQueue handle() const { return queue_; }
+
+  iree_hal_vulkan_tracing_context_t* tracing_context() {
+    return tracing_context_;
+  }
+  void set_tracing_context(iree_hal_vulkan_tracing_context_t* tracing_context) {
+    tracing_context_ = tracing_context;
+  }
+
   bool can_dispatch() const {
     return iree_all_bits_set(supported_categories_,
                              IREE_HAL_COMMAND_CATEGORY_DISPATCH);
@@ -52,19 +62,19 @@ class CommandQueue {
   virtual iree_status_t WaitIdle(iree_time_t deadline_ns) = 0;
 
  protected:
-  CommandQueue(VkDeviceHandle* logical_device, std::string name,
+  CommandQueue(VkDeviceHandle* logical_device,
                iree_hal_command_category_t supported_categories, VkQueue queue)
       : logical_device_(logical_device),
-        name_(std::move(name)),
         supported_categories_(supported_categories),
         queue_(queue) {
     iree_slim_mutex_initialize(&queue_mutex_);
   }
 
   VkDeviceHandle* logical_device_;
-  const std::string name_;
   const iree_hal_command_category_t supported_categories_;
 
+  iree_hal_vulkan_tracing_context_t* tracing_context_ = nullptr;
+
   // VkQueue needs to be externally synchronized.
   iree_slim_mutex_t queue_mutex_;
   VkQueue queue_ IREE_GUARDED_BY(queue_mutex_);

iree/hal/vulkan/direct_command_buffer.cc

@@ -36,6 +36,7 @@ typedef struct {
   iree_hal_command_buffer_mode_t mode;
   iree_hal_command_category_t allowed_categories;
   iree_hal_queue_affinity_t queue_affinity;
+  iree_hal_vulkan_tracing_context_t* tracing_context;
 
   VkCommandPoolHandle* command_pool;
   VkCommandBuffer handle;
@@ -68,6 +69,7 @@ iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
     iree_hal_queue_affinity_t queue_affinity,
+    iree_hal_vulkan_tracing_context_t* tracing_context,
     iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(logical_device);
@@ -98,6 +100,7 @@ iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
     command_buffer->mode = mode;
     command_buffer->allowed_categories = command_categories;
     command_buffer->queue_affinity = queue_affinity;
+    command_buffer->tracing_context = tracing_context;
     command_buffer->command_pool = command_pool;
     command_buffer->handle = handle;
     command_buffer->syms = logical_device->syms().get();
@@ -564,6 +567,15 @@ static iree_status_t iree_hal_vulkan_direct_command_buffer_dispatch(
   iree_hal_vulkan_direct_command_buffer_t* command_buffer =
       iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer);
 
+  iree_hal_vulkan_source_location_t source_location;
+  iree_hal_vulkan_native_executable_entry_point_source_location(
+      executable, entry_point, &source_location);
+  IREE_VULKAN_TRACE_ZONE_BEGIN_EXTERNAL(
+      command_buffer->tracing_context, command_buffer->handle,
+      source_location.file_name.data, source_location.file_name.size,
+      source_location.line, source_location.func_name.data,
+      source_location.func_name.size, NULL, 0);
+
   // Get the compiled and linked pipeline for the specified entry point and
   // bind it to the command buffer.
   VkPipeline pipeline_handle = VK_NULL_HANDLE;
@@ -576,6 +588,9 @@ static iree_status_t iree_hal_vulkan_direct_command_buffer_dispatch(
   command_buffer->syms->vkCmdDispatch(command_buffer->handle, workgroup_x,
                                       workgroup_y, workgroup_z);
 
+  IREE_VULKAN_TRACE_ZONE_END(command_buffer->tracing_context,
+                             command_buffer->handle);
+
   return iree_ok_status();
 }
 
@@ -587,6 +602,15 @@ static iree_status_t iree_hal_vulkan_direct_command_buffer_dispatch_indirect(
   iree_hal_vulkan_direct_command_buffer_t* command_buffer =
       iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer);
 
+  iree_hal_vulkan_source_location_t source_location;
+  iree_hal_vulkan_native_executable_entry_point_source_location(
+      executable, entry_point, &source_location);
+  IREE_VULKAN_TRACE_ZONE_BEGIN_EXTERNAL(
+      command_buffer->tracing_context, command_buffer->handle,
+      source_location.file_name.data, source_location.file_name.size,
+      source_location.line, source_location.func_name.data,
+      source_location.func_name.size, NULL, 0);
+
   // Get the compiled and linked pipeline for the specified entry point and
   // bind it to the command buffer.
   VkPipeline pipeline_handle = VK_NULL_HANDLE;
@@ -602,6 +626,9 @@ static iree_status_t iree_hal_vulkan_direct_command_buffer_dispatch_indirect(
   command_buffer->syms->vkCmdDispatchIndirect(
       command_buffer->handle, workgroups_device_buffer, workgroups_offset);
 
+  IREE_VULKAN_TRACE_ZONE_END(command_buffer->tracing_context,
+                             command_buffer->handle);
+
   return iree_ok_status();
 }
 

iree/hal/vulkan/direct_command_buffer.h

@@ -18,6 +18,7 @@
 #include "iree/hal/api.h"
 #include "iree/hal/vulkan/descriptor_pool_cache.h"
 #include "iree/hal/vulkan/handle_util.h"
+#include "iree/hal/vulkan/tracing.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -30,6 +31,7 @@ iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
     iree_hal_queue_affinity_t queue_affinity,
+    iree_hal_vulkan_tracing_context_t* tracing_context,
     iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
     iree_hal_command_buffer_t** out_command_buffer);
 

iree/hal/vulkan/direct_command_queue.cc

@@ -26,10 +26,9 @@ namespace hal {
 namespace vulkan {
 
 DirectCommandQueue::DirectCommandQueue(
-    VkDeviceHandle* logical_device, std::string name,
+    VkDeviceHandle* logical_device,
     iree_hal_command_category_t supported_categories, VkQueue queue)
-    : CommandQueue(logical_device, std::move(name), supported_categories,
-                   queue) {}
+    : CommandQueue(logical_device, supported_categories, queue) {}
 
 DirectCommandQueue::~DirectCommandQueue() = default;
 
@@ -134,6 +133,7 @@ iree_status_t DirectCommandQueue::WaitIdle(iree_time_t deadline_ns) {
     iree_status_t status =
         VK_RESULT_TO_STATUS(syms()->vkQueueWaitIdle(queue_), "vkQueueWaitIdle");
     iree_slim_mutex_unlock(&queue_mutex_);
+    iree_hal_vulkan_tracing_context_collect(tracing_context(), VK_NULL_HANDLE);
     return status;
   }
 
@@ -191,6 +191,8 @@ iree_status_t DirectCommandQueue::WaitIdle(iree_time_t deadline_ns) {
 
   syms()->vkDestroyFence(*logical_device_, fence, logical_device_->allocator());
 
+  iree_hal_vulkan_tracing_context_collect(tracing_context(), VK_NULL_HANDLE);
+
   return status;
 }
 

iree/hal/vulkan/direct_command_queue.h

@@ -25,7 +25,7 @@ namespace vulkan {
 // Command queue implementation directly maps to VkQueue.
 class DirectCommandQueue final : public CommandQueue {
  public:
-  DirectCommandQueue(VkDeviceHandle* logical_device, std::string name,
+  DirectCommandQueue(VkDeviceHandle* logical_device,
                      iree_hal_command_category_t supported_categories,
                      VkQueue queue);
   ~DirectCommandQueue() override;

iree/hal/vulkan/dynamic_symbol_tables.h

@@ -113,7 +113,7 @@ namespace vulkan {
   DEV_PFN(EXCLUDED, vkCmdPushDescriptorSetWithTemplateKHR)              \
   DEV_PFN(EXCLUDED, vkCmdReserveSpaceForCommandsNVX)                    \
   DEV_PFN(REQUIRED, vkCmdResetEvent)                                    \
-  DEV_PFN(EXCLUDED, vkCmdResetQueryPool)                                \
+  DEV_PFN(REQUIRED, vkCmdResetQueryPool)                                \
   DEV_PFN(EXCLUDED, vkCmdResolveImage)                                  \
   DEV_PFN(EXCLUDED, vkCmdSetBlendConstants)                             \
   DEV_PFN(EXCLUDED, vkCmdSetCheckpointNV)                               \
@@ -174,7 +174,7 @@ namespace vulkan {
   DEV_PFN(EXCLUDED, vkCreateObjectTableNVX)                             \
   DEV_PFN(REQUIRED, vkCreatePipelineCache)                              \
   DEV_PFN(REQUIRED, vkCreatePipelineLayout)                             \
-  DEV_PFN(EXCLUDED, vkCreateQueryPool)                                  \
+  DEV_PFN(REQUIRED, vkCreateQueryPool)                                  \
   DEV_PFN(EXCLUDED, vkCreateRayTracingPipelinesNV)                      \
   DEV_PFN(EXCLUDED, vkCreateRenderPass)                                 \
   DEV_PFN(EXCLUDED, vkCreateRenderPass2KHR)                             \
@@ -207,7 +207,7 @@ namespace vulkan {
   DEV_PFN(REQUIRED, vkDestroyPipeline)                                  \
   DEV_PFN(REQUIRED, vkDestroyPipelineCache)                             \
   DEV_PFN(REQUIRED, vkDestroyPipelineLayout)                            \
-  DEV_PFN(EXCLUDED, vkDestroyQueryPool)                                 \
+  DEV_PFN(REQUIRED, vkDestroyQueryPool)                                 \
   DEV_PFN(EXCLUDED, vkDestroyRenderPass)                                \
   DEV_PFN(EXCLUDED, vkDestroySampler)                                   \
   DEV_PFN(EXCLUDED, vkDestroySamplerYcbcrConversion)                    \
@@ -228,7 +228,7 @@ namespace vulkan {
   DEV_PFN(REQUIRED, vkGetBufferMemoryRequirements)                      \
   DEV_PFN(EXCLUDED, vkGetBufferMemoryRequirements2)                     \
   DEV_PFN(EXCLUDED, vkGetBufferMemoryRequirements2KHR)                  \
-  DEV_PFN(EXCLUDED, vkGetCalibratedTimestampsEXT)                       \
+  DEV_PFN(OPTIONAL, vkGetCalibratedTimestampsEXT)                       \
   DEV_PFN(EXCLUDED, vkGetDescriptorSetLayoutSupport)                    \
   DEV_PFN(EXCLUDED, vkGetDescriptorSetLayoutSupportKHR)                 \
   DEV_PFN(EXCLUDED, vkGetDeviceGroupPeerMemoryFeatures)                 \
@@ -255,7 +255,7 @@ namespace vulkan {
   DEV_PFN(EXCLUDED, vkGetMemoryHostPointerPropertiesEXT)                \
   DEV_PFN(EXCLUDED, vkGetPastPresentationTimingGOOGLE)                  \
   DEV_PFN(REQUIRED, vkGetPipelineCacheData)                             \
-  DEV_PFN(EXCLUDED, vkGetQueryPoolResults)                              \
+  DEV_PFN(REQUIRED, vkGetQueryPoolResults)                              \
   DEV_PFN(EXCLUDED, vkGetRayTracingShaderGroupHandlesNV)                \
   DEV_PFN(EXCLUDED, vkGetRefreshCycleDurationGOOGLE)                    \
   DEV_PFN(EXCLUDED, vkGetRenderAreaGranularity)                         \
@@ -278,7 +278,8 @@ namespace vulkan {
   DEV_PFN(REQUIRED, vkResetDescriptorPool)                              \
   DEV_PFN(REQUIRED, vkResetEvent)                                       \
   DEV_PFN(REQUIRED, vkResetFences)                                      \
-  DEV_PFN(EXCLUDED, vkResetQueryPoolEXT)                                \
+  DEV_PFN(OPTIONAL, vkResetQueryPool)                                   \
+  DEV_PFN(OPTIONAL, vkResetQueryPoolEXT)                                \
   DEV_PFN(OPTIONAL, vkSetDebugUtilsObjectNameEXT)                       \
   DEV_PFN(OPTIONAL, vkSetDebugUtilsObjectTagEXT)                        \
   DEV_PFN(REQUIRED, vkSetEvent)                                         \
@@ -322,7 +323,7 @@ namespace vulkan {
   INS_PFN(EXCLUDED, vkGetDisplayPlaneCapabilities2KHR)                  \
   INS_PFN(EXCLUDED, vkGetDisplayPlaneCapabilitiesKHR)                   \
   INS_PFN(EXCLUDED, vkGetDisplayPlaneSupportedDisplaysKHR)              \
-  INS_PFN(EXCLUDED, vkGetPhysicalDeviceCalibrateableTimeDomainsEXT)     \
+  INS_PFN(OPTIONAL, vkGetPhysicalDeviceCalibrateableTimeDomainsEXT)     \
   INS_PFN(EXCLUDED, vkGetPhysicalDeviceCooperativeMatrixPropertiesNV)   \
   INS_PFN(EXCLUDED, vkGetPhysicalDeviceDisplayPlaneProperties2KHR)      \
   INS_PFN(EXCLUDED, vkGetPhysicalDeviceDisplayPlanePropertiesKHR)       \

iree/hal/vulkan/extensibility_util.cc

@@ -206,6 +206,12 @@ iree_hal_vulkan_populate_enabled_device_extensions(
     } else if (strcmp(extension_name,
                       VK_KHR_TIMELINE_SEMAPHORE_EXTENSION_NAME) == 0) {
       extensions.timeline_semaphore = true;
+    } else if (strcmp(extension_name, VK_EXT_HOST_QUERY_RESET_EXTENSION_NAME) ==
+               0) {
+      extensions.host_query_reset = true;
+    } else if (strcmp(extension_name,
+                      VK_EXT_CALIBRATED_TIMESTAMPS_EXTENSION_NAME) == 0) {
+      extensions.calibrated_timestamps = true;
     }
   }
   return extensions;
@@ -222,5 +228,11 @@ iree_hal_vulkan_infer_enabled_device_extensions(
   if (device_syms->vkSignalSemaphore || device_syms->vkSignalSemaphoreKHR) {
     extensions.timeline_semaphore = true;
   }
+  if (device_syms->vkResetQueryPoolEXT) {
+    extensions.host_query_reset = true;
+  }
+  if (device_syms->vkGetCalibratedTimestampsEXT) {
+    extensions.calibrated_timestamps = true;
+  }
   return extensions;
 }

iree/hal/vulkan/extensibility_util.h

@@ -83,6 +83,10 @@ typedef struct {
   bool push_descriptors : 1;
   // VK_KHR_timeline_semaphore is enabled.
   bool timeline_semaphore : 1;
+  // VK_EXT_host_query_reset is enabled.
+  bool host_query_reset : 1;
+  // VK_EXT_calibrated_timestamps is enabled.
+  bool calibrated_timestamps : 1;
 } iree_hal_vulkan_device_extensions_t;
 
 // Returns a bitfield with all of the provided extension names.