diff --git a/backends/cadence/aot/BUCK b/backends/cadence/aot/BUCK index 765f7b7f1e9..d41fba2e71f 100644 --- a/backends/cadence/aot/BUCK +++ b/backends/cadence/aot/BUCK @@ -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", @@ -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", @@ -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", @@ -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", diff --git a/backends/cadence/aot/fuse_ops.py b/backends/cadence/aot/fuse_ops.py index 428a64bf428..32cded0e1fe 100644 --- a/backends/cadence/aot/fuse_ops.py +++ b/backends/cadence/aot/fuse_ops.py @@ -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: @@ -593,23 +591,7 @@ class FuseCascadedViewOps(_SharedFuseCascadedViewOps): 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, @@ -618,192 +600,17 @@ class FuseQuantDequantToRequantizePass(FuseOpPairsAcrossBranchesPass): 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, + super().__init__( + allow_requantize=allow_requantize, + force_quant_dequant_fusion=force_quant_dequant_fusion, ) - # 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} - ), - ) - if modified: - graph_module.graph.eliminate_dead_code() - graph_module.recompile() - return PassResult(graph_module, True) - return PassResult(graph_module, False) class FuseMulScalarIntoDequantPass(RemoveOrReplacePassInterface): @@ -867,7 +674,6 @@ def targets(self) -> list[EdgeOpOverload]: 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 diff --git a/backends/cadence/aot/remove_ops.py b/backends/cadence/aot/remove_ops.py index 78144cf583c..a39f24114cf 100644 --- a/backends/cadence/aot/remove_ops.py +++ b/backends/cadence/aot/remove_ops.py @@ -6,7 +6,7 @@ # pyre-strict -from typing import cast, List, Optional, Sequence, Set, Type +from typing import cast, List, Optional, Set, Type # Import these for the cadence function signatures. import executorch.backends.cadence.aot.ops_registrations # noqa: F401 @@ -19,11 +19,36 @@ set_arg, ) from executorch.backends.cadence.aot.simplify_ops import SimplifySliceOpPass -from executorch.backends.cadence.aot.utils import get_edge_overload_packet -from executorch.backends.transforms.remove_clone_ops import RemoveCloneOpsTransform +from executorch.backends.transforms.remove_alias_copy_op import ( + RemoveAliasCopyOpPass as _SharedRemoveAliasCopyOpPass, +) +from executorch.backends.transforms.remove_branched_quant_dequant import ( + RemoveBranchedQuantDequant as _SharedRemoveBranchedQuantDequant, +) +from executorch.backends.transforms.remove_cat_from_slice_copy import ( + RemoveCatFromSliceCopyPass as _SharedRemoveCatFromSliceCopyPass, +) +from executorch.backends.transforms.remove_clone_ops_transform_imported import ( + RemoveCloneOpsTransformImported as _SharedRemoveCloneOpsTransformImported, +) +from executorch.backends.transforms.remove_nop_as_strided_copy_op import ( + RemoveNopAsStridedCopyOpPass as _SharedRemoveNopAsStridedCopyOpPass, +) +from executorch.backends.transforms.remove_nop_expand_op import ( + RemoveNopExpandOpPass as _SharedRemoveNopExpandOpPass, +) +from executorch.backends.transforms.remove_nop_slice_or_view_op import ( + RemoveNopSliceOrViewOpPass as _SharedRemoveNopSliceOrViewOpPass, +) +from executorch.backends.transforms.remove_permute_before_mean import ( + RemovePermuteBeforeMeanPass as _SharedRemovePermuteBeforeMeanPass, +) from executorch.backends.transforms.remove_permutes_around_elementwise_ops import ( RemovePermutesAroundElementwiseOps as _SharedRemovePermutesAroundElementwiseOps, ) +from executorch.backends.transforms.remove_zero_sized_cat_args import ( + RemoveZeroSizedCatArgsPass as _SharedRemoveZeroSizedCatArgsPass, +) from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( ReplaceSqueezeAndUnsqueezeWithViewPass as _SharedReplaceSqueezeAndUnsqueezeWithViewPass, ) @@ -32,11 +57,8 @@ from executorch.exir.pass_base import ( ExportedProgramPassBase, ExportedProgramPassResult, - ExportPass, PassResult, ) -from executorch.exir.pass_manager import PassManager, PassType -from executorch.exir.passes import dead_code_elimination_pass from torch.export import ExportedProgram from torch.export.graph_signature import InputKind, OutputKind from torch.fx.node import Node @@ -44,14 +66,8 @@ from torch.utils import _pytree as pytree -class RemoveCloneOpsTransformImported(ExportPass): - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - finalize_passes: List[PassType] = [ - RemoveCloneOpsTransform(eliminate_quant_dequant_pairs=False), - ] - result = PassManager(passes=finalize_passes)(graph_module) - dead_code_elimination_pass(result.graph_module) - return result +class RemoveCloneOpsTransformImported(_SharedRemoveCloneOpsTransformImported): + pass class RemoveDetachCopyPass(RemoveOrReplacePassInterface): @@ -75,79 +91,12 @@ class RemoveRedundantOps: ] -class RemoveZeroSizedCatArgsPass(RemoveOrReplacePassInterface): - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.cat.default] - - def maybe_remove_or_replace(self, node: Node) -> bool: - # Get the cat inputs (first argument is a list of tensors) - cat_inputs_arg = node.args[0] - - # Assert that cat_inputs_arg is iterable - assert isinstance( - cat_inputs_arg, (list, tuple) - ), "cat_inputs_arg must be a sequence type" - - # Filter out zero-sized tensors - cat_inputs: list[Node] = [] - for arg in cat_inputs_arg: - if isinstance(arg, Node) and arg.meta.get("val") is not None: - if arg.meta["val"].numel() > 0: - cat_inputs.append(arg) - - # If all tensors were empty, create a full op with the right shape - if not cat_inputs: - empty_shape = node.meta["val"].shape - dtype = node.meta["val"].dtype - # Create a new full node - with node.graph.inserting_before(node): - full_node = node.graph.call_function( - exir_ops.edge.aten.full.default, - args=(tuple(empty_shape), 0), - kwargs={"dtype": dtype}, - ) - full_node.meta = node.meta.copy() - node.replace_all_uses_with(full_node) - return True - - # If only one tensor remains, replace with it - if len(cat_inputs) == 1: - node.replace_all_uses_with(cat_inputs[0]) - return True - - # If the number of inputs changed, update the cat args - if len(cat_inputs) < len(cat_inputs_arg): - # Update the first argument with filtered inputs - new_args = list(node.args) - new_args[0] = cat_inputs - node.args = tuple(new_args) - return True - - # No changes needed - return False - +class RemoveZeroSizedCatArgsPass(_SharedRemoveZeroSizedCatArgsPass): + pass -class RemoveNopExpandOpPass(RemoveOrReplacePassInterface): - """ - For an expand op, if the operator shape matches the expand shape, then the - expand is a nop. - """ - @property - def targets(self) -> list[EdgeOpOverload]: - return [ - exir_ops.edge.aten.expand_copy.default, - exir_ops.edge.aten.expand.default, - ] - - def maybe_remove_or_replace(self, node: Node) -> bool: - input_node = node.args[0] - assert isinstance(input_node, Node) - if input_node.meta["val"].shape == node.meta["val"].shape: - node.replace_all_uses_with(input_node) - return True - return False +class RemoveNopExpandOpPass(_SharedRemoveNopExpandOpPass): + pass class RemoveToOpsPass(RemoveOrReplacePassInterface): @@ -191,61 +140,12 @@ def maybe_remove_or_replace(self, node: Node) -> bool: return True -class RemoveNopSliceOrViewOpPass(RemoveOrReplacePassInterface): - """ - Remove slice ops that are more like views, and view ops that do not change the shape - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [ - exir_ops.edge.aten.slice_copy.Tensor, - exir_ops.edge.aten.view_copy.default, - ] - - def maybe_remove_or_replace(self, node: Node) -> bool: - changed = False - input_node = node.args[0] - assert isinstance(input_node, Node) - if input_node.meta["val"].shape == node.meta["val"].shape: - node.replace_all_uses_with(input_node) - changed = True - - return changed - - -class RemoveNopAsStridedCopyOpPass(RemoveOrReplacePassInterface): - """Remove as_strided_copy ops that preserve every logical element.""" - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.as_strided_copy.default] - - def maybe_remove_or_replace(self, node: Node) -> bool: - input_node = node.args[0] - assert isinstance(input_node, Node) - input_value = cast(torch.Tensor, input_node.meta["val"]) - size = get_arg(node, "size", list[int]) - stride = get_arg(node, "stride", list[int]) - storage_offset = get_arg(node, "storage_offset", Optional[int]) +class RemoveNopSliceOrViewOpPass(_SharedRemoveNopSliceOrViewOpPass): + pass - if tuple(size) != tuple(input_value.shape) or len(stride) != len(size): - return False - if ( - storage_offset is not None - and storage_offset != input_value.storage_offset() - ): - return False - if not all( - dim_size == 1 or output_stride == input_stride - for dim_size, output_stride, input_stride in zip( - size, stride, input_value.stride() - ) - ): - return False - node.replace_all_uses_with(input_node) - return True +class RemoveNopAsStridedCopyOpPass(_SharedRemoveNopAsStridedCopyOpPass): + pass class RemoveNopLinalgVectorNormOpPass(RemoveOrReplacePassInterface): @@ -300,21 +200,8 @@ def maybe_remove_or_replace(self, node: Node) -> bool: return True -class RemoveAliasCopyOpPass(RemoveOrReplacePassInterface): - """ - - alias_copy is a no-op and can be removed. - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.alias_copy.default] - - def maybe_remove_or_replace(self, node: Node) -> bool: - input_node = node.args[0] - assert isinstance(input_node, Node) - node.replace_all_uses_with(input_node) - return True +class RemoveAliasCopyOpPass(_SharedRemoveAliasCopyOpPass): + pass class RemoveNopRequantizeOpPass(RemoveOrReplacePassInterface): @@ -416,217 +303,15 @@ def maybe_remove_or_replace(self, node: Node) -> bool: return False -class RemovePermuteBeforeMeanPass(RemoveOrReplacePassInterface): - """Remove or sink permute ops that precede mean reductions through unary chains. - - When a permute feeds into a mean (possibly through unary ops like - dequantize/quantize), two optimizations apply: - - 1. If non-reduced dims maintain their relative order and positions, the - permute is fully removed and the mean's reduction dims are remapped. - 2. Otherwise, the permute is moved after the mean so it operates on - smaller data. - - Cost model - ---------- - Let S_in = input size in bytes, S_out = output size in bytes, - R = S_in / S_out (reduction ratio), C_c = contiguous mean compute - cost, C_s = strided mean compute cost, and Delta = C_s - C_c. - - Original graph (permute -> mean): - Cost_orig = 3*S_in + S_out + C_c - Breakdown: permute reads and writes S_in (2*S_in), mean reads - S_in and writes S_out. - - Case 1 -- full removal (mean with remapped dims, no permute): - Cost_remove = S_in + S_out + C_s - Profitable when: Delta < 2*S_in - The strided access penalty must be less than twice the full - input tensor I/O (the eliminated permute cost). - - Case 2 -- reorder (mean with remapped dims -> small permute): - Cost_reorder = S_in + 3*S_out + C_s - Profitable when: Delta < 2*(S_in - S_out) = 2*S_in*(R-1)/R - At R = 4 the threshold is 1.5*S_in; at R = 16 it approaches - 2*S_in, converging to the full-removal bound as S_out becomes - negligible. - - Full removal always dominates reorder when both are structurally - possible (Cost_remove < Cost_reorder since S_out < 3*S_out), so - removal is applied without a cost gate. - - Additionally, if the original permute does not place the reduction - dims as the trailing (innermost) dimensions, the mean is already - strided in the original graph. Removing or sinking the permute - saves the permute I/O (2*S_in) while the mean performance can only - improve or stay the same. This makes the transformation - unconditionally profitable without needing the cost model. - """ - - _UNARY_TARGETS: frozenset[EdgeOpOverload] = frozenset( - { +class RemovePermuteBeforeMeanPass(_SharedRemovePermuteBeforeMeanPass): + _UNARY_TARGETS: frozenset[EdgeOpOverload] = ( + _SharedRemovePermuteBeforeMeanPass._UNARY_TARGETS + | { exir_ops.edge.cadence.dequantize_per_tensor.default, exir_ops.edge.cadence.quantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - exir_ops.edge.aten.clone.default, - exir_ops.edge.aten.relu.default, - exir_ops.edge.aten.neg.default, - exir_ops.edge.aten.abs.default, } ) - _MIN_REDUCTION_RATIO_FOR_REORDER: int = 4 - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.mean.dim] - - def _find_permute_through_unary_chain(self, mean_node: Node) -> Optional[Node]: - """Walk backward from mean through single-user unary ops to find a permute.""" - current = mean_node.args[0] - if not isinstance(current, Node): - return None - while True: - if current.target == exir_ops.edge.aten.permute_copy.default: - return current - if current.target not in self._UNARY_TARGETS: - return None - if len(current.users) != 1: - return None - parent = current.args[0] - if not isinstance(parent, Node): - return None - current = parent - - @staticmethod - def _reduction_dims_are_trailing(reduction_dims: list[int], ndim: int) -> bool: - """Check whether all reduction dims are the trailing (innermost) dims.""" - canonical = sorted(d % ndim for d in reduction_dims) - return canonical == list(range(ndim - len(canonical), ndim)) - - def _is_reorder_profitable( - self, - reduction_dims: list[int], - new_reduction_dims: list[int], - ndim: int, - input_shape: torch.Size, - ) -> bool: - """Determine whether sinking the permute past the mean is profitable. - - Three cases, in order of evaluation: - - 1. The original permute does not place the reduction dims as trailing - (innermost) dimensions. The mean is already strided, so removing - the permute saves 2*S_in with no mean degradation. Always - profitable. - - 2. The new (remapped) reduction dims are trailing in the original - layout. The mean stays contiguous after the transformation - (Delta ~ 0). Always profitable. - - 3. The mean becomes strided. The reorder is profitable when - Delta < 2*S_in*(R-1)/R. We approximate this by requiring - R >= _MIN_REDUCTION_RATIO_FOR_REORDER, ensuring the I/O savings - from operating on a smaller tensor outweigh the strided access - penalty. - """ - if not self._reduction_dims_are_trailing(reduction_dims, ndim): - return True - - if self._reduction_dims_are_trailing(new_reduction_dims, ndim): - return True - - reduction_ratio = 1 - canonical_reduction = set(new_reduction_dims) - for d in canonical_reduction: - reduction_ratio *= input_shape[d] - - return reduction_ratio >= self._MIN_REDUCTION_RATIO_FOR_REORDER - - def maybe_remove_or_replace(self, node: Node) -> bool: - reduction_dims = get_arg(node, "dim", list[int]) - - permute_node = self._find_permute_through_unary_chain(node) - if permute_node is None: - return False - - perm = get_arg(permute_node, "dims", list[int]) - ndim = len(perm) - - if len(permute_node.users) != 1: - return False - - permute_input = permute_node.args[0] - assert isinstance(permute_input, Node) - - new_reduction_dims = [perm[d] for d in reduction_dims] - keepdim = get_arg(node, "keepdim", bool) - - # Determine if the permute can be fully removed (post-mean permute - # would be a no-op) vs needing to be sunk after the mean. - canonical_reduction = set(new_reduction_dims) - if keepdim: - can_remove = all( - perm[d] == d for d in range(ndim) if d not in canonical_reduction - ) - else: - non_reduced_in_perm_order = [ - d for d in perm if d not in canonical_reduction - ] - can_remove = non_reduced_in_perm_order == sorted(non_reduced_in_perm_order) - - # Full removal is almost always profitable. Reorder requires a - # tighter cost bound; verify via the cost model before proceeding. - if not can_remove: - input_shape = permute_input.meta["val"].shape - if not self._is_reorder_profitable( - reduction_dims, new_reduction_dims, ndim, input_shape - ): - return False - - # Rewire: the permute's single user (either the mean itself or the - # first unary op in the chain) should read from the permute's input. - permute_user = next(iter(permute_node.users)) - permute_user.replace_input_with(permute_node, permute_input) - node.args = (node.args[0], new_reduction_dims) + node.args[2:] - - # Re-derive the mean's meta since its reduction dims changed. - fake_args = pytree.tree_map( - lambda x: x.meta["val"] if isinstance(x, Node) else x, - node.args, - ) - node.meta["val"] = node.target(*fake_args) # pyre-ignore[29] - - if not can_remove: - # Compute the post-mean permute on the reduced output. - if keepdim: - post_perm = list(perm) - else: - non_reduced_original = sorted( - d for d in range(ndim) if d not in canonical_reduction - ) - non_reduced_permuted = [d for d in perm if d not in canonical_reduction] - post_perm = [ - non_reduced_original.index(d) for d in non_reduced_permuted - ] - - graph = node.graph - with graph.inserting_after(node): - new_permute = graph.create_node( - "call_function", - exir_ops.edge.aten.permute_copy.default, - args=(node, post_perm), - ) - new_permute.meta["val"] = exir_ops.edge.aten.permute_copy.default( - node.meta["val"], post_perm - ) - for user in list(node.users): - if user is not new_permute: - user.replace_input_with(node, new_permute) - - return True - class ReplaceSqueezeAndUnsqueezeWithViewPassImported( _SharedReplaceSqueezeAndUnsqueezeWithViewPass @@ -776,15 +461,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: return PassResult(graph_module, False) -class RemoveBranchedQuantDequant(PassBase): - """ - This pass looks for adjacent quant and dequant nodes with identical - parameters, where the quant node has other users in addition to the - dequant. The quant and dequant pair would be removed by the - FuseQuantDequantToRequantizePass if not for the multiple users. This pass - removes just the dequant node by connecting it to the quant's parent node - """ - +class RemoveBranchedQuantDequant(_SharedRemoveBranchedQuantDequant): quantize_op_packets: set[EdgeOpOverloadPacket] = { exir_ops.edge.cadence.quantize_per_tensor, exir_ops.edge.quantized_decomposed.quantize_per_tensor, @@ -794,105 +471,9 @@ class RemoveBranchedQuantDequant(PassBase): exir_ops.edge.quantized_decomposed.dequantize_per_tensor, } - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - modified = self.remove_branched( - graph_module, self.quantize_op_packets, self.dequantize_op_packets - ) - modified |= self.remove_branched( - graph_module, self.dequantize_op_packets, self.quantize_op_packets - ) - - if modified: - graph_module.graph.eliminate_dead_code() - graph_module.recompile() - - return PassResult(graph_module, modified) - - def remove_branched( - self, - graph_module: torch.fx.GraphModule, - producer_pkts: set[EdgeOpOverloadPacket], - consumer_pkts: set[EdgeOpOverloadPacket], - ) -> bool: - modified = False - for node in graph_module.graph.nodes: - if ( - node.op != "call_function" - or not isinstance(node.target, EdgeOpOverload) - or get_edge_overload_packet(node.target) not in producer_pkts - ): - continue - - if len(node.users) < 2: - continue - - for user in node.users: - if ( - not isinstance(user.target, EdgeOpOverload) - or get_edge_overload_packet(user.target) not in consumer_pkts - ): - continue - - # check qparams match - if node.args[1:] != user.args[1:]: - continue - - user.replace_all_uses_with(node.args[0]) - modified = True - return modified - - -class RemoveCatFromSliceCopyPass(RemoveOrReplacePassInterface): - """ - Simplifies cat->slice_copy chains where one of the cat inputs can be directly passed - to the slice_copy. - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.slice_copy.Tensor] - - def maybe_remove_or_replace(self, node: Node) -> bool: - cat_node = get_arg(node, "input", Node) - slice_dim = get_arg(node, "dim", int) - start_idx = get_arg(node, "start", Optional[int]) - end_idx = get_arg(node, "end", Optional[int]) - step = get_arg(node, "step", int) - - if cat_node.target != exir_ops.edge.aten.cat.default or step != 1: - return False - - # Make sure cat and slice happens on the same dimension. - cat_dim = get_arg(cat_node, "dim", int) - if cat_dim != slice_dim: - return False - - # Canonicalize slice indices. - cat_output_shape = cat_node.meta["val"].shape - if start_idx is None: - start_idx = 0 - elif start_idx < 0: - start_idx += cat_output_shape[cat_dim] - if end_idx is None or end_idx > cat_output_shape[cat_dim]: - end_idx = cat_output_shape[cat_dim] - elif end_idx < 0: - end_idx += cat_output_shape[cat_dim] - - offset = 0 - for cat_input_node in get_arg(cat_node, "tensors", Sequence[Node]): - cat_input_shape = cat_input_node.meta["val"].shape - - # Check if the slice range overlaps with the cat input range. - if offset <= start_idx and end_idx <= offset + cat_input_shape[cat_dim]: - node.replace_input_with(cat_node, cat_input_node) - set_arg(node, "start", start_idx - offset) - set_arg(node, "end", end_idx - offset) - return True - - offset += cat_input_shape[cat_dim] - - return False +class RemoveCatFromSliceCopyPass(_SharedRemoveCatFromSliceCopyPass): + pass class CommonRemovePasses: diff --git a/backends/cadence/aot/reorder_ops.py b/backends/cadence/aot/reorder_ops.py index 4ee7bb193d6..23f7759b572 100644 --- a/backends/cadence/aot/reorder_ops.py +++ b/backends/cadence/aot/reorder_ops.py @@ -9,629 +9,79 @@ # This file contains all the functions that reorder ops in the graph module. -from collections import defaultdict -from math import prod -from typing import Callable, cast, DefaultDict, List, Optional, Tuple +# The Cadence pass subclasses below reference torch.ops.cadence.* at class +# definition, so the op library must be registered first. +import executorch.backends.cadence.aot.ops_registrations # noqa: F401 import torch import torch.fx -from executorch.backends.cadence.aot.compiler_utils import get_placeholders, get_shape -from executorch.backends.cadence.aot.pass_utils import ( - get_arg, - get_overload_packet, - RemoveOrReplacePassInterface, +from executorch.backends.cadence.aot.compiler_utils import get_placeholders +from executorch.backends.cadence.aot.pass_utils import RemoveOrReplacePassInterface +from executorch.backends.transforms.advance_quantize_op_above_def_chain import ( + AdvanceQuantizeOpAboveDefChainPass as _SharedAdvanceQuantizeOpAboveDefChainPass, +) +from executorch.backends.transforms.advance_quantize_op_above_def_in_branch import ( + AdvanceQuantizeOpAboveDefInBranchPass as _SharedAdvanceQuantizeOpAboveDefInBranchPass, +) +from executorch.backends.transforms.move_permute_after_concat import ( + MovePermuteAfterConcat as _SharedMovePermuteAfterConcat, +) +from executorch.backends.transforms.move_slice_before_permute import ( + MoveSliceBeforePermutePass as _SharedMoveSliceBeforePermutePass, +) +from executorch.backends.transforms.move_slice_before_view import ( + MoveSliceBeforeViewPass as _SharedMoveSliceBeforeViewPass, +) +from executorch.backends.transforms.postpone_dequantize_op_below_use_chain import ( + PostponeDequantizeOpBelowUseChainPass as _SharedPostponeDequantizeOpBelowUseChainPass, ) -from executorch.backends.cadence.aot.utils import get_edge_overload_packet from executorch.backends.transforms.postpone_permute_below_squeeze_view import ( PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView as _SharedPostponePermuteOpBelowSqueezeOrUnsqueezeLikeView, ) +from executorch.backends.transforms.propagate_slice import ( + PropagateSlice as _SharedPropagateSlice, +) +from executorch.backends.transforms.split_dequantized_cat import ( + SplitDequantizedCatPass as _SharedSplitDequantizedCatPass, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload -from executorch.exir.pass_base import ExportPass, PassResult -from executorch.exir.tensor import num_bytes_from_shape_and_dtype - -# A list of ops that can be trivially quantized -trivially_quantizable_ops_overloadpkt = { - exir_ops.edge.aten.chunk, - exir_ops.edge.aten.clone, - exir_ops.edge.aten.contiguous, - exir_ops.edge.aten.expand_copy, - exir_ops.edge.aten.permute_copy, - exir_ops.edge.aten.select_copy, - exir_ops.edge.aten.slice_copy, - exir_ops.edge.aten.squeeze_copy, - exir_ops.edge.aten.transpose_copy, - exir_ops.edge.aten.unfold_copy, - exir_ops.edge.aten.unsqueeze_copy, - exir_ops.edge.aten.view_copy, - torch.ops.aten.chunk, - torch.ops.aten.clone, - torch.ops.aten.contiguous, - torch.ops.aten.expand_copy, - torch.ops.aten.permute, - torch.ops.aten.permute_copy, - torch.ops.aten.select_copy, - torch.ops.aten.slice, - torch.ops.aten.slice_copy, - torch.ops.aten.squeeze, - torch.ops.aten.squeeze_copy, - torch.ops.aten.transpose, - torch.ops.aten.transpose_copy, - torch.ops.aten.unsqueeze, - torch.ops.aten.unsqueeze_copy, - torch.ops.aten.view, - torch.ops.aten.view_copy, -} - -# slice-equivalent ops -slice_or_select_overloadpkt = { - torch.ops.aten.slice_copy, - torch.ops.aten.select_copy, - exir_ops.edge.aten.slice_copy, - exir_ops.edge.aten.select_copy, -} - -supported_quantize_ops_overloadpkt = { - exir_ops.edge.quantized_decomposed.quantize_per_tensor, - exir_ops.edge.quantized_decomposed.quantize_per_channel, - exir_ops.edge.cadence.quantize_per_tensor, -} - - -class AdvanceQuantizeOpAboveDefInBranchPass(ExportPass): - """ - If the graph is branched with the following pattern: - I = ... - S1 = slice(I) - Q1 = quantize(S1) - S2 = slice(I) - Q2 = quantize(S2) - S3 = slice(I) - Q3 = quantize(S3) - ... - such that the elements in the slices S1 + S2 + S3 is greater than I, - we can advance the quantize above their defs (i.e., all the slice nodes), - and reorder the pattern to the following: - I = ... - Q1 = quantize(I) - S1 = slice(Q1) - Q1 = requantize(S1) - S2 = slice(Q1) - Q2 = requantize(S2) - S3 = slice(Q1) - Q3 = requantize(S3) - ... - Note that the other passes won't do this transformation because they expect - a linear chain of def-use, which is not true here; the uses of I are - branched. - """ - - def __init__(self): - super().__init__() - self.graph_module = None - - # Starting at node, iterate through its successors, bypassing any trivially - # quantizable op. If all the descendents are quantize ops, return them. - def get_descendent_quant_ops(self, node: torch.fx.Node) -> List[torch.fx.Node]: - # The list of quant ops that are descendents of node, such that the only - # nodes in the path from node --> quant are trivially quantizable ops. - descendent_quant_ops = [] - # The list of trivially quantizable ops in the path from node --> quant op. - trivial_quantized_ops = [] - - users = list(node.users.keys()) - while users: - user = users.pop(0) - user_target = get_overload_packet(user.target) - # Record a quant op successor - if user_target in { - torch.ops.quantized_decomposed.quantize_per_tensor, - exir_ops.edge.quantized_decomposed.quantize_per_tensor, - torch.ops.cadence.quantize_per_tensor, - exir_ops.edge.cadence.quantize_per_tensor, - }: - descendent_quant_ops.append(user) - # If the successor is a trivially quantizable op, consider its users - # instead. - elif user_target in trivially_quantizable_ops_overloadpkt: - trivial_quantized_ops.append(user) - users.extend(list(user.users.keys())) - # Otherwise all successors of node are not quant op, so break the loop. - else: - descendent_quant_ops.clear() - break - - # If all the nodes in trivial_quantize_ops of the node were slice ops, - # ensure that the advance is still profitable. - if descendent_quant_ops and all( - get_overload_packet(x.target) in slice_or_select_overloadpkt - for x in trivial_quantized_ops - ): - # Profitability metric: the sum of all the output slices must be - # more than half the input node size. - slice_sizes = [ - prod(list(y)) - for x in trivial_quantized_ops - if (y := get_shape(self.graph_module, x)) is not None - ] - node_shape = get_shape(self.graph_module, node) - node_size = prod(list(node_shape)) if node_shape is not None else 0 - if node_size >= 2 * sum(slice_sizes): - descendent_quant_ops.clear() - - return descendent_quant_ops - - def advance_quantize_op(self, graph_module: torch.fx.GraphModule): - graph = graph_module.graph - for node in graph.nodes: - # We are only interested in call functions and placeholders - if node.op not in {"placeholder", "call_function"}: - continue - # If the node is trivially quantizable, skip it - if ( - get_overload_packet(node.target) - in trivially_quantizable_ops_overloadpkt - ): - continue - # Get the descendent quant ops that are connected to the current - # node via trivially quantizable ops. - descendent_quant_ops = self.get_descendent_quant_ops(node) - if not descendent_quant_ops: - continue - - # Get the insertion point below which we need to insert anything. - # if node is a placeholder, we will only insert a new node after - # all the placeholders in the graph. - insertion_pt = ( - get_placeholders(graph)[-1] if node.op == "placeholder" else node - ) - - # If the node only has a single quant op as descendent, we can - # simply hoist the quant op below the current node as its single - # child. - if len(descendent_quant_ops) == 1: - quant_node = descendent_quant_ops.pop() - # Replace the uses of quant node with its predecessor - quant_node.replace_all_uses_with(quant_node.args[0]) # pyre-fixme[6] - # Hoist the quant node after the current node. Make sure that - # the insertion is after placeholders - with graph.inserting_after(insertion_pt): - dom_quant_args = (node,) + quant_node.args[1:] - dom_quant_node = graph.call_function( - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default - ) - dom_quant_node.meta = node.meta - node.replace_all_uses_with(dom_quant_node) - dom_quant_node.args = dom_quant_args - graph.erase_node(quant_node) - continue - - # Otherwise we have multiple quant descendents. Requantizing them from - # one hoisted quantize is only lossless if they all share qparams, so - # skip when they differ. - quant_args = { - tuple(quant_node.args[1:]) for quant_node in descendent_quant_ops - } - if len(quant_args) > 1: - continue - (rep_args,) = quant_args - - # Create a new quant node that dominates all the nodes in - # descendent_quant_ops. Make sure that the insertion is after - # all the placeholders. - with graph.inserting_after(insertion_pt): - dom_quant_args = (node,) + rep_args - dom_quant_node = graph.call_function( - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default - ) - dom_quant_node.meta = node.meta - node.replace_all_uses_with(dom_quant_node) - dom_quant_node.args = dom_quant_args - - # Finally, convert each of the quant node to a dequant/quant pair that - # requantizes the data flowing through dom_quant_node. - # TODO: Once requantize is implemented for PT2, replace the - # dequant/quant pair here with a single requantize node - for quant_node in descendent_quant_ops: - with graph.inserting_before(quant_node): - dequant_node = graph.call_function( - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ) - dequant_node.args = (quant_node.args[0],) + rep_args - quant_node.args = (dequant_node,) + quant_node.args[1:] - - graph_module.recompile() - graph_module.graph.eliminate_dead_code() - - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - self.graph_module = graph_module - self.advance_quantize_op(graph_module) - result = super().call(graph_module) - return result - - -class AdvanceQuantizeOpAboveDefChainPass(ExportPass): - """ - Advances a quantize op above data-movement ops to reduce data volume. - - Handles two cases: - - 1. Linear chain: if the input to a quantize op is a chain of trivially - quantizable ops (view, transpose, permute, slice), rewrite - data_movement(fp32) -> quantize to quantize -> data_movement(quantized) - so the data movement operates on smaller quantized tensors. - - 2. Cat: if the input to a quantize op is a cat with a single user (the - quantize), advance the quantize above the cat by quantizing each cat - input individually. A later pass can clean up any redundant - dequant-quant pairs on the inputs. - - 3. Caller-supplied ops: advance the quantize above a value-preserving op by - quantizing selected direct floating-point tensor inputs. The caller is - responsible for supplying only inputs that support the quantized dtype. - - For the cat case, SplitDequantizedCatPass should run first to ensure - each cat has at most one quantize consumer. - """ - - def __init__( - self, - extra_quantizable_ops: dict[EdgeOpOverload, tuple[int, ...]] | None = None, - ) -> None: - super().__init__() - self.graph_module = None - self._extra_quantizable_ops = extra_quantizable_ops or {} - - # Return true if advancing the quantize node is feasible - def advancing_feasible(self, quant_node: torch.fx.Node): - assert quant_node.op == "call_function" and len(quant_node.args) >= 1 - # Get the input of the quant node. Only proceed if it's a torch node. - inp = quant_node.args[0] - if not isinstance(inp, torch.fx.Node): - return False - - # Return false if the input to the quantize node is (1) not trivially - # quantizable, or (2) has more than one user. - inp_users = list(inp.users.keys()) - inp_overloadpkt = None - if isinstance(inp.target, EdgeOpOverload): - inp_overloadpkt = get_edge_overload_packet(inp.target) - else: - inp_overloadpkt = get_overload_packet(inp.target) - - if ( - inp_overloadpkt not in trivially_quantizable_ops_overloadpkt - or len(inp_users) != 1 - ): - return False - - # Advancing quantize op above slice nodes is tricky. If we advance the - # quantize node above slice, then we will quantize the input to the slice - # op, which can be expensive. We only bypass nop slice at present. - if inp_overloadpkt in slice_or_select_overloadpkt: - sliced_tensor = inp.args[0] - assert isinstance(sliced_tensor, torch.fx.Node) - slice_input_shape = get_shape(self.graph_module, sliced_tensor) - slice_output_shape = get_shape(self.graph_module, inp) - # If we could not glean the shapes, or the slice op is a nop, bail - if ( - slice_output_shape is None - or slice_input_shape is None - or prod(list(slice_output_shape)) < prod(list(slice_input_shape)) - ): - return False - # All the conditions satisfied, we advance. - return True - - def _advance_above_cat( - self, quant_node: torch.fx.Node, cat_node: torch.fx.Node - ) -> None: - """Advance a quantize op above a cat by quantizing each cat input.""" - graph = quant_node.graph - quant_params = quant_node.args[1:] - - cat_inputs = cat_node.args[0] - assert isinstance(cat_inputs, (list, tuple)) - - new_inputs: list[torch.fx.Node] = [] - for inp in cat_inputs: - # cat concatenates tensors, so every input must be a node. - assert isinstance(inp, torch.fx.Node) - - with graph.inserting_before(cat_node): - new_quant = graph.call_function( - # pyre-ignore[6] - quant_node.target, - args=(inp, *quant_params), - ) - # This copies the fp32 input's meta, so meta["val"] keeps the - # fp32 dtype rather than the quantized output dtype. That's fine: - # nothing in this pass reads dtype from meta (only shape, which - # is correct), and call() re-runs super().call() to re-propagate - # fake tensors, making meta dtype-consistent before we return. - new_quant.meta = inp.meta.copy() - new_inputs.append(new_quant) - - dim = get_arg(cat_node, "dim", int) - with graph.inserting_before(quant_node): - new_cat = graph.call_function( - # pyre-ignore[6] - cat_node.target, - args=(new_inputs, dim), - ) - new_cat.meta = quant_node.meta.copy() - - quant_node.replace_all_uses_with(new_cat) - graph.erase_node(quant_node) - - def _advance_above_extra_quantizable_op( - self, - quant_node: torch.fx.Node, - op_node: torch.fx.Node, - quantizable_input_indices: tuple[int, ...], - ) -> bool: - graph = quant_node.graph - new_op_args = list(op_node.args) - quantized_input = False - - for index in quantizable_input_indices: - assert 0 <= index < len(op_node.args) - arg = op_node.args[index] - assert isinstance(arg, torch.fx.Node) - value = arg.meta["val"] - assert isinstance(value, torch.Tensor) - if not value.dtype.is_floating_point: - continue - - quant_args = list(quant_node.args) - quant_args[0] = arg - with graph.inserting_before(op_node): - new_quant = graph.call_function( - # pyre-ignore[6] - quant_node.target, - args=tuple(quant_args), - kwargs=quant_node.kwargs, - ) - # We will correct the dtype when we run - # ExportPass call at the end. - new_quant.meta = arg.meta.copy() - new_op_args[index] = new_quant - quantized_input = True - - if not quantized_input: - return False - - with graph.inserting_before(quant_node): - new_op = graph.call_function( - # pyre-ignore[6] - op_node.target, - args=tuple(new_op_args), - kwargs=op_node.kwargs, - ) - new_op.meta = quant_node.meta.copy() - - quant_node.replace_all_uses_with(new_op) - graph.erase_node(quant_node) - return True - def advance_quantize_op(self, graph_module: torch.fx.GraphModule) -> bool: - graph = graph_module.graph - modified = False - for node in reversed(graph.nodes): - if get_overload_packet(node.target) not in ( - exir_ops.edge.quantized_decomposed.quantize_per_tensor, - torch.ops.quantized_decomposed.quantize_per_tensor, - exir_ops.edge.cadence.quantize_per_tensor, - torch.ops.cadence.quantize_per_tensor, - ): - continue - - inp = node.args[0] - if ( - isinstance(inp, torch.fx.Node) - and get_overload_packet(inp.target) - in (exir_ops.edge.aten.cat, torch.ops.aten.cat) - and len(inp.users) == 1 - ): - self._advance_above_cat(node, inp) - modified = True - continue - - if ( - isinstance(inp, torch.fx.Node) - and inp.target in self._extra_quantizable_ops - and len(inp.users) == 1 - and self._advance_above_extra_quantizable_op( - node, - inp, - self._extra_quantizable_ops[inp.target], - ) - ): - modified = True - continue - - if not self.advancing_feasible(node): - continue - - trivially_quantizable_op = node.args[0] - # The input to the quant node must now be the input to the trivially - # quantizable op. - quant_args = list(node.args) - quant_args[0] = trivially_quantizable_op.args[0] - - # Insert the new quant node with updated args before the current - # quant node. - with graph.inserting_before(node): - quant_node = graph.call_function(node.target, args=tuple(quant_args)) - quant_node.meta = node.meta - # Move the trivially quantizable node after the quant node - with graph.inserting_after(node): - tq_args = list(trivially_quantizable_op.args) - tq_args[0] = quant_node - tq_node = graph.call_function( - trivially_quantizable_op.target, - args=tuple(tq_args), - kwargs=trivially_quantizable_op.kwargs, - ) - tq_node.meta = trivially_quantizable_op.meta - # Replace all uses of node with newly created tq_node - node.replace_all_uses_with(tq_node) - # We can safely remove the quant node and trivially quantizable op - graph.erase_node(node) - graph.erase_node(trivially_quantizable_op) - modified = True - - return modified - - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - self.graph_module = graph_module - modified = self.advance_quantize_op(graph_module) - if modified: - graph_module.recompile() - graph_module.graph.eliminate_dead_code() - return super().call(graph_module) - - return PassResult(graph_module, False) - - -class PostponeDequantizeOpBelowUseChainPass(ExportPass): - """ - If the consumer of dequantize is a linear chain of view, transpose, permute, - or slice ops that are trivially quantized, we can convert the pattern - dequantize(int8/uint8) -> view/transpose/permute/slice(fp32) to - view/transpose/permute/slice(int8/uint8) -> dequantize(int8/uint8) - The benefit of such reordering is that the view/transpose/permute/slice - will move far less data. - """ - - def __init__(self): - super().__init__() - self.graph_module = None - - # Return true if postponing the dequantize node is feasible - def postponing_feasible(self, dequant_node: torch.fx.Node): - users = list(dequant_node.users.keys()) - # Check if the dequantize op has a single user, and that user is - # trivially quantizable. - trivially_quantizable_users = all( - get_overload_packet(user.target) in trivially_quantizable_ops_overloadpkt - for user in users - ) - if len(users) == 1: - return trivially_quantizable_users - - # Otherwise check if all the users are slice op - if not all( - get_overload_packet(user.target) in slice_or_select_overloadpkt - for user in users - ): - return False - - quantized_branches = [] - for user in users: - slice_users = list(user.users) - quantized_branches.append( - bool(slice_users) - and all( - slice_user.op == "call_function" - and get_overload_packet(slice_user.target) - in supported_quantize_ops_overloadpkt - for slice_user in slice_users - ) - ) - - # Preserve the existing fallback for forks whose branches all requantize. - if all(quantized_branches): - return True - - dequant_shape = get_shape(self.graph_module, dequant_node) - if dequant_shape is None: - return False - - dequant_bytes = num_bytes_from_shape_and_dtype(dequant_shape, torch.float32) - surviving_dequant_bytes = 0 - for user, quantized_branch in zip(users, quantized_branches): - if quantized_branch: - continue - slice_shape = get_shape(self.graph_module, user) - if slice_shape is None: - return False - # Nop slices are removed later and do not add another materialized copy. - if prod(list(slice_shape)) == prod(list(dequant_shape)): - continue - surviving_dequant_bytes += num_bytes_from_shape_and_dtype( - slice_shape, torch.float32 - ) - - return surviving_dequant_bytes <= dequant_bytes - - def postpone_dequantize_op(self, graph_module: torch.fx.GraphModule) -> bool: - # Different supported dequant ops have their own default variants - packet_to_overload_map = { - exir_ops.edge.quantized_decomposed.dequantize_per_tensor: "default", - torch.ops.cadence.dequantize_per_tensor: "default", - exir_ops.edge.quantized_decomposed.dequantize_per_channel: "default", - exir_ops.edge.cadence.dequantize_per_tensor: "default", - } - graph = graph_module.graph - modified = False - for node in graph.nodes: - overload_packet = get_overload_packet(node.target) - if ( - overload_packet not in packet_to_overload_map.keys() - or not self.postponing_feasible(node) - ): - continue - - for user in node.users: - with graph.inserting_after(user): - dequant_node = graph.call_function( - getattr( - overload_packet, packet_to_overload_map[overload_packet] - ), - args=(user, *node.args[1:]), - ) - dequant_node.meta = user.meta.copy() - # Remove meta["debug_handle"] on new node if it exists. - # Reassign it at the caller level by calling generate_missing_debug_handles - dequant_node.meta.pop("debug_handle", None) - user.replace_all_uses_with(dequant_node) - dequant_node.args = (user, *node.args[1:]) - - pred = node.args[0] - node.replace_all_uses_with(pred) - graph.erase_node(node) - modified = True - - if modified: - graph_module.recompile() - return modified - - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - # The logic in postpone_dequantize_op that handles branching checks the shape - # of the dequant node, which isn't available if that node was already postponed - # in the same pass invokation. The shape information is recreated by tracing in - # super().call(), meaning that every branch in the graph that we wish to postpone - # dequant past requires retracing. We iterate the pass until it no longer modifies - # the graph (up to 3 times max, to avoid potential infinite loops) - self.graph_module = graph_module - iter_count = 0 - local_modified = False - overall_modified = False +class AdvanceQuantizeOpAboveDefInBranchPass( + _SharedAdvanceQuantizeOpAboveDefInBranchPass +): + quantize_op_packets: set[object] = { + torch.ops.quantized_decomposed.quantize_per_tensor, + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + torch.ops.cadence.quantize_per_tensor, + exir_ops.edge.cadence.quantize_per_tensor, + } - while local_modified or iter_count == 0: - local_modified = self.postpone_dequantize_op(self.graph_module) - overall_modified |= local_modified - if local_modified: - self.graph_module = super().call(self.graph_module).graph_module +class AdvanceQuantizeOpAboveDefChainPass(_SharedAdvanceQuantizeOpAboveDefChainPass): + quantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + torch.ops.quantized_decomposed.quantize_per_tensor, + exir_ops.edge.cadence.quantize_per_tensor, + torch.ops.cadence.quantize_per_tensor, + } - iter_count += 1 - if iter_count == 3: - break - return PassResult(self.graph_module, overall_modified) +class PostponeDequantizeOpBelowUseChainPass( + _SharedPostponeDequantizeOpBelowUseChainPass +): + quantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + exir_ops.edge.quantized_decomposed.quantize_per_channel, + exir_ops.edge.cadence.quantize_per_tensor, + } + dequantize_packet_to_overload: dict[object, str] = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor: "default", + torch.ops.cadence.dequantize_per_tensor: "default", + exir_ops.edge.quantized_decomposed.dequantize_per_channel: "default", + exir_ops.edge.cadence.dequantize_per_tensor: "default", + } class SinkOpsCloserToUsePass(RemoveOrReplacePassInterface): @@ -774,674 +224,36 @@ class PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView( pass -class MovePermuteAfterConcat(RemoveOrReplacePassInterface): - """Move matching single-use permutes after a cat. - - For example, - - cat( - [ - permute(x[80, 16, 32], [1, 2, 0]), # [16, 32, 80] - permute(y[1, 16, 32], [1, 2, 0]), # [16, 32, 1] - ], - dim=2, - ) # [16, 32, 81] - - becomes - - permute(cat([x, y], dim=0), [1, 2, 0]) - # cat: [81, 16, 32], output: [16, 32, 81] - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.cat.default] - - @staticmethod - def _match_inputs( - cat_node: torch.fx.Node, - ) -> Optional[Tuple[List[torch.fx.Node], Tuple[int, ...]]]: - inputs = get_arg( - cat_node, - "tensors", - # pyre-ignore[6] - list[torch.fx.Node] | tuple[torch.fx.Node, ...], - ) - # Moving one reused permute after cat makes it process the larger output. - if inputs and all(inp is inputs[0] for inp in inputs): - return None - - unpermuted_inputs: List[torch.fx.Node] = [] - common_permutation: Optional[Tuple[int, ...]] = None - for permute_node in inputs: - if ( - permute_node.target != exir_ops.edge.aten.permute_copy.default - or len(permute_node.users) != 1 - ): - return None - - permute_input = cast(torch.fx.Node, permute_node.args[0]) - dims = get_arg( - permute_node, - "dims", - # pyre-ignore[6] - list[int] | tuple[int, ...], - ) - rank = len(dims) - permutation = tuple(dim % rank for dim in dims) - if common_permutation is None: - common_permutation = permutation - elif permutation != common_permutation: - return None - unpermuted_inputs.append(permute_input) - - return unpermuted_inputs, cast(Tuple[int, ...], common_permutation) - - def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: - match = self._match_inputs(node) - if match is None: - return False - inputs, permutation = match - - output_cat_dim = get_arg(node, "dim", int) % len(permutation) - input_cat_dim = permutation[output_cat_dim] - graph = node.graph - - with graph.inserting_before(node): - new_cat = graph.call_function( - exir_ops.edge.aten.cat.default, - args=(inputs, input_cat_dim), - ) - new_cat.meta["val"] = exir_ops.edge.aten.cat.default( - [inp.meta["val"] for inp in inputs], input_cat_dim - ) - new_permute = graph.call_function( - exir_ops.edge.aten.permute_copy.default, - args=(new_cat, list(permutation)), - ) - new_permute.meta = node.meta.copy() - - node.replace_all_uses_with(new_permute) - return True - - -class MoveSliceBeforePermutePass(RemoveOrReplacePassInterface): - """Move slice_copy ops before permute_copy to reduce permute data volume. - - Rewrites permute(input, perm) -> slice(dim=D) into - slice(input, dim=perm[D]) -> permute(sliced, perm), so the permute - operates on a smaller tensor. - - Scans slice nodes and matches upstream permutes. This also handles - chained cases (permute -> slice -> slice) in one pass: each slice - independently checks its input for a permute. - - Cost model: dim-0 slices are nop-eligible (zero-copy pointer offset - after MakeSliceAndCatDimOutermostPass). Moving such a slice loses the - nop, so we only move it when the permute savings outweigh the nop loss, - i.e. when the slice removes more than half the data (full > 2 * sliced). - Non-dim-0 slices have no nop opportunity, so any permute savings is - pure win. - """ - - STRIDED_SLICE_COST_FACTOR: int = 2 - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.slice_copy.Tensor] - - def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: - permute_node = get_arg(node, "input", torch.fx.Node) - if permute_node.target != exir_ops.edge.aten.permute_copy.default: - return False - - if len(permute_node.users) != 1: - return False - - perm = cast(list[int], permute_node.args[1]) - permute_input = permute_node.args[0] - assert isinstance(permute_input, torch.fx.Node) - - slice_dim = get_arg(node, "dim", int) - - full_size = prod(permute_node.meta["val"].shape) - sliced_size = prod(node.meta["val"].shape) - if slice_dim == 0 and full_size <= self.STRIDED_SLICE_COST_FACTOR * sliced_size: - return False - - new_dim = perm[slice_dim] - graph = node.graph - - with graph.inserting_before(permute_node): - new_slice_args = ( - permute_input, - new_dim, - get_arg(node, "start"), - get_arg(node, "end"), - get_arg(node, "step", int), - ) - new_slice = graph.create_node( - "call_function", - exir_ops.edge.aten.slice_copy.Tensor, - args=new_slice_args, - ) - new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( - permute_input.meta["val"], *new_slice_args[1:] - ) - new_permute = graph.create_node( - "call_function", - exir_ops.edge.aten.permute_copy.default, - args=(new_slice, perm), - ) - new_permute.meta["val"] = exir_ops.edge.aten.permute_copy.default( - new_slice.meta["val"], perm - ) - - node.replace_all_uses_with(new_permute) - return True - - -class MoveSliceBeforeViewPass(RemoveOrReplacePassInterface): - """Move a slice_copy above a view_copy when the slice is re-expressible as a - single slice on one dim of the pre-view tensor. - - Rewrites view(x) -> slice(dim=d, start, end, step) into - slice(x, dim=d', start', end', step') -> view(sliced, slice_out_shape), so the - slice lands directly on x. This may be useful in attention patterns, where - we view outputs of a large linear into a new shape where the number of - attention heads are the last dim, and we need to run independent computation - per head. Moving the slice before the view can allow us to then directly slice - the constant linear weights. - - A view is a contiguous reshape: it never moves or reorders elements, it only - re-groups the shared row-major index space into different dims. A slice keeps - an arithmetic progression of indices (start, start+step, ...) along one viewed - dim, and that progression collapses back to a *single* slice on one pre-view - dim exactly when the row-major strides line up. ``_derive_pre_view_slice`` - handles the three cases that qualify: - - * untouched dim: the viewed dim is left unchanged by the view -- same size - and same inner stride as some pre-view dim -- so the slice copies over - verbatim (any step). - * contiguous: the viewed dim and a pre-view dim span the same flat extent - (a split's outermost factor, or a merge that aligns), so a contiguous - (step==1) slice maps to a contiguous pre-view slice. - * strided: the viewed dim is an innermost factor of a pre-view dim - (identical inner stride) selected width-1, so it maps to a strided - pre-view slice with step == the viewed dim's size. - - Everything else -- middle factors, wider strided selections -- is block-strided - (runs separated by gaps), which no single slice can express, so it is left - unchanged. - - Each slice is handled independently, so a view that fans out to several slices - is rewritten one slice at a time and the now-dead view is removed by dead-code - elimination -- there is no single-user requirement on the view. - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.slice_copy.Tensor] - - def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: - view_node = get_arg(node, "input", torch.fx.Node) - if view_node.target != exir_ops.edge.aten.view_copy.default: - return False - - x_node = get_arg(view_node, "input", torch.fx.Node) - pre_view_shape = tuple(x_node.meta["val"].shape) - post_view_shape = tuple(view_node.meta["val"].shape) - if 0 in pre_view_shape or 0 in post_view_shape: - return False - - dim = get_arg(node, "dim", int) - if dim < 0: - dim += len(post_view_shape) - post_view_size = post_view_shape[dim] - - bounds = self._normalize_slice(node, post_view_size) - if bounds is None: - return False - start, stop, step = bounds - - # The slice's own output shape gives the selected-element count along the - # sliced dim directly -- it is exactly output_shape[dim]. - slice_out_shape = tuple(node.meta["val"].shape) - post_view_count = slice_out_shape[dim] - if post_view_count == 0: - return False - - # Row-major stride of the sliced viewed dim, and of every pre-view dim. - post_view_stride = prod(post_view_shape[dim + 1 :]) - pre_view_strides = self._row_major_strides(pre_view_shape) - - derived = self._derive_pre_view_slice( - pre_view_shape, - pre_view_strides, - post_view_stride, - post_view_size, - start, - stop, - step, - post_view_count, - ) - if derived is None: - return False - pre_view_dim, pre_view_start, pre_view_stop, pre_view_step = derived - - graph = node.graph - with graph.inserting_before(node): - new_slice_args = ( - x_node, - pre_view_dim, - pre_view_start, - pre_view_stop, - pre_view_step, - ) - new_slice = graph.create_node( - "call_function", - exir_ops.edge.aten.slice_copy.Tensor, - args=new_slice_args, - ) - new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( - x_node.meta["val"], *new_slice_args[1:] - ) - new_view = graph.create_node( - "call_function", - exir_ops.edge.aten.view_copy.default, - args=(new_slice, list(slice_out_shape)), - ) - new_view.meta["val"] = exir_ops.edge.aten.view_copy.default( - new_slice.meta["val"], list(slice_out_shape) - ) - - node.replace_all_uses_with(new_view) - return True - - @staticmethod - def _row_major_strides(shape: tuple[int, ...]) -> list[int]: - """Row-major (contiguous) strides for ``shape``.""" - strides = [1] * len(shape) - acc = 1 - for i in range(len(shape) - 1, -1, -1): - strides[i] = acc - acc *= shape[i] - return strides - - def _normalize_slice( - self, node: torch.fx.Node, post_view_size: int - ) -> Optional[tuple[int, int, int]]: - """Resolve the slice to concrete, clamped ``(start, stop, step)`` ints, or - None if the bounds are dynamic or the step is non-positive (neither of - which this pass handles).""" - step = get_arg(node, "step") - - if not isinstance(step, int): - return None - - if step <= 0: - return None - - raw_start = get_arg(node, "start") - raw_stop = get_arg(node, "end") - - # Make sure raw_start/raw_stop are not symbolic. - if (raw_start is not None and not isinstance(raw_start, int)) or ( - raw_stop is not None and not isinstance(raw_stop, int) - ): - return None - - start = 0 if raw_start is None else raw_start - stop = post_view_size if raw_stop is None else raw_stop - if start < 0: - start += post_view_size - if stop < 0: - stop += post_view_size - start = max(0, min(start, post_view_size)) - stop = max(0, min(stop, post_view_size)) - return start, stop, step - - def _derive_pre_view_slice( - self, - pre_view_shape: tuple[int, ...], - pre_view_strides: list[int], - post_view_stride: int, - post_view_size: int, - start: int, - stop: int, - step: int, - post_view_count: int, - ) -> tuple[int, int, int, int] | None: - """Return ``(dim, start, stop, step)`` for the single pre-view-tensor slice - equivalent to slicing the viewed dim, or None if no single pre-view slice - reproduces it. - - Both shapes index the same row-major flat space, so the sliced viewed dim - (size ``post_view_size``, inner stride ``post_view_stride``) lines up with - one pre-view dim (size ``pre_view_size``, inner stride ``pre_view_stride``) - in one of three ways. - """ - for pre_view_dim, (pre_view_stride, pre_view_size) in enumerate( - zip(pre_view_strides, pre_view_shape) - ): - # Untouched: the viewed dim is identical to this pre-view dim (same - # size and same inner stride), so the slice applies verbatim, any step. - if pre_view_stride == post_view_stride and pre_view_size == post_view_size: - return pre_view_dim, start, stop, step - - # Contiguous: the viewed dim and this pre-view dim span the same flat - # extent (same period), and the selected band aligns to this dim's - # boundaries. A contiguous (step==1) viewed slice - # [start, start+post_view_count) is the flat band [start* - # post_view_stride, (start+post_view_count)*post_view_stride), a - # contiguous slice on this pre-view dim iff both ends are multiples of - # its stride. - if ( - step == 1 - and post_view_size * post_view_stride == pre_view_size * pre_view_stride - ): - flat_start = start * post_view_stride - flat_stop = (start + post_view_count) * post_view_stride - if ( - flat_start % pre_view_stride == 0 - and flat_stop % pre_view_stride == 0 - ): - return ( - pre_view_dim, - flat_start // pre_view_stride, - flat_stop // pre_view_stride, - 1, - ) - - # Strided is the ONLY way the reshape itself introduces a stride, and - # it requires a width-1 selection (post_view_count == 1): the viewed - # dim is an innermost factor of this pre-view dim (identical inner - # stride), so fixing that single factor index and letting the rest of - # the pre-view dim run yields a uniform stride equal to the viewed dim's - # size. Any wider selection (post_view_count > 1) of an inner factor - # leaves runs separated by gaps -- block-strided, not a single slice -- - # so width-1 is required. - if ( - post_view_count == 1 - and post_view_size > 1 - and pre_view_stride == post_view_stride - and pre_view_size % post_view_size == 0 - ): - pre_view_count = pre_view_size // post_view_size - pre_view_stop = start + (pre_view_count - 1) * post_view_size + 1 - return pre_view_dim, start, pre_view_stop, post_view_size - return None - - -class PropagateSlice(RemoveOrReplacePassInterface): - """Propagate slice_copy before element-wise ops when the cost model - indicates it reduces total data movement. - - Supported ops (extensible via dispatch table): - - quantize_per_tensor: unary element-wise - - dequantize_per_tensor: unary element-wise - - add.Tensor: binary with broadcast — slices non-broadcasting inputs - - mul.Tensor: binary with broadcast — slices non-broadcasting inputs - - Callers can provide per-target filters for additional legality constraints. - Handles any slice dim and any step size. - """ - - def __init__( - self, - additional_unary_targets: Optional[list[EdgeOpOverload]] = None, - additional_binary_targets: Optional[list[EdgeOpOverload]] = None, - target_filters: Optional[ - dict[EdgeOpOverload, Callable[[torch.fx.Node], bool]] - ] = None, - ) -> None: - super().__init__() - unary_targets = [ - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - exir_ops.edge.cadence.quantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - exir_ops.edge.cadence.dequantize_per_tensor.default, - *(additional_unary_targets or []), - ] - binary_targets = [ - exir_ops.edge.aten.add.Tensor, - exir_ops.edge.aten.mul.Tensor, - *(additional_binary_targets or []), - ] - self._dispatch: dict[ - EdgeOpOverload, - tuple[ - Callable[[torch.fx.Node, torch.fx.Node], bool], - Callable[[torch.fx.Node, torch.fx.Node], bool], - ], - ] = {} - self._target_filters = target_filters or {} - for t in unary_targets: - self._dispatch[t] = ( - self._should_swap_elementwise, - self._swap_elementwise_slice, - ) - - for t in binary_targets: - self._dispatch[t] = ( - self._should_swap_binary_elementwise, - self._swap_binary_elementwise_slice, - ) - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.slice_copy.Tensor] - - def _should_swap_elementwise( - self, op_node: torch.fx.Node, slice_node: torch.fx.Node - ) -> bool: - full_size = prod(op_node.meta["val"].shape) - sliced_size = prod(slice_node.meta["val"].shape) - return sliced_size < full_size - - def _swap_elementwise_slice( - self, op_node: torch.fx.Node, slice_node: torch.fx.Node - ) -> bool: - op_input = op_node.args[0] - assert isinstance(op_input, torch.fx.Node) - graph = slice_node.graph - - slice_dim = get_arg(slice_node, "dim", int) - slice_start = get_arg(slice_node, "start") - slice_end = get_arg(slice_node, "end") - slice_step = get_arg(slice_node, "step", int) - - with graph.inserting_before(op_node): - new_slice = graph.call_function( - exir_ops.edge.aten.slice_copy.Tensor, - args=(op_input, slice_dim, slice_start, slice_end, slice_step), - ) - new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( - op_input.meta["val"], slice_dim, slice_start, slice_end, slice_step - ) - - new_args = list(op_node.args) - new_args[0] = new_slice - target = cast(EdgeOpOverload, op_node.target) - new_op = graph.call_function( - target, - args=tuple(new_args), - kwargs=op_node.kwargs, - ) - new_op.meta["val"] = target( - new_slice.meta["val"], - *[ - a.meta["val"] if isinstance(a, torch.fx.Node) else a - for a in new_args[1:] - ], - **{ - k: v.meta["val"] if isinstance(v, torch.fx.Node) else v - for k, v in op_node.kwargs.items() - }, - ) - - slice_node.replace_all_uses_with(new_op) - graph.erase_node(slice_node) - graph.erase_node(op_node) - return True - - def _should_swap_binary_elementwise( - self, op_node: torch.fx.Node, slice_node: torch.fx.Node - ) -> bool: - lhs, rhs = op_node.args[0], op_node.args[1] - assert isinstance(lhs, torch.fx.Node) and isinstance(rhs, torch.fx.Node) - if lhs.meta["val"].shape == rhs.meta["val"].shape: - return False - full_size = prod(op_node.meta["val"].shape) - sliced_size = prod(slice_node.meta["val"].shape) - return sliced_size < full_size - - def _swap_binary_elementwise_slice( - self, op_node: torch.fx.Node, slice_node: torch.fx.Node - ) -> bool: - lhs, rhs = op_node.args[0], op_node.args[1] - assert isinstance(lhs, torch.fx.Node) and isinstance(rhs, torch.fx.Node) - graph = slice_node.graph - - slice_dim = get_arg(slice_node, "dim", int) - slice_start = get_arg(slice_node, "start") - slice_end = get_arg(slice_node, "end") - slice_step = get_arg(slice_node, "step", int) - - output_shape = op_node.meta["val"].shape - output_dim = slice_dim % len(output_shape) - - new_args = list(op_node.args) - with graph.inserting_before(op_node): - for i, inp in enumerate([lhs, rhs]): - input_shape = inp.meta["val"].shape - # Broadcasting aligns operand dimensions to the right of the output. - input_dim = output_dim - (len(output_shape) - len(input_shape)) - if ( - input_dim >= 0 - and input_shape[input_dim] == output_shape[output_dim] - ): - new_slice = graph.call_function( - exir_ops.edge.aten.slice_copy.Tensor, - args=(inp, input_dim, slice_start, slice_end, slice_step), - ) - new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( - inp.meta["val"], - input_dim, - slice_start, - slice_end, - slice_step, - ) - new_args[i] = new_slice - - target = cast(EdgeOpOverload, op_node.target) - new_op = graph.call_function( - target, - args=tuple(new_args), - kwargs=op_node.kwargs, - ) - new_op.meta["val"] = target( - *[ - a.meta["val"] if isinstance(a, torch.fx.Node) else a - for a in new_args - ], - **{ - k: v.meta["val"] if isinstance(v, torch.fx.Node) else v - for k, v in op_node.kwargs.items() - }, - ) - - slice_node.replace_all_uses_with(new_op) - graph.erase_node(slice_node) - graph.erase_node(op_node) - return True - - def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: - parent = get_arg(node, "input", torch.fx.Node) - if len(parent.users) != 1: - return False - if not isinstance(parent.target, EdgeOpOverload): - return False - - entry = self._dispatch.get(parent.target) - if entry is None: - return False - target_filter = self._target_filters.get(parent.target) - if target_filter is not None and not target_filter(parent): - return False - - should_swap, do_swap = entry - return should_swap(parent, node) and do_swap(parent, node) - - -_QUANT_OVERLOAD_PACKETS = { - exir_ops.edge.quantized_decomposed.quantize_per_tensor, - exir_ops.edge.cadence.quantize_per_tensor, -} - -_DEQUANT_OVERLOAD_PACKETS = { - exir_ops.edge.quantized_decomposed.dequantize_per_tensor, - exir_ops.edge.cadence.dequantize_per_tensor, -} - - -class SplitDequantizedCatPass(RemoveOrReplacePassInterface): - """Split a cat node so that quantize consumers get their own copy. +class MovePermuteAfterConcat(_SharedMovePermuteAfterConcat): + pass - Fires when a cat has all floating-point inputs, at least one dequantize - input, and at least one quantize consumer. Quant consumers are grouped - by matching qparams; each group receives a dedicated duplicate of the - cat node. Non-quant consumers stay on the original cat, whose - semantics are unchanged. - A later pass (e.g. AdvanceQuantizeOpAboveDefChainPass extended for cat) - can then hoist each quant above its single-consumer cat copy without - affecting the non-quant paths. - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.cat.default] - - def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: - cat_inputs = node.args[0] - if not isinstance(cat_inputs, (list, tuple)): - return False - - has_dequant_input = False - for inp in cat_inputs: - assert isinstance(inp, torch.fx.Node) - val = inp.meta["val"] - if val is None or not val.is_floating_point(): - return False - if get_overload_packet(inp.target) in _DEQUANT_OVERLOAD_PACKETS: - has_dequant_input = True +class MoveSliceBeforePermutePass(_SharedMoveSliceBeforePermutePass): + pass - if not has_dequant_input: - return False - quant_groups: DefaultDict[Tuple, List[torch.fx.Node]] = defaultdict(list) - for user in list(node.users.keys()): - if get_overload_packet(user.target) in _QUANT_OVERLOAD_PACKETS: - quant_groups[user.args[1:]].append(user) +class MoveSliceBeforeViewPass(_SharedMoveSliceBeforeViewPass): + pass - if not quant_groups: - return False - graph = node.graph - dim = get_arg(node, "dim", int) - for quant_consumers in quant_groups.values(): - with graph.inserting_after(node): - dup_cat = graph.call_function( - exir_ops.edge.aten.cat.default, - args=(list(cat_inputs), dim), - ) - dup_cat.meta = node.meta.copy() +class PropagateSlice(_SharedPropagateSlice): + quant_unary_targets: list[EdgeOpOverload] = [ + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.cadence.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + exir_ops.edge.cadence.dequantize_per_tensor.default, + ] - for q_node in quant_consumers: - q_node.replace_input_with(node, dup_cat) - return True +class SplitDequantizedCatPass(_SharedSplitDequantizedCatPass): + quantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + exir_ops.edge.cadence.quantize_per_tensor, + } + dequantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor, + exir_ops.edge.cadence.dequantize_per_tensor, + } # The following class consolidates functions to reoder ops (i.e., either hoist diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index bc04f749491..bd11ee81b79 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/backends/cadence/aot/replace_ops.py @@ -13,9 +13,8 @@ import logging import math -import operator from operator import neg -from typing import cast, Dict, Optional, Sequence +from typing import cast, Dict, Sequence import torch import torch.fx @@ -32,6 +31,12 @@ from executorch.backends.transforms.replace_scalar_with_tensor import ( ReplaceScalarWithTensorArgPass, ) +from executorch.backends.transforms.replace_select_with_view_op import ( + ReplaceSelectWithViewOpPass as _SharedReplaceSelectWithViewOpPass, +) +from executorch.backends.transforms.replace_split_with_slice import ( + ReplaceSplitWithSlicePass as _SharedReplaceSplitWithSlicePass, +) from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( ReplaceSqueezeAndUnsqueezeWithViewPass as _SharedReplaceSqueezeAndUnsqueezeWithViewPass, ) @@ -224,41 +229,8 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: return True -class ReplaceSelectWithViewOpPass(RemoveOrReplacePassInterface): - """ - If the size along the select dim is 1, then the select op can be replaced - by view op. - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.select_copy.int] - - def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: - # Get the input tensor and shapes - in_tensor_node = node.args[0] - assert isinstance(in_tensor_node, torch.fx.Node) - in_shape = in_tensor_node.meta["val"].shape - out_shape = node.meta["val"].shape - - # Get the select dimension - select_dim = node.args[1] - assert isinstance(select_dim, int) - select_dim = select_dim if select_dim >= 0 else select_dim + len(in_shape) - - if in_shape[select_dim] == 1: - # Replace with view op with the new shape - with node.graph.inserting_before(node): - new_node = node.graph.call_function( - exir_ops.edge.aten.view_copy.default, - args=(node.args[0], list(out_shape)), - ) - # Important to copy metadata - new_node.meta = node.meta - node.replace_all_uses_with(new_node) - return True - - return False +class ReplaceSelectWithViewOpPass(_SharedReplaceSelectWithViewOpPass): + pass class ReplaceMMWithAddMMPass(RemoveOrReplacePassInterface): @@ -2179,75 +2151,8 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: # Adapted from fbcode/pyspeech/opt_passes/replace_ops.py -class ReplaceSplitWithSlicePass(RemoveOrReplacePassInterface): - """ - split_with_sizes() delegates to slice() op, so perform this replacement here. - This avoids the expense of delegation from ATen. - """ - - @property - def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.split_with_sizes_copy.default] - - def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: - # All the users of this split_with_sizes op must be getitem ops - if any(user.target != operator.getitem for user in node.users): - return False - - # Get the slice dim and extent for each split - slice_ops = self._get_split_sizes(node) - if slice_ops is None: - return False - - graph = node.graph - - # Go over each getitem user, and replace it with slice op - for user in list(node.users.keys()): - assert user.target == operator.getitem - item_idx = int(user.args[1]) - assert item_idx < len(slice_ops) - cur_slice = slice_ops[item_idx] - with graph.inserting_before(user): - cur_slice_node = graph.call_function( - exir_ops.edge.aten.slice_copy.Tensor, - (node.args[0], cur_slice[0], cur_slice[1], cur_slice[2], 1), - ) - # Metadata copy important - cur_slice_node.meta = user.meta - user.replace_all_uses_with(cur_slice_node) - - # Return True to indicate the split node should be removed - return True - - def _get_split_sizes(self, node: torch.fx.Node) -> Optional[list[tuple[int, ...]]]: - """For split_with_sizes, return the slice dim and extent for each split.""" - # Parse the args of the split_with_sizes op - tensor_arg, split_sizes = node.args[0:2] - assert isinstance(tensor_arg, torch.fx.Node) - - # Get shape from node metadata - val = tensor_arg.meta.get("val") - if val is None: - return None - in_shape = val.shape - - split_dim = 0 if len(node.args) < 3 else node.args[2] - - # Canonicalize the split dimension - assert isinstance(split_dim, int) - split_dim = split_dim if split_dim >= 0 else len(in_shape) + split_dim - - # Create the slice op args corresponding to each split - slice_ops = [] - split_start = 0 - assert isinstance(split_sizes, list) - for split_size in split_sizes: - split_end = split_start + split_size - slice_args = (split_dim, split_start, split_end) - slice_ops.append(slice_args) - split_start = split_end - - return slice_ops +class ReplaceSplitWithSlicePass(_SharedReplaceSplitWithSlicePass): + pass class ReplacePowWithMulPass(RemoveOrReplacePassInterface): diff --git a/backends/transforms/advance_quantize_op_above_def_chain.py b/backends/transforms/advance_quantize_op_above_def_chain.py new file mode 100644 index 00000000000..5e029e9ac3f --- /dev/null +++ b/backends/transforms/advance_quantize_op_above_def_chain.py @@ -0,0 +1,270 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from math import prod + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + get_edge_overload_packet, + get_overload_packet, + get_shape, +) +from executorch.backends.transforms.quantize_reorder_utils import ( + slice_or_select_overloadpkt, + trivially_quantizable_ops_overloadpkt, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportPass, PassResult + + +class AdvanceQuantizeOpAboveDefChainPass(ExportPass): + """ + Advances a quantize op above data-movement ops to reduce data volume. + + Handles two cases: + + 1. Linear chain: if the input to a quantize op is a chain of trivially + quantizable ops (view, transpose, permute, slice), rewrite + data_movement(fp32) -> quantize to quantize -> data_movement(quantized) + so the data movement operates on smaller quantized tensors. + + 2. Cat: if the input to a quantize op is a cat with a single user (the + quantize), advance the quantize above the cat by quantizing each cat + input individually. A later pass can clean up any redundant + dequant-quant pairs on the inputs. + + 3. Caller-supplied ops: advance the quantize above a value-preserving op by + quantizing selected direct floating-point tensor inputs. The caller is + responsible for supplying only inputs that support the quantized dtype. + + For the cat case, SplitDequantizedCatPass should run first to ensure + each cat has at most one quantize consumer. + """ + + quantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + torch.ops.quantized_decomposed.quantize_per_tensor, + } + + def __init__( + self, + extra_quantizable_ops: dict[EdgeOpOverload, tuple[int, ...]] | None = None, + ) -> None: + super().__init__() + self.graph_module = None + self._extra_quantizable_ops = extra_quantizable_ops or {} + + # Return true if advancing the quantize node is feasible + def advancing_feasible(self, quant_node: torch.fx.Node): + assert quant_node.op == "call_function" and len(quant_node.args) >= 1 + # Get the input of the quant node. Only proceed if it's a torch node. + inp = quant_node.args[0] + if not isinstance(inp, torch.fx.Node): + return False + + # Return false if the input to the quantize node is (1) not trivially + # quantizable, or (2) has more than one user. + inp_users = list(inp.users.keys()) + inp_overloadpkt = None + if isinstance(inp.target, EdgeOpOverload): + inp_overloadpkt = get_edge_overload_packet(inp.target) + else: + inp_overloadpkt = get_overload_packet(inp.target) + + if ( + inp_overloadpkt not in trivially_quantizable_ops_overloadpkt + or len(inp_users) != 1 + ): + return False + + # Advancing quantize op above slice nodes is tricky. If we advance the + # quantize node above slice, then we will quantize the input to the slice + # op, which can be expensive. We only bypass nop slice at present. + if inp_overloadpkt in slice_or_select_overloadpkt: + sliced_tensor = inp.args[0] + assert isinstance(sliced_tensor, torch.fx.Node) + slice_input_shape = get_shape(self.graph_module, sliced_tensor) + slice_output_shape = get_shape(self.graph_module, inp) + # If we could not glean the shapes, or the slice op is a nop, bail + if ( + slice_output_shape is None + or slice_input_shape is None + or prod(list(slice_output_shape)) < prod(list(slice_input_shape)) + ): + return False + + # All the conditions satisfied, we advance. + return True + + def _advance_above_cat( + self, quant_node: torch.fx.Node, cat_node: torch.fx.Node + ) -> None: + """Advance a quantize op above a cat by quantizing each cat input.""" + graph = quant_node.graph + quant_params = quant_node.args[1:] + + cat_inputs = cat_node.args[0] + assert isinstance(cat_inputs, (list, tuple)) + + new_inputs: list[torch.fx.Node] = [] + for inp in cat_inputs: + # cat concatenates tensors, so every input must be a node. + assert isinstance(inp, torch.fx.Node) + + with graph.inserting_before(cat_node): + new_quant = graph.call_function( + # pyre-ignore[6] + quant_node.target, + args=(inp, *quant_params), + ) + # This copies the fp32 input's meta, so meta["val"] keeps the + # fp32 dtype rather than the quantized output dtype. That's fine: + # nothing in this pass reads dtype from meta (only shape, which + # is correct), and call() re-runs super().call() to re-propagate + # fake tensors, making meta dtype-consistent before we return. + new_quant.meta = inp.meta.copy() + new_inputs.append(new_quant) + + dim = get_arg(cat_node, "dim", int) + with graph.inserting_before(quant_node): + new_cat = graph.call_function( + # pyre-ignore[6] + cat_node.target, + args=(new_inputs, dim), + ) + new_cat.meta = quant_node.meta.copy() + + quant_node.replace_all_uses_with(new_cat) + graph.erase_node(quant_node) + + def _advance_above_extra_quantizable_op( + self, + quant_node: torch.fx.Node, + op_node: torch.fx.Node, + quantizable_input_indices: tuple[int, ...], + ) -> bool: + graph = quant_node.graph + new_op_args = list(op_node.args) + quantized_input = False + + for index in quantizable_input_indices: + assert 0 <= index < len(op_node.args) + arg = op_node.args[index] + assert isinstance(arg, torch.fx.Node) + value = arg.meta["val"] + assert isinstance(value, torch.Tensor) + if not value.dtype.is_floating_point: + continue + + quant_args = list(quant_node.args) + quant_args[0] = arg + with graph.inserting_before(op_node): + new_quant = graph.call_function( + # pyre-ignore[6] + quant_node.target, + args=tuple(quant_args), + kwargs=quant_node.kwargs, + ) + # We will correct the dtype when we run + # ExportPass call at the end. + new_quant.meta = arg.meta.copy() + new_op_args[index] = new_quant + quantized_input = True + + if not quantized_input: + return False + + with graph.inserting_before(quant_node): + new_op = graph.call_function( + # pyre-ignore[6] + op_node.target, + args=tuple(new_op_args), + kwargs=op_node.kwargs, + ) + new_op.meta = quant_node.meta.copy() + + quant_node.replace_all_uses_with(new_op) + graph.erase_node(quant_node) + return True + + def advance_quantize_op(self, graph_module: torch.fx.GraphModule) -> bool: + graph = graph_module.graph + modified = False + for node in reversed(graph.nodes): + if get_overload_packet(node.target) not in self.quantize_op_packets: + continue + + inp = node.args[0] + if ( + isinstance(inp, torch.fx.Node) + and get_overload_packet(inp.target) + in (exir_ops.edge.aten.cat, torch.ops.aten.cat) + and len(inp.users) == 1 + ): + self._advance_above_cat(node, inp) + modified = True + continue + + if ( + isinstance(inp, torch.fx.Node) + and inp.target in self._extra_quantizable_ops + and len(inp.users) == 1 + and self._advance_above_extra_quantizable_op( + node, + inp, + self._extra_quantizable_ops[inp.target], + ) + ): + modified = True + continue + + if not self.advancing_feasible(node): + continue + + trivially_quantizable_op = node.args[0] + # The input to the quant node must now be the input to the trivially + # quantizable op. + quant_args = list(node.args) + quant_args[0] = trivially_quantizable_op.args[0] + + # Insert the new quant node with updated args before the current + # quant node. + with graph.inserting_before(node): + quant_node = graph.call_function(node.target, args=tuple(quant_args)) + quant_node.meta = node.meta + # Move the trivially quantizable node after the quant node + with graph.inserting_after(node): + tq_args = list(trivially_quantizable_op.args) + tq_args[0] = quant_node + tq_node = graph.call_function( + trivially_quantizable_op.target, + args=tuple(tq_args), + kwargs=trivially_quantizable_op.kwargs, + ) + tq_node.meta = trivially_quantizable_op.meta + # Replace all uses of node with newly created tq_node + node.replace_all_uses_with(tq_node) + # We can safely remove the quant node and trivially quantizable op + graph.erase_node(node) + graph.erase_node(trivially_quantizable_op) + modified = True + + return modified + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + self.graph_module = graph_module + modified = self.advance_quantize_op(graph_module) + if modified: + graph_module.recompile() + graph_module.graph.eliminate_dead_code() + return super().call(graph_module) + + return PassResult(graph_module, False) diff --git a/backends/transforms/advance_quantize_op_above_def_in_branch.py b/backends/transforms/advance_quantize_op_above_def_in_branch.py new file mode 100644 index 00000000000..18a6ee3f3d7 --- /dev/null +++ b/backends/transforms/advance_quantize_op_above_def_in_branch.py @@ -0,0 +1,196 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from math import prod +from typing import List + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_overload_packet, + get_placeholders, + get_shape, +) +from executorch.backends.transforms.quantize_reorder_utils import ( + slice_or_select_overloadpkt, + trivially_quantizable_ops_overloadpkt, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + + +class AdvanceQuantizeOpAboveDefInBranchPass(ExportPass): + """ + If the graph is branched with the following pattern: + I = ... + S1 = slice(I) + Q1 = quantize(S1) + S2 = slice(I) + Q2 = quantize(S2) + S3 = slice(I) + Q3 = quantize(S3) + ... + such that the elements in the slices S1 + S2 + S3 is greater than I, + we can advance the quantize above their defs (i.e., all the slice nodes), + and reorder the pattern to the following: + I = ... + Q1 = quantize(I) + S1 = slice(Q1) + Q1 = requantize(S1) + S2 = slice(Q1) + Q2 = requantize(S2) + S3 = slice(Q1) + Q3 = requantize(S3) + ... + Note that the other passes won't do this transformation because they expect + a linear chain of def-use, which is not true here; the uses of I are + branched. + """ + + quantize_op_packets: set[object] = { + torch.ops.quantized_decomposed.quantize_per_tensor, + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + } + + def __init__(self): + super().__init__() + self.graph_module = None + + # Starting at node, iterate through its successors, bypassing any trivially + # quantizable op. If all the descendents are quantize ops, return them. + def get_descendent_quant_ops(self, node: torch.fx.Node) -> List[torch.fx.Node]: + # The list of quant ops that are descendents of node, such that the only + # nodes in the path from node --> quant are trivially quantizable ops. + descendent_quant_ops = [] + # The list of trivially quantizable ops in the path from node --> quant op. + trivial_quantized_ops = [] + + users = list(node.users.keys()) + while users: + user = users.pop(0) + user_target = get_overload_packet(user.target) + # Record a quant op successor + if user_target in self.quantize_op_packets: + descendent_quant_ops.append(user) + # If the successor is a trivially quantizable op, consider its users + # instead. + elif user_target in trivially_quantizable_ops_overloadpkt: + trivial_quantized_ops.append(user) + users.extend(list(user.users.keys())) + # Otherwise all successors of node are not quant op, so break the loop. + else: + descendent_quant_ops.clear() + break + + # If all the nodes in trivial_quantize_ops of the node were slice ops, + # ensure that the advance is still profitable. + if descendent_quant_ops and all( + get_overload_packet(x.target) in slice_or_select_overloadpkt + for x in trivial_quantized_ops + ): + # Profitability metric: the sum of all the output slices must be + # more than half the input node size. + slice_sizes = [ + prod(list(y)) + for x in trivial_quantized_ops + if (y := get_shape(self.graph_module, x)) is not None + ] + node_shape = get_shape(self.graph_module, node) + node_size = prod(list(node_shape)) if node_shape is not None else 0 + if node_size >= 2 * sum(slice_sizes): + descendent_quant_ops.clear() + + return descendent_quant_ops + + def advance_quantize_op(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + for node in graph.nodes: + # We are only interested in call functions and placeholders + if node.op not in {"placeholder", "call_function"}: + continue + # If the node is trivially quantizable, skip it + if ( + get_overload_packet(node.target) + in trivially_quantizable_ops_overloadpkt + ): + continue + # Get the descendent quant ops that are connected to the current + # node via trivially quantizable ops. + descendent_quant_ops = self.get_descendent_quant_ops(node) + if not descendent_quant_ops: + continue + + # Get the insertion point below which we need to insert anything. + # if node is a placeholder, we will only insert a new node after + # all the placeholders in the graph. + insertion_pt = ( + get_placeholders(graph)[-1] if node.op == "placeholder" else node + ) + + # If the node only has a single quant op as descendent, we can + # simply hoist the quant op below the current node as its single + # child. + if len(descendent_quant_ops) == 1: + quant_node = descendent_quant_ops.pop() + # Replace the uses of quant node with its predecessor + quant_node.replace_all_uses_with(quant_node.args[0]) # pyre-fixme[6] + # Hoist the quant node after the current node. Make sure that + # the insertion is after placeholders + with graph.inserting_after(insertion_pt): + dom_quant_args = (node,) + quant_node.args[1:] + dom_quant_node = graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + ) + dom_quant_node.meta = node.meta + node.replace_all_uses_with(dom_quant_node) + dom_quant_node.args = dom_quant_args + graph.erase_node(quant_node) + continue + + # Otherwise we have multiple quant descendents. Requantizing them from + # one hoisted quantize is only lossless if they all share qparams, so + # skip when they differ. + quant_args = { + tuple(quant_node.args[1:]) for quant_node in descendent_quant_ops + } + if len(quant_args) > 1: + continue + (rep_args,) = quant_args + + # Create a new quant node that dominates all the nodes in + # descendent_quant_ops. Make sure that the insertion is after + # all the placeholders. + with graph.inserting_after(insertion_pt): + dom_quant_args = (node,) + rep_args + dom_quant_node = graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + ) + dom_quant_node.meta = node.meta + node.replace_all_uses_with(dom_quant_node) + dom_quant_node.args = dom_quant_args + + # Finally, convert each of the quant node to a dequant/quant pair that + # requantizes the data flowing through dom_quant_node. + # TODO: Once requantize is implemented for PT2, replace the + # dequant/quant pair here with a single requantize node + for quant_node in descendent_quant_ops: + with graph.inserting_before(quant_node): + dequant_node = graph.call_function( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + ) + dequant_node.args = (quant_node.args[0],) + rep_args + quant_node.args = (dequant_node,) + quant_node.args[1:] + + graph_module.recompile() + graph_module.graph.eliminate_dead_code() + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + self.graph_module = graph_module + self.advance_quantize_op(graph_module) + result = super().call(graph_module) + return result diff --git a/backends/transforms/fuse_quant_dequant_to_requantize.py b/backends/transforms/fuse_quant_dequant_to_requantize.py new file mode 100644 index 00000000000..722afab71b4 --- /dev/null +++ b/backends/transforms/fuse_quant_dequant_to_requantize.py @@ -0,0 +1,240 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import cast, Optional + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + FuseOpPairsAcrossBranchesPass, + get_arg, +) +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 torch.utils import _pytree as pytree + + +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, + } + + quantize_op_packets: set[EdgeOpOverloadPacket] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + } + dequantize_op_packets: set[EdgeOpOverloadPacket] = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor, + } + # Op used to fuse a dequantize-quantize pair whose qparams differ. There is + # no backend-neutral requantize op, so a backend that wants such pairs fused + # must set this; otherwise only pairs with matching qparams are removed. + requantize_op: Optional[EdgeOpOverload] = None + + def __init__( + self, allow_requantize: bool = False, force_quant_dequant_fusion: bool = False + ) -> None: + super().__init__() + if allow_requantize and self.requantize_op is None: + raise ValueError( + f"{type(self).__name__} has no requantize_op, so it cannot " + "run with allow_requantize=True" + ) + 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( + cast(EdgeOpOverload, self.requantize_op), + 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"] = cast(EdgeOpOverload, self.requantize_op)( + 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} + ), + ) + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) + return PassResult(graph_module, False) diff --git a/backends/transforms/move_permute_after_concat.py b/backends/transforms/move_permute_after_concat.py new file mode 100644 index 00000000000..a08bc91442a --- /dev/null +++ b/backends/transforms/move_permute_after_concat.py @@ -0,0 +1,109 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import cast, List, Optional, Tuple + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +class MovePermuteAfterConcat(RemoveOrReplacePassInterface): + """Move matching single-use permutes after a cat. + + For example, + + cat( + [ + permute(x[80, 16, 32], [1, 2, 0]), # [16, 32, 80] + permute(y[1, 16, 32], [1, 2, 0]), # [16, 32, 1] + ], + dim=2, + ) # [16, 32, 81] + + becomes + + permute(cat([x, y], dim=0), [1, 2, 0]) + # cat: [81, 16, 32], output: [16, 32, 81] + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.cat.default] + + @staticmethod + def _match_inputs( + cat_node: torch.fx.Node, + ) -> Optional[Tuple[List[torch.fx.Node], Tuple[int, ...]]]: + inputs = get_arg( + cat_node, + "tensors", + # pyre-ignore[6] + list[torch.fx.Node] | tuple[torch.fx.Node, ...], + ) + # Moving one reused permute after cat makes it process the larger output. + if inputs and all(inp is inputs[0] for inp in inputs): + return None + + unpermuted_inputs: List[torch.fx.Node] = [] + common_permutation: Optional[Tuple[int, ...]] = None + for permute_node in inputs: + if ( + permute_node.target != exir_ops.edge.aten.permute_copy.default + or len(permute_node.users) != 1 + ): + return None + + permute_input = cast(torch.fx.Node, permute_node.args[0]) + dims = get_arg( + permute_node, + "dims", + # pyre-ignore[6] + list[int] | tuple[int, ...], + ) + rank = len(dims) + permutation = tuple(dim % rank for dim in dims) + if common_permutation is None: + common_permutation = permutation + elif permutation != common_permutation: + return None + unpermuted_inputs.append(permute_input) + + return unpermuted_inputs, cast(Tuple[int, ...], common_permutation) + + def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: + match = self._match_inputs(node) + if match is None: + return False + inputs, permutation = match + + output_cat_dim = get_arg(node, "dim", int) % len(permutation) + input_cat_dim = permutation[output_cat_dim] + graph = node.graph + + with graph.inserting_before(node): + new_cat = graph.call_function( + exir_ops.edge.aten.cat.default, + args=(inputs, input_cat_dim), + ) + new_cat.meta["val"] = exir_ops.edge.aten.cat.default( + [inp.meta["val"] for inp in inputs], input_cat_dim + ) + new_permute = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(new_cat, list(permutation)), + ) + new_permute.meta = node.meta.copy() + + node.replace_all_uses_with(new_permute) + return True diff --git a/backends/transforms/move_slice_before_permute.py b/backends/transforms/move_slice_before_permute.py new file mode 100644 index 00000000000..888c00c5638 --- /dev/null +++ b/backends/transforms/move_slice_before_permute.py @@ -0,0 +1,95 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from math import prod +from typing import cast + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +class MoveSliceBeforePermutePass(RemoveOrReplacePassInterface): + """Move slice_copy ops before permute_copy to reduce permute data volume. + + Rewrites permute(input, perm) -> slice(dim=D) into + slice(input, dim=perm[D]) -> permute(sliced, perm), so the permute + operates on a smaller tensor. + + Scans slice nodes and matches upstream permutes. This also handles + chained cases (permute -> slice -> slice) in one pass: each slice + independently checks its input for a permute. + + Cost model: dim-0 slices are nop-eligible (zero-copy pointer offset + after MakeSliceAndCatDimOutermostPass). Moving such a slice loses the + nop, so we only move it when the permute savings outweigh the nop loss, + i.e. when the slice removes more than half the data (full > 2 * sliced). + Non-dim-0 slices have no nop opportunity, so any permute savings is + pure win. + """ + + STRIDED_SLICE_COST_FACTOR: int = 2 + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.slice_copy.Tensor] + + def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: + permute_node = get_arg(node, "input", torch.fx.Node) + if permute_node.target != exir_ops.edge.aten.permute_copy.default: + return False + + if len(permute_node.users) != 1: + return False + + perm = cast(list[int], permute_node.args[1]) + permute_input = permute_node.args[0] + assert isinstance(permute_input, torch.fx.Node) + + slice_dim = get_arg(node, "dim", int) + + full_size = prod(permute_node.meta["val"].shape) + sliced_size = prod(node.meta["val"].shape) + if slice_dim == 0 and full_size <= self.STRIDED_SLICE_COST_FACTOR * sliced_size: + return False + + new_dim = perm[slice_dim] + graph = node.graph + + with graph.inserting_before(permute_node): + new_slice_args = ( + permute_input, + new_dim, + get_arg(node, "start"), + get_arg(node, "end"), + get_arg(node, "step", int), + ) + new_slice = graph.create_node( + "call_function", + exir_ops.edge.aten.slice_copy.Tensor, + args=new_slice_args, + ) + new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( + permute_input.meta["val"], *new_slice_args[1:] + ) + new_permute = graph.create_node( + "call_function", + exir_ops.edge.aten.permute_copy.default, + args=(new_slice, perm), + ) + new_permute.meta["val"] = exir_ops.edge.aten.permute_copy.default( + new_slice.meta["val"], perm + ) + + node.replace_all_uses_with(new_permute) + return True diff --git a/backends/transforms/move_slice_before_view.py b/backends/transforms/move_slice_before_view.py new file mode 100644 index 00000000000..c8f4dc823d7 --- /dev/null +++ b/backends/transforms/move_slice_before_view.py @@ -0,0 +1,251 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from math import prod +from typing import Optional + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +class MoveSliceBeforeViewPass(RemoveOrReplacePassInterface): + """Move a slice_copy above a view_copy when the slice is re-expressible as a + single slice on one dim of the pre-view tensor. + + Rewrites view(x) -> slice(dim=d, start, end, step) into + slice(x, dim=d', start', end', step') -> view(sliced, slice_out_shape), so the + slice lands directly on x. This may be useful in attention patterns, where + we view outputs of a large linear into a new shape where the number of + attention heads are the last dim, and we need to run independent computation + per head. Moving the slice before the view can allow us to then directly slice + the constant linear weights. + + A view is a contiguous reshape: it never moves or reorders elements, it only + re-groups the shared row-major index space into different dims. A slice keeps + an arithmetic progression of indices (start, start+step, ...) along one viewed + dim, and that progression collapses back to a *single* slice on one pre-view + dim exactly when the row-major strides line up. ``_derive_pre_view_slice`` + handles the three cases that qualify: + + * untouched dim: the viewed dim is left unchanged by the view -- same size + and same inner stride as some pre-view dim -- so the slice copies over + verbatim (any step). + * contiguous: the viewed dim and a pre-view dim span the same flat extent + (a split's outermost factor, or a merge that aligns), so a contiguous + (step==1) slice maps to a contiguous pre-view slice. + * strided: the viewed dim is an innermost factor of a pre-view dim + (identical inner stride) selected width-1, so it maps to a strided + pre-view slice with step == the viewed dim's size. + + Everything else -- middle factors, wider strided selections -- is block-strided + (runs separated by gaps), which no single slice can express, so it is left + unchanged. + + Each slice is handled independently, so a view that fans out to several slices + is rewritten one slice at a time and the now-dead view is removed by dead-code + elimination -- there is no single-user requirement on the view. + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.slice_copy.Tensor] + + def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: + view_node = get_arg(node, "input", torch.fx.Node) + if view_node.target != exir_ops.edge.aten.view_copy.default: + return False + + x_node = get_arg(view_node, "input", torch.fx.Node) + pre_view_shape = tuple(x_node.meta["val"].shape) + post_view_shape = tuple(view_node.meta["val"].shape) + if 0 in pre_view_shape or 0 in post_view_shape: + return False + + dim = get_arg(node, "dim", int) + if dim < 0: + dim += len(post_view_shape) + post_view_size = post_view_shape[dim] + + bounds = self._normalize_slice(node, post_view_size) + if bounds is None: + return False + start, stop, step = bounds + + # The slice's own output shape gives the selected-element count along the + # sliced dim directly -- it is exactly output_shape[dim]. + slice_out_shape = tuple(node.meta["val"].shape) + post_view_count = slice_out_shape[dim] + if post_view_count == 0: + return False + + # Row-major stride of the sliced viewed dim, and of every pre-view dim. + post_view_stride = prod(post_view_shape[dim + 1 :]) + pre_view_strides = self._row_major_strides(pre_view_shape) + + derived = self._derive_pre_view_slice( + pre_view_shape, + pre_view_strides, + post_view_stride, + post_view_size, + start, + stop, + step, + post_view_count, + ) + if derived is None: + return False + pre_view_dim, pre_view_start, pre_view_stop, pre_view_step = derived + + graph = node.graph + with graph.inserting_before(node): + new_slice_args = ( + x_node, + pre_view_dim, + pre_view_start, + pre_view_stop, + pre_view_step, + ) + new_slice = graph.create_node( + "call_function", + exir_ops.edge.aten.slice_copy.Tensor, + args=new_slice_args, + ) + new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( + x_node.meta["val"], *new_slice_args[1:] + ) + new_view = graph.create_node( + "call_function", + exir_ops.edge.aten.view_copy.default, + args=(new_slice, list(slice_out_shape)), + ) + new_view.meta["val"] = exir_ops.edge.aten.view_copy.default( + new_slice.meta["val"], list(slice_out_shape) + ) + + node.replace_all_uses_with(new_view) + return True + + @staticmethod + def _row_major_strides(shape: tuple[int, ...]) -> list[int]: + """Row-major (contiguous) strides for ``shape``.""" + strides = [1] * len(shape) + acc = 1 + for i in range(len(shape) - 1, -1, -1): + strides[i] = acc + acc *= shape[i] + return strides + + def _normalize_slice( + self, node: torch.fx.Node, post_view_size: int + ) -> Optional[tuple[int, int, int]]: + """Resolve the slice to concrete, clamped ``(start, stop, step)`` ints, or + None if the bounds are dynamic or the step is non-positive (neither of + which this pass handles).""" + step = get_arg(node, "step") + + if not isinstance(step, int): + return None + + if step <= 0: + return None + + raw_start = get_arg(node, "start") + raw_stop = get_arg(node, "end") + + # Make sure raw_start/raw_stop are not symbolic. + if (raw_start is not None and not isinstance(raw_start, int)) or ( + raw_stop is not None and not isinstance(raw_stop, int) + ): + return None + + start = 0 if raw_start is None else raw_start + stop = post_view_size if raw_stop is None else raw_stop + if start < 0: + start += post_view_size + if stop < 0: + stop += post_view_size + start = max(0, min(start, post_view_size)) + stop = max(0, min(stop, post_view_size)) + return start, stop, step + + def _derive_pre_view_slice( + self, + pre_view_shape: tuple[int, ...], + pre_view_strides: list[int], + post_view_stride: int, + post_view_size: int, + start: int, + stop: int, + step: int, + post_view_count: int, + ) -> tuple[int, int, int, int] | None: + """Return ``(dim, start, stop, step)`` for the single pre-view-tensor slice + equivalent to slicing the viewed dim, or None if no single pre-view slice + reproduces it. + + Both shapes index the same row-major flat space, so the sliced viewed dim + (size ``post_view_size``, inner stride ``post_view_stride``) lines up with + one pre-view dim (size ``pre_view_size``, inner stride ``pre_view_stride``) + in one of three ways. + """ + for pre_view_dim, (pre_view_stride, pre_view_size) in enumerate( + zip(pre_view_strides, pre_view_shape) + ): + # Untouched: the viewed dim is identical to this pre-view dim (same + # size and same inner stride), so the slice applies verbatim, any step. + if pre_view_stride == post_view_stride and pre_view_size == post_view_size: + return pre_view_dim, start, stop, step + + # Contiguous: the viewed dim and this pre-view dim span the same flat + # extent (same period), and the selected band aligns to this dim's + # boundaries. A contiguous (step==1) viewed slice + # [start, start+post_view_count) is the flat band [start* + # post_view_stride, (start+post_view_count)*post_view_stride), a + # contiguous slice on this pre-view dim iff both ends are multiples of + # its stride. + if ( + step == 1 + and post_view_size * post_view_stride == pre_view_size * pre_view_stride + ): + flat_start = start * post_view_stride + flat_stop = (start + post_view_count) * post_view_stride + if ( + flat_start % pre_view_stride == 0 + and flat_stop % pre_view_stride == 0 + ): + return ( + pre_view_dim, + flat_start // pre_view_stride, + flat_stop // pre_view_stride, + 1, + ) + + # Strided is the ONLY way the reshape itself introduces a stride, and + # it requires a width-1 selection (post_view_count == 1): the viewed + # dim is an innermost factor of this pre-view dim (identical inner + # stride), so fixing that single factor index and letting the rest of + # the pre-view dim run yields a uniform stride equal to the viewed dim's + # size. Any wider selection (post_view_count > 1) of an inner factor + # leaves runs separated by gaps -- block-strided, not a single slice -- + # so width-1 is required. + if ( + post_view_count == 1 + and post_view_size > 1 + and pre_view_stride == post_view_stride + and pre_view_size % post_view_size == 0 + ): + pre_view_count = pre_view_size // post_view_size + pre_view_stop = start + (pre_view_count - 1) * post_view_size + 1 + return pre_view_dim, start, pre_view_stop, post_view_size + return None diff --git a/backends/transforms/permute_pass_utils.py b/backends/transforms/permute_pass_utils.py index a4a0f7c7a5b..a3ad1406a0d 100644 --- a/backends/transforms/permute_pass_utils.py +++ b/backends/transforms/permute_pass_utils.py @@ -14,7 +14,7 @@ from abc import abstractmethod from collections import deque -from typing import cast, List, Optional, Type, TypeVar, Union +from typing import Callable, cast, List, Optional, Type, TypeVar, Union import torch import torch.fx @@ -22,12 +22,27 @@ 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 ExportPass, PassResult +from torch._ops import OpOverloadPacket from torch.fx import Node from torch.fx.node import Argument T = TypeVar("T") +def get_overload_packet( + op: Union[Callable[..., str], str], +) -> Union[OpOverloadPacket, EdgeOpOverloadPacket, None]: + return ( + get_edge_overload_packet(op) + if isinstance(op, EdgeOpOverload) + else getattr(op, "overloadpacket", None) + ) + + +def get_placeholders(graph: torch.fx.Graph) -> List[torch.fx.Node]: + return list(filter(lambda x: x.op == "placeholder", graph.nodes)) + + def get_edge_overload_packet(edge_op: EdgeOpOverload) -> EdgeOpOverloadPacket: edge_op_namespace, edge_op_name = ( edge_op.namespace, diff --git a/backends/transforms/postpone_dequantize_op_below_use_chain.py b/backends/transforms/postpone_dequantize_op_below_use_chain.py new file mode 100644 index 00000000000..93fb4d1f28a --- /dev/null +++ b/backends/transforms/postpone_dequantize_op_below_use_chain.py @@ -0,0 +1,165 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from math import prod + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_overload_packet, + get_shape, +) +from executorch.backends.transforms.quantize_reorder_utils import ( + slice_or_select_overloadpkt, + trivially_quantizable_ops_overloadpkt, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.tensor import num_bytes_from_shape_and_dtype + + +class PostponeDequantizeOpBelowUseChainPass(ExportPass): + """ + If the consumer of dequantize is a linear chain of view, transpose, permute, + or slice ops that are trivially quantized, we can convert the pattern + dequantize(int8/uint8) -> view/transpose/permute/slice(fp32) to + view/transpose/permute/slice(int8/uint8) -> dequantize(int8/uint8) + The benefit of such reordering is that the view/transpose/permute/slice + will move far less data. + """ + + quantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + exir_ops.edge.quantized_decomposed.quantize_per_channel, + } + dequantize_packet_to_overload: dict[object, str] = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor: "default", + exir_ops.edge.quantized_decomposed.dequantize_per_channel: "default", + } + + def __init__(self): + super().__init__() + self.graph_module = None + + # Return true if postponing the dequantize node is feasible + def postponing_feasible(self, dequant_node: torch.fx.Node): + users = list(dequant_node.users.keys()) + # Check if the dequantize op has a single user, and that user is + # trivially quantizable. + trivially_quantizable_users = all( + get_overload_packet(user.target) in trivially_quantizable_ops_overloadpkt + for user in users + ) + if len(users) == 1: + return trivially_quantizable_users + + # Otherwise check if all the users are slice op + if not all( + get_overload_packet(user.target) in slice_or_select_overloadpkt + for user in users + ): + return False + + quantized_branches = [] + for user in users: + slice_users = list(user.users) + quantized_branches.append( + bool(slice_users) + and all( + slice_user.op == "call_function" + and get_overload_packet(slice_user.target) + in self.quantize_op_packets + for slice_user in slice_users + ) + ) + + # Preserve the existing fallback for forks whose branches all requantize. + if all(quantized_branches): + return True + + dequant_shape = get_shape(self.graph_module, dequant_node) + if dequant_shape is None: + return False + + dequant_bytes = num_bytes_from_shape_and_dtype(dequant_shape, torch.float32) + surviving_dequant_bytes = 0 + for user, quantized_branch in zip(users, quantized_branches): + if quantized_branch: + continue + slice_shape = get_shape(self.graph_module, user) + if slice_shape is None: + return False + # Nop slices are removed later and do not add another materialized copy. + if prod(list(slice_shape)) == prod(list(dequant_shape)): + continue + surviving_dequant_bytes += num_bytes_from_shape_and_dtype( + slice_shape, torch.float32 + ) + + return surviving_dequant_bytes <= dequant_bytes + + def postpone_dequantize_op(self, graph_module: torch.fx.GraphModule) -> bool: + packet_to_overload_map = self.dequantize_packet_to_overload + graph = graph_module.graph + modified = False + for node in graph.nodes: + overload_packet = get_overload_packet(node.target) + if ( + overload_packet not in packet_to_overload_map.keys() + or not self.postponing_feasible(node) + ): + continue + + for user in node.users: + with graph.inserting_after(user): + dequant_node = graph.call_function( + getattr( + overload_packet, packet_to_overload_map[overload_packet] + ), + args=(user, *node.args[1:]), + ) + dequant_node.meta = user.meta.copy() + # Remove meta["debug_handle"] on new node if it exists. + # Reassign it at the caller level by calling generate_missing_debug_handles + dequant_node.meta.pop("debug_handle", None) + user.replace_all_uses_with(dequant_node) + dequant_node.args = (user, *node.args[1:]) + + pred = node.args[0] + node.replace_all_uses_with(pred) + graph.erase_node(node) + modified = True + + if modified: + graph_module.recompile() + return modified + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + # The logic in postpone_dequantize_op that handles branching checks the shape + # of the dequant node, which isn't available if that node was already postponed + # in the same pass invokation. The shape information is recreated by tracing in + # super().call(), meaning that every branch in the graph that we wish to postpone + # dequant past requires retracing. We iterate the pass until it no longer modifies + # the graph (up to 3 times max, to avoid potential infinite loops) + self.graph_module = graph_module + iter_count = 0 + local_modified = False + overall_modified = False + + while local_modified or iter_count == 0: + local_modified = self.postpone_dequantize_op(self.graph_module) + overall_modified |= local_modified + + if local_modified: + self.graph_module = super().call(self.graph_module).graph_module + + iter_count += 1 + if iter_count == 3: + break + + return PassResult(self.graph_module, overall_modified) diff --git a/backends/transforms/propagate_slice.py b/backends/transforms/propagate_slice.py new file mode 100644 index 00000000000..e1733553c97 --- /dev/null +++ b/backends/transforms/propagate_slice.py @@ -0,0 +1,222 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from math import prod +from typing import Callable, cast, Optional + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +class PropagateSlice(RemoveOrReplacePassInterface): + """Propagate slice_copy before element-wise ops when the cost model + indicates it reduces total data movement. + + Supported ops (extensible via dispatch table): + - quantize_per_tensor: unary element-wise + - dequantize_per_tensor: unary element-wise + - add.Tensor: binary with broadcast — slices non-broadcasting inputs + - mul.Tensor: binary with broadcast — slices non-broadcasting inputs + + Callers can provide per-target filters for additional legality constraints. + Handles any slice dim and any step size. + """ + + quant_unary_targets: list[EdgeOpOverload] = [ + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ] + + def __init__( + self, + additional_unary_targets: Optional[list[EdgeOpOverload]] = None, + additional_binary_targets: Optional[list[EdgeOpOverload]] = None, + target_filters: Optional[ + dict[EdgeOpOverload, Callable[[torch.fx.Node], bool]] + ] = None, + ) -> None: + super().__init__() + unary_targets = [ + *self.quant_unary_targets, + *(additional_unary_targets or []), + ] + binary_targets = [ + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.mul.Tensor, + *(additional_binary_targets or []), + ] + self._dispatch: dict[ + EdgeOpOverload, + tuple[ + Callable[[torch.fx.Node, torch.fx.Node], bool], + Callable[[torch.fx.Node, torch.fx.Node], bool], + ], + ] = {} + self._target_filters = target_filters or {} + for t in unary_targets: + self._dispatch[t] = ( + self._should_swap_elementwise, + self._swap_elementwise_slice, + ) + + for t in binary_targets: + self._dispatch[t] = ( + self._should_swap_binary_elementwise, + self._swap_binary_elementwise_slice, + ) + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.slice_copy.Tensor] + + def _should_swap_elementwise( + self, op_node: torch.fx.Node, slice_node: torch.fx.Node + ) -> bool: + full_size = prod(op_node.meta["val"].shape) + sliced_size = prod(slice_node.meta["val"].shape) + return sliced_size < full_size + + def _swap_elementwise_slice( + self, op_node: torch.fx.Node, slice_node: torch.fx.Node + ) -> bool: + op_input = op_node.args[0] + assert isinstance(op_input, torch.fx.Node) + graph = slice_node.graph + + slice_dim = get_arg(slice_node, "dim", int) + slice_start = get_arg(slice_node, "start") + slice_end = get_arg(slice_node, "end") + slice_step = get_arg(slice_node, "step", int) + + with graph.inserting_before(op_node): + new_slice = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + args=(op_input, slice_dim, slice_start, slice_end, slice_step), + ) + new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( + op_input.meta["val"], slice_dim, slice_start, slice_end, slice_step + ) + + new_args = list(op_node.args) + new_args[0] = new_slice + target = cast(EdgeOpOverload, op_node.target) + new_op = graph.call_function( + target, + args=tuple(new_args), + kwargs=op_node.kwargs, + ) + new_op.meta["val"] = target( + new_slice.meta["val"], + *[ + a.meta["val"] if isinstance(a, torch.fx.Node) else a + for a in new_args[1:] + ], + **{ + k: v.meta["val"] if isinstance(v, torch.fx.Node) else v + for k, v in op_node.kwargs.items() + }, + ) + + slice_node.replace_all_uses_with(new_op) + graph.erase_node(slice_node) + graph.erase_node(op_node) + return True + + def _should_swap_binary_elementwise( + self, op_node: torch.fx.Node, slice_node: torch.fx.Node + ) -> bool: + lhs, rhs = op_node.args[0], op_node.args[1] + assert isinstance(lhs, torch.fx.Node) and isinstance(rhs, torch.fx.Node) + if lhs.meta["val"].shape == rhs.meta["val"].shape: + return False + full_size = prod(op_node.meta["val"].shape) + sliced_size = prod(slice_node.meta["val"].shape) + return sliced_size < full_size + + def _swap_binary_elementwise_slice( + self, op_node: torch.fx.Node, slice_node: torch.fx.Node + ) -> bool: + lhs, rhs = op_node.args[0], op_node.args[1] + assert isinstance(lhs, torch.fx.Node) and isinstance(rhs, torch.fx.Node) + graph = slice_node.graph + + slice_dim = get_arg(slice_node, "dim", int) + slice_start = get_arg(slice_node, "start") + slice_end = get_arg(slice_node, "end") + slice_step = get_arg(slice_node, "step", int) + + output_shape = op_node.meta["val"].shape + output_dim = slice_dim % len(output_shape) + + new_args = list(op_node.args) + with graph.inserting_before(op_node): + for i, inp in enumerate([lhs, rhs]): + input_shape = inp.meta["val"].shape + # Broadcasting aligns operand dimensions to the right of the output. + input_dim = output_dim - (len(output_shape) - len(input_shape)) + if ( + input_dim >= 0 + and input_shape[input_dim] == output_shape[output_dim] + ): + new_slice = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + args=(inp, input_dim, slice_start, slice_end, slice_step), + ) + new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( + inp.meta["val"], + input_dim, + slice_start, + slice_end, + slice_step, + ) + new_args[i] = new_slice + + target = cast(EdgeOpOverload, op_node.target) + new_op = graph.call_function( + target, + args=tuple(new_args), + kwargs=op_node.kwargs, + ) + new_op.meta["val"] = target( + *[ + a.meta["val"] if isinstance(a, torch.fx.Node) else a + for a in new_args + ], + **{ + k: v.meta["val"] if isinstance(v, torch.fx.Node) else v + for k, v in op_node.kwargs.items() + }, + ) + + slice_node.replace_all_uses_with(new_op) + graph.erase_node(slice_node) + graph.erase_node(op_node) + return True + + def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: + parent = get_arg(node, "input", torch.fx.Node) + if len(parent.users) != 1: + return False + if not isinstance(parent.target, EdgeOpOverload): + return False + + entry = self._dispatch.get(parent.target) + if entry is None: + return False + target_filter = self._target_filters.get(parent.target) + if target_filter is not None and not target_filter(parent): + return False + + should_swap, do_swap = entry + return should_swap(parent, node) and do_swap(parent, node) diff --git a/backends/transforms/quantize_reorder_utils.py b/backends/transforms/quantize_reorder_utils.py new file mode 100644 index 00000000000..ba7a6d47115 --- /dev/null +++ b/backends/transforms/quantize_reorder_utils.py @@ -0,0 +1,51 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import torch +from executorch.exir.dialects._ops import ops as exir_ops + +# A list of ops that can be trivially quantized +trivially_quantizable_ops_overloadpkt = { + exir_ops.edge.aten.chunk, + exir_ops.edge.aten.clone, + exir_ops.edge.aten.contiguous, + exir_ops.edge.aten.expand_copy, + exir_ops.edge.aten.permute_copy, + exir_ops.edge.aten.select_copy, + exir_ops.edge.aten.slice_copy, + exir_ops.edge.aten.squeeze_copy, + exir_ops.edge.aten.transpose_copy, + exir_ops.edge.aten.unfold_copy, + exir_ops.edge.aten.unsqueeze_copy, + exir_ops.edge.aten.view_copy, + torch.ops.aten.chunk, + torch.ops.aten.clone, + torch.ops.aten.contiguous, + torch.ops.aten.expand_copy, + torch.ops.aten.permute, + torch.ops.aten.permute_copy, + torch.ops.aten.select_copy, + torch.ops.aten.slice, + torch.ops.aten.slice_copy, + torch.ops.aten.squeeze, + torch.ops.aten.squeeze_copy, + torch.ops.aten.transpose, + torch.ops.aten.transpose_copy, + torch.ops.aten.unsqueeze, + torch.ops.aten.unsqueeze_copy, + torch.ops.aten.view, + torch.ops.aten.view_copy, +} + +# slice-equivalent ops +slice_or_select_overloadpkt = { + torch.ops.aten.slice_copy, + torch.ops.aten.select_copy, + exir_ops.edge.aten.slice_copy, + exir_ops.edge.aten.select_copy, +} diff --git a/backends/transforms/remove_alias_copy_op.py b/backends/transforms/remove_alias_copy_op.py new file mode 100644 index 00000000000..5dc41f5c7c3 --- /dev/null +++ b/backends/transforms/remove_alias_copy_op.py @@ -0,0 +1,31 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from executorch.backends.transforms.permute_pass_utils import ( + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx.node import Node + + +class RemoveAliasCopyOpPass(RemoveOrReplacePassInterface): + """ + + alias_copy is a no-op and can be removed. + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.alias_copy.default] + + def maybe_remove_or_replace(self, node: Node) -> bool: + input_node = node.args[0] + assert isinstance(input_node, Node) + node.replace_all_uses_with(input_node) + return True diff --git a/backends/transforms/remove_branched_quant_dequant.py b/backends/transforms/remove_branched_quant_dequant.py new file mode 100644 index 00000000000..f57c2d0ecc9 --- /dev/null +++ b/backends/transforms/remove_branched_quant_dequant.py @@ -0,0 +1,80 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import get_edge_overload_packet +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 torch.fx.passes.infra.pass_base import PassBase + + +class RemoveBranchedQuantDequant(PassBase): + """ + This pass looks for adjacent quant and dequant nodes with identical + parameters, where the quant node has other users in addition to the + dequant. The quant and dequant pair would be removed by the + FuseQuantDequantToRequantizePass if not for the multiple users. This pass + removes just the dequant node by connecting it to the quant's parent node + """ + + quantize_op_packets: set[EdgeOpOverloadPacket] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + } + dequantize_op_packets: set[EdgeOpOverloadPacket] = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor, + } + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = self.remove_branched( + graph_module, self.quantize_op_packets, self.dequantize_op_packets + ) + modified |= self.remove_branched( + graph_module, self.dequantize_op_packets, self.quantize_op_packets + ) + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + + return PassResult(graph_module, modified) + + def remove_branched( + self, + graph_module: torch.fx.GraphModule, + producer_pkts: set[EdgeOpOverloadPacket], + consumer_pkts: set[EdgeOpOverloadPacket], + ) -> bool: + modified = False + for node in graph_module.graph.nodes: + if ( + node.op != "call_function" + or not isinstance(node.target, EdgeOpOverload) + or get_edge_overload_packet(node.target) not in producer_pkts + ): + continue + + if len(node.users) < 2: + continue + + for user in node.users: + if ( + not isinstance(user.target, EdgeOpOverload) + or get_edge_overload_packet(user.target) not in consumer_pkts + ): + continue + + # check qparams match + if node.args[1:] != user.args[1:]: + continue + + user.replace_all_uses_with(node.args[0]) + modified = True + + return modified diff --git a/backends/transforms/remove_cat_from_slice_copy.py b/backends/transforms/remove_cat_from_slice_copy.py new file mode 100644 index 00000000000..73a56e15c32 --- /dev/null +++ b/backends/transforms/remove_cat_from_slice_copy.py @@ -0,0 +1,70 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import Optional, Sequence + +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + RemoveOrReplacePassInterface, + set_arg, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx.node import Node + + +class RemoveCatFromSliceCopyPass(RemoveOrReplacePassInterface): + """ + Simplifies cat->slice_copy chains where one of the cat inputs can be directly passed + to the slice_copy. + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.slice_copy.Tensor] + + def maybe_remove_or_replace(self, node: Node) -> bool: + cat_node = get_arg(node, "input", Node) + slice_dim = get_arg(node, "dim", int) + start_idx = get_arg(node, "start", Optional[int]) + end_idx = get_arg(node, "end", Optional[int]) + step = get_arg(node, "step", int) + + if cat_node.target != exir_ops.edge.aten.cat.default or step != 1: + return False + + # Make sure cat and slice happens on the same dimension. + cat_dim = get_arg(cat_node, "dim", int) + if cat_dim != slice_dim: + return False + + # Canonicalize slice indices. + cat_output_shape = cat_node.meta["val"].shape + if start_idx is None: + start_idx = 0 + elif start_idx < 0: + start_idx += cat_output_shape[cat_dim] + if end_idx is None or end_idx > cat_output_shape[cat_dim]: + end_idx = cat_output_shape[cat_dim] + elif end_idx < 0: + end_idx += cat_output_shape[cat_dim] + + offset = 0 + for cat_input_node in get_arg(cat_node, "tensors", Sequence[Node]): + cat_input_shape = cat_input_node.meta["val"].shape + + # Check if the slice range overlaps with the cat input range. + if offset <= start_idx and end_idx <= offset + cat_input_shape[cat_dim]: + node.replace_input_with(cat_node, cat_input_node) + set_arg(node, "start", start_idx - offset) + set_arg(node, "end", end_idx - offset) + return True + + offset += cat_input_shape[cat_dim] + + return False diff --git a/backends/transforms/remove_clone_ops_transform_imported.py b/backends/transforms/remove_clone_ops_transform_imported.py new file mode 100644 index 00000000000..3a9c19262bf --- /dev/null +++ b/backends/transforms/remove_clone_ops_transform_imported.py @@ -0,0 +1,26 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import List + +import torch +import torch.fx +from executorch.backends.transforms.remove_clone_ops import RemoveCloneOpsTransform +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.pass_manager import PassManager, PassType +from executorch.exir.passes import dead_code_elimination_pass + + +class RemoveCloneOpsTransformImported(ExportPass): + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + finalize_passes: List[PassType] = [ + RemoveCloneOpsTransform(eliminate_quant_dequant_pairs=False), + ] + result = PassManager(passes=finalize_passes)(graph_module) + dead_code_elimination_pass(result.graph_module) + return result diff --git a/backends/transforms/remove_nop_as_strided_copy_op.py b/backends/transforms/remove_nop_as_strided_copy_op.py new file mode 100644 index 00000000000..67ac4e983e0 --- /dev/null +++ b/backends/transforms/remove_nop_as_strided_copy_op.py @@ -0,0 +1,52 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import cast, Optional + +import torch +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx.node import Node + + +class RemoveNopAsStridedCopyOpPass(RemoveOrReplacePassInterface): + """Remove as_strided_copy ops that preserve every logical element.""" + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.as_strided_copy.default] + + def maybe_remove_or_replace(self, node: Node) -> bool: + input_node = node.args[0] + assert isinstance(input_node, Node) + input_value = cast(torch.Tensor, input_node.meta["val"]) + size = get_arg(node, "size", list[int]) + stride = get_arg(node, "stride", list[int]) + storage_offset = get_arg(node, "storage_offset", Optional[int]) + + if tuple(size) != tuple(input_value.shape) or len(stride) != len(size): + return False + if ( + storage_offset is not None + and storage_offset != input_value.storage_offset() + ): + return False + if not all( + dim_size == 1 or output_stride == input_stride + for dim_size, output_stride, input_stride in zip( + size, stride, input_value.stride() + ) + ): + return False + + node.replace_all_uses_with(input_node) + return True diff --git a/backends/transforms/remove_nop_expand_op.py b/backends/transforms/remove_nop_expand_op.py new file mode 100644 index 00000000000..3fb23b110d3 --- /dev/null +++ b/backends/transforms/remove_nop_expand_op.py @@ -0,0 +1,36 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from executorch.backends.transforms.permute_pass_utils import ( + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx.node import Node + + +class RemoveNopExpandOpPass(RemoveOrReplacePassInterface): + """ + For an expand op, if the operator shape matches the expand shape, then the + expand is a nop. + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [ + exir_ops.edge.aten.expand_copy.default, + exir_ops.edge.aten.expand.default, + ] + + def maybe_remove_or_replace(self, node: Node) -> bool: + input_node = node.args[0] + assert isinstance(input_node, Node) + if input_node.meta["val"].shape == node.meta["val"].shape: + node.replace_all_uses_with(input_node) + return True + return False diff --git a/backends/transforms/remove_nop_slice_or_view_op.py b/backends/transforms/remove_nop_slice_or_view_op.py new file mode 100644 index 00000000000..f15d095b7f7 --- /dev/null +++ b/backends/transforms/remove_nop_slice_or_view_op.py @@ -0,0 +1,37 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from executorch.backends.transforms.permute_pass_utils import ( + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx.node import Node + + +class RemoveNopSliceOrViewOpPass(RemoveOrReplacePassInterface): + """ + Remove slice ops that are more like views, and view ops that do not change the shape + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [ + exir_ops.edge.aten.slice_copy.Tensor, + exir_ops.edge.aten.view_copy.default, + ] + + def maybe_remove_or_replace(self, node: Node) -> bool: + changed = False + input_node = node.args[0] + assert isinstance(input_node, Node) + if input_node.meta["val"].shape == node.meta["val"].shape: + node.replace_all_uses_with(input_node) + changed = True + + return changed diff --git a/backends/transforms/remove_permute_before_mean.py b/backends/transforms/remove_permute_before_mean.py new file mode 100644 index 00000000000..4bf0d415c86 --- /dev/null +++ b/backends/transforms/remove_permute_before_mean.py @@ -0,0 +1,230 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import Optional + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx.node import Node +from torch.utils import _pytree as pytree + + +class RemovePermuteBeforeMeanPass(RemoveOrReplacePassInterface): + """Remove or sink permute ops that precede mean reductions through unary chains. + + When a permute feeds into a mean (possibly through unary ops like + dequantize/quantize), two optimizations apply: + + 1. If non-reduced dims maintain their relative order and positions, the + permute is fully removed and the mean's reduction dims are remapped. + 2. Otherwise, the permute is moved after the mean so it operates on + smaller data. + + Cost model + ---------- + Let S_in = input size in bytes, S_out = output size in bytes, + R = S_in / S_out (reduction ratio), C_c = contiguous mean compute + cost, C_s = strided mean compute cost, and Delta = C_s - C_c. + + Original graph (permute -> mean): + Cost_orig = 3*S_in + S_out + C_c + Breakdown: permute reads and writes S_in (2*S_in), mean reads + S_in and writes S_out. + + Case 1 -- full removal (mean with remapped dims, no permute): + Cost_remove = S_in + S_out + C_s + Profitable when: Delta < 2*S_in + The strided access penalty must be less than twice the full + input tensor I/O (the eliminated permute cost). + + Case 2 -- reorder (mean with remapped dims -> small permute): + Cost_reorder = S_in + 3*S_out + C_s + Profitable when: Delta < 2*(S_in - S_out) = 2*S_in*(R-1)/R + At R = 4 the threshold is 1.5*S_in; at R = 16 it approaches + 2*S_in, converging to the full-removal bound as S_out becomes + negligible. + + Full removal always dominates reorder when both are structurally + possible (Cost_remove < Cost_reorder since S_out < 3*S_out), so + removal is applied without a cost gate. + + Additionally, if the original permute does not place the reduction + dims as the trailing (innermost) dimensions, the mean is already + strided in the original graph. Removing or sinking the permute + saves the permute I/O (2*S_in) while the mean performance can only + improve or stay the same. This makes the transformation + unconditionally profitable without needing the cost model. + """ + + _UNARY_TARGETS: frozenset[EdgeOpOverload] = frozenset( + { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.aten.clone.default, + exir_ops.edge.aten.relu.default, + exir_ops.edge.aten.neg.default, + exir_ops.edge.aten.abs.default, + } + ) + + _MIN_REDUCTION_RATIO_FOR_REORDER: int = 4 + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.mean.dim] + + def _find_permute_through_unary_chain(self, mean_node: Node) -> Optional[Node]: + """Walk backward from mean through single-user unary ops to find a permute.""" + current = mean_node.args[0] + if not isinstance(current, Node): + return None + while True: + if current.target == exir_ops.edge.aten.permute_copy.default: + return current + if current.target not in self._UNARY_TARGETS: + return None + if len(current.users) != 1: + return None + parent = current.args[0] + if not isinstance(parent, Node): + return None + current = parent + + @staticmethod + def _reduction_dims_are_trailing(reduction_dims: list[int], ndim: int) -> bool: + """Check whether all reduction dims are the trailing (innermost) dims.""" + canonical = sorted(d % ndim for d in reduction_dims) + return canonical == list(range(ndim - len(canonical), ndim)) + + def _is_reorder_profitable( + self, + reduction_dims: list[int], + new_reduction_dims: list[int], + ndim: int, + input_shape: torch.Size, + ) -> bool: + """Determine whether sinking the permute past the mean is profitable. + + Three cases, in order of evaluation: + + 1. The original permute does not place the reduction dims as trailing + (innermost) dimensions. The mean is already strided, so removing + the permute saves 2*S_in with no mean degradation. Always + profitable. + + 2. The new (remapped) reduction dims are trailing in the original + layout. The mean stays contiguous after the transformation + (Delta ~ 0). Always profitable. + + 3. The mean becomes strided. The reorder is profitable when + Delta < 2*S_in*(R-1)/R. We approximate this by requiring + R >= _MIN_REDUCTION_RATIO_FOR_REORDER, ensuring the I/O savings + from operating on a smaller tensor outweigh the strided access + penalty. + """ + if not self._reduction_dims_are_trailing(reduction_dims, ndim): + return True + + if self._reduction_dims_are_trailing(new_reduction_dims, ndim): + return True + + reduction_ratio = 1 + canonical_reduction = set(new_reduction_dims) + for d in canonical_reduction: + reduction_ratio *= input_shape[d] + + return reduction_ratio >= self._MIN_REDUCTION_RATIO_FOR_REORDER + + def maybe_remove_or_replace(self, node: Node) -> bool: + reduction_dims = get_arg(node, "dim", list[int]) + + permute_node = self._find_permute_through_unary_chain(node) + if permute_node is None: + return False + + perm = get_arg(permute_node, "dims", list[int]) + ndim = len(perm) + + if len(permute_node.users) != 1: + return False + + permute_input = permute_node.args[0] + assert isinstance(permute_input, Node) + + new_reduction_dims = [perm[d] for d in reduction_dims] + keepdim = get_arg(node, "keepdim", bool) + + # Determine if the permute can be fully removed (post-mean permute + # would be a no-op) vs needing to be sunk after the mean. + canonical_reduction = set(new_reduction_dims) + if keepdim: + can_remove = all( + perm[d] == d for d in range(ndim) if d not in canonical_reduction + ) + else: + non_reduced_in_perm_order = [ + d for d in perm if d not in canonical_reduction + ] + can_remove = non_reduced_in_perm_order == sorted(non_reduced_in_perm_order) + + # Full removal is almost always profitable. Reorder requires a + # tighter cost bound; verify via the cost model before proceeding. + if not can_remove: + input_shape = permute_input.meta["val"].shape + if not self._is_reorder_profitable( + reduction_dims, new_reduction_dims, ndim, input_shape + ): + return False + + # Rewire: the permute's single user (either the mean itself or the + # first unary op in the chain) should read from the permute's input. + permute_user = next(iter(permute_node.users)) + permute_user.replace_input_with(permute_node, permute_input) + node.args = (node.args[0], new_reduction_dims) + node.args[2:] + + # Re-derive the mean's meta since its reduction dims changed. + fake_args = pytree.tree_map( + lambda x: x.meta["val"] if isinstance(x, Node) else x, + node.args, + ) + node.meta["val"] = node.target(*fake_args) # pyre-ignore[29] + + if not can_remove: + # Compute the post-mean permute on the reduced output. + if keepdim: + post_perm = list(perm) + else: + non_reduced_original = sorted( + d for d in range(ndim) if d not in canonical_reduction + ) + non_reduced_permuted = [d for d in perm if d not in canonical_reduction] + post_perm = [ + non_reduced_original.index(d) for d in non_reduced_permuted + ] + + graph = node.graph + with graph.inserting_after(node): + new_permute = graph.create_node( + "call_function", + exir_ops.edge.aten.permute_copy.default, + args=(node, post_perm), + ) + new_permute.meta["val"] = exir_ops.edge.aten.permute_copy.default( + node.meta["val"], post_perm + ) + for user in list(node.users): + if user is not new_permute: + user.replace_input_with(node, new_permute) + + return True diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index bea6782ba5a..332abdd7101 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -7,6 +7,7 @@ # pyre-unsafe +import math from collections.abc import Callable from dataclasses import dataclass, field from typing import cast @@ -791,6 +792,36 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 if len(perm) != len(inp.meta["val"].shape): return False + # Validate: no node in the region may broadcast operands of different + # rank. Broadcasting right-aligns shapes, so once the boundary permutes + # are removed a lower-rank operand lines up against different axes of + # the others and the region silently computes a differently shaped + # result -- a [216, 1] that met a [1, 1, 216, 1] now meets a + # [1, 216, 1, 1] and yields 216 * 216 elements. The mismatch only + # surfaces later, as a downstream view whose recorded shape no longer + # fits its input. Same-rank broadcasting is safe: every operand carries + # the region's permutation, and broadcasting commutes with applying one + # permutation to all operands. numel-1 operands broadcast identically + # under any permutation, and cat never broadcasts. + for node in subgraph.nodes: + if ( + node in interleave_nodes + or node.target in self._VIEW_OPS + or node.target == exir_ops.edge.aten.cat.default + ): + continue + operand_ranks = { + len(shape) + for shape in ( + self._concrete_shape(inp) + for inp in node.all_input_nodes + if isinstance(inp.meta.get("val"), torch.Tensor) + ) + if shape is not None and math.prod(shape) > 1 + } + if len(operand_ranks) > 1: + return False + # Handle dimension related node arguments FIRST, before # bypassing permutes (which changes node inputs/metadata). for node in subgraph.nodes: diff --git a/backends/transforms/remove_zero_sized_cat_args.py b/backends/transforms/remove_zero_sized_cat_args.py new file mode 100644 index 00000000000..4a78052815d --- /dev/null +++ b/backends/transforms/remove_zero_sized_cat_args.py @@ -0,0 +1,67 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from executorch.backends.transforms.permute_pass_utils import ( + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx.node import Node + + +class RemoveZeroSizedCatArgsPass(RemoveOrReplacePassInterface): + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.cat.default] + + def maybe_remove_or_replace(self, node: Node) -> bool: + # Get the cat inputs (first argument is a list of tensors) + cat_inputs_arg = node.args[0] + + # Assert that cat_inputs_arg is iterable + assert isinstance( + cat_inputs_arg, (list, tuple) + ), "cat_inputs_arg must be a sequence type" + + # Filter out zero-sized tensors + cat_inputs: list[Node] = [] + for arg in cat_inputs_arg: + if isinstance(arg, Node) and arg.meta.get("val") is not None: + if arg.meta["val"].numel() > 0: + cat_inputs.append(arg) + + # If all tensors were empty, create a full op with the right shape + if not cat_inputs: + empty_shape = node.meta["val"].shape + dtype = node.meta["val"].dtype + # Create a new full node + with node.graph.inserting_before(node): + full_node = node.graph.call_function( + exir_ops.edge.aten.full.default, + args=(tuple(empty_shape), 0), + kwargs={"dtype": dtype}, + ) + full_node.meta = node.meta.copy() + node.replace_all_uses_with(full_node) + return True + + # If only one tensor remains, replace with it + if len(cat_inputs) == 1: + node.replace_all_uses_with(cat_inputs[0]) + return True + + # If the number of inputs changed, update the cat args + if len(cat_inputs) < len(cat_inputs_arg): + # Update the first argument with filtered inputs + new_args = list(node.args) + new_args[0] = cat_inputs + node.args = tuple(new_args) + return True + + # No changes needed + return False diff --git a/backends/transforms/replace_select_with_view_op.py b/backends/transforms/replace_select_with_view_op.py new file mode 100644 index 00000000000..07e1317a399 --- /dev/null +++ b/backends/transforms/replace_select_with_view_op.py @@ -0,0 +1,52 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +class ReplaceSelectWithViewOpPass(RemoveOrReplacePassInterface): + """ + If the size along the select dim is 1, then the select op can be replaced + by view op. + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.select_copy.int] + + def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: + # Get the input tensor and shapes + in_tensor_node = node.args[0] + assert isinstance(in_tensor_node, torch.fx.Node) + in_shape = in_tensor_node.meta["val"].shape + out_shape = node.meta["val"].shape + + # Get the select dimension + select_dim = node.args[1] + assert isinstance(select_dim, int) + select_dim = select_dim if select_dim >= 0 else select_dim + len(in_shape) + + if in_shape[select_dim] == 1: + # Replace with view op with the new shape + with node.graph.inserting_before(node): + new_node = node.graph.call_function( + exir_ops.edge.aten.view_copy.default, + args=(node.args[0], list(out_shape)), + ) + # Important to copy metadata + new_node.meta = node.meta + node.replace_all_uses_with(new_node) + return True + + return False diff --git a/backends/transforms/replace_split_with_slice.py b/backends/transforms/replace_split_with_slice.py new file mode 100644 index 00000000000..6272e70e9cf --- /dev/null +++ b/backends/transforms/replace_split_with_slice.py @@ -0,0 +1,89 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import operator +from typing import Optional + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +class ReplaceSplitWithSlicePass(RemoveOrReplacePassInterface): + """ + split_with_sizes() delegates to slice() op, so perform this replacement here. + This avoids the expense of delegation from ATen. + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.split_with_sizes_copy.default] + + def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: + # All the users of this split_with_sizes op must be getitem ops + if any(user.target != operator.getitem for user in node.users): + return False + + # Get the slice dim and extent for each split + slice_ops = self._get_split_sizes(node) + if slice_ops is None: + return False + + graph = node.graph + + # Go over each getitem user, and replace it with slice op + for user in list(node.users.keys()): + assert user.target == operator.getitem + item_idx = int(user.args[1]) + assert item_idx < len(slice_ops) + cur_slice = slice_ops[item_idx] + with graph.inserting_before(user): + cur_slice_node = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + (node.args[0], cur_slice[0], cur_slice[1], cur_slice[2], 1), + ) + # Metadata copy important + cur_slice_node.meta = user.meta + user.replace_all_uses_with(cur_slice_node) + + # Return True to indicate the split node should be removed + return True + + def _get_split_sizes(self, node: torch.fx.Node) -> Optional[list[tuple[int, ...]]]: + """For split_with_sizes, return the slice dim and extent for each split.""" + # Parse the args of the split_with_sizes op + tensor_arg, split_sizes = node.args[0:2] + assert isinstance(tensor_arg, torch.fx.Node) + + # Get shape from node metadata + val = tensor_arg.meta.get("val") + if val is None: + return None + in_shape = val.shape + + split_dim = 0 if len(node.args) < 3 else node.args[2] + + # Canonicalize the split dimension + assert isinstance(split_dim, int) + split_dim = split_dim if split_dim >= 0 else len(in_shape) + split_dim + + # Create the slice op args corresponding to each split + slice_ops = [] + split_start = 0 + assert isinstance(split_sizes, list) + for split_size in split_sizes: + split_end = split_start + split_size + slice_args = (split_dim, split_start, split_end) + slice_ops.append(slice_args) + split_start = split_end + + return slice_ops diff --git a/backends/transforms/split_dequantized_cat.py b/backends/transforms/split_dequantized_cat.py new file mode 100644 index 00000000000..61f0a31526e --- /dev/null +++ b/backends/transforms/split_dequantized_cat.py @@ -0,0 +1,86 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from collections import defaultdict +from typing import DefaultDict, List, Tuple + +import torch +import torch.fx +from executorch.backends.transforms.permute_pass_utils import ( + get_arg, + get_overload_packet, + RemoveOrReplacePassInterface, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +class SplitDequantizedCatPass(RemoveOrReplacePassInterface): + """Split a cat node so that quantize consumers get their own copy. + + Fires when a cat has all floating-point inputs, at least one dequantize + input, and at least one quantize consumer. Quant consumers are grouped + by matching qparams; each group receives a dedicated duplicate of the + cat node. Non-quant consumers stay on the original cat, whose + semantics are unchanged. + + A later pass (e.g. AdvanceQuantizeOpAboveDefChainPass extended for cat) + can then hoist each quant above its single-consumer cat copy without + affecting the non-quant paths. + """ + + quantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + } + dequantize_op_packets: set[object] = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor, + } + + @property + def targets(self) -> list[EdgeOpOverload]: + return [exir_ops.edge.aten.cat.default] + + def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: + cat_inputs = node.args[0] + if not isinstance(cat_inputs, (list, tuple)): + return False + + has_dequant_input = False + for inp in cat_inputs: + assert isinstance(inp, torch.fx.Node) + val = inp.meta["val"] + if val is None or not val.is_floating_point(): + return False + if get_overload_packet(inp.target) in self.dequantize_op_packets: + has_dequant_input = True + + if not has_dequant_input: + return False + + quant_groups: DefaultDict[Tuple, List[torch.fx.Node]] = defaultdict(list) + for user in list(node.users.keys()): + if get_overload_packet(user.target) in self.quantize_op_packets: + quant_groups[user.args[1:]].append(user) + + if not quant_groups: + return False + + graph = node.graph + dim = get_arg(node, "dim", int) + for quant_consumers in quant_groups.values(): + with graph.inserting_after(node): + dup_cat = graph.call_function( + exir_ops.edge.aten.cat.default, + args=(list(cat_inputs), dim), + ) + dup_cat.meta = node.meta.copy() + + for q_node in quant_consumers: + q_node.replace_input_with(node, dup_cat) + + return True diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 4dd160aa44d..4cd4db88fa6 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -920,3 +920,303 @@ def define_common_targets(): "fbsource//third-party/pypi/pytest:pytest", ], ) + + runtime.python_library( + name = "quantize_reorder_utils", + srcs = ["quantize_reorder_utils.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_library( + name = "fuse_quant_dequant_to_requantize", + srcs = ["fuse_quant_dequant_to_requantize.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + "//executorch/exir:pass_base", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_alias_copy_op", + srcs = ["remove_alias_copy_op.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_branched_quant_dequant", + srcs = ["remove_branched_quant_dequant.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + "//executorch/exir:pass_base", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_cat_from_slice_copy", + srcs = ["remove_cat_from_slice_copy.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_clone_ops_transform_imported", + srcs = ["remove_clone_ops_transform_imported.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/passes:lib", + "//executorch/exir:pass_base", + "//executorch/exir:pass_manager", + ":remove_clone_ops", + ], + ) + + runtime.python_library( + name = "remove_nop_expand_op", + srcs = ["remove_nop_expand_op.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_nop_slice_or_view_op", + srcs = ["remove_nop_slice_or_view_op.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_permute_before_mean", + srcs = ["remove_permute_before_mean.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_zero_sized_cat_args", + srcs = ["remove_zero_sized_cat_args.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "advance_quantize_op_above_def_in_branch", + srcs = ["advance_quantize_op_above_def_in_branch.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects:lib", + "//executorch/exir:pass_base", + ":permute_pass_utils", + ":quantize_reorder_utils", + ], + ) + + runtime.python_library( + name = "advance_quantize_op_above_def_chain", + srcs = ["advance_quantize_op_above_def_chain.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + "//executorch/exir:pass_base", + ":permute_pass_utils", + ":quantize_reorder_utils", + ], + ) + + runtime.python_library( + name = "postpone_dequantize_op_below_use_chain", + srcs = ["postpone_dequantize_op_below_use_chain.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects:lib", + "//executorch/exir:pass_base", + "//executorch/exir:tensor", + ":permute_pass_utils", + ":quantize_reorder_utils", + ], + ) + + runtime.python_library( + name = "move_permute_after_concat", + srcs = ["move_permute_after_concat.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "move_slice_before_permute", + srcs = ["move_slice_before_permute.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "move_slice_before_view", + srcs = ["move_slice_before_view.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "propagate_slice", + srcs = ["propagate_slice.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "split_dequantized_cat", + srcs = ["split_dequantized_cat.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "replace_select_with_view_op", + srcs = ["replace_select_with_view_op.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "replace_split_with_slice", + srcs = ["replace_split_with_slice.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/dialects:lib", + ":permute_pass_utils", + ], + ) + + runtime.python_library( + name = "remove_nop_as_strided_copy_op", + srcs = ["remove_nop_as_strided_copy_op.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/dialects:lib", + "//executorch/exir/dialects/edge:lib", + ":permute_pass_utils", + ], + ) diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 7d1f507ccf0..d5f16609c32 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1756,6 +1756,71 @@ def test_permutation_sink_view_terminal_broadcast_is_optimized(self) -> None: "permutation_sink_view_terminal_broadcast_is_optimized", ) + def test_numel_one_operand_does_not_block_region(self) -> None: + x_data = torch.randn(1, 4, 2, 3) + scale_data = torch.tensor(0.5) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + scale = builder.placeholder("scale", scale_data) + permute_in = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) + ) + mul = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(permute_in, scale) + ) + permute_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [0, 3, 1, 2]) + ) + builder.output([permute_out]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data, scale_data], + "numel_one_operand_does_not_block_region", + ) + + def test_cat_with_different_concat_sizes_does_not_block_region(self) -> None: + x_data = torch.randn(1, 4, 2, 2) + y_data = torch.randn(1, 6, 2, 2) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + y = builder.placeholder("y", y_data) + permute_x = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) + ) + permute_y = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(y, [0, 2, 3, 1]) + ) + cat = builder.call_operator( + op=exir_ops.edge.aten.cat.default, args=([permute_x, permute_y], 3) + ) + permute_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(cat, [0, 3, 1, 2]) + ) + builder.output([permute_out]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data, y_data], + "cat_with_different_concat_sizes_does_not_block_region", + ) + def test_split_sink_view_is_not_remapped_for_broadcast(self) -> None: x_data = torch.randn(1, 4, 2) sink_data = torch.randn(1, 1, 1, 8)