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
pytorch · May 29, 2024 · 757a6ad · 757a6ad
1 parent da73e8f
commit 757a6ad
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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} = "