[XNNPACK] Lower fp32 <-> bf16 casts to XNNPACK converts

backends/xnnpack/partition/config/generic_node_configs.py

@@ -667,7 +667,8 @@ class ToDimOrderCopyConfig(GenericNodePartitionerConfig):
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
-        Only support dim order conversion partitioning, not DType conversions
+        Partitions dim order conversions, and fp32 <-> fp16 and fp32 <-> bf16
+        dtype casts that keep the dim order.
         """
         if not self.check_common_constraints(node, ep):
             return False
@@ -678,19 +679,26 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         output_dtype = node.meta["val"].dtype
 
         if input_dtype != output_dtype:
-            # A dtype cast is lowered to an XNNPACK convert node, as long as
-            # the op does not change the dim order at the same time.
-            supported_cast = {input_dtype, output_dtype} <= {
-                torch.float32,
-                torch.float16,
-            }
+            # A dtype cast is lowered to an XNNPACK convert node. XNNPACK has
+            # no fp16 <-> bf16 convert, and bf16 also needs enable_bf16, which
+            # check_common_constraints enforces.
+            supported_cast = {input_dtype, output_dtype} in (
+                {torch.float32, torch.float16},
+                {torch.float32, torch.bfloat16},
+            )
+            if not supported_cast:
+                why(
+                    node,
+                    reason=f"dtype conversion from {input_dtype} to {output_dtype} is not supported",
+                )
+                return False
             same_dim_order = list(input_node.meta["val"].dim_order()) == list(
                 node.meta["val"].dim_order()
             )
-            if not (supported_cast and same_dim_order):
+            if not same_dim_order:
                 why(
                     node,
-                    reason=f"dtype conversion from {input_dtype} to {output_dtype} is not supported",
+                    reason=f"dtype conversion from {input_dtype} to {output_dtype} cannot also change the dim order",
                 )
                 return False
 

backends/xnnpack/test/ops/test_to_copy_dtype.py

@@ -7,11 +7,14 @@
 import unittest
 
 import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.backends.xnnpack.test.tester import Tester
+from executorch.backends.xnnpack.test.tester.tester import ToEdgeTransformAndLower
+from executorch.export.utils import is_fbcode
 
 
 class TestToCopyDtype(unittest.TestCase):
-    """fp32 <-> fp16 casts are lowered to XNNPACK convert nodes."""
+    """fp32 <-> fp16 and fp32 <-> bf16 casts are lowered to XNNPACK convert nodes."""
 
     def setUp(self):
         torch._dynamo.reset()
@@ -35,11 +38,15 @@ def forward(self, x):
         "executorch_exir_dialects_edge__ops_aten__to_copy_default",
     ]
 
-    def _run(self, module, inputs):
+    def _run(self, module, inputs, enable_bf16=False):
         (
             Tester(module.eval(), inputs)
             .export()
-            .to_edge_transform_and_lower()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    partitioners=[XnnpackPartitioner(enable_bf16=enable_bf16)]
+                )
+            )
             .check_not(self.cast_ops)
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .to_executorch()
@@ -100,3 +107,96 @@ def forward(self, x):
                 return x.to(memory_format=torch.channels_last).to(torch.float16)
 
         self._run(ChannelsLastToHalf(), (torch.randn(1, 3, 4, 4),))
+
+    class Bf16Inner(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(16, 8).to(torch.bfloat16)
+
+        def forward(self, x):
+            return self.linear(x.to(torch.bfloat16)).to(torch.float32)
+
+    class Bf16Upcast(torch.nn.Module):
+        """A bf16 model that computes a step in fp32 and casts back."""
+
+        def forward(self, x):
+            return (x.to(torch.float32) * 2.0).to(torch.bfloat16)
+
+    class Bf16SoftmaxUpcast(torch.nn.Module):
+        """Softmax computed in fp32 on a bf16 input, cast back to bf16."""
+
+        def forward(self, x):
+            return torch.softmax(x, -1, dtype=torch.float32).to(torch.bfloat16)
+
+    class ToBf16(torch.nn.Module):
+        def forward(self, x):
+            return x.to(torch.bfloat16)
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_fp32_io_around_bf16_linear(self):
+        self._run(self.Bf16Inner(), (torch.randn(4, 16),), enable_bf16=True)
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_upcast(self):
+        self._run(
+            self.Bf16Upcast(),
+            (torch.randn(2, 5, dtype=torch.bfloat16),),
+            enable_bf16=True,
+        )
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_bf16_softmax_upcast(self):
+        self._run(
+            self.Bf16SoftmaxUpcast(),
+            (torch.randn(2, 16, 32, dtype=torch.bfloat16),),
+            enable_bf16=True,
+        )
+
+    @unittest.skipIf(is_fbcode(), "skipping for fbcode")
+    def test_fp32_to_bf16_rounds_like_pytorch(self):
+        """
+        XNNPACK's convert must round fp32 to bf16 exactly as PyTorch does,
+        including values halfway between two bf16 numbers.
+        """
+        halfway = torch.tensor(
+            [1 + 2**-8, 1 + 3 * 2**-8, -(1 + 3 * 2**-8), 2**-100 * (1 + 3 * 2**-8)]
+        )
+        x = torch.cat([torch.randn(4096), halfway, torch.tensor([3.0e38, 0.0, -0.0])])
+        (
+            Tester(self.ToBf16().eval(), (x,))
+            .export()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    partitioners=[XnnpackPartitioner(enable_bf16=True)]
+                )
+            )
+            .check_not(self.cast_ops)
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(inputs=(x,), atol=0, rtol=0)
+        )
+
+    def test_bf16_cast_not_partitioned_without_enable_bf16(self):
+        (
+            Tester(self.Bf16Upcast().eval(), (torch.randn(2, 5, dtype=torch.bfloat16),))
+            .export()
+            .to_edge_transform_and_lower()
+            .check_count({self.cast_ops[0]: 2})
+        )
+
+    def test_fp16_bf16_cast_not_partitioned(self):
+        class HalfToBf16(torch.nn.Module):
+            def forward(self, x):
+                return x.to(torch.bfloat16)
+
+        (
+            Tester(HalfToBf16().eval(), (torch.randn(2, 5, dtype=torch.float16),))
+            .export()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    partitioners=[XnnpackPartitioner(enable_bf16=True)]
+                )
+            )
+            .check_not(["torch.ops.higher_order.executorch_call_delegate"])
+        )