diff --git a/backends/xnnpack/operators/op_static_resize_bilinear_2d.py b/backends/xnnpack/operators/op_static_resize_bilinear_2d.py index 96e68750785..83f3f0ea059 100644 --- a/backends/xnnpack/operators/op_static_resize_bilinear_2d.py +++ b/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 diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index f1f51b27b6c..835db7a764f 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/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, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index caf61f33c9b..e325f6c75bb 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/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 diff --git a/backends/xnnpack/passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/passes/channels_last_tagged_reshape_pass.py index 1816d6ccfcf..6a14b1679ed 100644 --- a/backends/xnnpack/passes/channels_last_tagged_reshape_pass.py +++ b/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, diff --git a/backends/xnnpack/passes/convert_to_upsample_bilinear2d.py b/backends/xnnpack/passes/convert_to_upsample_bilinear2d.py index 065f3254ce6..45956ee6f6f 100644 --- a/backends/xnnpack/passes/convert_to_upsample_bilinear2d.py +++ b/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), ) diff --git a/backends/xnnpack/test/ops/bilinear2d.py b/backends/xnnpack/test/ops/bilinear2d.py index d3c85350692..bf89e2196f7 100644 --- a/backends/xnnpack/test/ops/bilinear2d.py +++ b/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