ns for fx: skip shadowing for torch.cat, and also for nodes with only kwargs (pytorch#76561)

Summary:
Pull Request resolved: pytorch#76561

User model had syntax like `torch.cat(tensors=[x])`. This PR fixes two errors
to unbreak this in NS shadow model:
1. skip nodes which only have kwargs (instead of throwing an exception)
2. explicitly skip shadowing of `torch.cat` (since it's not supported anyways)

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

Reviewed By: hx89

Differential Revision: D36017356

Pulled By: vkuzo

fbshipit-source-id: 0da4840a62c2dac183f8294c2cec4fce262474b3
(cherry picked from commit 88409c1)

Author: vkuzo

test/quantization/fx/test_numeric_suite_fx.py

@@ -1171,33 +1171,6 @@ def forward(self, x):
             # that shadowing models with method calls does not crash.
             results_len=0)
 
-    @skipIfNoFBGEMM
-    def test_add_shadow_loggers_multiple_dtype_casts(self):
-        """
-        Verifies that for nodes where the first input arg is a list,
-        such as `cat`, we insert an individual dtype cast for each
-        arg of the list.
-        """
-        class M(nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, x):
-                x = torch.cat([x, x, x], dim=0)
-                return x
-
-        m = M().eval()
-        expected_occurrence = {
-            # 3 dequantize function calls from the 3 dtype casts for [x, x, x]
-            ns.call_module(torch.nn.Identity): 3,
-            # 1 dequantize method call for module output
-            ns.call_method("dequantize"): 1,
-        }
-        self._test_match_shadow_activations(
-            m, (torch.randn(4, 4),),
-            prepared_expected_node_occurrence=expected_occurrence,
-            results_len=1, compare_fp32_vs_fp32_prepared=False)
-
     @skipIfNoFBGEMM
     def test_shadow_activations_fqn(self):
         m = nn.Sequential(
@@ -1237,7 +1210,7 @@ def forward(self, x):
         m = M().eval()
         self._test_match_shadow_activations(
             m, (torch.randn(1, 1, 4, 4),),
-            results_len=2,
+            results_len=1,
             should_log_inputs=True)
 
     @skipIfNoFBGEMM
@@ -1954,13 +1927,27 @@ def test_fp16_shadows_fp32(self):
         mq = convert_fx(mp, is_reference=True)
         mq_shadows_m = add_shadow_loggers('a', mq, 'b', m, OutputLogger)
 
-    def test_mul_add_skips_shadowing(self):
+    def test_mul_add_cat_stack_skips_shadowing(self):
         class M(nn.Module):
             def forward(self, x):
                 x = x * x
                 x = torch.mul(x, x)
                 x = x + x
                 x = torch.add(x, x)
+                x = torch.cat([x])
+                x = torch.stack([x])
+                return x
+
+        m = M().eval()
+        self._test_match_shadow_activations(
+            m, (torch.randn(1, 1, 4, 4),),
+            results_len=0)
+
+    def test_op_with_only_kwargs_skips_shadowing(self):
+        class M(nn.Module):
+            def forward(self, x):
+                x = torch.cat(tensors=[x])
+                x = torch.stack(tensors=[x])
                 return x
 
         m = M().eval()
@@ -2101,7 +2088,7 @@ def test_sparsenn_shadow(self):
             x = torch.randn(2, 4)
             self._test_match_shadow_activations(
                 sparse_nn, (idx, offsets, x),
-                results_len=4,
+                results_len=3,
                 should_log_inputs=should_log_inputs)
 
     @skip_if_no_torchvision

torch/ao/ns/fx/graph_passes.py

@@ -602,6 +602,26 @@ def load_arg(a):
                 subgraph_a, ref_name, ref_node_type_a, ref_node_type_b = \
                     end_node_b_to_matched_subgraph_a_and_name[node_b]
 
+            if len(node_b.args) == 0:
+                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)}' +
+                    ', kwargs-only node not handled yet')
+                env_c[node_b.name] = graph_c.node_copy(node_b, load_arg)
+                continue
+
+            all_op_types_support_shadowing = (
+                op_type_supports_shadowing(subgraph_a.start_node) and
+                op_type_supports_shadowing(node_b)
+            )
+            if not all_op_types_support_shadowing:
+                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)}' +
+                    ', unsupported')
+                env_c[node_b.name] = graph_c.node_copy(node_b, load_arg)
+                continue
+
             # For both start_node and end_node verify that we know how to do
             # the dtype cast. If we do not, skip.
             node_input_type_a, node_output_type_a = \
@@ -626,18 +646,6 @@ def load_arg(a):
                 env_c[node_b.name] = graph_c.node_copy(node_b, load_arg)
                 continue
 
-            all_op_types_support_shadowing = (
-                op_type_supports_shadowing(subgraph_a.start_node) and
-                op_type_supports_shadowing(node_b)
-            )
-            if not all_op_types_support_shadowing:
-                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)}' +
-                    ', unsupported')
-                env_c[node_b.name] = graph_c.node_copy(node_b, load_arg)
-                continue
-
             # If we are shadowing from fp32 to int8, we need to insert
             # quantize_per_tensor call with qparams from the previous node.
             # Only do this if we are able to infer these qparams from the graph.

torch/ao/ns/fx/utils.py

@@ -492,7 +492,7 @@ def compute_cosine_similarity(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
 def op_type_supports_shadowing(node: Node) -> bool:
     if node.op == 'call_function':
-        if node.target in (torch.add, torch.mul, operator.add, operator.mul):
-            # shadowing for ops with two inputs is not implemented yet
+        if node.target in (torch.add, torch.mul, operator.add, operator.mul, torch.cat, torch.stack):
+            # shadowing for ops with multiple tensor inputs is not implemented yet
             return False
     return True