From e578110e97dbf2f090f5de96d5f83ca7d230e693 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 9 Nov 2026 16:52:33 -0500 Subject: [PATCH] Lower mean.dim beyond 4D pooling as an XNNPACK static mean ## Ulterior Motive Run RMSNorm in LeRobot policies on XNNPACK (#24667), where `mean.dim` fails to partition with "only supports averaging 4D tensors, got tensor of rank 3". ## Rationale **What**: Add an `XNNStaticMean` node backed by `xnn_define_static_reduce(xnn_reduce_mean)`, and lower every `mean.dim` that is not Global Average Pooling to it: ranks 1 to 6, any dims, and either `keepdim`. **Why**: XNNPACK only lowered a 4D mean over dims [2, 3] with `keepdim=True`. RMSNorm averages the last dim of a 3-D tensor, so its mean stayed on the portable kernel and split the delegate around every norm. ## Details Review order: the serialization schema and `XNNCompiler.cpp`, then `op_mean_dim.py` and `is_global_avg_pool_mean`, then `ChannelsLastTaggedReshapePass`, then `MeanDimConfig`, then the tests. - The 4D, dims [2, 3], `keepdim=True` case still lowers to Global Average Pooling on NHWC. For those programs the `.pte` is byte-identical to main (fp32, fp16, qs8, conv-bn-hardtanh-mean, and a CNN ending in `AdaptiveAvgPool2d(1)`), and the tests check which XNNPACK node each mean serializes to and how many transposes surround it. - Every other `mean.dim` reduces over NCHW dims, so `ChannelsLastTaggedReshapePass` gives it an NCHW input, and the visitor refuses an NHWC one as an internal layout error. After a convolution that adds layout copies inside the one delegate. For conv then `mean(1, keepdim=True)` then conv, the program has 4 transposes, as on main. For `y - y.mean(1, keepdim=True)` after a convolution (the centering in ConvNeXt's LayerNorm2d), the subtraction stays NHWC, so its full-size result is copied back to NCHW at the output: 4 transposes in one delegate, where main has 2 in two delegates around a portable mean. At 96 channels and 56x56 that is 0.75x at 1 thread, 0.85x at 4 and 1.12x at 16, and conv, mean, conv is 0.54x, 0.70x and 0.90x. Remapping the reduction axes to read NHWC directly would remove these copies, also for heads like `x.mean([2, 3])`; that is left for a follow-up. - `is_global_avg_pool_mean` takes the rank from the mean's own output, which with `keepdim=True` equals the input's, because inside `ChannelsLastTaggedReshapePass` the mean's input can be a newly inserted node without `"val"`. - The node uses a shared `_XNNStaticReduce` table, so other XNNPACK reductions can reuse it and `defineGenericReduceNode`. It is appended to `XNodeUnion`, so existing node IDs do not change. Practically every re-exported XNNPACK LLM program contains it, because RMSNorm (`examples/models/llama/norm.py`) and QK-norm lower to it, and those programs fail to load on runtimes built before this change, including released prebuilt runtimes, with `Unhandled node type` in the log. The release note should say so. - `ChannelsLastTaggedReshapePass` now reads the layout of a quantized graph input from the input's own tag. It used to test `meta["val"][0]`, a 3-D slice that is never channels_last, so a channels_last input behind a quantize and dequantize read as NCHW. A mean over NCHW dims of such an input then reduced the wrong axes, and on main a convolution of it got a transpose that fails at runtime. Both now lower correctly. - An empty dim list averages over every dim, as in PyTorch. - Latency against main, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: RMSNorm 1x128x2048 0.48x, 3-D last-dim mean 0.45x, RMSNorm then Linear 0.77x, 4 pre-norm residual MLP blocks 0.99x for decode (1 token) and 0.80x for prefill (64 tokens), conv then mean over channels 0.88x. - XNNPACK splits a reduction with more than one output into threadpool tasks, so a small mean with several outputs can cost more in XNNPACK than on the portable kernel when the pool has many threads. For RMSNorm on 10 shapes from 1x1x128 to 1x32x2048, every shape is 0.43x to 0.97x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.55x (1.5 to 2.7 us more), and the rest 0.44x to 0.93x. At 16 threads those small shapes are 3.7x to 5.7x (21 to 23 us more), and the rest 0.70x to 0.98x. Inside a model with neighbouring delegated ops the cost can remain: decode QK-norm between two linears, 8 heads of 64, goes from 11.8 us to 39.1 us at 16 threads, and 16 heads of 128 from 222 us to 241 us. A mean alone in its delegate pays the same wake-up (a bare 1x128x128 mean is 3.26x at 16 threads and 1.14x at 4), as does any op lowered alone, such as `add` (2.95x at 1x128x128 and 16 threads on main). These means are lowered like every other op the XNNPACK partitioner takes, rather than gated by size, which would be tuned to one thread count and treat static and dynamic exports differently. End to end, Qwen3-0.6B with 8da4w quantization (113 RMSNorm and QK-norm means, 368 delegates on main and 283 here) decodes, after over before in `executor_runner` with the LLM custom ops: on the same Apple M-series CPU 0.98x at 1 thread and 1.11x to 1.12x at the default 16 threads (97 to 87 tokens/s); on an iPhone 17 1.00x at 1 thread and 1.05x at the runner's default (101 to 96 tokens/s); on a Pixel 10 0.98x and 1.00x. The cost is the multi-head QK-norm means. Running small multi-output reductions without waking the pool belongs in XNNPACK and is left for a follow-up. - The LoRA CI scripts check the start of the 8da4w-quantized Qwen3 base model's greedy output with a regex. With RMSNorm and QK-norm now in XNNPACK, its first token can tip from `:\n` to `:` (`assistant: me` instead of `assistant:\n`): unquantized, main and this branch generate the same tokens with logits within 1e-4, but 8da4w re-rounds activations and the two candidates are within about 0.5 logits. The regex accepts exactly those two continuations, and the scripts print it when it does not match. - Tests cover 1-D to 6-D inputs, all dims, `keepdim=False`, fp16 and qs8 over the innermost and a middle dim, dynamic kept and reduced dims over 5 random shapes, RMSNorm, fp32 and qs8 convolutions followed by a mean, a mean over channels whose result meets the convolution output or feeds another convolution (as in ConvNeXt's LayerNorm2d), qs8 channels_last inputs to a mean and to a convolution, the serialized XNNPACK nodes (moved `xnnpack_node_types` from the layout pass tests to the XNNPACK `Tester`), and the layout copies the pass inserts. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code. --- .ci/scripts/test_lora.sh | 16 +- .ci/scripts/test_lora_multimethod.sh | 16 +- .../channels_last_tagged_reshape_pass.py | 38 +- backends/xnnpack/operators/op_mean_dim.py | 79 ++-- .../partition/config/generic_node_configs.py | 31 +- backends/xnnpack/runtime/XNNCompiler.cpp | 58 +++ .../xnnpack/serialization/runtime_schema.fbs | 10 + backends/xnnpack/serialization/schema.fbs | 10 + .../serialization/xnnpack_graph_schema.py | 14 + backends/xnnpack/test/ops/test_mean_dim.py | 369 +++++++++++++++--- .../test_channels_last_tagged_reshape.py | 114 ++++-- backends/xnnpack/test/tester/BUCK | 2 + backends/xnnpack/test/tester/tester.py | 53 ++- backends/xnnpack/utils/utils.py | 19 +- docs/source/backends/xnnpack/op-support.csv | 2 +- 15 files changed, 621 insertions(+), 210 deletions(-) diff --git a/.ci/scripts/test_lora.sh b/.ci/scripts/test_lora.sh index 2adb80e8bdf..535721b55fe 100644 --- a/.ci/scripts/test_lora.sh +++ b/.ci/scripts/test_lora.sh @@ -46,13 +46,18 @@ import re import sys text = pathlib.Path(sys.argv[1]).read_text() +# Under 8da4w, the base model's first token is a near-tie between ":\n" and ":" +# that tips with small numeric changes, such as RMSNorm and QK-norm moving into +# XNNPACK. Accept both observed continuations. pattern = re.compile( - r"^<\|im_start\|>user Calculate 15% of 80\?<\|im_end\|><\|im_start\|>assistant:\n" - r"(?:\n)+" + r"^<\|im_start\|>user Calculate 15% of 80\?<\|im_end\|><\|im_start\|>assistant:" + r"(?:\n(?:\n)+| me\n)" r"Okay, so I need to calculate 15% of 80\.", re.MULTILINE, ) -sys.exit(0 if pattern.match(text) else 1) +if not pattern.match(text): + print(f"Expected a match for: {pattern.pattern}") + sys.exit(1) PY } @@ -174,9 +179,6 @@ else fi ### QUANTIZATION & PROGRAM DATA SEPARATION ### -EXPECTED_QUANT_PREFIX="<|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant: - -Okay, so I need to calculate 15% of 80." EXPECTED_QUANT_LORA_PREFIX=" <|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant To calculate 15% of 80, we can multiply 80 by 15/100. @@ -238,11 +240,9 @@ NOW=$(date +"%H:%M:%S") echo "Finished at ${NOW}" RESULT=$(cat result.txt) if matches_base_response_prefix result.txt; then - echo "Expected result prefix: ${EXPECTED_QUANT_PREFIX}" echo "Actual result: ${RESULT}" echo "Test 3: Success" else - echo "Expected result prefix: ${EXPECTED_QUANT_PREFIX}" echo "Actual result: ${RESULT}" echo "Test 3: Failure; results not the same" cleanup_files diff --git a/.ci/scripts/test_lora_multimethod.sh b/.ci/scripts/test_lora_multimethod.sh index 811d847a255..68608fd3ae2 100755 --- a/.ci/scripts/test_lora_multimethod.sh +++ b/.ci/scripts/test_lora_multimethod.sh @@ -46,13 +46,18 @@ import re import sys text = pathlib.Path(sys.argv[1]).read_text() +# Under 8da4w, the base model's first token is a near-tie between ":\n" and ":" +# that tips with small numeric changes, such as RMSNorm and QK-norm moving into +# XNNPACK. Accept both observed continuations. pattern = re.compile( - r"^<\|im_start\|>user Calculate 15% of 80\?<\|im_end\|><\|im_start\|>assistant:\n" - r"(?:\n)+" + r"^<\|im_start\|>user Calculate 15% of 80\?<\|im_end\|><\|im_start\|>assistant:" + r"(?:\n(?:\n)+| me\n)" r"Okay, so I need to calculate 15% of 80\.", re.MULTILINE, ) -sys.exit(0 if pattern.match(text) else 1) +if not pattern.match(text): + print(f"Expected a match for: {pattern.pattern}") + sys.exit(1) PY } @@ -107,10 +112,6 @@ EXPECTED_LORA_PREFIX=" <|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant To calculate 15% of 80" -EXPECTED_BASE_PREFIX="<|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant: - -Okay, so I need to calculate 15% of 80." - ### TEST 1: Run lora_forward method ### NOW=$(date +"%H:%M:%S") echo "Test 1: Multimethod lora_forward. Starting at ${NOW}" @@ -151,7 +152,6 @@ if matches_base_response_prefix result_base.txt; then echo "Test 2 (base_forward): Success" else echo "Test 2 (base_forward): Failure" - echo "Expected result prefix: ${EXPECTED_BASE_PREFIX}" echo "Actual result: ${RESULT}" cleanup_files exit 1 diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index b7a62547c39..4a251a695ed 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -24,6 +24,7 @@ get_param_tensor, is_dtype_cast, is_foldable_broadcast_operand, + is_global_avg_pool_mean, is_param_node, ) from executorch.exir.dialects._ops import ops as exir_ops @@ -69,7 +70,6 @@ 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.aten.mean.dim, exir_ops.edge.aten.max_pool2d.default, exir_ops.edge.aten.amax.default, exir_ops.edge.aten.max.dim, @@ -137,32 +137,14 @@ def is_nhwc_node(node: torch.fx.Node) -> bool: quantize_node = node.args[0] if len(quantize_node.all_input_nodes) > 0: actual_node = quantize_node.args[0] - if actual_node.op == "placeholder": - return ChannelsLastTaggedReshapePass.is_nhwc_tensor( - actual_node.meta["val"][0] - ) - else: - return actual_node.meta.get( - ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False - ) + return actual_node.meta.get( + ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False + ) return node.meta.get(ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False) @staticmethod def is_nchw_node(node: torch.fx.Node) -> bool: - if is_dequant(node) and len(node.all_input_nodes) > 0: - quantize_node = node.args[0] - if len(quantize_node.all_input_nodes) > 0: - actual_node = quantize_node.args[0] - if actual_node.op == "placeholder": - return not ChannelsLastTaggedReshapePass.is_nhwc_tensor( - actual_node.meta["val"][0] - ) - else: - return not actual_node.meta.get( - ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False - ) - return not ChannelsLastTaggedReshapePass.is_nhwc_node(node) @staticmethod @@ -316,11 +298,17 @@ def layout_inputs(self, node: torch.fx.Node) -> List[torch.fx.Node]: ] def requires_nhwc_input(self, node: torch.fx.Node) -> bool: + if node.target == exir_ops.edge.aten.mean.dim: + return is_global_avg_pool_mean(node) return node.target in self.memory_sensitive_ops_nhwc def requires_nchw_inputs(self, node: torch.fx.Node) -> bool: if node.target == exir_ops.edge.aten.view_copy.default: return True + # A mean.dim that is not Global Average Pooling reduces over the dims + # as given, which index the NCHW layout. + if node.target == exir_ops.edge.aten.mean.dim: + return not is_global_avg_pool_mean(node) return node.target in self.memory_sensitive_ops_nchw @@ -546,12 +534,6 @@ def input_to_nhwc( # serializing graph, but don't do anything else here self.mark_as_nhwc_node(input_node) - if input_node.op == "placeholder": - if self._is_nhwc(input_node.meta["val"][0]): - return - elif ChannelsLastTaggedReshapePass.is_nhwc_node(input_node): - return - if self.input_dim_order(input_node, InputDimOrder.NHWC): return diff --git a/backends/xnnpack/operators/op_mean_dim.py b/backends/xnnpack/operators/op_mean_dim.py index fbf7995c650..f0feb232d10 100644 --- a/backends/xnnpack/operators/op_mean_dim.py +++ b/backends/xnnpack/operators/op_mean_dim.py @@ -5,9 +5,12 @@ # 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 +from typing import Dict import torch +from executorch.backends.xnnpack._passes.channels_last_tagged_reshape_pass import ( + ChannelsLastTaggedReshapePass, +) from executorch.backends.xnnpack.operators.node_visitor import ( check_or_raise, get_tensor_value, @@ -17,19 +20,24 @@ from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( XNNGlobalAvgPooling2d, XNNGraph, + XNNStaticMean, XNode, ) -from executorch.backends.xnnpack.utils.utils import normalize_mean_dims +from executorch.backends.xnnpack.utils.utils import ( + get_input_node, + is_global_avg_pool_mean, + normalize_mean_dims, +) from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_FLAG_KEEP_DIMS @register_node_visitor class MeanDim(NodeVisitor): """ - XNNPACK only supports the special case of mean.dim that can be lowered - to Global Average Pooling. The input tensor must be 4D, keepdim must be - True, and the reduced dimensions must normalize to [2, 3] (for example - [2, 3], [3, 2], [-1, -2], or [-2, -1]). + mean.dim over the innermost dimensions of a 4D tensor with keepdim=True + (dims normalizing to [2, 3]) is lowered to Global Average Pooling on an + NHWC input. Any other mean.dim is lowered to a static mean reduction on + the NCHW input. """ target = "aten.mean.dim" @@ -44,44 +52,39 @@ def define_node( vals_to_ids: Dict[torch.fx.Node, int], debug_handle: int, ) -> None: - self.define_nodes_tensor_inputs_outputs( - node, xnn_graph, vals_to_ids, convert_to_nhwc=True - ) - # input - input_id = vals_to_ids[cast(torch.fx.Node, node.args[0])] - - # output - output_id = vals_to_ids[node] - - input_shape = get_tensor_value(xnn_graph.xvalues[input_id]).dims - check_or_raise( - len(input_shape) == 4, "Require input to mean.dim be 4 dimensional" - ) - - # This visitor serializes mean.dim as Global Average Pooling, which has - # no field for an explicit dtype override. check_or_raise( node.kwargs.get("dtype") is None, "XNNPACK does not support mean.dim with dtype", ) - # mean dims - mean_dims = normalize_mean_dims(node.args[1], len(input_shape)) - check_or_raise( - sorted(mean_dims) == [2, 3], - "XNNPACK only supports mean.dim across the innermost dimensions", - ) - - # keep dims - check_or_raise( - len(node.args) == 3 and bool(node.args[2]), - "XNNPACK only supports mean.dim that keeps dims", + global_avg_pool = is_global_avg_pool_mean(node) + self.define_nodes_tensor_inputs_outputs( + node, xnn_graph, vals_to_ids, convert_to_nhwc=global_avg_pool ) + input_node = get_input_node(node, 0) + input_id = vals_to_ids[input_node] + output_id = vals_to_ids[node] - ser_node = XNode( - xnode_union=XNNGlobalAvgPooling2d( + if global_avg_pool: + xnode_union = XNNGlobalAvgPooling2d( input_id=input_id, output_id=output_id, flags=XNN_FLAG_KEEP_DIMS - ), - debug_handle=debug_handle, - ) + ) + else: + check_or_raise( + not input_node.meta.get( + ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False + ), + "Internal layout error: ChannelsLastTaggedReshapePass gave an NHWC " + "input to a mean.dim that reduces over NCHW dims", + ) + input_rank = len(get_tensor_value(xnn_graph.xvalues[input_id]).dims) + keepdim = len(node.args) >= 3 and bool(node.args[2]) + xnode_union = XNNStaticMean( + reduction_axes=sorted(normalize_mean_dims(node.args[1], input_rank)), + input_id=input_id, + output_id=output_id, + flags=XNN_FLAG_KEEP_DIMS if keepdim else 0, + ) + + ser_node = XNode(xnode_union=xnode_union, debug_handle=debug_handle) xnn_graph.xnodes.append(ser_node) diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index 1358d38b8fc..d5de7e25e79 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -714,44 +714,25 @@ class MeanDimConfig(GenericNodePartitionerConfig): def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: """ - Mean Dim currently only supports averaging 4D tensors across the innermost - dimensions + Mean Dim lowers to Global Average Pooling when it averages a 4D tensor + across the innermost dimensions with keepdim=True, and to a static mean + reduction otherwise. """ if not self.check_common_constraints(node, ep): return False - input_rank = get_input_node(node, 0).meta["val"].dim() - if input_rank != 4: - why( - node, - reason=f"mean.dim only supports averaging 4D tensors, got tensor of rank {input_rank}", - ) - return False - - # This path lowers mean.dim to XNNPACK Global Average Pooling, which - # cannot encode an explicit dtype override. + # Neither lowering can encode an explicit dtype override. if node.kwargs.get("dtype") is not None: why(node, reason="mean.dim does not support dtype") return False - keepdim = len(node.args) >= 3 and bool(node.args[2]) + input_node = get_input_node(node, 0) try: - dims = normalize_mean_dims(node.args[1], input_rank) + normalize_mean_dims(node.args[1], input_node.meta["val"].dim()) except ValueError as error: why(node, reason=f"mean.dim has invalid dims: {error}") return False - if sorted(dims) != [2, 3]: - why( - node, - reason="mean.dim only supports averaging 4D tensors across the innermost dimensions", - ) - return False - - if not keepdim: - why(node, reason="mean.dim only supports keepdim=True") - return False - return True def supported_precision_types(self) -> List[ConfigPrecisionType]: diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index b9797949a01..429a2527a23 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -2037,6 +2037,60 @@ Error defineGenericBinaryNode( node->debug_handle()); \ } +// Generic helper function for reduce operations +Error defineGenericReduceNode( + xnn_subgraph_t subgraph_ptr, + const std::unordered_map& remapped_ids, + const fb_xnnpack::_XNNStaticReduce* graph_node, + xnn_reduce_operator op_type, + fb_xnnpack::XNodeUnion node_type, + uint32_t debug_handle) noexcept { + ET_CHECK_OR_RETURN_ERROR( + graph_node->reduction_axes() != nullptr, + InvalidProgram, + "%s: reduction_axes is null", + fb_xnnpack::EnumNameXNodeUnion(node_type)); + std::vector reduction_axes = + flatbufferDimsToVector(graph_node->reduction_axes()); + + REMAP_ID(remapped_ids, graph_node->input_id(), reduce_input); + REMAP_ID(remapped_ids, graph_node->output_id(), reduce_output); + + xnn_status status = xnn_define_static_reduce_v2( + subgraph_ptr, + op_type, + reduction_axes.size(), + reduction_axes.data(), + reduce_input, + reduce_output, + graph_node->flags()); + ET_CHECK_OR_RETURN_ERROR( + status == xnn_status_success, + Internal, + "Failed to create %s node %i with code: %s", + fb_xnnpack::EnumNameXNodeUnion(node_type), + debug_handle, + xnn_status_to_string(status)); + + return Error::Ok; +} + +#define _DEFINE_REDUCE_NODE(name, op_type) \ + Error define##name##Node( \ + xnn_subgraph_t subgraph_ptr, \ + const std::unordered_map& remapped_ids, \ + const NodePtr node, \ + const fb_xnnpack::XNNGraph* graph) noexcept { \ + MAYBE_UNUSED(graph); \ + return defineGenericReduceNode( \ + subgraph_ptr, \ + remapped_ids, \ + node->xnode_union_as_XNN##name(), \ + op_type, \ + node->xnode_union_type(), \ + node->debug_handle()); \ + } + /* Returns the pointer to the defineNode function that handles the given XNode type @@ -2080,6 +2134,9 @@ _DEFINE_BINARY_NODE_WITH_MINMAX(Div, xnn_binary_divide) _DEFINE_BINARY_NODE_NO_PARAMS(Minimum, xnn_binary_minimum) _DEFINE_BINARY_NODE_NO_PARAMS(Maximum, xnn_binary_maximum) +// Reduce Ops +_DEFINE_REDUCE_NODE(StaticMean, xnn_reduce_mean) + DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) { switch (nodeType) { // Binary ops @@ -2125,6 +2182,7 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) { _DEFINE(MaxPooling2d) _DEFINE(Convert) _DEFINE(GlobalAvgPooling2d) + _DEFINE(StaticMean) _DEFINE(StaticReshape) _DEFINE(DepthToSpace2d) _DEFINE(ArgMaxPooling2d) diff --git a/backends/xnnpack/serialization/runtime_schema.fbs b/backends/xnnpack/serialization/runtime_schema.fbs index 62143ab8c62..c1db1a8b28d 100644 --- a/backends/xnnpack/serialization/runtime_schema.fbs +++ b/backends/xnnpack/serialization/runtime_schema.fbs @@ -161,6 +161,7 @@ union XNodeUnion { XNNCos: _XNNNode1x1, XNNApproxGelu: _XNNNode1x1, XNNDepthToSpace2d, + XNNStaticMean: _XNNStaticReduce, } union XValueUnion { @@ -306,6 +307,15 @@ table XNNDepthToSpace2d { flags: uint; } +// A reduction over static axes +// Not meant to be used directly +table _XNNStaticReduce { + reduction_axes: [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 a48f9bc4187..e42c10050fd 100644 --- a/backends/xnnpack/serialization/schema.fbs +++ b/backends/xnnpack/serialization/schema.fbs @@ -157,6 +157,7 @@ union XNodeUnion { XNNCos: _XNNNode1x1, XNNApproxGelu: _XNNNode1x1, XNNDepthToSpace2d, + XNNStaticMean: _XNNStaticReduce, } union XValueUnion { @@ -302,6 +303,15 @@ table XNNDepthToSpace2d { flags: uint; } +// A reduction over static axes +// Not meant to be used directly +table _XNNStaticReduce { + reduction_axes: [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 e6da4338dd0..e8bf2b64653 100644 --- a/backends/xnnpack/serialization/xnnpack_graph_schema.py +++ b/backends/xnnpack/serialization/xnnpack_graph_schema.py @@ -210,6 +210,19 @@ class XNNDepthToSpace2d: flags: int +@dataclass +class XNNStaticReduce: + reduction_axes: List[int] + input_id: int + output_id: int + flags: int + + +@dataclass +class XNNStaticMean(XNNStaticReduce): + pass + + @dataclass class XNNStaticSlice: num_dims: int @@ -436,6 +449,7 @@ class XNNScaledDotProductAttention: XNNCos, XNNApproxGelu, XNNDepthToSpace2d, + XNNStaticMean, ] diff --git a/backends/xnnpack/test/ops/test_mean_dim.py b/backends/xnnpack/test/ops/test_mean_dim.py index 731ac378649..5bf4e54c265 100644 --- a/backends/xnnpack/test/ops/test_mean_dim.py +++ b/backends/xnnpack/test/ops/test_mean_dim.py @@ -30,72 +30,326 @@ def forward(self, x): z = torch.mean(y, self.dims, keepdim=self.keepdim, dtype=self.dtype) return z - def _test_mean_dim(self, inputs, dims=(-1, -2)): + class RMSNorm(torch.nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.rand(dim)) + + def forward(self, x): + rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return x * rms * self.weight + + class ConvMeanDim(torch.nn.Module): + def __init__(self, dims, keepdim=True): + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3) + self.dims = dims + self.keepdim = keepdim + + def forward(self, x): + return torch.mean(self.conv(x), self.dims, keepdim=self.keepdim) + + class ConvSubtractMean(torch.nn.Module): + """The centering in ConvNeXt's LayerNorm2d, after a convolution.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3) + + def forward(self, x): + y = self.conv(x) + return y - y.mean(1, keepdim=True) + + class ConvMeanConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 16, 3) + self.conv2 = torch.nn.Conv2d(1, 8, 3) + + def forward(self, x): + return self.conv2(self.conv1(x).mean(1, keepdim=True)) + + @staticmethod + def _xnnpack_mean_nodes(global_avg_pool): + """ + For a mean of a contiguous graph input, Global Average Pooling is the + only lowering of a 4D mean over [2, 3] with keepdim, with its input + transposed to NHWC and its output back, as on main; every other mean is + one XNNStaticMean with no transpose. + """ + if global_avg_pool: + return { + "XNNGlobalAvgPooling2d": 1, + "XNNStaticMean": 0, + "XNNStaticTranspose": 2, + } + return { + "XNNGlobalAvgPooling2d": 0, + "XNNStaticMean": 1, + "XNNStaticTranspose": 0, + } + + def _test_mean_dim( + self, + inputs, + dims=(-1, -2), + keepdim=True, + dynamic_shapes=None, + global_avg_pool=False, + num_runs=1, + ): ( - Tester(self.MeanDim(dims), inputs) + Tester( + self.MeanDim(dims, keepdim=keepdim), + inputs, + dynamic_shapes=dynamic_shapes, + ) .export() .check_count({"torch.ops.aten.mean.dim": 1}) .to_edge_transform_and_lower() .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) .check_not(["executorch_exir_dialects_edge__ops_aten_mean_dim"]) + .check_xnnpack_node_count(self._xnnpack_mean_nodes(global_avg_pool)) .to_executorch() .serialize() - .run_method_and_compare_outputs() + .run_method_and_compare_outputs(num_runs=num_runs) + ) + + def _test_qs8_mean_dim( + self, inputs, dims, global_avg_pool=False, xnnpack_nodes=None + ): + """The input quantize and the output dequantize are the two converts.""" + if xnnpack_nodes is None: + xnnpack_nodes = self._xnnpack_mean_nodes(global_avg_pool) + ( + Tester(self.MeanDim(dims), inputs) + .quantize() + .export() + .check_node_count( + { + torch.ops.aten.mean.dim: 1, + torch.ops.quantized_decomposed.quantize_per_tensor.default: 3, + } + ) + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten_mean_dim", + "torch.ops.quantized_decomposed", + ] + ) + .check_xnnpack_node_count({**xnnpack_nodes, "XNNConvert": 2}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(qtol=1) ) def test_fp16_mean_dim(self): inputs = (torch.randn(1, 5, 4, 4).to(torch.float16),) - self._test_mean_dim(inputs) + self._test_mean_dim(inputs, global_avg_pool=True) def test_fp32_mean_dim(self): inputs = (torch.randn(1, 5, 4, 4),) - self._test_mean_dim(inputs) + self._test_mean_dim(inputs, global_avg_pool=True) def test_fp32_mean_dim_positive_dims(self): inputs = (torch.randn(1, 5, 4, 4),) - self._test_mean_dim(inputs, dims=(2, 3)) + self._test_mean_dim(inputs, dims=(2, 3), global_avg_pool=True) - def test_fp32_mean_dim_unsupported(self): - """ - XNNPack mean.dim implementation only supports innermost two dimensions. As such, - we expect it to fail to partition when dim=(3). - """ - inputs = (torch.randn(1, 5, 4, 4),) + def test_fp32_mean_dim_single_innermost(self): + inputs = (torch.randn(1, 16, 32, 32),) + self._test_mean_dim(inputs, dims=(3,)) + + def test_fp32_mean_dim_channels(self): + inputs = (torch.randn(1, 16, 32, 32),) + self._test_mean_dim(inputs, dims=(1,)) + + def test_fp32_mean_dim_keepdim_false(self): + inputs = (torch.randn(1, 16, 32, 32),) + self._test_mean_dim(inputs, keepdim=False) + + def test_fp32_mean_dim_last_dim_3d(self): + inputs = (torch.randn(2, 64, 128),) + self._test_mean_dim(inputs, dims=(-1,)) + + def test_fp16_mean_dim_last_dim_3d(self): + inputs = (torch.randn(2, 64, 128).to(torch.float16),) + self._test_mean_dim(inputs, dims=(-1,)) + + def test_fp16_mean_dim_middle_dim_3d(self): + inputs = (torch.randn(2, 64, 128).to(torch.float16),) + self._test_mean_dim(inputs, dims=(1,)) + + def test_fp32_mean_dim_3d_keepdim_false(self): + inputs = (torch.randn(2, 64, 128),) + self._test_mean_dim(inputs, dims=(-1, -2), keepdim=False) + + def test_fp32_mean_dim_1d(self): + inputs = (torch.randn(64),) + for keepdim in (True, False): + with self.subTest(keepdim=keepdim): + self._test_mean_dim(inputs, dims=(0,), keepdim=keepdim) + + def test_fp32_mean_dim_2d(self): + inputs = (torch.randn(128, 128),) + self._test_mean_dim(inputs, dims=(0,)) + + def test_fp32_mean_dim_6d(self): + inputs = (torch.randn(2, 2, 4, 4, 8, 16),) + self._test_mean_dim(inputs, dims=(1, 4)) + + def test_fp32_mean_dim_5d(self): + inputs = (torch.randn(2, 4, 8, 16, 16),) + self._test_mean_dim(inputs, dims=(1, 3)) + + def test_fp32_mean_dim_all_dims(self): + inputs = (torch.randn(2, 7, 16),) + for dims in (None, []): + for keepdim in (True, False): + with self.subTest(dims=dims, keepdim=keepdim): + self._test_mean_dim(inputs, dims=dims, keepdim=keepdim) + + def test_fp32_mean_dim_small_single_output(self): + inputs = (torch.randn(1, 1, 64),) + self._test_mean_dim(inputs, dims=(-1,)) + + def test_fp32_mean_dim_small_multi_output(self): + inputs = (torch.randn(2, 7, 16),) + self._test_mean_dim(inputs, dims=(-1,)) + + def test_fp32_mean_dim_last_dim_3d_dynamic(self): + inputs = (torch.randn(2, 7, 16),) + seq = torch.export.Dim("seq", min=1, max=32) + self._test_mean_dim(inputs, dims=(-1,), dynamic_shapes=({1: seq},), num_runs=5) + + def test_fp32_mean_dim_dynamic_reduced_dim(self): + inputs = (torch.randn(2, 7, 16),) + seq = torch.export.Dim("seq", min=1, max=32) + self._test_mean_dim(inputs, dims=(1,), dynamic_shapes=({1: seq},), num_runs=5) + + def test_fp32_rms_norm(self): + inputs = (torch.randn(2, 64, 128),) ( - Tester(self.MeanDim((3)), inputs) + Tester(self.RMSNorm(128), inputs) .export() .check_count({"torch.ops.aten.mean.dim": 1}) .to_edge_transform_and_lower() - .check_count({"executorch_exir_dialects_edge__ops_aten_mean_dim": 1}) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten_mean_dim", + "executorch_exir_dialects_edge__ops_aten_pow_Tensor_Scalar", + "executorch_exir_dialects_edge__ops_aten_rsqrt_default", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + ] + ) + .check_xnnpack_node_count({"XNNStaticMean": 1, "XNNGlobalAvgPooling2d": 0}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() ) - def test_fp32_mean_dim_unsupported_3d(self): + def test_fp32_conv_mean_dim(self): """ - XNNPack mean.dim implementation only supports 4D tensors. + The convolution output is NHWC inside the delegate. Global Average + Pooling reads it as is, and any other mean reads it back in NCHW. """ - inputs = (torch.randn(1, 5, 4),) - ( - Tester(self.MeanDim((-1, -2)), inputs) - .export() - .check_count({"torch.ops.aten.mean.dim": 1}) - .to_edge_transform_and_lower() - .check_count({"executorch_exir_dialects_edge__ops_aten_mean_dim": 1}) + inputs = (torch.randn(1, 3, 34, 34),) + cases = ( + ( + (-1, -2), + True, + { + "XNNGlobalAvgPooling2d": 1, + "XNNStaticMean": 0, + "XNNStaticTranspose": 2, + }, + ), + ((1,), True, {"XNNStaticMean": 1, "XNNStaticTranspose": 2}), + ((-1, -2), False, {"XNNStaticMean": 1, "XNNStaticTranspose": 2}), ) + for dims, keepdim, xnnpack_nodes in cases: + with self.subTest(dims=dims, keepdim=keepdim): + ( + Tester(self.ConvMeanDim(dims, keepdim), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not(["executorch_exir_dialects_edge__ops_aten_mean_dim"]) + .check_xnnpack_node_count(xnnpack_nodes) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) - def test_fp32_mean_dim_unsupported_keepdim_false(self): - inputs = (torch.randn(1, 5, 4, 4),) - ( - Tester(self.MeanDim((-1, -2), keepdim=False), inputs) - .export() - .check_count({"torch.ops.aten.mean.dim": 1}) - .to_edge_transform_and_lower() - .check_count({"executorch_exir_dialects_edge__ops_aten_mean_dim": 1}) - ) + def test_qs8_conv_mean_dim(self): + inputs = (torch.randn(1, 3, 34, 34),) + for dims, keepdim in (((1,), True), ((-1, -2), False)): + with self.subTest(dims=dims, keepdim=keepdim): + ( + Tester(self.ConvMeanDim(dims, keepdim).eval(), inputs) + .quantize() + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten_mean_dim", + "torch.ops.quantized_decomposed", + ] + ) + .check_xnnpack_node_count( + {"XNNStaticMean": 1, "XNNStaticTranspose": 2, "XNNConvert": 2} + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(qtol=1) + ) - def test_fp32_mean_dim_unsupported_dtype(self): - inputs = (torch.randn(1, 5, 4, 4),) + def test_conv_mean_meets_nhwc(self): + """ + A mean over channels after a convolution reads the NHWC convolution + output back in NCHW, and its result goes back to NHWC to meet the + convolution output or feed another convolution. With the input, the + mean's input and output, and the delegate output, that is 4 transposes; + main has 2 for ConvSubtractMean (3 in qs8), split over two delegates + around a portable mean, and 4 for ConvMeanConv. In qs8, the mean can + round one step (about 0.016 here) away from the reference, and that + difference reaches the output. + """ + inputs = (torch.randn(1, 3, 18, 18),) + for module in (self.ConvSubtractMean(), self.ConvMeanConv()): + for quantize in (False, True): + with self.subTest(type(module).__name__, quantize=quantize): + tester = Tester(module.eval(), inputs) + if quantize: + tester.quantize() + ( + tester.export() + .to_edge_transform_and_lower() + .check_count( + {"torch.ops.higher_order.executorch_call_delegate": 1} + ) + .check_not(["executorch_exir_dialects_edge__ops_aten_mean_dim"]) + .check_xnnpack_node_count( + {"XNNStaticMean": 1, "XNNStaticTranspose": 4} + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(atol=2e-2 if quantize else 1e-3) + ) + + def test_fp16_mean_dim_dtype_argument(self): + """ + An fp16 mean with an fp32 dtype passes the generic dtype checks, so + MeanDimConfig's own dtype check keeps it on the portable kernel. + """ + inputs = (torch.randn(2, 7, 16).to(torch.float16),) ( - Tester(self.MeanDim((-1, -2), dtype=torch.float64), inputs) + Tester(self.MeanDim((-1,), dtype=torch.float32), inputs) .export() .check_count({"torch.ops.aten.mean.dim": 1}) .to_edge_transform_and_lower() @@ -104,25 +358,26 @@ def test_fp32_mean_dim_unsupported_dtype(self): def test_qs8_mean_dim(self): inputs = (torch.randn(1, 5, 4, 4),) - ( - Tester(self.MeanDim((-1, -2)), inputs) - .quantize() - .export() - .check_node_count( - { - torch.ops.aten.mean.dim: 1, - torch.ops.quantized_decomposed.quantize_per_tensor.default: 3, - } - ) - .to_edge_transform_and_lower() - .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) - .check_not( - [ - "executorch_exir_dialects_edge__ops_aten_mean_dim", - "torch.ops.quantized_decomposed", - ] - ) - .to_executorch() - .serialize() - .run_method_and_compare_outputs(qtol=1) - ) + self._test_qs8_mean_dim(inputs, dims=(-1, -2), global_avg_pool=True) + + def test_qs8_mean_dim_last_dim_3d(self): + inputs = (torch.randn(2, 64, 128),) + self._test_qs8_mean_dim(inputs, dims=(-1,)) + + def test_qs8_mean_dim_channels_last_input(self): + """ + The quantized channels_last input is NHWC, so a mean over NCHW dims + reads it through a transpose, and Global Average Pooling reads it as is + and only transposes its output back. + """ + inputs = (torch.randn(1, 5, 4, 4).to(memory_format=torch.channels_last),) + for dims, xnnpack_nodes in ( + ((1,), {"XNNStaticMean": 1, "XNNStaticTranspose": 1}), + ((-1, -2), {"XNNGlobalAvgPooling2d": 1, "XNNStaticTranspose": 1}), + ): + with self.subTest(dims=dims): + self._test_qs8_mean_dim(inputs, dims=dims, xnnpack_nodes=xnnpack_nodes) + + def test_qs8_mean_dim_middle_dim_3d(self): + inputs = (torch.randn(2, 64, 128),) + self._test_qs8_mean_dim(inputs, dims=(1,)) 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 c91c9d56ae3..c49b800e78d 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -6,11 +6,7 @@ # 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 @@ -20,10 +16,6 @@ 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, ) @@ -34,7 +26,6 @@ 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 @@ -275,6 +266,32 @@ def test_qs8_channels_last_tagged_reshape_pass_conv_relu(self): .run_method_and_compare_outputs() ) + class AddConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(5, 8, 3) + + def forward(self, x): + return self.conv(x + x) + + def test_qs8_channels_last_input_add_conv(self): + """ + A quantized channels_last input is already NHWC, so the add and the + convolution read it without a transpose. + """ + x = torch.randn(1, 5, 8, 8).to(memory_format=torch.channels_last) + ( + Tester(self.AddConv().eval(), (x,)) + .quantize() + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_xnnpack_node_count({"XNNConv2d": 1, "XNNStaticTranspose": 0}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(qtol=1) + ) + class Conv2dBnHardtanhMeanSequenceModule(torch.nn.Module): def __init__(self): super().__init__() @@ -490,36 +507,6 @@ 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 _delegate_input_buffers(self, tester) -> set: """Names of the buffers the single delegate takes as inputs.""" exported_program = tester.get_artifact( @@ -540,7 +527,7 @@ def _delegate_input_buffers(self, tester) -> set: 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) + node_types = tester.xnnpack_node_types() self.assertEqual(node_types["XNNStaticTranspose"], 2, node_types) self.assertEqual(node_types["XNNStaticReshape"], 0, node_types) @@ -712,7 +699,7 @@ def test_fp32_channels_last_tagged_reshape_pass_4d_constant_operand(self): .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) + node_types = tester.xnnpack_node_types() 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() @@ -820,7 +807,7 @@ def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): .to_edge_transform_and_lower() .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) ) - node_types = self._xnnpack_node_types(tester) + node_types = tester.xnnpack_node_types() 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( @@ -919,6 +906,49 @@ def test_fp32_channels_last_tagged_reshape_pass_conv_bn_hardtanh_mean_seq(self): .run_method_and_compare_outputs() ) + class ConvMeanChannels(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3) + + def forward(self, x): + return torch.mean(self.conv(x), 1, keepdim=True) + + def test_fp32_channels_last_tagged_reshape_pass_conv_mean_channels(self): + # Copy #1 is for input to conv, nchw -> nhwc + # Copy #2 is for conv to mean, nhwc -> nchw + tester = ( + Tester(self.ConvMeanChannels().eval(), (torch.randn(1, 3, 8, 8),)) + .export() + .to_edge() + .run_passes(self.PassStage) + .check_count({self.to_copy_name: 2}) + ) + graph = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph + ) + (mean,) = [n for n in graph.nodes if n.target == exir_ops.edge.aten.mean.dim] + copy = mean.args[0] + self.assertEqual(copy.target, exir_ops.edge.aten._to_copy.default) + self.assertEqual(copy.kwargs["memory_format"], torch.contiguous_format) + tester.run_method_and_compare_outputs() + + class MeanLastDim(torch.nn.Module): + def forward(self, x): + return torch.mean(x, -1, keepdim=True) + + def test_fp32_channels_last_tagged_reshape_pass_mean_3d(self): + ( + Tester(self.MeanLastDim().eval(), (torch.randn(2, 7, 16),)) + .export() + .to_edge() + .run_passes(self.PassStage) + .check_not([self.to_copy_name]) + .run_method_and_compare_outputs() + ) + class Conv2dDynamicQuant(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/xnnpack/test/tester/BUCK b/backends/xnnpack/test/tester/BUCK index c9b0239de02..79c402c88af 100644 --- a/backends/xnnpack/test/tester/BUCK +++ b/backends/xnnpack/test/tester/BUCK @@ -17,6 +17,8 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/xnnpack/partition:xnnpack_partitioner", "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer", "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer_utils", + "//executorch/backends/xnnpack/serialization:xnnpack_schema", + "//executorch/backends/xnnpack/serialization:xnnpack_serializer", "//executorch/backends/xnnpack/utils:xnnpack_utils", "//executorch/devtools/visualization:visualization", "//executorch/exir:lib", diff --git a/backends/xnnpack/test/tester/tester.py b/backends/xnnpack/test/tester/tester.py index fc12da231c0..88b111f0fdf 100644 --- a/backends/xnnpack/test/tester/tester.py +++ b/backends/xnnpack/test/tester/tester.py @@ -5,7 +5,11 @@ # 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 Any, Callable, List, Optional, Sequence, Tuple, Type +import json +import os +import tempfile +from collections import Counter +from typing import Any, Callable, Dict, get_args, List, Optional, Sequence, Tuple, Type import executorch import executorch.backends.test.harness.stages as BaseStages @@ -23,8 +27,14 @@ from executorch.backends.xnnpack.quantizer.xnnpack_quantizer_utils import ( QuantizationConfig, ) +from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import XNodeUnion +from executorch.backends.xnnpack.serialization.xnnpack_graph_serialize import ( + _get_schema_bytes, + XNNHeader, +) from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config from executorch.exir import EdgeCompileConfig +from executorch.exir._serialize._flatbuffer import _flatc_decompile from executorch.exir.backend.partitioner import Partitioner from torch._export.pass_base import PassType from torchao.quantization.pt2e.quantizer import Quantizer @@ -132,3 +142,44 @@ def __init__( dynamic_shapes=dynamic_shapes, **kwargs, ) + + def xnnpack_node_types(self) -> Counter: + """Node types in the serialized XNNPACK graph of the single delegate + produced by to_edge_transform_and_lower.""" + graph_module = ( + self.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 check_xnnpack_node_count(self, expected: Dict[str, int]): + node_type_names = {t.__name__ for t in get_args(XNodeUnion)} + unknown = set(expected) - node_type_names + assert not unknown, f"Unknown XNNPACK node types: {sorted(unknown)}" + node_types = self.xnnpack_node_types() + for node_type, count in expected.items(): + assert ( + node_types[node_type] == count + ), f"Expected {count} {node_type} in the XNNPACK graph, got {node_types}" + return self diff --git a/backends/xnnpack/utils/utils.py b/backends/xnnpack/utils/utils.py index 223cb18bf7c..970db3712ea 100644 --- a/backends/xnnpack/utils/utils.py +++ b/backends/xnnpack/utils/utils.py @@ -75,10 +75,11 @@ def normalize_mean_dims(mean_dims: Sequence[int] | int | None, rank: int) -> Lis """Return mean dims as non-negative indices for the given rank.""" if rank <= 0: raise ValueError(f"Expected rank > 0, got {rank}") - if mean_dims is None: - return list(range(rank)) if isinstance(mean_dims, int): mean_dims = [mean_dims] + # None and an empty list both average over every dim. + if not mean_dims: + return list(range(rank)) normalized_dims = [] for dim in mean_dims: if dim < -rank or dim >= rank: @@ -87,6 +88,20 @@ def normalize_mean_dims(mean_dims: Sequence[int] | int | None, rank: int) -> Lis return normalized_dims +def is_global_avg_pool_mean(node: torch.fx.Node) -> bool: + """ + Whether a mean.dim node averages a 4D tensor across dims [2, 3] with + keepdim=True, which lowers to Global Average Pooling on an NHWC input. + """ + if not (len(node.args) >= 3 and bool(node.args[2])): + return False + # With keepdim the output has the input's rank. Reading it from the mean + # itself keeps this valid inside ChannelsLastTaggedReshapePass, which can + # replace the mean's input with a node that has no "val" yet. + rank = node.meta["val"].dim() + return rank == 4 and sorted(normalize_mean_dims(node.args[1], rank)) == [2, 3] + + def normalize_pool2d_args( node: torch.fx.Node, has_dilation: bool ) -> Tuple[List[int], List[int], List[int], List[int]]: diff --git a/docs/source/backends/xnnpack/op-support.csv b/docs/source/backends/xnnpack/op-support.csv index 1b2077ee13b..9bd121fa3d1 100644 --- a/docs/source/backends/xnnpack/op-support.csv +++ b/docs/source/backends/xnnpack/op-support.csv @@ -29,7 +29,7 @@ quantize\_: 8-bit dynamic activations log,"fp16, fp32",, max_pool2d,"fp16, fp32",,"stride ≤ kernel_size, ceil_mode only for static shapes" maximum,"fp16, fp32",, -mean,"fp16, fp32",,"4D tensors only; dims=[-2,-1] or [-1,-2]" +mean,"fp16, fp32",PT2E: static int8,no dtype argument minimum,"fp16, fp32",, mul,"fp16, fp32",PT2E: static int8, neg,"fp16, fp32",,