[mlir][vulkan-runner] Add basic timing for compute pipeline

[antiagainst]

Summary

This commit adds timestamp query commands in Vulkan runner's
compute pipeline to gain insights into how long it takes to
run the compute shader. This commit also adds timing from CPU
side for VkQueueSubmit and vkQueueWaitIdle.

[rriddle]

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

17

LLVM forbids the use of iostream.

https://llvm.org/docs/CodingStandards.html#include-iostream-is-forbidden

[benvanik]

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

721

another*

722

you want to insert this after the vkCmdDispatch below - the timestamp only latches the value after previous commands complete (in this case, just the timestamp reset/initial query)

[antiagainst]

Address comments

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

17

Changed to llvm::outs. I saw we are using iostream already in RunnerUtils.h and cblas_interface.h so wanted to be consistent there. Should those be fixed?

722

Thanks! I find the spec a bit difficult to parse: "latches the value of the timer when all previous commands have completed executing as far as the specified pipeline stage". I interpret it as the pipeline stage (VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT) controls when to write out the timestamp, not the placement of the vkCmdWriteTimestamp command. (Otherwise the pipeline stage specification does not really make sense.) But I do see there is a difference if I move the second vkCmdWriteTimestamp below vkCmdDispatch. So I'm a bit puzzled here.

[benvanik]

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

722

I believe that vkCmdWriteTimestamp acts as an implicit vkCmdPipelineBarrier with the srcStageMask set to what you pass in (BOTTOM_OF_PIPE), so it effectively partitions the command buffer into two synchronization scopes based on the order the commands were recorded.

[mehdi_amini]

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

17

Isn’t cblas_interface for generated code intended to be used a runtime for generated code, including outside of a JIT environment?

[antiagainst]

Address comments

[antiagainst]

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

17

Good point!

722

Done. I think your interpretation is correct. This makes me wanting to look into the driver implementation over this. :)

mlir/tools/mlir-vulkan-runner/VulkanRuntime.h

 
 private:
   //===--------------------------------------------------------------------===//
   // Pipeline creation methods.
   //===--------------------------------------------------------------------===//
 
   LogicalResult createInstance();
   LogicalResult createDevice();
-  LogicalResult getBestComputeQueue(const VkPhysicalDevice &physicalDevice);
+  LogicalResult getBestComputeQueue();
   LogicalResult createMemoryBuffers();
   LogicalResult createShaderModule();
   void initDescriptorSetLayoutBindingMap();
   LogicalResult createDescriptorSetLayout();
   LogicalResult createPipelineLayout();
   LogicalResult createComputePipeline();
   LogicalResult createDescriptorPool();
   LogicalResult allocateDescriptorSets();
   LogicalResult setWriteDescriptors();
   LogicalResult createCommandPool();
+  LogicalResult createQueryPool();
   LogicalResult createComputeCommandBuffer();
   LogicalResult submitCommandBuffersToQueue();
 
   //===--------------------------------------------------------------------===//
   // Helper methods.
   //===--------------------------------------------------------------------===//
 
   /// Maps storage class to a descriptor type.
   LogicalResult
   mapStorageClassToDescriptorType(spirv::StorageClass storageClass,
                                   VkDescriptorType &descriptorType);
 
   /// Maps storage class to buffer usage flags.
   LogicalResult
   mapStorageClassToBufferUsageFlag(spirv::StorageClass storageClass,
                                    VkBufferUsageFlagBits &bufferUsage);
 
   LogicalResult countDeviceMemorySize();
 
   //===--------------------------------------------------------------------===//
   // Vulkan objects.
   //===--------------------------------------------------------------------===//
 
-  VkInstance instance;
-  VkDevice device;
-  VkQueue queue;
+  VkInstance instance{VK_NULL_HANDLE};
+  VkPhysicalDevice physicalDevice{VK_NULL_HANDLE};
+  VkDevice device{VK_NULL_HANDLE};
+  VkQueue queue{VK_NULL_HANDLE};
 
   /// Specifies VulkanDeviceMemoryBuffers divided into sets.
   llvm::DenseMap<DescriptorSetIndex,
                  llvm::SmallVector<VulkanDeviceMemoryBuffer, 1>>
       deviceMemoryBufferMap;
 
   /// Specifies shader module.
-  VkShaderModule shaderModule;
+  VkShaderModule shaderModule{VK_NULL_HANDLE};
 
   /// Specifies layout bindings.
   llvm::DenseMap<DescriptorSetIndex,
                  llvm::SmallVector<VkDescriptorSetLayoutBinding, 1>>
       descriptorSetLayoutBindingMap;
 
   /// Specifies layouts of descriptor sets.
   llvm::SmallVector<VkDescriptorSetLayout, 1> descriptorSetLayouts;
