diff --git a/backends/arm/_passes/remove_noop_pass.py b/backends/arm/_passes/remove_noop_pass.py index b70ea23454e..b2816abd8b9 100644 --- a/backends/arm/_passes/remove_noop_pass.py +++ b/backends/arm/_passes/remove_noop_pass.py @@ -9,10 +9,11 @@ from typing import Any, Set, Type from executorch.backends.arm._passes import ArmOpTargetedPass +from executorch.backends.arm.constants import DQ_OPS, Q_OPS from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult, ProxyValue -from torch.fx import GraphModule +from torch.fx import GraphModule, Node logger = logging.getLogger(__name__) @@ -25,6 +26,10 @@ class RemoveNoopPass(ArmOpTargetedPass): exir_ops.edge.aten.cat.default, exir_ops.edge.aten.concatenate.default, ) + _resize_ops = ( + exir_ops.edge.aten.upsample_bilinear2d.vec, + exir_ops.edge.aten.upsample_nearest2d.vec, + ) target_ops = ( exir_ops.edge.dim_order_ops._clone_dim_order.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, @@ -32,8 +37,7 @@ class RemoveNoopPass(ArmOpTargetedPass): exir_ops.edge.aten.copy.default, exir_ops.edge.aten.detach_copy.default, *_single_input_concat_ops, - exir_ops.edge.aten.upsample_bilinear2d.vec, - exir_ops.edge.aten.upsample_nearest2d.vec, + *_resize_ops, exir_ops.backend.tosa.PAD.default, exir_ops.backend.tosa.SLICE.default, ) @@ -54,7 +58,103 @@ def _get_static_shape(value: Any) -> tuple[int, ...] | None: return None return tuple(value) + @staticmethod + def _is_identity_scale_factor(scale_factors: Any) -> bool: + if type(scale_factors) in (int, float): + return scale_factors == 1.0 + return ( + isinstance(scale_factors, (list, tuple)) + and len(scale_factors) == 2 + and all( + type(scale) in (int, float) and scale == 1.0 for scale in scale_factors + ) + ) + + def _resize_scale_factors(self, op, args: tuple[Any, ...]) -> Any: + return args[3] if op == exir_ops.edge.aten.upsample_bilinear2d.vec else args[2] + + def _is_identity_resize(self, op, args: tuple[Any, ...]) -> bool: + scale_factors = self._resize_scale_factors(op, args) + if self._is_identity_scale_factor(scale_factors): + return True + + output_spatial_size = self._get_static_shape(args[1]) + input_spatial_size = self._get_static_shape(args[0].data.shape[2:]) + return ( + output_spatial_size is not None + and input_spatial_size is not None + and input_spatial_size == output_spatial_size + and scale_factors is None + ) + + def _is_identity_resize_node(self, node: Node) -> bool: + if node.target not in self._resize_ops: + return False + input_node = node.args[0] + if not isinstance(input_node, Node): + return False + + scale_factors = self._resize_scale_factors(node.target, node.args) + if self._is_identity_scale_factor(scale_factors): + return True + + output_spatial_size = self._get_static_shape(node.args[1]) + input_spatial_size = self._get_static_shape(input_node.meta["val"].shape[2:]) + return ( + output_spatial_size is not None + and input_spatial_size is not None + and input_spatial_size == output_spatial_size + and scale_factors is None + ) + + def _is_removable_noop_node(self, node: Node) -> bool: + if node.target in self._single_input_concat_ops: + inputs = node.args[0] + return isinstance(inputs, (list, tuple)) and len(inputs) == 1 + + if node.target in ( + exir_ops.edge.dim_order_ops._clone_dim_order.default, + exir_ops.edge.aten.alias_copy.default, + exir_ops.edge.aten.detach_copy.default, + ): + return True + + if node.target == exir_ops.edge.aten.copy.default: + return True + + if node.target == exir_ops.edge.dim_order_ops._to_dim_order_copy.default: + input_node = node.args[0] + if not isinstance(input_node, Node): + return False + input_dtype = input_node.meta["val"].dtype + output_dtype = node.kwargs.get("dtype", input_dtype) + return input_dtype == output_dtype + + return False + + def _partition_would_be_empty_without_identity_resize( + self, graph_module: GraphModule + ) -> bool: + has_identity_resize = False + for node in graph_module.graph.nodes: + if ( + node.op != "call_function" + or node.target in Q_OPS + or node.target in DQ_OPS + ): + continue + if self._is_identity_resize_node(node): + has_identity_resize = True + continue + if self._is_removable_noop_node(node): + continue + return False + return has_identity_resize + def call(self, graph_module: GraphModule) -> PassResult: + if self._partition_would_be_empty_without_identity_resize(graph_module): + return PassResult(graph_module, False) + result = super().call(graph_module) # Removing a no-op can leave its shape operands without users. removed_dead_code = result.graph_module.graph.eliminate_dead_code() @@ -63,7 +163,7 @@ def call(self, graph_module: GraphModule) -> PassResult: result.graph_module.recompile() return PassResult(result.graph_module, result.modified or removed_dead_code) - def call_operator(self, op, args, kwargs, meta, updated=False): + def call_operator(self, op, args, kwargs, meta, updated=False): # noqa: C901 if op not in self.target_ops: return super().call_operator(op, args, kwargs, meta, updated) @@ -74,21 +174,8 @@ def call_operator(self, op, args, kwargs, meta, updated=False): return inputs[0] return super().call_operator(op, args, kwargs, meta, updated) - if op in ( - exir_ops.edge.aten.upsample_bilinear2d.vec, - exir_ops.edge.aten.upsample_nearest2d.vec, - ): - scale_factors = ( - args[3] if op == exir_ops.edge.aten.upsample_bilinear2d.vec else args[2] - ) - if ( - isinstance(scale_factors, (list, tuple)) - and len(scale_factors) == 2 - and all( - type(scale) in (int, float) and scale == 1.0 - for scale in scale_factors - ) - ): + if op in self._resize_ops: + if self._is_identity_resize(op, args): return args[0] return super().call_operator(op, args, kwargs, meta, updated) diff --git a/backends/arm/test/ops/test_upsample_bilinear2d.py b/backends/arm/test/ops/test_upsample_bilinear2d.py index 752986992ad..7196bc6844e 100644 --- a/backends/arm/test/ops/test_upsample_bilinear2d.py +++ b/backends/arm/test/ops/test_upsample_bilinear2d.py @@ -131,6 +131,11 @@ ), } +test_data_suite_Uxx_same_size = { + "rand_same_size": lambda: (torch.rand(2, 3, 5, 5), (5, 5), None, False), + "rand_same_scale": lambda: (torch.rand(2, 3, 5, 5), None, 1.0, False), +} + test_data_u55 = { "rand_double_size": lambda: (torch.rand(2, 4, 8, 3), (16, 6), None, True), } @@ -423,6 +428,24 @@ def test_upsample_bilinear2d_vec_u55_INT_UpsamplingBilinear2d_not_delegated( pipeline.run() +@common.parametrize("test_data", test_data_suite_Uxx_same_size) +def test_upsample_bilinear2d_vec_u85_INT_same_size( + test_data: torch.Tensor, +): + test_data, size, scale_factor, compare_outputs = test_data() + + pipeline = EthosU85PipelineINT[input_t1]( + InterpolateAlignCornersFalse(size, scale_factor), + (test_data,), + aten_op, + qtol=1, + use_to_edge_transform_and_lower=True, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + @common.parametrize("test_data", test_data_suite_Uxx) @common.XfailIfNoCorstone320 def test_upsample_bilinear2d_vec_u85_INT_Upsample(test_data: input_t1): diff --git a/backends/arm/test/ops/test_upsample_nearest2d.py b/backends/arm/test/ops/test_upsample_nearest2d.py index 1f05ad83c08..3d35f6f8633 100644 --- a/backends/arm/test/ops/test_upsample_nearest2d.py +++ b/backends/arm/test/ops/test_upsample_nearest2d.py @@ -17,6 +17,7 @@ from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, + EthosU85PipelineINT, OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, @@ -92,6 +93,11 @@ "rand_octuple_size": lambda: (torch.rand(1, 4, 8, 3), (64, 24), None, True), } +test_data_suite_u85_same_size = { + "rand_same_size": lambda: (torch.rand(2, 3, 5, 5), (5, 5), None, False), + "rand_same_scale": lambda: (torch.rand(2, 3, 5, 5), None, 1.0, False), +} + test_data_suite_dynamic = { # (test_name, test_data, size, scale_factor, compare_outputs) "rand_double_scale": lambda: (torch.rand(2, 4, 8, 3), None, 2.0, False), @@ -513,6 +519,23 @@ def test_upsample_nearest2d_vec_u55_INT_UpsamplingNearest2d( pipeline.run() +@common.parametrize("test_data", test_data_suite_u85_same_size) +def test_upsample_nearest2d_vec_u85_INT_same_size( + test_data: torch.Tensor, +): + test_data, size, scale_factor, compare_outputs = test_data() + + pipeline = EthosU85PipelineINT[input_t1]( + Interpolate(size, scale_factor), + (test_data,), + aten_op, + exir_op, + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + def test_upsample_nearest2d_vec_u55_INT_unsupported_scale_not_delegated(): # 2.25 rounds a 2x2 input to 4x4, which looks like a supported 2x resize from shapes alone. test_data = torch.rand(1, 4, 2, 2) diff --git a/backends/arm/test/passes/test_remove_data_layout_noops.py b/backends/arm/test/passes/test_remove_data_layout_noops.py index 7090cfbf008..a690f3cfb3a 100644 --- a/backends/arm/test/passes/test_remove_data_layout_noops.py +++ b/backends/arm/test/passes/test_remove_data_layout_noops.py @@ -25,7 +25,7 @@ ) from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops -from torch.export import export, ExportedProgram +from torch.export import Dim, export, ExportedProgram from torch.fx import Graph, GraphModule, Node @@ -119,7 +119,7 @@ def test_keep_multi_input_concat(): assert _count_target(result, exir_ops.edge.aten.cat.default) == 1 -def test_remove_identity_bilinear_upsample(): +def test_keep_identity_bilinear_upsample_when_only_partition_op(): graph = Graph() x = graph.placeholder("x") x.meta["val"] = torch.ones(1, 4, 768, 384) @@ -133,11 +133,10 @@ def test_remove_identity_bilinear_upsample(): result = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) - assert _count_target(result, exir_ops.edge.aten.upsample_bilinear2d.vec) == 0 - assert result.graph.output_node().args[0][0].op == "placeholder" + assert _count_target(result, exir_ops.edge.aten.upsample_bilinear2d.vec) == 1 -def test_remove_identity_nearest_upsample(): +def test_keep_identity_nearest_upsample_when_only_partition_op(): graph = Graph() x = graph.placeholder("x") x.meta["val"] = torch.ones(1, 4, 768, 384) @@ -151,8 +150,57 @@ def test_remove_identity_nearest_upsample(): result = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) - assert _count_target(result, exir_ops.edge.aten.upsample_nearest2d.vec) == 0 - assert result.graph.output_node().args[0][0].op == "placeholder" + assert _count_target(result, exir_ops.edge.aten.upsample_nearest2d.vec) == 1 + + +def test_keep_identity_bilinear_upsample_with_only_removable_noops(): + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(1, 4, 768, 384) + upsample = _call( + graph, + exir_ops.edge.aten.upsample_bilinear2d.vec, + (x, None, False, [1.0, 1.0]), + torch.ones(1, 4, 768, 384), + ) + alias = _call( + graph, + exir_ops.edge.aten.alias_copy.default, + (upsample,), + torch.ones(1, 4, 768, 384), + ) + graph.output((alias,)) + + result = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) + + assert _count_target(result, exir_ops.edge.aten.upsample_bilinear2d.vec) == 1 + assert _count_target(result, exir_ops.edge.aten.alias_copy.default) == 1 + + +def test_remove_identity_bilinear_upsample_when_compute_op_survives(): + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(1, 4, 768, 384) + y = graph.placeholder("y") + y.meta["val"] = torch.ones(1, 4, 768, 384) + upsample = _call( + graph, + exir_ops.edge.aten.upsample_bilinear2d.vec, + (x, None, False, [1.0, 1.0]), + torch.ones(1, 4, 768, 384), + ) + add = _call( + graph, + exir_ops.edge.aten.add.Tensor, + (upsample, y), + torch.ones(1, 4, 768, 384), + ) + graph.output((add,)) + + result = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) + + assert _count_target(result, exir_ops.edge.aten.upsample_bilinear2d.vec) == 0 + assert _count_target(result, exir_ops.edge.aten.add.Tensor) == 1 def test_keep_bilinear_upsample_with_rounded_identity_shape(): @@ -172,6 +220,54 @@ def test_keep_bilinear_upsample_with_rounded_identity_shape(): assert _count_target(result, exir_ops.edge.aten.upsample_bilinear2d.vec) == 1 +class _ResizeNearestToExampleSpatialSizeModule(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate(x, size=(8, 3), mode="nearest") + + +class _ResizeBilinearToExampleSpatialSizeModule(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate( + x, size=(8, 3), mode="bilinear", align_corners=False + ) + + +@pytest.mark.parametrize( + "module, target", + [ + ( + _ResizeNearestToExampleSpatialSizeModule(), + exir_ops.edge.aten.upsample_nearest2d.vec, + ), + ( + _ResizeBilinearToExampleSpatialSizeModule(), + exir_ops.edge.aten.upsample_bilinear2d.vec, + ), + ], +) +def test_keep_upsample_matching_example_shape_with_dynamic_spatial_dims(module, target): + exported_program = export( + module, + (torch.ones(2, 4, 8, 3),), + dynamic_shapes={ + "x": { + 2: Dim("input_height", min=4, max=12), + 3: Dim("input_width", min=2, max=6), + } + }, + strict=True, + ) + edge_program = to_edge( + exported_program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + + assert _count_target(edge_program.graph_module, target) == 1 + result = _run_remove_noop(edge_program.graph_module) + + assert _count_target(result, target) == 1 + + class _IdentityUpsampleModule(torch.nn.Module): def forward(self, x): return torch.nn.functional.interpolate( @@ -179,7 +275,7 @@ def forward(self, x): ) -def test_remove_identity_bilinear_upsample_backend_pipeline(): +def test_keep_identity_bilinear_upsample_backend_pipeline_when_only_partition_op(): exported_program = export( _IdentityUpsampleModule(), (torch.ones(1, 4, 768, 384),), strict=True ) @@ -199,7 +295,7 @@ def test_remove_identity_bilinear_upsample_backend_pipeline(): TosaCompileSpec("TOSA-1.0+FP") ).transform_to_backend_pipeline(edge_program, edge_program.graph_module) - assert _count_target(graph_module, exir_ops.backend.tosa.RESIZE.default) == 0 + assert _count_target(graph_module, exir_ops.backend.tosa.RESIZE.default) == 1 def test_remove_full_slice_and_unused_shape_constants():