Replace node.meta source_fn with source_fn_stack

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__
             )