[ET-VK][runtime] Expose raw shader timestamps and dispatch arguments

backends/vulkan/runtime/api/Context.cpp

@@ -76,13 +76,15 @@ void Context::report_shader_dispatch_start(
     const GlobalWorkGrid& gwg,
     const LocalWorkGroup& lwg,
     const uint32_t dispatch_id) {
-  if (querypool_) {
+  // Dispatch encoding and mode changes hold cmd_mutex_, keeping this flag
+  // unchanged until the paired dispatch end.
+  if (querypool_ && shader_profiling_enabled_) {
     querypool_.shader_profile_begin(cmd_, dispatch_id, shader_name, gwg, lwg);
   }
 }
 
 void Context::report_shader_dispatch_end() {
-  if (querypool_) {
+  if (querypool_ && shader_profiling_enabled_) {
     querypool_.shader_profile_end(cmd_);
   }
 }

backends/vulkan/runtime/api/Context.h

@@ -66,6 +66,7 @@ class Context final {
   vkapi::FencePool fences_;
   // Diagnostics
   vkapi::QueryPool querypool_;
+  bool shader_profiling_enabled_ = true;
   // Command buffers submission
   std::mutex cmd_mutex_;
   vkapi::CommandBuffer cmd_;
@@ -137,6 +138,13 @@ class Context final {
    */
   void initialize_querypool();
 
+  void set_shader_profiling_enabled(bool enabled) {
+    // Call outside dispatch_lock(); mode changes wait for paired timestamp
+    // recording to finish under the command mutex.
+    std::lock_guard<std::mutex> cmd_lock(cmd_mutex_);
+    shader_profiling_enabled_ = enabled;
+  }
+
   /*
    * Encodes a vkResetQueryPool command to the current command buffer, and reset
    * the internal state of the querypool. If the querypool is not initialized

backends/vulkan/runtime/graph/ops/ExecuteNode.h

@@ -82,6 +82,10 @@ class ExecuteNode {
     return name_;
   }
 
+  inline const std::vector<ArgGroup>& args() const {
+    return args_;
+  }
+
  protected:
   uint32_t node_id_;
   const ResizeFunction resize_fn_;

backends/vulkan/runtime/vk_api/QueryPool.cpp

@@ -198,6 +198,45 @@ std::vector<ShaderResult> QueryPool::get_shader_timestamp_data() {
   return shader_result;
 }
 
+std::vector<ShaderTimestamp> QueryPool::get_shader_timestamp_ticks() {
+  std::lock_guard<std::mutex> lock(mutex_);
+  if (querypool_ == VK_NULL_HANDLE || num_queries_ == 0) {
+    return {};
+  }
+  for (const ShaderDuration& entry : shader_durations_) {
+    if (entry.start_query_idx >= num_queries_ ||
+        entry.end_query_idx >= num_queries_) {
+      return {};
+    }
+  }
+  std::vector<uint64_t> ticks(num_queries_);
+  // Without PARTIAL_BIT, any unavailable query makes the whole call NOT_READY.
+  const VkResult status = vkGetQueryPoolResults(
+      device_,
+      querypool_,
+      0u,
+      num_queries_,
+      sizeof(uint64_t) * ticks.size(),
+      ticks.data(),
+      sizeof(uint64_t),
+      VK_QUERY_RESULT_64_BIT);
+  if (status == VK_NOT_READY) {
+    return {};
+  }
+  VK_CHECK(status);
+  std::vector<ShaderTimestamp> results;
+  results.reserve(shader_durations_.size());
+  for (const ShaderDuration& entry : shader_durations_) {
+    results.push_back(
+        {entry.kernel_name,
+         ticks.at(entry.start_query_idx),
+         ticks.at(entry.end_query_idx),
+         {{entry.gwg[0u], entry.gwg[1u], entry.gwg[2u]},
+          {entry.lwg.x(), entry.lwg.y(), entry.lwg.z()}}});
+  }
+  return results;
+}
+
 std::string QueryPool::generate_string_report() {
   std::lock_guard<std::mutex> lock(mutex_);
 

backends/vulkan/runtime/vk_api/QueryPool.h

@@ -27,8 +27,8 @@ namespace vkcompute {
 namespace vkapi {
 
 struct ShaderMetadata final {
-  const uint32_t gwg[3];
-  const uint32_t lwg[3];
+  uint32_t gwg[3];
+  uint32_t lwg[3];
 };
 
 struct ShaderResult final {
@@ -39,6 +39,13 @@ struct ShaderResult final {
   const ShaderMetadata metadata;
 };
 
+struct ShaderTimestamp final {
+  std::string kernel_name;
+  uint64_t start_ticks;
+  uint64_t end_ticks;
+  ShaderMetadata metadata;
+};
+
 struct QueryPoolConfig final {
   uint32_t max_query_count = VULKAN_QUERY_POOL_SIZE;
   uint32_t initial_reserve_size = 256u;
@@ -111,6 +118,9 @@ class QueryPool final {
   void extract_results();
 
   std::vector<ShaderResult> get_shader_timestamp_data();
+  // Read after fence completion; unfinished or unavailable captures return no
+  // samples and never wait for GPU work.
+  std::vector<ShaderTimestamp> get_shader_timestamp_ticks();
   void print_results(const bool tsv_format = false);
   unsigned long get_total_shader_ns(std::string kernel_name);
   unsigned long get_mean_shader_ns(std::string kernel_name);

backends/vulkan/test/api_abstractions/VulkanTest.h

@@ -0,0 +1,26 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+#include <executorch/runtime/platform/runtime.h>
+
+#include <gtest/gtest.h>
+
+namespace vkcompute {
+
+class VulkanTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    executorch::runtime::runtime_init();
+    if (!api::available()) {
+      GTEST_SKIP() << "Vulkan is unavailable on this test runner";
+    }
+  }
+};
+
+} // namespace vkcompute

backends/vulkan/test/api_abstractions/profiling_api_test.cpp

@@ -0,0 +1,180 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+#include <executorch/backends/vulkan/test/api_abstractions/VulkanTest.h>
+#include <executorch/runtime/platform/runtime.h>
+
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <chrono>
+#include <functional>
+#include <future>
+#include <thread>
+#include <vector>
+
+namespace vkcompute {
+namespace {
+
+using ProfilingApiTest = VulkanTest;
+
+TEST(ProfilingValueTest, TimestampVectorsSupportStandardAlgorithms) {
+  const std::vector<vkapi::ShaderTimestamp> samples = {
+      {"first", 1, 2, {{8, 1, 1}, {4, 1, 1}}},
+      {"second", 3, 4, {{16, 1, 1}, {8, 1, 1}}}};
+  auto reordered = samples;
+  std::reverse(reordered.begin(), reordered.end());
+  EXPECT_EQ(reordered.front().kernel_name, "second");
+  reordered.erase(reordered.begin());
+  reordered.insert(reordered.begin(), samples.back());
+  reordered.back() = samples.front();
+  reordered.back().metadata.gwg[0] = 32;
+  EXPECT_EQ(reordered.back().metadata.gwg[0], 32);
+  EXPECT_EQ(samples.front().metadata.gwg[0], 8);
+}
+
+bool timestamps_available(ComputeGraph& graph) {
+  uint32_t family_count = 0;
+  const VkPhysicalDevice physical =
+      graph.context()->adapter_ptr()->physical_handle();
+  vkGetPhysicalDeviceQueueFamilyProperties(physical, &family_count, nullptr);
+  std::vector<VkQueueFamilyProperties> families(family_count);
+  vkGetPhysicalDeviceQueueFamilyProperties(
+      physical, &family_count, families.data());
+  return families.at(graph.context()->queue().family_index).timestampValidBits >
+      0 &&
+      graph.context()->adapter_ptr()->timestamp_period() > 0;
+}
+
+void disable_profiling(
+    api::Context& context,
+    std::promise<void>& started,
+    std::promise<void>& completed) {
+  started.set_value();
+  context.set_shader_profiling_enabled(false);
+  completed.set_value();
+}
+
+TEST_F(ProfilingApiTest, EmptyQueryPoolHasNoTimestampTicks) {
+  ComputeGraph graph{GraphConfig{}};
+  EXPECT_TRUE(
+      graph.context()->querypool().get_shader_timestamp_ticks().empty());
+}
+
+TEST_F(ProfilingApiTest, ProfilingModeChangesWaitForCommandRecording) {
+  ComputeGraph graph{GraphConfig{}};
+  auto& context = *graph.context();
+  std::promise<void> started;
+  std::promise<void> completed;
+  auto entered = started.get_future();
+  auto finished = completed.get_future();
+  auto recording = context.dispatch_lock();
+  std::thread worker(
+      disable_profiling,
+      std::ref(context),
+      std::ref(started),
+      std::ref(completed));
+  entered.wait();
+  const auto status = finished.wait_for(std::chrono::milliseconds(100));
+  recording.unlock();
+  worker.join();
+  EXPECT_EQ(status, std::future_status::timeout);
+}
+
+TEST_F(ProfilingApiTest, UnfinishedAndUnsubmittedCapturesHaveNoTimestampTicks) {
+  ComputeGraph probe{GraphConfig{}};
+  if (!timestamps_available(probe)) {
+    GTEST_SKIP() << "GPU timestamps are unavailable";
+  }
+  for (const bool complete : {false, true}) {
+    SCOPED_TRACE(complete);
+    GraphConfig config;
+    config.enable_querypool = true;
+    ComputeGraph graph(config);
+    auto& context = *graph.context();
+    context.initialize_querypool();
+    {
+      auto resetting = context.dispatch_lock();
+      context.set_cmd();
+      context.cmd_reset_querypool();
+      vkapi::VulkanFence fence(context.device());
+      context.submit_cmd_to_gpu(fence.get_submit_handle());
+      fence.wait();
+    }
+    context.flush();
+    auto recording = context.dispatch_lock();
+    context.set_cmd();
+    context.report_shader_dispatch_start(
+        "unsubmitted",
+        GlobalWorkGrid({1, 1, 1}, kTextureExtentsWorkGrid),
+        LocalWorkGroup(1, 1, 1));
+    if (complete) {
+      context.report_shader_dispatch_end();
+    }
+    EXPECT_TRUE(context.querypool().get_shader_timestamp_ticks().empty());
+  }
+}
+
+TEST_F(ProfilingApiTest, TimestampTogglePreservesExecutionAndArgumentMetadata) {
+  executorch::runtime::runtime_init();
+  ComputeGraph probe{GraphConfig{}};
+  const bool has_timestamps = timestamps_available(probe);
+  for (const bool enabled : {true, false, true}) {
+    SCOPED_TRACE(enabled);
+    GraphConfig config;
+    config.enable_querypool = has_timestamps;
+    config.execute_threshold_node_count = 64;
+    ComputeGraph graph(config);
+    graph.context()->set_shader_profiling_enabled(enabled);
+    const auto a =
+        graph.add_input_tensor({1, 1, 1, 16}, vkapi::kFloat, utils::kBuffer);
+    const auto b =
+        graph.add_input_tensor({1, 1, 1, 16}, vkapi::kFloat, utils::kBuffer);
+    const auto out =
+        graph.add_tensor({1, 1, 1, 16}, vkapi::kFloat, utils::kBuffer);
+    VK_GET_OP_FN("aten.add.Tensor")
+    (graph, {a.value, b.value, graph.add_scalar(1.0), out});
+    const auto& args = graph.execute_nodes().back()->args();
+    ASSERT_EQ(args.size(), 2);
+    EXPECT_EQ(args[0].refs, (std::vector<ValueRef>{out}));
+    EXPECT_EQ(args[0].access, vkapi::kWrite);
+    EXPECT_EQ(args[1].refs, (std::vector<ValueRef>{a.value, b.value}));
+    EXPECT_EQ(args[1].access, vkapi::kRead);
+    const auto ostage = graph.set_output_tensor(out);
+    graph.prepare();
+    graph.prepack();
+    graph.propagate_resize();
+    std::vector<float> input(16, 0.25f);
+    graph.maybe_cast_and_copy_into_staging(
+        a.staging, input.data(), input.size(), vkapi::kFloat);
+    graph.maybe_cast_and_copy_into_staging(
+        b.staging, input.data(), input.size(), vkapi::kFloat);
+    graph.execute();
+    std::vector<float> actual(16);
+    graph.maybe_cast_and_copy_from_staging(
+        ostage, actual.data(), actual.size(), vkapi::kFloat);
+    EXPECT_EQ(actual, std::vector<float>(16, 0.5f));
+    const auto ticks =
+        graph.context()->querypool().get_shader_timestamp_ticks();
+    if (enabled && has_timestamps) {
+      ASSERT_EQ(ticks.size(), graph.execute_nodes().size());
+      const auto target =
+          std::find_if(ticks.begin(), ticks.end(), [](const auto& entry) {
+            return entry.kernel_name.find("binary_add_") == 0;
+          });
+      ASSERT_NE(target, ticks.end());
+      EXPECT_GE(target->end_ticks, target->start_ticks);
+      EXPECT_GT(target->metadata.gwg[0], 0);
+      EXPECT_GT(target->metadata.lwg[0], 0);
+    } else {
+      EXPECT_TRUE(ticks.empty());
+    }
+  }
+}
+
+} // namespace
+} // namespace vkcompute

backends/vulkan/test/api_abstractions/targets.bzl

@@ -5,6 +5,28 @@ def define_common_targets(is_fbcode = False):
     if is_fbcode:
         return
 
+    runtime.cxx_library(
+        name = "test_utils",
+        exported_headers = ["VulkanTest.h"],
+        platforms = get_platforms(),
+        exported_deps = [
+            "//executorch/backends/vulkan:vulkan_graph_runtime",
+            "//third-party/googletest:gtest_main",
+        ],
+        visibility = ["PUBLIC"],
+    )
+
+    runtime.cxx_test(
+        name = "profiling_api_test",
+        srcs = ["profiling_api_test.cpp"],
+        platforms = get_platforms(),
+        deps = [
+            ":test_utils",
+            "//executorch/backends/vulkan:vulkan_graph_runtime",
+            "//third-party/googletest:gtest_main",
+        ],
+    )
+
     runtime.cxx_test(
         name = "dispatch_grid_test",
         srcs = ["dispatch_grid_test.cpp"],