RegCount NITER calculation (#4159)

Summary:
Pull Request resolved: #4159

This adds an internal implementation of https://github.com/microsoft/ArchProbe.

This stack introduces a kernel that can be used to get the number of available registers on a mobile GPU by gradually increasing the number of accessed elements and detecting dramatic drops in performance. See [this paper ](https://www.microsoft.com/en-us/research/uploads/prod/2022/02/mobigpu_mobicom22_camera.pdf), page 4, for more information.

This first diff gets the number of iterations (NITER) that can run in 1000us, to be used in the following tests.

The kernel looks like the following for any K number of registers:

  float reg_data0 = float(niter) + 0;
  float reg_data1 = float(niter) + 1;
  ...
  float reg_dataK = float(niter) + K;

  int i = 0;
  for (; i < niter; ++i) {
    reg_data0 *= reg_dataK;
    reg_data1 *= reg_data0;
    reg_data2 *= reg_data1;
    ...
    reg_dataK *= reg_data(K-1);
  }

  i = i >> 31;

  buffer_out.data[0 * i] = reg_data0;
  buffer_out.data[1 * i] = reg_data1;
  ...
  buffer_out.data[K * i] = reg_dataK;

Differential Revision: D59405012

Author: estebanpadilla

backends/vulkan/runtime/vk_api/QueryPool.cpp

@@ -248,5 +248,22 @@ unsigned long QueryPool::get_total_shader_ns(std::string kernel_name) {
   }
   return 0;
 }
+
+unsigned long QueryPool::get_mean_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);
+      total_ns += exec_duration_ns.count();
+      count++;
+    }
+  }
+  if (count == 0) {
+    return 0;
+  }
+  return total_ns / count;
+}
 } // namespace vkapi
 } // namespace vkcompute

backends/vulkan/runtime/vk_api/QueryPool.h

@@ -102,6 +102,7 @@ class QueryPool final {
   std::string generate_string_report();
   void print_results();
   unsigned long get_total_shader_ns(std::string kernel_name);
+  unsigned long get_mean_shader_ns(std::string kernel_name);
 
   operator bool() const {
     return querypool_ != VK_NULL_HANDLE;

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/glsl/reg_count.glsl

@@ -0,0 +1,39 @@
+/*
+ * 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() {
+
+  $for k in range(int(NREG)):
+     float reg_data${k} = float(NITER) + ${k};
+
+  int i = 0;
+  for (; i < NITER; ++i) {
+    reg_data0 *= reg_data${int(NREG)-1};
+    $for k in range(1, int(NREG)):
+      reg_data${k} *= reg_data${k-1};
+  }
+  i = i >> 31;
+
+  $for k in range(int(NREG)):
+    out_buff.data[${k} * i] = reg_data${k};
+}