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()