Arm backend: Remove get_output_nodes from runner_utils. (pytorch#13417)

A graph has only one output node containing a list of output tensors.
Remove the use of this function to better reflect this.


Signed-off-by: Adrian Lundell <[email protected]>

Author: AdrianLundell

backends/arm/test/runner_utils.py

@@ -130,28 +130,8 @@ def get_input_quantization_params(
     return quant_params
 
 
-def get_output_nodes(program: ExportedProgram) -> list[Node]:
-    """
-    Get output node to this model.
-
-    Args:
-        program (ExportedProgram): The program to get the output nodes from.
-    Returns:
-        The nodes that are the outputs of the 'program'.
-    """
-    output_nodes = []
-    for node in program.graph.nodes:
-        if node.op == "output":
-            for output in node.args[0]:
-                output_nodes.append(output)
-    if len(output_nodes) == 0:
-        raise RuntimeError("No output nodes found.")
-    else:
-        return output_nodes
-
-
 def get_output_quantization_params(
-    output_nodes: list[Node],
+    output_node: Node,
 ) -> dict[Node, QuantizationParams | None]:
     """
     Get output QuantizationParams from a program.
@@ -164,7 +144,7 @@ def get_output_quantization_params(
         RuntimeError if no output quantization parameters are found.
     """
     quant_params = {}
-    for node in output_nodes:
+    for node in output_node.args[0]:
         if node.target == torch.ops.quantized_decomposed.dequantize_per_tensor.default:
             quant_params[node] = QuantizationParams(
                 node_name=node.args[0].name,
@@ -411,9 +391,9 @@ def run_corstone(
             f"Corstone simulation failed:\ncmd: {' '.join(command_args)}\nlog: \n {result_stdout}\n{result.stderr.decode()}"
         )
 
-    output_nodes = get_output_nodes(exported_program)
     output_np = []
-    for i, node in enumerate(output_nodes):
+    output_node = exported_program.graph_module.graph.output_node()
+    for i, node in enumerate(output_node.args[0]):
         output_shape = node.meta["val"].shape
         output_dtype = node.meta["val"].dtype
         tosa_ref_output = np.fromfile(

backends/arm/test/tester/analyze_output_utils.py

@@ -10,7 +10,6 @@
 from executorch.backends.arm.arm_backend import get_intermediate_path
 from executorch.backends.arm.test.runner_utils import (
     get_input_quantization_params,
-    get_output_nodes,
     get_output_quantization_params,
 )
 
@@ -254,9 +253,9 @@ def dump_error_output(
     export_stage = tester.stages.get(StageType.EXPORT, None)
     quantize_stage = tester.stages.get(StageType.QUANTIZE, None)
     if export_stage is not None and quantize_stage is not None:
-        output_nodes = get_output_nodes(export_stage.artifact)
+        output_node = export_stage.artifact.graph_module.output_node()
         qp_input = get_input_quantization_params(export_stage.artifact)
-        qp_output = get_output_quantization_params(output_nodes)
+        qp_output = get_output_quantization_params(output_node)
         logger.error(f"Input QuantArgs: {qp_input}")
         logger.error(f"Output QuantArgs: {qp_output}")
 

backends/arm/test/tester/arm_tester.py

@@ -48,7 +48,6 @@
 from executorch.backends.arm.test.runner_utils import (
     dbg_tosa_fb_to_json,
     get_elf_path,
-    get_output_nodes,
     get_output_quantization_params,
     get_target_board,
     run_target,
@@ -484,9 +483,8 @@ def run_method_and_compare_outputs(
             reference_stage = self.stages[StageType.INITIAL_MODEL]
 
         exported_program = self.stages[StageType.EXPORT].artifact
-        output_nodes = get_output_nodes(exported_program)
-
-        output_qparams = get_output_quantization_params(output_nodes)
+        output_node = exported_program.graph_module.graph.output_node()
+        output_qparams = get_output_quantization_params(output_node)
 
         quantization_scales = []
         for node in output_qparams: