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: Esteban Padilla Cerdio

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/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))])}
+}