From bdf14d2fa8d9cb047eb022997ec59f443e447ad3 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Tue, 27 Oct 2026 09:03:25 +0100 Subject: [PATCH 01/11] [XNNPACK] Keep NHWC across scalar and low-rank broadcast operands Profiling a conv encoder-decoder (MI-GAN inpainting, fp16) on XNNPACK showed ~25% of the time spent in NCHW<->NHWC transposes around elementwise ops with a scalar or low-rank second operand. In ChannelsLastTaggedReshapePass, single-element operands of add/sub/mul/div/maximum/minimum (a lifted Python scalar such as the gain in x * 1.4142) no longer take part in the layout decision, and rank < 4 operands that broadcast over the trailing dims (a [H, W] map added to [N, C, H, W]) are reshaped to 4D with leading ones so they follow the NHWC operand instead of dragging the activation back to NCHW. Dequantized operands keep their rank, since a view between the dequantize and the op would carry no quant params, and operands with symbolic dims keep the NCHW path, since XNNStaticReshape cannot express their shape. --- .../channels_last_tagged_reshape_pass.py | 97 +++++++++- .../test_channels_last_tagged_reshape.py | 165 ++++++++++++++++++ 2 files changed, 255 insertions(+), 7 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 5238c9689ee..41b3df9b83a 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -22,6 +22,7 @@ from executorch.backends.xnnpack.utils.utils import is_param_node from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import PassResult +from torch.fx.experimental.symbolic_shapes import has_free_symbols class InputDimOrder(Enum): @@ -76,7 +77,9 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.linear.default, } - # Broadcasting binary ops whose operands must share one memory format. A + # Broadcasting binary ops whose operands must share one memory format. + # Single-element operands take no part in the layout decision, and rank < 4 + # operands are reshaped to 4D so they can follow an NHWC operand. A # dynamic-quant input_to_nhwc may blanket-replace a shared activation # (replace_all_uses_with), switching one operand of an already-processed # binary op to NHWC while its other (e.g. per-channel constant) operand @@ -86,6 +89,8 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.mul.Tensor, exir_ops.edge.aten.sub.Tensor, exir_ops.edge.aten.div.Tensor, + exir_ops.edge.aten.maximum.default, + exir_ops.edge.aten.minimum.default, } # Tag which is added to a node's meta to indicate that it uses NHWC format. @@ -173,6 +178,72 @@ def _is_nhwc_tensor(tensor: torch.Tensor) -> bool: def _is_nhwc(self, tensor: torch.Tensor) -> bool: return ChannelsLastTaggedReshapePass._is_nhwc_tensor(tensor) + @staticmethod + def static_tensor_val(node: torch.fx.Node) -> Optional[torch.Tensor]: + val = node.meta.get("val", None) + if isinstance(val, torch.Tensor) and not has_free_symbols(val.shape): + return val + return None + + @staticmethod + def is_single_element(node: torch.fx.Node) -> bool: + val = ChannelsLastTaggedReshapePass.static_tensor_val(node) + return val is not None and val.numel() == 1 + + @staticmethod + def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: + # A symbolic shape cannot be expressed in XNNStaticReshape, so such an + # operand takes the NCHW path. + val = ChannelsLastTaggedReshapePass.static_tensor_val(node) + return ( + val is not None + and 0 < val.dim() < 4 + and val.numel() > 1 + and not ChannelsLastTaggedReshapePass.is_nhwc_node(node) + # A view between dequantize and the op would carry no qparams. + and not is_dequant(node) + ) + + def unsqueeze_to_4d( + self, + graph_module: torch.fx.GraphModule, + input_node: torch.fx.Node, + target_node: torch.fx.Node, + ) -> torch.fx.Node: + """ + Reshape a rank < 4 operand of `target_node` to 4D by prepending ones, + which is how broadcasting aligns it, so it can be converted to NHWC. + """ + val = input_node.meta["val"] + new_shape = [1] * (4 - val.dim()) + list(val.shape) + with graph_module.graph.inserting_before(target_node): + view = self.create_call_function_node( + graph_module=graph_module, + target=exir_ops.edge.aten.view_copy.default, + args=(input_node, new_shape), + ) + view.meta["val"] = val.reshape(new_shape) + # Left untagged on purpose: it is NCHW, and the view_copy visitor + # treats the presence of the tag key as "permute the shape". + target_node.replace_input_with(input_node, view) + return view + + def layout_inputs(self, node: torch.fx.Node) -> List[torch.fx.Node]: + """ + Inputs that take part in the layout decision of `node`. Single-element + operands of a broadcasting binary op (e.g. the scalar in `x * 1.4142`) + broadcast the same way in either layout, so they must not force an + NHWC activation back to NCHW. + """ + return [ + input_node + for input_node in node.all_input_nodes + if not ( + node.target in self.broadcast_binary_ops + and self.is_single_element(input_node) + ) + ] + def requires_nhwc_input(self, node: torch.fx.Node) -> bool: return node.target in self.memory_sensitive_ops_nhwc @@ -567,28 +638,40 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 else: # The node can have inputs in any format (but all must be the # same format) + layout_inputs = self.layout_inputs(node) is_or_isnt_nhwc_node = [ ChannelsLastTaggedReshapePass.is_nhwc_node(input_node) - for input_node in node.all_input_nodes + for input_node in layout_inputs ] - if all(is_or_isnt_nhwc_node): + if is_or_isnt_nhwc_node and all(is_or_isnt_nhwc_node): # All inputs are nhwc so this node's output is nhwc too self.mark_as_nhwc_node(node) elif any(is_or_isnt_nhwc_node): # There is at least one node of each format, so we need to # convert some of them to another format such that all # are the same. + # Reshape lower-rank operands (e.g. a [H, W] map added to + # [N, C, H, W]) to 4D so they can follow the NHWC operand. + if node.target in self.broadcast_binary_ops: + layout_inputs = [ + ( + self.unsqueeze_to_4d(graph_module, input_node, node) + if self.is_lower_rank_broadcast_operand(input_node) + else input_node + ) + for input_node in layout_inputs + ] # If all nodes can be converted to nhwc, then convert them # all to nhwc. Otherwise, convert all to nchw if all( self.can_be_converted_to_nhwc(input_node) - for input_node in node.all_input_nodes + for input_node in layout_inputs ): - for input_node in node.all_input_nodes: + for input_node in layout_inputs: self.input_to_nhwc(graph_module, input_node, node) self.mark_as_nhwc_node(node) else: - for input_node in node.all_input_nodes: + for input_node in layout_inputs: self.input_to_nchw(graph_module, input_node, node) graph_module.recompile() @@ -612,7 +695,7 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 or node.target not in ChannelsLastTaggedReshapePass.broadcast_binary_ops ): continue - input_nodes = node.all_input_nodes + input_nodes = self.layout_inputs(node) if len(input_nodes) != 2: continue layouts = [ diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index bfe2ce862c9..0b5a0f5b4c9 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -290,6 +290,171 @@ def forward(self, x): x = torch.mean(x, (-1, -2), keepdim=True) return x + class ConvScalarMulConv(torch.nn.Module): + """conv -> x * scalar -> clamp -> conv: the scalar must not pull the + activation back to NCHW between the two convolutions.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + + def forward(self, x): + x = self.conv1(x) + x = torch.clamp(x * 1.4142, -256.0, 256.0) + return self.conv2(x) + + def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_stays_nhwc(self): + ( + Tester(self.ConvScalarMulConv().eval(), (torch.randn(1, 3, 8, 8),)) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC and output -> NCHW only + .check_count({self.to_copy_name: 2}) + .run_method_and_compare_outputs() + ) + + class ConvNoiseAddConv(torch.nn.Module): + """conv -> x + noise[H, W] -> conv: the rank-2 operand is reshaped to + [1, 1, H, W] and converted, instead of converting the activation.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("noise", torch.randn(8, 8)) + + def forward(self, x): + x = self.conv1(x) + x = x + self.noise + return self.conv2(x) + + class ConvChannelScaleConv(torch.nn.Module): + """conv -> x * scale[C, 1, 1] -> conv: a per-channel operand is + reshaped to [1, C, 1, 1] and converted.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("scale", torch.randn(8, 1, 1)) + + def forward(self, x): + x = self.conv1(x) + x = x * self.scale + return self.conv2(x) + + class ConvNoiseSubConv(torch.nn.Module): + """conv -> noise[H, W] - x -> conv: the lower-rank operand is the + left argument.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("noise", torch.randn(8, 8)) + + def forward(self, x): + x = self.conv1(x) + x = self.noise - x + return self.conv2(x) + + def _test_lower_rank_operand_stays_nhwc(self, module): + ( + Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC, the 4D view of the operand -> NHWC, output -> NCHW + .check_count( + { + self.to_copy_name: 3, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + } + ) + .run_method_and_compare_outputs() + ) + + def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_stays_nhwc(self): + self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseAddConv()) + + def test_fp32_channels_last_tagged_reshape_pass_per_channel_operand_stays_nhwc( + self, + ): + self._test_lower_rank_operand_stays_nhwc(self.ConvChannelScaleConv()) + + def test_fp32_channels_last_tagged_reshape_pass_lower_rank_left_operand_stays_nhwc( + self, + ): + self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseSubConv()) + + def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): + # The dequantized operand is left at its own rank: a view between + # dequantize and the op would carry no quant params and fail at load. + for module in (self.ConvNoiseAddConv(), self.ConvChannelScaleConv()): + ( + Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) + .quantize() + .export() + .to_edge() + .run_passes(self.PassStage) + .check_not( + ["executorch_exir_dialects_edge__ops_aten_view_copy_default"] + ) + .run_method_and_compare_outputs() + ) + + class ConvDynamicNoiseAddConv(torch.nn.Module): + """conv -> x + noise[H, W] -> conv with dynamic H and W: the operand + cannot be reshaped by XNNStaticReshape, so the add takes the NCHW path.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + + def forward(self, x, noise): + x = self.conv1(x) + x = x + noise + return self.conv2(x) + + def test_fp32_channels_last_tagged_reshape_pass_dynamic_lower_rank_operand(self): + height = torch.export.Dim("height", min=4, max=32) + width = torch.export.Dim("width", min=4, max=32) + dynamic_shapes = ({2: height, 3: width}, {0: height, 1: width}) + inputs = (torch.randn(1, 3, 8, 8), torch.randn(8, 8)) + ( + Tester( + self.ConvDynamicNoiseAddConv().eval(), + inputs, + dynamic_shapes=dynamic_shapes, + ) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC, conv1 output -> NCHW, add output -> NHWC, + # output -> NCHW + .check_count({self.to_copy_name: 4}) + .check_not(["executorch_exir_dialects_edge__ops_aten_view_copy_default"]) + .run_method_and_compare_outputs() + ) + ( + Tester( + self.ConvDynamicNoiseAddConv().eval(), + inputs, + dynamic_shapes=dynamic_shapes, + ) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs( + inputs=(torch.randn(1, 3, 12, 6), torch.randn(12, 6)) + ) + ) + def test_fp32_channels_last_tagged_reshape_pass_conv_bn_hardtanh_mean_seq(self): # Copy #1 is for input to conv, nchw -> nhwc # Copy #2 is for conv to _native_batch_norm_legit_no_training, nhwc -> nchw From 7f54e521d539d56f39b7a5417c3795ba4ca2b9a7 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Thu, 29 Oct 2026 11:23:10 +0100 Subject: [PATCH 02/11] [XNNPACK] Fold constant broadcast operands to NHWC ahead of time A single-element operand of rank > 4 made the op output rank > 4 while tagged NHWC, which failed at serialization. Such an operand now takes part in the layout decision again, so the op takes the NCHW path. A rank < 4 constant operand with a single user is reshaped to 4D in the exported program and tagged NHWC, so the serializer writes it in channels_last and the graph carries no reshape or transpose for it. Non-constant operands and constants with several users keep the view_copy + _to_copy path. Tests cover the 5D single-element operand, [H, W], [C, 1, 1] and [C, H, W] constants, a constant shared by two ops and a [H, W] graph input, each through to_edge_transform_and_lower, to_executorch, serialize and run. --- .../channels_last_tagged_reshape_pass.py | 52 +++++- .../test_channels_last_tagged_reshape.py | 167 ++++++++++++++++-- 2 files changed, 198 insertions(+), 21 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 41b3df9b83a..e4a6b3e1491 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -19,7 +19,7 @@ is_tagged_as_implicit_q_dq, tag_as_implicit_q_dq, ) -from executorch.backends.xnnpack.utils.utils import is_param_node +from executorch.backends.xnnpack.utils.utils import get_param_tensor, is_param_node from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import PassResult from torch.fx.experimental.symbolic_shapes import has_free_symbols @@ -187,8 +187,10 @@ def static_tensor_val(node: torch.fx.Node) -> Optional[torch.Tensor]: @staticmethod def is_single_element(node: torch.fx.Node) -> bool: + # A rank > 4 operand makes the op output rank > 4, which cannot be + # NHWC, so such an operand must still take part in the layout decision. val = ChannelsLastTaggedReshapePass.static_tensor_val(node) - return val is not None and val.numel() == 1 + return val is not None and val.numel() == 1 and val.dim() <= 4 @staticmethod def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: @@ -204,6 +206,44 @@ def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: and not is_dequant(node) ) + def is_foldable_constant(self, node: torch.fx.Node) -> bool: + """ + A lifted constant with a single user can be reshaped and laid out as + NHWC ahead of time instead of at every inference. + """ + return ( + node.op == "placeholder" + and is_param_node(self.exported_program, node) + and len(node.users) == 1 + ) + + def constant_to_4d_nhwc(self, input_node: torch.fx.Node) -> torch.fx.Node: + """ + Reshape a rank < 4 constant operand to 4D in place and tag it NHWC, so + the serializer writes its data in channels_last and the graph carries + no reshape or transpose for it. + """ + val = input_node.meta["val"] + new_shape = [1] * (4 - val.dim()) + list(val.shape) + data = get_param_tensor(self.exported_program, input_node) + assert data is not None + data_4d = data.reshape(new_shape).contiguous() + signature = self.exported_program.graph_signature + if torch._export.utils.is_param(self.exported_program, input_node): + name = signature.inputs_to_parameters[input_node.name] + self.exported_program.state_dict[name] = torch.nn.Parameter( + data_4d, requires_grad=False + ) + elif torch._export.utils.is_buffer(self.exported_program, input_node): + name = signature.inputs_to_buffers[input_node.name] + self.exported_program.state_dict[name] = data_4d + else: + name = signature.inputs_to_lifted_tensor_constants[input_node.name] + self.exported_program.constants[name] = data_4d + input_node.meta["val"] = val.reshape(new_shape) + self.mark_as_nhwc_node(input_node) + return input_node + def unsqueeze_to_4d( self, graph_module: torch.fx.GraphModule, @@ -214,6 +254,8 @@ def unsqueeze_to_4d( Reshape a rank < 4 operand of `target_node` to 4D by prepending ones, which is how broadcasting aligns it, so it can be converted to NHWC. """ + if self.is_foldable_constant(input_node): + return self.constant_to_4d_nhwc(input_node) val = input_node.meta["val"] new_shape = [1] * (4 - val.dim()) + list(val.shape) with graph_module.graph.inserting_before(target_node): @@ -668,7 +710,11 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 for input_node in layout_inputs ): for input_node in layout_inputs: - self.input_to_nhwc(graph_module, input_node, node) + # A constant folded to NHWC above needs no copy. + if not ChannelsLastTaggedReshapePass.is_nhwc_node( + input_node + ): + self.input_to_nhwc(graph_module, input_node, node) self.mark_as_nhwc_node(node) else: for input_node in layout_inputs: diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 0b5a0f5b4c9..218c1b900e5 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -315,9 +315,44 @@ def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_stays_nhwc(self): .run_method_and_compare_outputs() ) + class ConvMul5dScalar(torch.nn.Module): + """conv -> x * buf[1, 1, 1, 1, 1]: the single-element operand has + rank 5, so the mul output is 5D and cannot be NHWC. The operand takes + part in the layout decision and the mul takes the NCHW path.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + self.register_buffer("buf", torch.full((1, 1, 1, 1, 1), 1.5)) + + def forward(self, x): + return self.conv(x) * self.buf + + def test_fp32_channels_last_tagged_reshape_pass_5d_single_element_operand(self): + inputs = (torch.randn(1, 3, 8, 8),) + ( + Tester(self.ConvMul5dScalar().eval(), inputs) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC and conv output -> NCHW; the 5D mul output stays NCHW + .check_count({self.to_copy_name: 2}) + .run_method_and_compare_outputs() + ) + ( + Tester(self.ConvMul5dScalar().eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + class ConvNoiseAddConv(torch.nn.Module): - """conv -> x + noise[H, W] -> conv: the rank-2 operand is reshaped to - [1, 1, H, W] and converted, instead of converting the activation.""" + """conv -> x + noise[H, W] -> conv: the constant operand is reshaped to + [1, 1, H, W] and laid out as NHWC ahead of time, instead of converting + the activation.""" def __init__(self): super().__init__() @@ -331,8 +366,8 @@ def forward(self, x): return self.conv2(x) class ConvChannelScaleConv(torch.nn.Module): - """conv -> x * scale[C, 1, 1] -> conv: a per-channel operand is - reshaped to [1, C, 1, 1] and converted.""" + """conv -> x * scale[C, 1, 1] -> conv: a per-channel constant is + reshaped to [1, C, 1, 1] ahead of time.""" def __init__(self): super().__init__() @@ -360,34 +395,129 @@ def forward(self, x): x = self.noise - x return self.conv2(x) - def _test_lower_rank_operand_stays_nhwc(self, module): + class ConvChwBiasAddConv(torch.nn.Module): + """conv -> x + bias[C, H, W] -> conv: a constant whose NHWC data + layout differs from NCHW, so the ahead-of-time transpose is + observable in the output.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("bias", torch.randn(8, 8, 8)) + + def forward(self, x): + x = self.conv1(x) + x = x + self.bias + return self.conv2(x) + + def _test_constant_lower_rank_operand_folded_to_nhwc(self, module): + inputs = (torch.randn(1, 3, 8, 8),) ( - Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) + Tester(module.eval(), inputs) .export() .to_edge() .run_passes(self.PassStage) - # input -> NHWC, the 4D view of the operand -> NHWC, output -> NCHW - .check_count( - { - self.to_copy_name: 3, - "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - } - ) + # input -> NHWC and output -> NCHW only: the constant is reshaped + # and transposed ahead of time, not in the graph + .check_count({self.to_copy_name: 2}) + .check_not(["executorch_exir_dialects_edge__ops_aten_view_copy_default"]) + .run_method_and_compare_outputs() + ) + ( + Tester(module.eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() .run_method_and_compare_outputs() ) def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_stays_nhwc(self): - self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseAddConv()) + self._test_constant_lower_rank_operand_folded_to_nhwc(self.ConvNoiseAddConv()) def test_fp32_channels_last_tagged_reshape_pass_per_channel_operand_stays_nhwc( self, ): - self._test_lower_rank_operand_stays_nhwc(self.ConvChannelScaleConv()) + self._test_constant_lower_rank_operand_folded_to_nhwc( + self.ConvChannelScaleConv() + ) def test_fp32_channels_last_tagged_reshape_pass_lower_rank_left_operand_stays_nhwc( self, ): - self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseSubConv()) + self._test_constant_lower_rank_operand_folded_to_nhwc(self.ConvNoiseSubConv()) + + def test_fp32_channels_last_tagged_reshape_pass_chw_operand_stays_nhwc(self): + self._test_constant_lower_rank_operand_folded_to_nhwc(self.ConvChwBiasAddConv()) + + class ConvSharedNoiseAddConv(torch.nn.Module): + """conv -> x + noise[H, W] -> conv -> x + noise: the constant has two + users, so it keeps its rank and each use is reshaped in the graph.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("noise", torch.randn(8, 8)) + + def forward(self, x): + x = self.conv1(x) + self.noise + return self.conv2(x) + self.noise + + def test_fp32_channels_last_tagged_reshape_pass_shared_lower_rank_operand(self): + inputs = (torch.randn(1, 3, 8, 8),) + ( + Tester(self.ConvSharedNoiseAddConv().eval(), inputs) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC, each 4D view of the operand -> NHWC, output -> NCHW + .check_count( + { + self.to_copy_name: 4, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 2, + } + ) + .run_method_and_compare_outputs() + ) + ( + Tester(self.ConvSharedNoiseAddConv().eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + + def test_fp32_channels_last_tagged_reshape_pass_input_lower_rank_operand(self): + # A non-constant [H, W] operand is reshaped and converted in the graph. + inputs = (torch.randn(1, 3, 8, 8), torch.randn(8, 8)) + ( + Tester(self.ConvDynamicNoiseAddConv().eval(), inputs) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC, the 4D view of the operand -> NHWC, output -> NCHW + .check_count( + { + self.to_copy_name: 3, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + } + ) + .run_method_and_compare_outputs() + ) + ( + Tester(self.ConvDynamicNoiseAddConv().eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): # The dequantized operand is left at its own rank: a view between @@ -406,8 +536,9 @@ def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): ) class ConvDynamicNoiseAddConv(torch.nn.Module): - """conv -> x + noise[H, W] -> conv with dynamic H and W: the operand - cannot be reshaped by XNNStaticReshape, so the add takes the NCHW path.""" + """conv -> x + noise[H, W] -> conv with noise as a graph input. With + dynamic H and W the operand cannot be reshaped by XNNStaticReshape, so + the add takes the NCHW path.""" def __init__(self): super().__init__() From 40026539a7a014903c156aaf9962ccf2d85d4d9c Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 15:55:26 +0100 Subject: [PATCH 03/11] [XNNPACK] Fold non-persistent buffer operands into the program's constants Reshaping a lower-rank constant operand to 4D wrote the data to the state dict regardless of what kind of constant it was, so a non-persistent buffer, which the exported program keeps in its constants, ended up with a stray state dict entry while the serializer still read the original tensor. Use set_param_tensor, which routes each kind to its own table. --- .../channels_last_tagged_reshape_pass.py | 17 ++++---------- .../test_channels_last_tagged_reshape.py | 23 +++++++++++++++++++ 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index e4a6b3e1491..e41443ccba7 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -10,6 +10,7 @@ from typing import List, Optional, Tuple import torch +from executorch.backends.transforms.utils import set_param_tensor from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.backends.xnnpack.utils.quant_utils import ( extract_qdq_affine_op_args_for_decomposed_ops, @@ -227,19 +228,9 @@ def constant_to_4d_nhwc(self, input_node: torch.fx.Node) -> torch.fx.Node: new_shape = [1] * (4 - val.dim()) + list(val.shape) data = get_param_tensor(self.exported_program, input_node) assert data is not None - data_4d = data.reshape(new_shape).contiguous() - signature = self.exported_program.graph_signature - if torch._export.utils.is_param(self.exported_program, input_node): - name = signature.inputs_to_parameters[input_node.name] - self.exported_program.state_dict[name] = torch.nn.Parameter( - data_4d, requires_grad=False - ) - elif torch._export.utils.is_buffer(self.exported_program, input_node): - name = signature.inputs_to_buffers[input_node.name] - self.exported_program.state_dict[name] = data_4d - else: - name = signature.inputs_to_lifted_tensor_constants[input_node.name] - self.exported_program.constants[name] = data_4d + set_param_tensor( + self.exported_program, input_node, data.reshape(new_shape).contiguous() + ) input_node.meta["val"] = val.reshape(new_shape) self.mark_as_nhwc_node(input_node) return input_node diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 218c1b900e5..938efbbb7bb 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -452,6 +452,29 @@ def test_fp32_channels_last_tagged_reshape_pass_lower_rank_left_operand_stays_nh def test_fp32_channels_last_tagged_reshape_pass_chw_operand_stays_nhwc(self): self._test_constant_lower_rank_operand_folded_to_nhwc(self.ConvChwBiasAddConv()) + class ConvNonPersistentNoiseAddConv(torch.nn.Module): + """conv -> x + noise[H, W] -> conv with a non-persistent buffer, which + the exported program stores among its constants, not in the state + dict.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("noise", torch.randn(8, 8), persistent=False) + + def forward(self, x): + x = self.conv1(x) + x = x + self.noise + return self.conv2(x) + + def test_fp32_channels_last_tagged_reshape_pass_non_persistent_buffer_operand_stays_nhwc( + self, + ): + self._test_constant_lower_rank_operand_folded_to_nhwc( + self.ConvNonPersistentNoiseAddConv() + ) + class ConvSharedNoiseAddConv(torch.nn.Module): """conv -> x + noise[H, W] -> conv -> x + noise: the constant has two users, so it keeps its rank and each use is reshaped in the graph.""" From 3db7aff388d84e01b1887a8c93010ead6edd7078 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 16:10:46 +0100 Subject: [PATCH 04/11] [XNNPACK] Decide the layout before reshaping lower-rank operands The main traversal reshaped a lower-rank operand to 4D, and folded a constant one in place, before checking whether the op can go NHWC at all, and the re-converge loop after it did not reshape such operands. Move the decision into one converge_layouts helper used by both: it checks the other operands first, reshapes only on the NHWC path, and leaves the graph and the constants alone on the NCHW path. --- .../channels_last_tagged_reshape_pass.py | 71 +++++++++---------- 1 file changed, 35 insertions(+), 36 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index e41443ccba7..1403d102620 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -235,6 +235,39 @@ def constant_to_4d_nhwc(self, input_node: torch.fx.Node) -> torch.fx.Node: self.mark_as_nhwc_node(input_node) return input_node + def converge_layouts( + self, + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + layout_inputs: List[torch.fx.Node], + ) -> None: + """ + Bring operands of mixed layouts to one: NHWC if every operand can be + converted, NCHW otherwise. A lower-rank operand of a broadcasting op + (e.g. a [H, W] map added to [N, C, H, W]) can follow the NHWC operand + once reshaped to 4D, so the layout is decided on the other operands + first and only then is the graph or the constant touched. + """ + lower_rank = [ + node.target in self.broadcast_binary_ops + and self.is_lower_rank_broadcast_operand(input_node) + for input_node in layout_inputs + ] + if all( + is_lower_rank or self.can_be_converted_to_nhwc(input_node) + for input_node, is_lower_rank in zip(layout_inputs, lower_rank) + ): + for input_node, is_lower_rank in zip(layout_inputs, lower_rank): + if is_lower_rank: + input_node = self.unsqueeze_to_4d(graph_module, input_node, node) + # A constant folded to NHWC needs no copy. + if not ChannelsLastTaggedReshapePass.is_nhwc_node(input_node): + self.input_to_nhwc(graph_module, input_node, node) + self.mark_as_nhwc_node(node) + else: + for input_node in layout_inputs: + self.input_to_nchw(graph_module, input_node, node) + def unsqueeze_to_4d( self, graph_module: torch.fx.GraphModule, @@ -683,33 +716,7 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 # There is at least one node of each format, so we need to # convert some of them to another format such that all # are the same. - # Reshape lower-rank operands (e.g. a [H, W] map added to - # [N, C, H, W]) to 4D so they can follow the NHWC operand. - if node.target in self.broadcast_binary_ops: - layout_inputs = [ - ( - self.unsqueeze_to_4d(graph_module, input_node, node) - if self.is_lower_rank_broadcast_operand(input_node) - else input_node - ) - for input_node in layout_inputs - ] - # If all nodes can be converted to nhwc, then convert them - # all to nhwc. Otherwise, convert all to nchw - if all( - self.can_be_converted_to_nhwc(input_node) - for input_node in layout_inputs - ): - for input_node in layout_inputs: - # A constant folded to NHWC above needs no copy. - if not ChannelsLastTaggedReshapePass.is_nhwc_node( - input_node - ): - self.input_to_nhwc(graph_module, input_node, node) - self.mark_as_nhwc_node(node) - else: - for input_node in layout_inputs: - self.input_to_nchw(graph_module, input_node, node) + self.converge_layouts(graph_module, node, layout_inputs) graph_module.recompile() @@ -741,15 +748,7 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 ] if layouts[0] == layouts[1]: continue - if all( - self.can_be_converted_to_nhwc(input_node) for input_node in input_nodes - ): - for input_node in input_nodes: - self.input_to_nhwc(graph_module, input_node, node) - self.mark_as_nhwc_node(node) - else: - for input_node in input_nodes: - self.input_to_nchw(graph_module, input_node, node) + self.converge_layouts(graph_module, node, input_nodes) reconverged = True if reconverged: From 7e817a33e3ec9c5524ab3b84cc7214e916c36211 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 16:20:56 +0100 Subject: [PATCH 05/11] [XNNPACK] Test scalar and lower-rank operands of every broadcasting binary op Cover div, maximum and minimum, which had no test, with one model parameterized over the six ops: the scalar case runs through serialization as well. --- .../test_channels_last_tagged_reshape.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 938efbbb7bb..48c75f026b0 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -315,6 +315,68 @@ def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_stays_nhwc(self): .run_method_and_compare_outputs() ) + class ConvBinaryOpConv(torch.nn.Module): + """conv -> op(x, operand) -> conv for any of the broadcasting binary + ops the pass handles.""" + + def __init__(self, op, operand): + super().__init__() + self.op = op + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("operand", operand) + + def forward(self, x): + x = self.conv1(x) + x = self.op(x, self.operand) + return self.conv2(x) + + broadcast_binary_ops = ( + torch.add, + torch.mul, + torch.sub, + torch.div, + torch.maximum, + torch.minimum, + ) + + def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_all_binary_ops( + self, + ): + inputs = (torch.randn(1, 3, 8, 8),) + for op in self.broadcast_binary_ops: + with self.subTest(op=op.__name__): + module = self.ConvBinaryOpConv(op, torch.tensor(1.4142)).eval() + ( + Tester(module, inputs) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC and output -> NCHW only + .check_count({self.to_copy_name: 2}) + .run_method_and_compare_outputs() + ) + ( + Tester(module, inputs) + .export() + .to_edge_transform_and_lower() + .check_count( + {"torch.ops.higher_order.executorch_call_delegate": 1} + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + + def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_all_binary_ops( + self, + ): + for op in self.broadcast_binary_ops: + with self.subTest(op=op.__name__): + self._test_constant_lower_rank_operand_folded_to_nhwc( + self.ConvBinaryOpConv(op, torch.rand(8, 8) + 0.5) + ) + class ConvMul5dScalar(torch.nn.Module): """conv -> x * buf[1, 1, 1, 1, 1]: the single-element operand has rank 5, so the mul output is 5D and cannot be NHWC. The operand takes From 38bf48cea184f17c56914e408812b8bec9a6f3c3 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 16:21:26 +0100 Subject: [PATCH 06/11] [XNNPACK] Format the binary op test --- .../xnnpack/test/passes/test_channels_last_tagged_reshape.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 48c75f026b0..832686c69d7 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -360,9 +360,7 @@ def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_all_binary_ops( Tester(module, inputs) .export() .to_edge_transform_and_lower() - .check_count( - {"torch.ops.higher_order.executorch_call_delegate": 1} - ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) .to_executorch() .serialize() .run_method_and_compare_outputs() From 56c09e063b7399588a0ee04c7e525fed356270dc Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 16:23:43 +0100 Subject: [PATCH 07/11] [XNNPACK] Use a non-square spatial shape in the lower-rank operand tests With H == W a mixed-up transpose keeps the shape, so the fold tests now run on 5x7 inputs and constants. --- .../test_channels_last_tagged_reshape.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 832686c69d7..c95d68ef00f 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -372,7 +372,7 @@ def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_all_binary_op for op in self.broadcast_binary_ops: with self.subTest(op=op.__name__): self._test_constant_lower_rank_operand_folded_to_nhwc( - self.ConvBinaryOpConv(op, torch.rand(8, 8) + 0.5) + self.ConvBinaryOpConv(op, torch.rand(5, 7) + 0.5) ) class ConvMul5dScalar(torch.nn.Module): @@ -418,7 +418,7 @@ def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(8, 8)) + self.register_buffer("noise", torch.randn(5, 7)) def forward(self, x): x = self.conv1(x) @@ -448,7 +448,7 @@ def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(8, 8)) + self.register_buffer("noise", torch.randn(5, 7)) def forward(self, x): x = self.conv1(x) @@ -464,7 +464,7 @@ def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("bias", torch.randn(8, 8, 8)) + self.register_buffer("bias", torch.randn(8, 5, 7)) def forward(self, x): x = self.conv1(x) @@ -472,7 +472,7 @@ def forward(self, x): return self.conv2(x) def _test_constant_lower_rank_operand_folded_to_nhwc(self, module): - inputs = (torch.randn(1, 3, 8, 8),) + inputs = (torch.randn(1, 3, 5, 7),) ( Tester(module.eval(), inputs) .export() @@ -521,7 +521,7 @@ def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(8, 8), persistent=False) + self.register_buffer("noise", torch.randn(5, 7), persistent=False) def forward(self, x): x = self.conv1(x) @@ -543,14 +543,14 @@ def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(8, 8)) + self.register_buffer("noise", torch.randn(5, 7)) def forward(self, x): x = self.conv1(x) + self.noise return self.conv2(x) + self.noise def test_fp32_channels_last_tagged_reshape_pass_shared_lower_rank_operand(self): - inputs = (torch.randn(1, 3, 8, 8),) + inputs = (torch.randn(1, 3, 5, 7),) ( Tester(self.ConvSharedNoiseAddConv().eval(), inputs) .export() @@ -577,7 +577,7 @@ def test_fp32_channels_last_tagged_reshape_pass_shared_lower_rank_operand(self): def test_fp32_channels_last_tagged_reshape_pass_input_lower_rank_operand(self): # A non-constant [H, W] operand is reshaped and converted in the graph. - inputs = (torch.randn(1, 3, 8, 8), torch.randn(8, 8)) + inputs = (torch.randn(1, 3, 5, 7), torch.randn(5, 7)) ( Tester(self.ConvDynamicNoiseAddConv().eval(), inputs) .export() @@ -607,7 +607,7 @@ def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): # dequantize and the op would carry no quant params and fail at load. for module in (self.ConvNoiseAddConv(), self.ConvChannelScaleConv()): ( - Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) + Tester(module.eval(), (torch.randn(1, 3, 5, 7),)) .quantize() .export() .to_edge() @@ -637,7 +637,7 @@ def test_fp32_channels_last_tagged_reshape_pass_dynamic_lower_rank_operand(self) height = torch.export.Dim("height", min=4, max=32) width = torch.export.Dim("width", min=4, max=32) dynamic_shapes = ({2: height, 3: width}, {0: height, 1: width}) - inputs = (torch.randn(1, 3, 8, 8), torch.randn(8, 8)) + inputs = (torch.randn(1, 3, 5, 7), torch.randn(5, 7)) ( Tester( self.ConvDynamicNoiseAddConv().eval(), From 6b3b9f2e6a7cede16f398920a71e6202447b88f9 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 16:31:20 +0100 Subject: [PATCH 08/11] [XNNPACK] Test the fold for every kind of constant A parameter, a persistent and a non-persistent buffer and a lifted tensor constant are stored in different tables of the exported program, so run the lower-rank fold test over all four. --- .../test_channels_last_tagged_reshape.py | 48 ++++++++----------- 1 file changed, 21 insertions(+), 27 deletions(-) diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index c95d68ef00f..33bc4ad3edb 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -412,13 +412,24 @@ def test_fp32_channels_last_tagged_reshape_pass_5d_single_element_operand(self): class ConvNoiseAddConv(torch.nn.Module): """conv -> x + noise[H, W] -> conv: the constant operand is reshaped to [1, 1, H, W] and laid out as NHWC ahead of time, instead of converting - the activation.""" + the activation. `kind` picks how the constant is held by the module, + which decides where the exported program stores it.""" - def __init__(self): + def __init__(self, kind="buffer"): super().__init__() self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(5, 7)) + noise = torch.randn(5, 7) + if kind == "buffer": + self.register_buffer("noise", noise) + elif kind == "non_persistent_buffer": + self.register_buffer("noise", noise, persistent=False) + elif kind == "parameter": + self.noise = torch.nn.Parameter(noise, requires_grad=False) + elif kind == "tensor_constant": + self.noise = noise + else: + raise ValueError(kind) def forward(self, x): x = self.conv1(x) @@ -495,7 +506,13 @@ def _test_constant_lower_rank_operand_folded_to_nhwc(self, module): ) def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_stays_nhwc(self): - self._test_constant_lower_rank_operand_folded_to_nhwc(self.ConvNoiseAddConv()) + # Every kind of constant: state dict buffer and parameter, and the + # non-persistent buffer and lifted tensor held among the constants. + for kind in ("buffer", "non_persistent_buffer", "parameter", "tensor_constant"): + with self.subTest(kind=kind): + self._test_constant_lower_rank_operand_folded_to_nhwc( + self.ConvNoiseAddConv(kind) + ) def test_fp32_channels_last_tagged_reshape_pass_per_channel_operand_stays_nhwc( self, @@ -512,29 +529,6 @@ def test_fp32_channels_last_tagged_reshape_pass_lower_rank_left_operand_stays_nh def test_fp32_channels_last_tagged_reshape_pass_chw_operand_stays_nhwc(self): self._test_constant_lower_rank_operand_folded_to_nhwc(self.ConvChwBiasAddConv()) - class ConvNonPersistentNoiseAddConv(torch.nn.Module): - """conv -> x + noise[H, W] -> conv with a non-persistent buffer, which - the exported program stores among its constants, not in the state - dict.""" - - def __init__(self): - super().__init__() - self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) - self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(5, 7), persistent=False) - - def forward(self, x): - x = self.conv1(x) - x = x + self.noise - return self.conv2(x) - - def test_fp32_channels_last_tagged_reshape_pass_non_persistent_buffer_operand_stays_nhwc( - self, - ): - self._test_constant_lower_rank_operand_folded_to_nhwc( - self.ConvNonPersistentNoiseAddConv() - ) - class ConvSharedNoiseAddConv(torch.nn.Module): """conv -> x + noise[H, W] -> conv -> x + noise: the constant has two users, so it keeps its rank and each use is reshaped in the graph.""" From 8bfcbf8b6a9882fb1e5f25c2321765dba38cc456 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 16:53:42 +0100 Subject: [PATCH 09/11] [XNNPACK] Delegate constant operands of broadcasting binary ops Only convolution and linear pulled their constants into the partition, so a constant operand of add, mul, sub, div, maximum or minimum reached the delegate as a runtime input and the ahead-of-time fold never ran on the lowered path: the serialized graph still reshaped and transposed it at every inference. The binary op configs now delegate such constants with the op. The end-to-end tests decode the delegate's XNNPACK graph and check that only the input and output transposes remain. --- .../partition/config/generic_node_configs.py | 30 ++++++++-- .../test_channels_last_tagged_reshape.py | 60 ++++++++++++++++--- 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index d445b5ca8ac..54306b91059 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -178,6 +178,24 @@ def get_node_and_deps( return deps + quantized_deps +class BroadcastBinaryConfig(GenericNodePartitionerConfig): + """ + A binary op whose constant operands are delegated with it, so they are + serialized as constant data and a lower-rank one can be laid out as NHWC + ahead of time instead of being reshaped and transposed at every inference. + """ + + def get_node_and_deps( + self, node: torch.fx.Node, ep: ExportedProgram + ) -> List[torch.fx.Node]: + deps = super().get_node_and_deps(node, ep) + return deps + [ + input_node + for input_node in node.all_input_nodes + if is_param_node(ep, input_node) + ] + + class QuantizedPerTensorConfig(GenericNodePartitionerConfig): target_name = "quantize_per_tensor.default" @@ -209,7 +227,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] -class AddConfig(GenericNodePartitionerConfig): +class AddConfig(BroadcastBinaryConfig): target_name = "add.Tensor" def __init__(self, **kwargs): @@ -353,7 +371,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] -class DivConfig(GenericNodePartitionerConfig): +class DivConfig(BroadcastBinaryConfig): target_name = "div.Tensor" def supported_precision_types(self) -> List[ConfigPrecisionType]: @@ -433,14 +451,14 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: ) -class MulConfig(GenericNodePartitionerConfig): +class MulConfig(BroadcastBinaryConfig): target_name = "mul.Tensor" def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] -class MaximumConfig(GenericNodePartitionerConfig): +class MaximumConfig(BroadcastBinaryConfig): target_name = "maximum.default" def supported_precision_types(self) -> List[ConfigPrecisionType]: @@ -683,7 +701,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] -class MinimumConfig(GenericNodePartitionerConfig): +class MinimumConfig(BroadcastBinaryConfig): target_name = "minimum.default" def supported_precision_types(self) -> List[ConfigPrecisionType]: @@ -817,7 +835,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32] -class SubConfig(GenericNodePartitionerConfig): +class SubConfig(BroadcastBinaryConfig): target_name = "sub.Tensor" def supported_precision_types(self) -> List[ConfigPrecisionType]: diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 33bc4ad3edb..55611990d9f 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -6,7 +6,11 @@ # pyre-unsafe +import json +import os +import tempfile import unittest +from collections import Counter import torch from executorch.backends.test.harness.stages.stage import StageType @@ -16,6 +20,10 @@ from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( get_symmetric_quantization_config, ) +from executorch.backends.xnnpack.serialization.xnnpack_graph_serialize import ( + _get_schema_bytes, + XNNHeader, +) from executorch.backends.xnnpack.test.test_xnnpack_utils_classes import ( OpSequencesAddConv2d, ) @@ -26,6 +34,7 @@ is_quant, is_tagged_as_implicit_q_dq, ) +from executorch.exir._serialize._flatbuffer import _flatc_decompile from executorch.exir.dialects._ops import ops as exir_ops @@ -356,15 +365,14 @@ def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_all_binary_ops( .check_count({self.to_copy_name: 2}) .run_method_and_compare_outputs() ) - ( + tester = ( Tester(module, inputs) .export() .to_edge_transform_and_lower() .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) - .to_executorch() - .serialize() - .run_method_and_compare_outputs() ) + self._assert_only_io_transposes(tester) + tester.to_executorch().serialize().run_method_and_compare_outputs() def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_all_binary_ops( self, @@ -482,6 +490,43 @@ def forward(self, x): x = x + self.bias return self.conv2(x) + def _xnnpack_node_types(self, tester) -> Counter: + """Node types in the serialized XNNPACK graph of the single delegate + produced by to_edge_transform_and_lower.""" + graph_module = ( + tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + .exported_program() + .graph_module + ) + (delegate,) = [ + node + for node in graph_module.graph.nodes + if node.target == torch.ops.higher_order.executorch_call_delegate + ] + payload = getattr(graph_module, delegate.args[0].target).processed_bytes + header = XNNHeader.from_bytes(payload[: XNNHeader.EXPECTED_LENGTH]) + flatbuffer = payload[ + header.flatbuffer_offset : header.flatbuffer_offset + header.flatbuffer_size + ] + with tempfile.TemporaryDirectory() as d: + schema_path = os.path.join(d, "schema.fbs") + with open(schema_path, "wb") as f: + f.write(_get_schema_bytes()) + bin_path = os.path.join(d, "graph.bin") + with open(bin_path, "wb") as f: + f.write(flatbuffer) + _flatc_decompile(d, schema_path, bin_path) + with open(os.path.join(d, "graph.json")) as f: + graph = json.load(f) + return Counter(node["xnode_union_type"] for node in graph["xnodes"]) + + def _assert_only_io_transposes(self, tester): + """The delegate transposes the input to NHWC and the output back, and + nothing else: no reshape or transpose for the operand.""" + node_types = self._xnnpack_node_types(tester) + self.assertEqual(node_types["XNNStaticTranspose"], 2, node_types) + self.assertEqual(node_types["XNNStaticReshape"], 0, node_types) + def _test_constant_lower_rank_operand_folded_to_nhwc(self, module): inputs = (torch.randn(1, 3, 5, 7),) ( @@ -495,15 +540,14 @@ def _test_constant_lower_rank_operand_folded_to_nhwc(self, module): .check_not(["executorch_exir_dialects_edge__ops_aten_view_copy_default"]) .run_method_and_compare_outputs() ) - ( + tester = ( Tester(module.eval(), inputs) .export() .to_edge_transform_and_lower() .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) - .to_executorch() - .serialize() - .run_method_and_compare_outputs() ) + self._assert_only_io_transposes(tester) + tester.to_executorch().serialize().run_method_and_compare_outputs() def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_stays_nhwc(self): # Every kind of constant: state dict buffer and parameter, and the From 9c7544eacf27e1d8f20f79af5f2d73dc82466d71 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Fri, 30 Oct 2026 17:00:49 +0100 Subject: [PATCH 10/11] [XNNPACK] Check the NCHW path and the lowered graph in the qs8 lower-rank test The op must stay NCHW and the delegate must carry no reshape; the test now asserts both and runs the serialized program. --- .../test_channels_last_tagged_reshape.py | 54 +++++++++++++++---- 1 file changed, 43 insertions(+), 11 deletions(-) diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 55611990d9f..a05e52293e4 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -643,18 +643,50 @@ def test_fp32_channels_last_tagged_reshape_pass_input_lower_rank_operand(self): def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): # The dequantized operand is left at its own rank: a view between # dequantize and the op would carry no quant params and fail at load. - for module in (self.ConvNoiseAddConv(), self.ConvChannelScaleConv()): - ( - Tester(module.eval(), (torch.randn(1, 3, 5, 7),)) - .quantize() - .export() - .to_edge() - .run_passes(self.PassStage) - .check_not( - ["executorch_exir_dialects_edge__ops_aten_view_copy_default"] + # The op takes the NCHW path instead. + inputs = (torch.randn(1, 3, 5, 7),) + for module, op in ( + (self.ConvNoiseAddConv(), exir_ops.edge.aten.add.Tensor), + (self.ConvChannelScaleConv(), exir_ops.edge.aten.mul.Tensor), + ): + with self.subTest(module=type(module).__name__): + tester = ( + Tester(module.eval(), inputs) + .quantize() + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC, conv1 output -> NCHW, op output -> NHWC, + # output -> NCHW + .check_count({self.to_copy_name: 4}) + .check_not( + ["executorch_exir_dialects_edge__ops_aten_view_copy_default"] + ) + .run_method_and_compare_outputs() + ) + graph_module = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module + ) + (binary_op,) = [ + node for node in graph_module.graph.nodes if node.target == op + ] + self.assertFalse(ChannelsLastTaggedReshapePass.is_nhwc_node(binary_op)) + + tester = ( + Tester(module.eval(), inputs) + .quantize() + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + ) + node_types = self._xnnpack_node_types(tester) + self.assertEqual(node_types["XNNStaticTranspose"], 4, node_types) + self.assertEqual(node_types["XNNStaticReshape"], 0, node_types) + tester.to_executorch().serialize().run_method_and_compare_outputs( + qtol=1 ) - .run_method_and_compare_outputs() - ) class ConvDynamicNoiseAddConv(torch.nn.Module): """conv -> x + noise[H, W] -> conv with noise as a graph input. With From a013bac0f51cfd2c43e44e03d04eb1332505314e Mon Sep 17 00:00:00 2001 From: borys3kk Date: Mon, 2 Nov 2026 11:12:22 +0100 Subject: [PATCH 11/11] Delegate only foldable constants with broadcasting binary ops BroadcastBinaryConfig pulled every constant operand into the delegate, which froze mutated buffers, forced 4D constants onto the NCHW path and duplicated constants shared with other partitions. It now delegates only operands the pass can fold: non-mutated, static, rank < 4, more than one element. The conditions live in one helper shared with the pass. --- .../channels_last_tagged_reshape_pass.py | 8 +- .../partition/config/generic_node_configs.py | 10 +- .../test_channels_last_tagged_reshape.py | 139 ++++++++++++++++++ backends/xnnpack/utils/utils.py | 24 +++ 4 files changed, 173 insertions(+), 8 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 7c47ccafd43..2df6906cb51 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -23,6 +23,7 @@ from executorch.backends.xnnpack.utils.utils import ( get_param_tensor, is_dtype_cast, + is_foldable_broadcast_operand, is_param_node, ) from executorch.exir.dialects._ops import ops as exir_ops @@ -213,12 +214,11 @@ def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: def is_foldable_constant(self, node: torch.fx.Node) -> bool: """ - A lifted constant with a single user can be reshaped and laid out as - NHWC ahead of time instead of at every inference. + A constant with a single user can be reshaped and laid out as NHWC + ahead of time instead of at every inference. """ return ( - node.op == "placeholder" - and is_param_node(self.exported_program, node) + is_foldable_broadcast_operand(self.exported_program, node) and len(node.users) == 1 ) diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index 0642f2ba03b..e6cd46addb9 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -24,6 +24,7 @@ from executorch.backends.xnnpack.utils.utils import ( get_input_node, get_param_tensor, + is_foldable_broadcast_operand, is_param_node, normalize_mean_dims, normalize_pool2d_args, @@ -180,9 +181,10 @@ def get_node_and_deps( class BroadcastBinaryConfig(GenericNodePartitionerConfig): """ - A binary op whose constant operands are delegated with it, so they are - serialized as constant data and a lower-rank one can be laid out as NHWC - ahead of time instead of being reshaped and transposed at every inference. + A binary op whose rank < 4 constant operands are delegated with it, so + they are serialized as constant data and laid out as NHWC ahead of time + instead of being reshaped and transposed at every inference. Any other + constant operand stays a delegate input, as for every other op. """ def get_node_and_deps( @@ -192,7 +194,7 @@ def get_node_and_deps( return deps + [ input_node for input_node in node.all_input_nodes - if is_param_node(ep, input_node) + if is_foldable_broadcast_operand(ep, input_node) ] diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index a05e52293e4..c91c9d56ae3 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -520,6 +520,23 @@ def _xnnpack_node_types(self, tester) -> Counter: graph = json.load(f) return Counter(node["xnode_union_type"] for node in graph["xnodes"]) + def _delegate_input_buffers(self, tester) -> set: + """Names of the buffers the single delegate takes as inputs.""" + exported_program = tester.get_artifact( + StageType.TO_EDGE_TRANSFORM_AND_LOWER + ).exported_program() + (delegate,) = [ + node + for node in exported_program.graph_module.graph.nodes + if node.target == torch.ops.higher_order.executorch_call_delegate + ] + inputs_to_buffers = exported_program.graph_signature.inputs_to_buffers + return { + inputs_to_buffers[arg.name] + for arg in delegate.args[1:] + if arg.name in inputs_to_buffers + } + def _assert_only_io_transposes(self, tester): """The delegate transposes the input to NHWC and the output back, and nothing else: no reshape or transpose for the operand.""" @@ -613,6 +630,128 @@ def test_fp32_channels_last_tagged_reshape_pass_shared_lower_rank_operand(self): .run_method_and_compare_outputs() ) + class ConvMutatedStateMulConv(torch.nn.Module): + """conv -> x * state[H, W] -> conv, with the state updated after each + inference: a mutated buffer is not constant, so it is neither folded + nor delegated as constant data.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("state", torch.randn(5, 7)) + + def forward(self, x): + x = self.conv1(x) + x = x * self.state + self.state.add_(1.0) + return self.conv2(x) + + def test_fp32_channels_last_tagged_reshape_pass_mutated_buffer_operand(self): + inputs = (torch.randn(1, 3, 5, 7),) + tester = ( + Tester(self.ConvMutatedStateMulConv().eval(), inputs) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC, the 4D view of the state -> NHWC, output -> NCHW + .check_count( + { + self.to_copy_name: 3, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + } + ) + ) + state = tester.get_artifact(StageType.RUN_PASSES).exported_program() + self.assertEqual(tuple(state.state_dict["state"].shape), (5, 7)) + tester = ( + Tester(self.ConvMutatedStateMulConv().eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + ) + # Read at every inference, so it stays a delegate input. + self.assertIn("state", self._delegate_input_buffers(tester)) + + class ConvScale4dMulConv(torch.nn.Module): + """conv -> x * scale[1, C, 1, 1] -> conv: a 4D constant is already + broadcast-aligned, so it stays a delegate input and is converted to + NHWC in the graph, as before.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("scale", torch.randn(1, 8, 1, 1)) + + def forward(self, x): + x = self.conv1(x) + x = x * self.scale + return self.conv2(x) + + def test_fp32_channels_last_tagged_reshape_pass_4d_constant_operand(self): + inputs = (torch.randn(1, 3, 5, 7),) + ( + Tester(self.ConvScale4dMulConv().eval(), inputs) + .export() + .to_edge() + .run_passes(self.PassStage) + # On the whole program the constant is tagged NCHW, so the op + # takes the NCHW path: input -> NHWC, conv1 output -> NCHW, + # op output -> NHWC, output -> NCHW + .check_count({self.to_copy_name: 4}) + .check_not(["executorch_exir_dialects_edge__ops_aten_view_copy_default"]) + .run_method_and_compare_outputs() + ) + # In the delegate the constant is an input: input -> NHWC, + # scale -> NHWC, output -> NCHW + tester = ( + Tester(self.ConvScale4dMulConv().eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + ) + self.assertIn("scale", self._delegate_input_buffers(tester)) + node_types = self._xnnpack_node_types(tester) + self.assertEqual(node_types["XNNStaticTranspose"], 3, node_types) + self.assertEqual(node_types["XNNStaticReshape"], 0, node_types) + tester.to_executorch().serialize().run_method_and_compare_outputs() + + class ConvNoiseAddConvCumsum(torch.nn.Module): + """conv -> x + noise[H, W] -> conv, and cumsum(noise) outside the + delegate: the delegate gets its own folded copy of the constant and + the non-delegated consumer keeps the original.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("noise", torch.randn(5, 7)) + + def forward(self, x): + x = self.conv1(x) + x = x + self.noise + return self.conv2(x), torch.cumsum(self.noise, 0) + + def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_shared_with_non_delegated_op( + self, + ): + inputs = (torch.randn(1, 3, 5, 7),) + tester = ( + Tester(self.ConvNoiseAddConvCumsum().eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count( + { + "torch.ops.higher_order.executorch_call_delegate": 1, + "executorch_exir_dialects_edge__ops_aten_cumsum_default": 1, + } + ) + ) + self.assertEqual(self._delegate_input_buffers(tester), set()) + self._assert_only_io_transposes(tester) + tester.to_executorch().serialize().run_method_and_compare_outputs() + def test_fp32_channels_last_tagged_reshape_pass_input_lower_rank_operand(self): # A non-constant [H, W] operand is reshaped and converted in the graph. inputs = (torch.randn(1, 3, 5, 7), torch.randn(5, 7)) diff --git a/backends/xnnpack/utils/utils.py b/backends/xnnpack/utils/utils.py index 2d66ef75aea..38a08f3dcfc 100644 --- a/backends/xnnpack/utils/utils.py +++ b/backends/xnnpack/utils/utils.py @@ -20,6 +20,7 @@ is_lifted_tensor_constant, is_param, ) +from torch.fx.experimental.symbolic_shapes import has_free_symbols from torchao.quantization.pt2e.utils import _is_conv_node, _is_conv_transpose_node @@ -144,6 +145,29 @@ def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool: ) +def is_foldable_broadcast_operand( + exp_prog: ExportedProgram, node: torch.fx.Node +) -> bool: + """ + A rank < 4 constant operand of a broadcasting op that ChannelsLastTaggedReshapePass + can reshape to 4D and lay out as NHWC ahead of time. A mutated buffer is + not constant, and a symbolic shape cannot be expressed in XNNStaticReshape. + """ + if node.op != "placeholder" or not is_param_node(exp_prog, node): + return False + signature = exp_prog.graph_signature + buffer = signature.inputs_to_buffers.get(node.name) + if buffer in signature.buffers_to_mutate.values(): + return False + val = node.meta.get("val", None) + return ( + isinstance(val, torch.Tensor) + and not has_free_symbols(val.shape) + and 0 < val.dim() < 4 + and val.numel() > 1 + ) + + def get_param_tensor( exp_prog: ExportedProgram, node: torch.fx.Node ) -> Optional[torch.Tensor]: