Update the qspec sharing logic for concat op

src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py

@@ -886,9 +886,12 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
     def generate_qspec_sharing_constraints(
         cls, node: torch.fx.Node, qspecs: ProvisionalQSpecMap
     ) -> list[Constraint]:
-        """Inputs must observe together unless the output is per-channel along the
-        concat axis, in which case each keeps its own scale and they need only
-        agree on ``DTYPE``.
+        """Return the constraints tying the cat's input and output quantizers together.
+
+        They share one quantizer only if every slot (inputs and output) is per-tensor, or
+        per-channel along a non-concat axis. Otherwise each input keeps its own quantizer,
+        since the inputs can differ in size along the concat axis, and they only agree on
+        ``DTYPE``.
         """
         output_slot, input_slots = _shared_op_boundary_slots(node)
         if not _any_input_slot_populated(input_slots, qspecs):
@@ -897,14 +900,15 @@ def generate_qspec_sharing_constraints(
         all_slots = frozenset([output_slot, *input_slots])
 
         concat_dim = _concat_dim(node)
-        output_fields = qspecs.get(output_slot, ProvisionalQSpec()).fields
-        output_granularity = output_fields.get(FieldName.GRANULARITY)
-
-        # A per-tensor granularity has ``axis is None``.
-        granularity_axis = output_granularity.value.axis if output_granularity is not None else None
-        per_channel_along_concat_axis = granularity_axis is not None and _same_axis(
-            granularity_axis, concat_dim, node
-        )
+        per_channel_along_concat_axis = False
+        for slot in all_slots:
+            granularity = qspecs.get(slot, ProvisionalQSpec()).fields.get(FieldName.GRANULARITY)
+            if granularity is None or granularity.value is None:
+                continue
+            # A per-tensor granularity has ``axis is None``.
+            granularity_axis = granularity.value.axis
+            if granularity_axis is not None and _same_axis(granularity_axis, concat_dim, node):
+                per_channel_along_concat_axis = True
         if not per_channel_along_concat_axis:
             # Sharing an instance subsumes sharing fields, since the merged group
             # reconciles all of them.

tests/quantization/test_annotation_pattern_registry.py

@@ -402,38 +402,23 @@ def forward(self, x):
         return torch.cat([x, xt], dim=0)  # concat along batch, not channel/spatial
 
 
-class SimpleCatModel(nn.Module):
-    """Simple model with concat operation."""
+class CatOfParametersModel(nn.Module):
+    """Concat of two parameters with different sizes along the concat axis.
 
-    def __init__(self):
-        super().__init__()
-        self.linear1 = torch.nn.Linear(2, 2, bias=False)
-        self.linear2 = torch.nn.Linear(2, 2, bias=False)
-        self.linear3 = torch.nn.Linear(2, 2, bias=False)
-
-    def forward(self, x1, x2):
-        x1 = self.linear1(x1)
-        x2 = self.linear2(x2)
-        x = torch.cat([x1, x2])
-        x = self.linear3(x)
-        return x
-
-
-class SimpleConcatModel(nn.Module):
-    """Simple model with concat operation."""
+    Similar to token concat done in models such as SAM 2's mask decoder:
+    ``torch.cat([iou_token.weight, mask_tokens.weight], dim=0)``.
+    """
 
     def __init__(self):
         super().__init__()
-        self.linear1 = torch.nn.Linear(2, 2, bias=False)
-        self.linear2 = torch.nn.Linear(2, 2, bias=False)
-        self.linear3 = torch.nn.Linear(2, 2, bias=False)
+        self.token1 = nn.Embedding(1, 4)
+        self.token2 = nn.Embedding(3, 4)
+        self.linear = nn.Linear(4, 4)
 
-    def forward(self, x1, x2):
-        x1 = self.linear1(x1)
-        x2 = self.linear2(x2)
-        x = torch.concat([x1, x2])
-        x = self.linear3(x)
-        return x
+    def forward(self, x):
+        tokens = torch.cat([self.token1.weight, self.token2.weight], dim=0)
+        out = self.linear(tokens + x)
+        return out
 
 
 class AddLinearMulModel(nn.Module):
@@ -1981,69 +1966,127 @@ def test_shared_observer_preserves_per_block_when_pool_preserves_shape(
         expected_scale_shape = (2, 4, *spatial[:-1], spatial[-1] // 2)
         assert tuple(input_fq.qparams_calculator.scale.shape) == expected_scale_shape
 
-    def test_cat(self):
+    @pytest.mark.parametrize("cat_fn", [torch.cat, torch.concat], ids=["cat", "concat"])
+    def test_cat_per_tensor_shares_quantizers(self, cat_fn):
         """
         Given a model with
 
-        input1 -> relu1
-                        \
-                          ->  Cat -> Linear
-        input2 -> relu2 /
+        input1 -> linear1
+                          \
+                            -> cat -> linear3
+        input2 -> linear2 /
 
