helper function to print node info (#512)

Summary:
Pull Request resolved: #512

Add a helper function to print the node info given the graph and a node
```
print(inspect_node(graph, node))

_param_constant1 error_msg:  Here is the node in the graph module:
graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %_param_constant0 : [num_users=1] = get_attr[target=_param_constant0]
--> %_param_constant1 : [num_users=1] = get_attr[target=_param_constant1]
    %aten_convolution_default : [num_users=2] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%arg0_1, %_param_constant0, %_param_constant1, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %_param_constant2 : [num_users=1] = get_attr[target=_param_constant2]
    %_param_constant3 : [num_users=1] = get_attr[target=_param_constant3]
    %aten_convolution_default_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_convolution_default, %_param_constant2, %_param_constant3, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%aten_convolution_default, %aten_convolution_default_1), kwargs = {})
    %_param_constant4 : [num_users=1] = get_attr[target=_param_constant4]
    %_param_constant5 : [num_users=1] = get_attr[target=_param_constant5]
    %aten_convolution_default_2 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_add_tensor, %_param_constant4, %_param_constant5, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %aten_gelu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.gelu.default](args = (%aten_convolution_default_2,), kwargs = {})
    return [aten_gelu_default]
This node _param_constant1 has metadata of:
The node stacktrace:
Traceback (most recent call last):
    File "/tmp/ipykernel_1204253/3382880687.py", line 7, in forward
return self.test_model(x)
    File "/mnt/xarfuse/uid-25337/7b86ad0c-seed-nspid4026532987_cgpid2707357-ns-4026532984/torch/nn/modules/module.py", line 1528, in _call_impl
return forward_call(*args, **kwargs)
    File "/tmp/ipykernel_1204253/712280972.py", line 10, in forward
a = self.conv1(x)
```

Reviewed By: tarun292

Differential Revision: D49715284

fbshipit-source-id: 0e901c7c5b8be074261c5f4fc4f8dc72d7c23db5

Author: cccclai

exir/common.py

@@ -92,43 +92,6 @@ def format_schema_name(schema: torch._C.FunctionSchema) -> str:
     return schema.name
 
 
-def add_cursor_to_graph(graph: torch.fx.Graph, finding_node: torch.fx.Node) -> str:
-    """
-    Insert a cursor at the node location in the fx.Graph.
-    e.g:
-    # graph():
-    #   %x : [#users=1] = placeholder[target=x]
-    #   %param : [#users=1] = get_attr[target=param]
-    #   %add : [#users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {})
-    # --> %linear : [#users=1] = call_module[target=linear](args = (%add,), kwargs = {})
-    #   %clamp : [#users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0})
-    #   return clamp
-
-    This is mostly used for error reporting
-    """
-
-    new_graph = copy.deepcopy(graph)
-
-    found_at = -1
-    for ix, node in enumerate(graph.nodes):
-        if node == finding_node:
-            found_at = ix
-
-    # This is heavily based on __str__ method of fx.Graph
-    def _format_graph(graph: torch.fx.Graph, offending_node_idx: int) -> str:
-        s = "graph():"
-        for ix, node in enumerate(graph.nodes):
-            node_str = node.format_node()
-            if node_str:
-                if ix != offending_node_idx:
-                    s += "\n    " + node_str
-                else:
-                    s += "\n--> " + node_str
-        return s
-
-    return _format_graph(new_graph, found_at)
-
-
 @contextmanager
 def override_logger(
     newLevel: int = logging.DEBUG,

exir/emit/TARGETS

@@ -19,7 +19,6 @@ python_library(
     deps = [
         "fbsource//third-party/pypi/typing-extensions:typing-extensions",
         "//caffe2:torch",
-        "//executorch/exir:common",
         "//executorch/exir:delegate",
         "//executorch/exir:error",
         "//executorch/exir:memory",

exir/emit/_emitter.py

@@ -30,7 +30,6 @@
 # pyre-strict
 import ctypes
 import operator
-import re
 import typing
 from dataclasses import dataclass, field
 from typing import Callable, cast, Dict, List, Mapping, Optional, Tuple, Union
@@ -39,14 +38,13 @@
 import executorch.extension.pytree as ex_pytree
 import torch
 import torch.fx
-from executorch.exir.common import add_cursor_to_graph
 from executorch.exir.delegate import executorch_call_delegate, is_lowered_module
 from executorch.exir.dialects.backend._ops import BackendOpOverload
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.error import ExportError, ExportErrorType, InternalError
 from executorch.exir.operator.convert import is_out_variant
 from executorch.exir.passes.executorch_prim_ops_registry import is_sym_op
-from executorch.exir.print_program import pretty_print_stacktraces
+from executorch.exir.print_program import _stacktrace_to_framelist, inspect_node
 from executorch.exir.schema import (
     BackendDelegate,
     BackendDelegateDataReference,
@@ -62,8 +60,6 @@
     DoubleList,
     EValue,
     ExecutionPlan,
-    Frame,
-    FrameList,
     FreeCall,
     Instruction,
     Int,
@@ -238,44 +234,12 @@ def __init__(
             int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]
         ] = {}
 
-    def _stacktrace_to_framelist(self, stacktrace: str) -> FrameList:
-        """Creates a frame list from a stacktrace string."""
-        pattern = r'File "(.*?)", line (\d+), in (.*?)\n'
-        matches = re.findall(pattern, stacktrace)
-        mapped_frame_list = [
-            Frame(
-                filename=match[0],
-                lineno=int(match[1]),
-                name=match[2],
-                context=stacktrace.split("\n")[i * 2 + 1].strip(),
-            )
-            for i, match in enumerate(matches)
-        ]
-        return FrameList(mapped_frame_list)
-
     def _emit_node_specific_error(self, node: torch.fx.Node, err_msg: str) -> str:
         """Returns 'err_msg' with node specific information attached."""
-        graph_str_with_cursor = add_cursor_to_graph(self.graph_module.graph, node)
-
-        error_msg = (
-            f"Failed with error: {str(err_msg)}\n"
-            f"Here is the failing node in the graph module:\n"
-            f"{graph_str_with_cursor}\n"
-            f"This node {self.node} has metadata of:\n"
+        err_msg = f"Failed with error: {str(err_msg)}\n" + inspect_node(
+            self.graph_module.graph, node
         )
-
-        # Node spec error message
-        if hasattr(self.node.meta, "spec"):
-            error_msg += f"The node spec:\n{self.node.meta['spec']}\n"
-
-        # Stacktrace error message
-        if hasattr(self.node.meta, "stack_trace"):
-            framelist = self._stacktrace_to_framelist(self.node.meta["stack_trace"])
-            error_msg += (
-                f"The node stacktrace:\n{pretty_print_stacktraces(framelist)}\n"
-            )
-
-        return error_msg
+        return err_msg
 
     def _internal_assert_emitter(
         self, pred: bool, node: torch.fx.Node, assert_msg: str
@@ -1133,7 +1097,7 @@ def _get_empty_tensor_evalue() -> EValue:
             stack_trace = self.node.meta["stack_trace"]
             chain_stacktrace = self.chain.stacktrace or []
 
-            chain_stacktrace.append(self._stacktrace_to_framelist(stack_trace))
+            chain_stacktrace.append(_stacktrace_to_framelist(stack_trace))
             self._internal_assert_emitter(
                 len(chain_stacktrace) == len(self.chain.instructions),
                 self.node,

exir/print_program.py

@@ -6,21 +6,24 @@
 
 # pyre-strict
 
+import copy
+import re
 import reprlib
 from dataclasses import fields
 from enum import IntEnum
 from typing import Any, List
 
 import torch
-
 from executorch.exir.error import ExportError, ExportErrorType, InternalError
+
 from executorch.exir.schema import (
     Bool,
     BoolList,
     DelegateCall,
     Double,
     DoubleList,
     EValue,
+    Frame,
     FrameList,
     FreeCall,
     Int,
@@ -302,3 +305,109 @@ def pretty_print_stacktraces(obj: FrameList) -> str:
         pretty += f"{frame.context} \n"
     pretty += "\n"
     return pretty
+
+
+def add_cursor_to_graph(graph: torch.fx.Graph, finding_node: torch.fx.Node) -> str:
+    """
+    Insert a cursor at the node location in the fx.Graph.
+    e.g:
+    # graph():
+    #   %x : [#users=1] = placeholder[target=x]
+    #   %param : [#users=1] = get_attr[target=param]
+    #   %add : [#users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {})
+    # --> %linear : [#users=1] = call_module[target=linear](args = (%add,), kwargs = {})
+    #   %clamp : [#users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0})
+    #   return clamp
+
+    This is mostly used for error reporting
+    """
+
+    new_graph = copy.deepcopy(graph)
+
+    found_at = -1
+    for ix, node in enumerate(graph.nodes):
+        if node == finding_node:
+            found_at = ix
+
+    # This is heavily based on __str__ method of fx.Graph
+    def _format_graph(graph: torch.fx.Graph, offending_node_idx: int) -> str:
+        s = "graph():"
+        for ix, node in enumerate(graph.nodes):
+            node_str = node.format_node()
+            if node_str:
+                if ix != offending_node_idx:
+                    s += "\n    " + node_str
+                else:
+                    s += "\n--> " + node_str
+        return s
+
+    return _format_graph(new_graph, found_at)
+
+
+def _stacktrace_to_framelist(stacktrace: str) -> FrameList:
+    """Creates a frame list from a stacktrace string."""
+    pattern = r'File "(.*?)", line (\d+), in (.*?)\n'
+    matches = re.findall(pattern, stacktrace)
+    mapped_frame_list = [
+        Frame(
+            filename=match[0],
+            lineno=int(match[1]),
+            name=match[2],
+            context=stacktrace.split("\n")[i * 2 + 1].strip(),
+        )
+        for i, match in enumerate(matches)
+    ]
+    return FrameList(mapped_frame_list)
+
+
+def inspect_node(graph: torch.fx.Graph, node: torch.fx.Node) -> str:
+    """
+    Inspect a node by highlighting the node in the graph as well as the stacktrace.
+
+    Args:
+        graph: The graph containing the node
+        node: The node to be inspected
+
+    Return: A string. An example output is:
+
+    _param_constant0 error_msg:  Here is the failing node in the graph module:
+    graph():
+        %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
+    --> %_param_constant0 : [num_users=1] = get_attr[target=_param_constant0]
+        %_param_constant1 : [num_users=1] = get_attr[target=_param_constant1]
+        %aten_convolution_default : [num_users=2] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%arg0_1, %_param_constant0, %_param_constant1, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
+        %_param_constant2 : [num_users=1] = get_attr[target=_param_constant2]
+        %_param_constant3 : [num_users=1] = get_attr[target=_param_constant3]
+        %aten_convolution_default_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_convolution_default, %_param_constant2, %_param_constant3, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
+        %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%aten_convolution_default, %aten_convolution_default_1), kwargs = {})
+        %_param_constant4 : [num_users=1] = get_attr[target=_param_constant4]
+        %_param_constant5 : [num_users=1] = get_attr[target=_param_constant5]
+        %aten_convolution_default_2 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_add_tensor, %_param_constant4, %_param_constant5, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
+        %aten_gelu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.gelu.default](args = (%aten_convolution_default_2,), kwargs = {})
+        return [aten_gelu_default]
+    This node _param_constant0 has metadata of:
+    The node stacktrace:
+    Traceback (most recent call last):
+        File "/tmp/ipykernel_1204253/3382880687.py", line 7, in forward
+    return self.test_model(x)
+        File "/mnt/xarfuse/uid-25337/7b86ad0c-seed-nspid4026532987_cgpid2707357-ns-4026532984/torch/nn/modules/module.py", line 1528, in _call_impl
+    return forward_call(*args, **kwargs)
+        File "/tmp/ipykernel_1204253/712280972.py", line 10, in forward
+    a = self.conv1(x)
+
+    """
+    graph_str_with_cursor = add_cursor_to_graph(graph, node)
+    error_msg = (
+        f"Here is the node in the graph module:\n"
+        f"{graph_str_with_cursor}\n"
+        f"This node {node} has metadata of:\n"
+    )
+    # Node spec error message
+    if hasattr(node.meta, "spec"):
+        error_msg += f"The node spec:\n{node.meta['spec']}\n"
+
+    # Stacktrace error message
+    if "stack_trace" in node.meta:
+        framelist = _stacktrace_to_framelist(node.meta["stack_trace"])
+        error_msg += f"The node stacktrace:\n{pretty_print_stacktraces(framelist)}\n"
+    return error_msg

exir/tests/TARGETS

@@ -265,6 +265,7 @@ python_unittest(
     deps = [
         "//caffe2:torch",
         "//executorch/exir:common",
+        "//executorch/exir:print_program",
     ],
 )
 
@@ -430,3 +431,15 @@ python_unittest(
         "//executorch/exir/passes:lib",
     ],
 )
+
+python_unittest(
+    name = "print_program",
+    srcs = [
+        "test_print_program.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/exir:print_program",
+    ],
+)

exir/tests/test_common.py

@@ -12,11 +12,8 @@
 import torch
 import torch.fx
 
-from executorch.exir.common import (
-    add_cursor_to_graph,
-    extract_out_arguments,
-    get_schema_for_operators,
-)
+from executorch.exir.common import extract_out_arguments, get_schema_for_operators
+from executorch.exir.print_program import add_cursor_to_graph
 
 
 class TestExirCommon(unittest.TestCase):

exir/tests/test_print_program.py

@@ -0,0 +1,51 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch import exir
+from executorch.exir.print_program import inspect_node
+
+
+class TestPrintProgram(unittest.TestCase):
+    def test_inspect_node(self) -> None:
+        class TestModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv1 = torch.nn.Conv2d(32, 32, 1)
+                self.conv2 = torch.nn.Conv2d(32, 32, 1)
+                self.conv3 = torch.nn.Conv2d(32, 32, 1)
+                self.gelu = torch.nn.GELU()
+
+            def forward(self, x: torch.Tensor):
+                a = self.conv1(x)
+                b = self.conv2(a)
+                c = self.conv3(a + b)
+                return self.gelu(c)
+
+        class WrapModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.test_model = TestModel()
+
+            def forward(self, x):
+                return self.test_model(x)
+
+        warp_model = WrapModule()
+        example_inputs = (torch.rand(1, 32, 16, 16),)
+
+        exir_exported_program = exir.capture(warp_model, example_inputs).to_edge()
+        number_of_stack_trace = 0
+        for node in exir_exported_program.exported_program.graph.nodes:
+            node_info = inspect_node(exir_exported_program.exported_program.graph, node)
+            self.assertRegexpMatches(node_info, r".*-->.*")
+            if "stack_trace" in node.meta:
+                self.assertRegexpMatches(
+                    node_info, r".*Traceback \(most recent call last\)\:.*"
+                )
+                number_of_stack_trace = number_of_stack_trace + 1
+        self.assertGreaterEqual(number_of_stack_trace, 1)