[XNNPACK] Fold zero constant_pad_nd into convolutions, lower fp32<->fp16 casts as convert

backends/xnnpack/_passes/__init__.py

@@ -28,6 +28,9 @@
 from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
+from executorch.backends.xnnpack._passes.fuse_pad_into_conv_pass import (
+    FusePadIntoConvPass,
+)
 from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass
 from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
 from executorch.backends.xnnpack._passes.propagate_custom_meta_pass import (
@@ -90,6 +93,7 @@ def __init__(
                 PReLUReshapePass,
                 ChannelsLastTaggedReshapePass,
                 RemoveRedundantCopyPass,
+                FusePadIntoConvPass,
             ]
         else:
             self.passes = passes

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -173,6 +173,18 @@ def _is_nhwc_tensor(tensor: torch.Tensor) -> bool:
     def _is_nhwc(self, tensor: torch.Tensor) -> bool:
         return ChannelsLastTaggedReshapePass._is_nhwc_tensor(tensor)
 
+    @staticmethod
+    def is_dtype_cast(node: torch.fx.Node) -> bool:
+        """A _to_copy that changes dtype (fp32 <-> fp16) keeps the layout of
+        its input and is serialized as a convert node, not a transpose."""
+        src = node.args[0]
+        return (
+            isinstance(src, torch.fx.Node)
+            and "val" in src.meta
+            and "val" in node.meta
+            and node.meta["val"].dtype != src.meta["val"].dtype
+        )
+
     def requires_nhwc_input(self, node: torch.fx.Node) -> bool:
         return node.target in self.memory_sensitive_ops_nhwc
 
@@ -562,7 +574,10 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                 # The node requires nchw inputs
                 for input_node in node.all_input_nodes:
                     self.input_to_nchw(graph_module, input_node, node)
-            elif node.target == exir_ops.edge.aten._to_copy.default:
+            elif (
+                node.target == exir_ops.edge.aten._to_copy.default
+                and not self.is_dtype_cast(node)
+            ):
                 self.tag_node(node)
             else:
                 # The node can have inputs in any format (but all must be the

backends/xnnpack/_passes/fuse_pad_into_conv_pass.py

@@ -0,0 +1,101 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import cast, List, Tuple
+
+import torch
+from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
+from executorch.backends.xnnpack.utils.quant_utils import is_dequant, is_quant
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
+
+
+class FusePadIntoConvPass(XNNPACKPass):
+    """
+    Folds a zero-valued constant_pad_nd on the two spatial dims into the
+    convolution that consumes it. XNNPACK convolutions take independent
+    top / right / bottom / left padding, so asymmetric pads such as
+    F.pad(x, (2, 1, 2, 1)) before a 4x4 filter fold as well, and the padded
+    copy of the activation (a full extra pass over the tensor, noticeably slow
+    in NHWC) disappears.
+
+    The graph is left intact so shape metadata stays valid: the pad node is
+    tagged FOLDED_PAD_TAG and its visitor aliases the pad output to the pad
+    input instead of emitting a node, and the convolution is tagged with the
+    combined (top, right, bottom, left) padding that its visitor serializes.
+
+    Must run after ChannelsLastTaggedReshapePass: if that pass had to insert
+    a layout conversion between the pad and the convolution, the pad's only
+    user is no longer the convolution and nothing is fused.
+    """
+
+    FOLDED_PAD_TAG = "XNN_FOLDED_PAD"
+    FUSED_PADDING_TAG = "XNN_CONV_PADDING_TRBL"
+
+    @staticmethod
+    def get_fused_padding(node: torch.fx.Node):
+        return node.meta.get(FusePadIntoConvPass.FUSED_PADDING_TAG, None)
+
+    @staticmethod
+    def is_folded_pad(node: torch.fx.Node) -> bool:
+        return node.meta.get(FusePadIntoConvPass.FOLDED_PAD_TAG, False)
+
+    def _foldable(self, pad: torch.fx.Node) -> Tuple[bool, List[int]]:
+        """Returns (ok, [left, right, top, bottom]) for a 4D spatial zero pad."""
+        if len(pad.users) != 1:
+            return False, [0, 0, 0, 0]
+        if len(pad.args) > 2 and cast(float, pad.args[2]) != 0.0:
+            return False, [0, 0, 0, 0]
+        src = pad.args[0]
+        if not isinstance(src, torch.fx.Node) or "val" not in src.meta:
+            return False, [0, 0, 0, 0]
+        if src.meta["val"].dim() != 4:
+            return False, [0, 0, 0, 0]
+        amounts = list(cast(List[int], pad.args[1]))
+        if len(amounts) % 2 or len(amounts) > 4 or any(a < 0 for a in amounts):
+            return False, [0, 0, 0, 0]
+        # Quantized graphs route the pad through q/dq pairs (InsertPadQDQPass);
+        # leave those alone.
+        if is_dequant(src) or any(is_quant(u) for u in pad.users):
+            return False, [0, 0, 0, 0]
+        amounts = amounts + [0] * (4 - len(amounts))  # [left, right, top, bottom]
+        return True, amounts
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        modified = False
+        for pad in graph_module.graph.nodes:
+            if (
+                pad.op != "call_function"
+                or pad.target != exir_ops.edge.aten.constant_pad_nd.default
+            ):
+                continue
+            ok, (left, right, top, bottom) = self._foldable(pad)
+            if not ok:
+                continue
+            conv = next(iter(pad.users))
+            if (
+                conv.op != "call_function"
+                or conv.target != exir_ops.edge.aten.convolution.default
+                or conv.args[0] is not pad
+                or conv.args[6]  # transposed convolution pads differently
+                or self.get_fused_padding(conv) is not None
+            ):
+                continue
+            conv_padding = list(cast(List[int], conv.args[4]))
+            if len(conv_padding) == 1:
+                conv_padding = conv_padding + conv_padding
+            if len(conv_padding) != 2:
+                continue
+            conv.meta[self.FUSED_PADDING_TAG] = (
+                conv_padding[0] + top,
+                conv_padding[1] + right,
+                conv_padding[0] + bottom,
+                conv_padding[1] + left,
+            )
+            pad.meta[self.FOLDED_PAD_TAG] = True
+            modified = True
+        # Tags only, no graph change: the metadata stays as traced.
+        return PassResult(graph_module, modified)

backends/xnnpack/_passes/remove_redundant_copy_pass.py

@@ -25,9 +25,15 @@ def _try_remove_regular_redundant_to_copy(self, node, graph):
         """
         input_node = node.args[0]
 
+        # A dtype cast (fp32 <-> fp16) is not a layout copy: it is serialized
+        # as a convert node and must stay, whatever the tags around it say.
+        if ChannelsLastTaggedReshapePass.is_dtype_cast(node):
+            return False
+
         # Check if input is a to_copy with opposite memory format
         if (
             input_node.target == exir_ops.edge.aten._to_copy.default
+            and not ChannelsLastTaggedReshapePass.is_dtype_cast(input_node)
             and ChannelsLastTaggedReshapePass.is_nchw_node(input_node)
             != ChannelsLastTaggedReshapePass.is_nchw_node(node)
             and len(input_node.users) == 1

backends/xnnpack/operators/op_conv2d.py

@@ -9,6 +9,9 @@
 import torch
 from executorch.backends.transforms import get_shape
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
+from executorch.backends.xnnpack._passes.fuse_pad_into_conv_pass import (
+    FusePadIntoConvPass,
+)
 from executorch.backends.xnnpack.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
@@ -155,10 +158,17 @@ def define_node(
         check_or_raise(
             len(stride) == 2, "XNNPACK currently only supports 2D convolution"
         )
-        kwargs["padding_top"] = padding[0]
-        kwargs["padding_right"] = padding[1]
-        kwargs["padding_bottom"] = padding[0]
-        kwargs["padding_left"] = padding[1]
+        (
+            kwargs["padding_top"],
+            kwargs["padding_right"],
+            kwargs["padding_bottom"],
+            kwargs["padding_left"],
+        ) = FusePadIntoConvPass.get_fused_padding(node) or (
+            padding[0],
+            padding[1],
+            padding[0],
+            padding[1],
+        )
         kwargs["kernel_height"] = kernel_shape[2]
         kwargs["kernel_width"] = kernel_shape[3]
         kwargs["subsampling_height"] = stride[0]

backends/xnnpack/operators/op_static_constant_pad.py

@@ -8,6 +8,9 @@
 
 import torch
 
+from executorch.backends.xnnpack._passes.fuse_pad_into_conv_pass import (
+    FusePadIntoConvPass,
+)
 from executorch.backends.xnnpack.operators.node_visitor import (
     get_tensor_value,
     NodeVisitor,
@@ -39,6 +42,14 @@ def define_node(
         vals_to_ids: Dict[torch.fx.Node, int],
         debug_handle: int,
     ) -> None:
+        if FusePadIntoConvPass.is_folded_pad(node):
+            # Folded into the consuming convolution's padding: no node, the
+            # pad output is the pad input as far as XNNPACK is concerned.
+            input_node = get_input_node(node, 0)
+            self.define_tensor(input_node, xnn_graph, vals_to_ids)
+            vals_to_ids[node] = vals_to_ids[input_node]
+            return
+
         self.define_nodes_tensor_inputs_outputs(node, xnn_graph, vals_to_ids)
 
         # input

backends/xnnpack/operators/op_to_copy.py

@@ -16,6 +16,7 @@
 from executorch.backends.xnnpack.operators.quant_params import QuantParams
 
 from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    XNNConvert,
     XNNGraph,
     XNNStaticTranspose,
     XNode,
@@ -42,6 +43,25 @@ def define_node(
         vals_to_ids: Dict[torch.fx.Node, int],
         debug_handle: int,
     ) -> None:
+        input_node = get_input_node(node, 0)
+        if node.meta["val"].dtype != input_node.meta["val"].dtype:
+            # fp32 <-> fp16 cast (the partitioner admits no other dtype
+            # conversion and no dim-order change on the same node): an XNNPACK
+            # convert node, the layout follows the tags on the two tensors.
+            self.define_tensor(input_node, xnn_graph, vals_to_ids)
+            self.define_tensor(node, xnn_graph, vals_to_ids)
+            xnn_graph.xnodes.append(
+                XNode(
+                    xnode_union=XNNConvert(
+                        input_id=vals_to_ids[input_node],
+                        output_id=vals_to_ids[node],
+                        flags=0,
+                    ),
+                    debug_handle=debug_handle,
+                )
+            )
+            return
+
         memory_format_target = node.kwargs.get("memory_format", torch.contiguous_format)
         to_channels_last = bool(memory_format_target == torch.channels_last)
         to_contiguous = bool(memory_format_target == torch.contiguous_format)
@@ -50,7 +70,6 @@ def define_node(
             "Unsupported Memory Format for XNNPACK",
         )
 
-        input_node = get_input_node(node, 0)
         input_quant_params = QuantParams.from_inputs(input_node, self._exported_program)
         output_quant_params = QuantParams.from_outputs(node)
 

backends/xnnpack/partition/config/generic_node_configs.py

@@ -620,13 +620,20 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         input_dtype = input_node.meta["val"].dtype
         output_dtype = node.meta["val"].dtype
 
-        # Return False if doing dtype conversion
         if input_dtype != output_dtype:
-            why(
-                node,
-                reason=f"dtype conversion from {input_dtype} to {output_dtype} is not supported",
+            # fp32 <-> fp16 is lowered to an XNNPACK convert node, as long as
+            # the op does not change the dim order at the same time.
+            float_pair = {input_dtype, output_dtype} <= {torch.float32, torch.float16}
+            dim_order = node.kwargs.get("dim_order", None)
+            default_order = dim_order is None or list(dim_order) == list(
+                range(len(dim_order))
             )
-            return False
+            if not (float_pair and default_order):
+                why(
+                    node,
+                    reason=f"dtype conversion from {input_dtype} to {output_dtype} is not supported",
+                )
+                return False
 
         return True
 

backends/xnnpack/test/ops/test_to_copy_dtype.py

@@ -0,0 +1,79 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.backends.xnnpack.test.tester import Tester
+
+
+class TestToCopyDtype(unittest.TestCase):
+    """fp32 <-> fp16 casts are lowered to XNNPACK convert nodes."""
+
+    def setUp(self):
+        torch._dynamo.reset()
+
+    class Fp16Inner(torch.nn.Module):
+        """fp32 I/O contract around an fp16 model: the common deployment shape."""
+
+        def __init__(self):
+            super().__init__()
+            self.conv = torch.nn.Conv2d(3, 8, 3, padding=1).half()
+
+        def forward(self, x):
+            return self.conv(x.to(torch.float16)).to(torch.float32)
+
+    class CastOnly(torch.nn.Module):
+        def forward(self, x):
+            return (x.to(torch.float16) * 2.0).to(torch.float32)
+
+    cast_ops = [
+        "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default",
+        "executorch_exir_dialects_edge__ops_aten__to_copy_default",
+    ]
+
+    def _run(self, module, inputs):
+        (
+            Tester(module.eval(), inputs)
+            .export()
+            .to_edge_transform_and_lower()
+            .check_not(self.cast_ops)
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(atol=2e-2, rtol=2e-2)
+        )
+
+    def test_fp32_io_around_fp16_conv(self):
+        self._run(self.Fp16Inner(), (torch.randn(1, 3, 8, 8),))
+
+    def test_cast_only(self):
+        self._run(self.CastOnly(), (torch.randn(2, 5),))
+
+    def test_int_cast_not_partitioned(self):
+        class ToInt(torch.nn.Module):
+            def forward(self, x):
+                return x.to(torch.int32)
+
+        (
+            Tester(ToInt().eval(), (torch.randn(2, 5),))
+            .export()
+            .to_edge_transform_and_lower()
+            .check_not(["torch.ops.higher_order.executorch_call_delegate"])
+        )
+
+    def test_dtype_cast_with_dim_order_change_not_partitioned(self):
+        class ToHalfChannelsLast(torch.nn.Module):
+            def forward(self, x):
+                return x.to(torch.float16, memory_format=torch.channels_last)
+
+        (
+            Tester(ToHalfChannelsLast().eval(), (torch.randn(1, 3, 4, 4),))
+            .export()
+            .to_edge_transform_and_lower()
+            .check_not(["torch.ops.higher_order.executorch_call_delegate"])
+            .check_count({self.cast_ops[0]: 1})
+        )