Arm backend: Improve round lowering (pytorch#20379)

Quantized rounding was lowered through a decomposition. This change
replaces round with a lookup table. This adds support for round on U55


cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson
@robell @rascani

Signed-off-by: Oscar Andersson <oscar.andersson@arm.com>

Author: oscarandersson8218

backends/arm/_passes/arm_pass_manager.py

@@ -677,7 +677,6 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                     InsertInt32CastsAfterInt64PlaceholdersPass(tfa_pass=True),
                     DecomposeEmbeddingPass(tfa_pass=True),
                     DecomposeScaledDotProductAttentionPass(tfa_pass=True),
-                    DecomposeRoundPass(tfa_pass=True),
                     DecomposeLogitPass(tfa_pass=True),
                     PromoteBoolOperandsPass(tfa_pass=True),
                     DecomposeSignPass(tfa_pass=True),

backends/arm/_passes/decompose_round_pass.py

@@ -5,7 +5,6 @@
 
 from typing import Set, Type
 
-import torch
 from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
@@ -33,16 +32,6 @@ def _get_round_decomposition_ops(op) -> tuple[Op, Op, Op, Op, Op, Op, Op]:
             exir_ops.edge.aten.ceil.default,
             exir_ops.edge.aten.where.self,
         )
-    elif op == torch.ops.aten.round.default:
-        return (
-            torch.ops.aten.full.default,
-            torch.ops.aten.ge.Tensor,
-            torch.ops.aten.add.Scalar,
-            torch.ops.aten.sub.Scalar,
-            torch.ops.aten.floor.default,
-            torch.ops.aten.ceil.default,
-            torch.ops.aten.where.self,
-        )
     raise RuntimeError(f"Can't get round decomposition ops for op {op}")
 
 
@@ -65,11 +54,10 @@ class DecomposeRoundPass(ArmOpTargetedPass):
 
     target_ops = {
         exir_ops.edge.aten.round.default,
-        torch.ops.aten.round.default,
     }
 
     def call_operator(self, op, args, kwargs, meta, updated=False):
-        if op not in self.target_ops or not self.allowed_to_transform(meta):
+        if op not in self.target_ops or self._is_quantized_meta(meta):
             return super().call_operator(op, args, kwargs, meta, updated)
         x = args[0]
         input_dtype = x.node.meta["val"].dtype

backends/arm/_passes/insert_table_ops.py

@@ -58,6 +58,7 @@ class TableOps:
         exir_ops.edge.aten.acos.default: torch.acos,
         exir_ops.edge.aten.tan.default: torch.tan,
         exir_ops.edge.aten.silu.default: torch.nn.functional.silu,
+        exir_ops.edge.aten.round.default: torch.round,
     }
 
     # Targets that must be treated explicitly

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -128,6 +128,7 @@
     exir_ops.edge.aten.tan.default,
     exir_ops.edge.aten.silu.default,
     exir_ops.edge.aten.detach_copy.default,
+    exir_ops.edge.aten.round.default,
 }
 
 

backends/arm/quantizer/quantization_annotator.py

@@ -532,6 +532,7 @@ def _get_fixed_qparams_qspec(
     torch.ops.aten.selu.default,
     torch.ops.aten.celu.default,
     torch.ops.aten.floor.default,
+    torch.ops.aten.round.default,
     torch.ops.aten.log.default,
     torch.ops.aten.reciprocal.default,
     torch.ops.aten.rsqrt.default,

backends/arm/test/ops/test_round.py

@@ -6,7 +6,6 @@
 
 from typing import Tuple
 
-import pytest
 import torch
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
@@ -67,7 +66,6 @@ def test_round_tosa_INT(test_data: torch.Tensor):
 
 @common.parametrize("test_data", test_data_suite)
 @common.XfailIfNoCorstone300
-@pytest.mark.xfail(reason="where.self not supported on U55")
 def test_round_u55_INT(test_data: torch.Tensor):
     pipeline = EthosU55PipelineINT[input_t1](
         Round(),