Skip to content

Commit

Permalink
[ET-VK] Add workgroup sizes to ShaderData
Browse files Browse the repository at this point in the history
Pull Request resolved: #7875

Information is critical for shader optimization.
ghstack-source-id: 262751884
@exported-using-ghexport

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

Co-authored-by: jorgep31415 <jorgepineda140@gmail.com>
  • Loading branch information
pytorchbot and jorgep31415 authored Jan 24, 2025
1 parent 5cbfcdc commit 59aec10
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 2 deletions.
4 changes: 3 additions & 1 deletion backends/vulkan/runtime/VulkanBackend.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -602,7 +602,9 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
event_name.c_str(),
/* delegate_debug_id = */ -1,
r.start_time_ns,
r.end_time_ns);
r.end_time_ns,
(void*)(&r.metadata),
sizeof(r.metadata));
}
#endif // ET_EVENT_TRACER_ENABLED

Expand Down
12 changes: 11 additions & 1 deletion backends/vulkan/runtime/vk_api/QueryPool.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,7 @@ std::string stringize(const VkExtent3D& extents) {
<< "}";
return ss.str();
}

std::vector<ShaderResult> QueryPool::get_shader_timestamp_data() {
if (querypool_ == VK_NULL_HANDLE) {
return {};
Expand All @@ -188,7 +189,16 @@ std::vector<ShaderResult> QueryPool::get_shader_timestamp_data() {
.dispatch_id = entry.dispatch_id,
.start_time_ns = entry.start_time_ns,
.end_time_ns = entry.end_time_ns,
});
.metadata = ShaderMetadata{
.global_workgroup_size =
{entry.global_workgroup_size.width,
entry.global_workgroup_size.height,
entry.global_workgroup_size.depth},
.local_workgroup_size =
{entry.local_workgroup_size.width,
entry.local_workgroup_size.height,
entry.local_workgroup_size.depth},
}});
}
return shader_result;
}
Expand Down
6 changes: 6 additions & 0 deletions backends/vulkan/runtime/vk_api/QueryPool.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,17 @@
namespace vkcompute {
namespace vkapi {

struct ShaderMetadata final {
const uint32_t global_workgroup_size[3];
const uint32_t local_workgroup_size[3];
};

struct ShaderResult final {
const std::string kernel_name;
const uint32_t dispatch_id;
const uint64_t start_time_ns;
const uint64_t end_time_ns;
const ShaderMetadata metadata;
};

struct QueryPoolConfig final {
Expand Down

0 comments on commit 59aec10

Please sign in to comment.