Update on "[ET-VK] Introduce graph runtime shader library that enables dynamic shapes"

## Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
  ivec4 output_sizes;
  ivec4 other_sizes;
  float alpha;
}
```

Shaders will accept separate UBOs for each piece of tensor metadata:

```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
  ivec4 data;
}
out_sizes;

layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
  ivec4 data;
}
in_sizes;

layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
  ivec4 data;
}
other_sizes;

layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
  float data;
}
alpha;
```

Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.


## Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.

Constructing an `ExecuteNode` for binary operators is now

```
  graph.execute_nodes().emplace_back(new ExecuteNode(
      graph,
      api::shader_registry().get_shader_info(kernel_name.str()),
      global_size,
      local_size,
      {{out, api::MemoryAccessType::WRITE},
       {{arg1, arg2}, api::MemoryAccessType::READ}},
      {t_out.gpu_sizes_ubo(),
       t_in1.gpu_sizes_ubo(),
       t_in2.gpu_sizes_ubo(),
       graph.create_params_buffer(alpha_val)}))
```

instead of

```
ArithmeticParams block{
      get_size_as_ivec4(t_out),
      get_size_as_ivec4(t_in1),
      get_size_as_ivec4(t_in2),
      alpha_val,
  };
  api::UniformParamsBuffer params(graph.context(), block);

  graph.execute_nodes().emplace_back(new ExecuteNode(
      graph,
      shader,
      global_size,
      local_size,
      {{out, api::MemoryAccessType::WRITE},
       {{arg1, arg2}, api::MemoryAccessType::READ}},
      std::move(params)));
```

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/targets.bzl

@@ -37,6 +37,7 @@ def vulkan_spv_shader_lib(name, spv_filegroups, is_fbcode = False):
         srcs = [
             ":{}[{}.cpp]".format(genrule_name, name),
         ],
+        define_static_target = False,
         # 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`

examples/models/llama2/runner/targets.bzl

@@ -6,7 +6,7 @@ def _get_operator_lib(aten = False):
     elif runtime.is_oss:
         return ["//executorch/kernels/portable:generated_lib"]
     else:
-        return ["//executorch/kernels/portable:generated_lib", "//executorch/examples/models/llama2/custom_ops:custom_ops", "//executorch/examples/models/llama2/ops:generated_lib"]
+        return ["//executorch/configurations:optimized_native_cpu_ops", "//executorch/examples/models/llama2/custom_ops:custom_ops", "//executorch/examples/models/llama2/ops:generated_lib"]
 
 def define_common_targets():
     for aten in (True, False):

kernels/portable/test/targets.bzl

@@ -8,7 +8,7 @@ def define_common_targets():
     """
     define_supported_features_lib()
 
-    op_test(name = "op_allclose_test", aten_compatible = False)
+    op_test(name = "op_allclose_test")
     op_test(name = "op_div_test")
     op_test(name = "op_gelu_test")
     op_test(name = "op_mul_test")

kernels/test/targets.bzl

@@ -1,14 +1,14 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 load("@fbsource//xplat/executorch/kernels/test:util.bzl", "codegen_function_header_wrapper", "generated_op_test", "op_test")
 
-def _common_op_test(name, kernels, aten_compatible = True):
+def _common_op_test(name, kernels):
     """
     Defines test targets in format of <kernel>_op_<op-name>_test
     For ATen kernel testing, let's use portable functions.yaml for tested ops.
     """
     for kernel in kernels:
         deps = [":function_header_wrapper_{}".format(kernel)]
-        op_test(name, aten_compatible = aten_compatible, kernel_name = kernel, use_kernel_prefix = True, deps = deps)
+        op_test(name, kernel_name = kernel, use_kernel_prefix = True, deps = deps)
 
 def make_example_generated_op_test_target():
     """

kernels/test/util.bzl

@@ -1,7 +1,7 @@
 load("@fbsource//tools/build_defs:fbsource_utils.bzl", "is_xplat")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
-def op_test(name, deps = [], aten_compatible = True, kernel_name = "portable", use_kernel_prefix = False):
+def op_test(name, deps = [], kernel_name = "portable", use_kernel_prefix = False):
     """Defines a cxx_test() for an "op_*_test.cpp" file.
 
     Args:
@@ -11,8 +11,6 @@ def op_test(name, deps = [], aten_compatible = True, kernel_name = "portable", u
             under //kernels/<kernel>/...; e.g., "op_add_test" will depend on
             "//kernels/portable/cpu:op_add".
         deps: Optional extra deps to add to the cxx_test().
-        aten_compatible: If True, the operator under test is ATen-compatible
-            (i.e., appears in `functions.yaml`).
         kernel_name: The name string as in //executorch/kernels/<kernel_name>.
         use_kernel_prefix: If True, the target name is
             <kernel>_op_<operator-group-name>_test. Used by common kernel testing.