Sub, Sqrt, Pad

Differential Revision: D60492340

Pull Request resolved: #4533

Author: mcr229

backends/xnnpack/partition/config/__init__.py

@@ -20,6 +20,7 @@
     CatConfig,
     CeilConfig,
     ClampConfig,
+    ConstantPadConfig,
     DeQuantizedPerTensorConfig,
     DivConfig,
     FloorConfig,
@@ -40,6 +41,8 @@
     SigmoidConfig,
     SliceCopyConfig,
     SoftmaxConfig,
+    SquareRootConfig,
+    SubConfig,
     UpsampleBilinear2dConfig,
 )
 from executorch.backends.xnnpack.partition.config.node_configs import (
@@ -55,6 +58,7 @@
     # GEMM-like Configs
     AddmmConfig,
     LinearConfig,
+    ConstantPadConfig,
     ConvolutionConfig,
     # BatchNorm Config
     BatchNormConfig,
@@ -82,6 +86,8 @@
     SoftmaxConfig,
     SigmoidConfig,
     SliceCopyConfig,
+    SquareRootConfig,
+    SubConfig,
     PermuteConfig,
     # EluConfig, # Waiting for PyTorch Pin Update
     ReLUConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -382,3 +382,24 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
 
     def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
+
+
+class SquareRootConfig(GenericNodePartitionerConfig):
+    target_name = "sqrt.default"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+
+class ConstantPadConfig(GenericNodePartitionerConfig):
+    target_name = "constant_pad_nd.default"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+
+class SubConfig(GenericNodePartitionerConfig):
+    target_name = "sub.Tensor"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]

backends/xnnpack/test/ops/sqrt.py

@@ -25,9 +25,7 @@ def _test_sqrt(self, inputs):
             Tester(self.Sqrt(), inputs)
             .export()
             .check_count({"torch.ops.aten.sqrt.default": 1})
-            .to_edge()
-            .check_count({"executorch_exir_dialects_edge__ops_aten_sqrt_default": 1})
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(["executorch_exir_dialects_edge__ops_aten_sqrt_default"])
             .to_executorch()

backends/xnnpack/test/ops/static_constant_pad.py

@@ -88,11 +88,7 @@ def _test_static_constant_pad_functional(self, inputs):
             Tester(self.StaticConstantPadFunctional(), inputs)
             .export()
             .check_count({"torch.ops.aten.pad.default": 8})
-            .to_edge()
-            .check_count(
-                {"executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default": 8}
-            )
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 ["executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default"]
@@ -139,11 +135,7 @@ def forward(self, x):
             .export()
             .check_count({"torch.ops.aten.pad.default": 1})
             .check(["torch.ops.quantized_decomposed"])
-            .to_edge()
-            .check_count(
-                {"executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default": 1}
-            )
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 [
@@ -164,11 +156,7 @@ def test_qs8_static_constant_pad_2d(self):
             .export()
             .check_count({"torch.ops.aten.pad.default": 1})
             .check(["torch.ops.quantized_decomposed"])
-            .to_edge()
-            .check_count(
-                {"executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default": 1}
-            )
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 [

backends/xnnpack/test/ops/sub.py

@@ -32,9 +32,7 @@ def _test_sub(self, inputs):
             Tester(self.Sub(), inputs)
             .export()
             .check_count({"torch.ops.aten.sub.Tensor": 1})
-            .to_edge()
-            .check_count({"executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1})
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(["executorch_exir_dialects_edge__ops_aten_sub_Tensor"])
             .to_executorch()
@@ -62,9 +60,7 @@ def _test_qs8_sub(self):
             .export()
             .check_count({"torch.ops.aten.sub.Tensor": 1})
             .check(["torch.ops.quantized_decomposed"])
-            .to_edge()
-            .check_count({"executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1})
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 [
@@ -86,9 +82,7 @@ def _test_qs8_sub2(self):
             .export()
             .check_count({"torch.ops.aten.sub.Tensor": 1})
             .check(["torch.ops.quantized_decomposed"])
-            .to_edge()
-            .check_count({"executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1})
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 [
@@ -110,9 +104,7 @@ def _test_qs8_sub3(self):
             .export()
             .check_count({"torch.ops.aten.sub.Tensor": 1})
             .check(["torch.ops.quantized_decomposed"])
-            .to_edge()
-            .check_count({"executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1})
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 [
@@ -144,14 +136,7 @@ def forward(self, x, y):
                 }
             )
             .check(["torch.ops.quantized_decomposed"])
-            .to_edge()
-            .check_count(
-                {
-                    "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1,
-                    "executorch_exir_dialects_edge__ops_aten_relu_default": 1,
-                }
-            )
-            .partition()
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 [