Update on "[ET-VK] Fix OSS build + separate test build into its own CMakeLists.txt"

## Context

As title. In the next diff, the vulkan test binary will be added to CI.

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

[ghstack-poisoned]

Author: SS-JIA

.ci/scripts/setup-macos.sh

@@ -28,15 +28,18 @@ install_buck() {
   fi
 
   pushd .ci/docker
-
   # TODO(huydo): This is a one-off copy of buck2 2024-05-15 to unblock Jon and
   # re-enable ShipIt. It’s not ideal that upgrading buck2 will require a manual
   # update the cached binary on S3 bucket too. Let me figure out if there is a
   # way to correctly implement the previous setup of installing a new version of
   # buck2 only when it’s needed. AFAIK, the complicated part was that buck2
   # --version doesn't say anything w.r.t its release version, i.e. 2024-05-15.
   # See D53878006 for more details.
-  BUCK2=buck2-aarch64-apple-darwin.zst
+  #
+  # If you need to upgrade buck2 version on S3, please reach out to Dev Infra
+  # team for help.
+  BUCK2_VERSION=$(cat ci_commit_pins/buck2.txt)
+  BUCK2=buck2-aarch64-apple-darwin-${BUCK2_VERSION}.zst
   curl -s "https://ossci-macos.s3.amazonaws.com/${BUCK2}" -o "${BUCK2}"
 
   zstd -d "${BUCK2}" -o buck2

backends/vulkan/main.cpp

@@ -52,12 +52,15 @@ def __contains__(self, op):
 UNARY_OPS = [
     exir_ops.edge.aten.abs.default,
     exir_ops.edge.aten.clamp.default,
+    exir_ops.edge.aten.cos.default,
     exir_ops.edge.aten.exp.default,
     exir_ops.edge.aten.gelu.default,
     exir_ops.edge.aten.hardshrink.default,
     exir_ops.edge.aten.hardtanh.default,
+    exir_ops.edge.aten.neg.default,
     exir_ops.edge.aten.relu.default,
     exir_ops.edge.aten.sigmoid.default,
+    exir_ops.edge.aten.sin.default,
     exir_ops.edge.aten.sqrt.default,
     exir_ops.edge.aten.tanh.default,
 ]
@@ -84,6 +87,7 @@ def __contains__(self, op):
 ]
 
 NORMALIZATION_OPS = [
+    exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
     exir_ops.edge.aten.native_layer_norm.default,
 ]
 

backends/vulkan/partitioner/supported_ops.py

@@ -18,12 +18,16 @@
 #include <executorch/backends/vulkan/runtime/api/Command.h>
 #include <executorch/backends/vulkan/runtime/api/Pipeline.h>
 
