Implement upsample_nearest2d.vec op to pass codegen operator tests: scale factor supports only

Summary:
Implement upsample_nearest2d.vec op to pass codegen operator tests
- Scale factor supports only
- Renamed the glsl file to upsample_nearest2d since the caffe2 version has multiple upsample op versions
- Output sizes support and others will be following up in T190297757
- May follow up on the dynamic shape support next

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

Reviewed By: copyrightly

Differential Revision: D57698138

fbshipit-source-id: 8a4d1f5bd4565c52a0b919302ac202aede2a6b74

Author: Di Xu (AR)

backends/vulkan/runtime/graph/ops/glsl/upsample.glsl renamed to backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.glsl

@@ -20,6 +20,8 @@ layout(std430) buffer;
 ${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)}
 ${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)}
 ${layout_declare_ubo(2, "ivec3", "out_limits")}
+${layout_declare_ubo(3, "ivec2", "input_size")}
+${layout_declare_ubo(4, "vec2", "rev_scales")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -30,6 +32,8 @@ void main() {
     return;
   }
 
-  VEC4_T in_texel = texelFetch(t_in, pos, 0);
+  const ivec2 ipos = clamp(ivec2(pos.xy * rev_scales), ivec2(0), input_size);
+
+  VEC4_T in_texel = texelFetch(t_in, ivec3(ipos, pos.z), 0);
   imageStore(t_out, pos, in_texel);
 }

backends/vulkan/runtime/graph/ops/glsl/upsample.yaml renamed to backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.yaml

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-upsample:
+upsample_nearest2d:
   parameter_names_with_default_values:
     NDIM: 3
     DTYPE: float
@@ -15,4 +15,4 @@ upsample:
       - VALUE: half
       - VALUE: float
   shader_variants:
-    - NAME: upsample
+    - NAME: upsample_nearest2d

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

@@ -10,33 +10,56 @@
 
 #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;
-}
+// Executorch-Vulkan framework to add node
+// Args:
+//   in: will be converted from NCHW input tensor to 3D ARGB representation in
+//   openGL (via Executorch) output_sizes: optional 2D array of targetting
+//   output size of H and W dimensions. >= input sizes;
 
-void add_upsample_node(
+//      will be computed if only given the scale_factors.
+//   scale_factors: optional 2D array of scale factors for H and W dimensions.
+//      Will be computed if only given the output_sizes.
+void add_upsample_nearest2d_node(
     ComputeGraph& graph,
     const ValueRef in,
+    const ValueRef output_sizes,
+    const ValueRef scale_factors,
     const ValueRef out) {
-  ValueRef arg = prepack_if_tensor_ref(graph, in);
+  // TODO(T190297757) add supports for output_sizes
+  if (graph.val_is_none(output_sizes) && graph.val_is_none(scale_factors)) {
+    VK_THROW(
+        "Invalid input, must provide either output_sizes or scale_factors");
+  }
+  if (!graph.val_is_none(output_sizes) && !graph.val_is_none(scale_factors)) {
+    VK_THROW(
+        "Invalid input, must provide ONLY one of output_sizes or scale_factors");
+  }
+  auto scales = graph.get_double_list(scale_factors);
+
+  ValueRef arg_in = prepack_if_tensor_ref(graph, in);
+
+  vTensorPtr t_in = graph.get_tensor(in);
+  api::utils::uvec3 input_sizes = t_in->image_extents();
+
+  api::utils::ivec2 input_size = {
+      api::utils::safe_downcast<int32_t>(input_sizes.data[0]),
+      api::utils::safe_downcast<int32_t>(input_sizes.data[1])};
+  // Reverse scale factors that pre-computed before GLSL.
+  api::utils::vec2 rev_scales = {
+      api::utils::safe_downcast<float>(1.0 / scales->at(1)),
+      api::utils::safe_downcast<float>(1.0 / scales->at(0))};
 
   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");
+  std::string kernel_name("upsample_nearest2d");
   kernel_name.reserve(kShaderNameReserve);
 
   add_dtype_suffix(kernel_name, *t_out);
@@ -47,17 +70,18 @@ void add_upsample_node(
       global_size,
       local_size,
       // Inputs and Outputs
-      {{out, api::MemoryAccessType::WRITE}, {arg, api::MemoryAccessType::READ}},
+      {{out, api::MemoryAccessType::WRITE},
+       {arg_in, api::MemoryAccessType::READ}},
       // Shader params buffers
-      {t_out->texture_limits_ubo()},
+      {t_out->texture_limits_ubo(),
+       graph.create_params_buffer(input_size),
+       graph.create_params_buffer(rev_scales)},
       // 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]);
+  return add_upsample_nearest2d_node(graph, args[0], args[1], args[2], args[3]);
 }
 
 REGISTER_OPERATORS {

backends/vulkan/test/op_tests/cases.py

@@ -244,8 +244,12 @@ def get_native_layer_norm_inputs():
 def get_upsample_inputs():
     test_suite = VkTestSuite(
         [
-            # TODO(dixu): implement the basic upsample logic to have a meaningful test
+            # (input tensor shape, output 2D image size (H, W), output scaling factors)
             ((2, 2, 2, 2), None, [1, 1]),
+            ((1, 1, 2, 2), None, [2, 2]),
+            ((1, 1, 2, 2), None, [2, 4]),
+            ((1, 1, 2, 2), None, [4, 2]),
+            # TODO(T190297757) add supports for output_sizes
         ]
     )
     return test_suite