[ATen-vulkan] Implement global shader registry (pytorch#121088)

Summary:

X-link: pytorch/executorch#2222

## Context

This changeset updates Vulkan SPIR-V codegen to introduce a global SPIR-V shader registry and register shaders dynamically at static initialization time. This change makes it possible to define and link custom shader libraries to the ATen-Vulkan runtime.

Before:

* `gen_vulkan_spv.py` generated two files, `spv.h` and `spv.cpp` which would contain the definition and initialization of Vulkan shader registry variables. 


After:

* Introduce the `ShaderRegistry` class in `api/`, which encapsulates functionality of the `ShaderRegistry` class previously defined in the generated `spv.h` file
* Introduce a global shader registry (defined as a static variable in the `api::shader_registry() function`
* Define a `ShaderRegisterInit` class (taking inspiration from `TorchLibraryInit`) that allows for dynamic shader registration
* `gen_vulkan_spv.py` now only generates `spv.cpp`, which defines a static `ShaderRegisterInit` instance that triggers registration of the compiled shaders to the global shader registry.

Benefits:

* Cleaner code base; we no longer have `ShaderRegistry` defined in a generated file, and don't need a separate implementation file (`impl/Registry.*`) to handle shader lookup. All that logic now lives under `api/ShaderRegistry.*`
* Makes it possible to compile and link separate shader libraries, providing similar flexibility as defining and linking custom ATen operators

Test Plan: CI

Differential Revision: D54447700

Author: SS-JIA

aten/src/ATen/native/vulkan/api/ShaderRegistry.cpp

@@ -0,0 +1,61 @@
+#include <ATen/native/vulkan/api/ShaderRegistry.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+namespace api {
+
+bool ShaderRegistry::has_shader(const std::string& shader_name) {
+  const ShaderListing::const_iterator it = listings_.find(shader_name);
+  return it != listings_.end();
+}
+
+bool ShaderRegistry::has_dispatch(const std::string& op_name) {
+  const Registry::const_iterator it = registry_.find(op_name);
+  return it != registry_.end();
+}
+
+void ShaderRegistry::register_shader(ShaderInfo&& shader_info) {
+  if (has_shader(shader_info.kernel_name)) {
+    VK_THROW(
+        "Shader with name ", shader_info.kernel_name, "already registered");
+  }
+  listings_.emplace(shader_info.kernel_name, shader_info);
+}
+
+void ShaderRegistry::register_op_dispatch(
+    const std::string& op_name,
+    const DispatchKey key,
+    const std::string& shader_name) {
+  if (!has_dispatch(op_name)) {
+    registry_.emplace(op_name, Dispatcher());
+  }
+  const Dispatcher::const_iterator it = registry_[op_name].find(key);
+  if (it != registry_[op_name].end()) {
+    registry_[op_name][key] = shader_name;
+  } else {
+    registry_[op_name].emplace(key, shader_name);
+  }
+}
+
+const ShaderInfo& ShaderRegistry::get_shader_info(
+    const std::string& shader_name) {
+  const ShaderListing::const_iterator it = listings_.find(shader_name);
+
+  VK_CHECK_COND(
+      it != listings_.end(),
+      "Could not find ShaderInfo with name ",
+      shader_name);
+
+  return it->second;
+}
+
+ShaderRegistry& shader_registry() {
+  static ShaderRegistry registry;
+  return registry;
+}
+
+} // namespace api
+} // namespace vulkan
+} // namespace native
+} // namespace at

aten/src/ATen/native/vulkan/api/ShaderRegistry.h

@@ -0,0 +1,84 @@
+#pragma once
+
+// @lint-ignore-every CLANGTIDY facebook-hte-BadMemberName
+
+#ifdef USE_VULKAN_API
+
+#include <ATen/native/vulkan/api/Shader.h>
+
+#include <string>
+#include <unordered_map>
+
+#define VK_KERNEL(shader_name) \
+  ::at::native::vulkan::api::shader_registry().get_shader_info(#shader_name)
+
+namespace at {
+namespace native {
+namespace vulkan {
+namespace api {
+
+enum class DispatchKey : int8_t {
+  CATCHALL,
+  ADRENO,
+  MALI,
+  OVERRIDE,
+};
+
+class ShaderRegistry final {
+  using ShaderListing = std::unordered_map<std::string, ShaderInfo>;
+  using Dispatcher = std::unordered_map<DispatchKey, std::string>;
+  using Registry = std::unordered_map<std::string, Dispatcher>;
+
+  ShaderListing listings_;
+  Dispatcher dispatcher_;
+  Registry registry_;
+
+ public:
+  /*
+   * Check if the registry has a shader registered under the given name
+   */
+  bool has_shader(const std::string& shader_name);
+
+  /*
+   * Check if the registry has a dispatch registered under the given name
+   */
+  bool has_dispatch(const std::string& op_name);
+
+  /*
+   * Register a ShaderInfo to a given shader name
+   */
+  void register_shader(ShaderInfo&& shader_info);
+
+  /*
+   * Register a dispatch entry to the given op name
+   */
+  void register_op_dispatch(
+      const std::string& op_name,
+      const DispatchKey key,
+      const std::string& shader_name);
+
+  /*
+   * Given a shader name, return the ShaderInfo which contains the SPIRV binary
+   */
+  const ShaderInfo& get_shader_info(const std::string& shader_name);
+};
+
+class ShaderRegisterInit final {
+  using InitFn = void();
+
+ public:
+  ShaderRegisterInit(InitFn* init_fn) {
+    init_fn();
+  };
+};
+
+// The global shader registry is retrieved using this function, where it is
+// declared as a static local variable.
+ShaderRegistry& shader_registry();
+
+} // namespace api
+} // namespace vulkan
+} // namespace native
+} // namespace at
+
+#endif /* USE_VULKAN_API */

aten/src/ATen/native/vulkan/api/api.h

@@ -10,6 +10,7 @@
 #include <ATen/native/vulkan/api/Resource.h>
 #include <ATen/native/vulkan/api/Runtime.h>
 #include <ATen/native/vulkan/api/Shader.h>
+#include <ATen/native/vulkan/api/ShaderRegistry.h>
 #include <ATen/native/vulkan/api/Tensor.h>
 #include <ATen/native/vulkan/api/Utils.h>
 

aten/src/ATen/native/vulkan/impl/Common.h

@@ -3,12 +3,6 @@
 #ifdef USE_VULKAN_API
 
 #include <ATen/native/vulkan/api/api.h>
-#include <ATen/native/vulkan/impl/Registry.h>
-
-#define VK_KERNEL(shader_name) \
-  ::at::native::vulkan::get_shader_info(#shader_name)
-#define VK_LOOKUP_KERNEL(op_name) \
-  ::at::native::vulkan::look_up_shader_info(#op_name)
 
 namespace at {
 namespace native {

aten/src/ATen/native/vulkan/impl/Registry.cpp

@@ -10,11 +10,6 @@
 #include <ATen/native/vulkan/impl/Common.h>
 #include <ATen/native/vulkan/ops/Convert.h>
 
-#define VK_KERNEL(shader_name) \
-  ::at::native::vulkan::get_shader_info(#shader_name)
-#define VK_LOOKUP_KERNEL(op_name) \
-  ::at::native::vulkan::look_up_shader_info(#op_name)
-
 namespace at {
 namespace native {
 namespace vulkan {

aten/src/ATen/native/vulkan/impl/Registry.h

@@ -319,7 +319,7 @@ static api::ShaderInfo get_shader(
 
   switch (method) {
     case Conv2dSlidingWindow:
-      shader = VK_LOOKUP_KERNEL(conv2d);
+      shader = VK_KERNEL(conv2d);
       break;
     case Conv2dDepthwise:
       shader = VK_KERNEL(conv2d_dw);
@@ -335,7 +335,7 @@ static api::ShaderInfo get_shader(
       }
       break;
     case Conv2dPointwise:
-      shader = VK_LOOKUP_KERNEL(conv2d_pw);
+      shader = VK_KERNEL(conv2d_pw_output_tile_2x2);
       break;
   }
   return shader;

aten/src/ATen/native/vulkan/ops/Common.h

@@ -144,6 +144,9 @@ def get_glsl_paths():
         ],
     )
 
+def spv_shader_library():
+    pass
+
 # @lint-ignore BUCKRESTRICTEDSYNTAX
 IS_OSS = read_config("pt", "is_oss", "0") == "1"  # True for OSS BUCK build, and False for internal BUCK build
 
@@ -700,6 +703,43 @@ def gen_aten_libtorch_files(name, extra_params = [], compatible_with = [], apple
         apple_sdks = apple_sdks,
     )
 
+def vulkan_spv_shader_library(name, spv_filegroup):
+    genrule_cmd = [
+        "$(exe //xplat/caffe2/tools:gen_aten_vulkan_spv_bin)",
+        "--glsl-paths $(location {})".format(spv_filegroup),
+        "--output-path $OUT --env FLOAT_IMAGE_FORMAT={}".format(get_glsl_image_format()),
+        "--glslc-path=$(exe //xplat/caffe2/fb/vulkan/dotslash:glslc)",
+        "--tmp-dir-path=$TMP",
+    ]
+
+    genrule_name = "gen_{}_cpp".format(name)
+    fb_xplat_genrule(
+        name = "gen_{}_cpp".format(name),
+        outs = {
+            "{}.cpp".format(name): ["spv.cpp"],
+        },
+        cmd = " ".join(genrule_cmd),
+        default_outs = ["."],
+        labels = ["uses_dotslash"],
+    )
+
+    fb_xplat_cxx_library(
+        name = name,
+        srcs = [
+            ":{}[{}.cpp]".format(genrule_name, name),
+        ],
+        # Static initialization is used to register shaders to the global shader registry,
+        # therefore link_whole must be True to make sure unused symbols are not discarded.
+        # @lint-ignore BUCKLINT: Avoid `link_whole=True`
+        link_whole = True,
+        # Define a soname that can be used for dynamic loading in Java, Python, etc.
+        soname = "lib{}.$(ext)".format(name),
+        visibility = ["PUBLIC"],
+        exported_deps = [
+            "//xplat/caffe2:torch_vulkan_api",
+        ],
+    )
+
 def copy_metal(name, apple_sdks = None):
     cmd = []
     cmd_exe = []