Lower mean.dim beyond 4D pooling as an XNNPACK static mean

.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

@@ -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]: