[xnnpack] Fold float dtype conversions into the delegate

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -112,7 +112,9 @@ def mark_as_nchw_node(node: torch.fx.Node) -> None:
         node.meta[ChannelsLastTaggedReshapePass.XNN_NHWC_NODE] = False
 
     def tag_node(self, node: torch.fx.Node) -> None:
-        if node.kwargs["memory_format"] == torch.channels_last:
+        # A dtype-only _to_copy (e.g. x.float()) has no memory_format kwarg and
+        # does not change layout; treat it as contiguous (nchw).
+        if node.kwargs.get("memory_format") == torch.channels_last:
             self.mark_as_nhwc_node(node)
         else:
             self.mark_as_nchw_node(node)

backends/xnnpack/_passes/remove_redundant_copy_pass.py

@@ -25,6 +25,18 @@ def _try_remove_regular_redundant_to_copy(self, node, graph):
         """
         input_node = node.args[0]
 
+        # These removals assume the copies are memory-format-only (dtype
+        # preserving).
+        def _changes_dtype(n):
+            return (
+                getattr(n, "op", None) == "call_function"
+                and n.target == exir_ops.edge.aten._to_copy.default
+                and n.args[0].meta["val"].dtype != n.meta["val"].dtype
+            )
+
+        if _changes_dtype(node) or _changes_dtype(input_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

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,17 @@ def define_node(
         vals_to_ids: Dict[torch.fx.Node, int],
         debug_handle: int,
     ) -> None:
+        input_node = get_input_node(node, 0)
+        input_dtype = input_node.meta["val"].dtype
+        output_dtype = node.meta["val"].dtype
+
+        # A pure dtype conversion is serialized as an xnn_define_convert.
+        if input_dtype != output_dtype:
+            self._define_dtype_convert(
+                node, input_node, xnn_graph, vals_to_ids, 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)
@@ -89,3 +101,36 @@ def define_node(
             debug_handle=debug_handle,
         )
         xnn_graph.xnodes.append(ser_node)
+
+    def _define_dtype_convert(
+        self,
+        node: torch.fx.Node,
+        input_node: torch.fx.Node,
+        xnn_graph: XNNGraph,
+        vals_to_ids: Dict[torch.fx.Node, int],
+        debug_handle: int,
+    ) -> None:
+        # Input and output tensors keep their own (differing) dtypes; the
+        # convert node bridges them via xnn_define_convert at runtime.
+        self.define_tensor(
+            input_node,
+            xnn_graph,
+            vals_to_ids,
+            quant_params=QuantParams.from_inputs(input_node, self._exported_program),
+        )
+        self.define_tensor(
+            node,
+            xnn_graph,
+            vals_to_ids,
+            quant_params=QuantParams.from_outputs(node),
+        )
+
+        ser_node = XNode(
+            xnode_union=XNNConvert(
+                input_id=vals_to_ids[input_node],
+                output_id=vals_to_ids[node],
+                flags=0,
+            ),
+            debug_handle=debug_handle,
+        )
+        xnn_graph.xnodes.append(ser_node)

backends/xnnpack/partition/config/__init__.py

@@ -56,6 +56,7 @@
     SquareRootConfig,
     SubConfig,
     TanhConfig,
+    ToCopyConfig,
     ToDimOrderCopyConfig,
     UnsqueezeCopyConfig,
     UpsampleBilinear2dConfig,
@@ -114,6 +115,7 @@
     ReLUConfig,
     TanhConfig,
     ToDimOrderCopyConfig,
+    ToCopyConfig,
     SigmoidConfig,
     SinConfig,
     CosConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -23,6 +23,7 @@
 )
 from executorch.backends.xnnpack.utils.utils import (
     get_input_node,
+    is_param_node,
     normalize_mean_dims,
     normalize_pool2d_args,
 )
@@ -503,9 +504,17 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
 class ToDimOrderCopyConfig(GenericNodePartitionerConfig):
     target_name = "_to_dim_order_copy.default"
 
+    _SUPPORTED_DTYPE_CONVERSIONS = {
+        (torch.float32, torch.float16),
+        (torch.float16, torch.float32),
+        (torch.float32, torch.bfloat16),
+        (torch.bfloat16, torch.float32),
+    }
+
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
-        Only support dim order conversion partitioning, not DType conversions
+        Support dim-order conversion, plus the float dtype conversions XNNPACK
+        can fold into the delegate via xnn_define_convert.
         """
         if not self.check_common_constraints(node, ep):
             return False
@@ -515,8 +524,19 @@ 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:
+        # Dim-order-only conversion (no dtype change) is always supported.
+        if input_dtype == output_dtype:
+            return True
+
+        # Only fold dtype conversions of dynamic activations. Conversions of
+        # params/constants (e.g. weight.float()) should be constant-folded, not
+        # run as a runtime xnn_define_convert.
+        if is_param_node(ep, input_node):
+            why(node, reason="dtype conversion of a constant/param is not folded")
+            return False
+
+        # Otherwise only partition dtype conversions XNNPACK can serialize.
+        if (input_dtype, output_dtype) not in self._SUPPORTED_DTYPE_CONVERSIONS:
             why(
                 node,
                 reason=f"dtype conversion from {input_dtype} to {output_dtype} is not supported",
@@ -529,6 +549,11 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
 
 
+class ToCopyConfig(ToDimOrderCopyConfig):
+    # Non-dim-order form, used under EdgeCompileConfig(_skip_dim_order=True).
+    target_name = "_to_copy.default"
+
+
 class MeanDimConfig(GenericNodePartitionerConfig):
     target_name = "mean.dim"