Lower bf16 mean.dim to XNNPACK as an fp32 static mean

## Ulterior Motive

Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels.

## Rationale

**What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from pytorch#24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul.

**Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate.

## Details

Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new.

- `defineGenericReduceNode` adds the converts when the serialized input and output are bf16, next to `defineBf16MatMulResult`. The fp32 tensor definition is shared with the matmul path as `defineFp32Like`, and `convertBf16MatMulResult` becomes `convertBf16Result`. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048).
- A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it.
- bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741.
- Latency, 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: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741.
- A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main).
- Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries.
- An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`.
- `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16.

Test plan:

```
python -m pytest backends/xnnpack/test/ops/test_mean_dim.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

backends/xnnpack/operators/op_mean_dim.py

@@ -34,10 +34,11 @@
 @register_node_visitor
 class MeanDim(NodeVisitor):
     """
-    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.
+    A mean.dim that is not bf16, 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, which the runtime computes in
+    fp32 for bf16.
     """
 
     target = "aten.mean.dim"

backends/xnnpack/partition/config/generic_node_configs.py

@@ -710,13 +710,12 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
 
 class MeanDimConfig(GenericNodePartitionerConfig):
     target_name = "mean.dim"
-    supports_bf16 = False
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
         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.
+        that is not bf16 across the innermost dimensions with keepdim=True, and
+        to a static mean reduction otherwise, computed in fp32 for bf16.
         """
         if not self.check_common_constraints(node, ep):
             return False

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -818,10 +818,42 @@ const fb_xnnpack::XNNTensorValue* getSerializedTensorValue(
   return nullptr;
 }
 
+/*
+Define an fp32 value with the dims of a serialized bf16 tensor, for an op that
+XNNPACK only runs in fp32.
+*/
+Error defineFp32Like(
+    xnn_subgraph_t subgraph_ptr,
+    const fb_xnnpack::XNNTensorValue* value,
+    const NodePtr node,
+    uint32_t* fp32_id) noexcept {
+  // Like defineTensor(), treat a null dims vector as rank 0.
+  std::vector<size_t> dims;
+  if (value->dims() != nullptr) {
+    dims = flatbufferDimsToVector<size_t>(value->dims());
+  }
+  xnn_status status = xnn_define_tensor_value(
+      subgraph_ptr,
+      xnn_datatype_fp32,
+      dims.size(),
+      dims.data(),
+      /*data=*/nullptr,
+      /*external_id=*/XNN_INVALID_VALUE_ID,
+      /*flags=*/0,
+      fp32_id);
+  ET_CHECK_OR_RETURN_ERROR(
+      status == xnn_status_success,
+      Internal,
+      "Failed to define an fp32 value for bf16 node %i: %s",
+      node->debug_handle(),
+      xnn_status_to_string(status));
+  return Error::Ok;
+}
+
 /*
 XNNPACK multiplies bf16 matrices only into an fp32 result. When the serialized
 graph multiplies bf16 matrices into a bf16 output, define an fp32 value for the
-op to write instead, which convertBf16MatMulResult() then converts into the
+op to write instead, which convertBf16Result() then converts into the
 output, so the delegate boundary stays bf16. Otherwise the op writes the output
 itself.
 */
@@ -844,34 +876,15 @@ Error defineBf16MatMulResult(
       output->datatype() != DataType::xnn_datatype_bf16) {
     return Error::Ok;
   }
-  // Like defineTensor(), treat a null dims vector as rank 0.
-  std::vector<size_t> dims;
-  if (output->dims() != nullptr) {
-    dims = flatbufferDimsToVector<size_t>(output->dims());
-  }
-  xnn_status status = xnn_define_tensor_value(
-      subgraph_ptr,
-      xnn_datatype_fp32,
-      dims.size(),
-      dims.data(),
-      /*data=*/nullptr,
-      /*external_id=*/XNN_INVALID_VALUE_ID,
-      /*flags=*/0,
-      result_id);
-  ET_CHECK_OR_RETURN_ERROR(
-      status == xnn_status_success,
-      Internal,
-      "Failed to define the fp32 result of bf16 node %i: %s",
-      node->debug_handle(),
-      xnn_status_to_string(status));
-  return Error::Ok;
+  return defineFp32Like(subgraph_ptr, output, node, result_id);
 }
 
 /*
-Converts the fp32 result that defineBf16MatMulResult() defined into the bf16
-output. Does nothing when the op wrote the output itself.
+Converts the fp32 result of a bf16 node, as defined by defineBf16MatMulResult()
+or defineGenericReduceNode(), into the bf16 output. Does nothing when the op
+wrote the output itself.
 */
-Error convertBf16MatMulResult(
+Error convertBf16Result(
     xnn_subgraph_t subgraph_ptr,
     uint32_t result_id,
     uint32_t remapped_output_id,
@@ -934,7 +947,7 @@ Error defineFullyConnectedNode(
       node->debug_handle(),
       xnn_status_to_string(status));
 
-  return convertBf16MatMulResult(subgraph_ptr, fc_result, fc_output, node);
+  return convertBf16Result(subgraph_ptr, fc_result, fc_output, node);
 };
 
 /*
@@ -1793,7 +1806,7 @@ Error defineBatchMatrixMultiplyNode(
       node->debug_handle(),
       xnn_status_to_string(status));
 
-  return convertBf16MatMulResult(subgraph_ptr, bmm_result, bmm_out, node);
+  return convertBf16Result(subgraph_ptr, bmm_result, bmm_out, node);
 }
 
 /*
@@ -2043,36 +2056,61 @@ Error defineGenericReduceNode(
     const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
     const fb_xnnpack::_XNNStaticReduce* graph_node,
     xnn_reduce_operator op_type,
-    fb_xnnpack::XNodeUnion node_type,
-    uint32_t debug_handle) noexcept {
+    const NodePtr node,
+    const fb_xnnpack::XNNGraph* graph) noexcept {
+  const char* name = fb_xnnpack::EnumNameXNodeUnion(node->xnode_union_type());
   ET_CHECK_OR_RETURN_ERROR(
       graph_node->reduction_axes() != nullptr,
       InvalidProgram,
       "%s: reduction_axes is null",
-      fb_xnnpack::EnumNameXNodeUnion(node_type));
+      name);
   std::vector<int64_t> reduction_axes =
       flatbufferDimsToVector<int64_t>(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(
+  // XNNPACK has no bf16 reduction, so a bf16 input is reduced in fp32 and the
+  // result converted back, keeping the delegate boundary bf16.
+  uint32_t reduce_source = reduce_input;
+  uint32_t reduce_result = reduce_output;
+  const auto* input = getSerializedTensorValue(graph, graph_node->input_id());
+  const auto* output = getSerializedTensorValue(graph, graph_node->output_id());
+  xnn_status status;
+  if (input != nullptr && output != nullptr &&
+      input->datatype() == DataType::xnn_datatype_bf16 &&
+      output->datatype() == DataType::xnn_datatype_bf16) {
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        defineFp32Like(subgraph_ptr, input, node, &reduce_source));
+    status = xnn_define_convert(
+        subgraph_ptr, reduce_input, reduce_source, /*flags=*/0);
+    ET_CHECK_OR_RETURN_ERROR(
+        status == xnn_status_success,
+        Internal,
+        "Failed to convert the input of bf16 node %i to fp32: %s",
+        node->debug_handle(),
+        xnn_status_to_string(status));
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        defineFp32Like(subgraph_ptr, output, node, &reduce_result));
+  }
+
+  status = xnn_define_static_reduce_v2(
       subgraph_ptr,
       op_type,
       reduction_axes.size(),
       reduction_axes.data(),
-      reduce_input,
-      reduce_output,
+      reduce_source,
+      reduce_result,
       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,
+      name,
+      node->debug_handle(),
       xnn_status_to_string(status));
 
-  return Error::Ok;
+  return convertBf16Result(subgraph_ptr, reduce_result, reduce_output, node);
 }
 
 #define _DEFINE_REDUCE_NODE(name, op_type)                        \
@@ -2081,14 +2119,13 @@ Error defineGenericReduceNode(
       const std::unordered_map<uint32_t, uint32_t>& 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());                                    \
+        node,                                                     \
+        graph);                                                   \
   }
 
 /*

backends/xnnpack/test/ops/test_mean_dim.py

@@ -8,7 +8,9 @@
 import unittest
 
 import torch
-from executorch.backends.xnnpack.test.tester import Tester
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+from executorch.backends.xnnpack.test.tester import Tester, ToEdgeTransformAndLower
+from executorch.export.utils import is_fbcode
 
 
 class TestMeanDim(unittest.TestCase):
@@ -30,6 +32,15 @@ def forward(self, x):
                 z = torch.mean(y, self.dims, keepdim=self.keepdim, dtype=self.dtype)
             return z
 
+    class GraphInputMeanDim(torch.nn.Module):
+        def __init__(self, dims, keepdim=True):
+            super().__init__()
+            self.dims = dims
+            self.keepdim = keepdim
+
+        def forward(self, x):
+            return torch.mean(x, self.dims, keepdim=self.keepdim)
+
     class RMSNorm(torch.nn.Module):
         def __init__(self, dim, eps=1e-6):
             super().__init__()
@@ -73,22 +84,36 @@ def _test_mean_dim(
         dynamic_shapes=None,
         global_avg_pool=False,
         num_runs=1,
+        module=None,
+        xnnpack_nodes=None,
     ):
+        """
+        A bf16 mean serializes as a bf16 XNNStaticMean with no converts; the
+        runtime computes it in fp32.
+        """
+        bf16 = inputs[0].dtype == torch.bfloat16
+        module = module or self.MeanDim(dims, keepdim=keepdim)
+        if xnnpack_nodes is None:
+            xnnpack_nodes = self._xnnpack_mean_nodes(global_avg_pool)
+        if bf16:
+            module = module.to(torch.bfloat16)
+            xnnpack_nodes = {"XNNConvert": 0, **xnnpack_nodes}
+        tol = 1e-2 if bf16 else 1e-3
         (
-            Tester(
-                self.MeanDim(dims, keepdim=keepdim),
-                inputs,
-                dynamic_shapes=dynamic_shapes,
-            )
+            Tester(module, inputs, dynamic_shapes=dynamic_shapes)
             .export()
             .check_count({"torch.ops.aten.mean.dim": 1})
-            .to_edge_transform_and_lower()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    partitioners=[XnnpackPartitioner(enable_bf16=bf16)]
+                )
+            )
             .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))
+            .check_xnnpack_node_count(xnnpack_nodes)
             .to_executorch()
             .serialize()
-            .run_method_and_compare_outputs(num_runs=num_runs)
+            .run_method_and_compare_outputs(num_runs=num_runs, atol=tol, rtol=tol)
         )
 
     def _test_qs8_mean_dim(self, inputs, dims, global_avg_pool=False):
@@ -262,6 +287,90 @@ def test_qs8_conv_mean_dim(self):
                     .run_method_and_compare_outputs(qtol=1)
                 )
 
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_last_dim_3d(self):
+        inputs = (torch.randn(2, 64, 128).to(torch.bfloat16),)
+        self._test_mean_dim(inputs, dims=(-1,))
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_middle_dim_keepdim_false(self):
+        inputs = (torch.randn(2, 64, 128).to(torch.bfloat16),)
+        self._test_mean_dim(inputs, dims=(1,), keepdim=False)
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_global_avg_pool_shape(self):
+        """XNNPACK has no bf16 Global Average Pooling."""
+        for shape in ((1, 16, 32, 32), (1, 5, 4, 4)):
+            with self.subTest(shape=shape):
+                inputs = (torch.randn(*shape).to(torch.bfloat16),)
+                self._test_mean_dim(inputs)
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_channels_last_input(self):
+        """A channels_last input is copied to NCHW for the static mean."""
+        inputs = (
+            torch.randn(1, 16, 8, 8)
+            .to(torch.bfloat16)
+            .to(memory_format=torch.channels_last),
+        )
+        self._test_mean_dim(
+            inputs,
+            module=self.GraphInputMeanDim((-1, -2)),
+            xnnpack_nodes={
+                "XNNGlobalAvgPooling2d": 0,
+                "XNNStaticMean": 1,
+                "XNNStaticTranspose": 1,
+            },
+        )
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_small(self):
+        for shape in ((1, 1, 64), (2, 7, 16)):
+            with self.subTest(shape=shape):
+                inputs = (torch.randn(*shape).to(torch.bfloat16),)
+                self._test_mean_dim(inputs, dims=(-1,))
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_all_dims(self):
+        inputs = (torch.randn(2, 7, 16).to(torch.bfloat16),)
+        for keepdim in (True, False):
+            with self.subTest(keepdim=keepdim):
+                self._test_mean_dim(inputs, dims=None, keepdim=keepdim)
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_of_graph_input(self):
+        inputs = (torch.randn(2, 64, 128).to(torch.bfloat16),)
+        self._test_mean_dim(inputs, module=self.GraphInputMeanDim((-1,)))
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_mean_dim_dynamic(self):
+        seq = torch.export.Dim("seq", min=1, max=32)
+        for dims in ((-1,), (1,)):
+            with self.subTest(dims=dims):
+                inputs = (torch.randn(2, 7, 16).to(torch.bfloat16),)
+                self._test_mean_dim(
+                    inputs, dims=dims, dynamic_shapes=({1: seq},), num_runs=5
+                )
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_rms_norm(self):
+        inputs = (torch.randn(2, 64, 128).to(torch.bfloat16),)
+        # The one convert casts the eps constant to bf16.
+        self._test_mean_dim(
+            inputs,
+            module=self.RMSNorm(128),
+            xnnpack_nodes={"XNNStaticMean": 1, "XNNConvert": 1},
+        )
+
+    def test_bf16_mean_dim_not_partitioned_without_enable_bf16(self):
+        inputs = (torch.randn(2, 64, 128).to(torch.bfloat16),)
+        (
+            Tester(self.MeanDim((-1,)).to(torch.bfloat16), inputs)
+            .export()
+            .to_edge_transform_and_lower()
+            .check_count({"executorch_exir_dialects_edge__ops_aten_mean_dim": 1})
+        )
+
     def test_fp32_mean_dim_unsupported_dtype(self):
         inputs = (torch.randn(1, 5, 4, 4),)
         (