Arm backend: Add LUT-based leaky_relu lowering (pytorch#20837)

Add option to lower quantized leaky_relu through TOSA TABLE instead of
decomposing it into clamp, mul, and add. Add an Arm pass pipeline config
switch so callers can enable the table-based lowering rather than the
old decomposition.

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.py

@@ -82,12 +82,16 @@ def get_name(pass_) -> str:
             )
 
     def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False):
+        ops_without_quantized_fake_kernel = {
+            exir_ops.edge.aten.bmm.default,
+            exir_ops.edge.aten.leaky_relu.default,
+        }
         if (
-            op == exir_ops.edge.aten.bmm.default
+            op in ops_without_quantized_fake_kernel
             and isinstance(meta, NodeMetadata)
             and len(meta.data.get("input_qparams", {})) > 0
         ):
-            return self._call_quantized_bmm_without_fake_kernel(op, args, kwargs, meta)
+            return self._call_quantized_op_without_fake_kernel(op, args, kwargs, meta)
 
         if not updated:
             return super().call_operator(op, args, kwargs, meta)
@@ -101,7 +105,7 @@ def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False)
         new_meta["stack_trace"] = f"{old_stack_trace}\n{traceback.format_stack()[-2]}"
         return super().call_operator(op, args, kwargs, NodeMetadata(new_meta))
 
-    def _call_quantized_bmm_without_fake_kernel(
+    def _call_quantized_op_without_fake_kernel(
         self,
         op,
         args: tuple[ProxyValue, ...],

backends/arm/_passes/arm_pass_manager.py

@@ -164,7 +164,10 @@
 )
 from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
-from executorch.backends.arm.common.pipeline_config import SoftmaxDecompositionConfig
+from executorch.backends.arm.common.pipeline_config import (
+    LeakyReLULoweringConfig,
+    SoftmaxDecompositionConfig,
+)
 from executorch.backends.arm.tosa.specification import (
     tosa_spec_in_set,
     TosaLoweringContext,
@@ -292,7 +295,8 @@ def configure_skip_passes(self) -> tuple[type, ...]:
                 pass
             case SoftmaxDecompositionConfig.STABLE:
                 skip_set.add(DecomposeMaskedFillPass)
-
+        if config.leaky_relu == LeakyReLULoweringConfig.TABLE:
+            skip_set.add(DecomposeLeakyReLUPass)
         self._skip_pass_types = tuple(skip_set)
         skip_names = [skipped_pass.__name__ for skipped_pass in self._skip_pass_types]
         logger.debug(f"Passes in skip list: {skip_names}")
@@ -704,7 +708,6 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                     DecomposeEinsumPass(tfa_pass=True),
                     RewriteInplaceArithmeticPass(tfa_pass=True),
                     DecomposeAddSubAlphaPass(tfa_pass=True),
-                    DecomposeLeakyReLUPass(tfa_pass=True),
                     ConvertEluFamilyToEluPass(tfa_pass=True),
                     DecomposeGroupNormPass(tfa_pass=True),
                     DecomposeLayerNormPass(tfa_pass=True),
@@ -715,6 +718,12 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 ]
             )
 
