Replace node.meta source_fn with source_fn_stack (#210)

Summary:

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

X-link: pytorch/pytorch#108595

Test Plan: See added tests in test_higher_order_ops.py and modify existing test.

Differential Revision: D48984986

Pulled By: ydwu4

Author: ydwu4

backends/xnnpack/partition/xnnpack_partitioner.py

@@ -87,8 +87,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

@@ -90,8 +90,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