From 757a6ad212d780526243a954141a23dc08234b32 Mon Sep 17 00:00:00 2001
From: "Di Xu (AR)" <dixu@meta.com>
Date: Tue, 28 May 2024 21:31:45 -0700
Subject: [PATCH] Upsample add support for output_sizes argument

Summary:
Upsample operator add support for output_sizes argument

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

Reviewed By: copyrightly

Differential Revision: D57887635

fbshipit-source-id: 1e5d53bc11c79b022e31aed3bf3290207abf418f
---
 .../runtime/graph/ops/impl/Upsample.cpp       | 23 +++++++++++++++----
 backends/vulkan/test/op_tests/cases.py        |  4 +++-
 .../test/op_tests/utils/codegen_base.py       |  6 ++---
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp
index ddce4549c3..3646f1f2b1 100644
--- a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp
@@ -31,7 +31,6 @@ void add_upsample_nearest2d_node(
     const ValueRef output_sizes,
     const ValueRef scale_factors,
     const ValueRef out) {
-  // 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");
@@ -40,7 +39,6 @@ void add_upsample_nearest2d_node(
     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);
 
@@ -50,10 +48,25 @@ void add_upsample_nearest2d_node(
   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))};
+      api::utils::safe_downcast<float>(1.0),
+      api::utils::safe_downcast<float>(1.0)};
+
+  // Reverse scale factors that pre-computed before GLSL.
+  if (!graph.val_is_none(output_sizes)) {
+    auto output_size_ref = graph.get_int_list(output_sizes);
+    rev_scales = {
+        api::utils::safe_downcast<float>(
+            (float)input_size.data[0] / output_size_ref->at(1)),
+        api::utils::safe_downcast<float>(
+            (float)input_size.data[1] / output_size_ref->at(0))};
+
+  } else {
+    auto scales = graph.get_double_list(scale_factors);
+    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();
diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py
index 8754a79002..07abe277fe 100644
--- a/backends/vulkan/test/op_tests/cases.py
+++ b/backends/vulkan/test/op_tests/cases.py
@@ -249,7 +249,9 @@ def get_upsample_inputs():
             ((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
+            ((1, 1, 2, 2), [2, 2], None),
+            ((1, 1, 2, 2), [2, 4], None),
+            ((1, 1, 2, 2), [3, 2], None),
         ]
     )
     return test_suite
diff --git a/backends/vulkan/test/op_tests/utils/codegen_base.py b/backends/vulkan/test/op_tests/utils/codegen_base.py
index 238c8b594b..d4480ef3c5 100644
--- a/backends/vulkan/test/op_tests/utils/codegen_base.py
+++ b/backends/vulkan/test/op_tests/utils/codegen_base.py
@@ -146,10 +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":
+        elif cpp_type == OPT_AT_DOUBLE_ARRAY_REF and str(data) != "None":
             ret_str = f"std::vector<double> {arg.name} = "
+        elif cpp_type == OPT_AT_INT_ARRAY_REF and str(data) != "None":
+            ret_str = f"std::vector<int64_t> {arg.name} = "
         else:
             ret_str = f"{cpp_type} {arg.name} = "