[ET-VK] Add utility functions to get pipeline executable properties

Differential Revision: D69125488

Pull Request resolved: #8182

backends/vulkan/runtime/api/Context.cpp

@@ -8,6 +8,11 @@
 
 #include <executorch/backends/vulkan/runtime/api/Context.h>
 
+#ifdef VULKAN_DEBUG
+#include <iomanip>
+#include <iostream>
+#endif // VULKAN_DEBUG
+
 #ifndef VULKAN_DESCRIPTOR_POOL_SIZE
 #define VULKAN_DESCRIPTOR_POOL_SIZE 1024u
 #endif
@@ -261,5 +266,220 @@ Context* context() {
   return context.get();
 }
 
+#ifdef VULKAN_DEBUG
+
+#ifdef VK_KHR_pipeline_executable_properties
+
+VkPipeline Context::get_shader_pipeline(
+    const vkapi::ShaderInfo& shader,
+    const vkapi::SpecVarList& spec_constants) {
+  const uint32_t push_constants_size = 128u;
+
+  VkDescriptorSetLayout shader_layout =
+      shader_layout_cache().retrieve(shader.kernel_layout);
+  VkPipelineLayout pipeline_layout =
+      pipeline_layout_cache().retrieve(shader_layout, push_constants_size);
+
+  vkapi::SpecVarList spec_constants_full_list = {4u, 4u, 1u};
+  spec_constants_full_list.append(spec_constants);
+
+  VkPipeline pipeline = pipeline_cache().retrieve(
+      {pipeline_layout,
+       shader_cache().retrieve(shader),
+       spec_constants_full_list});
+
+  return pipeline;
+}
+
+std::vector<VkPipelineExecutablePropertiesKHR>
+Context::get_pipeline_executable_props(const VkPipeline pipeline) {
+  VkPipelineInfoKHR pipeline_info{
+      VK_STRUCTURE_TYPE_PIPELINE_INFO_KHR,
+      nullptr,
+      pipeline,
+  };
+
+  uint32_t shader_props_count = 0u;
+  vkGetPipelineExecutablePropertiesKHR(
+      device(), &pipeline_info, &shader_props_count, nullptr);
+
+  std::vector<VkPipelineExecutablePropertiesKHR> pipeline_props(
+      shader_props_count);
+  for (int i = 0; i < shader_props_count; i++) {
+    pipeline_props.at(i).sType =
+        VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_PROPERTIES_KHR;
+    pipeline_props.at(i).pNext = nullptr;
+  }
+  vkGetPipelineExecutablePropertiesKHR(
+      device(), &pipeline_info, &shader_props_count, pipeline_props.data());
+
+  return pipeline_props;
+}
+
+std::tuple<
+    std::vector<VkPipelineExecutableInternalRepresentationKHR>,
+    std::vector<std::vector<char>>>
+Context::get_shader_executable_irs(
+    const VkPipeline pipeline,
+    const uint32_t pipeline_exec_idx) {
+  VkPipelineExecutableInfoKHR exec_info{
+      VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_INFO_KHR,
+      nullptr,
+      pipeline,
+      pipeline_exec_idx,
+  };
+
+  uint32_t ir_count;
+  VK_CHECK(vkGetPipelineExecutableInternalRepresentationsKHR(
+      device(), &exec_info, &ir_count, nullptr));
+
+  std::vector<VkPipelineExecutableInternalRepresentationKHR> irs(ir_count);
+  for (int i = 0; i < ir_count; i++) {
+    irs.at(i).sType =
+        VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_INTERNAL_REPRESENTATION_KHR;
+    irs.at(i).pNext = nullptr;
+    irs.at(i).pData = nullptr;
+  }
+  VK_CHECK(vkGetPipelineExecutableInternalRepresentationsKHR(
+      device(), &exec_info, &ir_count, irs.data()));
+
+  std::vector<std::vector<char>> irs_data(ir_count);
+  for (int i = 0; i < ir_count; i++) {
+    irs_data.at(i).resize(irs.at(i).dataSize);
+    irs.at(i).pData = irs_data.at(i).data();
+  }
+  VK_CHECK(vkGetPipelineExecutableInternalRepresentationsKHR(
+      device(), &exec_info, &ir_count, irs.data()));
+
+  return std::make_tuple(irs, irs_data);
+}
+
+std::vector<VkPipelineExecutableStatisticKHR>
+Context::get_shader_executable_stats(
+    const VkPipeline pipeline,
+    const uint32_t pipeline_exec_idx) {
+  VkPipelineExecutableInfoKHR exec_info{
+      VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_INFO_KHR,
+      nullptr,
+      pipeline,
+      pipeline_exec_idx,
+  };
+
+  uint32_t stats_count;
+  VK_CHECK(vkGetPipelineExecutableStatisticsKHR(
+      device(), &exec_info, &stats_count, NULL));
+
+  std::vector<VkPipelineExecutableStatisticKHR> shader_stats(stats_count);
+  for (int i = 0; i < stats_count; i++) {
+    shader_stats.at(i).sType =
+        VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_STATISTIC_KHR;
+    shader_stats.at(i).pNext = nullptr;
+  }
+  vkGetPipelineExecutableStatisticsKHR(
+      device(), &exec_info, &stats_count, shader_stats.data());
+
+  return shader_stats;
+}
+
+std::ostream& operator<<(
+    std::ostream& os,
+    const VkPipelineExecutablePropertiesKHR& props) {
+  os << std::left << std::setw(10) << "name: " << props.name << std::endl;
+  os << std::left << std::setw(10) << "descr: " << props.description
+     << std::endl;
+  os << std::left << std::setw(10) << "subgroup: " << props.subgroupSize
+     << std::endl;
+
+  return os;
+}
+
+std::ostream& operator<<(
+    std::ostream& os,
+    const VkPipelineExecutableInternalRepresentationKHR& ir) {
+  os << std::left << std::setw(10) << "descr: " << ir.description << std::endl;
+  os << std::left << std::setw(10) << "isText: " << ir.isText << std::endl;
+  os << std::left << std::setw(10) << "size: " << ir.dataSize << std::endl;
+  if (ir.isText) {
+    os << "text:" << std::endl;
+    char* str = (char*)ir.pData;
+    os << str << std::endl;
+  }
+  return os;
+}
+
+std::ostream& operator<<(
+    std::ostream& os,
+    VkPipelineExecutableStatisticKHR& stat) {
+  os << stat.name << ": ";
+  switch (stat.format) {
+    case VK_PIPELINE_EXECUTABLE_STATISTIC_FORMAT_BOOL32_KHR:
+      os << (stat.value.b32 ? "true" : "false") << std::endl;
+      break;
+    case VK_PIPELINE_EXECUTABLE_STATISTIC_FORMAT_INT64_KHR:
+      os << stat.value.i64 << std::endl;
+      break;
+    case VK_PIPELINE_EXECUTABLE_STATISTIC_FORMAT_UINT64_KHR:
+      os << stat.value.u64 << std::endl;
+      break;
+    case VK_PIPELINE_EXECUTABLE_STATISTIC_FORMAT_FLOAT64_KHR:
+      os << stat.value.f64 << std::endl;
+      break;
+    default:
+      break;
+  }
+  os << "    " << stat.description << std::endl;
+  return os;
+}
+
+std::ostream& operator<<(
+    std::ostream& os,
+    std::vector<VkPipelineExecutableStatisticKHR>& shader_stats) {
+  for (int i = 0; i < shader_stats.size(); i++) {
+    VkPipelineExecutableStatisticKHR& stat = shader_stats.at(i);
+    os << stat;
+  }
+  return os;
+}
+
+void Context::print_shader_executable_properties(
+    const vkapi::ShaderInfo& shader,
+    const vkapi::SpecVarList& spec_constants) {
+  VkPipeline pipeline = get_shader_pipeline(shader, spec_constants);
+
+  std::vector<VkPipelineExecutablePropertiesKHR> pipeline_props_list =
+      get_pipeline_executable_props(pipeline);
+
+  VK_CHECK_COND(pipeline_props_list.size() == 1u);
+
+  std::cout << pipeline_props_list.at(0) << std::endl;
+
+  std::tuple<
+      std::vector<VkPipelineExecutableInternalRepresentationKHR>,
+      std::vector<std::vector<char>>>
+      irs_and_irs_data = get_shader_executable_irs(pipeline, 0u);
+
+  std::vector<VkPipelineExecutableInternalRepresentationKHR>& irs =
+      std::get<0>(irs_and_irs_data);
+
+  std::cout << "Found " << irs.size() << " IRs" << std::endl << std::endl;
+  for (int i = 0; i < irs.size(); i++) {
+    std::cout << "====== IR " << i << ": " << irs.at(i).name
+              << " ======" << std::endl;
+    std::cout << irs.at(i) << std::endl;
+  }
+
+  std::vector<VkPipelineExecutableStatisticKHR> shader_stats =
+      get_shader_executable_stats(pipeline, 0u);
+  std::cout << "Found " << shader_stats.size() << " Statistics" << std::endl;
+  if (shader_stats.size() > 0) {
+    std::cout << "====== Statistics: ======" << std::endl;
+    std::cout << shader_stats << std::endl;
+  }
+}
+
+#endif // VK_KHR_pipeline_executable_properties
+
+#endif // VULKAN_DEBUG
+
 } // namespace api
 } // namespace vkcompute

