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

sdk/edir/et_schema.py

@@ -34,7 +34,7 @@
 class RESERVED_METADATA_ARG(Enum):
     DEBUG_HANDLE = "debug_handle"
     MODULE_STACK = "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__
             )