Arm backend: Add TOSA (DW) Conv2d dialect operator (#14843)

Adds backend dialect operators for:
- Conv2D
- Depthwise Conv2D

Also removes AddBiasPass and moves that logic to new RewriteConv2dPass.

Signed-off-by: Oscar Andersson <[email protected]>

Author: oscarandersson8218

backends/arm/_passes/__init__.py

@@ -6,7 +6,6 @@
 
 from . import arm_pass_utils  # noqa
 from .arm_pass import ArmPass  # noqa  # usort: skip
-from .add_bias_pass import AddBiasPass  # noqa
 from .annotate_decomposed_matmul import AnnotateDecomposedMatmulPass  # noqa
 from .annotate_output_dim_order_pass import AnnotateOutputDimOrderPass  # noqa
 from .broadcast_args_pass import BroadcastArgsPass  # noqa
@@ -92,6 +91,7 @@
     ReplaceScalarWithTensorArgPassTOSABI,
     ReplaceScalarWithTensorArgPassTOSAMI,
 )
+from .rewrite_conv2d_pass import RewriteConv2dPass  # noqa
 from .rewrite_matmul import RewriteMatmulPass  # noqa
 from .rewrite_upsample import RewriteUpsamplePass  # noqa
 from .scalars_to_attribute_pass import ScalarsToAttributePass  # noqa

backends/arm/_passes/add_bias_pass.py

@@ -12,7 +12,6 @@
 
 import executorch.backends.arm.tosa.dialect  # noqa: unused
 from executorch.backends.arm._passes import (
-    AddBiasPass,
     AnnotateDecomposedMatmulPass,
     AnnotateOutputDimOrderPass,
     BroadcastArgsPass,
@@ -93,6 +92,7 @@
     ReplaceScalarWithTensorArgPassTOSABI,
     ReplaceScalarWithTensorArgPassTOSAMI,
     RetraceFoldedDtypesPass,
+    RewriteConv2dPass,
     RewriteMatmulPass,
     RewriteUpsamplePass,
     ScalarsToAttributePass,
@@ -207,13 +207,13 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(InsertTableOpsPass(exported_program))
         # If we have a conv2d with int16 activation split up into a convolution
         # and an addition, to work-around the lack of support for int48 in torch
-        # needs to happen before AddBiasPass, but after the table ops are inserted
+        # needs to happen before RewriteConv2dPass, but after the table ops are inserted
         # to be able to validate that conv2d has right dtype arguments.
         self.add_pass(DecomposeConv2dWithInt16ActivationPass())
-        self.add_pass(RewriteUpsamplePass())
-        self.add_pass(AddBiasPass(exported_program))
+        self.add_pass(RewriteConv2dPass(exported_program))
 
         self.add_pass(RewriteMatmulPass())
+        self.add_pass(RewriteUpsamplePass())
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))
         self.add_pass(ToTosaMemoryFormatPass(exported_program))
         self.add_pass(RemoveNoopPass())
@@ -297,9 +297,9 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
 
         self.add_pass(FuseViewCopyTransform())
         self.add_pass(FuseConstantArgsPass(exported_program))
+        self.add_pass(RewriteConv2dPass(exported_program))
         self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
         self.add_pass(RewriteUpsamplePass())
-        self.add_pass(AddBiasPass(exported_program))
         self.add_pass(InsertTableOpsPass(exported_program))
         self.add_pass(RewriteMatmulPass())
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))

backends/arm/_passes/arm_pass_manager.py

@@ -10,7 +10,7 @@
 
 from executorch.backends.arm._passes import ArmPass
 
-from executorch.backends.arm._passes.add_bias_pass import AddBiasPass
+from executorch.backends.arm._passes.rewrite_conv2d_pass import RewriteConv2dPass
 from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass
 
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -28,7 +28,10 @@ class Conv1dUnsqueezePass(ArmPass):
     3) squeeze the output back down to 3d.
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = {AddBiasPass, SizeAdjustInputPass}
+    _passes_required_after: Set[Type[ExportPass]] = {
+        RewriteConv2dPass,
+        SizeAdjustInputPass,
+    }
 
     def call_operator(self, op, args, kwargs, meta):
         if op != exir_ops.edge.aten.convolution.default:

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -8,9 +8,9 @@
 
 import torch
 from executorch.backends.arm._passes import ArmPass
-from executorch.backends.arm._passes.add_bias_pass import AddBiasPass
 from executorch.backends.arm._passes.arm_pass_utils import create_node
 from executorch.backends.arm._passes.quant_args import QuantArgs
+from executorch.backends.arm._passes.rewrite_conv2d_pass import RewriteConv2dPass
 
 from executorch.backends.transforms.utils import create_constant_placeholder
 from executorch.exir import ExportedProgram
@@ -42,7 +42,7 @@ class DecomposeCumsumPass(ArmPass):
     And the convolution is applied over dimension H.
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = {AddBiasPass}
+    _passes_required_after: Set[Type[ExportPass]] = {RewriteConv2dPass}
 
     def __init__(self, exported_program: ExportedProgram) -> None:
         super().__init__()