Arm backend: Extend fp64 partition check to more dtypes (#16703)

Check both inputs and outputs.
Specifically add check for bf16 operators.
Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

backends/arm/operator_support/tosa_supported_operators.py

@@ -1,4 +1,4 @@
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
+# Copyright 2024-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -282,7 +282,6 @@ def tosa_support_factory(
     # Negative checks: Remove nodes from partitioning
     negative_checks: list[OperatorSupportBase] = [
         CheckInt64InputsAndOutputs(exported_program, reporter),
-        CheckFloat64Inputs(exported_program, reporter),
         RankCheck(reporter, max_rank=MAX_RANK),
         *[
             reporter.wrap_check(check, f"Rejected by {check.__class__.__name__}")
@@ -293,6 +292,15 @@ def tosa_support_factory(
     if not tosa_spec.support_float():
         negative_checks.append(CheckArmQuantized(reporter))
         negative_checks.append(CheckProperQuantization(reporter))
+
+    disallowed_dtypes = [torch.float64]
+    if not tosa_spec.support_extension("bf16"):
+        disallowed_dtypes.append(torch.bfloat16)
+    negative_checks.append(
+        CheckDtypeInputsAndOutputs(
+            exported_program, reporter, disallowed_dtypes, tosa_spec
+        )
+    )
     if tosa_spec.is_U55_subset:
         negative_checks.append(EthosU55NotSupported(reporter))
         negative_checks.append(EthosU55DtypeSupport(reporter))
@@ -657,24 +665,26 @@ def is_node_supported(
         return True
 
 
-class CheckFloat64Inputs(OperatorSupportBase):
-    """Reject nodes with float64 inputs.
-
-    Useful as a negative check for specs that do not allow float64.
-
-    """
+class CheckDtypeInputsAndOutputs(OperatorSupportBase):
+    """Reject nodes with at least one disallowed dtype on inputs or outputs."""
 
     def __init__(
-        self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter
+        self,
+        exported_program: ExportedProgram,
+        reporter: WhyNoPartitionReporter,
+        disallowed_dtypes: list[torch.dtype],
+        tosa_spec: TosaSpecification,
     ):
         """Initialize the check with program context and reporter."""
         self.reporter = reporter
+        self.disallowed_dtypes = disallowed_dtypes
+        self.tosa_spec = tosa_spec
         super().__init__()
 
     def is_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
-        """Return True if no float64 inputs are present."""
+        """Return True if no disallowed dtypes are present on inputs or outputs."""
         if is_submodule_node(node):
             return True
         for input_node in (
@@ -683,10 +693,29 @@ def is_node_supported(
             if input_node.op != "get_attr"
         ):
             tensor = get_first_fake_tensor(input_node)
-            if tensor.dtype == torch.float64:
+            if tensor.dtype in self.disallowed_dtypes:
+                self.reporter.report_reject(
+                    node,
+                    f"Had {tensor.dtype} input {input_node.name} that is not supported by {self.tosa_spec}.",
+                )
+                return False
+
+        meta_val = node.meta["val"]
+        if isinstance(
+            meta_val, (Sequence, torch.fx.immutable_collections.immutable_list)
+        ):
+            outputs = meta_val
+        else:
+            outputs = (meta_val,)
+
+        for output in outputs:
+            if (
+                isinstance(output, FakeTensor)
+                and output.dtype in self.disallowed_dtypes
+            ):
                 self.reporter.report_reject(
                     node,
-                    f"Had float64 input {input_node.name} that couldn't be handled.",
+                    f"Had {output.dtype} output that is not supported by {self.tosa_spec}.",
                 )
                 return False
         return True