Don't Fold In Negative Scalar Multiplies Into Quant/Dequant

Summary:
## Context
The way we lowered the GDN model is we looked at the traced model, and then recreated "stages" of it(i.e created a graph for that stage), and lowered each one of those stages, making them go through the entire compiler pipeline. 

That’s usually fine for the ops/nodes in the middle, because they still have the surrounding nodes(producer consumer chain) and those get collapsed/changed in the same way, but in the boundaries, where there’s connections to other other nodes not in the subgraph, the lowering might end up producing a different result than what we expect in the actual graph. The quant/dequant insertion, calibration and FuseMul adjacency depend on the global context so it’ll look different. 

That's what happens here. 

## Problem

In D117100989, we do `N_c = -A_c` as `aten.full(-1.0) + aten.mul.Tensor`(elementwise kernel). However, before the multiply, there's a quantize, and so multiply gets combined with the quantize. The issue is that this is a negative, and so we're going to have a negative scale in the quantize/dequantize scales, because of `FuseMulTensorIntoDequantPass` and `FuseMulScalarIntoDequantPass` folding it in. 





## What this diff does

- Make all three multiply fusion passes stop when the derived scale is not positive.
- Leave the original multiply in the graph when fusion is not valid.
- Add focused tests for tensor-to-quant, tensor-from-dequant, and scalar-from-dequant with multipliers `0.0` and `-1.0`.
- Give the chunk-32 GDN test a test-only gate offset of `-5.0`. This keeps the 32-step cumulative `Gamma_c` values in the useful int16 calibration range. Chunk size 2 stays unchanged.
- Use measured tolerances for the three equivalent chunk-32 pre-trace transforms.

## Result

The full chunk-32 graph keeps one explicit `aten.mul.Tensor` with `aten.full(-1.0)` for `N_c = -A_c`. All quantize, dequantize, requantize, and packed prep scales in the dumped graph are finite and positive.

Reviewed By: ThomasJannaud

Differential Revision: D120871768

Author: RahulC7

backends/cadence/aot/fuse_ops.py

@@ -839,7 +839,10 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         assert isinstance(dequant_node.args[1], Number)
         assert isinstance(mul_node.args[1], Number)
         # pyre-ignore[58]: Unsupported operand *
-        new_deq_args[1] = dequant_node.args[1] * mul_node.args[1]
+        new_scale = dequant_node.args[1] * mul_node.args[1]
+        if new_scale <= 0:
+            return False
+        new_deq_args[1] = new_scale
 
         # Replace all uses of mul with the dequant node
         mul_node.replace_all_uses_with(dequant_node)
@@ -925,6 +928,8 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         if mul_scalar == 0 or old_scale == 0:
             return False
         new_scale = float(old_scale) / float(mul_scalar)
