Keep unsupported BF16 and dynamic reshapes out of XNNPACK (pytorch#23021)

Leave BF16 convolution and Conv1d/unsqueeze reshapes with multiple
dynamic
dimensions to ExecuTorch's native kernels.

Backend regression tests export and execute Conv1d, Conv2d, and
unsqueeze
without a model dependency. They check BF16 fallback, dynamic batch and
length,
and continued delegation with one dynamic dimension against PyTorch
outputs.

Authored with OpenAI Codex.

Author: mergennachin

backends/xnnpack/partition/config/gemm_configs.py

@@ -40,6 +40,7 @@
 from executorch.exir.backend.utils import WhyNoPartition
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
+from torch.fx.experimental.symbolic_shapes import free_symbols
 from torch.fx.passes.utils.source_matcher_utils import (
     get_source_partitions,
     SourcePartition,
@@ -363,6 +364,12 @@ 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.
@@ -371,6 +378,17 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
             why(node, "Only support 1D + 2D Conv")
             return False  # Only support 1D + 2D Conv
 
+        # Conv1d lowering inserts reshapes with at most one inferred dimension.
+        if (
+            conv_dim == 1
+            and sum(
+                bool(free_symbols(d)) for d in get_input_node(node, 0).meta["val"].shape
+            )
+            > 1
+        ):
+            why(node, "Conv1d reshapes only support one dynamic dimension")
+            return False
+
         weight_quant_params = QuantParams.from_weights(kernel_node, ep)
         groups = cast(int, node.args[8])
         is_transpose = node.args[6]

backends/xnnpack/partition/config/generic_node_configs.py

@@ -34,6 +34,7 @@
 from executorch.exir.backend.utils import is_shape_dynamic, WhyNoPartition
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.export import ExportedProgram
+from torch.fx.experimental.symbolic_shapes import free_symbols
 
 logger = logging.getLogger(__name__)
 why = WhyNoPartition(logger=logger)
@@ -900,4 +901,8 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
             why(node, reason="unsqueeze_copy only supported on the trailing dimension")
             return False
 
+        if sum(bool(free_symbols(d)) for d in node.meta["val"].shape) > 1:
+            why(node, reason="XNNPACK reshape only supports one dynamic dimension")
+            return False
+
         return True

backends/xnnpack/test/ops/test_conv1d.py

@@ -22,6 +22,57 @@ class TestConv1d(unittest.TestCase):
     def setUp(self):
         torch._dynamo.reset()
 
+    def test_dynamic_conv1d(self):
+        for dynamic_dims in ((0,), (2,), (0, 2)):
+            with self.subTest(dynamic_dims=dynamic_dims):
+                shapes = {
+                    0: torch.export.Dim("batch", min=1, max=4),
+                    2: torch.export.Dim("length", min=5, max=16),
+                }
+                tester = (
+                    Tester(
+                        torch.nn.Conv1d(2, 2, 3).eval(),
+                        (torch.randn(2, 2, 11),),
+                        dynamic_shapes=({dim: shapes[dim] for dim in dynamic_dims},),
+                    )
+                    .export()
+                    .to_edge_transform_and_lower()
+                    .check_count(
+                        {
+                            "torch.ops.higher_order.executorch_call_delegate": int(
+                                len(dynamic_dims) == 1
+                            )
+                        }
+                    )
+                    .to_executorch()
+                    .serialize()
+                )
+                for batch, length in ((1, 5), (4, 16), (2, 7)):
+                    inputs = torch.randn(
+                        batch if 0 in dynamic_dims else 2,
+                        2,
+                        length if 2 in dynamic_dims else 11,
+                    )
+                    tester.run_method_and_compare_outputs(inputs=(inputs,))
+
+    def test_bf16_conv1d_fallback(self):
+        (
+            Tester(
+                torch.nn.Conv1d(2, 2, 3).eval().to(torch.bfloat16),
+                (torch.randn(2, 2, 11, dtype=torch.bfloat16),),
+            )
+            .export()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    partitioners=[XnnpackPartitioner(enable_bf16=True)]
+                )
+            )
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 0})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(atol=0.01, rtol=0.01)
+        )
+
     class Conv1d(torch.nn.Module):
         def __init__(self, dtype: torch.dtype = torch.float):
             groups = 1

backends/xnnpack/test/ops/test_conv2d.py

@@ -220,6 +220,24 @@ class TestConv2d(unittest.TestCase):
     def setUp(self):
         torch._dynamo.reset()
 
+    def test_bf16_conv2d_fallback(self):
+        (
+            Tester(
+                torch.nn.Conv2d(2, 2, 3).eval().to(torch.bfloat16),
+                (torch.randn(2, 2, 7, 7, dtype=torch.bfloat16),),
+            )
+            .export()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    partitioners=[XnnpackPartitioner(enable_bf16=True)]
+                )
+            )
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 0})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(atol=0.01, rtol=0.01)
+        )
+
     def _test(
         self,
         m: torch.nn.Module,

backends/xnnpack/test/ops/test_unsqueeze.py

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.backends.xnnpack.test.tester import Tester
+
+
+class TestUnsqueeze(unittest.TestCase):
+    def test_dynamic_trailing_unsqueeze(self):
+        class Unsqueeze(torch.nn.Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.dim = dim
+
+            def forward(self, x):
+                return x.unsqueeze(self.dim)
+
+        for dim in (-1, 2):
+            for dynamic_dims in (1, 2):
+                with self.subTest(dim=dim, dynamic_dims=dynamic_dims):
+                    shapes = {
+                        i: torch.export.Dim(f"dim_{i}", min=1, max=8)
+                        for i in range(dynamic_dims)
+                    }
+                    tester = (
+                        Tester(
+                            Unsqueeze(dim),
+                            (torch.randn(2, 3),),
+                            dynamic_shapes=(shapes,),
+                        )
+                        .export()
+                        .to_edge_transform_and_lower()
+                        .check_count(
+                            {
+                                "torch.ops.higher_order.executorch_call_delegate": int(
+                                    dynamic_dims == 1
+                                )
+                            }
+                        )
+                        .to_executorch()
+                        .serialize()
+                    )
+                    for batch, width in ((1, 1), (8, 8), (4, 2)):
+                        inputs = torch.randn(batch, width if dynamic_dims == 2 else 3)
+                        tester.run_method_and_compare_outputs(inputs=(inputs,))