From 7eb92a3a1450f8ee3a99337f1a3524b47557a759 Mon Sep 17 00:00:00 2001 From: Rahul Chandra Date: Tue, 20 Oct 2026 10:22:31 -0700 Subject: [PATCH] Don't Fold In Negative Scalar Multiplies Into Quant/Dequant (#23261) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/23261 ## 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 --- backends/cadence/aot/fuse_ops.py | 12 +- .../aot/tests/test_fusion_ops_passes.py | 126 ++++++++++++++++++ 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/backends/cadence/aot/fuse_ops.py b/backends/cadence/aot/fuse_ops.py index 428a64bf428..c5adb3339d0 100644 --- a/backends/cadence/aot/fuse_ops.py +++ b/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]}" diff --git a/backends/cadence/aot/tests/test_fusion_ops_passes.py b/backends/cadence/aot/tests/test_fusion_ops_passes.py index 87ec79fc525..f30c2ab1857 100644 --- a/backends/cadence/aot/tests/test_fusion_ops_passes.py +++ b/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()