-  VkPipelineLayout pipelineLayout;
+  VkPipelineLayout pipelineLayout{VK_NULL_HANDLE};
 
   /// Specifies descriptor sets.
   llvm::SmallVector<VkDescriptorSet, 1> descriptorSets;
 
   /// Specifies a pool of descriptor set info, each descriptor set must have
   /// information such as type, index and amount of bindings.
   llvm::SmallVector<DescriptorSetInfo, 1> descriptorSetInfoPool;
-  VkDescriptorPool descriptorPool;
+  VkDescriptorPool descriptorPool{VK_NULL_HANDLE};
+
+  /// Timestamp query.
+  VkQueryPool queryPool{VK_NULL_HANDLE};
+  // Number of nonoseconds for timestamp to increase 1
+  float timestampPeriod{0.f};
 
   /// Computation pipeline.
-  VkPipeline pipeline;
-  VkCommandPool commandPool;
+  VkPipeline pipeline{VK_NULL_HANDLE};
+  VkCommandPool commandPool{VK_NULL_HANDLE};
   llvm::SmallVector<VkCommandBuffer, 1> commandBuffers;
 
   //===--------------------------------------------------------------------===//
   // Vulkan memory context.
   //===--------------------------------------------------------------------===//
 
   uint32_t queueFamilyIndex{0};
+  VkQueueFamilyProperties queueFamilyProperties{};
   uint32_t memoryTypeIndex{VK_MAX_MEMORY_TYPES};
   VkDeviceSize memorySize{0};
 
   //===--------------------------------------------------------------------===//
   // Vulkan execution context.
   //===--------------------------------------------------------------------===//
 
   NumWorkGroups numWorkGroups;

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

 //===- VulkanRuntime.cpp - MLIR Vulkan runtime ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file provides a library for running a module on a Vulkan device.
 // Implements a Vulkan runtime.
 //
 //===----------------------------------------------------------------------===//
 
 #include "VulkanRuntime.h"
 
