From f7bb1c54d516fc7b322e4b58b9784b7bfaf550bf Mon Sep 17 00:00:00 2001 From: borys3kk Date: Thu, 22 Oct 2026 12:33:26 +0200 Subject: [PATCH 1/9] [XNNPACK] Lower upsample_nearest2d x2 and keep NHWC across scalar / low-rank broadcast operands Profiling a conv encoder-decoder (MI-GAN inpainting, fp16) showed the graph split into 7-8 XNNPACK partitions around portable upsample_nearest2d ops and ~25% of the time spent in NCHW<->NHWC transposes. upsample_nearest2d.vec with an exact x2 scale on a static 4D input is now partitioned and lowered as concatenate4 of the input along channels followed by depth_to_space_2d(block_size=2). XNNPACK has no nearest resize and its static_broadcast is not implemented at run time. This adds an XNNDepthToSpace2d node to the flatbuffer schemas and XNNCompiler. Static quantization is supported: the quantizer annotates the op with shared input/output qparams and the intermediate value carries them. Other scales, including an explicit scale_factors that is not exactly 2.0, stay portable. 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. --- .../channels_last_tagged_reshape_pass.py | 88 ++++++++++++- backends/xnnpack/operators/__init__.py | 1 + .../operators/op_upsample_nearest2d.py | 124 ++++++++++++++++++ backends/xnnpack/partition/config/__init__.py | 2 + .../partition/config/generic_node_configs.py | 38 ++++++ .../xnnpack/quantizer/xnnpack_quantizer.py | 3 + .../quantizer/xnnpack_quantizer_utils.py | 35 +++++ backends/xnnpack/runtime/XNNCompiler.cpp | 34 +++++ .../xnnpack/serialization/runtime_schema.fbs | 8 ++ backends/xnnpack/serialization/schema.fbs | 8 ++ .../serialization/xnnpack_graph_schema.py | 9 ++ .../test/ops/test_upsample_nearest2d.py | 107 +++++++++++++++ .../test_channels_last_tagged_reshape.py | 115 ++++++++++++++++ 13 files changed, 565 insertions(+), 7 deletions(-) create mode 100644 backends/xnnpack/operators/op_upsample_nearest2d.py create mode 100644 backends/xnnpack/test/ops/test_upsample_nearest2d.py diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 5238c9689ee..15f156bc679 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -61,6 +61,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): memory_sensitive_ops_nhwc = { exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.upsample_bilinear2d.vec, + exir_ops.edge.aten.upsample_nearest2d.vec, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.max_pool2d.default, exir_ops.edge.aten.amax.default, @@ -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,63 @@ 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_single_element(node: torch.fx.Node) -> bool: + val = node.meta.get("val", None) + return isinstance(val, torch.Tensor) and val.numel() == 1 + + @staticmethod + def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: + val = node.meta.get("val", None) + return ( + isinstance(val, torch.Tensor) + 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 +629,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 +686,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/operators/__init__.py b/backends/xnnpack/operators/__init__.py index e06f337f9ee..4b7b93a8174 100644 --- a/backends/xnnpack/operators/__init__.py +++ b/backends/xnnpack/operators/__init__.py @@ -57,5 +57,6 @@ op_sub, op_tanh, op_to_copy, + op_upsample_nearest2d, op_view_copy, ) diff --git a/backends/xnnpack/operators/op_upsample_nearest2d.py b/backends/xnnpack/operators/op_upsample_nearest2d.py new file mode 100644 index 00000000000..6e96c24afab --- /dev/null +++ b/backends/xnnpack/operators/op_upsample_nearest2d.py @@ -0,0 +1,124 @@ +# 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 Dict + +import torch +from executorch.backends.transforms import get_shape +from executorch.backends.xnnpack.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.xnnpack.operators.quant_params import QuantParams +from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( + XNNConcatenate4, + XNNDepthToSpace2d, + XNNGraph, + XNNQuantizedTensorValue, + XNNTensorValue, + XNode, + XValue, +) +from executorch.backends.xnnpack.utils.utils import check_or_raise, get_input_node +from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_INVALID_VALUE_ID + + +@register_node_visitor +class UpsampleNearest2dVisitor(NodeVisitor): + """ + XNNPACK has no nearest resize, so x2 upsampling is lowered as four copies + of the input concatenated along channels, [N, H, W, 4C], followed by + depth_to_space_2d(block_size=2). Neither op changes values, so the + intermediate carries the input's quantization parameters. + """ + + target = "aten.upsample_nearest2d.vec" + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + xnn_graph: XNNGraph, + vals_to_ids: Dict[torch.fx.Node, int], + debug_handle: int, + ) -> None: + input_node = get_input_node(node, 0) + input_quant_params = QuantParams.from_inputs(input_node, self._exported_program) + self.define_tensor( + input_node, + xnn_graph, + vals_to_ids, + convert_to_nhwc=True, + quant_params=input_quant_params, + ) + self.define_tensor( + node, + xnn_graph, + vals_to_ids, + convert_to_nhwc=True, + quant_params=QuantParams.from_outputs(node), + ) + input_id = vals_to_ids[input_node] + output_id = vals_to_ids[node] + + in_shape = get_shape(input_node) + out_shape = get_shape(node) + check_or_raise(len(in_shape) == 4, "upsample_nearest2d expects a 4D input") + n, c, h, w = in_shape + check_or_raise( + out_shape[-2] == 2 * h and out_shape[-1] == 2 * w, + f"only x2 nearest upsampling is supported, got {in_shape} -> {out_shape}", + ) + + concat_id = len(xnn_graph.xvalues) + concat_value = XNNTensorValue( + datatype=self.get_serialized_dtype(input_quant_params, node), + num_dims=4, + dims=[n, h, w, 4 * c], + external_id=XNN_INVALID_VALUE_ID, + constant_buffer_idx=0, + flags=0, + id_out=concat_id, + ) + xnn_graph.xvalues.append( + XValue(xvalue_union=concat_value) + if input_quant_params is None + else XValue( + xvalue_union=XNNQuantizedTensorValue( + tensor_value=concat_value, + quant_params=self.get_quant_params(input_quant_params, xnn_graph), + ) + ) + ) + + xnn_graph.xnodes.append( + XNode( + xnode_union=XNNConcatenate4( + axis=3, + input1_id=input_id, + input2_id=input_id, + input3_id=input_id, + input4_id=input_id, + output_id=concat_id, + flags=0, + input5_id=XNN_INVALID_VALUE_ID, + ), + debug_handle=debug_handle, + ) + ) + xnn_graph.xnodes.append( + XNode( + xnode_union=XNNDepthToSpace2d( + block_size=2, + input_id=concat_id, + output_id=output_id, + flags=0, + ), + debug_handle=debug_handle, + ) + ) diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index 7775a674510..e0f4e560666 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/backends/xnnpack/partition/config/__init__.py @@ -60,6 +60,7 @@ ToDimOrderCopyConfig, UnsqueezeCopyConfig, UpsampleBilinear2dConfig, + UpsampleNearest2dConfig, ViewCopyConfig, ) from executorch.backends.xnnpack.partition.config.node_configs import ( @@ -125,6 +126,7 @@ SubConfig, UnsqueezeCopyConfig, UpsampleBilinear2dConfig, + UpsampleNearest2dConfig, ViewCopyConfig, # Quant/Dequant Op Configs QuantizedPerTensorConfig, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index d445b5ca8ac..0f80efcf138 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -503,6 +503,44 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]: return torch.ops.aten.upsample_bilinear2d.vec +class UpsampleNearest2dConfig(GenericNodePartitionerConfig): + target_name = "upsample_nearest2d.vec" + + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + """ + Lowered as concatenate4 + depth_to_space_2d, which only expresses + exact x2 scaling of a static 4D tensor. + """ + if not self.check_common_constraints(node, ep): + return False + + if is_shape_dynamic(node): + why(node, reason="dynamic output sizes are not supported") + return False + + in_shape = get_input_node(node, 0).meta["val"].shape + out_shape = node.meta["val"].shape + if len(in_shape) != 4: + why(node, reason="only 4D inputs are supported") + return False + if out_shape[-2] != 2 * in_shape[-2] or out_shape[-1] != 2 * in_shape[-1]: + why(node, reason="only x2 nearest upsampling is supported") + return False + # PyTorch indexes with an explicit scale as given, so e.g. 2.004 + # doubles the size but picks different source pixels than x2. + scale_factors = node.args[2] if len(node.args) > 2 else None + if scale_factors is not None and any(s != 2.0 for s in scale_factors): + why(node, reason="only an exact x2 scale factor is supported") + return False + return True + + def supported_precision_types(self) -> List[ConfigPrecisionType]: + return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] + + def get_original_aten(self) -> Optional[torch._ops.OpOverload]: + return torch.ops.aten.upsample_nearest2d.vec + + class ExpConfig(GenericNodePartitionerConfig): target_name = "exp.default" diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer.py b/backends/xnnpack/quantizer/xnnpack_quantizer.py index 3c82a65ad71..cd2b4eb2f6c 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer.py @@ -272,6 +272,8 @@ class QuantPattern: CAT_TARGETS = {torch.ops.aten.cat.default} +UPSAMPLE_NEAREST2D_TARGETS = {torch.ops.aten.upsample_nearest2d.vec} + class XNNPACKQuantizer(Quantizer): supported_config_and_operators = _get_supported_config_and_operators() @@ -292,6 +294,7 @@ class XNNPACKQuantizer(Quantizer): QuantPattern("mul_relu", False, False, MUL_TARGETS), QuantPattern("mul", False, False, MUL_TARGETS), QuantPattern("cat", False, False, CAT_TARGETS), + QuantPattern("upsample_nearest2d", False, False, UPSAMPLE_NEAREST2D_TARGETS), ] def __init__(self) -> None: diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index a207d77674d..8755ea0c52b 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -727,6 +727,41 @@ def _annotate_adaptive_avg_pool2d( return annotated_partitions +@register_annotator("upsample_nearest2d") +def _annotate_upsample_nearest2d( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + """Nearest upsampling only copies values, so its output shares the + input's quantization parameters.""" + annotated_partitions = [] + for node in gm.graph.nodes: + if node.target != torch.ops.aten.upsample_nearest2d.vec: + continue + if _is_annotated([node]) or (filter_fn and not filter_fn(node)): + continue + + input_act = node.args[0] + assert isinstance(input_act, Node) + if ( + Q_ANNOTATION_KEY not in input_act.meta + or not input_act.meta[Q_ANNOTATION_KEY]._annotated + or input_act.meta[Q_ANNOTATION_KEY].output_qspec is None + ): + input_act_qspec = get_input_act_qspec(quantization_config) + else: + input_act_qspec = SharedQuantizationSpec(input_act) + + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map={input_act: input_act_qspec}, + output_qspec=SharedQuantizationSpec((input_act, node)), + _annotated=True, + ) + annotated_partitions.append([node]) + return annotated_partitions + + def _is_input_large_scalar(node: Node, gm: torch.fx.GraphModule): """Check if input is a large scalar value. So that we can skip quantization for the node since histc op (in HistogramObserver) only works for values up to certain upper bound diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 74a5eec4a91..c3b88bdc445 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -1642,6 +1642,39 @@ Error defineConcatenate5Node( return Error::Ok; } +/* +Defines serialized depth_to_space_2d node into the subgraph, +using the remapped ids to map the serialized ids, +to the new ids generated when defining the tensor value +*/ +Error defineDepthToSpace2dNode( + xnn_subgraph_t subgraph_ptr, + const std::unordered_map& remapped_ids, + const NodePtr node, + const fb_xnnpack::XNNGraph* graph) noexcept { + MAYBE_UNUSED(graph); + + auto graph_node = node->xnode_union_as_XNNDepthToSpace2d(); + + REMAP_ID(remapped_ids, graph_node->input_id(), d2s_input); + REMAP_ID(remapped_ids, graph_node->output_id(), d2s_output); + + xnn_status status = xnn_define_depth_to_space_2d( + subgraph_ptr, + graph_node->block_size(), + d2s_input, + d2s_output, + graph_node->flags()); + ET_CHECK_OR_RETURN_ERROR( + status == xnn_status_success, + Internal, + "Failed to create depth_to_space node %i with code: %s", + node->debug_handle(), + xnn_status_to_string(status)); + + return Error::Ok; +} + /* Defines serialized static_slice node into the subgraph, using the remapped ids to map the serialized ids, @@ -2052,6 +2085,7 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) { _DEFINE(Convert) _DEFINE(GlobalAvgPooling2d) _DEFINE(StaticReshape) + _DEFINE(DepthToSpace2d) _DEFINE(ArgMaxPooling2d) _DEFINE(Concatenate2) _DEFINE(Concatenate3) diff --git a/backends/xnnpack/serialization/runtime_schema.fbs b/backends/xnnpack/serialization/runtime_schema.fbs index 888384e032e..62143ab8c62 100644 --- a/backends/xnnpack/serialization/runtime_schema.fbs +++ b/backends/xnnpack/serialization/runtime_schema.fbs @@ -160,6 +160,7 @@ union XNodeUnion { XNNCopy: _XNNNode1x1, XNNCos: _XNNNode1x1, XNNApproxGelu: _XNNNode1x1, + XNNDepthToSpace2d, } union XValueUnion { @@ -298,6 +299,13 @@ table XNNStaticReshape { flags: uint; } +table XNNDepthToSpace2d { + block_size: uint; + input_id: uint; + output_id: uint; + flags: uint; +} + table XNNStaticSlice { num_dims:uint; offsets:[uint]; diff --git a/backends/xnnpack/serialization/schema.fbs b/backends/xnnpack/serialization/schema.fbs index 96e4256f078..a48f9bc4187 100644 --- a/backends/xnnpack/serialization/schema.fbs +++ b/backends/xnnpack/serialization/schema.fbs @@ -156,6 +156,7 @@ union XNodeUnion { XNNCopy: _XNNNode1x1, XNNCos: _XNNNode1x1, XNNApproxGelu: _XNNNode1x1, + XNNDepthToSpace2d, } union XValueUnion { @@ -294,6 +295,13 @@ table XNNStaticReshape { flags: uint; } +table XNNDepthToSpace2d { + block_size: uint; + input_id: uint; + output_id: uint; + flags: uint; +} + table XNNStaticSlice { num_dims:uint; offsets:[uint]; diff --git a/backends/xnnpack/serialization/xnnpack_graph_schema.py b/backends/xnnpack/serialization/xnnpack_graph_schema.py index 360034c1987..e6da4338dd0 100644 --- a/backends/xnnpack/serialization/xnnpack_graph_schema.py +++ b/backends/xnnpack/serialization/xnnpack_graph_schema.py @@ -202,6 +202,14 @@ class XNNStaticTranspose: flags: int +@dataclass +class XNNDepthToSpace2d: + block_size: int + input_id: int + output_id: int + flags: int + + @dataclass class XNNStaticSlice: num_dims: int @@ -427,6 +435,7 @@ class XNNScaledDotProductAttention: XNNCopy, XNNCos, XNNApproxGelu, + XNNDepthToSpace2d, ] diff --git a/backends/xnnpack/test/ops/test_upsample_nearest2d.py b/backends/xnnpack/test/ops/test_upsample_nearest2d.py new file mode 100644 index 00000000000..a4a105e6b09 --- /dev/null +++ b/backends/xnnpack/test/ops/test_upsample_nearest2d.py @@ -0,0 +1,107 @@ +# 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 TestUpsampleNearest2d(unittest.TestCase): + def setUp(self): + torch._dynamo.reset() + + class Nearest2x(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + + class Nearest2xBySize(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate( + x, size=(x.shape[2] * 2, x.shape[3] * 2), mode="nearest" + ) + + class Nearest3x(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate(x, scale_factor=3.0, mode="nearest") + + class NearestInexact2x(torch.nn.Module): + """Doubles the size, but PyTorch indexes with 1/2.004, so the x2 + duplication lowering would not match eager.""" + + def forward(self, x): + return torch.nn.functional.interpolate( + x, scale_factor=2.004, mode="nearest" + ) + + class ConvNearestConv(torch.nn.Module): + """The pattern that used to split XNNPACK partitions in conv decoders.""" + + 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.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + return self.conv2(x) + + edge_op = "executorch_exir_dialects_edge__ops_aten_upsample_nearest2d_vec" + + def _test_lowered(self, module, inputs, quant=False): + tester = Tester(module, inputs) + if quant: + tester.quantize() + tester.export().to_edge_transform_and_lower().check_not([self.edge_op]) + if quant: + tester.check_not(["torch.ops.quantized_decomposed"]) + ( + tester.check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(qtol=2 if quant else 0) + ) + + def test_fp32_upsample_nearest2d_scale_factor(self): + self._test_lowered(self.Nearest2x(), (torch.randn(2, 3, 4, 5),)) + + def test_fp32_upsample_nearest2d_size(self): + self._test_lowered(self.Nearest2xBySize(), (torch.randn(1, 5, 7, 3),)) + + def test_fp16_upsample_nearest2d(self): + self._test_lowered( + self.Nearest2x(), (torch.randn(1, 6, 8, 8).to(torch.float16),) + ) + + def test_fp32_conv_upsample_conv_single_partition(self): + self._test_lowered(self.ConvNearestConv(), (torch.randn(1, 3, 16, 16),)) + + def test_qs8_conv_upsample_conv_single_partition(self): + self._test_lowered( + self.ConvNearestConv(), (torch.randn(1, 3, 16, 16),), quant=True + ) + + def test_fp32_upsample_nearest2d_x3_not_partitioned(self): + ( + Tester(self.Nearest3x(), (torch.randn(1, 3, 4, 4),)) + .export() + .to_edge_transform_and_lower() + .check_count({self.edge_op: 1}) + .check_not(["torch.ops.higher_order.executorch_call_delegate"]) + ) + + def test_fp32_upsample_nearest2d_inexact_scale_not_partitioned(self): + ( + Tester(self.NearestInexact2x(), (torch.randn(1, 3, 200, 200),)) + .export() + .to_edge_transform_and_lower() + .check_count({self.edge_op: 1}) + .check_not(["torch.ops.higher_order.executorch_call_delegate"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) 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..53bb958df06 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,121 @@ 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() + ) + 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 d6e8816d15989d3f77205aa3c22c2dcdc29ecb34 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Wed, 21 Oct 2026 15:27:22 +0200 Subject: [PATCH 2/9] Make delete_constant_placeholder edit the graph signature in place Backend passes keep a reference to the ExportedProgram they were built from (XNNPACKPass.exported_program, for one), while ExportedProgramPassManager runs them on a shallow copy of it and rebuilds the signature of that copy afterwards. delete_constant_placeholder replaced the original program's signature object, so the copy kept the old one and its verifier failed with Parameter is not in the state dict. create_constant_placeholder already inserts into the shared input_specs list in place; do the same for deletion. Tests cover a deletion made inside a pass run through the manager for every constant placeholder kind. Authored with Claude Code. --- ...test_create_delete_constant_placeholder.py | 58 +++++++++++++++++++ backends/transforms/utils.py | 15 ++--- 2 files changed, 63 insertions(+), 10 deletions(-) diff --git a/backends/transforms/test/test_create_delete_constant_placeholder.py b/backends/transforms/test/test_create_delete_constant_placeholder.py index 9444e6a4b63..4c4bc9f7822 100644 --- a/backends/transforms/test/test_create_delete_constant_placeholder.py +++ b/backends/transforms/test/test_create_delete_constant_placeholder.py @@ -10,6 +10,8 @@ ) from executorch.exir import to_edge from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from executorch.exir.pass_manager import ExportedProgramPassManager from torch.export import export from torch.export.graph_signature import InputKind @@ -204,3 +206,59 @@ def test_create_delete_non_persistent_buffer(): def test_create_delete_constant_tensor(): _test_create_delete(InputKind.CONSTANT_TENSOR) + + +def _test_delete_inside_a_pass(kind: InputKind, persistent_buffer: bool = None): + """Backend passes keep a reference to the program they were built from, + while ExportedProgramPassManager runs them on a shallow copy of it. The + deletion has to edit the shared signature in place, or the copy that the + manager verifies and returns still lists the deleted placeholder.""" + module = EmptyNetwork() + exported_program = export(module, args=module.test_data, strict=True) + exported_program = to_edge(exported_program).exported_program() + graph = exported_program.graph_module.graph + + const_name = "test_node" + input_node = list(graph.nodes)[0] + with graph.inserting_before(input_node): + const_node = create_constant_placeholder( + exp_program=exported_program, + graph=graph, + kind=kind, + name=const_name, + data=torch.ones(1), + persistent_buffer=persistent_buffer, + ) + assert len(exported_program.graph_signature.input_specs) == 2 + + def delete_unused_constant(graph_module: torch.fx.GraphModule) -> PassResult: + delete_constant_placeholder(exported_program, const_node) + return PassResult(graph_module, True) + + result = ExportedProgramPassManager([delete_unused_constant])(exported_program) + transformed = result.exported_program + signature = transformed.graph_signature + + assert len(signature.input_specs) == 1 + assert const_name not in signature.inputs_to_parameters + assert const_name not in signature.inputs_to_buffers + assert const_name not in signature.inputs_to_lifted_tensor_constants + assert len(transformed.state_dict) == 0 + assert len(transformed.constants) == 0 + assert transformed.module()(torch.zeros(1)) == 0 + + +def test_delete_inside_a_pass_parameter(): + _test_delete_inside_a_pass(InputKind.PARAMETER) + + +def test_delete_inside_a_pass_persistent_buffer(): + _test_delete_inside_a_pass(InputKind.BUFFER, True) + + +def test_delete_inside_a_pass_non_persistent_buffer(): + _test_delete_inside_a_pass(InputKind.BUFFER, False) + + +def test_delete_inside_a_pass_constant_tensor(): + _test_delete_inside_a_pass(InputKind.CONSTANT_TENSOR) diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index d3c50c692b1..0ecfb99c8c0 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -266,16 +266,11 @@ def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Nod f"Cannot delete input node {node.name} since it is not a parameter, a buffer, nor a lifted tensor constant." ) - # Remove input from graph signature - input_specs = [ - spec - for spec in exp_program.graph_signature.input_specs - if spec.arg.name != node.name - ] - new_graph_signature = ExportGraphSignature( - input_specs, exp_program.graph_signature.output_specs - ) - exp_program._graph_signature = new_graph_signature + # Edit the spec list in place: ExportedProgramPassManager runs passes on a + # shallow copy of the program, which would not see a replaced signature. + input_specs = exp_program.graph_signature.input_specs + for spec in [spec for spec in input_specs if spec.arg.name == node.name]: + input_specs.remove(spec) # Remove node from graph node.graph.erase_node(node) From b929250ab985f4c71dea3962b37e2eadbf0ca340 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Tue, 20 Oct 2026 16:14:35 +0200 Subject: [PATCH 3/9] [XNNPACK] Lower depthwise transposed convs as polyphase depthwise convs + depth_to_space XNNPACK runs transposed convolutions through its subconvolution IGEMM path one group at a time. With one channel per group (a depthwise ConvTranspose2d, the natural way to write a StyleGAN upfirdn x2 upsampling filter) that is far slower than the depthwise conv microkernels: on MI-GAN the twelve 4x4 / stride-2 filters more than double the forward. DecomposeDepthwiseConvTransposePass rewrites such a node, when it is a clean x s upscale ((H-1)s - 2p + k + output_padding == sH, dilation 1, every phase's taps inside the input), into s*s stride-1 depthwise convs on the low-resolution input (phase (ay, ax) correlates with every s-th tap of the kernel, through a possibly asymmetric zero pad), a channel concatenation in (ay, ax) order and a block-major depth_to_space_2d. That last op is XNNPACK's native channel order (aten.pixel_shuffle is channel-major), so it is a new et_xnnpack custom op created only after partitioning, with a reference implementation for tracing and tests, a node visitor emitting XNNDepthToSpace2d, and an NHWC tag in the channels-last pass. The pass runs before FuseActivationPass, so an activation after the transposed conv is never fused into the node it removes. Every convolution is planned before any is rewritten, so deconvs feeding each other work, and the phase weights keep the original weight's input kind (parameter, buffer or constant). Deleting the folded weight relies on delete_constant_placeholder editing the graph signature in place. Authored with Claude Code. --- backends/xnnpack/BUCK | 11 + backends/xnnpack/_passes/BUCK | 1 + backends/xnnpack/_passes/__init__.py | 4 + .../channels_last_tagged_reshape_pass.py | 2 + .../decompose_depthwise_conv_transpose.py | 218 ++++++++++++++ backends/xnnpack/custom_ops_lib.py | 43 +++ backends/xnnpack/operators/BUCK | 1 + backends/xnnpack/operators/__init__.py | 1 + .../xnnpack/operators/op_depth_to_space.py | 70 +++++ ...test_decompose_depthwise_conv_transpose.py | 267 ++++++++++++++++++ 10 files changed, 618 insertions(+) create mode 100644 backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py create mode 100644 backends/xnnpack/custom_ops_lib.py create mode 100644 backends/xnnpack/operators/op_depth_to_space.py create mode 100644 backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py diff --git a/backends/xnnpack/BUCK b/backends/xnnpack/BUCK index 1e12548bc4c..f26b7fb6b0c 100644 --- a/backends/xnnpack/BUCK +++ b/backends/xnnpack/BUCK @@ -45,3 +45,14 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/xnnpack/utils:xnnpack_utils", ], ) + +fbcode_target(_kind = runtime.python_library, + name = "custom_ops_lib", + srcs = [ + "custom_ops_lib.py", + ], + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + ], +) diff --git a/backends/xnnpack/_passes/BUCK b/backends/xnnpack/_passes/BUCK index 2edf294fa4e..081b0ff31b2 100644 --- a/backends/xnnpack/_passes/BUCK +++ b/backends/xnnpack/_passes/BUCK @@ -11,6 +11,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/transforms:addmm_mm_to_linear", "//executorch/backends/transforms:remove_clone_ops", "//executorch/backends/transforms:lib", + "//executorch/backends/xnnpack:custom_ops_lib", "//executorch/backends/xnnpack/partition:partitioner_graphs", "//executorch/backends/xnnpack/serialization:xnnpack_schema", "//executorch/backends/xnnpack/utils:xnnpack_utils", diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index f55144395af..383734d466a 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/backends/xnnpack/_passes/__init__.py @@ -26,6 +26,9 @@ ) from executorch.backends.xnnpack._passes.decompose_batch_norm import DecomposeBatchNorm from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate +from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( + DecomposeDepthwiseConvTransposePass, +) 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.insert_pad_qdq import InsertPadQDQPass @@ -82,6 +85,7 @@ def __init__( ConstPropPass, FuseBatchNormPass, DecomposeBatchNorm, + DecomposeDepthwiseConvTransposePass, FuseActivationPass, DecomposeConcatenate, RemoveGetItemPass, diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 15f156bc679..31e2e71214b 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -9,6 +9,7 @@ from enum import Enum from typing import List, Optional, Tuple +import executorch.backends.xnnpack.custom_ops_lib # noqa: F401 registers et_xnnpack ops import torch from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.backends.xnnpack.utils.quant_utils import ( @@ -62,6 +63,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.upsample_bilinear2d.vec, exir_ops.edge.aten.upsample_nearest2d.vec, + exir_ops.edge.et_xnnpack.depth_to_space_2d.default, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.max_pool2d.default, exir_ops.edge.aten.amax.default, diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py new file mode 100644 index 00000000000..875d7de5f11 --- /dev/null +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -0,0 +1,218 @@ +# 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 math +from typing import cast, List, Optional, Tuple + +import executorch.backends.xnnpack.custom_ops_lib # noqa: F401 registers et_xnnpack ops +import torch +from executorch.backends.transforms.utils import ( + create_constant_placeholder, + delete_constant_placeholder, + get_param_tensor, + is_param_node, +) +from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult + + +class DecomposeDepthwiseConvTransposePass(XNNPACKPass): + """ + Lower a depthwise transposed conv (groups == in_channels == out_channels) + with stride s > 1 into s*s depthwise convs on the low-resolution input, a + channel concatenation and a block-major depth_to_space_2d. + + XNNPACK runs transposed convs through its subconvolution IGEMM path, one + group at a time; with one channel per group that is far slower than its + depthwise conv microkernels. Polyphase decomposition: output phase a of a + stride-s transposed conv is an ordinary stride-1 correlation of the input + with every s-th tap of the kernel, + + y[s*m + a] = sum_t x[m - t] * w[s*t + a + p], 0 <= s*t + a + p < k, + + per axis. The s*s phase outputs concatenated along channels in (ay, ax) + order are exactly the block-major input depth_to_space_2d expects. Each + phase reads the input through a (possibly asymmetric) zero pad. + + Only applied when the transposed conv is a clean x s upscale, i.e. + (H - 1)*s - 2p + k + output_padding == s*H on both axes with the same s, + dilation is 1 and every phase's taps stay inside the input; any other + transposed conv is left to XNNPACK's own node. + """ + + def call(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + # Plan every conv before rewriting any: the nodes created by + # _decompose have no meta["val"] until the retrace below. + plans = [] + for node in graph.nodes: + if ( + node.op == "call_function" + and node.target == exir_ops.edge.aten.convolution.default + ): + plan = self._plan(node) + if plan is not None: + plans.append((node, plan)) + for node, plan in plans: + self._decompose(graph, node, *plan) + + modified = bool(plans) + if modified: + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, modified) + + # (t_min, t_max) per phase for one axis, None when the axis + # does not decompose. + @staticmethod + def _axis_phases( + size: int, kernel: int, stride: int, padding: int, output_padding: int + ) -> Optional[List[Tuple[int, int]]]: + if stride < 2: + return None + if (size - 1) * stride - 2 * padding + kernel + output_padding != size * stride: + return None + phases = [] + for a in range(stride): + # integer t with 0 <= stride*t + a + padding <= kernel - 1 + t_min = math.ceil(-(a + padding) / stride) + t_max = math.floor((kernel - 1 - a - padding) / stride) + if t_min > 0 or t_max < 0: + return None + phases.append((t_min, t_max)) + return phases + + def _plan(self, node: torch.fx.Node): + ( + x, + weight, + _bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) = node.args + if not transposed: + return None + if not is_param_node(self.exported_program, weight): + return None + x_shape = list(x.meta["val"].shape) + w_shape = list(weight.meta["val"].shape) + if len(x_shape) != 4 or len(w_shape) != 4: + return None + channels = x_shape[1] + # transposed conv weight is (in, out / groups, kh, kw) + if groups != channels or w_shape[0] != channels or w_shape[1] != 1: + return None + if any(d != 1 for d in dilation): + return None + # ATen broadcasts a single spatial value over both spatial dims. + stride, padding, output_padding = ( + p * 2 if len(p) == 1 else p + for p in cast(List[List[int]], [stride, padding, output_padding]) + ) + if stride[0] != stride[1]: + return None + phases_y = self._axis_phases( + x_shape[2], w_shape[2], stride[0], padding[0], output_padding[0] + ) + phases_x = self._axis_phases( + x_shape[3], w_shape[3], stride[1], padding[1], output_padding[1] + ) + if phases_y is None or phases_x is None: + return None + return stride[0], padding, phases_y, phases_x + + def _decompose( + self, + graph: torch.fx.Graph, + node: torch.fx.Node, + s: int, + padding: List[int], + phases_y: List[Tuple[int, int]], + phases_x: List[Tuple[int, int]], + ) -> None: + x, weight_node, bias_node = node.args[0], node.args[1], node.args[2] + groups = node.args[8] + weight = get_param_tensor(self.exported_program, weight_node) + assert weight is not None + weight_spec = next( + spec + for spec in self.exported_program.graph_signature.input_specs + if spec.arg.name == weight_node.name + ) + + phase_convs = [] + for ay, (ty_min, ty_max) in enumerate(phases_y): + for ax, (tx_min, tx_max) in enumerate(phases_x): + # taps of phase (ay, ax) as a stride-1 correlation kernel: + # W[q] = w[s * (t_max - q) + a + p] + iy = torch.tensor( + [ + s * (ty_max - q) + ay + padding[0] + for q in range(ty_max - ty_min + 1) + ] + ) + ix = torch.tensor( + [ + s * (tx_max - q) + ax + padding[1] + for q in range(tx_max - tx_min + 1) + ] + ) + phase_weight = weight[:, :, iy][:, :, :, ix].contiguous() + with graph.inserting_before(weight_node): + phase_weight_node = create_constant_placeholder( + exp_program=self.exported_program, + graph=graph, + kind=weight_spec.kind, + persistent_buffer=weight_spec.persistent, + name=f"{weight_node.name}_phase_{ay}_{ax}", + data=phase_weight, + ) + # (left, right, top, bottom) + pad = [tx_max, -tx_min, ty_max, -ty_min] + with graph.inserting_before(node): + conv_input = x + if any(pad): + conv_input = graph.create_node( + "call_function", + exir_ops.edge.aten.constant_pad_nd.default, + (x, pad, 0.0), + ) + phase_convs.append( + graph.create_node( + "call_function", + exir_ops.edge.aten.convolution.default, + ( + conv_input, + phase_weight_node, + bias_node, + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + groups, + ), + ) + ) + + with graph.inserting_before(node): + cat = graph.create_node( + "call_function", exir_ops.edge.aten.cat.default, (phase_convs, 1) + ) + d2s = graph.create_node( + "call_function", + exir_ops.edge.et_xnnpack.depth_to_space_2d.default, + (cat, s), + ) + node.replace_all_uses_with(d2s) + graph.erase_node(node) + if len(weight_node.users) == 0: + delete_constant_placeholder(self.exported_program, weight_node) diff --git a/backends/xnnpack/custom_ops_lib.py b/backends/xnnpack/custom_ops_lib.py new file mode 100644 index 00000000000..5c393b54238 --- /dev/null +++ b/backends/xnnpack/custom_ops_lib.py @@ -0,0 +1,43 @@ +# 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. + +""" +Custom ops that XNNPACK passes introduce after partitioning. They never reach +the partitioner: a pass creates them inside an already lowered subgraph and a +node visitor serializes them, the reference implementation below only serves +tracing and tests. +""" + +import torch.library + +namespace = "et_xnnpack" +lib = torch.library.Library(namespace, "DEF") + +####################### +## depth_to_space_2d ## +####################### + + +def depth_to_space_2d_impl(x: torch.Tensor, block_size: int) -> torch.Tensor: + """ + Block-major depth to space, the channel order XNNPACK's (and TFLite's) + DEPTH_TO_SPACE uses: input channel (by * block_size + bx) * C + c lands at + output pixel (h * block_size + by, w * block_size + bx), channel c. + aten.pixel_shuffle is the channel-major variant, c * block_size**2 + by * + block_size + bx, and does not map onto XNNPACK's node directly. + """ + n, c_in, h, w = x.shape + bs = block_size + c = c_in // (bs * bs) + y = x.reshape(n, bs, bs, c, h, w) + y = y.permute(0, 3, 4, 1, 5, 2) # n, c, h, by, w, bx + return y.reshape(n, c, h * bs, w * bs) + + +name = "depth_to_space_2d" +lib.define(f"{name}(Tensor x, int block_size) -> Tensor") +lib.impl(name, depth_to_space_2d_impl, "CompositeExplicitAutograd") +depth_to_space_2d_op = torch.ops.et_xnnpack.depth_to_space_2d diff --git a/backends/xnnpack/operators/BUCK b/backends/xnnpack/operators/BUCK index fdc2ed1db49..a27c5d14026 100644 --- a/backends/xnnpack/operators/BUCK +++ b/backends/xnnpack/operators/BUCK @@ -8,6 +8,7 @@ fbcode_target(_kind = runtime.python_library, srcs = glob(["*.py"]), visibility = ["PUBLIC"], deps = [ + "//executorch/backends/xnnpack:custom_ops_lib", "//executorch/backends/xnnpack/utils:xnnpack_utils", "//executorch/exir:graph_module", "//executorch/exir/backend:backend_details", diff --git a/backends/xnnpack/operators/__init__.py b/backends/xnnpack/operators/__init__.py index 4b7b93a8174..3ee54158f7b 100644 --- a/backends/xnnpack/operators/__init__.py +++ b/backends/xnnpack/operators/__init__.py @@ -19,6 +19,7 @@ op_clone, op_conv2d, op_cos, + op_depth_to_space, op_div, op_dynamic_dequantize_ops, op_dynamic_quantize_ops, diff --git a/backends/xnnpack/operators/op_depth_to_space.py b/backends/xnnpack/operators/op_depth_to_space.py new file mode 100644 index 00000000000..1d913128515 --- /dev/null +++ b/backends/xnnpack/operators/op_depth_to_space.py @@ -0,0 +1,70 @@ +# 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, Dict + +import executorch.backends.xnnpack.custom_ops_lib # noqa: F401 registers et_xnnpack ops +import torch +from executorch.backends.transforms import get_shape +from executorch.backends.xnnpack.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( + XNNDepthToSpace2d, + XNNGraph, + XNode, +) +from executorch.backends.xnnpack.utils.utils import check_or_raise, get_input_node + + +@register_node_visitor +class DepthToSpace2dVisitor(NodeVisitor): + """ + et_xnnpack.depth_to_space_2d, the block-major depth to space that XNNPACK + provides natively. Introduced by DecomposeDepthwiseConvTransposePass, see + custom_ops_lib.py for the channel order. + """ + + target = "et_xnnpack.depth_to_space_2d.default" + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + xnn_graph: XNNGraph, + vals_to_ids: Dict[torch.fx.Node, int], + debug_handle: int, + ) -> None: + input_node = get_input_node(node, 0) + self.define_tensor(input_node, xnn_graph, vals_to_ids, convert_to_nhwc=True) + self.define_tensor(node, xnn_graph, vals_to_ids, convert_to_nhwc=True) + + block_size = cast(int, node.args[1]) + in_shape = get_shape(input_node) + out_shape = get_shape(node) + check_or_raise( + len(in_shape) == 4 + and block_size >= 1 + and in_shape[1] % (block_size * block_size) == 0 + and out_shape[2] == in_shape[2] * block_size + and out_shape[3] == in_shape[3] * block_size, + f"depth_to_space_2d(block_size={block_size}) cannot map {in_shape} -> {out_shape}", + ) + + xnn_graph.xnodes.append( + XNode( + xnode_union=XNNDepthToSpace2d( + block_size=block_size, + input_id=vals_to_ids[input_node], + output_id=vals_to_ids[node], + flags=0, + ), + debug_handle=debug_handle, + ) + ) diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py new file mode 100644 index 00000000000..b1ffc54eb04 --- /dev/null +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -0,0 +1,267 @@ +# 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 math +import unittest + +import torch +from executorch.backends.test.harness.stages import StageType +from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( + DecomposeDepthwiseConvTransposePass, +) +from executorch.backends.xnnpack.test.tester import RunPasses, Tester +from executorch.exir.dialects._ops import ops as exir_ops + + +class TestDecomposeDepthwiseConvTransposePass(unittest.TestCase): + PassStage = RunPasses([DecomposeDepthwiseConvTransposePass]) + + conv = "executorch_exir_dialects_edge__ops_aten_convolution_default" + pad = "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default" + cat = "executorch_exir_dialects_edge__ops_aten_cat_default" + d2s = "executorch_exir_dialects_edge__ops_et_xnnpack_depth_to_space_2d_default" + + def setUp(self): + torch._dynamo.reset() + + class DwDeconv(torch.nn.Module): + def __init__( + self, kernel, stride, padding, output_padding=0, bias=True, channels=8 + ): + super().__init__() + self.deconv = torch.nn.ConvTranspose2d( + channels, + channels, + kernel, + stride=stride, + padding=padding, + output_padding=output_padding, + groups=channels, + bias=bias, + ) + with torch.no_grad(): + self.deconv.weight.normal_() + if bias: + self.deconv.bias.normal_() + + def forward(self, x): + return self.deconv(x) + + class QuantizedDwDeconv(torch.nn.Module): + # The XNNPACK quantizer never annotates grouped transposed convs, so + # build the dequantized weight by hand. + def __init__(self): + super().__init__() + self.register_buffer( + "weight", torch.randint(-128, 127, (8, 1, 4, 4), dtype=torch.int8) + ) + + def forward(self, x): + weight = torch.ops.quantized_decomposed.dequantize_per_tensor( + self.weight, 0.02, 0, -128, 127, torch.int8 + ) + return torch.nn.functional.conv_transpose2d( + x, weight, stride=2, padding=1, groups=8 + ) + + def _check_pass(self, module, inputs, counts, tol=1e-3, dynamic_shapes=None): + return ( + Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes) + .export() + .to_edge() + .run_passes(self.PassStage) + .check_count(counts) + .run_method_and_compare_outputs(atol=tol, rtol=tol) + ) + + def _check_lowered(self, module, inputs, tol=1e-3, dynamic_shapes=None): + ( + Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not([self.conv]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(num_runs=3, atol=tol, rtol=tol) + ) + + def test_upfirdn_4x4_stride2(self): + # StyleGAN-style x2 upsampling filter: 4 phases of 2x2 taps, each with + # a one-sided pad. + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(1, 8, 6, 10),) + self._check_pass( + module, inputs, {self.conv: 4, self.pad: 4, self.cat: 1, self.d2s: 1} + ) + self._check_lowered(module, inputs) + + def test_no_bias_fp16(self): + module = self.DwDeconv(kernel=4, stride=2, padding=1, bias=False).half() + inputs = (torch.randn(1, 8, 5, 7).half(),) + # fp16 accumulation order differs from the reference: 1 ulp at |y| ~ 10 + self._check_pass(module, inputs, {self.conv: 4, self.d2s: 1}, tol=1e-2) + self._check_lowered(module, inputs, tol=1e-2) + + def test_3x3_stride2_output_padding(self): + # k=3, p=1, output_padding=1 keeps the x2 scaling; the phases have 1 + # and 2 taps per axis. + module = self.DwDeconv(kernel=3, stride=2, padding=1, output_padding=1) + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(module, inputs, {self.conv: 4, self.cat: 1, self.d2s: 1}) + self._check_lowered(module, inputs) + + def test_2x2_stride2_no_padding(self): + # Every phase is a single tap, no pads at all. + module = self.DwDeconv(kernel=2, stride=2, padding=0) + inputs = (torch.randn(1, 8, 4, 6),) + self._check_pass( + module, inputs, {self.conv: 4, self.pad: 0, self.cat: 1, self.d2s: 1} + ) + self._check_lowered(module, inputs) + + def test_stride3(self): + module = self.DwDeconv(kernel=3, stride=3, padding=0, channels=4) + inputs = (torch.randn(1, 4, 4, 5),) + self._check_pass(module, inputs, {self.conv: 9, self.cat: 1, self.d2s: 1}) + self._check_lowered(module, inputs) + + def test_conv_before_deconv(self): + # The phase pads read the conv's output rather than a graph input. + module = torch.nn.Sequential( + torch.nn.Conv2d(8, 8, 3, padding=1), + self.DwDeconv(kernel=4, stride=2, padding=1), + ) + inputs = (torch.randn(1, 8, 6, 10),) + tester = self._check_pass( + module, inputs, {self.conv: 5, self.pad: 4, self.cat: 1, self.d2s: 1} + ) + graph = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph + ) + for node in graph.nodes: + if node.target == exir_ops.edge.aten.constant_pad_nd.default: + self.assertEqual( + node.args[0].target, exir_ops.edge.aten.convolution.default + ) + self._check_lowered(module, inputs) + + def test_two_deconvs_in_a_row(self): + # The second deconv's input is the first one's depth_to_space_2d. + module = torch.nn.Sequential( + self.DwDeconv(kernel=4, stride=2, padding=1), + self.DwDeconv(kernel=4, stride=2, padding=1), + ) + inputs = (torch.randn(1, 8, 4, 6),) + self._check_pass( + module, inputs, {self.conv: 8, self.pad: 8, self.cat: 2, self.d2s: 2} + ) + self._check_lowered(module, inputs) + + def test_batch_gt_1(self): + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(3, 8, 6, 10),) + self._check_pass( + module, inputs, {self.conv: 4, self.pad: 4, self.cat: 1, self.d2s: 1} + ) + self._check_lowered(module, inputs) + + def test_single_element_spatial_params(self): + class Deconv(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(8, 1, 4, 4)) + + def forward(self, x): + return torch.ops.aten.convolution.default( + x, self.weight, None, [2], [1], [1], True, [0], 8 + ) + + inputs = (torch.randn(1, 8, 6, 10),) + self._check_pass(Deconv(), inputs, {self.conv: 4, self.pad: 4, self.d2s: 1}) + self._check_lowered(Deconv(), inputs) + + def test_deconv_relu(self): + # The activation follows depth_to_space_2d instead of fusing into the + # removed transposed conv. + module = torch.nn.Sequential( + self.DwDeconv(kernel=4, stride=2, padding=1), torch.nn.ReLU() + ) + inputs = (torch.randn(1, 8, 6, 10),) + self._check_lowered(module, inputs) + + def test_dynamic_shape(self): + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(1, 8, 6, 10),) + dynamic_shapes = ( + { + 2: torch.export.Dim("height", min=2, max=16), + 3: torch.export.Dim("width", min=2, max=16), + }, + ) + self._check_pass( + module, + inputs, + {self.conv: 4, self.pad: 4, self.cat: 1, self.d2s: 1}, + dynamic_shapes=dynamic_shapes, + ) + self._check_lowered(module, inputs, dynamic_shapes=dynamic_shapes) + + def test_quantized_is_left_alone(self): + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(self.QuantizedDwDeconv(), inputs, {self.conv: 1, self.d2s: 0}) + + def test_not_a_clean_upscale_is_left_alone(self): + # k=4, p=0: output is 2H + 2, not a x2 upscale. + module = self.DwDeconv(kernel=4, stride=2, padding=0) + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(module, inputs, {self.conv: 1, self.d2s: 0}) + + def test_grouped_but_not_depthwise_is_left_alone(self): + module = torch.nn.ConvTranspose2d(8, 8, 4, stride=2, padding=1, groups=4) + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(module, inputs, {self.conv: 1, self.d2s: 0}) + + def test_phase_math_matches_conv_transpose2d(self): + # The decomposition in plain torch, independent of the lowering. + import executorch.backends.xnnpack.custom_ops_lib as lib + + for k, s, p, o in [ + (4, 2, 1, 0), + (3, 2, 1, 1), + (2, 2, 0, 0), + (3, 3, 0, 0), + (5, 2, 2, 1), + ]: + c, h, w = 3, 5, 6 + x = torch.randn(1, c, h, w) + weight = torch.randn(c, 1, k, k) + ref = torch.nn.functional.conv_transpose2d( + x, weight, stride=s, padding=p, output_padding=o, groups=c + ) + phases = [] + for ay in range(s): + for ax in range(s): + taps = [] + for a in (ay, ax): + t_min = math.ceil(-(a + p) / s) + t_max = math.floor((k - 1 - a - p) / s) + taps.append((t_min, t_max)) + (ty_min, ty_max), (tx_min, tx_max) = taps + iy = torch.tensor( + [s * (ty_max - q) + ay + p for q in range(ty_max - ty_min + 1)] + ) + ix = torch.tensor( + [s * (tx_max - q) + ax + p for q in range(tx_max - tx_min + 1)] + ) + wp = weight[:, :, iy][:, :, :, ix] + xp = torch.nn.functional.pad(x, (tx_max, -tx_min, ty_max, -ty_min)) + phases.append(torch.nn.functional.conv2d(xp, wp, groups=c)) + out = lib.depth_to_space_2d_op(torch.cat(phases, 1), s) + self.assertEqual(out.shape, ref.shape, (k, s, p, o)) + self.assertTrue(torch.allclose(out, ref, atol=1e-5), (k, s, p, o)) From 743c690f953855f56bfb3ff4a56de7e33b948c9f Mon Sep 17 00:00:00 2001 From: borys3kk Date: Mon, 26 Oct 2026 13:21:24 +0100 Subject: [PATCH 4/9] Remove the deleted spec with a slice assignment --- backends/transforms/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 0ecfb99c8c0..fe36ff9ca69 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -269,8 +269,7 @@ def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Nod # Edit the spec list in place: ExportedProgramPassManager runs passes on a # shallow copy of the program, which would not see a replaced signature. input_specs = exp_program.graph_signature.input_specs - for spec in [spec for spec in input_specs if spec.arg.name == node.name]: - input_specs.remove(spec) + input_specs[:] = [spec for spec in input_specs if spec.arg.name != node.name] # Remove node from graph node.graph.erase_node(node) From 1c8a70d5a5d7a1d4950d7701abfbf22020436211 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Tue, 27 Oct 2026 08:48:53 +0100 Subject: [PATCH 5/9] Use a Software Mansion copyright header on the new files --- backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py | 2 +- backends/xnnpack/custom_ops_lib.py | 2 +- backends/xnnpack/operators/op_depth_to_space.py | 2 +- backends/xnnpack/operators/op_upsample_nearest2d.py | 2 +- backends/xnnpack/test/ops/test_upsample_nearest2d.py | 2 +- .../test/passes/test_decompose_depthwise_conv_transpose.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py index 875d7de5f11..94fac0a506c 100644 --- a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/custom_ops_lib.py b/backends/xnnpack/custom_ops_lib.py index 5c393b54238..290adfa6d84 100644 --- a/backends/xnnpack/custom_ops_lib.py +++ b/backends/xnnpack/custom_ops_lib.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/operators/op_depth_to_space.py b/backends/xnnpack/operators/op_depth_to_space.py index 1d913128515..0c7a2ae2baa 100644 --- a/backends/xnnpack/operators/op_depth_to_space.py +++ b/backends/xnnpack/operators/op_depth_to_space.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/operators/op_upsample_nearest2d.py b/backends/xnnpack/operators/op_upsample_nearest2d.py index 6e96c24afab..ec408a65e74 100644 --- a/backends/xnnpack/operators/op_upsample_nearest2d.py +++ b/backends/xnnpack/operators/op_upsample_nearest2d.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/test/ops/test_upsample_nearest2d.py b/backends/xnnpack/test/ops/test_upsample_nearest2d.py index a4a105e6b09..e70c4d4af72 100644 --- a/backends/xnnpack/test/ops/test_upsample_nearest2d.py +++ b/backends/xnnpack/test/ops/test_upsample_nearest2d.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py index b1ffc54eb04..3cc2c3d8aec 100644 --- a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the From 445b1d50f5e87319a823f017d41d4ad5f330372f Mon Sep 17 00:00:00 2001 From: borys3kk Date: Tue, 27 Oct 2026 09:07:08 +0100 Subject: [PATCH 6/9] Drop the channels-last broadcast operand changes that moved out of #24113 --- .../channels_last_tagged_reshape_pass.py | 87 ++----------- .../test_channels_last_tagged_reshape.py | 115 ------------------ 2 files changed, 7 insertions(+), 195 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 31e2e71214b..656d1214627 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -79,9 +79,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.linear.default, } - # 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 + # Broadcasting binary ops whose operands must share one memory format. 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 @@ -91,8 +89,6 @@ 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. @@ -180,63 +176,6 @@ 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_single_element(node: torch.fx.Node) -> bool: - val = node.meta.get("val", None) - return isinstance(val, torch.Tensor) and val.numel() == 1 - - @staticmethod - def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: - val = node.meta.get("val", None) - return ( - isinstance(val, torch.Tensor) - 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 @@ -631,40 +570,28 @@ 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 layout_inputs + for input_node in node.all_input_nodes ] - if is_or_isnt_nhwc_node and all(is_or_isnt_nhwc_node): + if 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 layout_inputs + for input_node in node.all_input_nodes ): - for input_node in layout_inputs: + for input_node in node.all_input_nodes: self.input_to_nhwc(graph_module, input_node, node) self.mark_as_nhwc_node(node) else: - for input_node in layout_inputs: + for input_node in node.all_input_nodes: self.input_to_nchw(graph_module, input_node, node) graph_module.recompile() @@ -688,7 +615,7 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 or node.target not in ChannelsLastTaggedReshapePass.broadcast_binary_ops ): continue - input_nodes = self.layout_inputs(node) + input_nodes = node.all_input_nodes 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 53bb958df06..bfe2ce862c9 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -290,121 +290,6 @@ 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() - ) - 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 c22288b72ef0c73a8bc76efc5cc35a2a41d0e220 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Thu, 29 Oct 2026 12:16:52 +0100 Subject: [PATCH 7/9] Key the phase weight placeholders by stride and padding as well create_constant_placeholder returns an existing placeholder of the same name, so two depthwise deconvs sharing one weight but using different stride or padding ended up with the first one's phase kernels. Include both in the name and add tests for a shared weight with different and with identical parameters. --- .../decompose_depthwise_conv_transpose.py | 9 ++- ...test_decompose_depthwise_conv_transpose.py | 57 +++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py index 94fac0a506c..986615932fe 100644 --- a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -166,13 +166,20 @@ def _decompose( ] ) phase_weight = weight[:, :, iy][:, :, :, ix].contiguous() + # The phase taps depend on s and padding too, so a weight + # shared by differently configured deconvs must not reuse + # the placeholders create_constant_placeholder finds by name. + phase_name = ( + f"{weight_node.name}_s{s}_p{padding[0]}x{padding[1]}" + f"_phase_{ay}_{ax}" + ) with graph.inserting_before(weight_node): phase_weight_node = create_constant_placeholder( exp_program=self.exported_program, graph=graph, kind=weight_spec.kind, persistent_buffer=weight_spec.persistent, - name=f"{weight_node.name}_phase_{ay}_{ax}", + name=phase_name, data=phase_weight, ) # (left, right, top, bottom) diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py index 3cc2c3d8aec..8ff95e4d22c 100644 --- a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -163,6 +163,63 @@ def test_two_deconvs_in_a_row(self): ) self._check_lowered(module, inputs) + class SharedWeightDeconvs(torch.nn.Module): + def __init__(self, kernel, stride, paddings, output_paddings): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(8, 1, kernel, kernel)) + self.stride = stride + self.configs = list(zip(paddings, output_paddings)) + + def forward(self, x): + return sum( + torch.nn.functional.conv_transpose2d( + x, + self.weight, + stride=self.stride, + padding=padding, + output_padding=output_padding, + groups=8, + ) + for padding, output_padding in self.configs + ) + + def _count_phase_placeholders(self, tester): + graph = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph + ) + return sum( + 1 + for node in graph.nodes + if node.op == "placeholder" and "_phase_" in node.name + ) + + def test_shared_weight_different_padding(self): + # Both deconvs are clean x3 upscales of the same weight, but their + # phase taps differ, so they must not share the phase placeholders. + module = self.SharedWeightDeconvs( + kernel=9, stride=3, paddings=(3, 4), output_paddings=(0, 2) + ) + inputs = (torch.randn(1, 8, 5, 6),) + tester = self._check_pass( + module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} + ) + self.assertEqual(self._count_phase_placeholders(tester), 18) + self._check_lowered(module, inputs) + + def test_shared_weight_same_params(self): + # Identical deconvs on one weight share one set of phase placeholders. + module = self.SharedWeightDeconvs( + kernel=9, stride=3, paddings=(3, 3), output_paddings=(0, 0) + ) + inputs = (torch.randn(1, 8, 5, 6),) + tester = self._check_pass( + module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} + ) + self.assertEqual(self._count_phase_placeholders(tester), 9) + self._check_lowered(module, inputs) + def test_batch_gt_1(self): module = self.DwDeconv(kernel=4, stride=2, padding=1) inputs = (torch.randn(3, 8, 6, 10),) From 850c9763dee4d7952ef5b81bb09457c1783a5f97 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Mon, 2 Nov 2026 12:49:19 +0100 Subject: [PATCH 8/9] Keep the external-constant tag on the phase weights create_constant_placeholder only sets meta["val"], so the phase weights lost the delegate_constant_tag of the weight they replace and landed in the .pte instead of the .ptd file. --- .../xnnpack/_passes/decompose_depthwise_conv_transpose.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py index 986615932fe..c41e0be5eeb 100644 --- a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -182,6 +182,11 @@ def _decompose( name=phase_name, data=phase_weight, ) + # create_constant_placeholder only sets meta["val"]; keep the + # external-constant tag so the phase weights follow the + # original weight into its .ptd file. + if "custom" in weight_node.meta: + phase_weight_node.meta["custom"] = dict(weight_node.meta["custom"]) # (left, right, top, bottom) pad = [tx_max, -tx_min, ty_max, -ty_min] with graph.inserting_before(node): From 4ffc1bad8f1d1ac1f2b942af15750bbb03ee577f Mon Sep 17 00:00:00 2001 From: borys3kk Date: Mon, 2 Nov 2026 12:49:20 +0100 Subject: [PATCH 9/9] Test that the pass runs during lowering, its order against FuseActivationPass, and external constants XNNPACK delegates an fp32 depthwise transposed conv on its own, so the lowered tests passed whether or not the decomposition happened. _check_lowered now counts the _decompose calls and checks the serialized graph for depth_to_space_2d and no transposed conv. test_deconv_relu runs the default pass list and checks that the relu reads the depth_to_space_2d output and that the phase pads fold. A new test lowers with external constants and expects every constant in model.ptd. --- ...test_decompose_depthwise_conv_transpose.py | 178 ++++++++++++++++-- 1 file changed, 166 insertions(+), 12 deletions(-) diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py index 8ff95e4d22c..5bba9cc4479 100644 --- a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -4,16 +4,39 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import json import math +import os +import tempfile import unittest +from collections import Counter +from unittest import mock +import executorch.backends.test.harness.stages as BaseStages import torch from executorch.backends.test.harness.stages import StageType +from executorch.backends.xnnpack._passes import XNNPACKPassManager from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( DecomposeDepthwiseConvTransposePass, ) +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.serialization.xnnpack_graph_serialize import ( + _get_schema_bytes, + XNNHeader, +) from executorch.backends.xnnpack.test.tester import RunPasses, Tester +from executorch.exir._serialize import _deserialize_pte_binary +from executorch.exir._serialize._flatbuffer import _flatc_decompile from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.external_constants_pass import ( + delegate_external_constants_pass_unlifted, +) +from executorch.extension.flat_tensor.serialize.serialize import ( + _deserialize_to_flat_tensor, +) class TestDecomposeDepthwiseConvTransposePass(unittest.TestCase): @@ -77,17 +100,68 @@ def _check_pass(self, module, inputs, counts, tol=1e-3, dynamic_shapes=None): .run_method_and_compare_outputs(atol=tol, rtol=tol) ) - def _check_lowered(self, module, inputs, tol=1e-3, dynamic_shapes=None): + @staticmethod + def _xnnpack_node_types(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: + # flatc writes unbounded clamp limits as a bare inf + text = ( + f.read().replace("-inf", "-Infinity").replace(" inf", " Infinity") + ) + graph = json.loads(text) + return Counter(node["xnode_union_type"] for node in graph["xnodes"]) + + def _check_lowered( + self, module, inputs, tol=1e-3, dynamic_shapes=None, num_deconvs=1 + ): + # XNNPACK delegates an fp32 depthwise transposed conv on its own, so + # one delegate and matching outputs do not prove the pass ran during + # lowering: count its calls and check the serialized graph. + tester = Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes).export() + with mock.patch.object( + DecomposeDepthwiseConvTransposePass, + "_decompose", + autospec=True, + side_effect=DecomposeDepthwiseConvTransposePass._decompose, + ) as decompose: + tester.to_edge_transform_and_lower() + self.assertEqual(decompose.call_count, num_deconvs) + tester.check_count( + {"torch.ops.higher_order.executorch_call_delegate": 1} + ).check_not([self.conv]) + node_types = self._xnnpack_node_types(tester) + self.assertEqual(node_types["XNNConvTranspose2d"], 0, node_types) + self.assertEqual(node_types["XNNDepthToSpace2d"], num_deconvs, node_types) ( - Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes) - .export() - .to_edge_transform_and_lower() - .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) - .check_not([self.conv]) - .to_executorch() + tester.to_executorch() .serialize() .run_method_and_compare_outputs(num_runs=3, atol=tol, rtol=tol) ) + return tester def test_upfirdn_4x4_stride2(self): # StyleGAN-style x2 upsampling filter: 4 phases of 2x2 taps, each with @@ -161,7 +235,7 @@ def test_two_deconvs_in_a_row(self): self._check_pass( module, inputs, {self.conv: 8, self.pad: 8, self.cat: 2, self.d2s: 2} ) - self._check_lowered(module, inputs) + self._check_lowered(module, inputs, num_deconvs=2) class SharedWeightDeconvs(torch.nn.Module): def __init__(self, kernel, stride, paddings, output_paddings): @@ -206,7 +280,7 @@ def test_shared_weight_different_padding(self): module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} ) self.assertEqual(self._count_phase_placeholders(tester), 18) - self._check_lowered(module, inputs) + self._check_lowered(module, inputs, num_deconvs=2) def test_shared_weight_same_params(self): # Identical deconvs on one weight share one set of phase placeholders. @@ -218,7 +292,7 @@ def test_shared_weight_same_params(self): module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} ) self.assertEqual(self._count_phase_placeholders(tester), 9) - self._check_lowered(module, inputs) + self._check_lowered(module, inputs, num_deconvs=2) def test_batch_gt_1(self): module = self.DwDeconv(kernel=4, stride=2, padding=1) @@ -245,13 +319,93 @@ def forward(self, x): def test_deconv_relu(self): # The activation follows depth_to_space_2d instead of fusing into the - # removed transposed conv. + # removed transposed conv, which relies on the pass running before + # FuseActivationPass in the default pass list. The phase pads read the + # NHWC output of the leading conv and fold into the phase convs through + # FusePadIntoConvPass at the end of that list. module = torch.nn.Sequential( - self.DwDeconv(kernel=4, stride=2, padding=1), torch.nn.ReLU() + torch.nn.Conv2d(8, 8, 3, padding=1), + self.DwDeconv(kernel=4, stride=2, padding=1), + torch.nn.ReLU(), ) inputs = (torch.randn(1, 8, 6, 10),) + default_passes = XNNPACKPassManager( + Tester(module.eval(), inputs).export().get_artifact() + ).passes + self.assertLess( + default_passes.index(DecomposeDepthwiseConvTransposePass), + default_passes.index(FuseActivationPass), + ) + tester = ( + Tester(module.eval(), inputs) + .export() + .to_edge() + .run_passes(RunPasses(default_passes)) + .check_count({self.conv: 5, self.pad: 4, self.d2s: 1}) + ) + graph = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph + ) + (relu,) = [ + n for n in graph.nodes if n.target == exir_ops.edge.aten.relu.default + ] + self.assertEqual( + relu.args[0].target, exir_ops.edge.et_xnnpack.depth_to_space_2d.default + ) + convs = [ + n for n in graph.nodes if n.target == exir_ops.edge.aten.convolution.default + ] + self.assertEqual( + [FuseActivationPass.get_fused_activation(n) for n in convs], [None] * 5 + ) + self.assertEqual( + sum(FusePadIntoConvPass.get_fused_padding(n) is not None for n in convs), 4 + ) + self.assertEqual( + sum( + FusePadIntoConvPass.is_folded_pad(n) + for n in graph.nodes + if n.target == exir_ops.edge.aten.constant_pad_nd.default + ), + 4, + ) self._check_lowered(module, inputs) + class ExportWithExternalConstants(BaseStages.Export): + def run(self, artifact, inputs): + tagged = torch.export.export(artifact, inputs, strict=True).module() + delegate_external_constants_pass_unlifted( + module=tagged, gen_tag_fn=lambda _: "model" + ) + self.exported_program = torch.export.export(tagged, inputs, strict=True) + + def test_external_constants(self): + # The phase weights inherit the original weight's external-constant + # tag, so every constant ends up in model.ptd and none in the .pte. + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(1, 8, 6, 10),) + tester = Tester(module.eval(), inputs) + tester.export(self.ExportWithExternalConstants()) + tester.to_edge_transform_and_lower().check_count( + {"torch.ops.higher_order.executorch_call_delegate": 1} + ).to_executorch() + exec = tester.get_artifact() + self.assertIsNone(_deserialize_pte_binary(exec.buffer).named_data) + self.assertEqual(list(exec._tensor_data), ["model"]) + ptd = _deserialize_to_flat_tensor(bytes(exec._tensor_data["model"])) + self.assertEqual(len(ptd.named_data), 5) # 4 phase weights + bias + + from executorch.extension.pybindings import portable_lib as runtime + + loaded = runtime._load_for_executorch_from_buffer( + exec.buffer, bytes(exec._tensor_data["model"]) + ) + self.assertTrue( + torch.allclose(loaded.forward(inputs)[0], module(*inputs), atol=1e-3) + ) + def test_dynamic_shape(self): module = self.DwDeconv(kernel=4, stride=2, padding=1) inputs = (torch.randn(1, 8, 6, 10),)