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

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

@@ -711,13 +711,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

@@ -946,11 +946,11 @@ Error defineFullyConnectedNode(
 };
 
 /*
-XNNPACK defines softmax only on fp32 and fp16 tensors. When a serialized
-softmax value is bf16, define an fp32 value for the softmax to use instead.
-Otherwise the softmax uses the value itself.
+XNNPACK defines some ops, such as softmax and reductions, only on fp32 and
+fp16 tensors. When a serialized value of such an op is bf16, define an fp32
+value for the op to use instead. Otherwise the op uses the value itself.
 */
-Error defineBf16SoftmaxFp32Value(
+Error defineFp32ValueForBf16(
     xnn_subgraph_t subgraph_ptr,
     const fb_xnnpack::XNNGraph* graph,
     uint32_t serialized_id,
@@ -980,15 +980,15 @@ Error defineSoftmaxNode(
   REMAP_ID(remapped_ids, graph_node->output_id(), sm_output);
 
   uint32_t softmax_input = XNN_INVALID_VALUE_ID;
-  ET_CHECK_OK_OR_RETURN_ERROR(defineBf16SoftmaxFp32Value(
+  ET_CHECK_OK_OR_RETURN_ERROR(defineFp32ValueForBf16(
       subgraph_ptr,
       graph,
       graph_node->input_id(),
       sm_input,
       node,
       &softmax_input));
   uint32_t softmax_output = XNN_INVALID_VALUE_ID;
-  ET_CHECK_OK_OR_RETURN_ERROR(defineBf16SoftmaxFp32Value(
+  ET_CHECK_OK_OR_RETURN_ERROR(defineFp32ValueForBf16(
       subgraph_ptr,
       graph,
       graph_node->output_id(),
@@ -2089,36 +2089,57 @@ 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);
 
+  uint32_t reduce_source = XNN_INVALID_VALUE_ID;
+  ET_CHECK_OK_OR_RETURN_ERROR(defineFp32ValueForBf16(
+      subgraph_ptr,
+      graph,
+      graph_node->input_id(),
+      reduce_input,
+      node,
+      &reduce_source));
+  uint32_t reduce_result = XNN_INVALID_VALUE_ID;
+  ET_CHECK_OK_OR_RETURN_ERROR(defineFp32ValueForBf16(
+      subgraph_ptr,
+      graph,
+      graph_node->output_id(),
+      reduce_output,
+      node,
+      &reduce_result));
+
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      convertIfDistinct(subgraph_ptr, reduce_input, reduce_source, node));
+
   xnn_status 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 convertIfDistinct(subgraph_ptr, reduce_result, reduce_output, node);
 }
 
 #define _DEFINE_REDUCE_NODE(name, op_type)                        \
@@ -2127,14 +2148,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__()
@@ -98,22 +109,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(
@@ -342,6 +367,90 @@ def test_conv_mean_meets_nhwc(self):
                         .run_method_and_compare_outputs(atol=2e-2 if quantize else 1e-3)
                     )
 
+    @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_fp16_mean_dim_dtype_argument(self):
         """
         An fp16 mean with an fp32 dtype passes the generic dtype checks, so

backends/xnnpack/test/test_xnnpack_partitioner.py

@@ -178,10 +178,6 @@ def test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16(self):
         same modules are delegated in fp32, so bf16 is what keeps them out.
         """
 
-        class Mean(torch.nn.Module):
-            def forward(self, x):
-                return x.mean((-1, -2), keepdim=True)
-
         def lower(module, x):
             return to_edge_transform_and_lower(
                 export(module, (x,)),
@@ -195,7 +191,6 @@ def lower(module, x):
                 torch.nn.Conv2d(4, 4, 3), torch.nn.BatchNorm2d(4)
             ),
             "max_pool2d": torch.nn.MaxPool2d(2),
-            "mean": Mean(),
             "upsample_bilinear2d": torch.nn.Upsample(scale_factor=2, mode="bilinear"),
         }
         x = torch.randn(1, 4, 8, 8)

backends/xnnpack/utils/utils.py

@@ -92,9 +92,12 @@ 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.
+    XNNPACK has no bf16 Global Average Pooling, so this is False for bf16.
     """
     if not (len(node.args) >= 3 and bool(node.args[2])):
         return False
+    if node.meta["val"].dtype == torch.bfloat16:
+        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.

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",PT2E: static int8,no dtype argument
+mean,"fp16, fp32",PT2E: static int8,"no dtype argument; with enable_bf16, a bf16 mean is computed in fp32"
 minimum,"fp16, fp32",,
 mul,"fp16, fp32",PT2E: static int8,
 neg,"fp16, fp32",,