Skip to content

Commit fc7294b

Browse files
ydwu4facebook-github-bot
authored andcommitted
Replace node.meta source_fn with source_fn_stack
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 Differential Revision: D48984986 Pulled By: ydwu4
1 parent dd57cc2 commit fc7294b

File tree

3 files changed

+8
-6
lines changed

3 files changed

+8
-6
lines changed

backends/xnnpack/partition/xnnpack_partitioner.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,8 @@ def __init__(
8787
assert len(self.constraints)
8888

8989
def check_common_constraints(self, node) -> bool:
90-
if self.unsupported_modules and "source_fn" in node.meta:
91-
return not node.meta["source_fn"][1] in self.unsupported_modules
90+
if self.unsupported_modules and "source_fn_stack" in node.meta:
91+
return not node.meta["source_fn_stack"][-1][1] in self.unsupported_modules
9292

9393
return True
9494

backends/xnnpack/passes/tag_implicit_q_dq_pass.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,8 @@ def is_supported_quant_op(self, node: torch.fx.Node) -> bool:
9090

9191
def is_supported_quant_module(self, node: torch.fx.Node) -> bool:
9292
is_supported = (
93-
"source_fn" in node.meta
94-
and node.meta["source_fn"][1] in SUPPORTED_IMPLICIT_Q_DQ_MODULES_SET
93+
"source_fn_stack" in node.meta
94+
and node.meta["source_fn_stack"][-1][1] in SUPPORTED_IMPLICIT_Q_DQ_MODULES_SET
9595
)
9696
if is_supported and self.is_supported_quant_op(node):
9797
raise RuntimeError(

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -247,12 +247,14 @@ the following metadata fields:
247247
{'self_linear': ('self.linear', <class 'torch.nn.Linear'>), 'self_sequential': ('self.sequential', <class 'torch.nn.Sequential'>)}
248248
```
249249

250-
* `node.meta["source_fn"]` contains the torch function or the leaf
250+
* `node.meta["source_fn_stack"]` contains the stack of torch function or the leaf
251251
`torch.nn.Module` class this node was called from before decomposition. For
252252
example, a node containing the `addmm` op from a `torch.nn.Linear` module call
253253
would contain `torch.nn.Linear` in their `source_fn`, and a node containing
254254
the `addmm` op from a `torch.nn.functional.Linear` module call would contain
255-
`torch.nn.functional.Linear` in their `source_fn`.
255+
`torch.nn.functional.Linear` in their `source_fn`. The stack records the higher order
256+
operator stack that this source_fn belongs to. For example, if a `torch.nn.Linear` module
257+
call is within the true branch of `cond`, then the stack will contain `['cond', 'torch.nn.Linear']`.
256258

257259

258260
### placeholder

0 commit comments

Comments
 (0)