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

## Context

Leverage the [VK_KHR_pipeline_executable_properties](https://registry.khronos.org/vulkan/specs/latest/man/html/VK_KHR_pipeline_executable_properties.html) extension to view shader compilation stats and potentially compiled IR that is exposed by the driver. This information can be useful during shader optimization to get more insight into what the compiled shader is actually doing.

Differential Revision: [D69125488](https://our.internmc.facebook.com/intern/diff/D69125488/)

ghstack-source-id: 264624416
Pull Request resolved: #8182

Co-authored-by: Stephen Jia <[email protected]>

Author: pytorchbot

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,