From d9ab5dc24185fbeed768865601985d5dd5fdaa5f Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Tue, 4 Feb 2025 09:06:49 -0800 Subject: [PATCH] [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: https://github.com/pytorch/executorch/pull/8182 --- backends/vulkan/runtime/api/Context.cpp | 220 ++++++++++++++++++ backends/vulkan/runtime/api/Context.h | 30 +++ backends/vulkan/runtime/vk_api/Adapter.cpp | 3 + backends/vulkan/runtime/vk_api/Pipeline.cpp | 8 +- .../vulkan/test/vulkan_compute_api_test.cpp | 9 + 5 files changed, 269 insertions(+), 1 deletion(-) diff --git a/backends/vulkan/runtime/api/Context.cpp b/backends/vulkan/runtime/api/Context.cpp index f425859935d..8178ada3a45 100644 --- a/backends/vulkan/runtime/api/Context.cpp +++ b/backends/vulkan/runtime/api/Context.cpp @@ -8,6 +8,11 @@ #include +#ifdef VULKAN_DEBUG +#include +#include +#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 +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 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, + std::vector>> +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 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> 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 +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 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& 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 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, + std::vector>> + irs_and_irs_data = get_shader_executable_irs(pipeline, 0u); + + std::vector& 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 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 diff --git a/backends/vulkan/runtime/api/Context.h b/backends/vulkan/runtime/api/Context.h index 0c199c24cc4..8bbcf79b45c 100644 --- a/backends/vulkan/runtime/api/Context.h +++ b/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 get_pipeline_executable_props( + const VkPipeline pipeline); + + std::tuple< + std::vector, + std::vector>> + get_shader_executable_irs( + const VkPipeline pipeline, + const uint32_t pipeline_exec_idx = 0u); + + std::vector 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(); diff --git a/backends/vulkan/runtime/vk_api/Adapter.cpp b/backends/vulkan/runtime/vk_api/Adapter.cpp index ec30650ba06..4102417436c 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.cpp +++ b/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 enabled_device_extensions; diff --git a/backends/vulkan/runtime/vk_api/Pipeline.cpp b/backends/vulkan/runtime/vk_api/Pipeline.cpp index 3856d406c24..0c66a085ad9 100644 --- a/backends/vulkan/runtime/vk_api/Pipeline.cpp +++ b/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 diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 604ad26588d..84e9826ec87 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/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 get_reference_strides( const std::vector& sizes, const utils::GPUMemoryLayout layout,