Upsample Bilinear 2d

Differential Revision: D60323281

Pull Request resolved: #4535

Author: mcr229

backends/xnnpack/operators/op_static_resize_bilinear_2d.py

@@ -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.
 
-from typing import cast, Dict, List
+from typing import cast, Dict
 
 import torch
 from executorch.backends.xnnpack.operators.node_visitor import (
@@ -23,7 +23,7 @@
 
 @register_node_visitor
 class StaticResizeBilinear2DVisitor(NodeVisitor):
-    target = "aten.upsample_bilinear2d.default"
+    target = "aten.upsample_bilinear2d.vec"
 
     def __init__(self, *args) -> None:
         super().__init__(*args)
@@ -44,7 +44,7 @@ def define_node(
         # output
         output_id = vals_to_ids[node]
 
-        new_size = cast(List[int], node.args[1])
+        new_size = node.meta["val"].shape[-2:]
 
         flags = XNN_FLAG_ALIGN_CORNERS if cast(bool, node.args[2]) else 0
 

backends/xnnpack/partition/config/__init__.py

@@ -32,6 +32,7 @@
     ReLUConfig,
     SigmoidConfig,
     SoftmaxConfig,
+    UpsampleBilinear2dConfig,
 )
 from executorch.backends.xnnpack.partition.config.node_configs import (
     BatchNormConfig,
@@ -66,6 +67,7 @@
     PermuteConfig,
     # EluConfig, # Waiting for PyTorch Pin Update
     ReLUConfig,
+    UpsampleBilinear2dConfig,
     # Quantization Op Configs
     QuantizedPerTensorConfig,
     DeQuantizedPerTensorConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -261,3 +261,13 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
 
     def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
         return torch.ops.aten.max_pool2d.default
+
+
+class UpsampleBilinear2dConfig(GenericNodePartitionerConfig):
+    target_name = "upsample_bilinear2d.vec"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+    def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
+        return torch.ops.aten.upsample_bilinear2d.vec

backends/xnnpack/passes/channels_last_tagged_reshape_pass.py

@@ -44,7 +44,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass):
     # Set of ops that require memory format to be channels last (NHWC)
     memory_sensitive_ops_nhwc = {
         exir_ops.edge.aten.convolution.default,
-        exir_ops.edge.aten.upsample_bilinear2d.default,
+        exir_ops.edge.aten.upsample_bilinear2d.vec,
         exir_ops.edge.aten.mean.dim,
         exir_ops.edge.aten.max_pool2d.default,
         exir_ops.edge.aten.amax.default,

backends/xnnpack/passes/convert_to_upsample_bilinear2d.py

@@ -36,7 +36,7 @@ def create_upsample_bilinear_2d(
         with graph_module.graph.inserting_before(output):
             upsample_node = graph_module.graph.create_node(
                 "call_function",
-                exir_ops.edge.aten.upsample_bilinear2d.default,
+                exir_ops.edge.aten.upsample_bilinear2d.vec,
                 # TODO(T166527012): Using output_h and output_w here only works with static shapes
                 args=(input_node, [output_h, output_w], align_corners, None),
             )

backends/xnnpack/test/ops/bilinear2d.py

@@ -83,8 +83,7 @@ def test_fp32_static_resize_bilinear2d(self):
         (
             Tester(self.StaticResizeBilinear2dModule(), example_inputs)
             .export()
-            .to_edge()
-            .partition()
+            .to_edge_transform_and_lower()
             .check_not(self.ops)
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .to_executorch()
@@ -97,8 +96,7 @@ def test_fp32_static_resize_bilinear2d_with_align_corners(self):
         (
             Tester(self.StaticResizeBilinear2dModuleWithAlignCorners(), example_inputs)
             .export()
-            .to_edge()
-            .partition()
+            .to_edge_transform_and_lower()
             .check_not(self.ops)
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .to_executorch()
@@ -112,13 +110,7 @@ def test_fp32_static_resize_bilinear2d_antialiased(self):
         (
             Tester(self.Bilinear2dAntiAlias(), example_inputs)
             .export()
-            .to_edge()
-            .check_count(
-                {
-                    "executorch_exir_dialects_edge__ops_aten__upsample_bilinear2d_aa_default": 2
-                }
-            )
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count(
                 {
                     "executorch_exir_dialects_edge__ops_aten__upsample_bilinear2d_aa_default": 2