+        if new_scale <= 0:
+            return False
 
         logging.debug(
             f"Fused {node} and {full_node} into {quant_node}. Updated scale from {quant_node.args[1]} to {new_scale}"
@@ -984,7 +989,10 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         assert isinstance(deq_node.args[1], Number)
         assert isinstance(mplier_node.args[1], Number)
         # pyre-ignore[58]: Unsupported operand *
-        new_deq_args[1] = deq_node.args[1] * mplier_node.args[1]
+        new_scale = deq_node.args[1] * mplier_node.args[1]
+        if new_scale <= 0:
+            return False
+        new_deq_args[1] = new_scale
 
         logging.debug(
             f"Fused {node} and {mplier_node} into {deq_node}. Updated scale from {deq_node.args[1]} to {new_deq_args[1]}"

backends/cadence/aot/tests/test_fusion_ops_passes.py

@@ -792,6 +792,51 @@ def test_fuse_mul_into_dequant(self) -> None:
                 deq_scale = node.args[1]
         self.assertEqual(deq_scale, DEQUANT_SCALE * FULL_VALUE)
 
+    @parameterized.expand([("zero", 0.0), ("negative", -1.0)])
+    def test_keep_non_positive_mul_tensor_after_dequant(
+        self, _name: str, mul_value: float
+    ) -> None:
+        input_shape: Final[List[int]] = [4, 32]
+        dequant_scale = 1.5
+
+        builder = GraphBuilder()
+        x = builder.placeholder(
+            "x", torch.randint(0, 255, input_shape, dtype=torch.uint8)
+        )
+        dequant = builder.call_operator(
+            op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x, dequant_scale, 0, 0, 255, torch.uint8),
+        )
+        full = builder.call_operator(
+            op=exir_ops.edge.aten.full.default,
+            args=(input_shape, mul_value),
+        )
+        mul = builder.call_operator(
+            op=exir_ops.edge.aten.mul.Tensor,
+            args=(dequant, full),
+        )
+        builder.output([mul])
+
+        result = cast(
+            PassResult,
+            FuseMulTensorIntoDequantPass()(builder.get_graph_module()),
+        )
+
+        self.assertFalse(result.modified)
+        self.check_op_counts(
+            result.graph_module,
+            expected_op_counts={
+                exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: 1,
+                exir_ops.edge.aten.full.default: 1,
+                exir_ops.edge.aten.mul.Tensor: 1,
+            },
+        )
+        dequant_nodes = result.graph_module.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+        )
+        self.assertEqual(dequant_nodes[0].args[1], dequant_scale)
+
     def test_fuse_mul_into_dequant_no_match(self) -> None:
         """
         Test that FuseMulTensorIntoDequantPass does NOT modify the graph
@@ -880,6 +925,45 @@ def test_fuse_mul_scalar_into_dequant(self) -> None:
                 deq_scale = node.args[1]
         self.assertEqual(deq_scale, dequant_scale * mul_value)
 
+    @parameterized.expand([("zero", 0.0), ("negative", -1.0)])
+    def test_keep_non_positive_mul_scalar_after_dequant(
+        self, _name: str, mul_value: float
+    ) -> None:
+        dequant_scale = 0.006
+
+        builder = GraphBuilder()
+        x = builder.placeholder(
+            "x", torch.randint(-128, 127, (2, 3, 4), dtype=torch.int8)
+        )
+        dequant = builder.call_operator(
+            op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x, dequant_scale, 5, -128, 127, torch.int8),
+        )
+        mul = builder.call_operator(
+            op=exir_ops.edge.aten.mul.Scalar,
+            args=(dequant, mul_value),
+        )
+        builder.output([mul])
+
+        result = cast(
+            PassResult,
+            FuseMulScalarIntoDequantPass()(builder.get_graph_module()),
+        )
+
+        self.assertFalse(result.modified)
+        self.check_op_counts(
+            result.graph_module,
+            expected_op_counts={
+                exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: 1,
+                exir_ops.edge.aten.mul.Scalar: 1,
+            },
+        )
+        dequant_nodes = result.graph_module.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+        )
+        self.assertEqual(dequant_nodes[0].args[1], dequant_scale)
+
     def test_fuse_mul_into_quant(self) -> None:
         quant_scale = 5
         mul_value = 10
@@ -931,6 +1015,48 @@ def test_fuse_mul_into_quant(self) -> None:
             new_quant_scale = node.args[1]
             self.assertEqual(new_quant_scale, quant_scale / mul_value)
 
+    @parameterized.expand([("zero", 0.0), ("negative", -1.0)])
+    def test_keep_non_positive_mul_tensor_before_quant(
+        self, _name: str, mul_value: float
+    ) -> None:
+        quant_scale = 5.0
+
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(4, 32, dtype=torch.float32))
+        full = builder.call_operator(
+            op=exir_ops.edge.aten.full.default,
+            args=([1], mul_value),
+        )
+        mul = builder.call_operator(
+            op=exir_ops.edge.aten.mul.Tensor,
+            args=(x, full),
+        )
+        quant = builder.call_operator(
+            op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(mul, quant_scale, 7, 0, 255, torch.uint8),
+        )
+        builder.output([quant])
+
+        result = cast(
+            PassResult,
+            FuseMulTensorIntoQuantPass()(builder.get_graph_module()),
+        )
+
+        self.assertFalse(result.modified)
+        self.check_op_counts(
+            result.graph_module,
+            expected_op_counts={
+                exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: 1,
+                exir_ops.edge.aten.full.default: 1,
+                exir_ops.edge.aten.mul.Tensor: 1,
+            },
+        )
+        quant_nodes = result.graph_module.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        )
+        self.assertEqual(quant_nodes[0].args[1], quant_scale)
+
     def test_fuse_then_transpose_pass(self) -> None:
         # Create a graph with full -> transpose -> permute -> view.
         builder = GraphBuilder()