RegCount NITER calculation (pytorch#4159)

Summary: Pull Request resolved: pytorch#4159

Differential Revision: D59405012

Author: estebanpadilla

backends/vulkan/runtime/api/vk_api/QueryPool.cpp

@@ -239,14 +239,20 @@ void QueryPool::print_results() {
 }
 
 unsigned long QueryPool::get_total_shader_ns(std::string kernel_name) {
+  uint64_t total_ns = 0;
+  uint32_t count = 0;
   for (ShaderDuration& entry : shader_durations_) {
     if (entry.kernel_name == kernel_name) {
       std::chrono::duration<size_t, std::nano> exec_duration_ns(
           entry.execution_duration_ns);
-      return exec_duration_ns.count();
+      total_ns += exec_duration_ns.count();
+      count++;
     }
   }
-  return 0;
+  if (count == 0) {
+    return 0;
+  }
+  return total_ns / count;
 }
 } // namespace vkapi
 } // namespace vkcompute

backends/vulkan/runtime/api/vk_api/QueryPool.h

@@ -103,6 +103,10 @@ class QueryPool final {
   void print_results();
   unsigned long get_total_shader_ns(std::string kernel_name);
 
+  uint32_t max_query_count() {
+    return config_.max_query_count;
+  }
+
   operator bool() const {
     return querypool_ != VK_NULL_HANDLE;
   }

backends/vulkan/tools/.gitignore

@@ -0,0 +1 @@
+/ArchProbe/

backends/vulkan/tools/gpuinfo/TARGETS

@@ -0,0 +1,49 @@
+load("@fbcode_macros//build_defs:native_rules.bzl", "buck_filegroup")
+load("@fbsource//tools/build_defs:fb_xplat_cxx_binary.bzl", "fb_xplat_cxx_binary")
+load(
+    "@fbsource//tools/build_defs:platform_defs.bzl",
+    "ANDROID",
+)
+load(
+    "@fbsource//xplat/executorch/backends/vulkan:targets.bzl",
+    "vulkan_spv_shader_lib",
+)
+
+oncall("executorch")
+
+buck_filegroup(
+    name = "gpuinfo_shaders",
+    srcs = glob([
+        "glsl/*",
+    ]),
+    visibility = [
+        "PUBLIC",
+    ],
+)
+
+vulkan_spv_shader_lib(
+    name = "gpuinfo_shader_lib",
+    spv_filegroups = {
+        ":gpuinfo_shaders": "glsl",
+    },
+)
+
+fb_xplat_cxx_binary(
+    name = "vulkan_gpuinfo",
+    srcs = glob([
+        "**/*.cpp",
+    ]),
+    headers = glob([
+        "**/*.h",
+    ]),
+    header_namespace = "/include",
+    include_directories = ["/include"],
+    platforms = ANDROID,
+    raw_headers = glob([
+        "**/*.h",
+    ]),
+    deps = [
+        ":gpuinfo_shader_lib",
+        "//executorch/backends/vulkan:vulkan_graph_runtime",
+    ],
+)

backends/vulkan/tools/gpuinfo/build_gpuinfo.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+# 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.
+
+buck2 build //xplat/executorch/backends/vulkan/tools/gpuinfo:vulkan_gpuinfo --target-platforms=ovr_config//platform/android:arm64-fbsource --show-output -c ndk.static_linking=true -c ndk.debug_info_level=1 -c executorch.event_tracer_enabled=true
+adb push ~/fbsource/buck-out/v2/gen/fbsource/f292370ff0e539e9/xplat/executorch/backends/vulkan/tools/gpuinfo/__vulkan_gpuinfo__/vulkan_gpuinfo /data/local/tmp

backends/vulkan/tools/gpuinfo/glsl/reg_count.glsl

@@ -0,0 +1,36 @@
+/*
+ * 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.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+layout(std430) buffer;
+
+layout(set = 0, binding = 0) buffer PRECISION restrict writeonly Buffer {
+  float data[];
+}
+out_buff;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+layout(constant_id = 3) const int NITER = 1;
+
+void main() {
+
+  ${"\n".join(["float reg_data"+str(i)+" = float(NITER) + "+str(i)+";" for i in range(int(NREG))])}
+
+  int i = 0;
+  for (; i < NITER; ++i) {
+    reg_data0 *= reg_data${int(NREG)-1};
+    ${"\n".join(["reg_data"+str(i)+" *= reg_data"+str(i-1)+";" for i in range(1, int(NREG))])}
+  }
+  i = i >> 31;
+
+  ${"\n".join(["out_buff.data["+str(i)+" * i] = reg_data"+str(i)+";" for i in range(int(NREG))])}
+}