Changing logic to deal with graphs with derived quantization spec (#16357)

Summary:
Pull Request resolved: #16357

We want to add a test for `default_addmm_A8W8` to fully finish testing `CadenceDefaultQuantizer`. However there are a couple changes we need to make to the testing function.

## Change 1: We allow passing `None` in the vec of `QuantizationSpec`

This is because  the addmm op has 3 inputs: `bias`, `mat1`, `mat2`. The bias uses a `DerivedQuantizationSpec`, which is dynamically constructed with references to the actual graph nodes (`mat1` and `mat2`). We can't construct an identical `DerivedQuantizationSpec` in the test because we'd need to reference the exact same node objects that the quantizer creates internally. Since we can't compare it directly, we use `None` to skip validation for that input. If `mat1` and `mat2` are quantized correctly, the derived bias spec will be correct too.

https://www.internalfb.com/code/fbsource/[2cfdb40fd8b628da2f46366115516408cfb9f50f]/xplat/executorch/backends/cadence/aot/quantizer/patterns.py?lines=91-103

## Change 2: We changed how we iterate through `input_qspec_map`

 `input_qspec_map` is a dictionary mapping input nodes to their `qspecs`. The iteration order depends on insertion order, which follows how the quantizer processes `PartitionAnchors`.

Each `QuantizationPattern` implements a `get_anchors()` method that returns a `PartitionAnchors` describing which arguments are inputs, weights, biases and nodes.  This is relevant because for `addmm`, the `PartitionAnchors` lists them as `inputs=[(node, 1)], weights=[(node, 2)], biases=[(node,  0, ...)]. ` So the map might iterate in order `mat1, mat2, bias` (args indices 1, 2, 0) rather than `bias, mat1, mat2` (args indices 0, 1, 2).

This means that our previous way of iterating wouldn't work. Thus, we now use the following way to iterate:

```
for input_node, input_qspec in annotation.input_qspec_map.items():
      // Find the index of this input node in the op's args
      arg_index = None
      for i, arg in enumerate(op_node.args):
          if arg is input_node:
              arg_index = i
              break
      self.assertIsNotNone(
          arg_index,
          f"Input node {input_node} not found in op_node.args",
      )
      # Skip comparison if expected qspec is None (e.g., for DerivedQuantizationSpec)
      if expected_input_qspecs[arg_index] is not None:
          self.assertEqual(
              input_qspec,
              expected_input_qspecs[arg_index],
              f"Input qspec mismatch at arg index {arg_index}",
          )
```
The new code looks up which argument index each input_node corresponds to by searching through `op_node.args`, rather than assuming the enumeration index i matches the argument position.

Reviewed By: hsharma35

Differential Revision: D88955761

Author: RahulC7

backends/cadence/aot/tests/test_quantizer_ops.py

@@ -61,14 +61,15 @@
 # Test case definitions for quantizer annotation tests.
 # Format: (name, graph_builder_fn, quantizer_instance, target_op, expected_output_qspec, expected_input_qspecs)
 # Adding a new quantizer test only requires adding a tuple to this list.
+# Note: Use None in expected_input_qspecs to skip comparison for that input (e.g., for DerivedQuantizationSpec).
 QUANTIZER_ANNOTATION_TEST_CASES: list[
     tuple[
         str,
         GraphBuilderFn,
         CadenceQuantizer,
         OpOverload,
         QuantizationSpec,
-        list[QuantizationSpec],
+        list[QuantizationSpec | None],
     ]
 ] = [
     (
@@ -189,6 +190,16 @@
         # For relu: only input_activation
         [qconfig_A8W8.input_activation],
     ),
+    (
+        "default_addmm_A8W8",
+        lambda self: self._build_addmm_graph(),
+        CadenceDefaultQuantizer(),
+        torch.ops.aten.addmm.default,
+        qconfig_A8W8.output_activation,
+        # For addmm: [bias (DerivedQuantizationSpec), mat1, mat2]
+        # Use None to skip comparison for bias since it's a DerivedQuantizationSpec
+        [None, qconfig_A8W8.input_activation, qconfig_A8W8.weight],
+    ),
 ]
 
 # Derive the set of tested quantizer classes from the test cases.
@@ -406,6 +417,31 @@ def _build_relu_graph(self) -> tuple[torch.fx.GraphModule, torch.fx.Node]:
         self.assertEqual(len(relu_nodes), 1, "Should find exactly one relu node")
         return gm, relu_nodes[0]
 
+    def _build_addmm_graph(self) -> tuple[torch.fx.GraphModule, torch.fx.Node]:
+        """Build a simple graph with an addmm operation."""
+        builder = GraphBuilder()
+        # addmm: bias + (mat1 @ mat2)
+        # args: (bias, mat1, mat2)
+        bias = builder.placeholder("bias", torch.randn(5))
+        mat1 = builder.placeholder("mat1", torch.randn(1, 10))
+        mat2 = builder.placeholder("mat2", torch.randn(10, 5))
+        addmm = builder.call_operator(
+            op=torch.ops.aten.addmm.default,
+            args=(bias, mat1, mat2),
+            meta=NodeMetadata(
+                {"source_fn_stack": [("addmm", torch.ops.aten.addmm.default)]}
+            ),
+        )
+        builder.output([addmm])
+        gm = builder.get_graph_module()
+
+        addmm_nodes = gm.graph.find_nodes(
+            op="call_function",
+            target=torch.ops.aten.addmm.default,
+        )
+        self.assertEqual(len(addmm_nodes), 1, "Should find exactly one addmm node")
+        return gm, addmm_nodes[0]
+
     @parameterized.expand(QUANTIZER_ANNOTATION_TEST_CASES)
     def test_quantizer_annotation(
         self,
@@ -414,7 +450,7 @@ def test_quantizer_annotation(
         quantizer: CadenceQuantizer,
         target: OpOverload,
         expected_output_qspec: QuantizationSpec,
-        expected_input_qspecs: list[QuantizationSpec],
+        expected_input_qspecs: list[QuantizationSpec | None],
     ) -> None:
         """Parameterized test for quantizer annotations."""
         gm, op_node = graph_builder_fn(self)
@@ -429,21 +465,24 @@ def test_quantizer_annotation(
 
         # Verify input annotations
         self.assertEqual(len(annotation.input_qspec_map), len(expected_input_qspecs))
-        for i, (input_node, input_qspec) in enumerate(
-            annotation.input_qspec_map.items()
-        ):
-            expected_arg = op_node.args[i]
-            assert isinstance(expected_arg, torch.fx.Node)
-            self.assertEqual(
-                input_node,
-                expected_arg,
-                f"Input node mismatch at index {i}",
-            )
-            self.assertEqual(
-                input_qspec,
-                expected_input_qspecs[i],
-                f"Input qspec mismatch at index {i}",
+        for input_node, input_qspec in annotation.input_qspec_map.items():
+            # Find the index of this input node in the op's args
+            arg_index = None
+            for i, arg in enumerate(op_node.args):
+                if arg is input_node:
+                    arg_index = i
+                    break
+            self.assertIsNotNone(
+                arg_index,
+                f"Input node {input_node} not found in op_node.args",
             )
+            # Skip comparison if expected qspec is None (e.g., for DerivedQuantizationSpec)
+            if expected_input_qspecs[arg_index] is not None:
+                self.assertEqual(
+                    input_qspec,
+                    expected_input_qspecs[arg_index],
+                    f"Input qspec mismatch at arg index {arg_index}",
+                )
 
     def test_all_quantizers_have_annotation_tests(self) -> None:
         """Ensure every CadenceQuantizer subclass is either tested or explicitly excluded."""