From f56f520b913de2ce7dda89520995d10901bef7ef Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 08:45:40 -0800
Subject: [PATCH 1/2] Keep ops above 6 dims out of XNNPACK and fix quantized
 ops it leaves out (#24380)

## What is wrong today

XNNPACK supports tensors with at most 6 dimensions, but the XNNPACK
partitioner, which decides which ops run on XNNPACK, does not check
this. A permute or another op on a 7-D or 8-D tensor then lowers fine
but fails to load:

```
Failed to define tensor 0 with code: xnn_status_unsupported_parameter
```

Some vision models hit this in a "patchify" step that cuts an image into
patches:

```python
y = x.reshape(2, 6, 8, 2, 16, 8, 2, 16).permute(0, 2, 5, 3, 6, 1, 4, 7)
y = y.reshape(2, 64, -1) * 2
```

With the 3-D convolution kernel from #24385, a statically quantized
`Conv3d` is worse: it fails to load or, without a fused ReLU, returns
one wrong value repeated everywhere, with no error. XNNPACK has no 3-D
convolution, so the partitioner leaves the conv out, but only after
marking the quantize and dequantize nodes around it as folded into it.
Those nodes then land in delegates that skip them, so nothing writes the
outputs. A statically quantized `Conv2d` or `Linear` whose output is not
quantized hits the same bug, even without #24385.

## What this change does

The partitioner now leaves out any op with an input or output above 6
dimensions, and it checks every node that a partition takes. A float
linear on such an input still keeps its 2-D matrix multiply on XNNPACK.

Quantize and dequantize nodes are now marked as folded only once the op
they belong to is really partitioned, so a rejected op keeps them as
separate ops.

`force_non_static_weights_for_f32_linear` no longer changes how a
quantized linear is partitioned. With it, a quantized linear decomposed
into a matrix multiply (for example with `to_backend`) used to lose its
quantize and dequantize nodes: a static one failed to lower or load, and
a dynamic one ran as a float matrix multiply. Now it is partitioned as
without the flag.

The XNNPACK quantizer now leaves 3-D convolutions in float. The model
file gets larger, but there are no extra quantize and dequantize steps
and no quantization error. They run on the CPU kernels, which support
3-D convolution once #24385 lands. Without it, such a model fails with
"Expect input tensor to be 3-D or 4-D".

Still not covered: a linear on such an input, quantized with torchao's
affine ops. Those ops have no portable kernels, so the model now fails
in `to_executorch()` instead of at load. And a linear on such an input
quantized with the XNNPACK quantizer runs outside XNNPACK, because its
quantized input has too many dims.

## Testing

- `test_permute.py` checks which ops are delegated for 6-D and 7-D
permutes, and for a linear on an 8-D input: float (also with a ReLU),
static and dynamic quantization, each in both partitioner modes.
- New tests cover a statically quantized Conv2d and Linear with a float
output, a statically and a dynamically quantized linear lowered through
`to_backend` with `force_non_static_weights_for_f32_linear`, and 3-D
convolutions that the quantizer leaves unquantized.
- The per-channel transposed conv test now also runs the convs that
XNNPACK rejects. Before this change they returned all zeros.
- Each new test fails without the part of the change it covers.
- On macOS arm64, the patchify example and other 7-D and 8-D cases load
and match eager PyTorch. A Conv3d model quantized with the XNNPACK
quantizer matches eager PyTorch with #24385, and fails without it.

(cherry picked from commit 1549787b25fe3fd398aad1976e96f0f970d30379)
---
 .../xnnpack/partition/config/gemm_configs.py  | 21 +++--
 .../partition/config/generic_node_configs.py  | 15 +---
 .../partition/config/xnnpack_config.py        | 45 ++++++++--
 .../xnnpack/quantizer/xnnpack_quantizer.py    |  2 -
 .../quantizer/xnnpack_quantizer_utils.py      | 12 ++-
 backends/xnnpack/test/BUCK                    |  1 +
 backends/xnnpack/test/ops/test_conv2d.py      | 76 ++++++++++++++--
 backends/xnnpack/test/ops/test_linear.py      | 54 +++++++++++
 backends/xnnpack/test/ops/test_permute.py     | 89 ++++++++++++++++++-
 9 files changed, 274 insertions(+), 41 deletions(-)

diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py
index 4df6c4f1141..b60740adbd9 100644
--- a/backends/xnnpack/partition/config/gemm_configs.py
+++ b/backends/xnnpack/partition/config/gemm_configs.py
@@ -25,7 +25,6 @@
     is_per_tensor,
     is_qparam,
     is_quant,
-    tag_as_implicit_q_dq,
 )
 from executorch.backends.xnnpack.utils.utils import (
     get_input_node,
@@ -139,11 +138,6 @@ def get_deps(
         valid_deps = valid_bias and valid_weight and valid_act and valid_output
         deps = list(chain(bias_deps, weight_deps, act_deps, output_deps))
 
-        # Tag q/dq nodes as implicit q/dq nodes
-        for dep in deps:
-            if is_dequant(dep) or is_quant(dep):
-                tag_as_implicit_q_dq(dep)
-
         return valid_deps, deps
 
     def _get_weight_deps(
@@ -224,9 +218,12 @@ def _get_output_deps(
             node_users = list(node.users.keys())
             if len(node_users) == 1:
                 n_output = node_users[0]
+                # After a decomposed linear, the activation can have more dims than
+                # XNNPACK allows even when the 2-D matmul does not. Leave it out then.
                 if (
                     n_output.op == "call_function"
                     and format_target_name(n_output.target.__name__) in self.fused_acts
+                    and self._check_tensor_ranks(n_output)
                 ):
                     gemm_deps.append(n_output)
 
@@ -566,17 +563,23 @@ def find_partition_args(input_node):
         node.args = tuple(fake_args)
         node.users = src_partition.output_nodes[0].users
         valid_deps, deps = super().get_deps(node, ep)
+        _, precision = self._overwrite_precision(node)
 
         # Reset addmm node back to old args and users
         node.args = old_args
         node.users = old_users
 
-        # When using force_non_static_weights_for_f32_linear, we want to get_deps to overwrite the source partition nodes.
-        # Else we want to be greedy.
+        # A float linear with force_non_static_weights_for_f32_linear takes only the
+        # deps inside the source partition. Else we want to be greedy, but leave out
+        # views XNNPACK cannot hold.
         ret_deps = (
             list(set(deps) & set(src_partition.nodes))
             if self.force_non_static_weights_for_f32_linear
-            else list(set(deps) | set(src_partition.nodes))
+            and precision == ConfigPrecisionType.FP32
+            else list(
+                set(deps)
+                | {n for n in src_partition.nodes if self._check_tensor_ranks(n)}
+            )
         )
 
         return valid_deps, ret_deps
diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py
index d445b5ca8ac..9394934e0c4 100644
--- a/backends/xnnpack/partition/config/generic_node_configs.py
+++ b/backends/xnnpack/partition/config/generic_node_configs.py
@@ -16,11 +16,7 @@
     ConfigPrecisionType,
     XNNPartitionerConfig,
 )
-from executorch.backends.xnnpack.utils.quant_utils import (
-    is_dequant,
-    is_quant,
-    tag_as_implicit_q_dq,
-)
+from executorch.backends.xnnpack.utils.quant_utils import is_dequant, is_quant
 from executorch.backends.xnnpack.utils.utils import (
     get_input_node,
     get_param_tensor,
@@ -164,15 +160,6 @@ def get_node_and_deps(
                 # Expected node --> fused_act (optional) --> dequant
                 return deps
 
-            # Tag input nodes (dq nodes) as implicit q/dq nodes
-            for dq_input in node.all_input_nodes:
-                if is_dequant(dq_input):
-                    tag_as_implicit_q_dq(dq_input)
-
-            # Tag node_output (q node) as an implicit q/dq node
-            if is_quant(node_output):
-                tag_as_implicit_q_dq(node_output)
-
             quantized_deps.append(node_output)
 
         return deps + quantized_deps
diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py
index e0a575e000a..b5267d2405a 100644
--- a/backends/xnnpack/partition/config/xnnpack_config.py
+++ b/backends/xnnpack/partition/config/xnnpack_config.py
@@ -14,7 +14,9 @@
     is_dequant,
     is_qparam,
     is_quant,
+    tag_as_implicit_q_dq,
 )
+from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_MAX_TENSOR_DIMS
 from executorch.exir.backend.canonical_partitioners.config_partitioner import (
     format_target_name,
     PartitionerConfig,
@@ -66,14 +68,31 @@ def get_partition(
         with the name of the XNNPartitionerConfig class which return this set of nodes.
         This enforces that all partitions returned by XNNPartitioner configs are labeled
         with the partitioner config which returned them
+
+        Returns an empty list if any of the nodes has more dims than XNNPACK supports.
+        Otherwise marks every quantize and dequantize node in the partition, other
+        than the given node, as implicit, so that it is folded into the op it feeds
+        or follows.
         """
         partitioned_nodes = self.get_node_and_deps(node, ep)
+        # A partition can take nodes that their own configs reject.
+        for n in partitioned_nodes:
+            if not self._check_tensor_ranks(n):
+                why(node, reason=f"its partition takes {n}, which has too many dims")
+                return []
+
+        # Fold only now, so a rejected node keeps its q/dq nodes as separate ops. A
+        # q/dq node partitioned on its own stays a separate op as well.
+        for n in partitioned_nodes:
+            if n is not node and (is_quant(n) or is_dequant(n)):
+                tag_as_implicit_q_dq(n)
+
         # label partitioned nodes with the name of the partitioner config
-        for node in partitioned_nodes:
-            if "xnn_partitioner_config" in node.meta:
-                node.meta["xnn_partitioner_config"].append(self.__class__.__name__)
+        for n in partitioned_nodes:
+            if "xnn_partitioner_config" in n.meta:
+                n.meta["xnn_partitioner_config"].append(self.__class__.__name__)
             else:
-                node.meta["xnn_partitioner_config"] = [self.__class__.__name__]
+                n.meta["xnn_partitioner_config"] = [self.__class__.__name__]
 
         return partitioned_nodes
 
@@ -94,7 +113,9 @@ def get_node_and_deps(
     ) -> List[torch.fx.Node]:
         """
         Takes in a node and its exported program and returns a list of nodes
-        and its dependencies that need to be partitioned together
+        and its dependencies that need to be partitioned together. Quantize and
+        dequantize nodes in the list, other than the node itself, are folded into
+        the partition's op once get_partition takes it.
 
         Args:
             node: Node to be partitioned
@@ -152,6 +173,20 @@ def check_common_constraints(
             why(node, reason="invalid dtype")
             return False
 
+        return self._check_tensor_ranks(node)
+
+    def _check_tensor_ranks(self, node: torch.fx.Node) -> bool:
+        vals = [arg.meta.get("val") for arg in node.all_input_nodes]
+        node_val = node.meta.get("val")
+        vals.extend(node_val if isinstance(node_val, (list, tuple)) else [node_val])
+        for val in vals:
+            if isinstance(val, torch.Tensor) and val.dim() > XNN_MAX_TENSOR_DIMS:
+                why(
+                    node,
+                    reason=f"XNNPACK supports at most {XNN_MAX_TENSOR_DIMS} dims, got {val.dim()}",
+                )
+                return False
+
         return True
 
     def _check_inputs_are_valid_dtypes(self, node, valid_dtypes):
diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer.py b/backends/xnnpack/quantizer/xnnpack_quantizer.py
index 3c82a65ad71..347c6b21972 100644
--- a/backends/xnnpack/quantizer/xnnpack_quantizer.py
+++ b/backends/xnnpack/quantizer/xnnpack_quantizer.py
@@ -256,8 +256,6 @@ class QuantPattern:
     torch.ops.aten.conv_transpose1d.default,
     torch.ops.aten.conv_transpose2d,
     torch.ops.aten.conv_transpose2d.input,
-    torch.ops.aten.conv_transpose3d,
-    torch.ops.aten.conv_transpose3d.input,
 }
 
 LINEAR_TARGETS = {
diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py
index a207d77674d..4c95901796e 100644
--- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py
+++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py
@@ -239,6 +239,11 @@ def _annotate_linear_relu(
     return annotated_partitions
 
 
+def _is_1d_or_2d_conv(weight_shape: Optional[torch.Size]) -> bool:
+    # XNNPACK runs no other convs.
+    return weight_shape is None or len(weight_shape) <= 4
+
+
 def _do_annotate_conv(
     gm: torch.fx.GraphModule,
     quantization_config: Optional[QuantizationConfig],
@@ -273,9 +278,11 @@ def _do_annotate_conv(
         assert isinstance(weight, Node)
         weight_qspec = get_weight_qspec(quantization_config)
         num_groups = get_groups_from_conv(conv_node)
+        weight_shape = getattr(weight.meta.get("val", None), "shape", None)
 
         # skip if transposed conv has more than 1 group
         skip = skip or (is_conv_transpose and num_groups != 1)
+        skip = skip or not _is_1d_or_2d_conv(weight_shape)
 
         if is_conv_transpose:
             # transposed convs per output channel quantization
@@ -290,8 +297,6 @@ def _do_annotate_conv(
 
         # Only annotate dynamically quantized conv if it's 2D and not depthwise
         if is_dynamic:
-            weight_val = weight.meta.get("val", None)
-            weight_shape = getattr(weight_val, "shape", None)
             # Skip if not a 4D weight tensor (i.e. not conv2d)
             skip = skip or (weight_shape is not None and len(weight_shape) != 4)
             # Skip if depthwise (default to groups=1 since it's not an arg)
@@ -355,6 +360,7 @@ def _do_annotate_conv_relu(
         assert isinstance(weight, Node)
         weight_qspec = get_weight_qspec(quantization_config)
         groups = get_groups_from_conv(conv_node)
+        weight_shape = getattr(weight.meta.get("val", None), "shape", None)
         if is_conv_transpose:
             # transposed convs per output channel quantization
             weight_qspec = change_quantization_config(weight_qspec, ch_axis=1)
@@ -370,7 +376,7 @@ def _do_annotate_conv_relu(
         if _is_annotated(partition):
             continue
 
-        if is_conv_transpose and groups != 1:
+        if (is_conv_transpose and groups != 1) or not _is_1d_or_2d_conv(weight_shape):
             continue
 
         if filter_fn and any(not filter_fn(n) for n in partition):
diff --git a/backends/xnnpack/test/BUCK b/backends/xnnpack/test/BUCK
index 0eb2af1bf9b..3e28be745d0 100644
--- a/backends/xnnpack/test/BUCK
+++ b/backends/xnnpack/test/BUCK
@@ -96,6 +96,7 @@ fbcode_target(_kind = runtime.python_test,
         "test_xnnpack_utils.py",
     ],
     supports_static_listing = False,
+    preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         "fbsource//third-party/pypi/parameterized:parameterized",
         "//executorch/backends/test/harness:tester",
diff --git a/backends/xnnpack/test/ops/test_conv2d.py b/backends/xnnpack/test/ops/test_conv2d.py
index 0db84852422..2daacfd44cf 100644
--- a/backends/xnnpack/test/ops/test_conv2d.py
+++ b/backends/xnnpack/test/ops/test_conv2d.py
@@ -6,6 +6,7 @@
 
 import itertools
 import unittest
+from dataclasses import replace
 from typing import Optional
 
 import torch
@@ -355,6 +356,34 @@ def test_fp32_conv3d_single_element_stride_doesnt_partition(self) -> None:
             .run_method_and_compare_outputs()
         )
 
+    def test_qs8_quantizer_skips_conv3d(self) -> None:
+        # XNNPACK has no conv3d support, so its quantizer leaves 3d convs in float,
+        # with or without a fused relu. The same model in 2d gets its activations
+        # and weights quantized. The weights are quantized ahead of time, so only
+        # their dequantize nodes remain.
+        quant_config = get_symmetric_quantization_config(is_per_channel=True)
+        ops = "torch.ops.quantized_decomposed"
+        for dims, conv, conv_transpose, quantized_activations, quantized_weights in (
+            (2, torch.nn.Conv2d, torch.nn.ConvTranspose2d, 4, 3),
+            (3, torch.nn.Conv3d, torch.nn.ConvTranspose3d, 0, 0),
+        ):
+            m = torch.nn.Sequential(
+                conv(2, 2, 2), torch.nn.ReLU(), conv(2, 2, 2), conv_transpose(2, 2, 2)
+            )
+            (
+                Tester(m.eval(), (torch.randn([1, 2] + [4] * dims),))
+                .quantize(Quantize(quantization_config=quant_config))
+                .export()
+                .check_count(
+                    {
+                        f"{ops}.quantize_per_tensor.default": quantized_activations,
+                        f"{ops}.dequantize_per_tensor.default": quantized_activations,
+                        f"{ops}.quantize_per_channel.default": 0,
+                        f"{ops}.dequantize_per_channel.default": quantized_weights,
+                    }
+                )
+            )
+
     def test_fp32_conv2d_permute(self) -> None:
         for transpose in (True, False):
             for perm_order in list(itertools.permutations([0, 1, 2, 3])):
@@ -375,6 +404,26 @@ def test_qs8_conv2d_per_channel(self) -> None:
                 quant_config=get_symmetric_quantization_config(is_per_channel=True),
             )
 
+    def test_qs8_conv2d_float_output(self) -> None:
+        # XNNPACK leaves out a quantized conv whose output is not quantized. The
+        # quantize and dequantize of its input then run in a delegate of their own.
+        quant_config = replace(
+            get_symmetric_quantization_config(is_per_channel=True),
+            output_activation=None,
+        )
+        m = Conv2d()
+        (
+            Tester(m.eval(), m.get_inputs())
+            .quantize(Quantize(quantization_config=quant_config))
+            .export()
+            .to_edge_transform_and_lower()
+            .check(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     def test_qs8_conv2d_even_kernel_same_padding(self) -> None:
         # An even-kernel 'same'-padding conv decomposes into
         # dequant -> constant_pad_nd -> convolution. The pad and conv are pulled
@@ -832,14 +881,27 @@ def get_inputs(self):
                         conv_count=1,
                     )
                 else:
-                    with self.assertRaises(RuntimeError):
-                        self._test(
-                            PerChannelConvTranspose2d(
-                                3 * groups, 5 * groups, groups, ch_axis
-                            ),  # ch_axis=0
-                            quant_config=None,
-                            conv_count=1,
+                    # XNNPACK leaves this conv out, and the quantize and dequantize
+                    # nodes around it then run in delegates of their own.
+                    m = PerChannelConvTranspose2d(
+                        3 * groups, 5 * groups, groups, ch_axis
+                    )
+                    (
+                        Tester(m.eval(), m.get_inputs())
+                        .export()
+                        .to_edge_transform_and_lower()
+                        .check(
+                            [
+                                "executorch_exir_dialects_edge__ops_aten_convolution_default"
+                            ]
                         )
+                        .check_count(
+                            {"torch.ops.higher_order.executorch_call_delegate": 2}
+                        )
+                        .to_executorch()
+                        .serialize()
+                        .run_method_and_compare_outputs(qtol=1)
+                    )
 
     def test_fp32_tconv_output_padding(self):
         """Test transposed convolution with non-zero output padding."""
diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py
index e309bc514a7..41e7cf751e3 100644
--- a/backends/xnnpack/test/ops/test_linear.py
+++ b/backends/xnnpack/test/ops/test_linear.py
@@ -6,6 +6,7 @@
 # pyre-unsafe
 
 import unittest
+from dataclasses import replace
 
 from itertools import product
 from typing import Callable, Dict, List, Optional, Tuple
@@ -33,6 +34,13 @@
 
 from torch.export.graph_signature import ExportGraphSignature, InputKind
 
+# Buck preloads the quantized kernels instead of shipping this module.
+try:
+    import executorch.extension.pybindings.portable_lib  # noqa[F401]
+    import executorch.kernels.quantized  # noqa[F401]
+except ImportError:
+    pass
+
 try:
     from torchao.quantization.granularity import PerGroup
     from torchao.quantization.quant_api import (
@@ -839,6 +847,26 @@ def test_qs8_linear(self):
                     quant_type="per_tensor",
                 )
 
+    def test_qs8_linear_float_output(self):
+        # XNNPACK leaves out a quantized linear whose output is not quantized. The
+        # quantize and dequantize of its input then run in a delegate of their own.
+        quant_config = replace(
+            get_symmetric_quantization_config(is_per_channel=True),
+            output_activation=None,
+        )
+        module = BaseLinear(use_bias=True)
+        (
+            Tester(module.eval(), module.get_inputs(rank=2))
+            .quantize(Quantize(quantization_config=quant_config))
+            .export()
+            .to_edge_transform_and_lower()
+            .check(["executorch_exir_dialects_edge__ops_aten_addmm_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     # Tests for q[dp]8-f16-qc8w
     def test_qd8_f16_per_channel_linear(self):
         self._test_qd8_per_channel_linear(dtype=torch.half)
@@ -1079,6 +1107,32 @@ def test_linear_qd8_as_fp32(self):
                 },
             )
 
+    def test_quantized_linear_with_force_non_static_weights_for_f32_linear(self):
+        # The flag is for float linears. A quantized linear that is decomposed into
+        # a matmul keeps its quantize and dequantize nodes and is delegated whole.
+        partitioner = XnnpackPartitioner(force_non_static_weights_for_f32_linear=True)
+        for is_dynamic in (False, True):
+            module = BaseLinear(use_bias=True)
+            quant_config = get_symmetric_quantization_config(
+                is_per_channel=True, is_dynamic=is_dynamic
+            )
+            (
+                Tester(module.eval(), module.get_inputs(rank=2))
+                .quantize(Quantize(quantization_config=quant_config))
+                .export()
+                .to_edge()
+                .partition(Partition(partitioner=partitioner))
+                .check_not(
+                    [
+                        "executorch_exir_dialects_edge__ops_aten_addmm_default",
+                        "quantized_decomposed",
+                    ]
+                )
+                .to_executorch()
+                .serialize()
+                .run_method_and_compare_outputs(atol=5e-2)
+            )
+
     def test_linear_with_force_non_static_weights_for_f32_linear(self):
         def check_signature(
             signature: ExportGraphSignature,
diff --git a/backends/xnnpack/test/ops/test_permute.py b/backends/xnnpack/test/ops/test_permute.py
index e0171f2178f..055108ab472 100644
--- a/backends/xnnpack/test/ops/test_permute.py
+++ b/backends/xnnpack/test/ops/test_permute.py
@@ -7,7 +7,22 @@
 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.quantizer.xnnpack_quantizer import (
+    get_symmetric_quantization_config,
+)
+from executorch.backends.xnnpack.test.tester import (
+    Quantize,
+    Tester,
+    ToEdgeTransformAndLower,
+)
+
+# Buck preloads the quantized kernels instead of shipping this module.
+try:
+    import executorch.extension.pybindings.portable_lib  # noqa[F401]
+    import executorch.kernels.quantized  # noqa[F401]
+except ImportError:
+    pass
 
 
 class TestPermute(unittest.TestCase):
@@ -83,6 +98,78 @@ def test_fp32_permute_copy(self):
             .run_method_and_compare_outputs()
         )
 
+    class LinearPermute8d(torch.nn.Module):
+        def __init__(self, relu=False):
+            super().__init__()
+            self.linear = torch.nn.Linear(4, 4)
+            self.relu = relu
+
+        def forward(self, x):
+            y = self.linear(x.reshape(2, 2, 2, 2, 2, 2, 2, 4))
+            if self.relu:
+                y = torch.relu(y)
+            return y.permute(0, 2, 5, 3, 6, 1, 4, 7).reshape(2, 64, -1) * 2
+
+    def test_fp32_permute_max_rank(self):
+        # XNNPACK allows at most 6 dims, so a 6-D permute is delegated and a 7-D
+        # one is not.
+        delegate = "torch.ops.higher_order.executorch_call_delegate"
+        permute = "executorch_exir_dialects_edge__ops_aten_permute_copy_default"
+        for rank, delegates in ((6, 1), (7, 0)):
+            dims = list(reversed(range(rank)))
+            (
+                Tester(self.Permute(dims), (torch.randn([2] * rank),))
+                .export()
+                .to_edge_transform_and_lower()
+                .check_count({delegate: delegates, permute: 1 - delegates})
+                .to_executorch()
+                .serialize()
+                .run_method_and_compare_outputs()
+            )
+
+    def _test_permute_above_max_rank(self, quant_config=None, relu=False):
+        # A delegate that holds a tensor above 6 dims fails to load, so the 8-D
+        # views and permute stay outside and the mul is delegated. The linear's
+        # 2-D matmul is delegated too, unless its activation is quantized in 8-D.
+        for force_non_static_weights in (False, True):
+            tester = Tester(self.LinearPermute8d(relu).eval(), (torch.randn(2, 64, 4),))
+            if quant_config is not None:
+                tester.quantize(Quantize(quantization_config=quant_config))
+            partitioner = XnnpackPartitioner(
+                force_non_static_weights_for_f32_linear=force_non_static_weights
+            )
+            tester.export().to_edge_transform_and_lower(
+                ToEdgeTransformAndLower([partitioner])
+            )
+            addmm = ["executorch_exir_dialects_edge__ops_aten_addmm_default"]
+            if quant_config is None:
+                tester.check_not(addmm).check_count(
+                    {"torch.ops.higher_order.executorch_call_delegate": 2}
+                )
+            else:
+                tester.check(addmm)
+            (
+                tester.check_count(
+                    {"executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1}
+                )
+                .check_not(["executorch_exir_dialects_edge__ops_aten_mul_Tensor"])
+                .to_executorch()
+                .serialize()
+                .run_method_and_compare_outputs(qtol=1)
+            )
+
+    def test_fp32_permute_above_max_rank(self):
+        for relu in (False, True):
+            self._test_permute_above_max_rank(relu=relu)
+
+    def test_qs8_permute_above_max_rank(self):
+        self._test_permute_above_max_rank(get_symmetric_quantization_config())
+
+    def test_qd8_permute_above_max_rank(self):
+        self._test_permute_above_max_rank(
+            get_symmetric_quantization_config(is_per_channel=True, is_dynamic=True)
+        )
+
     def test_qs8_permute(self):
         inputs = (torch.randn(1, 1, 4, 4),)
         (

From ba703c08efb0d284326b4bfc1aadbdb656f9e7bd Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 16:11:51 -0800
Subject: [PATCH 2/2] XNNPACK: fix enable_bf16 load failures and BatchNorm
 folding into a 3-D linear output (#24464)

### What is wrong today

`XnnpackPartitioner(enable_bf16=True)` lets XNNPACK run bf16 nodes.
Without it, bf16 matrix multiplies run on the single-threaded portable
kernels, which is very slow. With it, a program can fail to load:

1. The partitioner hands XNNPACK bf16 nodes that XNNPACK cannot define
in bf16: softmax, mean, average and max pooling, bilinear upsampling,
and batch norm that does not fold into a linear. Loading fails with
`xnn_status_invalid_parameter`.
2. Batch matrix multiply (`bmm`) is defined with a bf16 output, but
XNNPACK only multiplies bf16 matrices into an fp32 result. Loading fails
with `Failed to create BMM node`.

A bf16 attention block written with `matmul` and `softmax` hits both.

### What this change does

- A config can now say it does not support bf16 (`supports_bf16 =
False`), so its bf16 nodes stay on the portable kernels and the
partitioner logs why. Softmax, mean, both poolings, bilinear upsampling
and convolution set it. The default is `True`: XNNPACK rewrites most
bf16 ops to fp32 internally, and at the pinned commit only these reject
bf16 when defined. A bf16 run of the other float ops found no other
failure.
- A bf16 batch norm is kept in XNNPACK only when
`FuseBatchNormPass.can_fuse` says it folds into the linear before it.
- `can_fuse` now also requires the batch norm to be the only user of the
linear output, and the linear output to be 2-D. Without these checks, a
bf16 batch norm could be kept but left unfused, so loading failed, and a
batch norm after a 3-D linear output was folded along the wrong dim,
which gave wrong results in fp32 and bf16.
- A bf16 `bmm` now writes an fp32 result and converts it to bf16, as the
linear layer already did.

With this change, the SmolVLA robot policy (a vision-language-action
model) in bf16 with `enable_bf16=True` loads, and then runs one action
chunk in about 11 seconds on 16 x86 cores, instead of about 9 minutes
with the default partitioner.

### Tests

- `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16`: each op,
and a convolution followed by a batch norm, is delegated in fp32 but not
in bf16, and the bf16 program matches eager.
- `test_enable_bf16_fuses_batch_norm_into_a_linear`: a bf16 linear
followed by `BatchNorm1d` runs fully in XNNPACK. When the linear output
is also used elsewhere, the batch norm stays out.
- `test_enable_bf16_runs_an_attention_block`: the linear layers and both
`bmm`s run in XNNPACK, the softmax stays outside.
- `test_bf16_bmm` in `test_bmm.py`, and
`test_fp32_linear_batch_norm_on_3d_output` in
`test_batch_norm_fusion.py`.

All of these fail before this change, except the fused batch norm case,
which guards the fusion that already worked. Ran the XNNPACK
partitioner, op and pass tests on Linux x86_64. The only failures, in
the per-channel int4 linear tests, also happen without this change.

(cherry picked from commit 7bf0303039c6190b6081ac52811a7355609c9629)
---
 backends/xnnpack/_passes/fuse_batch_norm.py   |  12 ++
 .../xnnpack/partition/config/gemm_configs.py  |   7 +-
 .../partition/config/generic_node_configs.py  |   5 +
 .../xnnpack/partition/config/node_configs.py  |  12 ++
 .../partition/config/xnnpack_config.py        |  24 ++-
 backends/xnnpack/runtime/XNNCompiler.cpp      | 149 +++++++++++-------
 backends/xnnpack/test/ops/test_bmm.py         |  23 ++-
 .../test/passes/test_batch_norm_fusion.py     |  18 +++
 .../xnnpack/test/test_xnnpack_partitioner.py  | 136 ++++++++++++++++
 9 files changed, 319 insertions(+), 67 deletions(-)

diff --git a/backends/xnnpack/_passes/fuse_batch_norm.py b/backends/xnnpack/_passes/fuse_batch_norm.py
index 76ecce91585..9e383a77c5b 100644
--- a/backends/xnnpack/_passes/fuse_batch_norm.py
+++ b/backends/xnnpack/_passes/fuse_batch_norm.py
@@ -102,6 +102,12 @@ def can_fuse(  # noqa: C901
 
         is_conv = input_node.target == exir_ops.edge.aten.convolution.default
 
+        # Fusing changes what the input node computes, so nothing else may use it.
+        if len(input_node.users) != 1:
+            if why:
+                why(bn, "Batch norm must be the only user of the input node.")
+            return False
+
         # All users of the batch_norm node must be getitem ops.
         # batch_norm returns a 3-element tuple.
         # Each user must only access the first element of the tuple.
@@ -141,6 +147,12 @@ def can_fuse(  # noqa: C901
                 if why:
                     why("Convolution input must be rank 3 or 4.")
                 return False
+        # BatchNorm normalizes dim 1, which holds the linear's output features only
+        # when the linear output is 2-D.
+        elif "val" not in input_node.meta or len(input_node.meta["val"].shape) != 2:
+            if why:
+                why(bn, "Linear output must be rank 2.")
+            return False
 
         return True
 
diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py
index b60740adbd9..7fc6022b970 100644
--- a/backends/xnnpack/partition/config/gemm_configs.py
+++ b/backends/xnnpack/partition/config/gemm_configs.py
@@ -342,6 +342,7 @@ def supported_precision_types(self):
 
 class ConvolutionConfig(GEMMConfig):
     target_name = "convolution.default"
+    supports_bf16 = False
 
     def __init__(self, **kwargs):
         super().__init__(
@@ -361,12 +362,6 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
 
         kernel_node = get_input_node(node, 1)
         kernel_shape = get_shape(kernel_node)
-        if any(
-            value.meta["val"].dtype == torch.bfloat16
-            for value in (node, kernel_node, get_input_node(node, 0))
-        ):
-            why(node, "XNNPACK does not support BF16 convolution")
-            return False
         # The weight rank is the only reliable indicator of the conv dimensionality.
         # stride, padding and dilation may each be a single value that ATen
         # broadcasts over every spatial dim.
diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py
index 9394934e0c4..f53530909c3 100644
--- a/backends/xnnpack/partition/config/generic_node_configs.py
+++ b/backends/xnnpack/partition/config/generic_node_configs.py
@@ -233,6 +233,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
 
 class AvgPoolingConfig(GenericNodePartitionerConfig):
     target_name = "avg_pool2d.default"
+    supports_bf16 = False
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
@@ -359,6 +360,7 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
 
 class SoftmaxConfig(GenericNodePartitionerConfig):
     target_name = "_softmax.default"
+    supports_bf16 = False
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
@@ -436,6 +438,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
 
 class MaxPool2dConfig(GenericNodePartitionerConfig):
     target_name = "max_pool2d.default"
+    supports_bf16 = False
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
@@ -470,6 +473,7 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
 
 class UpsampleBilinear2dConfig(GenericNodePartitionerConfig):
     target_name = "upsample_bilinear2d.vec"
+    supports_bf16 = False
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
@@ -623,6 +627,7 @@ 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:
         """
diff --git a/backends/xnnpack/partition/config/node_configs.py b/backends/xnnpack/partition/config/node_configs.py
index bffe87d30c5..782a336053b 100644
--- a/backends/xnnpack/partition/config/node_configs.py
+++ b/backends/xnnpack/partition/config/node_configs.py
@@ -17,6 +17,7 @@
 )
 from executorch.backends.xnnpack.utils.utils import is_param_node
 from executorch.exir.backend.utils import WhyNoPartition
+from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
 
 logger = logging.getLogger(__name__)
@@ -40,6 +41,17 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
             why(node, f"BatchNorm cannot be decomposed or fused with {input_node}")
             return False
 
+        # Unless it fuses into a linear, a BatchNorm becomes a depthwise convolution
+        # or folds into a convolution, and the pinned XNNPACK defines neither in bf16.
+        if input_node.meta["val"].dtype == torch.bfloat16 and not (
+            can_fuse and input_node.target == exir_ops.edge.aten.linear.default
+        ):
+            why(
+                node,
+                "XNNPACK supports a bf16 BatchNorm only when it fuses into a linear",
+            )
+            return False
+
         return True
 
     def get_node_and_deps(
diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py
index b5267d2405a..e1bc5a4ab0a 100644
--- a/backends/xnnpack/partition/config/xnnpack_config.py
+++ b/backends/xnnpack/partition/config/xnnpack_config.py
@@ -35,6 +35,13 @@ class ConfigPrecisionType(Enum):
     DYNAMIC_QUANT = 3
 
 
+def _input_and_output_values(node: torch.fx.Node) -> List[object]:
+    values = [arg.meta.get("val") for arg in node.all_input_nodes]
+    node_val = node.meta.get("val")
+    values.extend(node_val if isinstance(node_val, (list, tuple)) else [node_val])
+    return values
+
+
 class XNNPartitionerConfig(PartitionerConfig):
     """
     Base partitioner config for XNNPACK Partitioner Configs. Base wrapper class
@@ -45,6 +52,11 @@ class XNNPartitionerConfig(PartitionerConfig):
     types they want to enable
     """
 
+    # At the pinned XNNPACK commit, some ops are defined only for fp32 and fp16
+    # tensors. Their configs set this to False, so their bf16 nodes stay on the
+    # portable kernels instead of failing when the program loads.
+    supports_bf16 = True
+
     def __init__(self, **kwargs):
         super().__init__()
         self.enabled_precision_types = self.supported_precision_types()
@@ -168,6 +180,13 @@ def check_common_constraints(
             why(node, reason="not enabled precision types")
             return False
 
+        if not self.supports_bf16 and any(
+            isinstance(val, torch.Tensor) and val.dtype == torch.bfloat16
+            for val in _input_and_output_values(node)
+        ):
+            why(node, reason="XNNPACK does not support bf16 for this op")
+            return False
+
         has_valid_dtypes = self._check_node_has_valid_dtype(node)
         if not has_valid_dtypes:
             why(node, reason="invalid dtype")
@@ -176,10 +195,7 @@ def check_common_constraints(
         return self._check_tensor_ranks(node)
 
     def _check_tensor_ranks(self, node: torch.fx.Node) -> bool:
-        vals = [arg.meta.get("val") for arg in node.all_input_nodes]
-        node_val = node.meta.get("val")
-        vals.extend(node_val if isinstance(node_val, (list, tuple)) else [node_val])
-        for val in vals:
+        for val in _input_and_output_values(node):
             if isinstance(val, torch.Tensor) and val.dim() > XNN_MAX_TENSOR_DIMS:
                 why(
                     node,
diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp
index 74a5eec4a91..8ae0dcb2e3b 100644
--- a/backends/xnnpack/runtime/XNNCompiler.cpp
+++ b/backends/xnnpack/runtime/XNNCompiler.cpp
@@ -818,6 +818,74 @@ const fb_xnnpack::XNNTensorValue* getSerializedTensorValue(
   return nullptr;
 }
 
+/*
+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
+output, so the delegate boundary stays bf16. Otherwise the op writes the output
+itself.
+*/
+Error defineBf16MatMulResult(
+    xnn_subgraph_t subgraph_ptr,
+    const fb_xnnpack::XNNGraph* graph,
+    uint32_t input1_id,
+    uint32_t input2_id,
+    uint32_t output_id,
+    uint32_t remapped_output_id,
+    const NodePtr node,
+    uint32_t* result_id) noexcept {
+  *result_id = remapped_output_id;
+  const auto* input1 = getSerializedTensorValue(graph, input1_id);
+  const auto* input2 = getSerializedTensorValue(graph, input2_id);
+  const auto* output = getSerializedTensorValue(graph, output_id);
+  if (input1 == nullptr || input2 == nullptr || output == nullptr ||
+      input1->datatype() != DataType::xnn_datatype_bf16 ||
+      input2->datatype() != DataType::xnn_datatype_bf16 ||
+      output->datatype() != DataType::xnn_datatype_bf16) {
+    return Error::Ok;
+  }
+  std::vector<size_t> 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;
+}
+
+/*
+Converts the fp32 result that defineBf16MatMulResult() defined into the bf16
+output. Does nothing when the op wrote the output itself.
+*/
+Error convertBf16MatMulResult(
+    xnn_subgraph_t subgraph_ptr,
+    uint32_t result_id,
+    uint32_t remapped_output_id,
+    const NodePtr node) noexcept {
+  if (result_id == remapped_output_id) {
+    return Error::Ok;
+  }
+  xnn_status status = xnn_define_convert(
+      subgraph_ptr, result_id, remapped_output_id, /*flags=*/0);
+  ET_CHECK_OR_RETURN_ERROR(
+      status == xnn_status_success,
+      Internal,
+      "Failed to convert the result of bf16 node %i to bf16: %s",
+      node->debug_handle(),
+      xnn_status_to_string(status));
+  return Error::Ok;
+}
+
 /*
 Define serialized linear(fully-connected) node into the subgraph using
 the remapped ids to map the serialized ids, to the new ids generated
@@ -828,8 +896,6 @@ Error defineFullyConnectedNode(
     const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
     const NodePtr node,
     const fb_xnnpack::XNNGraph* graph) noexcept {
-  MAYBE_UNUSED(graph);
-
   auto graph_node = node->xnode_union_as_XNNFullyConnected();
   std::pair<float, float> min_max = getOutputMinMax(node);
   REMAP_ID(remapped_ids, graph_node->input1_id(), fc_input1);
@@ -837,43 +903,16 @@ Error defineFullyConnectedNode(
   REMAP_ID(remapped_ids, graph_node->bias_id(), fc_bias);
   REMAP_ID(remapped_ids, graph_node->output_id(), fc_output);
 
-  // XNNPACK only provides a bf16 fully-connected of type bf16_bf16_f32:
-  // bf16 activation x bf16 weight -> fp32 output. When the serialized graph
-  // asks for a bf16 output (e.g. a fully bf16 model), define the FC with an
-  // fp32 intermediate output and append a convert (fp32 -> bf16) so the
-  // delegate boundary stays bf16.
-  const auto* in_tv = getSerializedTensorValue(graph, graph_node->input1_id());
-  const auto* filt_tv =
-      getSerializedTensorValue(graph, graph_node->filter_id());
-  const auto* out_tv = getSerializedTensorValue(graph, graph_node->output_id());
-  const bool needs_bf16_output_convert = in_tv != nullptr &&
-      filt_tv != nullptr && out_tv != nullptr &&
-      in_tv->datatype() == DataType::xnn_datatype_bf16 &&
-      filt_tv->datatype() == DataType::xnn_datatype_bf16 &&
-      out_tv->datatype() == DataType::xnn_datatype_bf16;
-
-  uint32_t fc_compute_output = fc_output;
-  if (needs_bf16_output_convert) {
-    std::vector<size_t> out_dims =
-        flatbufferDimsToVector<size_t>(out_tv->dims());
-    uint32_t intermediate_id = XNN_INVALID_VALUE_ID;
-    xnn_status ts = xnn_define_tensor_value(
-        subgraph_ptr,
-        xnn_datatype_fp32,
-        out_dims.size(),
-        out_dims.data(),
-        /*data=*/nullptr,
-        /*external_id=*/XNN_INVALID_VALUE_ID,
-        /*flags=*/0,
-        &intermediate_id);
-    ET_CHECK_OR_RETURN_ERROR(
-        ts == xnn_status_success,
-        Internal,
-        "Failed to define fp32 intermediate for bf16 linear node %i: %s",
-        node->debug_handle(),
-        xnn_status_to_string(ts));
-    fc_compute_output = intermediate_id;
-  }
+  uint32_t fc_result = XNN_INVALID_VALUE_ID;
+  ET_CHECK_OK_OR_RETURN_ERROR(defineBf16MatMulResult(
+      subgraph_ptr,
+      graph,
+      graph_node->input1_id(),
+      graph_node->filter_id(),
+      graph_node->output_id(),
+      fc_output,
+      node,
+      &fc_result));
 
   xnn_status status = xnn_define_fully_connected(
       subgraph_ptr,
@@ -882,7 +921,7 @@ Error defineFullyConnectedNode(
       fc_input1,
       fc_filter,
       fc_bias,
-      fc_compute_output,
+      fc_result,
       graph_node->flags());
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
@@ -891,18 +930,7 @@ Error defineFullyConnectedNode(
       node->debug_handle(),
       xnn_status_to_string(status));
 
-  if (needs_bf16_output_convert) {
-    xnn_status cs = xnn_define_convert(
-        subgraph_ptr, fc_compute_output, fc_output, /*flags=*/0);
-    ET_CHECK_OR_RETURN_ERROR(
-        cs == xnn_status_success,
-        Internal,
-        "Failed to define bf16 output convert for linear node %i: %s",
-        node->debug_handle(),
-        xnn_status_to_string(cs));
-  }
-
-  return Error::Ok;
+  return convertBf16MatMulResult(subgraph_ptr, fc_result, fc_output, node);
 };
 
 /*
@@ -1702,15 +1730,24 @@ Error defineBatchMatrixMultiplyNode(
     const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
     const NodePtr node,
     const fb_xnnpack::XNNGraph* graph) noexcept {
-  MAYBE_UNUSED(graph);
-
   auto graph_node = node->xnode_union_as_XNNBatchMatrixMultiply();
   REMAP_ID(remapped_ids, graph_node->input1_id(), bmm_in1);
   REMAP_ID(remapped_ids, graph_node->input2_id(), bmm_in2);
   REMAP_ID(remapped_ids, graph_node->output_id(), bmm_out);
 
+  uint32_t bmm_result = XNN_INVALID_VALUE_ID;
+  ET_CHECK_OK_OR_RETURN_ERROR(defineBf16MatMulResult(
+      subgraph_ptr,
+      graph,
+      graph_node->input1_id(),
+      graph_node->input2_id(),
+      graph_node->output_id(),
+      bmm_out,
+      node,
+      &bmm_result));
+
   xnn_status status = xnn_define_batch_matrix_multiply(
-      subgraph_ptr, bmm_in1, bmm_in2, bmm_out, graph_node->flags());
+      subgraph_ptr, bmm_in1, bmm_in2, bmm_result, graph_node->flags());
 
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
@@ -1719,7 +1756,7 @@ Error defineBatchMatrixMultiplyNode(
       node->debug_handle(),
       xnn_status_to_string(status));
 
-  return Error::Ok;
+  return convertBf16MatMulResult(subgraph_ptr, bmm_result, bmm_out, node);
 }
 
 /*
diff --git a/backends/xnnpack/test/ops/test_bmm.py b/backends/xnnpack/test/ops/test_bmm.py
index a029738e771..49d0112a19c 100644
--- a/backends/xnnpack/test/ops/test_bmm.py
+++ b/backends/xnnpack/test/ops/test_bmm.py
@@ -7,7 +7,8 @@
 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
 
 
 class TestBMM(unittest.TestCase):
@@ -47,3 +48,23 @@ def test_fp32_bmm(self):
             torch.randn(2, 4, 6),
         )
         self._test_bmm(inputs)
+
+    def test_bf16_bmm(self):
+        inputs = (
+            torch.randn(2, 3, 4).to(torch.bfloat16),
+            torch.randn(2, 4, 6).to(torch.bfloat16),
+        )
+        (
+            Tester(self.BMM(), inputs)
+            .export()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    partitioners=[XnnpackPartitioner(enable_bf16=True)]
+                )
+            )
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .check_not(["executorch_exir_dialects_edge__ops_aten_bmm_default"])
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(atol=0.01, rtol=0.01)
+        )
diff --git a/backends/xnnpack/test/passes/test_batch_norm_fusion.py b/backends/xnnpack/test/passes/test_batch_norm_fusion.py
index 87c7fa397ef..519cf94243a 100644
--- a/backends/xnnpack/test/passes/test_batch_norm_fusion.py
+++ b/backends/xnnpack/test/passes/test_batch_norm_fusion.py
@@ -135,6 +135,24 @@ def test_fp32_linear_batch_norm_fusion(self):
                 .run_method_and_compare_outputs()
             )
 
+    def test_fp32_linear_batch_norm_on_3d_output(self):
+        """
+        BatchNorm1d normalizes dim 1. When the linear output is 3-D, dim 1 is not the
+        linear's output features, so the batch norm must not fold into the linear.
+        """
+        model = torch.nn.Sequential(torch.nn.Linear(8, 4), torch.nn.BatchNorm1d(4))
+        model[1].running_mean.uniform_(-1, 1)
+        model[1].running_var.uniform_(0.5, 2)
+        (
+            Tester(model.eval(), (torch.randn(2, 4, 8),))
+            .export()
+            .to_edge_transform_and_lower()
+            .check_count({self.bn_name: 0})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     def test_fp32_conv3d_batch_norm_doesnt_partition(self):
         """
         Conv3d is not currently supported by XNNPACK. We also don't support standalone
diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py
index 894fab4098f..be775e5d784 100644
--- a/backends/xnnpack/test/test_xnnpack_partitioner.py
+++ b/backends/xnnpack/test/test_xnnpack_partitioner.py
@@ -19,6 +19,14 @@
 from torch.export import export
 
 
+def _ops_outside_delegates(program):
+    return [
+        node.target.__name__
+        for node in program.exported_program().graph.nodes
+        if node.op == "call_function" and hasattr(node.target, "__name__")
+    ]
+
+
 class TestXnnpackPartitioner(unittest.TestCase):
     """Test cases for XnnpackPartitioner functionality and deprecation warnings."""
 
@@ -161,3 +169,131 @@ def example_inputs(self):
         fwd2_et = executorch_module.run_method("forward_2", example_inputs)
         self.assertTrue(torch.allclose(fwd1_eager, fwd1_et[0], 1e-3))
         self.assertTrue(torch.allclose(fwd2_eager, fwd2_et[0], 1e-3))
+
+    def test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16(self):
+        """
+        XNNPACK refuses these ops on bf16 tensors when the program loads, so even
+        with enable_bf16 their bf16 nodes must run on the portable kernels. The
+        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)
+
+        class Softmax(torch.nn.Module):
+            def forward(self, x):
+                return torch.softmax(x, -1)
+
+        def lower(module, x):
+            return to_edge_transform_and_lower(
+                export(module, (x,)),
+                partitioner=[XnnpackPartitioner(enable_bf16=True)],
+            )
+
+        modules = {
+            "avg_pool2d": torch.nn.AvgPool2d(2),
+            "batch_norm": torch.nn.BatchNorm2d(4),
+            "conv_batch_norm": torch.nn.Sequential(
+                torch.nn.Conv2d(4, 4, 3), torch.nn.BatchNorm2d(4)
+            ),
+            "max_pool2d": torch.nn.MaxPool2d(2),
+            "mean": Mean(),
+            "softmax": Softmax(),
+            "upsample_bilinear2d": torch.nn.Upsample(scale_factor=2, mode="bilinear"),
+        }
+        x = torch.randn(1, 4, 8, 8)
+        for name, module in modules.items():
+            with self.subTest(name):
+                module.eval()
+                self.assertIn(
+                    "executorch_call_delegate", _ops_outside_delegates(lower(module, x))
+                )
+                module, bf16_x = module.to(torch.bfloat16), x.to(torch.bfloat16)
+                program = lower(module, bf16_x)
+                self.assertNotIn(
+                    "executorch_call_delegate", _ops_outside_delegates(program)
+                )
+                output = _load_for_executorch_from_buffer(
+                    program.to_executorch().buffer
+                ).forward((bf16_x,))[0]
+                torch.testing.assert_close(output, module(bf16_x))
+
+    def test_enable_bf16_fuses_batch_norm_into_a_linear(self):
+        """
+        XNNPACK cannot define a bf16 BatchNorm on its own, but one that follows a
+        linear folds into it, so the whole bf16 block still runs in XNNPACK. It
+        does not fold when the linear output has another user, so then it stays
+        on the portable kernels.
+        """
+
+        class LinearBatchNorm(torch.nn.Module):
+            def __init__(self, reuse_linear_output):
+                super().__init__()
+                self.linear = torch.nn.Linear(8, 4)
+                self.batch_norm = torch.nn.BatchNorm1d(4)
+                self.batch_norm.running_mean.uniform_(-1, 1)
+                self.batch_norm.running_var.uniform_(0.5, 2)
+                self.reuse_linear_output = reuse_linear_output
+
+            def forward(self, x):
+                y = self.linear(x)
+                z = self.batch_norm(y)
+                return z + y if self.reuse_linear_output else z
+
+        x = torch.randn(2, 8, dtype=torch.bfloat16)
+        delegated = {"executorch_call_delegate", "getitem"}
+        for reuse_linear_output, ops_outside in (
+            (False, delegated),
+            (True, delegated | {"aten._native_batch_norm_legit_no_training.default"}),
+        ):
+            with self.subTest(reuse_linear_output=reuse_linear_output):
+                module = LinearBatchNorm(reuse_linear_output).eval().to(torch.bfloat16)
+                program = to_edge_transform_and_lower(
+                    export(module, (x,)),
+                    partitioner=[XnnpackPartitioner(enable_bf16=True)],
+                )
+                self.assertEqual(set(_ops_outside_delegates(program)), ops_outside)
+                output = _load_for_executorch_from_buffer(
+                    program.to_executorch().buffer
+                ).forward((x,))[0]
+                torch.testing.assert_close(output, module(x), atol=1e-2, rtol=1e-2)
+
+    def test_enable_bf16_runs_an_attention_block(self):
+        """
+        The linear layers and both batch matrix multiplications of a bf16 attention
+        block run in XNNPACK, and the program loads and matches eager bf16.
+        """
+
+        class Attention(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.q = torch.nn.Linear(16, 16)
+                self.k = torch.nn.Linear(16, 16)
+                self.v = torch.nn.Linear(16, 16)
+                self.o = torch.nn.Linear(16, 16, bias=False)
+
+            def forward(self, x):
+                scores = self.q(x) @ self.k(x).transpose(1, 2) * 0.25
+                return self.o(torch.softmax(scores, dim=-1) @ self.v(x))
+
+        module = Attention().eval().to(torch.bfloat16)
+        x = torch.randn(2, 8, 16, dtype=torch.bfloat16)
+        program = to_edge_transform_and_lower(
+            export(module, (x,)), partitioner=[XnnpackPartitioner(enable_bf16=True)]
+        )
+        ops = _ops_outside_delegates(program)
+        # The expand copies that matmul adds split the block into 3 delegates.
+        self.assertEqual(ops.count("executorch_call_delegate"), 3)
+        self.assertIn("aten._softmax.default", ops)
+        matmuls = {
+            "aten.addmm.default",
+            "aten.bmm.default",
+            "aten.linear.default",
+            "aten.mm.default",
+        }
+        self.assertEqual(set(ops) & matmuls, set())
+        output = _load_for_executorch_from_buffer(
+            program.to_executorch().buffer
+        ).forward((x,))[0]
+        torch.testing.assert_close(output, module(x), atol=1e-2, rtol=1e-2)