Stop incorrect encoding propagation through non-grid-preserving ops

Signed-off-by: Kyunggeun Lee <kyunggeu@qti.qualcomm.com>

Author: kyunggeunlee

TrainingExtensions/torch/src/python/aimet_torch/experimental/quantsim_utils.py

@@ -125,8 +125,10 @@ def _set_src_qtzr(x: Product, consumer: Op, qtzr: AffineQuantizerBase):
             if i < len(qmodule.output_quantizers) and qmodule.output_quantizers[i]:
                 qmodule.output_quantizers[i] = qtzr
 
-        if not qmodule or _is_math_invariant_op(qmodule):
-            # 1. There is no qmodule associated with the graph node ``producer``, or
+        if (not qmodule and producer.is_grid_preserving_op()) or _is_math_invariant_op(
+            qmodule
+        ):
+            # 1. There is no qmodule associated with the grid-preserving ``producer``, or
             # 2. qmodule is a math invariant op (reshape, permute, etc).
             # In these cases, propagate encoding further to the ancestors
             for input in producer.inputs:

TrainingExtensions/torch/test/python/v2/quantsim/test_quantsim.py

@@ -1500,6 +1500,34 @@ def test_htp_interpolation_tie_encodings(self, module_factory, input_factory):
         out_enc.pop("name")
         assert inp_enc == out_enc
 
+    def test_tie_encodings_functional_add(self):
+        """
+        Given: Functional operator between Conv and Relu
+        When: Create quantsim with HTP V81 config
+        Then: Conv and Relu output quantizers should NOT be shared
+        """
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv1 = torch.nn.Conv2d(3, 3, 3)
+                self.conv2 = torch.nn.Conv2d(3, 3, 3)
+                self.relu = torch.nn.ReLU()
+
+            def forward(self, x):
+                x1 = self.conv1(x)
+                x2 = self.conv2(x)
+                return self.relu(x1 + x2)
+
+        model = Model()
+        x = torch.randn(1, 3, 10, 10)
+        sim = QuantizationSimModel(model, x, config_file="htp_v81")
+
+        # There should be no shared quantizers
+        assert list(sim.model.modules(remove_duplicate=True)) == list(
+            sim.model.modules(remove_duplicate=False)
+        )
+
     def test_conv_relu_supergroup(self, tmp_path: pathlib.Path):
         """
         When: Create quantsim with HTP V69 config or lower