ns for fx: skip shadowing ops if copy subgraph is not implemented (pytorch#76663)

Summary:
Pull Request resolved: pytorch#76663

Subgraph copy does not handle all edge cases. It's high eng time
to handle them all, and currently an unhandled edge case crashes
the script.

This PR adds a function to check if the subgraph copy is supported,
and skips shadowing if it is not supported. This way the model
can still go through the shadowing APIs without an exception.

Test Plan:
```
python test/test_quantization.py -k FXNumericSuite
```

Reviewed By: hx89

Differential Revision: D36069304

Pulled By: vkuzo

fbshipit-source-id: 6b38b8d8e43396a4cf2373b247223a19d451d096
(cherry picked from commit e2322ca)

Author: vkuzo

test/quantization/fx/test_numeric_suite_fx.py

@@ -1955,6 +1955,23 @@ def forward(self, x):
             m, (torch.randn(1, 1, 4, 4),),
             results_len=0)
 
+    def test_unsupported_op_copy_skips_shadowing(self):
+        """
+        Copying a `call_function` node is not implemented, test that this
+        does not crash shadowing but instead skips the node.
+        """
+        class M(nn.Module):
+            def forward(self, x):
+                # the second argument leads to attempting to copy a
+                # call_function node
+                x = F.layer_norm(x, x.shape[1:])
+                return x
+
+        m = M().eval()
+        self._test_match_shadow_activations(
+            m, (torch.randn(1, 1, 4, 4),),
+            results_len=0)
+
 
 class TestFXNumericSuiteCoreAPIsModels(FXNumericSuiteQuantizationTestCase):
     """

torch/ao/ns/fx/graph_passes.py

@@ -360,6 +360,53 @@ def _copy_node_from_a_to_c(
         raise AssertionError(
             f"handling of node {node_a.format_node()} with op {node_a.op} is not implemented")
 
+def _can_insert_copy_of_subgraph_a(
+    subgraph_a: NSSubgraph,
+    gm_a: GraphModule,
+    num_non_param_args_node_a: int,
+) -> bool:
+    """
+    This function returns `False` if the input subgraph cannot be copied by
+    `_insert_copy_of_subgraph_a_after_input_node_c`. This usually means
+    that there is a corner case logic for which copy is not yet implemented.
+    """
+    # populate the list of nodes we need to check
+    nodes = []
+    cur_node = subgraph_a.end_node
+    while cur_node != subgraph_a.start_node:
+        nodes.append(cur_node)
+        cur_node = cur_node.args[0]  # type: ignore[assignment]
+    nodes.append(cur_node)
+    nodes.reverse()
+
+    def _can_insert(node_a_arg, gm_a):
+        if isinstance(node_a_arg, Node):
+            arg_a = return_first_non_observer_node(node_a_arg, gm_a)
+            if arg_a.op == 'call_method':
+                return arg_a.target in ('dequantize', 'to')
+            elif arg_a.op == 'get_attr':
+                return True
+            else:
+                return False
+        elif isinstance(node_a_arg, (list, tuple)):
+            for el in node_a_arg:
+                if not isinstance(el, Node):
+                    return False
+        return True
+
+    # For each node, check if we handle the copy behavior. This follows the
+    # logic in `_insert_copy_of_subgraph_a_after_input_node_c`.
+    for node_a_arg in nodes[0].args[num_non_param_args_node_a:]:
+        if not _can_insert(node_a_arg, gm_a):
+            return False
+
+    for node in nodes[1:]:
+        for node_a_arg in node.args[1:]:
+            if not _can_insert(node_a_arg, gm_a):
+                return False
+
+    return True
+
 def _insert_copy_of_subgraph_a_after_input_node_c(
     input_node_c: Union[Node, List[Node]],
     input_node_c_2: Optional[Union[Node, List[Node]]],
@@ -663,6 +710,16 @@ def load_arg(a):
                     env_c[node_b.name] = graph_c.node_copy(node_b, load_arg)
                     continue
 
+            num_non_param_args_node_a = \
+                get_number_of_non_param_args(subgraph_a.start_node, gm_a)
+            if not _can_insert_copy_of_subgraph_a(subgraph_a, gm_a, num_non_param_args_node_a):
+                print(
+                    f'skipping shadow loggers for node_b: {get_target_type_str(node_b, gm_b)}' +
+                    f', start_node_a: {get_target_type_str(subgraph_a.start_node, gm_a)}' +
+                    ', unhandled logic in subgraph copy')
+                env_c[node_b.name] = graph_c.node_copy(node_b, load_arg)
+                continue
+
             fqn_base_a = _maybe_get_fqn(subgraph_a.base_op_node, gm_a)
             fqn_base_b = _maybe_get_fqn(subgraph_b.base_op_node, gm_b)