Arm backend: Add docstrings to tosa/backend.py (#15552)

Signed-off-by: Sebastian Larsson <[email protected]>

Author: Sebastian-Larsson

backends/arm/tosa/backend.py

@@ -2,13 +2,19 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+"""Provide TOSA backend entry points for the Arm ExecuTorch integration.
 
+Implement the Ahead-of-Time (AoT) preprocessing path that lowers an
+``ExportedProgram`` to a TOSA flatbuffer using Arm's lowering pipeline. Use
+this module either as a standalone backend that produces a TOSA artifact or as
+part of a composed pipeline for hardware backends that consume TOSA as an
+intermediate form.
+
+Use ``TOSABackend.preprocess`` to return the serialized TOSA flatbuffer that
+subsequent stages (for example, JIT or hardware-specific compilers) consume.
+
+"""
 
-#
-# Main implementation of AoT flow to partition and preprocess for Arm target
-# backends. Converts via TOSA as an intermediate form supported by AoT and
-# JIT compiler flows.
-#
 import logging
 import tempfile
 from collections import deque
@@ -89,7 +95,16 @@ def _sort_key(t: Node) -> int:
 
 
 def arm_get_first_delegation_tag(graph_module) -> str:
-    """Get the first delegation tag from the graph_module or return empty string."""
+    """Return the first delegation tag from the FX graph.
+
+    Args:
+        graph_module: FX GraphModule produced by the Arm passes.
+
+    Returns:
+        str: The first non-empty delegation tag found on any node, or an empty
+        string if none is present.
+
+    """
     for node in graph_module.graph.nodes:
         tag = node.meta.get("delegation_tag")
         if tag:
@@ -101,10 +116,11 @@ def arm_get_first_delegation_tag(graph_module) -> str:
 
 @final
 class TOSABackend(BackendDetails):
-    """
-    BackendDetails subclass for lowering to TOSA.
-    Is used either by itself to get to a TOSA representation, or with composition
-    to be used as a separate step to target TOSA compliant hardware.
+    """Provide a backend for lowering programs to TOSA.
+
+    Use this class standalone to produce a TOSA representation, or as part of a
+    composed pipeline for hardware backends that consume TOSA.
+
     """
 
     @staticmethod
@@ -118,6 +134,31 @@ def _preprocess(  # noqa: C901
         edge_program: ExportedProgram,
         compile_spec: TosaCompileSpec,
     ) -> PreprocessResult:
+        """Lower an exported program to a TOSA flatbuffer.
+
+        Apply Arm transformation passes to ``edge_program``, then walk the
+        transformed FX graph to emit a TOSA graph via the serializer. When
+        requested in ``compile_spec``, write additional debug artifacts.
+
+        Args:
+            edge_program (ExportedProgram): Program to lower to TOSA.
+            compile_spec (List[CompileSpec]): Backend options. Recognized keys:
+                - output_format: Must be "tosa".
+                - tosa_spec: Target TOSA version/capabilities.
+                - debug_artifact_path: Directory for debug outputs.
+                - compile_flags: Optional backend flags.
+                - dump_debug_info: Enable extra debug JSON dump.
+
+        Returns:
+            PreprocessResult: Result containing processed_bytes with the
+            serialized TOSA flatbuffer.
+
+        Raises:
+            ValueError: If output_format is not "tosa" or the TOSA
+                specification is missing from compile_spec.
+            RuntimeError: If an unsupported FX node type is encountered.
+
+        """
         # if a debug/test build capture output files from TOSA stage
         artifact_path = compile_spec.get_intermediate_path()
         tosa_spec = compile_spec.tosa_spec
@@ -258,6 +299,15 @@ def filter_tosa_compile_specs(
         the TOSA flatbuffer representation as an intermediate step. The TOSA
         flatbuffer can then be consumed by the backend targetting specific
         hardware.
+
+        Args:
+            compile_spec (ArmCompileSpec): Compile specification that may
+                include both TOSA and hardware-specific options.
+
+        Returns:
+            TosaCompileSpec: TOSA-only specification ready for
+                ``TOSABackend.preprocess``.
+
         """
 
         return (