From d5c2889355be02fbdddbb9c4818343f274c43d5f Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Wed, 29 Jul 2026 08:04:39 -0700 Subject: [PATCH] [xnnpack] Fold float dtype conversions into the delegate Partition and serialize fp16/bf16 <-> fp32 dtype-conversion copies into the XNNPACK delegate as xnn_define_convert, instead of leaving them as portable _to_copy ops that fragment the graph. This lets a mixed-precision model (e.g. a bf16 model whose RMSNorm computes in fp32) lower as cleanly as its fp32 counterpart: on google/gemma-3-1b-it bf16, delegate segments drop 446 -> 211 and portable dtype converts 364 -> 46, matching the fp32 structure, with identical outputs (argmax unchanged). Changes: - partition/config/generic_node_configs.py: ToDimOrderCopyConfig now allows the XNNPACK-serializable float conversions (fp32<->fp16, fp32<->bf16) for dynamic activations (constant/param converts are left for constant folding). Add ToCopyConfig (target _to_copy.default) so the same folding applies when the graph is lowered with EdgeCompileConfig(_skip_dim_order=True). - partition/config/__init__.py: register ToCopyConfig. - operators/op_to_copy.py: emit XNNConvert for a dtype-changing copy (memory format copies still emit XNNStaticTranspose). Two backend passes assumed _to_copy is always memory-format-only; partitioning dtype converts exposed both: - _passes/channels_last_tagged_reshape_pass.py: tag_node used node.kwargs['memory_format'] directly, which KeyErrors on a dtype-only copy (no memory_format kwarg). Default to contiguous (a dtype convert keeps layout). - _passes/remove_redundant_copy_pass.py: redundant-copy removal compared only memory format, so it would elide a genuine bf16<->fp32 convert as if it were a redundant layout copy, silently dropping the conversion. Skip copies that change dtype. Note: this is a graph-cleanliness/correctness change. On x86 it is latency neutral for bf16 (the bf16 cost there is intrinsic XNNPACK fp32-fallback for non-GEMM ops, not the converts); the win is a clean lowering that benefits memory-bound / ARM targets. --- .../channels_last_tagged_reshape_pass.py | 4 +- .../_passes/remove_redundant_copy_pass.py | 12 +++++ backends/xnnpack/operators/op_to_copy.py | 45 +++++++++++++++++++ backends/xnnpack/partition/config/__init__.py | 2 + .../partition/config/generic_node_configs.py | 31 +++++++++++-- 5 files changed, 90 insertions(+), 4 deletions(-) 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"