XNNPack: Ignore Alpha!=1 for add nodes (fixes: #11683)

Author: abhinaykukkadapu

backends/xnnpack/partition/config/generic_node_configs.py

@@ -9,6 +9,7 @@
 import logging
 from typing import cast, List, Optional
 
+import numpy as np
 import torch
 from executorch.backends.xnnpack.partition.config.xnnpack_config import (
     ConfigPrecisionType,
@@ -106,6 +107,13 @@ def __init__(self, **kwargs):
     def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
 
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        # No support for add nodes with alpha != 1
+        if "alpha" in node.kwargs and not np.isclose(node.kwargs["alpha"], 1.0, atol=1e-9, rtol=1e-9):
+            why(node, reason="Add node doesn't support alpha != 1")
+            return False
+        return True
+
 
 class ReLUConfig(GenericNodePartitionerConfig):
     target_name = "relu.default"
@@ -397,7 +405,7 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
                 node,
                 reason=f"mean.dim only supports averaging 4D tensors, got tensor of rank {output_dims}",
             )
-            return False
+          return False
         return True
 
     def supported_precision_types(self) -> List[ConfigPrecisionType]:

backends/xnnpack/test/ops/test_add.py

@@ -240,3 +240,27 @@ def forward(self, x, z):
             .serialize()
             .run_method_and_compare_outputs()
         )
+
+    class AddWithAlpha(torch.nn.Module):
+        def forward(self, x, y):
+            # node with alpha = 1.0 will be partitioned
+            out1 = torch.add(x, y, alpha=1)
+            # node with alpha != 1.0 will not be partitioned
+            out2 = torch.add(x, y, alpha=2)
+            return out1, out2
+
+    def test_add_with_alpha(self):
+        inputs = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 4))
+        (
+            Tester(self.AddWithAlpha(), inputs)
+            .export()
+            .check_count({"torch.ops.aten.add.Tensor": 2})
+            .to_edge_transform_and_lower()
+            # unpartitioned node
+            .check_count({"executorch_exir_dialects_edge__ops_aten_add_Tensor": 1})
+            # partitioned node
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )

backends/xnnpack/test/ops/test_linear.py

@@ -551,8 +551,9 @@ def test_qd8_f32_per_channel_shared_dq_chain(self):
                 inputs,
                 dynamic_shapes=None,
                 is_per_channel=True,
-                linear_count=2,
+                linear_count=3,  # 2 linear + 1 add
                 uses_bias=use_bias,
+                atol=1e-1,
             )
 
     def _test_qd8_per_channel_linear(self, dtype: torch.dtype = torch.float):
@@ -872,7 +873,7 @@ def test_qd8_per_channel_linear_parallel(self):
                 ),
                 inputs,
                 dynamic_shapes=dynamic_shapes,
-                linear_count=2,
+                linear_count=3,  # 2 linear + 1 add
                 is_per_channel=True,
                 uses_bias=True,
             )