Integrate a placeholder upsample_nearest2d.vec to Vulkan codegen operator tests (#3711)

Summary:
Pull Request resolved: #3711

 Integrate a placeholder upsample_nearest2d.vec to Vulkan codegen operator tests
- Right now just implement the codegen to support opt_int_array and opt_double_array.
- Using a scaling factor = 1 and just passing the input to pass the unit tests

bypass-github-export-checks
bypass-github-pytorch-ci-checks
bypass-github-executorch-ci-checks

Reviewed By: jorgep31415

Differential Revision: D57643709

fbshipit-source-id: 9d7f177890c035a5c7b44bf0a9a7e439d6a16da8

Author: Di Xu (AR)

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

@@ -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.yaml

@@ -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/impl/Upsample.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/test/op_tests/cases.py

@@ -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(
         [
@@ -796,4 +806,5 @@ 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(),
 }

backends/vulkan/test/op_tests/utils/codegen.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

backends/vulkan/test/op_tests/utils/codegen_base.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>"
@@ -142,6 +144,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 +162,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: