Add get_all_users_of function to GraphManipulation (pytorch#45216)

Summary:
This PR adds get_all_users_of function. The function returns all the users of a specific node. A test unit is also added.

Pull Request resolved: pytorch#45216

Reviewed By: ezyang

Differential Revision: D23883572

Pulled By: scottxu0730

fbshipit-source-id: 3eb68a411c3c6db39ed2506c9cb7bb7337520ee4

Author: scottxu0730

test/test_fx.py

@@ -6,6 +6,7 @@
 import copy
 from pathlib import Path
 from torch.fx import symbolic_trace, Proxy, Node, GraphModule, Tracer, Graph
+from torch.fx.experimental import GraphManipulation
 
 from torch.fx.proxy import TraceError
 
@@ -619,6 +620,26 @@ def forward(self, x):
         with self.assertRaisesRegex(AssertionError, message):
             traced(torch.rand(4, 3))
 
+    def test_get_all_users_of(self):
+        graph : torch.fx.Graph = torch.fx.Graph()
+        a : torch.fx.Node = graph.create_node('placeholder', 'x')
+        b : torch.fx.Node = graph.create_node('call_module', 'linear_mod', args=(a,))
+        c : torch.fx.Node = graph.create_node('get_attr', 'y_attr')
+        d : torch.fx.Node = graph.create_node('call_function', operator.add, args=(b, c))
+        graph.output(d)
+        linear_mod : torch.nn.Module = torch.nn.Linear(3, 4)
+        add_param : torch.Tensor = torch.rand(3, 4)
+        gm : torch.fx.GraphModule = torch.fx.GraphModule(
+            {'linear_mod': linear_mod, 'y_attr' : add_param}, graph)
+        expected_uses: Dict[int, List[int]] = {
+            0: [1],
+            1: [3],
+            2: [3],
+            3: []
+        }
+        for i, node in enumerate(graph.nodes):
+            user_indexes = GraphManipulation.get_all_users_of(gm, i)
+            assert user_indexes == expected_uses[i]
 
 if __name__ == '__main__':
     run_tests()

torch/fx/experimental/GraphManipulation.py

@@ -0,0 +1,31 @@
+from typing import List
+from torch.fx.graph_module import GraphModule
+from typing import Any
+from torch.fx.node import Node
+
+"""find_use is used to find out if the node is another node's arg or kwargs."""
+def find_use(arg: Any, node: Node) -> bool:
+    if isinstance(arg, (tuple, list)):
+        return any(find_use(elem, node) for elem in arg)
+    elif isinstance(arg, dict):
+        return any(find_use(v, node) for k, v in arg.items())
+    elif isinstance(arg, slice):
+        return any([find_use(arg.start, node), find_use(arg.stop, node), find_use(arg.step, node)])
+    elif isinstance(arg, Node):
+        return arg is node
+    else:
+        return False
+
+def get_all_users_of(fx_module: GraphModule, index: int) -> List[int]:
+    """Given the graph(fx_module) and an index, return a list of all node indexes that use this node"""
+    graph = fx_module.graph
+    current_node = graph.nodes[index]
+    user_indexes: List[int] = []
+    """if the node A is in node B's args, then B is the user of A
+       go through all the nodes, if the input node in any node's args,
+       then that node is the input node's user
+    """
+    for i, n in enumerate(graph.nodes):
+        if find_use(n.args, current_node) or find_use(n.kwargs, current_node):
+            user_indexes.append(i)
+    return user_indexes