Arm backend: Add docstrings to operator_support/tosa_supported_operators (#15777)

This file contains the base class `SupportedTOSAOperatorCheck`, which
all operator support check inherit from. The documentation for it is
applicable to those derived classes as well. The derived classes should
document members etc. that are not part of the base class.

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>
Co-authored-by: Zingo Andersen <zingo.andersen@arm.com>

Author: Sebastian-Larsson

backends/arm/operator_support/tosa_supported_operators.py

@@ -2,6 +2,12 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+"""Provide operator-support checks and registries for TOSA delegation.
+
+Define a base check class, a registry/dispatcher, and several generic checks
+used by the TOSA partitioner to decide if FX nodes are eligible for delegation.
+
+"""
 
 
 import itertools
@@ -46,31 +52,65 @@
 
 
 class SupportedTOSAOperatorCheck(OperatorSupportBase):
-    """
-    Supported OP for TOSA lowering
+    """Provide a base operator-support check for TOSA lowering.
+
+    Subclasses should implement :py:meth:`is_node_tosa_supported` and declare
+    the class attributes below to indicate what they support.
+
+    Attributes:
+        targets (list[OpOverload]): Operator overloads supported by this
+            check.
+        tosa_specs (list[TosaSpecification]): TOSA specs where the check is
+            applicable.
+
     """
 
     def __init__(self, tosa_spec: TosaSpecification, reporter: WhyNoPartitionReporter):
+        """Initialize the check with a TOSA spec and reporter.
+
+        Args:
+            tosa_spec (TosaSpecification): Active TOSA specification.
+            reporter (WhyNoPartitionReporter): Reporter for rejection reasons.
+
+        """
         self.tosa_spec = tosa_spec
         self.reporter = reporter
 
-    # Should be populated by subclass implementation
+    # Class attributes populated by subclasses
     tosa_specs: list[TosaSpecification] = []
     targets: list[str] = []
 
     @final
     def is_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
+        """Return True if the node matches targets and subclass-specific checks.
+
+        Args:
+            submodules (typing.Mapping[str, torch.nn.Module]): Exported program
+                modules.
+            node (fx.Node): Node to evaluate.
+
+        Returns:
+            bool: True if both the target and TOSA-specific checks pass.
+
+        """
         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.
+        """Check if the node is lowerable under the given TOSA spec.
+
+        Args:
+            node (fx.Node): FX node to check.
+            tosa_spec (TosaSpecification): Active TOSA specification.
+
+        Returns:
+            bool: True if supported; otherwise, False.
+
         """
         raise NotImplementedError("SupportedTOSAOperatorCheck must be extended.")
 
@@ -83,10 +123,15 @@ def is_node_tosa_supported(
 
 
 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.
+    """Register an operator-support checker for one or more TOSA specs.
+
+    Decorate subclasses of :py:class:`SupportedTOSAOperatorCheck` so they are
+    picked up by the factory and partitioner for the specs declared in their
+    ``tosa_specs`` class attribute.
+
+    Args:
+        checker (Type[SupportedTOSAOperatorCheck]): Checker class to register.
+
     """
     for tosa_spec in checker.tosa_specs:
         _tosa_spec_support[tosa_spec].append(checker)
@@ -96,6 +141,15 @@ def register_tosa_support_check(checker: Type[SupportedTOSAOperatorCheck]):
 def get_registered_tosa_support_checks(
     tosa_spec: TosaSpecification,
 ) -> list[Type[SupportedTOSAOperatorCheck]]:
+    """Get all registered operator-support checkers for a given spec.
+
+    Args:
+        tosa_spec (TosaSpecification): TOSA spec to query.
+
+    Returns:
+        list[Type[SupportedTOSAOperatorCheck]]: Registered checker classes.
+
+    """
     if tosa_spec not in _tosa_spec_support:
         raise RuntimeError(
             f"TOSA specification not valid: {tosa_spec} not in {list(_tosa_spec_support.keys())}"
@@ -110,8 +164,21 @@ def tosa_support_factory(
     reporter: WhyNoPartitionReporter,
     additional_checks: Optional[Sequence[OperatorSupportBase]] = None,
 ) -> OperatorSupportBase:
-    """Generates an OperatorSupport class depending on the given `tosa_spec`.
-    Additional checks can be supplied to avoid partitioning additional nodes.
+    """Create an OperatorSupport composite for a TOSA spec.
+
+    Combine profile-specific positive checks, registered operator checks, and
+    negative checks into a single :py:class:`OperatorSupportBase` chain.
+
+    Args:
+        tosa_spec (TosaSpecification): Active TOSA specification.
+        exported_program (ExportedProgram): Program context for checks.
+        reporter (WhyNoPartitionReporter): Reporter for rejections.
+        additional_checks (Optional[Sequence[OperatorSupportBase]]): Extra
+            negative checks to apply.
+
+    Returns:
+        OperatorSupportBase: Composite checker for the given spec.
+
     """
     # Postive checks: Add nodes to partitioning
     positive_checks: list[OperatorSupportBase] = [
@@ -158,37 +225,45 @@ def tosa_support_factory(
 
 
 class TOSAProINTSupportList(OperatorSupportBase):
-    """
-    TOSA_PRO_INT_SupportList:
-        Ops supported in INT profile via native TOSA ops, decomposition/transformation, pre-compute, or TableOps.
-        Note that ops supported via pre-quantization decompositions are not included here.
+    """Provide the INT profile support list for TOSA.
+
+    TOSA_PRO_INT_SupportList enumerates ops supported in the INT profile via
+    native TOSA ops, decompositions, pre-compute steps, or TableOps.
+
+    Note:
+        Ops supported via pre-quantization decompositions are not included
+        here.
+
     """
 
     def is_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
-
+        """Return True if the node is in the INT profile support list."""
         return node.op == "call_function" and node.target in TOSA_PRO_INT_SupportList
 
 
 class TOSAProFPSupportList(OperatorSupportBase):
-    """
-    TOSA_PRO_FP_SupportList:
-        Ops supported in FP profile via native TOSA ops, decomposition/transformation, pre-compute
+    """Provide the FP profile support list for TOSA.
+
+    Includes ops supported natively, via decomposition/transformation, and pre-
+    compute.
+
     """
 
     def is_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
-
+        """Return True if the node is in the FP profile support list."""
         return node.op == "call_function" and node.target in TOSA_PRO_FP_SupportList
 
 
 class CheckProperQuantization(OperatorSupportBase):
-    """
-    For targeted nodes, check that it has been quantized as expected. In most cases this means that a pair of quantize
-    and dequantize nodes surrounds the node. This is neccessary for table operators and operators that need to rescale
-    activations.
+    """Ensure targeted nodes are properly quantized.
+
+    Verify that a pair of quantize/dequantize nodes surrounds targeted ops so
+    rescaling and table operators behave correctly.
+
     """
 
     targeted_ops = (
@@ -214,13 +289,28 @@ class CheckProperQuantization(OperatorSupportBase):
     )
 
     def __init__(self, reporter: WhyNoPartitionReporter):
+        """Initialize the check with a reporter."""
         self.reporter = reporter
 
     def _is_matmul_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ):
-        """
-        Find the matmul source partition containing this node and check that all its inputs and outputs are quantized.
+        """Check quantization for decomposed matmul partitions.
+
+        Handles an edge case where the quantized pipeline
+        `dq -> torch.matmul/operator.matmul -> q` decomposes into
+        `dq -> expand -> view -> aten.mm -> view -> q`.
+
+        Args:
+            submodules (Mapping[str, torch.nn.Module]): Map of child modules to
+                inspect for matmul partitions.
+            node (fx.Node): Node that should belong to a quantized matmul
+                partition.
+
+        Returns:
+            bool: True if the matched partition uses quantized inputs and
+                outputs.
+
         """
         for graph_module in submodules.values():
             graph_module = typing.cast(fx.GraphModule, graph_module)
@@ -269,6 +359,12 @@ def _is_matmul_node_supported(
     def is_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
+        """Return True if the node passes constant-cast and multi-output checks.
+
+        Ensures decomposition-specific matmul partitions keep quantized inputs
+        and outputs.
+
+        """
         output_quantized = False
         input_quantized = False
         if node.target not in self.targeted_ops:
@@ -320,21 +416,22 @@ def is_node_supported(
 
 
 class CheckInt64InputsAndOutputs(OperatorSupportBase):
-    """TOSA does not support int64 tensors so in general, ops with int64 inputs or outputs should not be partitioned.
-    There are however some exceptions:
-        - Nodes with int64 output can be partitioned if they are constant, within int32,
-            and all users cast to something else. In this case, the int64 tensor can safely be cast to int32 AOT.
-        - Nodes with int64 output can be partitioned if all users are getitem with non-int64 output.
-            In this case, there are multiple outputs and the int64 ones are not used.
-        - Nodes with int64 inputs can be partitioned if the inputs are constant placeholders, or constant
-            ops fulfilling the criteria above.
-    Note that we don't check placeholders here, they are partitioned based on whether their users are partitioned
-    or not.
+    """Reject general int64 tensors while allowing safe exceptions.
+
+    Exceptions are:
+        - Nodes with contant int64 output within int32 range that are cast away
+          from int64 by all users.
+        - Int64 output where all users are getitem nodes with non-int64 outputs.
+          In this case there are multiple outputs and the int64 output is unused.
+        - Nodes where all inputs are int64 constant placeholders or constant ops
+          that fulfill the above exceptions.
+
     """
 
     def __init__(
         self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter
     ):
+        """Initialize the check with program context and reporter."""
         self.input_names = [
             spec.arg.name
             for spec in exported_program.graph_signature.input_specs
@@ -356,6 +453,7 @@ def inside_int32_bounds(self, node: torch.fx.Node) -> bool:
     def is_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
+        """Return True when int64 use is absent or safe per exceptions."""
         if is_submodule_node(node):
             return True
         vals = node.meta["val"]
@@ -427,16 +525,23 @@ def is_node_supported(
 
 
 class CheckFloat64Inputs(OperatorSupportBase):
+    """Reject nodes with float64 inputs.
+
+    Useful as a negative check for specs that do not allow float64.
+
+    """
 
     def __init__(
         self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter
     ):
+        """Initialize the check with program context and reporter."""
         self.reporter = reporter
         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."""
         if is_submodule_node(node):
             return True
         for input_node in (
@@ -455,16 +560,18 @@ def is_node_supported(
 
 
 class RankCheck(OperatorSupportBase):
-    """Makes sure that nodes with input or output tensors with rank > max_rank are not partitioned"""
+    """Reject nodes with rank greater than ``max_rank``."""
 
     def __init__(self, reporter: WhyNoPartitionReporter, max_rank: int):
+        """Initialize the check with a reporter and maximum rank."""
         self.reporter = reporter
         self.max_rank = max_rank
         super().__init__()
 
     def is_node_supported(
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
+        """Return True if input/output tensor ranks are within the limit."""
         if is_submodule_node(node):
             return True
         input_nodes = (
@@ -509,20 +616,36 @@ def is_node_supported(
 
 
 class CondSupported(OperatorSupportBase):
-    """Checks whether the cond operator, and it's submodule args, should be partitioned."""
+    """Check whether cond operator and submodule args should be partitioned.
+
+    Applies control-flow extension constraints before allowing delegation.
+
+    """
 
     def __init__(
         self,
         exported_program: ExportedProgram,
         tosa_spec: TosaSpecification,
         reporter: WhyNoPartitionReporter,
     ):
+        """Initialize conditional support checks for TOSA delegation.
+
+        Args:
+            exported_program (ExportedProgram): Program containing the cond
+                submodules to inspect.
+            tosa_spec (TosaSpecification): TOSA specification used to validate
+                supported operators.
+            reporter (WhyNoPartitionReporter): Reporter that records rejection
+                reasons for unsupported nodes.
+
+        """
         self.exported_program = exported_program
         self.reporter = reporter
         self.tosa_spec = tosa_spec
         super().__init__()
 
     def _fully_partitioned(self, submodule: fx.GraphModule) -> bool:
+        """Check whether all call_function nodes share one delegation tag."""
         partition_tag = None
         for submodule_node in submodule.graph.nodes:
             if submodule_node.op == "call_function":
@@ -546,8 +669,10 @@ def _fully_partitioned(self, submodule: fx.GraphModule) -> bool:
         return True
 
     def _cond_submodules_fully_partitioned(self, node: fx.Node) -> bool:
-        """Returns whether the submodule arguments to a cond node were fully partitioned.
-        Updates "val" meta of the submodules if they are.
+        """Determine whether cond submodules were fully partitioned.
+
+        Update the "val" meta of the submodules when they are partitioned.
+
         """
         cond_submodules = (
             (
@@ -570,6 +695,7 @@ def _cond_submodules_fully_partitioned(self, node: fx.Node) -> bool:
     def is_node_supported(  # noqa: C901
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
     ) -> bool:
+        """Check whether a cond node and its submodules are partitionable."""
         if is_submodule_node(node):
             if not isinstance(self.tosa_spec, Tosa_1_00):
                 self.reporter.report_reject(