+#include "llvm/Support/Format.h"
+#include <chrono>
+
 using namespace mlir;
 
 void VulkanRuntime::setNumWorkGroups(const NumWorkGroups &numberWorkGroups) {
   numWorkGroups = numberWorkGroups;
 }
 
 void VulkanRuntime::setResourceStorageClassBindingMap(
     const ResourceStorageClassBindingMap &stClassData) {
LogicalResult VulkanRuntime::destroy() {
   // an application is responsible for destroying/freeing any Vulkan objects
   // that were created using that device as the first parameter of the
   // corresponding vkCreate* or vkAllocate* command."
   RETURN_ON_VULKAN_ERROR(vkDeviceWaitIdle(device), "vkDeviceWaitIdle");
 
   // Free and destroy.
   vkFreeCommandBuffers(device, commandPool, commandBuffers.size(),
                        commandBuffers.data());
+  vkDestroyQueryPool(device, queryPool, nullptr);
   vkDestroyCommandPool(device, commandPool, nullptr);
   vkFreeDescriptorSets(device, descriptorPool, descriptorSets.size(),
                        descriptorSets.data());
   vkDestroyDescriptorPool(device, descriptorPool, nullptr);
   vkDestroyPipeline(device, pipeline, nullptr);
   vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
   for (auto &descriptorSetLayout: descriptorSetLayouts) {
     vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);
LogicalResult VulkanRuntime::run() {
   // must have a layout binding attached into a descriptor set layout.
   // Each layout set must be binded into a pipeline layout.
   initDescriptorSetLayoutBindingMap();
   if (failed(createDescriptorSetLayout()) || failed(createPipelineLayout()) ||
       // Each descriptor set must be allocated from a descriptor pool.
       failed(createComputePipeline()) || failed(createDescriptorPool()) ||
       failed(allocateDescriptorSets()) || failed(setWriteDescriptors()) ||
       // Create command buffer.
-      failed(createCommandPool()) || failed(createComputeCommandBuffer())) {
+      failed(createCommandPool()) || failed(createQueryPool()) ||
+      failed(createComputeCommandBuffer())) {
     return failure();
   }
 
   // Get working queue.
   vkGetDeviceQueue(device, queueFamilyIndex, 0, &queue);
 
+  auto submitStart = std::chrono::high_resolution_clock::now();
   // Submit command buffer into the queue.
   if (failed(submitCommandBuffersToQueue()))
     return failure();
+  auto submitEnd = std::chrono::high_resolution_clock::now();
 
   RETURN_ON_VULKAN_ERROR(vkQueueWaitIdle(queue), "vkQueueWaitIdle");
+  auto execEnd = std::chrono::high_resolution_clock::now();
+
+  auto submitDuration = std::chrono::duration_cast<std::chrono::microseconds>(
+      submitEnd - submitStart);
+  auto execDuration = std::chrono::duration_cast<std::chrono::microseconds>(
+      execEnd - submitEnd);
+
+  if (queryPool != VK_NULL_HANDLE) {
+    uint64_t timestamps[2];
+    RETURN_ON_VULKAN_ERROR(
+        vkGetQueryPoolResults(
+            device, queryPool, /*firstQuery=*/0, /*queryCount=*/2,
+            /*dataSize=*/sizeof(timestamps),
+            /*pData=*/reinterpret_cast<void *>(timestamps),
+            /*stride=*/sizeof(uint64_t),
+            VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT),
+        "vkGetQueryPoolResults");
+    float microsec = (timestamps[1] - timestamps[0]) * timestampPeriod / 1000;
+    llvm::outs() << "Compute shader execution time: "
+                 << llvm::format("%0.3fus\n", microsec);
+  }
+
+  llvm::outs() << "Command buffer submit time: " << submitDuration.count()
+               << "us\nWait idle time: " << execDuration.count() << "us\n";
+
   return success();
 }
 
 LogicalResult VulkanRuntime::createInstance() {
   VkApplicationInfo applicationInfo = {};
   applicationInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
   applicationInfo.pNext = nullptr;
   applicationInfo.pApplicationName = "MLIR Vulkan runtime";
RETURN_ON_VULKAN_ERROR(vkEnumeratePhysicalDevices(instance,
                                                     &physicalDeviceCount,
                                                     physicalDevices.data()),
                          "vkEnumeratePhysicalDevices");
 
   RETURN_ON_VULKAN_ERROR(physicalDeviceCount ? VK_SUCCESS : VK_INCOMPLETE,
                          "physicalDeviceCount");
 
   // TODO(denis0x0D): find the best device.
-  const auto &physicalDevice = physicalDevices.front();
-  getBestComputeQueue(physicalDevice);
+  physicalDevice = physicalDevices.front();
+  if (failed(getBestComputeQueue()))
+    return failure();
 
   const float queuePrioritory = 1.0f;
   VkDeviceQueueCreateInfo deviceQueueCreateInfo = {};
   deviceQueueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
   deviceQueueCreateInfo.pNext = nullptr;
   deviceQueueCreateInfo.flags = 0;
   deviceQueueCreateInfo.queueFamilyIndex = queueFamilyIndex;
   deviceQueueCreateInfo.queueCount = 1;
LogicalResult VulkanRuntime::createDevice() {
   }
 
   RETURN_ON_VULKAN_ERROR(memoryTypeIndex == VK_MAX_MEMORY_TYPES ? VK_INCOMPLETE
                                                                 : VK_SUCCESS,
                          "invalid memoryTypeIndex");
   return success();
 }
 
-LogicalResult
-VulkanRuntime::getBestComputeQueue(const VkPhysicalDevice &physicalDevice) {
+LogicalResult VulkanRuntime::getBestComputeQueue() {
   uint32_t queueFamilyPropertiesCount = 0;
   vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice,
                                            &queueFamilyPropertiesCount, 0);
-  SmallVector<VkQueueFamilyProperties, 1> queueFamilyProperties(
-      queueFamilyPropertiesCount);
 
-  vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice,
-                                           &queueFamilyPropertiesCount,
-                                           queueFamilyProperties.data());
+  SmallVector<VkQueueFamilyProperties, 1> familyProperties(
+      queueFamilyPropertiesCount);
+  vkGetPhysicalDeviceQueueFamilyProperties(
+      physicalDevice, &queueFamilyPropertiesCount, familyProperties.data());
 
   // VK_QUEUE_COMPUTE_BIT specifies that queues in this queue family support
-  // compute operations.
+  // compute operations. Try to find a compute-only queue first if possible.
   for (uint32_t i = 0; i < queueFamilyPropertiesCount; ++i) {
-    const VkQueueFlags maskedFlags =
-        (~(VK_QUEUE_TRANSFER_BIT | VK_QUEUE_SPARSE_BINDING_BIT) &
-         queueFamilyProperties[i].queueFlags);
-
-    if (!(VK_QUEUE_GRAPHICS_BIT & maskedFlags) &&
-        (VK_QUEUE_COMPUTE_BIT & maskedFlags)) {
+    auto flags = familyProperties[i].queueFlags;
+    if ((flags & VK_QUEUE_COMPUTE_BIT) && !(flags & VK_QUEUE_GRAPHICS_BIT)) {
       queueFamilyIndex = i;
+      queueFamilyProperties = familyProperties[i];
       return success();
     }
   }
 
+  // Otherwise use a queue that can also support graphics.
   for (uint32_t i = 0; i < queueFamilyPropertiesCount; ++i) {
-    const VkQueueFlags maskedFlags =
-        (~(VK_QUEUE_TRANSFER_BIT | VK_QUEUE_SPARSE_BINDING_BIT) &
-         queueFamilyProperties[i].queueFlags);
-
-    if (VK_QUEUE_COMPUTE_BIT & maskedFlags) {
+    auto flags = familyProperties[i].queueFlags;
+    if ((flags & VK_QUEUE_COMPUTE_BIT)) {
       queueFamilyIndex = i;
+      queueFamilyProperties = familyProperties[i];
       return success();
     }
   }
 
   llvm::errs() << "cannot find valid queue";
   return failure();
 }
 

 }
 
 LogicalResult VulkanRuntime::createCommandPool() {
   VkCommandPoolCreateInfo commandPoolCreateInfo = {};
   commandPoolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
   commandPoolCreateInfo.pNext = nullptr;
   commandPoolCreateInfo.flags = 0;
   commandPoolCreateInfo.queueFamilyIndex = queueFamilyIndex;
-  RETURN_ON_VULKAN_ERROR(
-      vkCreateCommandPool(device, &commandPoolCreateInfo, 0, &commandPool),
+  RETURN_ON_VULKAN_ERROR(vkCreateCommandPool(device, &commandPoolCreateInfo,
+                                             /*pAllocator=*/nullptr,
+                                             &commandPool),
                          "vkCreateCommandPool");
   return success();
 }
 
+LogicalResult VulkanRuntime::createQueryPool() {
+  // Return directly if timestamp query is not supported.
+  if (queueFamilyProperties.timestampValidBits == 0)
+    return success();
+
+  // Get timestamp period for this physical device.
+  VkPhysicalDeviceProperties deviceProperties = {};
+  vkGetPhysicalDeviceProperties(physicalDevice, &deviceProperties);
+  timestampPeriod = deviceProperties.limits.timestampPeriod;
+
+  // Create query pool.
+  VkQueryPoolCreateInfo queryPoolCreateInfo = {};
+  queryPoolCreateInfo.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO;
+  queryPoolCreateInfo.pNext = nullptr;
+  queryPoolCreateInfo.flags = 0;
+  queryPoolCreateInfo.queryType = VK_QUERY_TYPE_TIMESTAMP;
+  queryPoolCreateInfo.queryCount = 2;
+  queryPoolCreateInfo.pipelineStatistics = 0;
+  RETURN_ON_VULKAN_ERROR(vkCreateQueryPool(device, &queryPoolCreateInfo,
+                                           /*pAllocator=*/nullptr, &queryPool),
+                         "vkCreateQueryPool");
+
+  return success();
+}
+
 LogicalResult VulkanRuntime::createComputeCommandBuffer() {
   VkCommandBufferAllocateInfo commandBufferAllocateInfo = {};
-  VkCommandBuffer commandBuffer;
   commandBufferAllocateInfo.sType =
       VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
   commandBufferAllocateInfo.pNext = nullptr;
   commandBufferAllocateInfo.commandPool = commandPool;
   commandBufferAllocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
   commandBufferAllocateInfo.commandBufferCount = 1;
+
+  VkCommandBuffer commandBuffer;
   RETURN_ON_VULKAN_ERROR(vkAllocateCommandBuffers(device,
                                                   &commandBufferAllocateInfo,
                                                   &commandBuffer),
                          "vkAllocateCommandBuffers");
 
   VkCommandBufferBeginInfo commandBufferBeginInfo = {};
   commandBufferBeginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
   commandBufferBeginInfo.pNext = nullptr;
   commandBufferBeginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
   commandBufferBeginInfo.pInheritanceInfo = nullptr;
 
   // Commands begin.
   RETURN_ON_VULKAN_ERROR(
       vkBeginCommandBuffer(commandBuffer, &commandBufferBeginInfo),
       "vkBeginCommandBuffer");
 
-  // Commands.
+  if (queryPool != VK_NULL_HANDLE)
+    vkCmdResetQueryPool(commandBuffer, queryPool, 0, 2);
+
   vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
   vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE,
                           pipelineLayout, 0, descriptorSets.size(),
                           descriptorSets.data(), 0, 0);
+  // Get a timestamp before invoking the compute shader.
+  if (queryPool != VK_NULL_HANDLE)
+    vkCmdWriteTimestamp(commandBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
+                        queryPool, 0);
   vkCmdDispatch(commandBuffer, numWorkGroups.x, numWorkGroups.y,
                 numWorkGroups.z);
+  // Get another timestamp after invoking the compute shader.
+  if (queryPool != VK_NULL_HANDLE)
+    vkCmdWriteTimestamp(commandBuffer, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
+                        queryPool, 1);
 
   // Commands end.
   RETURN_ON_VULKAN_ERROR(vkEndCommandBuffer(commandBuffer),
                          "vkEndCommandBuffer");
 
   commandBuffers.push_back(commandBuffer);
   return success();
 }