Refactor TOSA support to use chain flow

This seems to match the intention of OperatorSupportBase
and increases the flexibility of our support flow.
New checks if different kinds are simply added
to `tosa_support_factory`

Signed-off-by: Erik Lundell <[email protected]>
Change-Id: I4f5f1669959f35c5e2770da396a352e1730cb3e0

Author: Erik-Lundell

backends/arm/arm_partitioner.py

@@ -14,7 +14,7 @@
     ArmBackend,
 )  # usort: skip
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
-    TOSASupportedOperators,
+    tosa_support_factory,
 )
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.exir.backend.compile_spec_schema import CompileSpec
@@ -72,7 +72,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 
         capability_partitioner = CapabilityBasedPartitioner(
             exported_program.graph_module,
-            TOSASupportedOperators(tosa_spec),
+            tosa_support_factory(tosa_spec),
             allows_single_node_partition=True,
         )
         partition_list = capability_partitioner.propose_partitions()

backends/arm/operator_support/convolution_support.py

@@ -24,7 +24,7 @@ class ConvolutionSupported(SupportedTOSAOperatorCheck):
         TosaSpecification.create_from_string("TOSA-0.80+MI"),
     ]
 
-    def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
+    def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
 
         # Not implemented
         transposed = cast(bool, node.args[6])

backends/arm/operator_support/pool_2d_support.py

@@ -43,7 +43,7 @@ class AvgPool2dSupported(SupportedTOSAOperatorCheck):
         TosaSpecification.create_from_string("TOSA-0.80+MI"),
     ]
 
-    def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
+    def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
         if not (isinstance(tosa_spec, Tosa_0_80) and tosa_spec.is_U55_subset):
             return True
 
@@ -73,7 +73,7 @@ class MaxPool2dSupported(SupportedTOSAOperatorCheck):
         TosaSpecification.create_from_string("TOSA-0.80+MI"),
     ]
 
-    def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
+    def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
         if not (isinstance(tosa_spec, Tosa_0_80) and tosa_spec.is_U55_subset):
             return True
 

backends/arm/operator_support/reduce_sum_support.py

@@ -23,7 +23,7 @@ class SumSupported(SupportedTOSAOperatorCheck):
         TosaSpecification.create_from_string("TOSA-0.80+MI"),
     ]
 
-    def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
+    def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
         if not (isinstance(tosa_spec, Tosa_0_80) and tosa_spec.is_U55_subset):
             return True
 

backends/arm/operator_support/right_shift_support.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-2025 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.
@@ -29,7 +29,7 @@ class RightShiftSupported(SupportedTOSAOperatorCheck):
         TosaSpecification.create_from_string("TOSA-0.80+MI"),
     ]
 
-    def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
+    def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
 
         # TODO MLETORCH-525 Remove warning
         if isinstance(tosa_spec, Tosa_0_80) and tosa_spec.is_U55_subset:

backends/arm/operator_support/to_copy_support.py

@@ -70,7 +70,9 @@ def _merge_supported_types(
     )
     POSSIBLE_TYPE_CONVERSIONS = {torch.int64: torch.int32}
 
-    def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification) -> bool:
+    def is_node_tosa_supported(
+        self, node: fx.Node, tosa_spec: TosaSpecification
+    ) -> bool:
         assert node.target in self.targets
 
         if tosa_spec not in self.tosa_specs:

backends/arm/operator_support/tosa_supported_operators.py

@@ -6,70 +6,77 @@
 # pyre-unsafe
 
 import operator
-from typing import Type
+from typing import final, Type
 
 import torch.fx as fx
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.exir.dialects._ops import ops as exir_ops
-from torch.fx.passes.operator_support import OperatorSupportBase
+from torch.fx.passes.operator_support import any_chain, OperatorSupportBase
 
 
-class SupportedTOSAOperatorCheck:
+class SupportedTOSAOperatorCheck(OperatorSupportBase):
     """
     Supported OP for TOSA lowering
     """
 
