From 1eb9309677a62e7a1d228a9d2338e95f3e0ed995 Mon Sep 17 00:00:00 2001 From: Sicheng Stephen Jia Date: Tue, 4 Feb 2025 14:16:12 -0500 Subject: [PATCH] [ET-VK] Add utility functions to get pipeline executable properties Differential Revision: D69125488 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 f425859935..8178ada3a4 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 0c199c24cc..8bbcf79b45 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 ec30650ba0..4102417436 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 3856d406c2..0c66a085ad 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 604ad26588..84e9826ec8 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,