Enhance "from_node" node meta to track source recursively (#142066)

yushangdi · facebook-github-bot · commit 448eb8d4bf61 · 2024-12-09T11:44:28.000-08:00
Summary:

Change the "from_node" node meta format to be able to track the provenance of nodes recursively.

The new "from_node" format is a a list node NodeSource. The old "from_node" is a list of tuples `(node.name, node.target)`

```
class NodeSource:
    class NodeInfo:
        def __init__(self, name: str, target: str, graph_id: int):
            self.name = name
            self.target = target
            self.graph_id = graph_id

    pass_name: str
    action: Optional["NodeSourceAction"]
    from_node: List["NodeSource"]
    node_info: Optional["NodeInfo"]
```


We also add a "pass_name" field to indicate which pass created the node from the source. For now the "pass_name" is only used in fx.Interpreter.

These changes will be used for the inductor provenance tracking.

Test Plan:
```
buck2 run 'fbcode//mode/dev-nosan' fbcode//caffe2/test:test_export -- -r test_unflatten_multiple_graphs_state
buck run mode/dev-nosan  fbcode//caffe2/test:fx -- -r node_source
buck run mode/dev-nosan  fbcode//caffe2/test:fx -- -r graph_provenance
buck run mode/dev-nosan  fbcode//caffe2/test/quantization:test_quantization -- -r test_simple_metadata_porting
```

Reviewed By: avikchaudhuri

Differential Revision: D66737916
diff --git a/test/fx/test_fx_traceback.py b/test/fx/test_fx_traceback.py
@@ -0,0 +1,172 @@
+# Owner(s): ["module: fx"]
+import json
+
+import torch
+from torch._inductor.compile_fx import aot_export_module
+from torch.fx.traceback import get_graph_provenance_json, NodeSource, NodeSourceAction
+from torch.testing._internal.common_utils import TestCase
+
+
+class TestFXNodeSource(TestCase):
+    def test_node_source(self):
+        node_source = NodeSource(
+            node=None, pass_name="test_pass", action=NodeSourceAction.CREATE
+        )
+        self.assertExpectedInline(
+            node_source.print_readable().strip(),
+            """(name=, pass_name=test_pass, action=create, graph_id=-1)""",
+        )
+        dummy_source_dict = {
+            "name": "",
+            "target": "",
+            "pass_name": "test_pass",
+            "action": NodeSourceAction.CREATE,
+            "graph_id": -1,
+            "from_node": [],
+        }
+        self.assertEqual(
+            node_source.to_dict(),
+            dummy_source_dict,
+        )
+
+        # Dummy node
+        node = torch.fx.Node(
+            graph=torch.fx.Graph(),
+            name="add",
+            op="call_function",
+            target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
+            args=(torch.tensor(3), torch.tensor(4)),
+            kwargs={},
+        )
+        node.meta["from_node"] = [node_source]
+
+        graph_id = id(node.graph)
+        node_source = NodeSource(
+            node=node, pass_name="test_pass", action=NodeSourceAction.CREATE
+        )
+        self.assertExpectedInline(
+            node_source.print_readable().strip(),
+            f"""\
+(name=add, pass_name=test_pass, action=create, graph_id={graph_id})
+    (name=, pass_name=test_pass, action=create, graph_id=-1)""",
+        )
+        self.assertEqual(
+            node_source.to_dict(),
+            {
+                "name": "add",
+                "target": "aten.add.Tensor",
+                "pass_name": "test_pass",
+                "action": NodeSourceAction.CREATE,
+                "graph_id": graph_id,
+                "from_node": [dummy_source_dict],
+            },
+        )
+
+    def test_graph_provenance(self):
+        def check_node_source(node_source_dict, name, pass_name, action):
+            self.assertEqual(node_source_dict["name"], name)
+            self.assertEqual(node_source_dict["pass_name"], pass_name)
+            self.assertEqual(node_source_dict["action"], action)
+
+        def get_first_node_source_and_check(node_source_dict):
+            """
+            Get the first node source from the from_node list.
+            """
+            self.assertEqual(len(node_source_dict["from_node"]), 1)
+            return node_source_dict["from_node"][0]
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fc1 = torch.nn.Linear(10, 16)
+                self.relu = torch.nn.ReLU()
+                self.fc2 = torch.nn.Linear(16, 1)
+                self.sigmoid = torch.nn.Sigmoid()
+
+            def forward(self, x):
+                x = self.fc1(x)
+                x = self.relu(x)
+                x = self.fc2(x)
+                x = self.sigmoid(x)
+                return (x,)
+
+        model = Model()
+        example_inputs = (torch.randn(8, 10),)
+        ep = torch.export.export(
+            model,
+            example_inputs,
+        )
+        gm = ep.module()
+        provenance = get_graph_provenance_json(gm.graph)
+        provenance = json.loads(provenance)
+        self.assertEqual(
+            set(provenance.keys()), {"relu", "linear", "sigmoid", "linear_1"}
+        )
+
+        # Check node "linear" is created from node "x" in PropagateUnbackedSymInts
+        key_provenance = provenance["linear"]
+        self.assertEqual(len(key_provenance), 1)
+        key_provenance = key_provenance[0]
+        check_node_source(
+            key_provenance,
+            "x",
+            "Interpreter_PropagateUnbackedSymInts",
+            NodeSourceAction.CREATE,
+        )
+
+        # Check node "x" is then created from another node "x" in FlattenInputOutputSignature
+        key_provenance = get_first_node_source_and_check(key_provenance)
+        check_node_source(
+            key_provenance,
+            "x",
+            "Interpreter_FlattenInputOutputSignature",
+            NodeSourceAction.CREATE,
+        )
+
+        gm, graph_signature = aot_export_module(
+            gm,
+            example_inputs,
+            trace_joint=False,
+        )
+
+        provenance = get_graph_provenance_json(gm.graph)
+        provenance = json.loads(provenance)
+
+        self.assertEqual(
+            set(provenance.keys()), {"t", "addmm", "relu", "t_1", "addmm_1", "sigmoid"}
+        )
+        for key in ["t", "addmm"]:
+            # The node provenance hierarchy should be:
+            # t -> linear -> x -> x
+            #
+            # x -> y means x is created from y
+
+            key_provenance = provenance[key]
+            self.assertEqual(len(key_provenance), 1)
+            key_provenance = key_provenance[0]
+
+            # Check node "t" and "addmm" is created from node "linear" in PropagateUnbackedSymInts
+            check_node_source(
+                key_provenance,
+                "linear",
+                "Interpreter_PropagateUnbackedSymInts",
+                NodeSourceAction.CREATE,
+            )
+
+            # Check node "linear" is then created from node "x" in PropagateUnbackedSymInts
+            key_provenance = get_first_node_source_and_check(key_provenance)
+            check_node_source(
+                key_provenance,
+                "x",
+                "Interpreter_PropagateUnbackedSymInts",
+                NodeSourceAction.CREATE,
+            )
+
+            # Check node "x" is then created from another node "x" in FlattenInputOutputSignature
+            key_provenance = get_first_node_source_and_check(key_provenance)
+            check_node_source(
+                key_provenance,
+                "x",
+                "Interpreter_FlattenInputOutputSignature",
+                NodeSourceAction.CREATE,
+            )
diff --git a/test/quantization/pt2e/test_metadata_porting.py b/test/quantization/pt2e/test_metadata_porting.py
@@ -74,8 +74,8 @@ def _test_quant_tag_preservation_through_decomp(
                         f"from_node metadata is of type {type(from_node_meta)}, but expected list"
                     )
                 for meta in from_node_meta:
-                    node_target = meta[1]
-                    if node_target == from_node:
+                    node_target = meta.target
+                    if node_target == str(from_node):
                         node_tag = n.meta.get("quantization_tag", None)
                         if node_tag is None or tag != node_tag:
                             not_found_nodes += str(n.target) + ", "
diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py
@@ -587,7 +587,7 @@ def get_kernel_metadata(node_schedule, wrapper):
             key = str(node.meta["original_aten"]._overloadpacket)
             original_aten_dict[key].append(node.name)
         if "from_node" in node.meta:
-            key = node.meta["from_node"][0][0]
+            key = node.meta["from_node"][0].name
             from_node_dict[key].append(node.name)
     sort_str = "Topologically Sorted" if single_graph is not None else "Unsorted"
     metadata = (
diff --git a/torch/export/exported_program.py b/torch/export/exported_program.py
@@ -719,7 +719,7 @@ def _get_updated_module_call_graph(
     provenance: Dict[str, str] = {}
     for node in gm.graph.nodes:
         if history := node.meta.get("from_node", []):
-            provenance[history[-1][0]] = node.name
+            provenance[history[-1].name] = node.name
 
     # map old names to new names in module call signatures
     for entry in new_module_call_graph:
diff --git a/torch/fx/interpreter.py b/torch/fx/interpreter.py
@@ -204,7 +204,9 @@ def boxed_run(self, args_list):
 
     @contextmanager
     def _set_current_node(self, node):
-        with fx_traceback.set_current_meta(node):
+        with fx_traceback.set_current_meta(
+            node, f"Interpreter_{self.__class__.__name__}"
+        ):
             yield
 
     @compatibility(is_backward_compatible=True)
diff --git a/torch/fx/traceback.py b/torch/fx/traceback.py
@@ -1,9 +1,14 @@
 # mypy: allow-untyped-defs
+import copy
+import json
 import traceback
 from contextlib import contextmanager
-from typing import Any, Dict, List
+from enum import Enum
+from typing import Any, Dict, List, Optional
 
 from ._compatibility import compatibility
+from .graph import Graph
+from .node import Node
 
 
 __all__ = [
@@ -15,12 +20,99 @@
     "format_stack",
     "set_current_meta",
     "get_current_meta",
+    "NodeSource",
+    "NodeSourceAction",
+    "get_graph_provenance_json",
 ]
 
 current_meta: Dict[str, Any] = {}
 should_preserve_node_meta = False
 
 
+@compatibility(is_backward_compatible=False)
+class NodeSourceAction(str, Enum):
+    CREATE = "create"
+    REPLACE = "replace"
+
+
+@compatibility(is_backward_compatible=False)
+class NodeSource:
+    """
+    NodeSource is a data structure that contains the provenance information of a node.
+    If node `a` is created from node `b`, then `a.meta["from_node"]` may contain NodeSource(b).
+    """
+
+    class NodeInfo:
+        def __init__(self, name: str, target: str, graph_id: int):
+            self.name = name
+            self.target = target
+            self.graph_id = graph_id
+
+    pass_name: str
+    action: Optional["NodeSourceAction"]
+    from_node: List["NodeSource"]
+    node_info: Optional["NodeInfo"]
+
+    def __init__(
+        self,
+        node: Optional[Node],
+        pass_name: str = "",
+        action: Optional["NodeSourceAction"] = None,
+    ):
+        self.pass_name = pass_name
+        self.action = action
+        if node:
+            self.node_info = self.NodeInfo(
+                name=node.name, target=str(node.target), graph_id=id(node.graph)
+            )
+            self.from_node = (
+                copy.deepcopy(node.meta["from_node"])
+                if "from_node" in node.meta
+                else []
+            )
+        else:
+            self.node_info = None
+            self.from_node = []
+
+    @property
+    def name(self) -> str:
+        return self.node_info.name if self.node_info else ""
+
+    @property
+    def target(self) -> str:
+        return self.node_info.target if self.node_info else ""
+
+    @property
+    def graph_id(self) -> int:
+        return self.node_info.graph_id if self.node_info else -1
+
+    def __repr__(self):
+        return self.print_readable()
+
+    def print_readable(self, indent=0):
+        if indent > 9:
+            return ""
+        result = ""
+        result += (
+            " " * indent * 4
+            + f"(name={self.name}, pass_name={self.pass_name}, action={self.action}, graph_id={self.graph_id})\n"
+        )
+        for item in self.from_node:
+            result += item.print_readable(indent + 1)
+        return result
+
+    def to_dict(self) -> dict:
+        # Convert the object to a dictionary
+        return {
+            "name": self.name,
+            "target": self.target,
+            "graph_id": self.graph_id,
+            "pass_name": self.pass_name,
+            "action": self.action,
+            "from_node": [node.to_dict() for node in self.from_node],
+        }
+
+
 @compatibility(is_backward_compatible=False)
 @contextmanager
 def preserve_node_meta():
@@ -90,21 +182,20 @@ def has_preserved_node_meta() -> bool:
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def set_current_meta(node):
+def set_current_meta(node, pass_name=""):
     global current_meta
     if should_preserve_node_meta and node.meta:
         saved_meta = current_meta
         try:
             current_meta = node.meta.copy()
 
-            # Append (node.name, node.target) onto "from_node" for provenance tracking
-            if "from_node" not in current_meta:
-                current_meta["from_node"] = [(node.name, node.target)]
-            elif current_meta["from_node"][-1][0] != node.name:
-                current_meta["from_node"] = current_meta["from_node"] + [
-                    (node.name, node.target)
-                ]
-
+            # Update the "from_node" field in current_meta for provenance tracking.
+            # Instead of appending, overwrite the "from_node" field because current_meta
+            # will be assigned to the new node. The new NodeSource(node, ...) will
+            # include the information from the previous current_meta["from_node"].
+            current_meta["from_node"] = [
+                NodeSource(node, pass_name, NodeSourceAction.CREATE)
+            ]
             yield
         finally:
             current_meta = saved_meta
@@ -115,3 +206,19 @@ def set_current_meta(node):
 @compatibility(is_backward_compatible=False)
 def get_current_meta() -> Dict[str, Any]:
     return current_meta
+
+
+@compatibility(is_backward_compatible=False)
+def get_graph_provenance_json(graph: Graph) -> str:
+    """
+    Given an fx.Graph, return a json string that contains the provenance information of each node.
+    """
+    provenance_tracking_json = {}
+    for node in graph.nodes:
+        if node.op == "call_function":
+            provenance_tracking_json[node.name] = (
+                [source.to_dict() for source in node.meta["from_node"]]
+                if "from_node" in node.meta
+                else []
+            )
+    return json.dumps(provenance_tracking_json)