backends/vulkan/runtime/api/Context.h

@@ -228,6 +228,36 @@ class Context final {
       const bool final_use = false);
 
   void flush();
+
+#ifdef VULKAN_DEBUG
+
+#ifdef VK_KHR_pipeline_executable_properties
+
+  VkPipeline get_shader_pipeline(
+      const vkapi::ShaderInfo& shader,
+      const vkapi::SpecVarList& spec_constants);
+
+  std::vector<VkPipelineExecutablePropertiesKHR> get_pipeline_executable_props(
+      const VkPipeline pipeline);
+
+  std::tuple<
+      std::vector<VkPipelineExecutableInternalRepresentationKHR>,
+      std::vector<std::vector<char>>>
+  get_shader_executable_irs(
+      const VkPipeline pipeline,
+      const uint32_t pipeline_exec_idx = 0u);
+
+  std::vector<VkPipelineExecutableStatisticKHR> get_shader_executable_stats(
+      const VkPipeline pipeline,
+      const uint32_t pipeline_exec_idx = 0u);
+
+  void print_shader_executable_properties(
+      const vkapi::ShaderInfo& shader,
+      const vkapi::SpecVarList& spec_constants);
+
+#endif // VK_KHR_pipeline_executable_properties
+
+#endif // VULKAN_DEBUG
 };
 
 bool available();

