Skip to content

Don't Fold In Negative Scalar Multiplies Into Quant/Dequant (#23261) #23261

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
RahulC7 wants to merge 1 commit into from
Closed

Don't Fold In Negative Scalar Multiplies Into Quant/Dequant (#23261) #23261

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions backends/cadence/aot/fuse_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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}"
Expand Down Expand Up @@ -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]}"
Expand Down
126 changes: 126 additions & 0 deletions backends/cadence/aot/tests/test_fusion_ops_passes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down