+            if (
+                self.compile_spec._get_pass_pipeline_config().leaky_relu
+                is LeakyReLULoweringConfig.DECOMPOSE
+            ):
+                self.add_pass(DecomposeLeakyReLUPass(tfa_pass=True))
+
             # Scalars -> tensors
             self.add_passes(
                 [

backends/arm/_passes/decompose_leaky_relu_pass.py

@@ -51,23 +51,30 @@ class DecomposeLeakyReLUPass(ArmOpTargetedPass):
     check_allowed_to_transform = True
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in self.target_ops or not self.allowed_to_transform(meta):
+        if (
+            op not in self.target_ops
+            or not self.allowed_to_transform(meta)
+            or (self._is_quantized_meta(meta))
+        ):
             return super().call_operator(op, args, kwargs, meta)
 
         x = args[0]
         slope = args[1] if len(args) > 1 else 0.01
         clamp, mul, add = _get_leaky_relu_ops(op)
         op1 = super().call_operator(
-            op=clamp, args=(x, 0, None), kwargs=kwargs, meta=meta
+            op=clamp, args=(x, 0, None), kwargs=kwargs, meta=meta, updated=True
         )
         op2 = super().call_operator(
-            op=clamp, args=(x, None, 0), kwargs=kwargs, meta=meta
+            op=clamp, args=(x, None, 0), kwargs=kwargs, meta=meta, updated=True
         )
         op4 = super().call_operator(
             op=mul,
             args=(op2, super().call_scalar(slope, meta)),
             kwargs=kwargs,
             meta=meta,
+            updated=True,
+        )
+        op5 = super().call_operator(
+            op=add, args=(op1, op4), kwargs=kwargs, meta=meta, updated=True
         )
-        op5 = super().call_operator(op=add, args=(op1, op4), kwargs=kwargs, meta=meta)
         return op5

backends/arm/_passes/insert_table_ops.py

@@ -66,6 +66,7 @@ class TableOps:
         exir_ops.edge.aten.pow.Tensor_Scalar,
         exir_ops.edge.aten.gelu.default,
         exir_ops.edge.aten.elu.default,
+        exir_ops.edge.aten.leaky_relu.default,
         exir_ops.edge.aten.remainder.Scalar,
     }
 
@@ -107,6 +108,18 @@ def __getitem__(self, node: Node):
                     return lambda x: torch.ops.aten.elu.default(
                         x, input_alpha, scale, input_scale
                     ).flatten()
+                case exir_ops.edge.aten.leaky_relu.default:
+                    negative_slope = cast(
+                        float,
+                        (
+                            node.args[1]
+                            if len(node.args) > 1
+                            else node.kwargs.get("negative_slope", 0.01)
+                        ),
+                    )
+                    return lambda x: torch.nn.functional.leaky_relu(
+                        x, negative_slope=negative_slope
+                    ).flatten()
                 case exir_ops.edge.aten.remainder.Scalar:
                     divisor = cast(float | int, node.args[1])
                     return lambda x: torch.remainder(x, divisor).flatten()

backends/arm/common/pipeline_config.py

@@ -14,6 +14,11 @@ class SoftmaxDecompositionConfig(Enum):
     STABLE = auto()  # Stable softmax, no masked fill decomposition
 
 
+class LeakyReLULoweringConfig(Enum):
+    TABLE = auto()  # Lower quantized leaky_relu with TOSA TABLE
+    DECOMPOSE = auto()  # Lower leaky_relu into clamp, mul, and add
+
+
 @dataclass
 class QuantizeInfConfig:
     """Replacement values for infinities before quantization.
@@ -37,13 +42,16 @@ class ArmPassPipelineConfig:
 
     Args:
         softmax (SoftmaxDecompositionConfig): Softmax decomposition mode.
+        leaky_relu (LeakyReLULoweringConfig): Quantized leaky_relu lowering
+            mode.
         quantize_inf (QuantizeInfConfig): Values used when replacing
             infinities before quantization.
 
     Example:
         compile_spec.set_pass_pipeline_config(
             ArmPassPipelineConfig(
                 softmax=SoftmaxDecompositionConfig.STABLE,
+                leaky_relu=LeakyReLULoweringConfig.DECOMPOSE,
                 quantize_inf=QuantizeInfConfig(
                     neg_inf=-100.0,
                     pos_inf=100.0,
@@ -54,11 +62,13 @@ class ArmPassPipelineConfig:
     """
 
     softmax: SoftmaxDecompositionConfig = SoftmaxDecompositionConfig.MASKED
+    leaky_relu: LeakyReLULoweringConfig = LeakyReLULoweringConfig.DECOMPOSE
     quantize_inf: QuantizeInfConfig = field(default_factory=QuantizeInfConfig)
 
     def is_default(self) -> bool:
         return (
             self.softmax is SoftmaxDecompositionConfig.MASKED
+            and self.leaky_relu is LeakyReLULoweringConfig.DECOMPOSE
             and self.quantize_inf == QuantizeInfConfig()
         )
 

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -70,6 +70,7 @@
     exir_ops.edge.aten.gt.Scalar,
     exir_ops.edge.aten.le.Tensor,
     exir_ops.edge.aten.le.Scalar,
+    exir_ops.edge.aten.leaky_relu.default,
     exir_ops.edge.aten.lt.Tensor,
     exir_ops.edge.aten.lt.Scalar,
     exir_ops.edge.aten.mul.Tensor,

backends/arm/quantizer/quantization_annotator.py

@@ -545,6 +545,8 @@ def _get_fixed_qparams_qspec(
     torch.ops.aten.hardsigmoid.default,
     torch.ops.aten.hardswish.default,
     torch.ops.aten.hardswish_.default,
+    torch.ops.aten.leaky_relu.default,
+    torch.ops.aten.leaky_relu_.default,
     torch.ops.aten.full_like.default,
     torch.ops.aten.zeros_like.default,
     torch.ops.aten.pow.Tensor_Scalar,

backends/arm/test/misc/test_compile_spec.py

@@ -5,7 +5,10 @@
 
 import warnings
 
-from executorch.backends.arm.common.pipeline_config import SoftmaxDecompositionConfig
+from executorch.backends.arm.common.pipeline_config import (
+    LeakyReLULoweringConfig,
+    SoftmaxDecompositionConfig,
+)
 from executorch.backends.arm.ethosu import EthosUCompileSpec
 from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec
 from executorch.backends.arm.vgf import VgfCompileSpec
@@ -72,11 +75,18 @@ def test_compile_spec_vgf_no_quant():
         EthosUCompileSpec._from_list(spec_list)
 
 
-def test_compile_spec_vgf_uses_default_pipeline_config():
+def test_compile_spec_vgf_defaults_leaky_relu_to_decompose():
     compile_spec = VgfCompileSpec()
     pipeline_config = compile_spec._get_pass_pipeline_config()
 
-    assert pipeline_config.is_default()
+    assert pipeline_config.leaky_relu is LeakyReLULoweringConfig.DECOMPOSE
+
+
+def test_compile_spec_tosa_defaults_leaky_relu_to_decompose():
+    compile_spec = TosaCompileSpec("TOSA-1.0+INT")
+    pipeline_config = compile_spec._get_pass_pipeline_config()
+
+    assert pipeline_config.leaky_relu is LeakyReLULoweringConfig.DECOMPOSE
 
 
 def test_compile_spec_tosa_INT():

backends/arm/test/misc/test_pass_pipeline_config.py

@@ -13,9 +13,12 @@
 from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager
 from executorch.backends.arm.common.pipeline_config import (
     ArmPassPipelineConfig,
+    LeakyReLULoweringConfig,
     QuantizeInfConfig,
     SoftmaxDecompositionConfig,
 )
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
 from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec
 from executorch.backends.arm.tosa.specification import TosaSpecification
 from torch.export import export
@@ -35,6 +38,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+class ModuleWithLeakyReLU(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.leaky_relu(x, negative_slope=0.2)
+
+
+def _call_function_targets(graph_module: torch.fx.GraphModule) -> list[object]:
+    return [
+        node.target for node in graph_module.graph.nodes if node.op == "call_function"
+    ]
+
+
 def test_pipeline_config_override_outside_compile_spec():
     compile_spec = TosaCompileSpec(
         TosaSpecification.create_from_string("TOSA-1.00+INT")
@@ -110,3 +124,87 @@ def test_quant_inf_config_reaches_annotation_pipeline():
     )
 
     assert tensor_constant_values == [QUANT_NEG_INF, QUANT_POS_INF]
+
+
+def test_leaky_relu_config_serializes_roundtrip():
+    config = ArmPassPipelineConfig(
+        leaky_relu=LeakyReLULoweringConfig.DECOMPOSE,
+    )
+    assert config.is_default()
+
+    from_dict = ArmPassPipelineConfig.from_dict(config.to_dict())
+    assert from_dict.leaky_relu is LeakyReLULoweringConfig.DECOMPOSE
+
+    compile_spec = TosaCompileSpec(
+        TosaSpecification.create_from_string("TOSA-1.00+INT")
+    )
+    compile_spec.set_pass_pipeline_config(config)
+    roundtripped = TosaCompileSpec._from_list(compile_spec._to_list())
+
+    assert (
+        roundtripped._get_pass_pipeline_config().leaky_relu
+        is LeakyReLULoweringConfig.DECOMPOSE
+    )
+
+
+def test_leaky_relu_table_config_reaches_annotation_pipeline():
+    config = ArmPassPipelineConfig(
+        leaky_relu=LeakyReLULoweringConfig.TABLE,
+    )
+    compile_spec = TosaCompileSpec(
+        TosaSpecification.create_from_string("TOSA-1.00+INT")
+    )
+    compile_spec.set_pass_pipeline_config(config)
+    manager = ArmPassManager(compile_spec)
+    exported = export(ModuleWithLeakyReLU(), (torch.randn(4),), strict=True)
+
+    transformed = manager.transform_for_annotation_pipeline(exported.graph_module)
+    targets = _call_function_targets(transformed)
+
+    assert torch.ops.aten.leaky_relu.default in targets
+
+
+def test_leaky_relu_decompose_config_reaches_annotation_pipeline():
+    config = ArmPassPipelineConfig(
+        leaky_relu=LeakyReLULoweringConfig.DECOMPOSE,
+    )
+    compile_spec = TosaCompileSpec(
+        TosaSpecification.create_from_string("TOSA-1.00+INT")
+    )
+    compile_spec.set_pass_pipeline_config(config)
+    manager = ArmPassManager(compile_spec)
+    exported = export(ModuleWithLeakyReLU(), (torch.randn(4),), strict=True)
+
+    transformed = manager.transform_for_annotation_pipeline(exported.graph_module)
+    targets = _call_function_targets(transformed)
+
+    assert torch.ops.aten.leaky_relu.default not in targets
+    assert torch.ops.aten.clamp.default in targets
+    assert torch.ops.aten.mul.Tensor in targets
+    assert torch.ops.aten.add.Tensor in targets
+
+
+def test_leaky_relu_decompose_config_reaches_backend_pipeline():
+    config = ArmPassPipelineConfig(
+        leaky_relu=LeakyReLULoweringConfig.DECOMPOSE,
+    )
+    compile_spec = common.get_tosa_compile_spec("TOSA-1.00+INT")
+    compile_spec.set_pass_pipeline_config(config)
+    tester = ArmTester(
+        ModuleWithLeakyReLU(),
+        example_inputs=(torch.linspace(-1.0, 1.0, 16),),
+        compile_spec=compile_spec,
+    )
+
+    tester.quantize().export()
+    exported_program = tester.get_artifact()
+    graph_module = ArmPassManager(compile_spec).transform_to_backend_pipeline(
+        exported_program, exported_program.graph_module
+    )
+    graph_code = graph_module.code
+
+    assert "torch.ops.aten.leaky_relu.default" not in graph_code
+    assert "torch.ops.backend.tosa.TABLE.default" not in graph_code
+    assert "torch.ops.aten.clamp.default" in graph_code
+    assert "torch.ops.aten.mul.Tensor" in graph_code
+    assert "torch.ops.aten.add.Tensor" in graph_code

backends/arm/test/ops/test_leaky_relu.py

@@ -61,7 +61,6 @@ def test_leaky_relu_tosa_INT(test_data):
         [],
         use_to_edge_transform_and_lower=True,
     )
-    pipeline.add_stage_after("quantize", pipeline.tester.check_not, [aten_op])
     pipeline.run()
 
 
@@ -75,7 +74,6 @@ def test_leaky_relu_u55_INT(test_data):
         [],
         use_to_edge_transform_and_lower=True,
     )
-    pipeline.add_stage_after("quantize", pipeline.tester.check_not, [aten_op])
     pipeline.run()
 
 
@@ -89,7 +87,6 @@ def test_leaky_relu_u85_INT(test_data):
         [],
         use_to_edge_transform_and_lower=True,
     )
-    pipeline.add_stage_after("quantize", pipeline.tester.check_not, [aten_op])
     pipeline.run()
 
 
@@ -121,5 +118,4 @@ def test_leaky_relu_vgf_quant(test_data):
         use_to_edge_transform_and_lower=True,
         quantize=True,
     )
-    pipeline.add_stage_after("quantize", pipeline.tester.check_not, [aten_op])
     pipeline.run()