Lower mean.dim beyond 4D pooling as an XNNPACK static mean (pytorch#24741)

## Ulterior Motive

Run RMSNorm in LeRobot policies on XNNPACK (pytorch#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<think>`): 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.

Author: JakeStevens

.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"(?:<think>\n)+"
+    r"^<\|im_start\|>user Calculate 15% of 80\?<\|im_end\|><\|im_start\|>assistant:"
+    r"(?:\n(?:<think>\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:
-<think>
-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

.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"(?:<think>\n)+"
+    r"^<\|im_start\|>user Calculate 15% of 80\?<\|im_end\|><\|im_start\|>assistant:"
+    r"(?:\n(?:<think>\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:
-<think>
-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

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
 

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)

backends/xnnpack/partition/config/generic_node_configs.py

@@ -715,44 +715,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]: