Arm backend: Add MXFP support for Conv2d (#20421)

Add the possibility to convert `torch.nn.Conv2d` submodules to the
custom implemented MXFP counterpart `MXFPConv2dOp`.

Rewrite the MXFP Conv2d custom op into block-scaled TOSA ops and
fix the serializer, typing.

Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>
Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>

Author: martinlsm

backends/arm/_passes/__init__.py

@@ -171,6 +171,7 @@
 from .rewrite_le_lt_to_ge_gt_pass import RewriteLeLtToGeGtPass  # noqa
 from .rewrite_matmul import RewriteMatmulPass  # noqa
 from .rewrite_max_pool2d_pass import RewriteMaxPool2dPass  # noqa
+from .rewrite_mxfp_conv2d import RewriteMXFPConv2dPass  # noqa
 from .rewrite_mxfp_linear import RewriteMXFPLinearPass  # noqa
 from .rewrite_pad import RewritePadPass  # noqa
 from .rewrite_slice import RewriteSlicePass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -147,6 +147,7 @@
     RewriteLeLtToGeGtPass,
     RewriteMatmulPass,
     RewriteMaxPool2dPass,
+    RewriteMXFPConv2dPass,
     RewriteMXFPLinearPass,
     RewritePadPass,
     RewriteSlicePass,
@@ -612,6 +613,7 @@ def _tosa_pipeline(
                 RewriteMaxPool2dPass(),
                 DecomposeAdaptiveMaxPool2dPass(),
                 RewriteConvPass(exported_program),
+                RewriteMXFPConv2dPass(exported_program),
                 RewriteMXFPLinearPass(exported_program),
                 RewriteMatmulPass(),
                 RewritePadPass(),