-        Check that concat inputs and outputs share the same quantizers.
+        Check that with per-tensor quantization, the cat's inputs and output share one
+        quantizer, while linear3's output has its own.
+
+        torch.cat and torch.concat export to different aten ops (aten.cat, aten.concat),
+        each matched by its own entry in ConcatPattern.
         """
-        model = SimpleCatModel()
-        inp = (torch.randn(1, 2), torch.randn(1, 2))
+
+        class CatModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear1 = nn.Linear(2, 2, bias=False)
+                self.linear2 = nn.Linear(2, 2, bias=False)
+                self.linear3 = nn.Linear(2, 2, bias=False)
+
+            def forward(self, x1, x2):
+                x = cat_fn([self.linear1(x1), self.linear2(x2)])
+                out = self.linear3(x)
+                return out
 
         config = QuantizerConfig(
-            global_config=ModuleQuantizerConfig(op_state_spec=None, op_input_spec=None),
+            global_config=ModuleQuantizerConfig(
+                op_state_spec=None,
+                op_input_spec=None,
+                op_output_spec={"*": per_tensor_int8_qspec()},
+            ),
         )
+        quantizer = Quantizer(CatModel(), config)
+        prepared_model = quantizer.prepare(example_inputs=(torch.randn(1, 2), torch.randn(1, 2)))
 
-        quantizer = Quantizer(model, config)
-        prepared_model = quantizer.prepare(example_inputs=inp)
+        cat_nodes = [
+            node
+            for node in prepared_model.graph.nodes
+            if node.target in (torch.ops.aten.cat.default, torch.ops.aten.concat.default)
+        ]
+        assert len(cat_nodes) == 1
+        cat = cat_nodes[0]
+        cat_input1_quantizer, cat_input2_quantizer = (
+            getattr(prepared_model, input_node.target) for input_node in cat.args[0]
+        )
+        (cat_output_quantizer_node,) = cat.users
+        cat_output_quantizer = getattr(prepared_model, cat_output_quantizer_node.target)
+        (linear3,) = cat_output_quantizer_node.users
+        (linear3_output_quantizer_node,) = linear3.users
+        linear3_output_quantizer = getattr(prepared_model, linear3_output_quantizer_node.target)
+
+        assert cat_input1_quantizer is cat_input2_quantizer
+        assert cat_input1_quantizer is cat_output_quantizer
+        assert linear3_output_quantizer is not cat_output_quantizer
+        # No other quantizers: the one the cat shares, and linear3's output quantizer.
         quantizers = [
             module
             for module in prepared_model.modules()
             if isinstance(module, _DefaultFakeQuantizeImpl)
         ]
         assert len(quantizers) == 2
-        assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_1
-        assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_2
 
-        # Final linear output quantizer is not associated with the others
-        assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3
+    def test_cat_inputs_per_channel_along_concat_axis(self):
+        """
+        Concat inputs that are per-channel along the concat axis keep separate quantizers.
 
-    def test_concat(self):
+        The inputs are parameters named ``weight``, so the default config gives them the
+        per-channel (axis 0) weight spec, while the output is per-tensor. The inputs differ
+        in size along axis 0 ([1, 4] and [3, 4]), so one shared per-channel quantizer can't
+        observe both.
         """
-        Given a model with
+        model = CatOfParametersModel().eval()
+        quantizer = Quantizer(model, QuantizerConfig())
+        prepared_model = quantizer.prepare(example_inputs=(torch.randn(4, 4),))
 
-        input1 -> relu1
-                        \
-                          ->  Concat -> Linear
-        input2 -> relu2 /
+        cat_nodes = [
+            node for node in prepared_model.graph.nodes if node.target == torch.ops.aten.cat.default
+        ]
+        assert len(cat_nodes) == 1
+        cat = cat_nodes[0]
+        token1_quantizer, token2_quantizer = (
+            getattr(prepared_model, input_node.target) for input_node in cat.args[0]
+        )
+        assert token1_quantizer is not token2_quantizer
 
-        Check that concat inputs and outputs share the same quantizers.
+    def test_cat_inputs_per_channel_along_other_axis(self):
         """
-        model = SimpleConcatModel()
-        inp = (torch.randn(1, 2), torch.randn(1, 2))
+        Concat inputs that are per-channel along a non-concat axis share one quantizer.
 
+        The inputs ([1, 4] and [3, 4], concatenated along axis 0) are quantized per-channel
+        along axis 1, where they have the same size, so one quantizer can observe both.
+        """
+        model = CatOfParametersModel().eval()
+        weight_spec = QuantizationSpec(
+            dtype=torch.int8,
+            qscheme=QuantizationScheme.SYMMETRIC,
+            granularity=PerChannelGranularity(axis=1),
+        )
         config = QuantizerConfig(
-            global_config=ModuleQuantizerConfig(op_state_spec=None, op_input_spec=None),
+            global_config=ModuleQuantizerConfig(
+                op_state_spec={"weight": weight_spec},
+                op_input_spec={"*": per_tensor_int8_qspec()},
+                op_output_spec={"*": per_tensor_int8_qspec()},
+            ),
         )
-
         quantizer = Quantizer(model, config)
-        prepared_model = quantizer.prepare(example_inputs=inp)
-        quantizers = [
-            module
-            for module in prepared_model.modules()
-            if isinstance(module, _DefaultFakeQuantizeImpl)
-        ]
-        assert len(quantizers) == 2
-        assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_1
-        assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_2
+        prepared_model = quantizer.prepare(example_inputs=(torch.randn(4, 4),))
 
-        # Final linear output quantizer is not associated with the others
-        assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3
+        cat_nodes = [
+            node for node in prepared_model.graph.nodes if node.target == torch.ops.aten.cat.default
+        ]
+        assert len(cat_nodes) == 1
+        cat = cat_nodes[0]
+        token1_quantizer, token2_quantizer = (
+            getattr(prepared_model, input_node.target) for input_node in cat.args[0]
+        )
+        assert token1_quantizer is token2_quantizer
 
     @pytest.mark.parametrize(
         "granularity",