+#ifndef VULKAN_QUERY_POOL_SIZE
+#define VULKAN_QUERY_POOL_SIZE 4096u
+#endif
+
 namespace vkcompute {
 namespace api {
 
 struct QueryPoolConfig final {
-  uint32_t max_query_count;
-  uint32_t initial_reserve_size;
+  uint32_t max_query_count = VULKAN_QUERY_POOL_SIZE;
+  uint32_t initial_reserve_size = 256u;
 };
 
 struct ShaderDuration final {

backends/vulkan/runtime/api/QueryPool.h

@@ -18,12 +18,18 @@ unary_op:
     - NAME: clamp_int
       OPERATOR: clamp(X, A, B)
       DTYPE: int
+    - NAME: cos
+      OPERATOR: cos(X)
     - NAME: exp
       OPERATOR: exp(X)
     - NAME: gelu
       OPERATOR: 0.5 * X * (1 + tanh(sqrt(2 / 3.141593) * (X + 0.044715 * X * X * X)))
+    - NAME: neg
+      OPERATOR: -X
     - NAME: sigmoid
       OPERATOR: 1 / (1 + exp(-1 * X))
+    - NAME: sin
+      OPERATOR: sin(X)
     - NAME: sqrt
       OPERATOR: sqrt(X)
     - NAME: tanh

backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml

@@ -0,0 +1,43 @@
+/*
+ * 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
+
+#include "broadcasting_utils.h"
+#include "indexing_utils.h"
+
+#define PRECISION ${PRECISION}
+
+#define VEC4_T ${texel_type(DTYPE)}
+
+layout(std430) buffer;
+
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
+
+layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
+
+layout(set = 0, binding = 2) uniform PRECISION restrict OutLimits {
+  ivec3 out_limits;
+};
+
+layout(set = 0, binding = 3) uniform PRECISION restrict Sizes {
+  ivec4 sizes;
+};
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  if (any(greaterThanEqual(pos, out_limits))) {
+    return;
+  }
+
+  VEC4_T in_texel = texelFetch(image_in, pos, 0);
+  imageStore(image_out, pos, in_texel);
+}

backends/vulkan/runtime/graph/ops/glsl/upsample.glsl

@@ -0,0 +1,17 @@
+# 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.
+
+upsample:
+  parameter_names_with_default_values:
+    NDIM: 3
+    DTYPE: float
+    PACKING: C_packed
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: upsample

backends/vulkan/runtime/graph/ops/glsl/upsample.yaml

@@ -122,8 +122,11 @@ void gelu(ComputeGraph& graph, const std::vector<ValueRef>& args) {
 }
 
 DEFINE_ACTIVATION_FN(abs);
+DEFINE_ACTIVATION_FN(cos);
 DEFINE_ACTIVATION_FN(exp);
+DEFINE_ACTIVATION_FN(neg);
 DEFINE_ACTIVATION_FN(sigmoid);
+DEFINE_ACTIVATION_FN(sin);
 DEFINE_ACTIVATION_FN(sqrt);
 DEFINE_ACTIVATION_FN(tanh);
 DEFINE_CLAMP_FN(clamp);
@@ -134,11 +137,14 @@ DEFINE_HARDSHRINK_FN(hardshrink);
 REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.abs.default, abs);
   VK_REGISTER_OP(aten.clamp.default, clamp);
+  VK_REGISTER_OP(aten.cos.default, cos);
   VK_REGISTER_OP(aten.exp.default, exp);
   VK_REGISTER_OP(aten.gelu.default, gelu);
   VK_REGISTER_OP(aten.hardtanh.default, hardtanh);
+  VK_REGISTER_OP(aten.neg.default, neg);
   VK_REGISTER_OP(aten.relu.default, relu);
   VK_REGISTER_OP(aten.sigmoid.default, sigmoid);
+  VK_REGISTER_OP(aten.sin.default, sin);
   VK_REGISTER_OP(aten.sqrt.default, sqrt);
   VK_REGISTER_OP(aten.tanh.default, tanh);
   VK_REGISTER_OP(aten.hardshrink.default, hardshrink);

backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp

@@ -0,0 +1,67 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+void resize_upsample_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& extra_args) {
+  (void)graph;
+  (void)args;
+  (void)extra_args;
+}
+
+void add_upsample_node(
+    ComputeGraph& graph,
+    const ValueRef in,
+    const ValueRef out) {
+  ValueRef arg = prepack_if_tensor_ref(graph, in);
+
+  vTensorPtr t_out = graph.get_tensor(out);
+  api::utils::uvec3 global_size = t_out->image_extents();
+  api::utils::uvec3 local_size = adaptive_work_group_size(global_size);
+
+  std::string kernel_name("upsample");
+  kernel_name.reserve(kShaderNameReserve);
+
+  add_dtype_suffix(kernel_name, *t_out);
+
+  graph.execute_nodes().emplace_back(new ExecuteNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      global_size,
+      local_size,
+      // Inputs and Outputs
+      {{out, api::MemoryAccessType::WRITE}, {arg, api::MemoryAccessType::READ}},
+      // Shader params buffers
+      {t_out->texture_limits_ubo(), graph.create_params_buffer(0.5)},
+      // Specialization Constants
+      {},
+      // Resizing Logic
+      resize_upsample_node));
+}
+
+void upsample(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  return add_upsample_node(graph, args[0], args[3]);
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(aten.upsample_nearest2d.vec, upsample);
+}
+
+} // namespace vkcompute

backends/vulkan/runtime/graph/ops/impl/Upsample.cpp

@@ -241,6 +241,16 @@ def get_native_layer_norm_inputs():
     return test_suite
 
 
+def get_upsample_inputs():
+    test_suite = VkTestSuite(
+        [
+            # TODO(dixu): implement the basic upsample logic to have a meaningful test
+            ((2, 2, 2, 2), None, [1, 1]),
+        ]
+    )
+    return test_suite
+
+
 def get_full_inputs():
     test_suite = VkTestSuite(
         [
@@ -672,6 +682,8 @@ def get_unary_ops_inputs():
         ]
     )
     test_suite.storage_types = ["api::kTexture3D", "api::kBuffer"]
+    test_suite.atol = "1e-4"
+    test_suite.rtol = "1e-4"
     return test_suite
 
 
@@ -796,4 +808,8 @@ def get_gelu_inputs():
     "aten._native_batch_norm_legit_no_training.default": get_native_batch_norm_inputs(),
     "aten.gelu.default": get_gelu_inputs(),
     "aten.hardshrink.default": get_unary_ops_inputs(),
+    "aten.upsample_nearest2d.vec": get_upsample_inputs(),
+    "aten.sin.default": get_unary_ops_inputs(),
+    "aten.neg.default": get_unary_ops_inputs(),
+    "aten.cos.default": get_unary_ops_inputs(),
 }

backends/vulkan/test/op_tests/cases.py

@@ -17,6 +17,8 @@
     CppTestFileGen,
     DOUBLE,
     INT,
+    OPT_AT_DOUBLE_ARRAY_REF,
+    OPT_AT_INT_ARRAY_REF,
     OPT_AT_TENSOR,
     OPT_BOOL,
     OPT_DEVICE,
@@ -289,6 +291,16 @@ def create_value_for(self, ref: ValueRefList) -> str:  # noqa: C901
             ret_str += f"{self.graph}{self.dot}add_scalar<int64_t>"
             ret_str += f"({ref.src_cpp_name}.value());\n"
             return ret_str
+        elif (
+            ref.src_cpp_type == OPT_AT_DOUBLE_ARRAY_REF
+            or ref.src_cpp_type == OPT_AT_INT_ARRAY_REF
+        ):
+            ret_str = f"{cpp_type} {ref.name} = "
+            ret_str += f"!{ref.src_cpp_name}.has_value() ? "
+            ret_str += f"{self.graph}{self.dot}add_none() : "
+            ret_str += f"{self.graph}{self.dot}add_scalar_list"
+            ret_str += f"({ref.src_cpp_name}->vec());\n"
+            return ret_str
         elif ref.src_cpp_type == AT_TENSOR_LIST:
             assert ref.is_in, "AT_TENSOR_LIST must be an input"
             # This logic is a bit convoluted. We need to create a IOValueRef for
@@ -588,8 +600,8 @@ class GeneratedOpsTest_{op_name} : public ::testing::TestWithParam< ::std::tuple
   protected:
     ComputeGraph* graph;
     at::ScalarType test_dtype = at::kFloat;
-    float rtol = 1e-5;
-    float atol = 1e-5;
+    float rtol = {rtol};
+    float atol = {atol};
 
     void SetUp() override {{
         GraphConfig config;
@@ -639,6 +651,8 @@ def generate_fixture_cpp(self) -> str:
             op_name=self.op_name,
             check_fn=check_fn,
             prepacked_check_fn=prepacked_check_fn,
+            rtol=self.suite_def.rtol,
+            atol=self.suite_def.atol,
         )
 
     def gen_parameterization(self) -> str:

backends/vulkan/test/op_tests/utils/codegen.py

@@ -22,6 +22,8 @@
 BOOL = "bool"
 DOUBLE = "double"
 INT = "int64_t"
+OPT_AT_DOUBLE_ARRAY_REF = "::std::optional<at::ArrayRef<double>>"
+OPT_AT_INT_ARRAY_REF = "at::OptionalIntArrayRef"
 OPT_AT_TENSOR = "::std::optional<at::Tensor>"
 OPT_BOOL = "::std::optional<bool>"
 OPT_INT64 = "::std::optional<int64_t>"
@@ -45,6 +47,8 @@ def __init__(self, input_cases: List[Any]):
         self.prepacked_args: List[str] = []
         self.requires_prepack: bool = False
         self.dtypes: List[str] = ["at::kFloat", "at::kHalf"]
+        self.atol: str = "1e-5"
+        self.rtol: str = "1e-5"
 
     def supports_prepack(self):
         return len(self.prepacked_args) > 0
@@ -142,6 +146,10 @@ def create_input_data(self, arg: Argument, data: Any) -> str:  # noqa: C901
 
         if cpp_type == AT_INT_ARRAY_REF:
             ret_str = f"std::vector<int64_t> {arg.name} = "
+        elif (
+            cpp_type == OPT_AT_DOUBLE_ARRAY_REF or cpp_type == OPT_AT_INT_ARRAY_REF
+        ) and str(data) != "None":
+            ret_str = f"std::vector<double> {arg.name} = "
         else:
             ret_str = f"{cpp_type} {arg.name} = "
 
@@ -156,6 +164,11 @@ def create_input_data(self, arg: Argument, data: Any) -> str:  # noqa: C901
             ret_str += f"{data};"
         elif cpp_type == AT_INT_ARRAY_REF:
             ret_str += f"{init_list_str(data)};"
+        elif cpp_type == OPT_AT_DOUBLE_ARRAY_REF or cpp_type == OPT_AT_INT_ARRAY_REF:
+            if str(data) == "None":
+                ret_str += "std::nullopt;"
+            else:
+                ret_str += f"{init_list_str(data)};"
         elif cpp_type == BOOL:
             ret_str += f"{str(data).lower()};"
         elif cpp_type == INT: