diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index c74c35532b0..b782ba0382e 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/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) diff --git a/backends/xnnpack/_passes/remove_redundant_copy_pass.py b/backends/xnnpack/_passes/remove_redundant_copy_pass.py index b8f0a0dfbf0..9fc8b1f43e9 100644 --- a/backends/xnnpack/_passes/remove_redundant_copy_pass.py +++ b/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 diff --git a/backends/xnnpack/operators/op_to_copy.py b/backends/xnnpack/operators/op_to_copy.py index 2239524a518..9270b19b460 100644 --- a/backends/xnnpack/operators/op_to_copy.py +++ b/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) diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index c6c54f083d6..ab1838dd228 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/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, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index c7a3be5f65f..28af0e62c3d 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/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"