Replace node.meta source_fn with source_fn_stack (#210)

Summary:
Pull Request resolved: #210

A resubmit of pytorch/pytorch#108447. Copy over the descriptions:

This is a follow-up of the discussion in pytorch/pytorch#108356, where we want to repalce source_fn with source_fn_stack

Before this PR, for the following example:
```python
backend = EagerAndRecordGraphs()

torch.compile(backend=backend, fullgraph=True)
def cond_f(pred, pred2, x, y):
    def true_fn(pred2, x, y):
        return x + y

    def false_fn(pred2, x, y):
        def true_fn2(x, y):
            return x.sin() - y.cos()

        def false_fn2(x, y):
            return x.cos() - y.sin()

        return control_flow.cond(pred2, true_fn2, false_fn2, (x, y))

    return control_flow.cond(pred, true_fn, false_fn, (pred2, x, y))
```
The graph captured is shown below:
```python
class GraphModule(torch.nn.Module):
    def forward(self, L_pred_ : torch.Tensor, L_pred2_ : torch.Tensor, L_x_ : torch.Tensor, L_y_ : torch.Tensor):
        l_pred_ = L_pred_
        l_pred2_ = L_pred2_
        l_x_ = L_x_
        l_y_ = L_y_

        cond_true_1 = self.cond_true_1
        cond_false_1 = self.cond_false_1
        cond = torch.ops.higher_order.cond(l_pred_, cond_true_1, cond_false_1, [l_pred2_, l_x_, l_y_]);  l_pred_ = cond_true_1 = cond_false_1 = l_pred2_ = l_x_ = l_y_ = None
        return (cond,)

    class GraphModule(torch.nn.Module):
        def forward(self, l_pred2_, l_x_, l_y_):
            add = l_x_ + l_y_;  l_x_ = l_y_ = None
            return add

    class GraphModule(torch.nn.Module):
        def forward(self, l_pred2_, l_x_, l_y_):
            cond_true_0 = self.cond_true_0
            cond_false_0 = self.cond_false_0
            cond = torch.ops.higher_order.cond(l_pred2_, cond_true_0, cond_false_0, [l_x_, l_y_]);  l_pred2_ = cond_true_0 = cond_false_0 = l_x_ = l_y_ = None
            return cond

        class GraphModule(torch.nn.Module):
            def forward(self, l_x_, l_y_):
                sin = l_x_.sin();  l_x_ = None
                cos = l_y_.cos();  l_y_ = None
                sub = sin - cos;  sin = cos = None
                return sub

        class GraphModule(torch.nn.Module):
            def forward(self, l_x_, l_y_):
                cos = l_x_.cos();  l_x_ = None
                sin = l_y_.sin();  l_y_ = None
                sub = cos - sin;  cos = sin = None
                return sub
```
the source_fn for inner cond, sin, cos will be a (name, target) tuple:
```
('cond', <torch._ops.HigherOrderOperator object at xxx>)
('sin', 'sin')
('cos', 'cos')
('sub'. <built-in function sub>)
```

After this pr, the source_fn_stack will be a list of (name, target) tuple. The bottom of stack is the end of the list.
```
[('cond', <torch._ops.HigherOrderOperator object at xxx>), ('cond', <torch._ops.HigherOrderOperator object at xxx>)],
[('cond', <torch._ops.HigherOrderOperator object at xxx>), ('cond', <torch._ops.HigherOrderOperator object at xxx>), ('sin', 'sin')],
[('cond', <torch._ops.HigherOrderOperator object at xxx>), ('cond', <torch._ops.HigherOrderOperator object at xxx>), ('cos', 'cos')]
[('cond', <torch._ops.HigherOrderOperator object at xxx>), ('cond', <torch._ops.HigherOrderOperator object at xxx>), ('sub', <built-in function sub>)]
```

X-link: pytorch/pytorch#108595

Test Plan:
See added tests in test_higher_order_ops.py and modify existing test.
Also updated bin by running:
"buck2 run @//mode/dev-nosan fbcode//aibench/api:gen_test_files --config client.id=nuclide"

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov

Reviewed By: angelayi

Differential Revision: D48984986

Pulled By: ydwu4

fbshipit-source-id: 12beb1cee5c6a5b7dc976eda54baaa8f564a0f1e

Author: ydwu4

backends/xnnpack/partition/xnnpack_partitioner.py

@@ -94,8 +94,8 @@ def __init__(
         assert len(self.constraints)
 
     def check_common_constraints(self, node) -> bool:
-        if self.unsupported_modules and "source_fn" in node.meta:
-            return not node.meta["source_fn"][1] in self.unsupported_modules
+        if self.unsupported_modules and "source_fn_stack" in node.meta:
+            return not node.meta["source_fn_stack"][-1][1] in self.unsupported_modules
 
         return True
 

backends/xnnpack/passes/tag_implicit_q_dq_pass.py

@@ -89,8 +89,9 @@ def is_supported_quant_op(self, node: torch.fx.Node) -> bool:
 
     def is_supported_quant_module(self, node: torch.fx.Node) -> bool:
         is_supported = (
-            "source_fn" in node.meta
-            and node.meta["source_fn"][1] in SUPPORTED_IMPLICIT_Q_DQ_MODULES_SET
+            "source_fn_stack" in node.meta
+            and node.meta["source_fn_stack"][-1][1]
+            in SUPPORTED_IMPLICIT_Q_DQ_MODULES_SET
         )
         if is_supported and self.is_supported_quant_op(node):
             raise RuntimeError(

docs/website/docs/ir_spec/00_exir.md

@@ -247,12 +247,14 @@ the following metadata fields:
   {'self_linear': ('self.linear', <class 'torch.nn.Linear'>), 'self_sequential': ('self.sequential', <class 'torch.nn.Sequential'>)}
   ```
 
-* `node.meta["source_fn"]` contains the torch function or the leaf
+* `node.meta["source_fn_stack"]` contains the stack of torch function or the leaf
   `torch.nn.Module` class this node was called from before decomposition. For
   example, a node containing the `addmm` op from a `torch.nn.Linear` module call
   would contain `torch.nn.Linear` in their `source_fn`, and a node containing
   the `addmm` op from a `torch.nn.functional.Linear` module call would contain
-  `torch.nn.functional.Linear` in their `source_fn`.
+  `torch.nn.functional.Linear` in their `source_fn`. The stack records the higher order
+  operator stack that this source_fn belongs to. For example, if a `torch.nn.Linear` module
+  call is within the true branch of `cond`, then the stack will contain `['cond', 'torch.nn.Linear']`.
 
 
 ### placeholder

exir/backend/test/test_backends.py

@@ -819,7 +819,7 @@ def forward(self, a, x, b):
                     self.assertTrue(
                         user.op == "call_function" and user.target == operator.getitem
                     )
-                    self.assertTrue(user.meta.get("source_fn", None) is None)
+                    self.assertTrue(user.meta.get("source_fn_stack", None) is None)
                     self.assertTrue(user.meta.get("nn_module_stack", None) is None)
 
         executorch_prog = executorch_prog.to_executorch(
@@ -886,7 +886,7 @@ def forward(self, x, y):
                     self.assertTrue(
                         user.op == "call_function" and user.target == operator.getitem
                     )
-                    self.assertTrue(user.meta.get("source_fn", None) is None)
+                    self.assertTrue(user.meta.get("source_fn_stack", None) is None)
                     self.assertTrue(user.meta.get("nn_module_stack", None) is None)
 
         executorch_prog = executorch_prog.to_executorch(

exir/lowered_backend_module.py

@@ -489,7 +489,7 @@ def create_submodule_from_nodes(
         # (ex. source_fn, # nn_module_stack)
         for user_node in submodule_node.users:
             user_node.meta.pop("nn_module_stack", None)
-            user_node.meta.pop("source_fn", None)
+            user_node.meta.pop("source_fn_stack", None)
 
     erase_nodes(gm, sorted_nodes)
 

sdk/edir/et_schema.py

@@ -34,7 +34,7 @@
 class RESERVED_METADATA_ARG(Enum):
     DEBUG_HANDLE = "debug_handle"
     MODULE_STACK = "nn_module_stack"
-    SOURCE_FN = "source_fn"
+    SOURCE_FN_STACK = "source_fn_stack"
     MODULE_TYPE = "module_type"
     PROFILE_START_TIME = "profile_start_time"
     PROFILE_END_TIME = "profile_end_time"
@@ -463,9 +463,10 @@ def _update_module_mapping(
         metadata: Dict[str, Any],
     ):
         if (
-            source_fn := metadata.get("source_fn")
+            source_fn_stack := metadata.get("source_fn_stack")
         ) is not None and "nn_module_stack" in metadata:
             # (module name, module type)
+            source_fn = source_fn_stack[-1]
             module_type = (
                 source_fn[1] if isinstance(source_fn[1], str) else source_fn[1].__name__
             )