+    def __init__(self, tosa_spec: TosaSpecification):
+        self.tosa_spec = tosa_spec
+
     # Should be populated by subclass implementation
     tosa_specs: list[TosaSpecification] = []
     targets: list[str] = []
 
-    def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification) -> bool:
+    @final
+    def is_node_supported(self, submodules, node: fx.Node) -> bool:
+        if node.target not in self.targets:
+            return False
+        return self.is_node_tosa_supported(node, self.tosa_spec)
+
+    def is_node_tosa_supported(
+        self, node: fx.Node, tosa_spec: TosaSpecification
+    ) -> bool:
         """
         Checks if the fx.Node node is lowerable using the TOSA specification defined by tosa_spec.
-        To be implemented by subclasses targeting
         """
-        raise NotImplementedError("NodeVisitor must be extended.")
+        raise NotImplementedError("SupportedTOSAOperatorCheck must be extended.")
 
 
 # container for all SupportedTosaOperatorCheck classes
-_tosa_spec_dicts: dict[
-    TosaSpecification, dict[str, Type[SupportedTOSAOperatorCheck]]
-] = {
-    TosaSpecification.create_from_string("TOSA-0.80+BI"): {},
-    TosaSpecification.create_from_string("TOSA-0.80+MI"): {},
+_tosa_spec_support: dict[TosaSpecification, list[Type[SupportedTOSAOperatorCheck]]] = {
+    TosaSpecification.create_from_string("TOSA-0.80+BI"): [],
+    TosaSpecification.create_from_string("TOSA-0.80+MI"): [],
 }
 
 
-def register_tosa_support_check(checker):
+def register_tosa_support_check(checker: Type[SupportedTOSAOperatorCheck]):
     """
     Decorator to mark a subclass implmentation of SupportedTosaOperatorCheck
     to be registered for checking if a torch.fx.Node is lowerable given
     a TOSA specification.
     """
     for tosa_spec in checker.tosa_specs:
-        for target in checker.targets:
-            _tosa_spec_dicts[tosa_spec][target] = checker
+        _tosa_spec_support[tosa_spec].append(checker)
     return checker
 
 
 def get_registered_tosa_support_checks(
     tosa_spec: TosaSpecification,
-) -> dict[str, SupportedTOSAOperatorCheck]:
+) -> list[Type[SupportedTOSAOperatorCheck]]:
 
-    if tosa_spec not in _tosa_spec_dicts:
+    if tosa_spec not in _tosa_spec_support:
         raise RuntimeError
 
-    tosa_support_checks = {}
-    for target, tosa_check in _tosa_spec_dicts[tosa_spec].items():
-        tosa_support_checks[target] = tosa_check()
+    return _tosa_spec_support[tosa_spec]
 
-    return tosa_support_checks
 
+def tosa_support_factory(tosa_spec: TosaSpecification) -> OperatorSupportBase:
+    return any_chain(
+        BaseTOSASupportList(),
+        *(check(tosa_spec) for check in get_registered_tosa_support_checks(tosa_spec)),
+    )
 
-class TOSASupportedOperators(OperatorSupportBase):
-    def __init__(self, tosa_spec: TosaSpecification):
-        super().__init__()
-        self.tosa_spec = tosa_spec
+
+class BaseTOSASupportList(OperatorSupportBase):
 
     def is_node_supported(self, submodules, node: fx.Node) -> bool:
         supported = node.op == "call_function" and node.target in [
@@ -123,18 +130,4 @@ def is_node_supported(self, submodules, node: fx.Node) -> bool:
             exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
         ]
 
-        if not supported:
-            supported = self.is_node_supported_custom(node)
-
-        # Override partitioning based on pre partition passes
-        if "arm_override_partition" in node.meta:
-            supported = supported & node.meta["arm_override_partition"]
-            node.meta.pop("arm_override_partition")
-
         return supported
-
-    def is_node_supported_custom(self, node: fx.Node) -> bool:
-        tosa_checks = get_registered_tosa_support_checks(self.tosa_spec)
-        if node.target in tosa_checks.keys():
-            return tosa_checks[node.target].is_node_supported(node, self.tosa_spec)  # type: ignore[index]
-        return False