Arm backend: Fix bug in decompose_linear_pass (pytorch#13725)

The final reshape node of linear decomposition got incorrect
quantization parameters which lead to incorrect outputs when that
reshape could be fused at a later stage in the pass pipeline.

Signed-off-by: Oscar Andersson <[email protected]>

Author: oscarandersson8218

backends/arm/_passes/decompose_linear_pass.py

@@ -90,6 +90,13 @@ def call(self, graph_module):
                     kwargs={},
                     from_node=node,
                 )
+                # Quantization parameters are inherited from original linear node, but
+                # output reshape should use the linear node's output qparams for both input
+                # and output.
+                if "input_qparams" in output.meta:
+                    output.meta["input_qparams"] = output.meta.get(
+                        "output_qparams", None
+                    )
 
             node.replace_all_uses_with(output)
             graph_module.graph.erase_node(node)

backends/arm/test/passes/test_fuse_constant_ops_pass.py

@@ -12,7 +12,11 @@
     FuseConstantArgsPass,
 )
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import (
+    PassPipeline,
+    TosaPipelineFP,
+    TosaPipelineINT,
+)
 
 input_t = Tuple[torch.Tensor]  # Input x
 input_t2 = Tuple[torch.Tensor, torch.Tensor]
@@ -103,6 +107,22 @@ def forward(self, a, b):
         return torch.cat((a, b), dim=0)
 
 
+class LinearConst(torch.nn.Module):
+    """A linear layer that can be computed AOT"""
+
+    def __init__(self, in_out_features: int = 3, bias: bool = True):
+        super().__init__()
+        self.linear = torch.nn.Linear(in_out_features, in_out_features, bias=bias)
+        self.example_input = torch.rand(in_out_features, in_out_features)
+
+    def forward(self, x: torch.Tensor):
+        y = torch.full_like(x, 1.0)
+        return self.linear(y) + x
+
+    def get_example_input(self):
+        return self.example_input
+
+
 modules = {
     "fuse_parameter": FuseParameter(),
     "fuse_buffer": FuseBuffer(),
@@ -152,3 +172,30 @@ def test_fuse_const_ops_tosa_BI_cat(module: torch.nn.Module):
         passes_with_exported_program=[ComputeConstantOpsAOT, FuseConstantArgsPass],
     )
     pipeline.run()
+
+
+def test_linear_const_tosa_FP():
+    model = LinearConst()
+    example_input = model.get_example_input()
+    pipeline = TosaPipelineFP[input_t](
+        model,
+        (example_input,),
+        aten_op=[],
+        exir_op=[],
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()
+
+
+def test_linear_const_tosa_INT():
+    model = LinearConst()
+    example_input = model.get_example_input()
+    pipeline = TosaPipelineINT[input_t](
+        model,
+        (example_input,),
+        aten_op=[],
+        exir_op=[],
+        per_channel_quantization=False,
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()