backends/vulkan/runtime/vk_api/Adapter.cpp

@@ -82,6 +82,9 @@ VkDevice create_logical_device(
 #ifdef VK_KHR_shader_float16_int8
       VK_KHR_SHADER_FLOAT16_INT8_EXTENSION_NAME,
 #endif /* VK_KHR_shader_float16_int8 */
+#if defined(VK_KHR_pipeline_executable_properties) && defined(VULKAN_DEBUG)
+      VK_KHR_PIPELINE_EXECUTABLE_PROPERTIES_EXTENSION_NAME,
+#endif /* VK_KHR_pipeline_executable_properties */
   };
 
   std::vector<const char*> enabled_device_extensions;

backends/vulkan/runtime/vk_api/Pipeline.cpp

@@ -287,10 +287,16 @@ ComputePipeline::ComputePipeline(
       &specialization_info, // pSpecializationInfo
   };
 
+  VkPipelineCreateFlags flags = 0u;
+#if defined(VULKAN_DEBUG) && defined(VK_KHR_pipeline_executable_properties)
+  flags = VK_PIPELINE_CREATE_CAPTURE_STATISTICS_BIT_KHR |
+      VK_PIPELINE_CREATE_CAPTURE_INTERNAL_REPRESENTATIONS_BIT_KHR | flags;
+#endif /* VULKAN_DEBUG && VK_KHR_pipeline_executable_properties */
+
   const VkComputePipelineCreateInfo compute_pipeline_create_info{
       VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO, // sType
       nullptr, // pNext
-      0u, // flags
+      flags, // flags
       shader_stage_create_info, // stage
       descriptor.pipeline_layout, // layout
       VK_NULL_HANDLE, // basePipelineHandle

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -101,6 +101,15 @@ TEST_F(VulkanComputeAPITest, print_adapter) {
   std::cout << *(context()->adapter_ptr()) << std::endl;
 }
 
+#if defined(VULKAN_DEBUG) && defined(VK_KHR_pipeline_executable_properties)
+
+TEST_F(VulkanComputeAPITest, print_shader_executable_properties) {
+  context()->print_shader_executable_properties(
+      VK_KERNEL(binary_add_nobroadcast__test_half), {0});
+}
+
+#endif // VULKAN_DEBUG && VK_KHR_pipeline_executable_properties
+
 std::vector<int64_t> get_reference_strides(
     const std::vector<int64_t>& sizes,
     const utils::GPUMemoryLayout layout,