Add resize function for aten.upsample_nearest2d.vec operator (#4069)

Summary:
Pull Request resolved: #4069

aten.upsample_nearest2d.vec is an operator in OCR model (Neck module). We added support to not decompose this op when lowering to Vulkan delegation. Add resize function so that we only need to build the graph once to test input with different sizes.

Reviewed By: copyrightly, jorgep31415

Differential Revision: D59030889

Author: Yujie Hui

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

@@ -16,6 +16,29 @@
 
 namespace vkcompute {
 
+void resize_upsample_nearest2d_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& extra_args) {
+  vTensorPtr out = graph->get_tensor(args[0].refs[0]);
+  vTensorPtr self = graph->get_tensor(args[1].refs[0]);
+  std::vector<int64_t> out_sizes = self->sizes(); // NCHW
+
+  const ValueRef output_sizes = extra_args[0]; // HW
+  const ValueRef scale_factors = extra_args[1]; // HW
+  if (!graph->val_is_none(output_sizes)) {
+    IntListPtr output_size_ref = graph->get_int_list(output_sizes);
+    out_sizes.at(2) = output_size_ref->at(0);
+    out_sizes.at(3) = output_size_ref->at(1);
+  } else {
+    DoubleListPtr scales = graph->get_double_list(scale_factors);
+    out_sizes.at(2) *= scales->at(0);
+    out_sizes.at(3) *= scales->at(1);
+  }
+
+  out->virtual_resize(out_sizes);
+}
+
 // ExecuTorch-Vulkan framework to add node
 // Args:
 //   in: will be converted from NCHW input tensor to 3D ARGB representation in
@@ -87,7 +110,9 @@ void add_upsample_nearest2d_node(
        graph.create_params_buffer(input_size),
        graph.create_params_buffer(rev_scales)},
       // Specialization Constants
-      {}));
+      {},
+      resize_upsample_nearest2d_node,
+      {output_sizes, scale_factors}));
 }
 
 void upsample(ComputeGraph& graph, const std::vector<ValueRef>& args) {