Skip to content

Add fused_quant, a backend-neutral quantization frontend (#23091) #23091

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

Open
mcremon-meta wants to merge 1 commit into
base: main
Choose a base branch
from
+23,259 −2,061
Open

Add fused_quant, a backend-neutral quantization frontend (#23091) #23091

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
21 changes: 21 additions & 0 deletions backends/cadence/aot/BUCK
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,7 @@ fbcode_target(_kind = runtime.python_library,
"fuse_ops.py",
],
deps = [
"//executorch/backends/transforms:fuse_quant_dequant_to_requantize",
":compiler_utils",
":ops_registrations",
"//caffe2:torch",
Expand Down Expand Up @@ -284,6 +285,15 @@ fbcode_target(_kind = runtime.python_library,
"remove_ops.py",
],
deps = [
"//executorch/backends/transforms:remove_nop_as_strided_copy_op",
"//executorch/backends/transforms:remove_alias_copy_op",
"//executorch/backends/transforms:remove_branched_quant_dequant",
"//executorch/backends/transforms:remove_cat_from_slice_copy",
"//executorch/backends/transforms:remove_clone_ops_transform_imported",
"//executorch/backends/transforms:remove_nop_expand_op",
"//executorch/backends/transforms:remove_nop_slice_or_view_op",
"//executorch/backends/transforms:remove_permute_before_mean",
"//executorch/backends/transforms:remove_zero_sized_cat_args",
":fuse_ops",
":ops_registrations",
"//caffe2:torch",
Expand All @@ -305,6 +315,15 @@ fbcode_target(_kind = runtime.python_library,
"reorder_ops.py",
],
deps = [
":ops_registrations",
"//executorch/backends/transforms:advance_quantize_op_above_def_in_branch",
"//executorch/backends/transforms:advance_quantize_op_above_def_chain",
"//executorch/backends/transforms:postpone_dequantize_op_below_use_chain",
"//executorch/backends/transforms:move_permute_after_concat",
"//executorch/backends/transforms:move_slice_before_permute",
"//executorch/backends/transforms:move_slice_before_view",
"//executorch/backends/transforms:propagate_slice",
"//executorch/backends/transforms:split_dequantized_cat",
"//caffe2:torch",
"//executorch/backends/cadence/aot:compiler_utils",
"//executorch/backends/cadence/aot:pass_utils",
Expand All @@ -323,6 +342,8 @@ fbcode_target(_kind = runtime.python_library,
"replace_ops.py",
],
deps = [
"//executorch/backends/transforms:replace_select_with_view_op",
"//executorch/backends/transforms:replace_split_with_slice",
":pass_utils",
"//caffe2:torch",
"//executorch/backends/cadence/aot:compiler_utils",
Expand Down
226 changes: 16 additions & 210 deletions backends/cadence/aot/fuse_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
Expand Down Expand Up @@ -38,18 +38,16 @@
from executorch.backends.transforms.fuse_cascaded_view_ops import (
FuseCascadedViewOps as _SharedFuseCascadedViewOps,
)
from executorch.backends.transforms.fuse_quant_dequant_to_requantize import (
FuseQuantDequantToRequantizePass as _SharedFuseQuantDequantToRequantizePass,
)
from executorch.backends.transforms.fuse_transpose_or_permute_op_pairs_pass import (
FuseTransposeOrPermuteOpPairsPass as _SharedFuseTransposeOrPermuteOpPairsPass,
)
from executorch.backends.transforms.permute_pass_utils import (
FuseOpPairsAcrossBranchesPass,
)
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.dialects.edge._ops import EdgeOpOverload, EdgeOpOverloadPacket
from executorch.exir.pass_base import PassResult
from executorch.exir.passes.cse_pass import CSEPass
from torch.nn.utils.fusion import fuse_conv_bn_weights
from torch.utils import _pytree as pytree


def get_tensor_arg(node: torch.fx.Node, arg_name: str) -> torch.Tensor:
Expand Down Expand Up @@ -463,9 +461,9 @@

from executorch.exir.dialects.edge._ops import EdgeOpOverload

assert isinstance(
bn_node.target, EdgeOpOverload
), f"Expected EdgeOpOverload, got {type(bn_node.target)}"
assert isinstance(bn_node.target, EdgeOpOverload), (
f"Expected EdgeOpOverload, got {type(bn_node.target)}"
)

# Extract parameters by name (not by index for maintainability)
# The get_arg function handles normalization of positional args to kwargs
Expand Down Expand Up @@ -593,23 +591,7 @@
pass


class FuseQuantDequantToRequantizePass(FuseOpPairsAcrossBranchesPass):
"""
Fuse dequantize-quantize op pairs to a single requantize op.
For the special case where quant params match, this will remove
both dequant and quant ops.
"""

# A list of ops that can be bypassed when looking for a
# dequantize->quantize chain
bypass_ops: set[EdgeOpOverload] = {
exir_ops.edge.aten.slice_copy.Tensor,
exir_ops.edge.aten.view_copy.default,
exir_ops.edge.aten.clone.default,
exir_ops.edge.aten.transpose_copy.int,
exir_ops.edge.aten.permute_copy.default,
}

class FuseQuantDequantToRequantizePass(_SharedFuseQuantDequantToRequantizePass):
quantize_op_packets: set[EdgeOpOverloadPacket] = {
exir_ops.edge.cadence.quantize_per_tensor,
exir_ops.edge.quantized_decomposed.quantize_per_tensor,
Expand All @@ -618,192 +600,17 @@
exir_ops.edge.cadence.dequantize_per_tensor,
exir_ops.edge.quantized_decomposed.dequantize_per_tensor,
}
requantize_op: Optional[EdgeOpOverload] = (
exir_ops.edge.cadence.requantize.per_tensor
)

def __init__(
self, allow_requantize: bool = True, force_quant_dequant_fusion: bool = False
) -> None:
super().__init__()
self.allow_requantize: bool = allow_requantize
self.force_quant_dequant_fusion: bool = force_quant_dequant_fusion

def _pkg_name_match(self, node1: torch.fx.Node, node2: torch.fx.Node) -> bool:
# pyre-ignore[16]: Item `typing.Callable` has no attribute `_op`
return node1.target._op.namespace == node2.target._op.namespace

def can_fuse_for_chain(
self,
producer: torch.fx.Node,
consumer: torch.fx.Node,
consumer_op_packets: set[EdgeOpOverloadPacket],
) -> bool:
return super().can_fuse_for_chain(
producer, consumer, consumer_op_packets
) and self._pkg_name_match(producer, consumer)

def _create_requantize_node(
self,
in_tensor: torch.fx.Node,
in_scale: float,
in_zero_point: int,
out_scale: float,
out_zero_point: int,
out_dtype: torch.dtype,
graph: torch.fx.Graph,
) -> torch.fx.Node:
return graph.call_function(
exir_ops.edge.cadence.requantize.per_tensor,
args=(
in_tensor,
in_scale,
in_zero_point,
out_scale,
out_zero_point,
out_dtype,
),
)

def _quant_params_match(self, node1: torch.fx.Node, node2: torch.fx.Node) -> bool:
arg_names = ("scale", "zero_point", "quant_min", "quant_max", "dtype")
arg_types = (float, int, int, int, torch.dtype)
return all(
get_arg(node1, name, arg_type) == get_arg(node2, name, arg_type)
for name, arg_type in zip(arg_names, arg_types)
)

def check_ok_to_fuse(
self,
producer: torch.fx.Node,
consumers: list[torch.fx.Node],
) -> bool:
"""Check if all node-user pairs are nops or are ok to replace with requant."""
for rnode in consumers:
if self.allow_requantize or self._quant_params_match(producer, rnode):
# Cannot remove quant-dequant pair if quant params don't match and requantize
# is not allowed.
continue
return False
return True

def _get_bypassed_nodes(
self,
producer: torch.fx.Node,
removal_candidates: list[torch.fx.Node],
) -> list[torch.fx.Node]:
"""Return the bypassed ops between producer and consumers, in order."""
candidates = set(removal_candidates)
bypassed: set[torch.fx.Node] = set()
stack = list(producer.users)
while stack:
user = stack.pop()
if user in candidates or user in bypassed:
continue
bypassed.add(user)
stack.extend(user.users)

# The graph's node list is topologically sorted, so walking it forward
# from the producer refreshes each node after its predecessor.
ordered: list[torch.fx.Node] = []
cursor = producer.next
while bypassed and cursor.op != "root":
if cursor in bypassed:
bypassed.remove(cursor)
ordered.append(cursor)
cursor = cursor.next
return ordered

def fuse(
self,
node: torch.fx.Node,
removal_candidates: list[torch.fx.Node],
graph_module: torch.fx.GraphModule,
) -> None:
# Capture the bypassed ops before the rewire erases the chain's root.
bypassed_nodes = self._get_bypassed_nodes(node, removal_candidates)
node.replace_all_uses_with(cast(torch.fx.Node, node.args[0]))
graph_module.graph.erase_node(node)
# They now consume the producer's input, so their metadata describes an
# operand they no longer have. Refresh before get_fused_node reads it.
for bypassed in bypassed_nodes:
args, kwargs = pytree.tree_map_only(
torch.fx.Node,
lambda arg: arg.meta["val"],
(bypassed.args, bypassed.kwargs),
)
assert callable(bypassed.target)
bypassed.meta["val"] = bypassed.target(*args, **kwargs)
bypassed.meta["tensor_meta"] = None
for rnode in removal_candidates:
rnode.replace_all_uses_with(self.get_fused_node(node, rnode, graph_module))
graph_module.graph.erase_node(rnode)

def get_fused_node(
self,
producer: torch.fx.Node,
consumer: torch.fx.Node,
graph_module: torch.fx.GraphModule,
) -> torch.fx.Node:
in_scale = get_arg(producer, "scale", float)
in_zero_point = get_arg(producer, "zero_point", int)
out_scale = get_arg(consumer, "scale", float)
out_zero_point = get_arg(consumer, "zero_point", int)
out_dtype = get_arg(consumer, "dtype", torch.dtype)
if in_scale == out_scale and in_zero_point == out_zero_point:
# If the quant params match, we can remove both dequantize-quantize ops.
return cast(torch.fx.Node, consumer.args[0])

assert (
self.allow_requantize
), f"Found {producer=} {in_scale=} {in_zero_point=} | {consumer=} {out_scale=} {out_zero_point=}"

with graph_module.graph.inserting_before(consumer):
requantize_node = self._create_requantize_node(
in_tensor=cast(torch.fx.Node, consumer.args[0]),
in_scale=in_scale,
in_zero_point=in_zero_point,
out_scale=out_scale,
out_zero_point=out_zero_point,
out_dtype=out_dtype,
graph=graph_module.graph,
)
requantize_node.meta = consumer.meta.copy()
requantize_node.meta["val"] = exir_ops.edge.cadence.requantize.per_tensor(
cast(torch.fx.Node, consumer.args[0]).meta["val"],
cast(float, in_scale),
cast(int, in_zero_point),
cast(float, out_scale),
cast(int, out_zero_point),
cast(torch.dtype, out_dtype),
)
requantize_node.meta["tensor_meta"] = None
return requantize_node

def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
# Remove any dequantize op that has only quantize ops as its users.
modified = self.find_and_fuse(
graph_module,
producer_op_packets=self.dequantize_op_packets,
consumer_op_packets=self.quantize_op_packets,
bypass_ops=self.bypass_ops,
)
# Remove any quantize op that has only dequantze ops as its users.
modified |= self.find_and_fuse(
graph_module,
producer_op_packets=self.quantize_op_packets,
consumer_op_packets=self.dequantize_op_packets,
# Do not requantize for quantize-dequantize pairs as this is not guaranteed
# to be better for performance/memory.
# Only fuse if all users of quant are dequant.
bypass_ops=(
self.bypass_ops
if self.force_quant_dequant_fusion
else {exir_ops.edge.aten.view_copy.default}
),
super().__init__(
allow_requantize=allow_requantize,
force_quant_dequant_fusion=force_quant_dequant_fusion,
)
if modified:
graph_module.graph.eliminate_dead_code()
graph_module.recompile()
return PassResult(graph_module, True)
return PassResult(graph_module, False)


class FuseMulScalarIntoDequantPass(RemoveOrReplacePassInterface):
Expand Down Expand Up @@ -867,7 +674,6 @@
return [exir_ops.edge.aten.mul.Tensor]

def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:

mul_node = node
if len(mul_node.users) != 1:
return False
Expand Down Expand Up @@ -1147,9 +953,9 @@
def _apply_flat_inplace(self, graph_module) -> bool:
# Call the CSE pass on the main graph, which performs CSE without recursing.
result = CSEPass().call(graph_module)
assert (
result.graph_module is graph_module
), f"Only in-place modification is allowed, but got {result.graph_module}"
assert result.graph_module is graph_module, (
f"Only in-place modification is allowed, but got {result.graph_module}"
)
return result.modified


Expand Down
Loading
Loading