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..11ebc487d3c 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: @@ -463,9 +461,9 @@ def _extract_batchnorm_params( from executorch.exir.dialects.edge._ops import EdgeOpOverload - assert isinstance( - bn_node.target, EdgeOpOverload - ), f"Expected EdgeOpOverload, got {type(bn_node.target)}" + assert isinstance(bn_node.target, EdgeOpOverload), ( + f"Expected EdgeOpOverload, got {type(bn_node.target)}" + ) # Extract parameters by name (not by index for maintainability) # The get_arg function handles normalization of positional args to kwargs @@ -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, - ) - # Remove any quantize op that has only dequantze ops as its users. - modified |= self.find_and_fuse( - graph_module, - producer_op_packets=self.quantize_op_packets, - consumer_op_packets=self.dequantize_op_packets, - # Do not requantize for quantize-dequantize pairs as this is not guaranteed - # to be better for performance/memory. - # Only fuse if all users of quant are dequant. - bypass_ops=( - self.bypass_ops - if self.force_quant_dequant_fusion - else {exir_ops.edge.aten.view_copy.default} - ), + super().__init__( + allow_requantize=allow_requantize, + force_quant_dequant_fusion=force_quant_dequant_fusion, ) - if modified: - graph_module.graph.eliminate_dead_code() - graph_module.recompile() - return PassResult(graph_module, True) - return PassResult(graph_module, False) class FuseMulScalarIntoDequantPass(RemoveOrReplacePassInterface): @@ -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 @@ -1147,9 +953,9 @@ class HierarchicalCSEPass(HierarchicalInplacePassInterface): def _apply_flat_inplace(self, graph_module) -> bool: # Call the CSE pass on the main graph, which performs CSE without recursing. result = CSEPass().call(graph_module) - assert ( - result.graph_module is graph_module - ), f"Only in-place modification is allowed, but got {result.graph_module}" + assert result.graph_module is graph_module, ( + f"Only in-place modification is allowed, but got {result.graph_module}" + ) return result.modified 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 9542704c8ab..068d61b278b 100644 --- a/backends/cadence/aot/reorder_ops.py +++ b/backends/cadence/aot/reorder_ops.py @@ -9,626 +9,78 @@ # 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 the quant descendents. Cluster them into sets - # that have the same scale, zero_point, and dtype. We use quant_dict - # for the clustering - quant_dict: DefaultDict[Tuple, int] = defaultdict(int) - for quant_node in descendent_quant_ops: - quant_dict[quant_node.args[1:]] += 1 - rep_args = sorted(quant_dict.keys(), key=lambda x: x[1]).pop() - - # 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", - 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", + exir_ops.edge.quantized_decomposed.dequantize_per_channel: "default", + exir_ops.edge.cadence.dequantize_per_tensor: "default", + } class SinkOpsCloserToUsePass(RemoveOrReplacePassInterface): @@ -771,674 +223,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 f68b05857d2..5fcacb31f1a 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): @@ -533,9 +505,9 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: # the output of repeat will be a higher-dimensional tensor. We reshape # the input so that it has the same dimensionality as the output tensor. diff = len(repeats) - len(in_shape) - assert ( - diff >= 0 - ), "Repeat arg malformed: expected a repeat along each dimension of input tensor" + assert diff >= 0, ( + "Repeat arg malformed: expected a repeat along each dimension of input tensor" + ) graph = node.graph result_node = in_tensor @@ -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/fused_quant/BUCK b/backends/fused_quant/BUCK new file mode 100644 index 00000000000..0a503173691 --- /dev/null +++ b/backends/fused_quant/BUCK @@ -0,0 +1,6 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +fbcode_target(_kind = define_common_targets) diff --git a/backends/fused_quant/README.md b/backends/fused_quant/README.md new file mode 100644 index 00000000000..8deee74310d --- /dev/null +++ b/backends/fused_quant/README.md @@ -0,0 +1,87 @@ +# fused_quant + +`fused_quant` is a backend-neutral quantization frontend for ExecuTorch. It represents a +`dequantize -> aten op -> quantize` triple as a single `torch.ops.fused_quant.*` operator +that carries its quantization parameters inline, so the quantized graph stays explicit and +optimizable instead of being scattered across q/dq pairs. + +It is not a delegate. Each operator has a reference (eager) kernel that defines its +numerics, but no runtime kernel — a `fused_quant` op never reaches `to_executorch`. It +produces an edge-dialect graph that a backend can color, lower, and partition; anything +the backend does not claim is turned back into q/dq + ATen by `DecomposeFusedQuant` and +runs on the ordinary portable kernels. `test_compile.py` shows that decomposition tail +written out for the no-backend case. + +A node a backend colored but its partitioner then rejected is not decomposed by +`DecomposeFusedQuant` before `to_backend`. Run `DecomposeFusedQuant` again after +`to_backend` to catch those, then `AssertNoFusedQuantOps` before `to_executorch`. + +## Directory structure + +``` +backends/fused_quant/ +├── ops.py fused_quant operator schemas, reference and meta kernels +├── ops_utils.py shape helpers shared by the op definitions +├── graph_utils.py qparam / constant / ExportedProgram helpers +├── fuse_aten.py generic dequant -> op -> quant folding +├── pass_base.py ExportedProgram pass base + IterativePassGroup +├── passes.py the backend-neutral pass pipeline built from those +├── decompose_fused_quant.py reverse lowering back to q/dq + ATen +├── colorer.py ColorerBase (claims ops for a backend) +├── frontend.py trace / prepare +├── compile.py compile_to_fused_quant, the backend seam +├── quantizer/ torchao PT2E quantizer that also fuses +├── pre_quantize_passes/ rewrites applied before annotation +├── optimization_passes/ rewrites applied to the fused_quant edge graph +└── test/ +``` + +## Operator representation + +Every `fused_quant` op takes its tensor operands first, then one qparam block +(`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per quantizable edge, including +the output. A qparam block whose `scale` is `None` means that edge is left in float, so a +single schema covers fully quantized, partially quantized, and float instances of an op. + +Granularity is carried by the shape of the `scale` tensor rather than by separate op +variants. The qparams are always shaped to broadcast against the tensor they describe: +per-tensor is a singleton tensor (`scale.numel() == 1`), per-channel has exactly one +non-unary dimension sitting on the quantized axis, and per-group and blockwise tile the +tensor by the ratio of the two shapes (`tensor.shape[i] // scale.shape[i]` per dimension). +A per-tensor scale may be shaped `()` or all-ones; test granularity with +`QuantParamsStruct.is_per_tensor()` / `is_per_channel(tensor_shape)`, not the shape. + +## Operator contract + +Backends lower these ops, so the following is stable: + +- **Argument layout.** Tensor operands, then one qparam block per tensor operand in the + same order, then the `out` block, then the ATen op's remaining arguments in ATen + order. A block is always the five fields above, named `_`. +- **Field meaning.** For an input block, `dtype` is the dequantized float dtype; for the + `out` block, it is the quantized storage dtype. A `None` scale/zero point means the edge + is float. +- **Evolving an op.** Existing schemas are not edited in place. A new capability is a new + overload, or a new trailing argument with a default that preserves current behavior. +- **Ahead-of-time only.** `fused_quant` ops have no runtime kernels and must never be + serialized into a `.pte`. `AssertNoFusedQuantOps` enforces this. +## Quantizer + +`quantizer/quantizer.py` defines a torchao `Quantizer` that both annotates and fuses: + +- `QuantizerBase` — the interface: annotate nodes, then fuse the annotated ones. +- `OpQuantizer` — annotates one ATen op type and folds matched instances into a + `fused_quant` op. Takes a `QuantizationConfig` plus the schema names of its quantized + activations, weights, and passthrough operands. Most quantizers are instances of this + rather than subclasses; subclass it only when an op needs custom fusion. +- `NoopQuantizer` — claims nodes with an empty annotation so a later quantizer skips them, + leaving the op in float. +- `FusedQuantQuantizer` — composes the above. Annotation is first-match-wins, so a narrower + quantizer must be listed before the general one it overrides. + +`preserved_ops()` reports the ops the quantizer matches; `trace` keeps them intact +through `run_decompositions` and uses them to pick which pre-quantize rewrites to run. + +## Known gaps + +- The package depends on `torchao`, which is not yet a core ExecuTorch dependency. diff --git a/backends/fused_quant/__init__.py b/backends/fused_quant/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/fused_quant/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/fused_quant/colorer.py b/backends/fused_quant/colorer.py new file mode 100644 index 00000000000..33be2680822 --- /dev/null +++ b/backends/fused_quant/colorer.py @@ -0,0 +1,91 @@ +# 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. + +import logging +from abc import abstractmethod + +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from torch.export import ExportedProgram +from torch.fx import Node + +# Meta key a colorer writes to claim a node for its backend. ExecuTorch's +# partitioning infrastructure uses the same key, so a claim made here is the +# same claim the partitioner will act on. +DELEGATION_TAG = "delegation_tag" + + +def is_colored(node: Node) -> bool: + """Whether a backend has claimed this node. + + The generic fused_quant passes use this to leave claimed nodes alone; see + :class:`ColorerBase` for the contract. + """ + return DELEGATION_TAG in node.meta + + +class ColorerBase(ExportedProgramPassBase): + """Tags the nodes a backend will take, by writing ``delegation_tag`` meta. + + Coloring is the step between optimization and lowering: it decides which + fused_quant ops the backend claims, so the lowering pass knows what to rewrite + and the partitioner knows what to delegate. A backend subclasses this and + answers two questions -- which ops it can run (:meth:`supported_ops`) and + whether a given instance is within its limits (:meth:`is_legal`). + + The tag is a claim of ownership, and the generic passes honour it: + :class:`~executorch.backends.fused_quant.optimization_passes.to_channels_last.ToChannelsLast` + and + :class:`~executorch.backends.fused_quant.optimization_passes.to_channels_first.ToChannelsFirst` + will not relayout a colored node, and + :class:`~executorch.backends.fused_quant.decompose_fused_quant.DecomposeFusedQuant` + will not decompose one. So the usual pipeline is colorer -> ToChannelsFirst + -> DecomposeFusedQuant: the backend keeps what it claimed, and everything + else falls back to quantize/dequantize plus ATen. That is what makes + claiming an op sufficient to keep it intact -- a backend describes what it + supports, not what it does not. Passes test this with :func:`is_colored`. + + Colorers compose by priority: run them in order and the first to claim a + node keeps it. A node that already carries a ``delegation_tag`` is skipped, + never overwritten, so a later (lower-priority) colorer only sees what the + earlier ones left. + """ + + def __init__(self, delegation_tag: str) -> None: + super().__init__() + self.delegation_tag = delegation_tag + + @abstractmethod + def supported_ops(self) -> set[EdgeOpOverload]: + """Edge ops this backend can run, before per-node legality is applied.""" + ... + + def is_legal(self, exported_program: ExportedProgram, node: Node) -> bool: + """Whether this specific node is within the backend's limits. + + Defaults to True: an op listed in :meth:`supported_ops` with no further + constraints is always claimed. Override to reject instances whose + qparams, attributes, or users the backend cannot handle. + """ + return True + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + modified = False + graph = exported_program.graph_module.graph + + for target in self.supported_ops(): + for node in graph.find_nodes(op="call_function", target=target): + if is_colored(node): + continue + if self.is_legal(exported_program, node): + node.meta[DELEGATION_TAG] = self.delegation_tag + modified = True + else: + logging.info( + f"Could not color {node.name} to run on {self.delegation_tag}" + ) + + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/compile.py b/backends/fused_quant/compile.py new file mode 100644 index 00000000000..fb307c3553d --- /dev/null +++ b/backends/fused_quant/compile.py @@ -0,0 +1,85 @@ +# 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. + +from collections.abc import Sequence +from typing import Optional + +import torch +from executorch.backends.fused_quant.frontend import prepare, trace +from executorch.backends.fused_quant.passes import get_fused_quant_passes +from executorch.backends.fused_quant.quantizer.defaults import ( + make_fused_quant_quantizer, +) +from executorch.backends.fused_quant.quantizer.quantizer import FusedQuantQuantizer +from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge +from executorch.exir.pass_manager import ExportedProgramPassManager, PassType +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e + + +def compile_to_fused_quant( + model: torch.nn.Module, + inputs: Sequence[tuple[object, ...]], + quantizer: Optional[FusedQuantQuantizer] = None, + *, + passes: Optional[Sequence[PassType]] = None, + fold_quantize_into_mutable_buffers: bool = False, +) -> EdgeProgramManager: + """Quantize, fuse, and optimize ``model`` into the fused_quant edge graph. + + Runs trace -> prepare -> calibrate -> convert -> fuse -> export -> to_edge -> + ``passes``, and stops there. The result is backend-independent: a mixture of + ATen and fused_quant ops in edge dialect, before any coloring or lowering. + This is the seam a backend plugs into. + + From here a backend runs its own colorer, lowering, and partitioner. What it + does not claim must be handled by ``DecomposeFusedQuant``, which turns the + remaining fused_quant ops back into quantize/dequantize plus ATen; see + ``test_compile.py`` for that tail written out. + + Args: + inputs: The calibration set. ``inputs[0]`` is also the example input for + tracing. + quantizer: Defaults to :func:`make_fused_quant_quantizer`. + passes: Edge passes to run after ``to_edge``. Defaults to + :func:`get_fused_quant_passes`; pass an explicit list to use a + different pipeline, or an empty one to skip optimization entirely. + """ + if quantizer is None: + quantizer = make_fused_quant_quantizer() + if passes is None: + passes = get_fused_quant_passes() + + program = trace(model, inputs[0], quantizer) + graph_module = prepare(program, quantizer) + + for inp in inputs: + graph_module(*inp) + + graph_module = convert_pt2e( + graph_module, + fold_quantize_into_mutable_buffers=fold_quantize_into_mutable_buffers, + ) + + # Fusion runs before export/to_edge, on ATen IR. The quantizer already + # identified which ops to quantize, so fusion can match directly on the + # quantized_decomposed ops that convert introduced. + quantizer.fuse(graph_module) + + # _skip_dim_order keeps classic aten copies instead of dim_order_ops + # (_to_dim_order_copy), which not every runtime registers. + # + # The preserved ops are held back from decomposition precisely because the + # quantizer fuses them, and DecomposeFusedQuant re-emits them if a backend + # does not claim the fused op. Several (aten.linear, aten.silu, ...) are not + # core ATen, so the edge verifier needs to know they are expected. + edge_program_manager = to_edge( + torch.export.export(graph_module, inputs[0]), + compile_config=EdgeCompileConfig( + _skip_dim_order=True, + _core_aten_ops_exception_list=list(quantizer.preserved_ops()), + ), + ) + return edge_program_manager.transform(ExportedProgramPassManager(list(passes))) diff --git a/backends/fused_quant/decompose_fused_quant.py b/backends/fused_quant/decompose_fused_quant.py new file mode 100644 index 00000000000..cebf4308e94 --- /dev/null +++ b/backends/fused_quant/decompose_fused_quant.py @@ -0,0 +1,459 @@ +# 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. + +from __future__ import annotations + +import operator +from typing import cast + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.colorer import is_colored +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_from_node, + get_scale, + get_zero_point, + split_fused_arg_names, +) +from executorch.backends.fused_quant.ops import QuantParamsStruct +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import ( + add_constant, + compute_meta_val, + delete_constant_placeholder, + get_constant, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from torch import fx +from torch._C import Argument +from torch._ops import OpOverload +from torch._subclasses.fake_tensor import UnsupportedOperatorException +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +_OpOverload = EdgeOpOverload | OpOverload + + +def _copy_meta_and_compute(node: fx.Node, source: fx.Node) -> None: + node.meta = source.meta.copy() + node.meta["val"] = compute_meta_val(node) + + +def _aten_meta_from_fused( + aten_node: fx.Node, + fused_node: fx.Node, + out_qparams: QuantParamsStruct[fx.Node] | None, + float_dtype: torch.dtype, +) -> None: + """Set the meta of the ATen node standing in for a fused op. + + Some ATen ops (e.g. ``_masked_softmax``) have no fake kernel, so their meta + cannot be computed. The fused op's own fake kernel already has the shape; + only the dtype differs when the output was quantized. + """ + aten_node.meta = fused_node.meta.copy() + try: + aten_node.meta["val"] = compute_meta_val(aten_node) + except UnsupportedOperatorException: + fused_val = fused_node.meta["val"] + if isinstance(fused_val, (tuple, list)): + raise + aten_node.meta["val"] = ( + fused_val.to(float_dtype) if out_qparams is not None else fused_val + ) + + +def _qparam_granularity( + qparams: QuantParamsStruct[fx.Node], + tensor_value: torch.Tensor, +) -> str: + qparams.validate() + scale_value = qparams.scale.meta["val"] + if qparams.is_per_tensor(): + return "per_tensor" + if qparams.is_per_channel(tensor_value.shape): + return "per_channel" + raise ValueError( + "Cannot decompose grouped or blockwise fused quantization with scale shape " + f"{tuple(scale_value.shape)} for tensor shape {tuple(tensor_value.shape)}; " + "only per-tensor and per-channel quantization are supported" + ) + + +def _per_channel_qparam_args( + exported_program: ExportedProgram, + graph: fx.Graph, + qparams: QuantParamsStruct[fx.Node], +) -> tuple[fx.Node, fx.Node, int]: + """The 1-D scale/zero_point a per-channel q/dq op takes, plus its axis. + + Constant qparams become new 1-D constants rather than a view of the + full-rank one, so partitioners that require per-channel qparams to be + placeholders can consume the decomposed graph without a constant fold. + """ + axis = qparams.channel_axis() + assert axis is not None + flattened: list[fx.Node] = [] + for qparam in (qparams.scale, qparams.zero_point): + value = get_constant(exported_program, qparam) + if value is not None: + flattened.append( + add_constant( + exported_program, + f"{qparam.name}_flat", + value.reshape(-1), + qparam, + InputKind.CONSTANT_TENSOR, + ) + ) + continue + view = graph.call_function( + exir_ops.edge.aten.view_copy.default, + args=(qparam, [-1]), + ) + _copy_meta_and_compute(view, qparam) + flattened.append(view) + return flattened[0], flattened[1], axis + + +def _dequantize( + exported_program: ExportedProgram, + graph: fx.Graph, + tensor: fx.Node | None, + qparams: QuantParamsStruct[fx.Node] | None, +) -> fx.Node | None: + if tensor is None or qparams is None: + return tensor + + if _qparam_granularity(qparams, tensor.meta["val"]) == "per_tensor": + target = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + args = ( + tensor, + get_scale(exported_program, qparams), + get_zero_point(exported_program, qparams), + qparams.quant_min, + qparams.quant_max, + tensor.meta["val"].dtype, + ) + else: + scale, zero_point, axis = _per_channel_qparam_args( + exported_program, graph, qparams + ) + target = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default + args = ( + tensor, + scale, + zero_point, + axis, + qparams.quant_min, + qparams.quant_max, + tensor.meta["val"].dtype, + ) + dequantized = graph.call_function( + target, + args=args, + kwargs={"out_dtype": qparams.dtype}, + ) + _copy_meta_and_compute(dequantized, tensor) + return dequantized + + +def _quantize( + exported_program: ExportedProgram, + graph: fx.Graph, + tensor: fx.Node, + qparams: QuantParamsStruct[fx.Node] | None, + source: fx.Node, +) -> fx.Node: + if qparams is None: + return tensor + + if _qparam_granularity(qparams, tensor.meta["val"]) == "per_tensor": + target = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + args = ( + tensor, + get_scale(exported_program, qparams), + get_zero_point(exported_program, qparams), + qparams.quant_min, + qparams.quant_max, + qparams.dtype, + ) + else: + scale, zero_point, axis = _per_channel_qparam_args( + exported_program, graph, qparams + ) + target = exir_ops.edge.quantized_decomposed.quantize_per_channel.default + args = ( + tensor, + scale, + zero_point, + axis, + qparams.quant_min, + qparams.quant_max, + qparams.dtype, + ) + quantized = graph.call_function(target, args=args) + _copy_meta_and_compute(quantized, source) + return quantized + + +def _get_aten_target( + fused_target: _OpOverload, + fused_name: str, + input_count: int, + extra_names: list[str], +) -> EdgeOpOverload: + packet = getattr(exir_ops.edge.aten, fused_name, None) + if packet is None: + packet = getattr(exir_ops.edge.aten, fused_name.replace("_", ""), None) + if packet is None: + raise ValueError(f"No ATen operator matches fused_quant::{fused_name}") + + overload_name = fused_target._schema.overload_name + candidate_names = [overload_name, "default", "Tensor"] + candidate_names.extend(getattr(packet, "_overload_names", ())) + + for candidate_name in (name for name in dict.fromkeys(candidate_names) if name): + candidate = getattr(packet, candidate_name, None) + if candidate is None: + continue + arguments = candidate._schema.arguments + if ( + not any(argument.name == "out" for argument in arguments) + and len(arguments) == input_count + len(extra_names) + and all( + "Tensor" in str(argument.type) for argument in arguments[:input_count] + ) + and [argument.name for argument in arguments[input_count:]] == extra_names + ): + return candidate + raise ValueError(f"No functional ATen overload matches fused_quant::{fused_name}") + + +def _replace_multi_output( + exported_program: ExportedProgram, + node: fx.Node, + aten_node: fx.Node, + out_qparams: QuantParamsStruct[fx.Node] | None, +) -> None: + users = list(node.users) + if any( + user.target is not operator.getitem or not isinstance(user.args[1], int) + for user in users + ): + raise ValueError( + f"Multi-output fused op {node.target} must be consumed through getitem" + ) + + graph = node.graph + replacements: dict[int, fx.Node] = {} + for user in users: + output_index = cast(int, user.args[1]) + replacement = replacements.get(output_index) + if replacement is None: + replacement = graph.call_function( + operator.getitem, + args=(aten_node, output_index), + ) + _copy_meta_and_compute(replacement, user) + if output_index == 0: + replacement = _quantize( + exported_program, graph, replacement, out_qparams, user + ) + replacements[output_index] = replacement + user.replace_all_uses_with(replacement) + graph.erase_node(user) + graph.erase_node(node) + + +def _decompose_requantize( + exported_program: ExportedProgram, + node: fx.Node, +) -> None: + graph = node.graph + inp = get_arg(node, "inp", fx.Node) + inp_qparams = get_qparams_from_node(node, "inp") + out_qparams = get_qparams_from_node(node, "out") + assert inp_qparams is not None and out_qparams is not None + + with graph.inserting_before(node): + dequantized = _dequantize(exported_program, graph, inp, inp_qparams) + assert dequantized is not None + quantized = _quantize(exported_program, graph, dequantized, out_qparams, node) + node.replace_all_uses_with(quantized) + graph.erase_node(node) + + +def _is_schema_default(argument: Argument, value: fx.node.Argument) -> bool: + if not argument.has_default_value(): + return False + try: + return bool(value == argument.default_value) + except (RuntimeError, ValueError): + # Comparison on tensor-valued arguments is not a scalar bool. + return False + + +def _qparam_blocks( + node: fx.Node, +) -> list[tuple[QuantParamsStruct[fx.Node], torch.Tensor]]: + """Every qparams block of a fused node, with the tensor value it applies to.""" + input_names, _ = split_fused_arg_names(cast(_OpOverload, node.target)) + blocks: list[tuple[QuantParamsStruct[fx.Node], torch.Tensor]] = [] + for name in input_names: + qparams = get_qparams_from_node(node, name) + if qparams is not None: + blocks.append((qparams, get_arg(node, name, fx.Node).meta["val"])) + out_qparams = get_qparams_from_node(node, "out") + if out_qparams is not None: + output_value = node.meta["val"] + if isinstance(output_value, (tuple, list)): + output_value = output_value[0] + blocks.append((out_qparams, output_value)) + return blocks + + +def _decompose_fused_op(exported_program: ExportedProgram, node: fx.Node) -> None: + fused_target = cast(_OpOverload, node.target) + fused_name = fused_target._schema.name.split("::", 1)[1] + if fused_name == "requantize": + _decompose_requantize(exported_program, node) + return + + input_names, extra_names = split_fused_arg_names(fused_target) + input_tensors = [get_arg(node, name, fx.Node | None) for name in input_names] + input_qparams = [get_qparams_from_node(node, name) for name in input_names] + out_qparams = get_qparams_from_node(node, "out") + graph = node.graph + + with graph.inserting_before(node): + inputs: list[fx.node.Argument] = [] + for tensor, qparams in zip(input_tensors, input_qparams): + inputs.append(_dequantize(exported_program, graph, tensor, qparams)) + + extra_values = [get_arg(node, name) for name in extra_names] + aten_target = _get_aten_target( + fused_target, fused_name, len(inputs), extra_names + ) + args = list(inputs) + kwargs: dict[str, fx.node.Argument] = {} + for argument, value in zip( + aten_target._schema.arguments[len(inputs) :], + extra_values, + strict=True, + ): + if not argument.kwarg_only: + args.append(value) + elif not _is_schema_default(argument, value): + # A kwarg pinned to its default is one that torch.export would + # have left off the node entirely. Emitting it anyway makes the + # rebuilt node wider than the ATen op it is standing in for, and + # backends that forward kwargs verbatim choke on the extra + # argument -- e.g. alpha=1.0 reaching the two-operand tosa.SUB. + kwargs[argument.name] = value + aten_node = graph.call_function(aten_target, args=tuple(args), kwargs=kwargs) + float_dtype = next( + ( + cast(torch.Tensor, arg.meta["val"]).dtype + for arg in inputs + if isinstance(arg, fx.Node) + and cast(torch.Tensor, arg.meta["val"]).is_floating_point() + ), + torch.float32, + ) + _aten_meta_from_fused(aten_node, node, out_qparams, float_dtype) + + if isinstance(node.meta.get("val"), (tuple, list)): + _replace_multi_output(exported_program, node, aten_node, out_qparams) + return + + output = _quantize(exported_program, graph, aten_node, out_qparams, node) + + node.replace_all_uses_with(output) + graph.erase_node(node) + + +class DecomposeFusedQuant(ExportedProgramPassBase): + """Decompose fused-quant ops into per-tensor/channel dq, ATen, and q ops. + + Colored nodes are left alone: a backend that claimed an op is going to lower + it itself, and defusing it would take it away. See + :class:`executorch.backends.fused_quant.colorer.ColorerBase`. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph_module = exported_program.graph_module + graph = graph_module.graph + nodes = [ + node + for node in graph.nodes + if is_fused_quant_node(node) and not is_colored(node) + ] + if not nodes: + return ExportedProgramPassResult(exported_program, False) + + # Validate everything first so an unsupported node fails the pass before + # any of the graph has been rewritten. + qparam_placeholders: set[fx.Node] = set() + for node in nodes: + for qparams, value in _qparam_blocks(node): + _qparam_granularity(qparams, value) + qparam_placeholders.update( + qparam + for qparam in (qparams.scale, qparams.zero_point) + if qparam.op == "placeholder" + ) + + for node in nodes: + _decompose_fused_op(exported_program, node) + + graph.eliminate_dead_code() + for placeholder in qparam_placeholders: + if ( + not placeholder.users + and get_constant(exported_program, placeholder) is not None + ): + delete_constant_placeholder(exported_program, placeholder) + graph.lint() + graph_module.recompile() + return ExportedProgramPassResult(exported_program, True) + + def ensures(self, exported_program: ExportedProgram) -> None: + exported_program.validate() + + +def is_fused_quant_node(node: fx.Node) -> bool: + return ( + node.op == "call_function" + and isinstance(node.target, (EdgeOpOverload, OpOverload)) + and node.target.namespace == "fused_quant" + ) + + +class AssertNoFusedQuantOps(ExportedProgramPassBase): + """Fail if any fused_quant op is left in the graph. + + fused_quant ops are an ahead-of-time IR with no runtime kernels. Run this + after lowering, once every backend has taken what it claimed and + :class:`DecomposeFusedQuant` has decomposed the rest, so a leak surfaces here + rather than as a missing out-variant at ``to_executorch``. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + leaked = [ + node.name + for node in exported_program.graph_module.graph.nodes + if is_fused_quant_node(node) + ] + if leaked: + raise RuntimeError( + f"fused_quant ops remain after lowering: {leaked}. Colored nodes a " + "backend did not lower must be decomposed with DecomposeFusedQuant " + "after to_backend." + ) + return ExportedProgramPassResult(exported_program, False) diff --git a/backends/fused_quant/frontend.py b/backends/fused_quant/frontend.py new file mode 100644 index 00000000000..a064b5a7607 --- /dev/null +++ b/backends/fused_quant/frontend.py @@ -0,0 +1,92 @@ +# 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. + +import contextlib + +import torch +from executorch.backends.fused_quant.pre_quantize_passes.fold_batch_norm import ( + FoldBatchNorm, +) +from executorch.backends.fused_quant.pre_quantize_passes.fuse_add_softmax_into_masked_softmax import ( + FuseAddSoftmaxIntoMaskedSoftmax, +) +from executorch.backends.fused_quant.pre_quantize_passes.replace_batch_norm_with_conv import ( + ReplaceBatchNormWithConv, +) +from executorch.backends.fused_quant.quantizer.quantizer import FusedQuantQuantizer +from torch._inductor.decomposition import remove_decompositions +from torch.export import ExportedProgram +from torch.fx import GraphModule +from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e + + +def trace( + model: torch.nn.Module, + inputs: tuple[object, ...], + quantizer: FusedQuantQuantizer, + *, + is_qat: bool = False, +) -> ExportedProgram: + """Export ``model`` to core ATen and apply the pre-quantization rewrites. + + The quantizer matches whole ATen ops, so every op in + ``quantizer.preserved_ops()`` is held back from ``run_decompositions``. + ``aten._safe_softmax`` is always kept: it is what SDPA decomposes to, and + :class:`FuseAddSoftmaxIntoMaskedSoftmax` needs it whole to recompose. + + The rewrites must run after decomposition (they match on ATen ops) but + before annotation (the quantizer must see the rewritten ops). Which ones run + is derived from the quantizer: a pass is only useful if the quantizer can + consume what it produces. + + With ``is_qat`` the model is exported in training mode and batch norms are + kept, so the result can go to ``prepare_qat_pt2e``; otherwise it is exported + in eval mode with batch norms folded. Either way the caller's model is left + in the mode it was passed in. + """ + ops_to_keep = quantizer.preserved_ops() + + decomp_table = torch.export.default_decompositions() + remove_decompositions( + decomp_table, + [*ops_to_keep, torch.ops.aten._safe_softmax.default], + ) + was_training = model.training + model.train(is_qat) + try: + with contextlib.nullcontext() if is_qat else torch.no_grad(): + program = torch.export.export( + model, inputs, strict=True + ).run_decompositions(decomp_table) + finally: + model.train(was_training) + + if is_qat: + # Batch norms train with the model; folding them is a post-training step. + if torch.ops.aten._masked_softmax.default in ops_to_keep: + program = FuseAddSoftmaxIntoMaskedSoftmax()(program).exported_program + return program + + program = FoldBatchNorm()(program).exported_program + + # Any batch norm not folded into an upstream conv/linear becomes a + # convolution instead, so it can be quantized -- but only if the quantizer + # actually handles convolution. + if torch.ops.aten.convolution.default in ops_to_keep: + program = ReplaceBatchNormWithConv()(program).exported_program + + # Recompose add(mask) + softmax into a single _masked_softmax before + # annotation, so the very-negative mask sentinel never enters the quantized + # domain. Only worth doing if the quantizer matches _masked_softmax. + if torch.ops.aten._masked_softmax.default in ops_to_keep: + program = FuseAddSoftmaxIntoMaskedSoftmax()(program).exported_program + + return program + + +def prepare(program: ExportedProgram, quantizer: FusedQuantQuantizer) -> GraphModule: + """Annotate ``program`` with ``quantizer`` and insert observers.""" + return prepare_pt2e(program.module(), quantizer) diff --git a/backends/fused_quant/fuse_aten.py b/backends/fused_quant/fuse_aten.py new file mode 100644 index 00000000000..430d574308d --- /dev/null +++ b/backends/fused_quant/fuse_aten.py @@ -0,0 +1,402 @@ +# 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. + + +import operator +from operator import attrgetter +from typing import Any, cast, Optional + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.graph_utils import ( + is_affine_quant_node, + is_dequantize_node, + is_per_channel_quant_node, + is_per_tensor_quant_node, + is_quantize_node, +) +from executorch.backends.transforms.permute_pass_utils import get_arg +from torch import fx +from torch._export.utils import _detect_fake_mode_from_gm +from torch._ops import OpOverload + +_NULL_QPARAMS: list[torch.fx.node.Argument] = [None, None, torch.uint8, 0, 0] + +QParamsFlat = tuple[fx.Node, fx.Node, torch.dtype, int, int] + + +def _resolve_const_tensor( + node: fx.Node, owning_module: torch.nn.Module +) -> torch.Tensor: + """Resolve a frozen get_attr constant to its backing tensor.""" + assert node.op == "get_attr", ( + f"Expected a get_attr constant, got op={node.op} for {node.name}" + ) + return attrgetter(str(node.target))(owning_module) + + +def _add_scalar_constant( + graph: fx.Graph, value: float | int, dtype: torch.dtype, name: str +) -> fx.Node: + """Register a 0-dim scalar qparam as a frozen module attribute, return a get_attr. + + ``name`` is a preferred base; ``graph.get_attr`` uniquifies it against the + graph, and the attribute is registered under the resulting unique name. + """ + gm = graph.owning_module + assert gm is not None + # get_attr uniquifies node.name but not node.target, so sync target and + # register the attribute under the generated name. + const_node = graph.get_attr(name) + const_node.target = const_node.name + setattr(gm, const_node.name, torch.tensor(value, dtype=dtype)) + return const_node + + +def _create_qparams_from_node( + node: fx.Node, graph: fx.Graph, out_dtype: torch.dtype +) -> QParamsFlat: + """Extract QuantParams from a quantize or dequantize node. + + For per-tensor quantization, creates scalar tensors for scale and zero_point. + For per-channel quantization, uses the existing tensor nodes and extracts axis. + + Args: + node: A quantize or dequantize node (per-tensor or per-channel). + graph: The graph to insert new nodes into. + out_dtype: The dtype after applying quantization/dequantization + + Returns: + A QuantParamsNode tuple. + """ + quant_min = get_arg(node, "quant_min") + quant_max = get_arg(node, "quant_max") + + if is_affine_quant_node(node): + # Affine nodes already carry a full-rank scale/zero_point whose shape + # encodes the block layout (block_size[i] = tensor.shape[i] // + # scale.shape[i]) -- exactly QuantParamsStruct's convention -- so the + # existing scale/zero_point nodes pass straight through, with no axis and + # no reshape (this covers per-tensor/channel/group/blockwise uniformly). + scale_node = get_arg(node, "scale", fx.Node) + zp_node = get_arg(node, "zero_point", fx.Node) + # pyrefly: ignore [bad-return] + return (scale_node, zp_node, out_dtype, quant_min, quant_max) + + if is_per_tensor_quant_node(node): + scale_val = get_arg(node, "scale", float) + zp_val = get_arg(node, "zero_point", int) + with graph.inserting_before(node): + # Per-tensor scale/zero_point arrive as inline scalars; lift them to + # frozen constants so they match the per-channel qparam representation. + scale_node = _add_scalar_constant(graph, scale_val, torch.float32, "_scale") + zp_node = _add_scalar_constant(graph, zp_val, torch.int64, "_zero_point") + # pyrefly: ignore [bad-return] + return (scale_node, zp_node, out_dtype, quant_min, quant_max) + elif is_per_channel_quant_node(node): + # Per-channel: QuantParamsStruct encodes granularity by the scale's shape + # (no axis), so reshape the 1D [C] scale/zero_point to full rank -- + # [1, .., C, .., 1] with C at `axis` -- giving block_size 1 along the + # channel axis and full along the rest. + scale_node = get_arg(node, "scales", fx.Node) + zp_node = get_arg(node, "zero_points", fx.Node) + axis = get_arg(node, "axis", int) + input_node = get_arg(node, "input", fx.Node) + # The weight and scale/zero_point are frozen get_attr constants (convert + # folds the quantize into a frozen param); resolve them off the owning module. + owning_module = graph.owning_module + assert owning_module is not None + ndim = _resolve_const_tensor(input_node, owning_module).ndim + scale_tensor = _resolve_const_tensor(scale_node, owning_module) + zp_tensor = _resolve_const_tensor(zp_node, owning_module) + view_shape = [1] * ndim + view_shape[axis] = -1 + with graph.inserting_before(node): + fake_mode = _detect_fake_mode_from_gm(owning_module) + assert fake_mode, "Expected fake tensor mode!" + scale_view = graph.call_function( + torch.ops.aten.view.default, args=(scale_node, view_shape) + ) + zp_view = graph.call_function( + torch.ops.aten.view.default, args=(zp_node, view_shape) + ) + with fake_mode: + scale_view.meta["val"] = fake_mode.from_tensor( + scale_tensor.reshape(view_shape) + ) + zp_view.meta["val"] = fake_mode.from_tensor( + zp_tensor.reshape(view_shape) + ) + # pyrefly: ignore [bad-return] + return (scale_view, zp_view, out_dtype, quant_min, quant_max) + + raise ValueError(f"Unsupported quantize node: {node}") + + +def create_qparams_from_dequant_node( + dequant_node: fx.Node, + graph: fx.Graph, +) -> QParamsFlat: + """Extract QuantParams from a dequantize node.""" + assert is_dequantize_node(dequant_node) + # Affine dequant names its result dtype `output_dtype`; the quantized_decomposed + # per-tensor/channel ops name it `out_dtype`. + dtype_arg = "output_dtype" if is_affine_quant_node(dequant_node) else "out_dtype" + out_dtype = get_arg(dequant_node, dtype_arg, Optional[torch.dtype]) or torch.float32 + return _create_qparams_from_node(dequant_node, graph, out_dtype) + + +def create_qparams_from_quant_node( + quant_node: fx.Node, + graph: fx.Graph, +) -> QParamsFlat: + """Extract QuantParams from a quantize node.""" + assert is_quantize_node(quant_node) + # Affine quant names the quantized dtype `output_dtype`; quantized_decomposed + # per-tensor/channel ops name it `dtype`. + dtype_arg = "output_dtype" if is_affine_quant_node(quant_node) else "dtype" + out_dtype = get_arg(quant_node, dtype_arg, torch.dtype) + return _create_qparams_from_node(quant_node, graph, out_dtype) + + +def arg_names(target: OpOverload) -> list[str]: + """The op's schema argument names, in schema order. + + The first arg is exposed as self on some schemas but as input to get_arg + (torch.fx normalizes self -> input since self can't be a keyword-argument + name); this normalizes to the get_arg spelling so edge names resolve directly. + """ + return ["input" if a.name == "self" else a.name for a in target._schema.arguments] + + +def output_node(node: fx.Node, output_index: int = 0) -> Optional[fx.Node]: + """The node carrying output output_index of node. + + A single-output op *is* its output; a multi-output op (native_layer_norm, + max_pool2d_with_indices) unpacks each output via getitem(node, output_index). + Returns None when that output node is absent. + """ + if not isinstance(node.meta.get("val"), (tuple, list)): + return node + for user in node.users: + if user.target is operator.getitem and user.args[1] == output_index: + return user + return None + + +def _get_extra_args( + node: fx.Node, + tensor_input_names: set[str], + fused_target: OpOverload, +) -> list[torch.fx.node.Argument]: + """Args of node that are neither tensor inputs nor qparam-bearing. + + Everything whose (normalized) schema name is not a tensor input + (stride/padding/groups, alpha, a scalar operand, normalized_shape/weight/bias, + dim/mask_type, ...) is threaded into the fused op by matching schema arg name. + A name the fused op does not accept is an unsupported lowering and raises. + """ + fused_arg_names = set(arg_names(fused_target)) + + extra_args: list[torch.fx.node.Argument] = [] + for name in arg_names(cast(OpOverload, node.target)): + if name in tensor_input_names: + continue + if name not in fused_arg_names: + raise ValueError( + f"Unsupported ATen arg: {name} from {node.target} " + f"when fusing to {fused_target}" + ) + extra_args.append(get_arg(node, name)) + return extra_args + + +def _build_fused_args( + node: fx.Node, + dequant_inputs: dict[str, fx.Node], + quant_outputs: dict[int, fx.Node], + fused_target: OpOverload, + ordered_input_names: list[str], + output_indices: list[int], +) -> list[torch.fx.node.Argument]: + """Build the flat fused-op arg list: + (tensor_inputs..., per-input qparams..., per-output qparams..., extra_args...) + + Tensor inputs are emitted in schema order, one qparams block each (from the + feeding dequant, or null when the edge is unquantized), then one output qparams + block per entry in output_indices (ascending; from the consuming quantize, + or null), then the remaining args by name. + + Newly created qparams (scale/zero_point) nodes may not be topologically + ordered relative to the fused op; the caller relies on a later legalize_graph. + """ + graph = node.graph + tensor_inputs: list[torch.fx.node.Argument] = [] + input_qparams: list[Optional[QParamsFlat]] = [] + for name in ordered_input_names: + dequant = dequant_inputs.get(name) + if dequant is not None: + tensor_inputs.append(get_arg(dequant, "input")) + input_qparams.append(create_qparams_from_dequant_node(dequant, graph)) + else: + tensor_inputs.append(get_arg(node, name)) + input_qparams.append(None) + + out_qparams_blocks: list[Optional[QParamsFlat]] = [] + for output_index in output_indices: + quant_node = quant_outputs.get(output_index) + out_qparams_blocks.append( + create_qparams_from_quant_node(quant_node, graph) + if quant_node is not None + else None + ) + + extra_args = _get_extra_args(node, set(ordered_input_names), fused_target) + + fused_args_list: list[torch.fx.node.Argument] = [] + fused_args_list.extend(tensor_inputs) + for qp in [*input_qparams, *out_qparams_blocks]: + fused_args_list.extend(qp if qp is not None else _NULL_QPARAMS) + fused_args_list.extend(extra_args) + return fused_args_list + + +def _rewire_multi_output( + graph: fx.Graph, + node: fx.Node, + fused_op: fx.Node, + quant_outputs: dict[int, fx.Node], +) -> None: + """Rewire a multi-output op to its tuple-returning fused replacement. + + The fused op returns a same-arity tuple, so the existing getitem accessors are + repointed from node to fused_op. For each quantized output the getitem + now already yields the quantized dtype, so its downstream quantize is redundant + and is collapsed away. + """ + node.replace_all_uses_with(fused_op) + for quant_node in quant_outputs.values(): + getitem = quant_node.all_input_nodes[0] + getitem.meta["val"] = quant_node.meta["val"] + quant_node.replace_all_uses_with(getitem) + graph.erase_node(quant_node) + graph.erase_node(node) + + +def _rewire_single_output( + graph: fx.Graph, + node: fx.Node, + fused_op: fx.Node, + quant_outputs: dict[int, fx.Node], +) -> None: + """Rewire a single-output op: replace its quantize user (or the op) with fused_op.""" + quant_out = quant_outputs.get(0) + if quant_out is not None: + quant_out.replace_all_uses_with(fused_op) + graph.erase_node(quant_out) + else: + node.replace_all_uses_with(fused_op) + graph.erase_node(node) + + +def _fused_meta( + node: fx.Node, + quant_outputs: dict[int, fx.Node], +) -> dict[str, Any]: + meta = node.meta.copy() + output_val = meta["val"] + if isinstance(output_val, (tuple, list)): + output_values = list(output_val) + for output_index, quant_output in quant_outputs.items(): + output_values[output_index] = quant_output.meta["val"] + meta["val"] = ( + tuple(output_values) if isinstance(output_val, tuple) else output_values + ) + elif quant_outputs: + assert len(quant_outputs) == 1 + quant_output = next(iter(quant_outputs.values())) + meta["val"] = quant_output.meta["val"] + return meta + + +def fuse_aten( + node: fx.Node, + fused_target: OpOverload, + activation_names: tuple[str, ...], + weight_names: tuple[str, ...] = (), + other_names: tuple[str, ...] = (), + output_indices: tuple[int, ...] = (0,), +) -> fx.Node: + """Generic fusion for an ATen op whose fused replacement follows the + fused_quant arg convention: (tensor_inputs..., per-input qparams..., + per-output qparams..., extra_args...). + + Detects the dequant -> op -> quant pattern around node and replaces it with + a single fused_target node, which is returned. + """ + tensor_input_names = {*activation_names, *weight_names, *other_names} + # Tensor inputs in schema order (so the fused op's positional prefix lines up); + # this also drops any name the op's schema does not have. + ordered_input_names = [ + n for n in arg_names(cast(OpOverload, node.target)) if n in tensor_input_names + ] + output_indices = tuple(sorted(set(output_indices))) + + # Detect the dequant -> op -> quant structure. Every tensor input is checked for a + # feeding dequant. + dequant_inputs: dict[str, fx.Node] = {} + for name in ordered_input_names: + arg = get_arg(node, name) + if isinstance(arg, fx.Node) and is_dequantize_node(arg) and len(arg.users) == 1: + dequant_inputs[name] = arg + + quant_outputs: dict[int, fx.Node] = {} + for output_index in output_indices: + accessor = output_node(node, output_index) + if accessor is None or len(accessor.users) != 1: + continue + candidate = next(iter(accessor.users.keys())) + if is_quantize_node(candidate) and len(candidate.all_input_nodes) == 1: + quant_outputs[output_index] = candidate + + if not dequant_inputs and not quant_outputs: + raise ValueError( + f"fuse_aten: {node.target} node {node.name!r} has no quantized input " + "or output to fuse (expected a surrounding dequantize/quantize)." + ) + + graph = node.graph + fused_meta = _fused_meta(node, quant_outputs) + multi_output = isinstance(node.meta["val"], (tuple, list)) + + fused_args_list = _build_fused_args( + node, + dequant_inputs, + quant_outputs, + fused_target, + ordered_input_names, + list(output_indices), + ) + + # Insert the fused op before node. Newly created qparams nodes may land after + # it (created in front of downstream quantize nodes), leaving the graph + # topologically invalid here; the caller relegalizes once after all fusions. + with graph.inserting_before(node): + fused_op = graph.call_function(fused_target, args=tuple(fused_args_list)) + fused_op.meta = fused_meta + + if multi_output: + _rewire_multi_output(graph, node, fused_op, quant_outputs) + else: + _rewire_single_output(graph, node, fused_op, quant_outputs) + + # set(): the same dequant can feed two inputs (e.g. x * x), so dedup before + # erasing to avoid a redundant erase of an already-removed node. + for dequant_node in set(dequant_inputs.values()): + if not dequant_node.users: + graph.erase_node(dequant_node) + + return fused_op diff --git a/backends/fused_quant/graph_utils.py b/backends/fused_quant/graph_utils.py new file mode 100644 index 00000000000..d84d4b59de0 --- /dev/null +++ b/backends/fused_quant/graph_utils.py @@ -0,0 +1,151 @@ +# 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. + + +from typing import Optional + +import torch +from executorch.backends.fused_quant.ops import QuantParamsStruct +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import get_constant +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch import fx +from torch._ops import OpOverload +from torch.export import ExportedProgram + + +_OpOverload = EdgeOpOverload | OpOverload +_QPARAM_FIELDS = ("scale", "zero_point", "dtype", "quant_min", "quant_max") + + +def split_fused_arg_names(target: _OpOverload) -> tuple[list[str], list[str]]: + """Split a fused op's schema into quantized inputs and extra arguments.""" + names = [argument.name for argument in target._schema.arguments] + try: + first_qparam = next( + index for index, name in enumerate(names) if name.endswith("_scale") + ) + except StopIteration as error: + raise ValueError(f"Fused op {target} has no qparam blocks") from error + + input_names = names[:first_qparam] + cursor = first_qparam + for prefix in (*input_names, "out"): + expected = [f"{prefix}_{field}" for field in _QPARAM_FIELDS] + actual = names[cursor : cursor + len(_QPARAM_FIELDS)] + if actual != expected: + raise ValueError( + f"Fused op {target} does not follow the qparam argument convention: " + f"expected {expected}, got {actual}" + ) + cursor += len(_QPARAM_FIELDS) + return input_names, names[cursor:] + + +def is_quantize_node(node: fx.Node) -> bool: + return node.op == "call_function" and node.target in ( + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_channel.default, + torch.ops.torchao.quantize_affine.default, + ) + + +def is_dequantize_node(node: fx.Node) -> bool: + return node.op == "call_function" and node.target in ( + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_channel.default, + torch.ops.torchao.dequantize_affine.default, + ) + + +def is_affine_quant_node(node: fx.Node) -> bool: + """Check if a node is a torchao affine quantize/dequantize op. + + Affine nodes carry an explicit ``block_size`` and a full-rank scale/zero_point + whose shape already encodes the block layout -- exactly QuantParamsStruct's + convention, covering per-tensor/channel/group/blockwise uniformly with no axis. + """ + return node.target in ( + torch.ops.torchao.quantize_affine.default, + torch.ops.torchao.dequantize_affine.default, + ) + + +def is_per_tensor_quant_node(node: fx.Node) -> bool: + """Check if a quantize/dequantize node is per-tensor.""" + return node.target in ( + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + ) + + +def is_per_channel_quant_node(node: fx.Node) -> bool: + """Check if a quantize/dequantize node is per-channel.""" + return node.target in ( + torch.ops.quantized_decomposed.quantize_per_channel.default, + torch.ops.quantized_decomposed.dequantize_per_channel.default, + ) + + +def get_qparams_from_node( + node: fx.Node, prefix: str +) -> Optional[QuantParamsStruct[fx.Node]]: + """Extract QuantParamsStruct from a node's flat args by prefix. + + Returns None if both scale and zero_point are None (unquantized). + Raises ValueError if only one of scale/zero_point is None. + """ + scale = get_arg(node, f"{prefix}_scale", Optional[fx.Node]) + zero_point = get_arg(node, f"{prefix}_zero_point", Optional[fx.Node]) + if scale is None and zero_point is None: + return None + if (scale is None) != (zero_point is None): + raise ValueError( + f"{prefix}_scale and {prefix}_zero_point must both be None or both be " + f"provided on node {node.name}" + ) + assert scale is not None and zero_point is not None + return QuantParamsStruct( + scale=scale, + zero_point=zero_point, + dtype=get_arg(node, f"{prefix}_dtype", torch.dtype), + quant_min=get_arg(node, f"{prefix}_quant_min", int), + quant_max=get_arg(node, f"{prefix}_quant_max", int), + ) + + +def get_qparams_flat(node: fx.Node, prefix: str) -> list[torch.fx.node.Argument]: + """Get the 5 flat qparams values from a node for pass-through to another op.""" + return [ + get_arg(node, f"{prefix}_scale"), + get_arg(node, f"{prefix}_zero_point"), + get_arg(node, f"{prefix}_dtype"), + get_arg(node, f"{prefix}_quant_min"), + get_arg(node, f"{prefix}_quant_max"), + ] + + +def get_scale(ep: ExportedProgram, qp: "QuantParamsStruct[fx.Node]") -> float: + """Read a per-tensor scale value from its qparams. + + Per-tensor qparams are 0-dim lifted constants (placeholders) post-export; + resolve the backing tensor and return its Python scalar. + """ + value = get_constant(ep, qp.scale) + assert value is not None, ( + f"expected a lifted-constant qparam for {qp.scale.name}, got op={qp.scale.op}" + ) + return float(value.item()) + + +def get_zero_point(ep: ExportedProgram, qp: "QuantParamsStruct[fx.Node]") -> int: + """Read a per-tensor zero_point value from its qparams. See get_scale.""" + value = get_constant(ep, qp.zero_point) + assert value is not None, ( + f"expected a lifted-constant qparam for {qp.zero_point.name}, " + f"got op={qp.zero_point.op}" + ) + return int(value.item()) diff --git a/backends/fused_quant/ops.py b/backends/fused_quant/ops.py new file mode 100644 index 00000000000..939676e500a --- /dev/null +++ b/backends/fused_quant/ops.py @@ -0,0 +1,2659 @@ +# 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. + + +import dataclasses +from collections.abc import Sequence +from typing import Callable, Generic, Optional, TypeVar, Union + +import torch + +# Registers torch.ops.torchao.(de)quantize_affine used by QuantParamsStruct. +import torchao.quantization.quant_primitives # noqa: F401 +from executorch.backends.fused_quant.ops_utils import ( + compute_conv_out_shape, + compute_conv_out_shape_nhwc, +) +from executorch.exir.pass_base import ProxyValue +from torch import fx + +# Importing this registers the torch.ops.quantized_decomposed.* ops +# (quantize/dequantize per_tensor/per_channel) used by QuantParamsStruct. +from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib # noqa: F401 +from torch.library import Library, register_fake + +_lib = Library("fused_quant", "DEF") + +_TensorType = TypeVar( + "_TensorType", bound=Union[torch.Tensor, "fx.Node", "fx.Proxy", ProxyValue] +) + + +def _get_tensor_val( + t: Union[torch.Tensor, fx.Node, fx.Proxy, ProxyValue], +) -> torch.Tensor: + """Extract the tensor value. For fx.Node, gets the value from meta['val'].""" + if isinstance(t, fx.Node): + return t.meta["val"] + if isinstance(t, torch.Tensor): + return t + if isinstance(t, ProxyValue): + return t.data + if isinstance(t, fx.Proxy): + return t.node.meta["val"] + raise TypeError(f"Unsupported type: {type(t)}") + + +@dataclasses.dataclass(frozen=True) +class QuantParamsStruct(Generic[_TensorType]): + """Quantization parameters as a dataclass for convenient access. + + Type Parameters: + _TensorType: The type of scale and zero_point, either torch.Tensor or fx.Node. + + Granularity is encoded entirely by the scale's shape relative to the tensor + it quantizes -- there is no ``axis``. A scale is either: + - a singleton (``numel() == 1``): per-tensor; broadcast over the whole + tensor (block_size == tensor.shape); or + - a full-rank tensor (``ndim == tensor.ndim``) whose dims encode the tiling: + ``block_size[i] = tensor.shape[i] // scale.shape[i]``. This covers + per-channel (exactly one non-unary dim) and per-group / blockwise (more). + + Attributes: + scale: Scale factor (singleton or full-rank; see above). + zero_point: Zero point; must have the same shape as scale. + dtype: The quantized dtype for outputs / the dequantized output dtype for + inputs. + quant_min: Minimum quantized value. + quant_max: Maximum quantized value. + """ + + scale: _TensorType + zero_point: _TensorType + dtype: torch.dtype + quant_min: int + quant_max: int + + @classmethod + def maybe_from_flat_args( + cls, + scale: Optional[_TensorType], + zero_point: Optional[_TensorType], + dtype: torch.dtype, + quant_min: int, + quant_max: int, + ) -> Optional["QuantParamsStruct[_TensorType]"]: + """Construct from flat args, or return None if scale and zero_point are both None.""" + if scale is None and zero_point is None: + return None + if (scale is None) != (zero_point is None): + raise ValueError( + "scale and zero_point must both be None or both be provided, " + f"got scale={'None' if scale is None else 'Tensor'} and " + f"zero_point={'None' if zero_point is None else 'Tensor'}" + ) + assert scale is not None and zero_point is not None + return cls(scale, zero_point, dtype, quant_min, quant_max) + + @staticmethod + def _non_unary_dims(scale: torch.Tensor) -> list[int]: + """Indices of the scale dims whose size != 1 (the quantized/tiled dims).""" + return [i for i, dim in enumerate(scale.shape) if dim != 1] + + def is_per_tensor(self) -> bool: + """One scale for the whole tensor. + + The scale may be shaped ``()`` or all-ones; only its ``numel()`` matters. + A per-channel scale over a single channel is per-tensor by this test, + which is numerically the same thing. + """ + return _get_tensor_val(self.scale).numel() == 1 + + def _single_axis_block_size(self, tensor_shape: Sequence[int]) -> Optional[int]: + scale = _get_tensor_val(self.scale) + non_unary = self._non_unary_dims(scale) + if len(non_unary) != 1 or scale.ndim != len(tensor_shape): + return None + axis = non_unary[0] + return tensor_shape[axis] // scale.shape[axis] + + def is_per_channel(self, tensor_shape: Optional[Sequence[int]] = None) -> bool: + """One scale per slice along a single axis. + + From the scale shape alone, a scale with one non-unary dim is + per-channel. Pass the quantized tensor's shape to also require a block + size of 1 along that axis, which rules out a group scale such as + ``[1, K // g]`` over ``[1, K]``. + """ + if len(self._non_unary_dims(_get_tensor_val(self.scale))) != 1: + return False + return tensor_shape is None or self._single_axis_block_size(tensor_shape) == 1 + + def is_per_group(self, tensor_shape: Optional[Sequence[int]] = None) -> bool: + """Blockwise: several tiled axes, or (given ``tensor_shape``) one axis + tiled in blocks larger than 1.""" + non_unary = self._non_unary_dims(_get_tensor_val(self.scale)) + if len(non_unary) >= 2: + return True + if tensor_shape is None or len(non_unary) != 1: + return False + block_size = self._single_axis_block_size(tensor_shape) + return block_size is not None and block_size > 1 + + def channel_axis(self) -> Optional[int]: + """The per-channel axis, derived from the scale shape: + - ``0`` if the scale is a singleton (per-tensor; axis is irrelevant), + - the single non-unary dim for per-channel, + - ``None`` if more than one dim is non-unary (blockwise / per-group, + which has no single channel axis). + """ + non_unary = self._non_unary_dims(_get_tensor_val(self.scale)) + if len(non_unary) == 0: + return 0 + if len(non_unary) == 1: + return non_unary[0] + return None + + def validate(self) -> None: + """Validate that scale and zero_point have the same shape.""" + scale = _get_tensor_val(self.scale) + zero_point = _get_tensor_val(self.zero_point) + if scale.shape != zero_point.shape: + raise ValueError( + f"scale and zero_point must have the same shape, got {scale.shape} and {zero_point.shape}" + ) + + def _broadcast_scale_zp(self, ndim: int) -> tuple[torch.Tensor, torch.Tensor]: + """Return (scale, zero_point) at full rank ``ndim`` for affine quant. + + A singleton (per-tensor) is reshaped to all-ones so it broadcasts over the + whole tensor; a full-rank scale is returned as-is (its shape encodes the + block layout). Any other rank is rejected -- non-singleton scales must be + full-rank. + """ + scale = _get_tensor_val(self.scale) + zero_point = _get_tensor_val(self.zero_point) + if scale.numel() == 1: + return scale.reshape([1] * ndim), zero_point.reshape([1] * ndim) + if scale.ndim == ndim: + return scale, zero_point + raise ValueError( + f"scale rank {scale.ndim} must be a singleton or match the tensor rank " + f"{ndim}; per-channel/group scales must be full-rank so their shape " + "encodes the block layout" + ) + + @staticmethod + def _block_size(tensor: torch.Tensor, scale: torch.Tensor) -> list[int]: + """Derive the affine block_size: each tile is tensor.shape[i] // scale.shape[i].""" + block_size: list[int] = [] + for tdim, sdim in zip(tensor.shape, scale.shape): + if tdim % sdim != 0: + raise ValueError( + f"tensor dim {tdim} must be divisible by scale dim {sdim}" + ) + block_size.append(tdim // sdim) + return block_size + + def quantize(self, tensor: torch.Tensor) -> torch.Tensor: + """Quantize ``tensor`` via affine (block-wise) quantization. The block + layout is derived from the scale shape: block_size[i] = tensor.shape[i] // + scale.shape[i] (covering per-tensor/channel/group/element uniformly).""" + assert isinstance(self.scale, torch.Tensor) + assert isinstance(self.zero_point, torch.Tensor) + scale, zero_point = self._broadcast_scale_zp(tensor.ndim) + block_size = self._block_size(tensor, scale) + return torch.ops.torchao.quantize_affine( + tensor, + block_size, + scale, + zero_point, + self.dtype, + self.quant_min, + self.quant_max, + ) + + def dequantize(self, tensor: torch.Tensor) -> torch.Tensor: + """Dequantize ``tensor`` via affine (block-wise) dequantization.""" + assert isinstance(self.scale, torch.Tensor) + assert isinstance(self.zero_point, torch.Tensor) + scale, zero_point = self._broadcast_scale_zp(tensor.ndim) + block_size = self._block_size(tensor, scale) + return torch.ops.torchao.dequantize_affine( + tensor, + block_size, + scale, + zero_point, + tensor.dtype, + self.quant_min, + self.quant_max, + output_dtype=self.dtype, + ) + + +AnyQuantParamsStruct = Union[ + QuantParamsStruct[torch.Tensor], QuantParamsStruct[fx.Node] +] + + +def _validate_qparams(*qparams_list: Optional[AnyQuantParamsStruct]) -> None: + """Validate every qparams block that is present.""" + for qp in qparams_list: + if qp is not None: + qp.validate() + + +def get_out_dtype( + out_qparams: Optional[AnyQuantParamsStruct], default_dtype: torch.dtype +) -> torch.dtype: + if out_qparams is None: + return default_dtype + return out_qparams.dtype + + +def dequantized_dtype( + tensor: torch.Tensor, qparams: Optional[AnyQuantParamsStruct] +) -> torch.dtype: + """The float dtype ``tensor`` has once its input qparams are applied. + + For an input qparams block, ``dtype`` is the dequantized float dtype; an + unquantized operand is already float. + """ + return qparams.dtype if qparams is not None else tensor.dtype + + +def _maybe_quantize( + tensor: torch.Tensor, qparams: Optional[AnyQuantParamsStruct] +) -> torch.Tensor: + if qparams is None: + return tensor + return qparams.quantize(tensor) + + +def _maybe_dequantize( + tensor: torch.Tensor, qparams: Optional[AnyQuantParamsStruct] +) -> torch.Tensor: + if qparams is None: + return tensor + return qparams.dequantize(tensor) + + +_QP = Optional[QuantParamsStruct[torch.Tensor]] + + +def _make_qp( + scale: Optional[torch.Tensor], + zero_point: Optional[torch.Tensor], + dtype: torch.dtype, + quant_min: int, + quant_max: int, +) -> _QP: + return QuantParamsStruct.maybe_from_flat_args( + scale, zero_point, dtype, quant_min, quant_max + ) + + +def _permute_contiguous(tensor: torch.Tensor, dims: list[int]) -> torch.Tensor: + return tensor.permute(*dims).clone(memory_format=torch.contiguous_format) + + +def _permute_qparams( + qparams: _QP, + dims: list[int], +) -> _QP: + """Permute layout-sensitive qparams with their tensor.""" + if qparams is None or qparams.is_per_tensor(): + return qparams + assert isinstance(qparams.scale, torch.Tensor) + assert isinstance(qparams.zero_point, torch.Tensor) + return dataclasses.replace( + qparams, + scale=_permute_contiguous(qparams.scale, dims), + zero_point=_permute_contiguous(qparams.zero_point, dims), + ) + + +def _binary_op_impl( + inp: torch.Tensor, + other: torch.Tensor, + inp_qp: _QP, + other_qp: _QP, + out_qp: _QP, + fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], +) -> torch.Tensor: + _validate_qparams(inp_qp, other_qp, out_qp) + dq_inp = _maybe_dequantize(inp, inp_qp) + dq_other = _maybe_dequantize(other, other_qp) + out = fn(dq_inp, dq_other) + return _maybe_quantize(out, out_qp) + + +def _binary_op_meta( + inp: torch.Tensor, + other: torch.Tensor, + inp_qp: _QP, + other_qp: _QP, + out_qp: _QP, +) -> torch.Tensor: + _validate_qparams(inp_qp, other_qp, out_qp) + out_size = torch.broadcast_shapes(inp.shape, other.shape) + return inp.new_empty( + out_size, + dtype=get_out_dtype( + out_qp, + torch.promote_types( + dequantized_dtype(inp, inp_qp), dequantized_dtype(other, other_qp) + ), + ), + ) + + +def _unary_op_impl( + inp: torch.Tensor, + inp_qp: _QP, + out_qp: _QP, + fn: Callable[[torch.Tensor], torch.Tensor], +) -> torch.Tensor: + _validate_qparams(inp_qp, out_qp) + dq_inp = _maybe_dequantize(inp, inp_qp) + out = fn(dq_inp) + return _maybe_quantize(out, out_qp) + + +def _unary_op_meta( + inp: torch.Tensor, + inp_qp: _QP, + out_qp: _QP, +) -> torch.Tensor: + _validate_qparams(inp_qp, out_qp) + return inp.new_empty( + inp.shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ) + + +def _binary_scalar_op_impl( + inp: torch.Tensor, + inp_qp: _QP, + out_qp: _QP, + other: float, + fn: Callable[[torch.Tensor, float], torch.Tensor], +) -> torch.Tensor: + """Binary op with a scalar second operand. Only the tensor input and the + output carry quant params; the scalar is applied directly to the dequantized + input (it has no qparams of its own).""" + _validate_qparams(inp_qp, out_qp) + dq_inp = _maybe_dequantize(inp, inp_qp) + out = fn(dq_inp, other) + return _maybe_quantize(out, out_qp) + + +def _binary_scalar_op_meta( + inp: torch.Tensor, + inp_qp: _QP, + out_qp: _QP, +) -> torch.Tensor: + # A scalar second operand does not broadcast, so the output keeps inp's shape. + _validate_qparams(inp_qp, out_qp) + return inp.new_empty( + inp.shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ) + + +# ============================================================================= +# Op definitions +# ============================================================================= + + +@torch.library.custom_op("fused_quant::requantize", mutates_args=()) +def _requantize_impl( + inp: torch.Tensor, + inp_scale: torch.Tensor, + inp_zero_point: torch.Tensor, + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: torch.Tensor, + out_zero_point: torch.Tensor, + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + inp_qp = QuantParamsStruct( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = QuantParamsStruct( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + return out_qp.quantize(inp_qp.dequantize(inp)) + + +def _requantize_meta( + inp: torch.Tensor, + inp_scale: torch.Tensor, + inp_zero_point: torch.Tensor, + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: torch.Tensor, + out_zero_point: torch.Tensor, + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + inp_qp = QuantParamsStruct( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = QuantParamsStruct( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + return inp.new_empty(inp.shape, dtype=out_qp.dtype) + + +_requantize_impl.register_fake(_requantize_meta) + + +@torch.library.custom_op("fused_quant::bmm", mutates_args=()) +def _bmm_impl( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _binary_op_impl( + inp, + other, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.bmm, + ) + + +def _bmm_meta( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + other_qp = _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + + _validate_qparams(inp_qp, other_qp, out_qp) + + if inp.ndim != 3 or other.ndim != 3: + raise ValueError( + f"Input tensors must be 3D for bmm, got {inp.ndim}D and {other.ndim}D" + ) + if inp.shape[0] != other.shape[0]: + raise ValueError( + f"Input tensors must have the same batch dimension, got {inp.shape[0]} and {other.shape[0]}" + ) + if inp.shape[2] != other.shape[1]: + raise ValueError( + f"Input tensors must have the same inner dimension, got {inp.shape[2]} and {other.shape[1]}" + ) + + out_shape = (inp.shape[0], inp.shape[1], other.shape[2]) + return inp.new_empty( + out_shape, + dtype=get_out_dtype( + out_qp, + torch.promote_types( + dequantized_dtype(inp, inp_qp), dequantized_dtype(other, other_qp) + ), + ), + ) + + +_bmm_impl.register_fake(_bmm_meta) + + +@torch.library.custom_op("fused_quant::relu", mutates_args=()) +def _relu_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.nn.functional.relu, + ) + + +def _relu_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_relu_impl.register_fake(_relu_meta) + + +@torch.library.custom_op("fused_quant::hardswish", mutates_args=()) +def _hardswish_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.nn.functional.hardswish, + ) + + +def _hardswish_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_hardswish_impl.register_fake(_hardswish_meta) + + +@torch.library.custom_op("fused_quant::sigmoid", mutates_args=()) +def _sigmoid_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.sigmoid, + ) + + +def _sigmoid_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_sigmoid_impl.register_fake(_sigmoid_meta) + + +@torch.library.custom_op("fused_quant::tanh", mutates_args=()) +def _tanh_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.tanh, + ) + + +def _tanh_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_tanh_impl.register_fake(_tanh_meta) + + +@torch.library.custom_op("fused_quant::hard_tanh", mutates_args=()) +def _hard_tanh_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + min_val: float = -1.0, + max_val: float = 1.0, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=lambda x: torch.nn.functional.hardtanh(x, min_val, max_val), + ) + + +def _hard_tanh_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + min_val: float = -1.0, + max_val: float = 1.0, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_hard_tanh_impl.register_fake(_hard_tanh_meta) + + +@torch.library.custom_op("fused_quant::silu", mutates_args=()) +def _silu_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.nn.functional.silu, + ) + + +def _silu_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_silu_impl.register_fake(_silu_meta) + + +@torch.library.custom_op("fused_quant::hardsigmoid", mutates_args=()) +def _hardsigmoid_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.nn.functional.hardsigmoid, + ) + + +def _hardsigmoid_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_hardsigmoid_impl.register_fake(_hardsigmoid_meta) + + +@torch.library.custom_op("fused_quant::gelu", mutates_args=()) +def _gelu_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + approximate: str, +) -> torch.Tensor: + return _unary_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=lambda x: torch.nn.functional.gelu(x, approximate=approximate), + ) + + +def _gelu_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + approximate: str, +) -> torch.Tensor: + return _unary_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_gelu_impl.register_fake(_gelu_meta) + + +@torch.library.custom_op("fused_quant::add", mutates_args=()) +def _add_impl( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_op_impl( + inp, + other, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=lambda x, y: torch.add(x, y, alpha=alpha), + ) + + +def _add_meta( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_op_meta( + inp, + other, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_add_impl.register_fake(_add_meta) + + +@torch.library.custom_op("fused_quant::mul", mutates_args=()) +def _mul_impl( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _binary_op_impl( + inp, + other, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=torch.mul, + ) + + +def _mul_meta( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + return _binary_op_meta( + inp, + other, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_mul_impl.register_fake(_mul_meta) + + +@torch.library.custom_op("fused_quant::sub", mutates_args=()) +def _sub_impl( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_op_impl( + inp, + other, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + fn=lambda x, y: torch.sub(x, y, alpha=alpha), + ) + + +def _sub_meta( + inp: torch.Tensor, + other: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + other_scale: Optional[torch.Tensor], + other_zero_point: Optional[torch.Tensor], + other_dtype: torch.dtype, + other_quant_min: int, + other_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_op_meta( + inp, + other, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp( + other_scale, + other_zero_point, + other_dtype, + other_quant_min, + other_quant_max, + ), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_sub_impl.register_fake(_sub_meta) + + +# ============================================================================= +# Scalar binary ops +# +# Variants of add/mul/sub whose second operand is a Python scalar rather than a +# tensor (e.g. RMSNorm's `x + eps`). The scalar has no quant params, so the +# schema carries qparams only for the tensor input and the output. The fusion +# pass lays args out as (tensor_input, input_qparams..., output_qparams..., +# extra_args...), so the scalar (and `alpha` for add/sub) lands at the very end. +# +# These are registered as `.Scalar` overloads of the existing add/mul/sub ops +# (mirroring aten's `add.Tensor`/`add.Scalar` split) rather than the +# `@custom_op` decorator, which only ever defines a `.default` overload. The +# functional op therefore goes through the lower-level Library API: an explicit +# schema, a CompositeExplicitAutograd impl, and a separately registered fake. +# ============================================================================= + + +def _add_scalar_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + other: float, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_scalar_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + other, + fn=lambda x, y: torch.add(x, y, alpha=alpha), + ) + + +def _add_scalar_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + other: float, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_scalar_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_lib.define( + "add.Scalar(Tensor inp, " + "Tensor? inp_scale, Tensor? inp_zero_point, ScalarType inp_dtype, " + "SymInt inp_quant_min, SymInt inp_quant_max, " + "Tensor? out_scale, Tensor? out_zero_point, ScalarType out_dtype, " + "SymInt out_quant_min, SymInt out_quant_max, " + "Scalar other, float alpha = 1.0) -> Tensor" +) +_lib.impl("add.Scalar", _add_scalar_impl, "CompositeExplicitAutograd") +register_fake("fused_quant::add.Scalar", _add_scalar_meta, lib=_lib) + + +def _mul_scalar_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + other: float, +) -> torch.Tensor: + return _binary_scalar_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + other, + fn=torch.mul, + ) + + +def _mul_scalar_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + other: float, +) -> torch.Tensor: + return _binary_scalar_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_lib.define( + "mul.Scalar(Tensor inp, " + "Tensor? inp_scale, Tensor? inp_zero_point, ScalarType inp_dtype, " + "SymInt inp_quant_min, SymInt inp_quant_max, " + "Tensor? out_scale, Tensor? out_zero_point, ScalarType out_dtype, " + "SymInt out_quant_min, SymInt out_quant_max, " + "Scalar other) -> Tensor" +) +_lib.impl("mul.Scalar", _mul_scalar_impl, "CompositeExplicitAutograd") +register_fake("fused_quant::mul.Scalar", _mul_scalar_meta, lib=_lib) + + +def _sub_scalar_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + other: float, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_scalar_op_impl( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + other, + fn=lambda x, y: torch.sub(x, y, alpha=alpha), + ) + + +def _sub_scalar_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + other: float, + alpha: float = 1.0, +) -> torch.Tensor: + return _binary_scalar_op_meta( + inp, + _make_qp(inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max), + _make_qp(out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max), + ) + + +_lib.define( + "sub.Scalar(Tensor inp, " + "Tensor? inp_scale, Tensor? inp_zero_point, ScalarType inp_dtype, " + "SymInt inp_quant_min, SymInt inp_quant_max, " + "Tensor? out_scale, Tensor? out_zero_point, ScalarType out_dtype, " + "SymInt out_quant_min, SymInt out_quant_max, " + "Scalar other, float alpha = 1.0) -> Tensor" +) +_lib.impl("sub.Scalar", _sub_scalar_impl, "CompositeExplicitAutograd") +register_fake("fused_quant::sub.Scalar", _sub_scalar_meta, lib=_lib) + + +@torch.library.custom_op("fused_quant::linear", mutates_args=()) +def _linear_impl( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + bias_qp = _make_qp( + bias_scale, + bias_zero_point, + bias_dtype, + bias_quant_min, + bias_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + + _validate_qparams(inp_qp, weight_qp, bias_qp, out_qp) + + dq_inp = _maybe_dequantize(inp, inp_qp) + dq_weight = _maybe_dequantize(weight, weight_qp) + dq_bias = _maybe_dequantize(bias, bias_qp) if bias is not None else None + + dq_out = torch.nn.functional.linear(dq_inp, dq_weight, dq_bias) + return _maybe_quantize(dq_out, out_qp) + + +def _linear_meta( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, weight_qp, out_qp) + output_shape = (*inp.shape[:-1], weight.shape[0]) + return inp.new_empty( + output_shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ) + + +_linear_impl.register_fake(_linear_meta) + + +@torch.library.custom_op("fused_quant::convolution", mutates_args=()) +def _convolution_impl( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + stride: list[int], + padding: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + bias_qp = _make_qp( + bias_scale, + bias_zero_point, + bias_dtype, + bias_quant_min, + bias_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + + _validate_qparams(inp_qp, weight_qp, bias_qp, out_qp) + + num_spatial_dims = len(stride) + assert len(padding) == num_spatial_dims + assert len(dilation) == num_spatial_dims + assert len(output_padding) == num_spatial_dims + + dq_inp = _maybe_dequantize(inp, inp_qp) + dq_weight = _maybe_dequantize(weight, weight_qp) + dq_bias = _maybe_dequantize(bias, bias_qp) if bias is not None else None + + if num_spatial_dims not in (1, 2, 3): + raise ValueError( + f"Unsupported number of spatial dimensions: {num_spatial_dims}. " + "Only 1D, 2D, and 3D convolutions are supported." + ) + + dq_out = torch.ops.aten.convolution.default( + dq_inp, + dq_weight, + dq_bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + return _maybe_quantize(dq_out, out_qp) + + +def _convolution_meta( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + stride: list[int], + padding: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, weight_qp, out_qp) + output_shape = compute_conv_out_shape( + inp, + weight, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + return inp.new_empty( + output_shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ) + + +_convolution_impl.register_fake(_convolution_meta) + + +def _validate_conv_nd_args( + name: str, + spatial_dims: int, + inp: torch.Tensor, + weight: torch.Tensor, + stride: list[int], + padding: list[int], + dilation: list[int], +) -> None: + expected_rank = spatial_dims + 2 + if inp.ndim != expected_rank or weight.ndim != expected_rank: + raise ValueError( + f"fused_quant::{name} expects input and weight rank {expected_rank}, " + f"got {inp.ndim} and {weight.ndim}" + ) + for arg_name, values in ( + ("stride", stride), + ("padding", padding), + ("dilation", dilation), + ): + if len(values) != spatial_dims: + raise ValueError( + f"fused_quant::{name} expects {arg_name} to contain " + f"{spatial_dims} values, got {len(values)}" + ) + + +def _register_conv_nd(name: str, spatial_dims: int) -> torch.library.CustomOpDef: + def impl( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + stride: list[int], + padding: list[int], + dilation: list[int], + groups: int, + ) -> torch.Tensor: + _validate_conv_nd_args( + name, spatial_dims, inp, weight, stride, padding, dilation + ) + return _convolution_impl( + inp, + weight, + bias, + inp_scale, + inp_zero_point, + inp_dtype, + inp_quant_min, + inp_quant_max, + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + bias_scale, + bias_zero_point, + bias_dtype, + bias_quant_min, + bias_quant_max, + out_scale, + out_zero_point, + out_dtype, + out_quant_min, + out_quant_max, + stride, + padding, + dilation, + False, + [0] * spatial_dims, + groups, + ) + + def meta( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + stride: list[int], + padding: list[int], + dilation: list[int], + groups: int, + ) -> torch.Tensor: + _validate_conv_nd_args( + name, spatial_dims, inp, weight, stride, padding, dilation + ) + return _convolution_meta( + inp, + weight, + bias, + inp_scale, + inp_zero_point, + inp_dtype, + inp_quant_min, + inp_quant_max, + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + bias_scale, + bias_zero_point, + bias_dtype, + bias_quant_min, + bias_quant_max, + out_scale, + out_zero_point, + out_dtype, + out_quant_min, + out_quant_max, + stride, + padding, + dilation, + False, + [0] * spatial_dims, + groups, + ) + + op = torch.library.custom_op(f"fused_quant::{name}", impl, mutates_args=()) + op.register_fake(meta) + return op + + +_conv1d_impl: torch.library.CustomOpDef = _register_conv_nd("conv1d", 1) +_conv2d_impl: torch.library.CustomOpDef = _register_conv_nd("conv2d", 2) +_conv3d_impl: torch.library.CustomOpDef = _register_conv_nd("conv3d", 3) + + +def _max_pool2d_with_indices( + inp: torch.Tensor, + kernel_size: list[int], + stride: list[int], + padding: list[int], + dilation: list[int], + ceil_mode: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.aten.max_pool2d_with_indices.default( + inp, + kernel_size, + stride or kernel_size, + padding, + dilation, + ceil_mode, + ) + + +@torch.library.custom_op("fused_quant::max_pool2d_with_indices", mutates_args=()) +def _max_pool2d_with_indices_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + dilation: list[int], + ceil_mode: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + values, indices = _max_pool2d_with_indices( + _maybe_dequantize(inp, inp_qp), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ) + return _maybe_quantize(values, out_qp), indices + + +def _max_pool2d_with_indices_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + dilation: list[int], + ceil_mode: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + values, indices = _max_pool2d_with_indices( + torch.ones_like(inp, dtype=torch.float32), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ) + return ( + values.new_empty( + values.shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ), + indices.new_empty(indices.shape), + ) + + +_max_pool2d_with_indices_impl.register_fake(_max_pool2d_with_indices_meta) + + +@torch.library.custom_op( + "fused_quant::max_pool2d_with_indices_channels_last", mutates_args=() +) +def _max_pool2d_with_indices_channels_last_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + dilation: list[int], + ceil_mode: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + inp_nchw = _permute_contiguous(inp, [0, 3, 1, 2]) + inp_qp = _permute_qparams(inp_qp, [0, 3, 1, 2]) + values_nchw, indices_nchw = _max_pool2d_with_indices( + _maybe_dequantize(inp_nchw, inp_qp), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ) + values = _permute_contiguous(values_nchw, [0, 2, 3, 1]) + indices = _permute_contiguous(indices_nchw, [0, 2, 3, 1]) + return _maybe_quantize(values, out_qp), indices + + +def _max_pool2d_with_indices_channels_last_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + dilation: list[int], + ceil_mode: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + values_nchw, indices_nchw = _max_pool2d_with_indices( + inp.permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ) + values = values_nchw.permute(0, 2, 3, 1) + indices = indices_nchw.permute(0, 2, 3, 1) + return ( + values.new_empty( + values.shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ), + indices.new_empty(indices.shape), + ) + + +_max_pool2d_with_indices_channels_last_impl.register_fake( + _max_pool2d_with_indices_channels_last_meta +) + + +def _avg_pool2d( + inp: torch.Tensor, + kernel_size: list[int], + stride: list[int], + padding: list[int], + ceil_mode: bool, + count_include_pad: bool, + divisor_override: Optional[int], +) -> torch.Tensor: + return torch.ops.aten.avg_pool2d.default( + inp, + kernel_size, + stride or kernel_size, + padding, + ceil_mode, + count_include_pad, + divisor_override, + ) + + +@torch.library.custom_op("fused_quant::avg_pool2d", mutates_args=()) +def _avg_pool2d_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + ceil_mode: bool, + count_include_pad: bool, + divisor_override: Optional[int], +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + output = _avg_pool2d( + _maybe_dequantize(inp, inp_qp), + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + divisor_override, + ) + return _maybe_quantize(output, out_qp) + + +def _avg_pool2d_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + ceil_mode: bool, + count_include_pad: bool, + divisor_override: Optional[int], +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + + output = _avg_pool2d( + # Doesn't matter if we are passing in float32 input, we just need + # a valid dtype to pass in so we can get output shape, overriding + # correct dtype later with get_out_dtype + torch.ones_like(inp, dtype=torch.float32), + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + divisor_override, + ) + return output.new_empty( + output.shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ) + + +_avg_pool2d_impl.register_fake(_avg_pool2d_meta) + + +@torch.library.custom_op("fused_quant::avg_pool2d_channels_last", mutates_args=()) +def _avg_pool2d_channels_last_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + ceil_mode: bool, + count_include_pad: bool, + divisor_override: Optional[int], +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + inp_nchw = _permute_contiguous(inp, [0, 3, 1, 2]) + inp_qp = _permute_qparams(inp_qp, [0, 3, 1, 2]) + output_nchw = _avg_pool2d( + _maybe_dequantize(inp_nchw, inp_qp), + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + divisor_override, + ) + output = _permute_contiguous(output_nchw, [0, 2, 3, 1]) + return _maybe_quantize(output, out_qp) + + +def _avg_pool2d_channels_last_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + kernel_size: list[int], + stride: list[int], + padding: list[int], + ceil_mode: bool, + count_include_pad: bool, + divisor_override: Optional[int], +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + output_nchw = _avg_pool2d( + torch.ones_like(inp, dtype=torch.float32).permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + divisor_override, + ) + output = output_nchw.permute(0, 2, 3, 1) + return output.new_empty( + output.shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ) + + +_avg_pool2d_channels_last_impl.register_fake(_avg_pool2d_channels_last_meta) + + +@torch.library.custom_op("fused_quant::convolution_channels_last", mutates_args=()) +def _convolution_channels_last_impl( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + stride: list[int], + padding: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + bias_qp = _make_qp( + bias_scale, + bias_zero_point, + bias_dtype, + bias_quant_min, + bias_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + + _validate_qparams(inp_qp, weight_qp, bias_qp, out_qp) + + num_spatial_dims = len(stride) + assert len(padding) == num_spatial_dims + assert len(dilation) == num_spatial_dims + assert len(output_padding) == num_spatial_dims + + nhwc_to_nchw = [0, num_spatial_dims + 1] + list(range(1, num_spatial_dims + 1)) + nchw_to_nhwc = [0] + list(range(2, num_spatial_dims + 2)) + [1] + ohwi_to_oihw = [0, num_spatial_dims + 1] + list(range(1, num_spatial_dims + 1)) + + inp_nchw = _permute_contiguous(inp, nhwc_to_nchw) + weight_oihw = _permute_contiguous(weight, ohwi_to_oihw) + inp_qp = _permute_qparams(inp_qp, nhwc_to_nchw) + weight_qp = _permute_qparams(weight_qp, ohwi_to_oihw) + + dq_inp = _maybe_dequantize(inp_nchw, inp_qp) + dq_weight = _maybe_dequantize(weight_oihw, weight_qp) + dq_bias = _maybe_dequantize(bias, bias_qp) if bias is not None else None + + if num_spatial_dims not in (1, 2, 3): + raise ValueError( + f"Unsupported number of spatial dimensions: {num_spatial_dims}. " + "Only 1D, 2D, and 3D convolutions are supported." + ) + + dq_out_nchw = torch.ops.aten.convolution.default( + dq_inp, + dq_weight, + dq_bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + dq_out = _permute_contiguous(dq_out_nchw, nchw_to_nhwc) + return _maybe_quantize(dq_out, out_qp) + + +def _convolution_channels_last_meta( + inp: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + bias_scale: Optional[torch.Tensor], + bias_zero_point: Optional[torch.Tensor], + bias_dtype: torch.dtype, + bias_quant_min: int, + bias_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + stride: list[int], + padding: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, weight_qp, out_qp) + output_shape = compute_conv_out_shape_nhwc( + inp, + weight, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + return inp.new_empty( + output_shape, dtype=get_out_dtype(out_qp, dequantized_dtype(inp, inp_qp)) + ) + + +_convolution_channels_last_impl.register_fake(_convolution_channels_last_meta) + + +# ============================================================================= +# Multi-output ops +# ============================================================================= +# +# Multi-output ops return a tuple. Unlike the single-output ops above, only a +# subset of their outputs are quantizable, and the op schema only carries output +# qparams blocks for those outputs. ``native_layer_norm`` returns +# ``(out, mean, rstd)`` but only ``out`` (output 0) is ever quantized; ``mean`` +# and ``rstd`` are always returned in float. Accordingly the schema has exactly +# one output qparams block (for output 0), and the input activation is the only +# quantizable input -- ``weight``/``bias`` pass through in float. + + +@torch.library.custom_op("fused_quant::native_layer_norm", mutates_args=()) +def _native_layer_norm_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + normalized_shape: list[int], + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + + _validate_qparams(inp_qp, out_qp) + + dq_inp = _maybe_dequantize(inp, inp_qp) + out, mean, rstd = torch.ops.aten.native_layer_norm.default( + dq_inp, normalized_shape, weight, bias, eps + ) + # Only output 0 is (optionally) quantized; mean/rstd stay float. + out = _maybe_quantize(out, out_qp) + return out, mean, rstd + + +def _native_layer_norm_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + normalized_shape: list[int], + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + + # Mirror the impl exactly: the norm runs on the dequantized input, so out0 + # (when unquantized) and mean/rstd take their dtypes straight from + # aten.native_layer_norm. ``compute_dtype`` is the dequantized input's float + # dtype -- ``inp_qp.dtype`` (the upstream dequant's output dtype, not always + # fp32), or the input dtype when the input isn't quantized. mean/rstd are not + # unconditionally fp32: aten produces them in its accumulation dtype and may + # cast them back to the (compute) input dtype -- so defer to aten rather than + # assuming fp32. + compute_dtype = dequantized_dtype(inp, inp_qp) + dq_inp = inp.new_empty(inp.shape, dtype=compute_dtype) + out, mean, rstd = torch.ops.aten.native_layer_norm.default( + dq_inp, normalized_shape, weight, bias, eps + ) + # Only output 0 is (optionally) quantized; mean/rstd stay float. + if out_qp is not None: + out = out.new_empty(out.shape, dtype=out_qp.dtype) + return out, mean, rstd + + +_native_layer_norm_impl.register_fake(_native_layer_norm_meta) + + +# ``rms_norm`` is single-output (unlike ``native_layer_norm``): it returns just +# the normalized tensor. Only the input activation and the output are quantized; +# ``weight`` (the learned scale) passes through in float, and ``normalized_shape`` +# / ``eps`` are plain attributes. RMS norm has no bias. +@torch.library.custom_op("fused_quant::rms_norm", mutates_args=()) +def _rms_norm_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + normalized_shape: list[int], + weight: Optional[torch.Tensor], + eps: float, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + + _validate_qparams(inp_qp, out_qp) + + dq_inp = _maybe_dequantize(inp, inp_qp) + out = torch.ops.aten.rms_norm.default(dq_inp, normalized_shape, weight, eps) + return _maybe_quantize(out, out_qp) + + +def _rms_norm_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + normalized_shape: list[int], + weight: Optional[torch.Tensor], + eps: float, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + + # Output matches the input's shape. When the output is quantized its dtype is + # the quantized out dtype. Otherwise the op returns the dequantized result, + # whose float dtype is whatever dequantizing the input yields -- i.e. + # ``inp_qp.dtype`` (the upstream dequant's output dtype), which is not always + # fp32. A pipeline-interior op with stripped quantization has neither input nor + # output quantized; then the input is already float, so use its own dtype. + compute_dtype = dequantized_dtype(inp, inp_qp) + return inp.new_empty(inp.shape, dtype=get_out_dtype(out_qp, compute_dtype)) + + +_rms_norm_impl.register_fake(_rms_norm_meta) + + +# ``_masked_softmax`` recomposes attention's ``add(mask) + softmax`` (see +# FuseAddSoftmaxIntoMaskedSoftmax). Only the input activation (the attention +# scores) and the output are quantized; the bool ``mask`` passes through +# unquantized and ``dim`` / ``mask_type`` are plain attributes. Like rms_norm it +# has no C++ kernel -- backend is expected to lower it before to_executorch. +@torch.library.custom_op("fused_quant::_masked_softmax", mutates_args=()) +def _masked_softmax_impl( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + mask: torch.Tensor, + dim: int, + mask_type: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + + _validate_qparams(inp_qp, out_qp) + + dq_inp = _maybe_dequantize(inp, inp_qp) + out = torch.ops.aten._masked_softmax.default(dq_inp, mask, dim, mask_type) + return _maybe_quantize(out, out_qp) + + +def _masked_softmax_meta( + inp: torch.Tensor, + inp_scale: Optional[torch.Tensor], + inp_zero_point: Optional[torch.Tensor], + inp_dtype: torch.dtype, + inp_quant_min: int, + inp_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + mask: torch.Tensor, + dim: int, + mask_type: int, +) -> torch.Tensor: + inp_qp = _make_qp( + inp_scale, inp_zero_point, inp_dtype, inp_quant_min, inp_quant_max + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + _validate_qparams(inp_qp, out_qp) + # Softmax preserves the input shape. When the output is quantized its dtype is + # the quantized out dtype; otherwise it is the dequantized input's float dtype + # (inp_qp.dtype, not always fp32), or -- for a pipeline-interior op stripped by + # interior quantization stripping, where neither input nor output is quantized + # -- the already-float input's own dtype. + compute_dtype = dequantized_dtype(inp, inp_qp) + return inp.new_empty(inp.shape, dtype=get_out_dtype(out_qp, compute_dtype)) + + +_masked_softmax_impl.register_fake(_masked_softmax_meta) + + +@torch.library.custom_op("fused_quant::embedding", mutates_args=()) +def _embedding_impl( + weight: torch.Tensor, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + indices: torch.Tensor, + padding_idx: int = -1, + scale_grad_by_freq: bool = False, + sparse: bool = False, +) -> torch.Tensor: + # Embedding is a compression-only op (no compute): the quantized table is + # dequantized and the rows selected by `indices` are returned. The table + # (the quantized "input") carries a dequant qparams block that may be + # per-tensor, per-channel, or per-group (the group axis partitions the + # embedding dimension). The output qparams block is optional: it is usually + # absent (the gathered embeddings stay float), + # but the layout mirrors every other fused_quant op so the fusion and + # optimization passes (e.g. quant absorption) can treat outputs uniformly. + # + # padding_idx/scale_grad_by_freq/sparse mirror aten.embedding so the generic + # fusion can thread them through. They are training/gradient concerns with no + # effect on the inference forward (the gather), so none of them are honored: + # scale_grad_by_freq/sparse must be False and padding_idx must be the aten + # default -1 (no padding token). The asserts make any model that relies on + # padding_idx semantics fail loudly rather than silently ignoring them. + assert padding_idx == -1, ( + "fused_quant.embedding does not honor padding_idx; expected the aten " + f"default -1 (no padding token), got {padding_idx}" + ) + assert not scale_grad_by_freq, ( + "fused_quant.embedding does not support scale_grad_by_freq=True (training only)" + ) + assert not sparse, ( + "fused_quant.embedding does not support sparse=True (training only)" + ) + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + # embedding is compression-only: an all-float form is meaningless, so unlike + # other ops it requires the table or the output to be quantized. + if weight_qp is None and out_qp is None: + raise ValueError( + "At least one of the embedding table (weight) or output must be quantized" + ) + _validate_qparams(weight_qp, out_qp) + if weight.ndim != 2: + raise ValueError( + "embedding weight (table) must be 2D [num_embeddings, embedding_dim], " + f"got {weight.ndim}D" + ) + dq_weight = _maybe_dequantize(weight, weight_qp) + out = torch.nn.functional.embedding(indices, dq_weight) + return _maybe_quantize(out, out_qp) + + +def _embedding_meta( + weight: torch.Tensor, + weight_scale: Optional[torch.Tensor], + weight_zero_point: Optional[torch.Tensor], + weight_dtype: torch.dtype, + weight_quant_min: int, + weight_quant_max: int, + out_scale: Optional[torch.Tensor], + out_zero_point: Optional[torch.Tensor], + out_dtype: torch.dtype, + out_quant_min: int, + out_quant_max: int, + indices: torch.Tensor, + padding_idx: int = -1, + scale_grad_by_freq: bool = False, + sparse: bool = False, +) -> torch.Tensor: + weight_qp = _make_qp( + weight_scale, + weight_zero_point, + weight_dtype, + weight_quant_min, + weight_quant_max, + ) + out_qp = _make_qp( + out_scale, out_zero_point, out_dtype, out_quant_min, out_quant_max + ) + assert padding_idx == -1, ( + "fused_quant.embedding does not honor padding_idx; expected the aten " + f"default -1 (no padding token), got {padding_idx}" + ) + assert not scale_grad_by_freq, ( + "fused_quant.embedding does not support scale_grad_by_freq=True (training only)" + ) + assert not sparse, ( + "fused_quant.embedding does not support sparse=True (training only)" + ) + if weight_qp is None and out_qp is None: + raise ValueError( + "At least one of the embedding table (weight) or output must be quantized" + ) + _validate_qparams(weight_qp, out_qp) + if weight.ndim != 2: + raise ValueError( + "embedding weight (table) must be 2D [num_embeddings, embedding_dim], " + f"got {weight.ndim}D" + ) + # Output is the gathered rows: indices.shape + [embedding_dim]. Its dtype + # mirrors the impl: the quantized dtype when the output is quantized, else the + # table's dequantize output dtype (weight_qp.dtype), or the table dtype when + # the table is not dequantized -- never assumed to be float32. + table_dtype = dequantized_dtype(weight, weight_qp) + out_shape = (*indices.shape, weight.shape[1]) + return weight.new_empty(out_shape, dtype=get_out_dtype(out_qp, table_dtype)) + + +_embedding_impl.register_fake(_embedding_meta) diff --git a/backends/fused_quant/ops_utils.py b/backends/fused_quant/ops_utils.py new file mode 100644 index 00000000000..78efe80844d --- /dev/null +++ b/backends/fused_quant/ops_utils.py @@ -0,0 +1,104 @@ +# 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. + +import torch + + +def compute_conv_out_shape( + inp: torch.Tensor, + weight: torch.Tensor, + stride: list[int], + padding: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> list[int]: + """Compute output shape for NCHW convolution. + + Args: + inp: Input tensor in NCHW format (batch, channels, *spatial) + weight: Weight tensor in OIHW or IOHW format + stride: Stride for each spatial dimension + padding: Padding for each spatial dimension + dilation: Dilation for each spatial dimension + transposed: Whether this is a transposed convolution + output_padding: Additional output size for transposed convolution + groups: Number of blocked connections from inputs to outputs + + Returns: + Output shape in NCHW format [batch, out_channels, *spatial_out] + """ + batch_size = inp.shape[0] + out_channels = weight.shape[1] * groups if transposed else weight.shape[0] + spatial_dims = [] + for i in range(len(stride)): + input_size = inp.shape[i + 2] # Skip N, C dimensions + kernel_size = weight.shape[i + 2] + if transposed: + out_size = ( + (input_size - 1) * stride[i] + - 2 * padding[i] + + dilation[i] * (kernel_size - 1) + + output_padding[i] + + 1 + ) + else: + out_size = ( + input_size + 2 * padding[i] - dilation[i] * (kernel_size - 1) - 1 + ) // stride[i] + 1 + spatial_dims.append(out_size) + + return [batch_size, out_channels] + spatial_dims + + +def compute_conv_out_shape_nhwc( + inp: torch.Tensor, + weight: torch.Tensor, + stride: list[int], + padding: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> list[int]: + """Compute output shape for NHWC convolution. + + Args: + inp: Input tensor in NHWC format (batch, *spatial, channels) + weight: Weight tensor in OHWI or IHWO format + stride: Stride for each spatial dimension + padding: Padding for each spatial dimension + dilation: Dilation for each spatial dimension + transposed: Whether this is a transposed convolution + output_padding: Additional output size for transposed convolution + groups: Number of blocked connections from inputs to outputs + + Returns: + Output shape in NHWC format [batch, *spatial_out, out_channels] + """ + batch_size = inp.shape[0] + out_channels = weight.shape[-1] * groups if transposed else weight.shape[0] + num_spatial_dims = len(stride) + spatial_dims = [] + for i in range(num_spatial_dims): + input_size = inp.shape[i + 1] # Skip N, spatial dims start at index 1 + kernel_size = weight.shape[i + 1] # OHWI: spatial dims at indices 1..n + if transposed: + out_size = ( + (input_size - 1) * stride[i] + - 2 * padding[i] + + dilation[i] * (kernel_size - 1) + + output_padding[i] + + 1 + ) + else: + out_size = ( + input_size + 2 * padding[i] - dilation[i] * (kernel_size - 1) - 1 + ) // stride[i] + 1 + spatial_dims.append(out_size) + + return [batch_size] + spatial_dims + [out_channels] diff --git a/backends/fused_quant/optimization_passes/__init__.py b/backends/fused_quant/optimization_passes/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/fused_quant/optimization_passes/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/fused_quant/optimization_passes/constant_fold.py b/backends/fused_quant/optimization_passes/constant_fold.py new file mode 100644 index 00000000000..0da3dae80b8 --- /dev/null +++ b/backends/fused_quant/optimization_passes/constant_fold.py @@ -0,0 +1,25 @@ +# 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. + + +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch.export import ExportedProgram + + +class ConstantFold(ExportedProgramPassBase): + """Runs ``constant_prop_pass`` as a pass with a correct ``modified`` flag. + + ``constant_prop_pass`` folds pure-constant subgraphs but returns no signal for + whether it changed the graph; this wrapper derives ``modified`` from the node + count so it composes correctly in an iterative pass pipeline. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + n_before = len(exported_program.graph_module.graph.nodes) + folded = constant_prop_pass(exported_program) + modified = len(folded.graph_module.graph.nodes) != n_before + return ExportedProgramPassResult(folded, modified) diff --git a/backends/fused_quant/optimization_passes/fold_scalar_mul.py b/backends/fused_quant/optimization_passes/fold_scalar_mul.py new file mode 100644 index 00000000000..985ecf47a42 --- /dev/null +++ b/backends/fused_quant/optimization_passes/fold_scalar_mul.py @@ -0,0 +1,163 @@ +# 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. + +from __future__ import annotations + +import math + +import executorch.backends.fused_quant.ops # noqa: F401 +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_from_node, + get_scale, + get_zero_point, +) +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import compute_meta_val +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch.export import ExportedProgram + + +def _fold_into_dequant_quant(ep: ExportedProgram, node: fx.Node, scalar: float) -> bool: + """Fold scalar into dequant scale when both inp and out qparams are present. + + ``Q_out(c * DQ_inp(x))`` becomes ``DQ_{inp*c}(x) -> Q_out``. + """ + inp_qparams = get_qparams_from_node(node, "inp") + out_qparams = get_qparams_from_node(node, "out") + assert inp_qparams is not None and out_qparams is not None + + if not inp_qparams.is_per_tensor() or not out_qparams.is_per_tensor(): + return False + + dequant_scale = get_scale(ep, inp_qparams) * scalar + if dequant_scale <= 0.0 or not math.isfinite(dequant_scale): + return False + + inp = get_arg(node, "inp", fx.Node) + with node.graph.inserting_before(node): + dequant = node.graph.call_function( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=( + inp, + dequant_scale, + get_zero_point(ep, inp_qparams), + inp_qparams.quant_min, + inp_qparams.quant_max, + inp.meta["val"].dtype, + ), + kwargs={"out_dtype": inp_qparams.dtype}, + ) + dequant.meta = node.meta.copy() + dequant.meta["val"] = compute_meta_val(dequant) + + quant = node.graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=( + dequant, + get_scale(ep, out_qparams), + get_zero_point(ep, out_qparams), + out_qparams.quant_min, + out_qparams.quant_max, + out_qparams.dtype, + ), + ) + quant.meta = node.meta.copy() + quant.meta["val"] = compute_meta_val(quant) + + node.replace_all_uses_with(quant) + node.graph.erase_node(node) + return True + + +def _fold_into_quant(ep: ExportedProgram, node: fx.Node, scalar: float) -> bool: + """Fold scalar into quant scale when only out qparams are present. + + ``Q_out(float_x * c)`` becomes ``Q_{out/c}(float_x)`` by dividing the + quantize scale by the scalar: ``round(x * c / s) = round(x / (s / c))``. + """ + out_qparams = get_qparams_from_node(node, "out") + assert out_qparams is not None + + if not out_qparams.is_per_tensor(): + return False + + new_scale = get_scale(ep, out_qparams) / scalar + if new_scale <= 0.0 or not math.isfinite(new_scale): + return False + + inp = get_arg(node, "inp", fx.Node) + with node.graph.inserting_before(node): + quant = node.graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=( + inp, + new_scale, + get_zero_point(ep, out_qparams), + out_qparams.quant_min, + out_qparams.quant_max, + out_qparams.dtype, + ), + ) + quant.meta = node.meta.copy() + quant.meta["val"] = compute_meta_val(quant) + + node.replace_all_uses_with(quant) + node.graph.erase_node(node) + return True + + +def _fold_scalar_mul(ep: ExportedProgram, node: fx.Node) -> bool: + other = get_arg(node, "other") + if not isinstance(other, (int, float)): + return False + scalar = float(other) + if scalar <= 0.0 or not math.isfinite(scalar): + return False + + inp_qparams = get_qparams_from_node(node, "inp") + out_qparams = get_qparams_from_node(node, "out") + + if inp_qparams is not None and out_qparams is not None: + return _fold_into_dequant_quant(ep, node, scalar) + + if inp_qparams is None and out_qparams is not None: + return _fold_into_quant(ep, node, scalar) + + return False + + +class FoldScalarMul(ExportedProgramPassBase): + """Fold a positive scalar multiply into neighboring quantization ops. + + Handles two cases of ``fused_quant.mul.Scalar(x, c)``: + + 1. Both input and output qparams present (fully quantized): + ``Q_B(c * DQ_A(x))`` -> ``DQ_{A*c}(x) -> Q_B``, folding ``c`` into the + dequant scale so later quant absorption can consume it. + + 2. Only output qparams present (float input, quantized output): + ``Q_B(float_x * c)`` -> ``Q_{B/c}(float_x)``, folding ``c`` into the + quant scale by dividing it, eliminating the multiply entirely. + + Negative, zero, and non-finite scalars are left unchanged (they cannot be + represented by a positive quantization scale). Per-channel qparams are + also left unchanged. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + for node in graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.mul.Scalar + ): + modified |= _fold_scalar_mul(exported_program, node) + + if modified: + exported_program = constant_prop_pass(exported_program) + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/optimization_passes/fuse_add_into_linear.py b/backends/fused_quant/optimization_passes/fuse_add_into_linear.py new file mode 100644 index 00000000000..5f73f350128 --- /dev/null +++ b/backends/fused_quant/optimization_passes/fuse_add_into_linear.py @@ -0,0 +1,188 @@ +# 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. + + +import logging +from typing import Optional + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg +from executorch.backends.transforms.utils import ( + add_constant, + get_constant, + set_param_tensor, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +logger: logging.Logger = logging.getLogger(__name__) + +_PASSTHROUGH_TARGETS: set[EdgeOpOverload] = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.view_copy.default, +} + +_CONV_TARGETS: set[EdgeOpOverload] = { + exir_ops.edge.fused_quant.convolution.default, + exir_ops.edge.fused_quant.convolution_channels_last.default, +} + +_AFFINE_TARGETS: set[EdgeOpOverload] = { + exir_ops.edge.fused_quant.linear.default, + *_CONV_TARGETS, +} + + +def _find_affine_through_passthrough( + node: fx.Node, +) -> Optional[tuple[fx.Node, list[fx.Node]]]: + chain: list[fx.Node] = [] + current = node + while current.target in _PASSTHROUGH_TARGETS: + if len(current.users) != 1: + return None + chain.append(current) + prev = current.args[0] + if not isinstance(prev, fx.Node): + return None + current = prev + + if current.target not in _AFFINE_TARGETS: + return None + if current.target in _CONV_TARGETS and get_arg(current, "transposed", bool): + return None + if len(current.users) != 1: + return None + + return current, chain + + +def _has_per_tensor_output_qparams(ep: ExportedProgram, add_node: fx.Node) -> bool: + for field in ("out_scale", "out_zero_point"): + node = get_arg(add_node, field, Optional[fx.Node]) + if node is None: + continue + val = get_constant(ep, node) + if val is None or val.numel() > 1: + return False + return True + + +def _get_bias_increment(add_node: fx.Node) -> float: + other = get_arg(add_node, "other", float | int) + alpha = get_arg(add_node, "alpha", float | int) + return float(other) * alpha + + +def _get_new_bias( + ep: ExportedProgram, affine_node: fx.Node, increment: float +) -> Optional[torch.Tensor]: + bias_node = get_arg(affine_node, "bias", Optional[fx.Node]) + if bias_node is None: + weight_node = get_arg(affine_node, "weight", fx.Node) + out_channels = weight_node.meta["val"].shape[0] + return torch.full((out_channels,), increment, dtype=torch.float32) + + if len(bias_node.users) != 1: + return None + if get_arg(affine_node, "bias_scale", Optional[fx.Node]) is not None: + return None + + bias = get_constant(ep, bias_node) + if bias is None: + return None + return bias + bias.new_tensor(increment) + + +def _output_qparams_match_input( + ep: ExportedProgram, affine_node: fx.Node, add_node: fx.Node +) -> bool: + for affine_field, add_field in ( + ("out_scale", "inp_scale"), + ("out_zero_point", "inp_zero_point"), + ): + affine_qp = get_arg(affine_node, affine_field, Optional[fx.Node]) + add_qp = get_arg(add_node, add_field, Optional[fx.Node]) + if affine_qp is None or add_qp is None: + assert affine_qp is None and add_qp is None + continue + if affine_qp is add_qp: + continue + affine_val = get_constant(ep, affine_qp) + add_val = get_constant(ep, add_qp) + if affine_val is None or add_val is None: + return False + if not torch.equal(affine_val, add_val): + return False + return True + + +class FuseAddIntoLinear(ExportedProgramPassBase): + """Folds fused_quant.add.Scalar into linear or convolution bias.""" + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + + for add_node in graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.add.Scalar + ): + inp_node = get_arg(add_node, "inp", fx.Node) + result = _find_affine_through_passthrough(inp_node) + if result is None: + continue + affine_node, passthrough_chain = result + + if passthrough_chain and not _has_per_tensor_output_qparams( + exported_program, add_node + ): + continue + if not _output_qparams_match_input(exported_program, affine_node, add_node): + continue + if get_arg(add_node, "out_dtype") != get_arg(affine_node, "out_dtype"): + continue + + new_bias = _get_new_bias( + exported_program, affine_node, _get_bias_increment(add_node) + ) + if new_bias is None: + continue + + bias_node = get_arg(affine_node, "bias", Optional[fx.Node]) + if bias_node is None: + bias_node = add_constant( + exported_program, + f"{affine_node.name}_bias", + new_bias, + affine_node, + InputKind.BUFFER, + ) + set_arg(affine_node, "bias", bias_node) + else: + set_param_tensor(exported_program, bias_node, new_bias) + + for field in ( + "out_zero_point", + "out_scale", + "out_quant_min", + "out_quant_max", + ): + set_arg(affine_node, field, get_arg(add_node, field)) + + add_node.replace_all_uses_with(inp_node) + graph.erase_node(add_node) + modified = True + logger.info("Folded scalar add into affine bias: %s", affine_node.name) + + if modified: + exported_program = constant_prop_pass(exported_program) + + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/optimization_passes/fuse_mul_into_linear.py b/backends/fused_quant/optimization_passes/fuse_mul_into_linear.py new file mode 100644 index 00000000000..a51ee756ddd --- /dev/null +++ b/backends/fused_quant/optimization_passes/fuse_mul_into_linear.py @@ -0,0 +1,593 @@ +# 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. + + +import logging +import math +from collections.abc import Iterator +from dataclasses import dataclass +from typing import cast, Optional + +import torch +from executorch.backends.fused_quant.graph_utils import get_qparams_from_node +from executorch.backends.fused_quant.ops import QuantParamsStruct +from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg +from executorch.backends.transforms.utils import ( + get_constant, + set_param_tensor, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch.export import ExportedProgram + +logger: logging.Logger = logging.getLogger(__name__) + +_PASSTHROUGH_TARGETS: set[EdgeOpOverload] = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.view_copy.default, +} +_CONV_TARGETS: set[EdgeOpOverload] = { + exir_ops.edge.fused_quant.convolution.default, + exir_ops.edge.fused_quant.convolution_channels_last.default, +} +_AFFINE_TARGETS: set[EdgeOpOverload] = { + exir_ops.edge.fused_quant.linear.default, + *_CONV_TARGETS, +} + + +@dataclass(frozen=True) +class _ConstantUpdate: + node: fx.Node + value: torch.Tensor + + +@dataclass(frozen=True) +class _SignedWeightPlan: + weight_update: Optional[_ConstantUpdate] + zero_point_update: Optional[_ConstantUpdate] + + +@dataclass(frozen=True) +class _FoldPlan: + affine_node: fx.Node + passthrough_chain: list[fx.Node] + weight_scale_node: fx.Node + new_weight_scale: torch.Tensor + signed_weight_plan: _SignedWeightPlan + bias_node: Optional[fx.Node] + new_bias: Optional[torch.Tensor] + + +def _find_affine_through_passthrough( + node: fx.Node, +) -> Optional[tuple[fx.Node, list[fx.Node]]]: + """Walk backward from *node* through passthrough ops to find an affine op. + + Returns ``(affine_node, passthrough_chain)`` where *passthrough_chain* is the + list of passthrough nodes between the affine and *node* (empty when *node* is + the affine itself). Returns ``None`` if no affine is reachable through + single-user passthrough ops. + """ + + chain: list[fx.Node] = [] + current = node + while current.target in _PASSTHROUGH_TARGETS: + if len(current.users) != 1: + return None + chain.append(current) + prev = current.args[0] + if not isinstance(prev, fx.Node): + return None + current = prev + + if current.target not in _AFFINE_TARGETS: + return None + if current.target in _CONV_TARGETS and get_arg(current, "transposed", bool): + return None + + affine_node = current + if len(affine_node.users) != 1: + return None + + return affine_node, chain + + +def _affine_output_channel_axis(affine_node: fx.Node) -> int: + if affine_node.target == exir_ops.edge.fused_quant.convolution.default: + return 1 + return affine_node.meta["val"].ndim - 1 + + +def _channel_axis_after_view( + input_shape: torch.Size, + output_shape: torch.Size, + input_channel_axis: int, +) -> Optional[int]: + channel_size = input_shape[input_channel_axis] + prefix_size = math.prod(input_shape[:input_channel_axis]) + suffix_size = math.prod(input_shape[input_channel_axis + 1 :]) + candidates = [ + axis + for axis, size in enumerate(output_shape) + if size == channel_size + and math.prod(output_shape[:axis]) == prefix_size + and math.prod(output_shape[axis + 1 :]) == suffix_size + ] + return candidates[0] if len(candidates) == 1 else None + + +def _channel_axis_after_passthrough( + affine_node: fx.Node, + passthrough_chain: list[fx.Node], +) -> Optional[int]: + channel_axis = _affine_output_channel_axis(affine_node) + input_shape = affine_node.meta["val"].shape + for passthrough in reversed(passthrough_chain): + output_shape = passthrough.meta["val"].shape + if passthrough.target == exir_ops.edge.aten.permute_copy.default: + dims = [ + dim % len(input_shape) + for dim in get_arg(passthrough, "dims", list[int]) + ] + if len(dims) != len(input_shape) or sorted(dims) != list( + range(len(input_shape)) + ): + return None + channel_axis = dims.index(channel_axis) + else: + view_channel_axis = _channel_axis_after_view( + input_shape, + output_shape, + channel_axis, + ) + if view_channel_axis is None: + return None + channel_axis = view_channel_axis + input_shape = output_shape + return channel_axis + + +def _constant_to_channel_scale( + scale_tensor: torch.Tensor, + output_shape: torch.Size, + output_channel_axis: Optional[int], +) -> Optional[torch.Tensor]: + """Collapse a scalar or output-channel multiplier to one dimension.""" + if scale_tensor.ndim > len(output_shape): + return None + + aligned_shape = (1,) * (len(output_shape) - scale_tensor.ndim) + tuple( + scale_tensor.shape + ) + if any( + scale_size not in (1, output_size) + for scale_size, output_size in zip(aligned_shape, output_shape) + ): + return None + + channel_scale = scale_tensor.reshape(-1) + if channel_scale.is_complex() or not bool(torch.all(torch.isfinite(channel_scale))): + return None + if channel_scale.numel() == 1: + return channel_scale + if output_channel_axis is None: + return None + if any( + axis != output_channel_axis and size != 1 + for axis, size in enumerate(aligned_shape) + ): + return None + if aligned_shape[output_channel_axis] != output_shape[output_channel_axis]: + return None + return channel_scale + + +def _is_per_channel_qparam(qparam: torch.Tensor, out_channels: int) -> bool: + return ( + qparam.numel() == out_channels + and qparam.shape[0] == out_channels + and all(size == 1 for size in qparam.shape[1:]) + ) + + +def _reshape_channel_values( + channel_values: torch.Tensor, + ndim: int, + out_channels: int, +) -> Optional[torch.Tensor]: + if channel_values.numel() == 1: + return channel_values.reshape(()) + if channel_values.numel() != out_channels or ndim == 0: + return None + return channel_values.reshape([out_channels] + [1] * (ndim - 1)) + + +def _get_weight_scale_factor( + weight_scale: torch.Tensor, + channel_scale: torch.Tensor, + out_channels: int, +) -> Optional[torch.Tensor]: + if channel_scale.numel() != 1 and not _is_per_channel_qparam( + weight_scale, out_channels + ): + return None + return _reshape_channel_values(channel_scale, weight_scale.ndim, out_channels) + + +def _get_new_weight_scale( + weight_scale: torch.Tensor, + channel_scale: torch.Tensor, + out_channels: int, +) -> Optional[torch.Tensor]: + scale_factor = _get_weight_scale_factor( + weight_scale, + channel_scale.abs(), + out_channels, + ) + zero_mask = _get_weight_scale_factor( + weight_scale, + channel_scale == 0, + out_channels, + ) + if scale_factor is None or zero_mask is None: + return None + scale_factor = scale_factor.to(weight_scale) + zero_mask = zero_mask.to(device=weight_scale.device) + return torch.where( + zero_mask, + torch.ones_like(weight_scale), + weight_scale * scale_factor, + ) + + +def _update_integer_channels( + value: torch.Tensor, + channel_scale: torch.Tensor, + out_channels: int, + quantized_range_sum: int, +) -> Optional[torch.Tensor]: + if value.is_floating_point() or value.is_complex(): + return None + shaped_scale = _reshape_channel_values( + channel_scale, + value.ndim, + out_channels, + ) + if shaped_scale is None: + return None + shaped_scale = shaped_scale.to(device=value.device) + # Applying the same range reflection to q and z gives q' - z' = -(q - z). + reflected = (quantized_range_sum - value.to(torch.int64)).to(value.dtype) + updated = torch.where(shaped_scale < 0, reflected, value) + return torch.where(shaped_scale == 0, torch.zeros_like(updated), updated) + + +def _get_signed_weight_plan( + exported_program: ExportedProgram, + affine_node: fx.Node, + channel_scale: torch.Tensor, + out_channels: int, +) -> Optional[_SignedWeightPlan]: + if bool(torch.all(channel_scale > 0)): + return _SignedWeightPlan(None, None) + + weight_node = get_arg(affine_node, "weight", fx.Node) + zero_point_node = get_arg(affine_node, "weight_zero_point", fx.Node) + if len(weight_node.users) != 1 or len(zero_point_node.users) != 1: + return None + weight = get_constant(exported_program, weight_node) + zero_point = get_constant(exported_program, zero_point_node) + if weight is None or zero_point is None: + return None + if weight.ndim == 0 or weight.shape[0] != out_channels: + return None + if channel_scale.numel() != 1 and not _is_per_channel_qparam( + zero_point, out_channels + ): + return None + + quant_min = get_arg(affine_node, "weight_quant_min", int) + quant_max = get_arg(affine_node, "weight_quant_max", int) + if bool(torch.any(channel_scale == 0)) and not quant_min <= 0 <= quant_max: + return None + range_sum = quant_min + quant_max + new_weight = _update_integer_channels( + weight, channel_scale, out_channels, range_sum + ) + new_zero_point = _update_integer_channels( + zero_point, channel_scale, out_channels, range_sum + ) + if new_weight is None or new_zero_point is None: + return None + return _SignedWeightPlan( + _ConstantUpdate(weight_node, new_weight), + _ConstantUpdate(zero_point_node, new_zero_point), + ) + + +def _get_scaled_bias( + exported_program: ExportedProgram, + affine_node: fx.Node, + channel_scale: torch.Tensor, +) -> Optional[tuple[Optional[fx.Node], Optional[torch.Tensor]]]: + bias_node = get_arg(affine_node, "bias", Optional[fx.Node]) + if bias_node is None: + return None, None + if len(bias_node.users) != 1: + return None + if get_arg(affine_node, "bias_scale", Optional[fx.Node]) is not None: + return None + bias = get_constant(exported_program, bias_node) + if bias is None: + return None + return bias_node, bias * channel_scale.to(bias) + + +def _make_fold_plan( + exported_program: ExportedProgram, + data_node: fx.Node, + scale_tensor: torch.Tensor, +) -> Optional[_FoldPlan]: + affine_info = _find_affine_through_passthrough(data_node) + if affine_info is None: + return None + affine_node, passthrough_chain = affine_info + + weight_scale_node = get_arg(affine_node, "weight_scale", Optional[fx.Node]) + if weight_scale_node is None or len(weight_scale_node.users) != 1: + return None + weight_scale = get_constant(exported_program, weight_scale_node) + if weight_scale is None: + return None + + output_channel_axis = None + if scale_tensor.numel() != 1: + output_channel_axis = _channel_axis_after_passthrough( + affine_node, + passthrough_chain, + ) + channel_scale = _constant_to_channel_scale( + scale_tensor, + data_node.meta["val"].shape, + output_channel_axis, + ) + if channel_scale is None: + return None + + affine_channel_axis = _affine_output_channel_axis(affine_node) + out_channels = cast(torch.Tensor, affine_node.meta["val"]).shape[ + affine_channel_axis + ] + new_weight_scale = _get_new_weight_scale( + weight_scale, + channel_scale, + out_channels, + ) + if new_weight_scale is None: + return None + + signed_weight_plan = _get_signed_weight_plan( + exported_program, + affine_node, + channel_scale, + out_channels, + ) + if signed_weight_plan is None: + return None + + scaled_bias = _get_scaled_bias(exported_program, affine_node, channel_scale) + if scaled_bias is None: + return None + bias_node, new_bias = scaled_bias + return _FoldPlan( + affine_node, + passthrough_chain, + weight_scale_node, + new_weight_scale, + signed_weight_plan, + bias_node, + new_bias, + ) + + +def _apply_fold( + exported_program: ExportedProgram, + mul_node: fx.Node, + data_node: fx.Node, + plan: _FoldPlan, +) -> None: + set_param_tensor( + exported_program, + plan.weight_scale_node, + plan.new_weight_scale, + ) + if plan.signed_weight_plan.weight_update is not None: + set_param_tensor( + exported_program, + plan.signed_weight_plan.weight_update.node, + plan.signed_weight_plan.weight_update.value, + ) + if plan.signed_weight_plan.zero_point_update is not None: + set_param_tensor( + exported_program, + plan.signed_weight_plan.zero_point_update.node, + plan.signed_weight_plan.zero_point_update.value, + ) + if plan.bias_node is not None and plan.new_bias is not None: + set_param_tensor(exported_program, plan.bias_node, plan.new_bias) + + for field in ( + "out_zero_point", + "out_scale", + "out_dtype", + "out_quant_min", + "out_quant_max", + ): + set_arg(plan.affine_node, field, get_arg(mul_node, field)) + + new_dtype = mul_node.meta["val"].dtype + for retyped in (plan.affine_node, *plan.passthrough_chain): + val = cast(torch.Tensor, retyped.meta["val"]) + retyped.meta["val"] = val.to(new_dtype) + + mul_node.replace_all_uses_with(data_node) + mul_node.graph.erase_node(mul_node) + + +def _trace_to_constant( + ep: ExportedProgram, + node: fx.Node, +) -> Optional[tuple[torch.Tensor, fx.Node]]: + """Find the constant tensor backing a node. + + Handles: + - Direct placeholder node → returns its stored tensor, which may be + floating point or integer quantized storage when an upstream pass folded + the quantize node. + - quantize_per_tensor(placeholder, ...) → returns the floating-point source + constant and the quantize node. + - view_copy/permute_copy(constant) → applies the layout transform to the + constant value. + + Callers must use the consuming op's qparams to interpret integer storage. + Returns None if the pattern doesn't match. + """ + tensor = get_constant(ep, node) + if tensor is not None: + return tensor, node + + if node.target == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: + source = node.args[0] + assert isinstance(source, fx.Node) + tensor = get_constant(ep, source) + if tensor is not None: + return tensor, node + + if node.target in _PASSTHROUGH_TARGETS: + source = node.args[0] + assert isinstance(source, fx.Node) + traced = _trace_to_constant(ep, source) + if traced is None: + return None + tensor, _ = traced + if node.target == exir_ops.edge.aten.permute_copy.default: + dims = get_arg(node, "dims", list[int]) + return tensor.permute(dims).contiguous(), node + elif node.target == exir_ops.edge.aten.view_copy.default: + size = get_arg(node, "size", list[int]) + return tensor.reshape(size), node + else: + raise AssertionError(f"Unexpected constant passthrough: {node.target}") + + return None + + +def _resolve_constant_value( + ep: ExportedProgram, + op_node: fx.Node, + qparams_prefix: str, + constant_node: fx.Node, +) -> Optional[torch.Tensor]: + constant_result = _trace_to_constant(ep, constant_node) + if constant_result is None: + return None + constant = constant_result[0] + if constant.is_floating_point(): + return constant + + qparams = get_qparams_from_node(op_node, qparams_prefix) + if qparams is None: + return constant + scale = get_constant(ep, qparams.scale) + zero_point = get_constant(ep, qparams.zero_point) + if scale is None or zero_point is None: + return None + return QuantParamsStruct( + scale=scale, + zero_point=zero_point, + dtype=qparams.dtype, + quant_min=qparams.quant_min, + quant_max=qparams.quant_max, + ).dequantize(constant) + + +def _get_mul_data_and_scale( + ep: ExportedProgram, mul_node: fx.Node +) -> Optional[tuple[fx.Node, torch.Tensor]]: + inp_node = get_arg(mul_node, "inp", fx.Node) + if mul_node.target == exir_ops.edge.fused_quant.mul.Scalar: + return inp_node, torch.as_tensor(get_arg(mul_node, "other")) + + other_node = get_arg(mul_node, "other", fx.Node) + scale_tensor = _resolve_constant_value(ep, mul_node, "other", other_node) + if scale_tensor is not None: + return inp_node, scale_tensor + + scale_tensor = _resolve_constant_value(ep, mul_node, "inp", inp_node) + if scale_tensor is not None: + return other_node, scale_tensor + + return None + + +def _iter_mul_data_and_scales( + ep: ExportedProgram, graph: fx.Graph +) -> Iterator[tuple[fx.Node, fx.Node, torch.Tensor]]: + mul_nodes = [ + *graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.mul.default + ), + *graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.mul.Scalar + ), + ] + for mul_node in mul_nodes: + data_and_scale = _get_mul_data_and_scale(ep, mul_node) + if data_and_scale is not None: + data_node, scale_tensor = data_and_scale + if mul_node.meta["val"].shape == data_node.meta["val"].shape: + yield mul_node, data_node, scale_tensor + + +class FuseMulIntoLinear(ExportedProgramPassBase): + """Folds a fused_quant.mul with a constant operand into a preceding + fused_quant linear or convolution by scaling its weight_scale and bias. + + Matches both fused_quant.mul tensor and Scalar overloads: + - affine → fused_quant.mul(constant) + - affine → [permute/view]* → fused_quant.mul(constant) + + Scalar constants are supported with per-tensor or per-channel weight qparams. + Per-output-channel constants require private per-channel weight qparams and an + output-channel axis that can be tracked through the passthroughs. Negative + channels reflect the integer weights and zero points across the quantized range; + zero channels use weight 0, scale 1, and zero point 0. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + + for mul_node, data_node, scale_tensor in _iter_mul_data_and_scales( + exported_program, graph + ): + plan = _make_fold_plan(exported_program, data_node, scale_tensor) + if plan is None: + continue + _apply_fold(exported_program, mul_node, data_node, plan) + + modified = True + logger.info("Folded mul into affine: %s", plan.affine_node.name) + + if modified: + # Doesn't hurt to run, and it also eliminates dead code and recompiles + exported_program = constant_prop_pass(exported_program) + + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/optimization_passes/hoist_activations.py b/backends/fused_quant/optimization_passes/hoist_activations.py new file mode 100644 index 00000000000..cab2a675d9a --- /dev/null +++ b/backends/fused_quant/optimization_passes/hoist_activations.py @@ -0,0 +1,87 @@ +# 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. + + +import executorch.backends.fused_quant.ops # noqa +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassBase, PassResult +from torch import fx +from torch._ops import OpOverload + + +# Volume-preserving data-movement ops that only rearrange tensor layout +# without changing element values. +_DATA_MOVEMENT_OPS: frozenset[OpOverload] = frozenset( + { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.view_copy.default, + } +) + +# The elementwise activations this pass hoists. +_ACTIVATION_OPS: frozenset[OpOverload] = frozenset( + { + exir_ops.edge.fused_quant.relu.default, + exir_ops.edge.fused_quant.hardswish.default, + exir_ops.edge.fused_quant.sigmoid.default, + exir_ops.edge.fused_quant.tanh.default, + exir_ops.edge.fused_quant.hard_tanh.default, + exir_ops.edge.fused_quant.silu.default, + exir_ops.edge.fused_quant.hardsigmoid.default, + exir_ops.edge.fused_quant.gelu.default, + } +) + + +class HoistActivations(PassBase): + """Hoist element-wise activations above the data-movement ops feeding them. + + Turns chains like ``permute → relu`` into ``relu → permute`` by moving each + single-user data-movement op forward past the activation. Each relocated + data-movement op is retyped to the activation's output dtype. + """ + + def call(self, graph_module: fx.GraphModule) -> PassResult: + modified = False + graph = graph_module.graph + + for target in _ACTIVATION_OPS: + for activation_node in graph.find_nodes(op="call_function", target=target): + out_dtype = activation_node.meta["val"].dtype + + data_movement_node = activation_node.args[0] + while ( + data_movement_node.target in _DATA_MOVEMENT_OPS + and len(data_movement_node.users) == 1 + ): + input_node = data_movement_node.args[0] + + # Rewire connections: + # from: input_node -> data_movement_node -> activation_node -> consumers + # to: input_node -> activation_node -> data_movement_node -> consumers + activation_node.replace_all_uses_with(data_movement_node) + data_movement_node.replace_input_with(input_node, activation_node) + activation_node.replace_input_with(data_movement_node, input_node) + + # Move data_movement_node to sit right after activation_node. + # Its sole user is activation_node, so this is always valid. + activation_node.append(data_movement_node) + + # activation_node keeps its own output dtype but takes its new + # input's shape; data_movement_node keeps its shape but now + # carries activation_node's output dtype. + activation_node.meta["val"] = input_node.meta["val"].to(out_dtype) + data_movement_node.meta["val"] = data_movement_node.meta["val"].to( + out_dtype + ) + # Advance to the next candidate: whatever now feeds activation_node. + data_movement_node = activation_node.args[0] + modified = True + + if modified: + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/fused_quant/optimization_passes/prequantize_embedding.py b/backends/fused_quant/optimization_passes/prequantize_embedding.py new file mode 100644 index 00000000000..5e431d01313 --- /dev/null +++ b/backends/fused_quant/optimization_passes/prequantize_embedding.py @@ -0,0 +1,122 @@ +# 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. + + +import torch +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_from_node, + get_scale, + get_zero_point, +) +from executorch.backends.fused_quant.ops import QuantParamsStruct +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import ( + compute_meta_val, + get_constant, + set_param_tensor, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch.export import ExportedProgram + + +def _try_prequantize_embedding(ep: ExportedProgram, node: fx.Node) -> bool: + """Fold a quantized-table embedding into a plain int8 gather. + + A fused_quant.embedding with a quantized table (weight qparams) AND a + per-tensor quantized output is a pure constant transform of the table: each + gathered row is dequantized (group/per-channel) and then requantized + per-tensor. Because the gather commutes with the elementwise (de)quantize, + we apply that transform to the whole table at compile time and replace the op + with a plain aten.embedding over the resulting int8 table -- dropping the + scale/zero_point constants and all runtime (de)quant work. The result is + bit-exact with what the op would compute, and the gathered int8 values feed + downstream ops directly (their input qparams already equal the former output + qparams). + + Embeddings whose output stays float (no per-tensor output quant) are left + unchanged. + """ + weight_qp = get_qparams_from_node(node, "weight") + out_qp = get_qparams_from_node(node, "out") + if weight_qp is None or out_qp is None or not out_qp.is_per_tensor(): + return False + + weight_node = get_arg(node, "weight", fx.Node) + indices = get_arg(node, "indices", fx.Node) + + table = get_constant(ep, weight_node) + weight_scale = get_constant(ep, weight_qp.scale) + weight_zero_point = get_constant(ep, weight_qp.zero_point) + if table is None or weight_scale is None or weight_zero_point is None: + return False + + # Per-tensor output qparams are 0-dim lifted constants. + out_scale = get_scale(ep, out_qp) + out_zero_point = get_zero_point(ep, out_qp) + + # Recompute the table exactly as the op would at runtime: dequantize the + # (group/per-channel) table to float, then quantize it per-tensor. + table_qp = QuantParamsStruct( + scale=weight_scale, + zero_point=weight_zero_point, + dtype=weight_qp.dtype, + quant_min=weight_qp.quant_min, + quant_max=weight_qp.quant_max, + ) + output_qp = QuantParamsStruct( + scale=torch.tensor(out_scale), + zero_point=torch.tensor(out_zero_point), + dtype=out_qp.dtype, + quant_min=out_qp.quant_min, + quant_max=out_qp.quant_max, + ) + prequantized = output_qp.quantize(table_qp.dequantize(table)) + # Same shape/dtype (int8 table) as the original, so the placeholder's meta + # stays valid -- only the backing values change. + set_param_tensor(ep, weight_node, prequantized) + + graph = node.graph + with graph.inserting_before(node): + gather = graph.call_function( + exir_ops.edge.aten.embedding.default, + args=(weight_node, indices), + ) + gather.meta = node.meta + gather.meta["val"] = compute_meta_val(gather) + node.replace_all_uses_with(gather) + graph.erase_node(node) + return True + + +class PrequantizeEmbedding(ExportedProgramPassBase): + """Replace a quantized-table fused_quant.embedding with a plain int8 gather. + + Runs before coloring. For an embedding whose table is quantized and whose + output is per-tensor quantized (after QuantAbsorptionPass folds the consumer's + quantize into the op's output qparams), the table is pre-quantized at compile + time and the op becomes aten.embedding over the resulting int8 table. The + now-dead scale/zero_point constants are removed by constant_prop_pass. + + Embeddings that keep a float output are untouched, so a backend can still + lower them as a quantized-table gather. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + for node in graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.embedding.default + ): + if _try_prequantize_embedding(exported_program, node): + modified = True + + if modified: + exported_program = constant_prop_pass(exported_program) + + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/optimization_passes/quant_absorption.py b/backends/fused_quant/optimization_passes/quant_absorption.py new file mode 100644 index 00000000000..54025062273 --- /dev/null +++ b/backends/fused_quant/optimization_passes/quant_absorption.py @@ -0,0 +1,624 @@ +# 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. + + +import logging +import math +import operator +from typing import Optional + +import torch +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_from_node, + get_scale, + get_zero_point, + split_fused_arg_names, +) +from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg +from executorch.backends.transforms.utils import add_constant +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +logger: logging.Logger = logging.getLogger(__name__) + +_DEQUANT_PER_TENSOR: EdgeOpOverload = ( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default +) +_QUANT_PER_TENSOR: EdgeOpOverload = ( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +) + + +def _is_fused_quant_op(node: fx.Node) -> bool: + return ( + node.op == "call_function" + and isinstance(node.target, EdgeOpOverload) + and node.target.namespace == "fused_quant" + ) + + +def _has_per_tensor_out_qparams(node: fx.Node) -> bool: + """Check if a fused_quant op has per-tensor (non-None) output quantization. + Note: This is just to simplify this pass since this is the most common case, + but it is possible that the output uses per-channel quantization, and we can + still fuse downstream dq/quant ops into the fused_quant op. + """ + out_scale = get_arg(node, "out_scale", Optional[fx.Node]) + if not isinstance(out_scale, fx.Node): + return False + # Per-tensor output is a singleton scale (granularity is encoded by shape). + return out_scale.meta["val"].numel() == 1 + + +def _has_unquantized_out(node: fx.Node) -> bool: + """Check if a fused_quant op leaves its output unquantized (out_scale is None). + + Most fused ops fold their output quantize at fusion time, but some (notably + fused_quant.embedding) deliberately keep a null output block -- their result + is float and a downstream consumer quantizes it. + """ + return get_arg(node, "out_scale", Optional[fx.Node]) is None + + +# Pure data-movement ops that are transparent to the requantization being +# absorbed. This is only safe because the pass requires *per-tensor* input or +# output quantization: a scalar scale/zero_point applies uniformly to every +# element regardless of layout, so reshapes and permutes leave it untouched. We +# would need to be more careful if we wanted to support fusion of different +# types of quantization. +_PASSTHROUGH_TARGETS: set[EdgeOpOverload] = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.view_copy.default, +} + + +def _find_dequant_through_passthrough( + node: fx.Node, +) -> tuple[fx.Node | None, list[fx.Node]]: + """Walk single-user passthrough ops from node to find a dequant. + + Returns (dequant_node, passthrough_chain) if found, else (None, []). + The passthrough_chain is the ordered list of passthrough nodes between the + fused op and the dequant (node is the first element, the immediate + predecessor of the dequant is the last). After absorption the quant's users + consume passthrough_chain[-1], and every node in the chain carries the + fused op's requantized output, so their meta['val'] dtype must follow it. + """ + chain: list[fx.Node] = [] + cursor = node + while len(cursor.users) == 1: + child = next(iter(cursor.users.keys())) + chain.append(cursor) + if child.target == _DEQUANT_PER_TENSOR: + return child, chain + if child.target not in _PASSTHROUGH_TARGETS: + return None, [] + cursor = child + return None, [] + + +def _find_quant_through_input_passthrough( + node: fx.Node, +) -> tuple[fx.Node | None, list[fx.Node]]: + """Walk backward through single-user passthrough ops to find a quantize. + + ``node`` is a fused consumer's input. The returned chain is ordered from the + consumer toward the quantize, so ``chain[-1]`` is the passthrough whose input + must be rewired when the quantize is absorbed. + """ + chain: list[fx.Node] = [] + cursor = node + while cursor.target in _PASSTHROUGH_TARGETS: + if len(cursor.users) != 1: + return None, [] + chain.append(cursor) + predecessor = cursor.args[0] + if not isinstance(predecessor, fx.Node): + return None, [] + cursor = predecessor + return (cursor, chain) if cursor.target == _QUANT_PER_TENSOR else (None, []) + + +def _output_zero_node(node: fx.Node) -> fx.Node | None: + """The node carrying the fused op's output-0 value. + + Single-output fused ops produce their value directly, so it is ``node`` + itself. Multi-output ops (e.g. ``fused_quant.native_layer_norm`` returning + ``(out, mean, rstd)``) expose each output through a ``getitem`` accessor; + only output 0 is quantized (its qparams live in the op's ``out_scale`` / + ``out_zero_point`` block), so its value flows through ``getitem(node, 0)``. + Returns that accessor for a multi-output op, ``node`` for a single-output op, + or None when a multi-output op's output 0 isn't a single clean getitem + (nothing to absorb through). + """ + if not isinstance(node.meta.get("val"), (tuple, list)): + return node + getitems = [ + u for u in node.users if u.target == operator.getitem and u.args[1] == 0 + ] + return getitems[0] if len(getitems) == 1 else None + + +def _compose_output_qparams( + exported_program: ExportedProgram, + producer: fx.Node, + dequantize: fx.Node, + quantize: fx.Node, +) -> tuple[float, int] | None: + """Compose Q(producer) → DQ → Q into one zero-error affine quantizer. + + Let ``x`` be the producer's floating-point result, ``q_A`` its original + quantized output, and ``y`` the value after the possibly mismatched + dequantize: + + q_A = x / s_A + z_A + y = (q_A - z_D) * s_D + q_B = y / s_B + z_B + + Ignoring rounding and saturation, choosing + + s_R = s_A * s_B / s_D + z_R = z_B + (z_A - z_D) * s_D / s_B + + makes the replacement producer output ``q_R = Q(R, x)`` equal ``q_B``. + A standard affine quantizer can represent this composition only when + ``z_R`` is integral. + + R includes the effects of A, D, and B. If the dequantize has surviving + float users, DQ(B) undoes just the B part and recovers ``y``. + """ + producer_qparams = get_qparams_from_node(producer, "out") + assert producer_qparams is not None and producer_qparams.is_per_tensor() + if ( + producer_qparams.dtype != get_arg(dequantize, "dtype", torch.dtype) + or producer_qparams.quant_min != get_arg(dequantize, "quant_min", int) + or producer_qparams.quant_max != get_arg(dequantize, "quant_max", int) + ): + return None + + producer_scale = get_scale(exported_program, producer_qparams) + dequant_scale = get_arg(dequantize, "scale", float) + output_scale = get_arg(quantize, "scale", float) + if any( + not math.isfinite(scale) or scale <= 0 + for scale in (producer_scale, dequant_scale, output_scale) + ): + return None + + producer_zp = get_zero_point(exported_program, producer_qparams) + dequant_zp = get_arg(dequantize, "zero_point", int) + output_zp = get_arg(quantize, "zero_point", int) + replacement_scale = producer_scale * output_scale / dequant_scale + replacement_zp_float = ( + output_zp + (producer_zp - dequant_zp) * dequant_scale / output_scale + ) + replacement_zp = round(replacement_zp_float) + if ( + not math.isfinite(replacement_scale) + or replacement_scale <= 0 + or not math.isclose( + replacement_zp_float, replacement_zp, rel_tol=0.0, abs_tol=1e-6 + ) + or not get_arg(quantize, "quant_min", int) + <= replacement_zp + <= get_arg(quantize, "quant_max", int) + ): + return None + return replacement_scale, replacement_zp + + +def _compose_input_qparams( + exported_program: ExportedProgram, + dequantize: fx.Node, + quantize: fx.Node, + consumer: fx.Node, + input_name: str, +) -> tuple[float, int] | None: + """Compose DQ(D) → Q(B) → consumer DQ(C) into consumer DQ(R). + + Ignoring rounding and saturation, choosing + + s_R = s_D * s_C / s_B + z_R = z_D + (z_C - z_B) * s_B / s_D + + preserves the value seen by the consumer. A standard affine quantizer can + represent the composition only when ``z_R`` is integral. + """ + consumer_qparams = get_qparams_from_node(consumer, input_name) + assert consumer_qparams is not None and consumer_qparams.is_per_tensor() + if ( + get_arg(quantize, "quant_min", int) != consumer_qparams.quant_min + or get_arg(quantize, "quant_max", int) != consumer_qparams.quant_max + ): + return None + + dequant_scale = get_arg(dequantize, "scale", float) + quant_scale = get_arg(quantize, "scale", float) + consumer_scale = get_scale(exported_program, consumer_qparams) + if any( + not math.isfinite(scale) or scale <= 0 + for scale in (dequant_scale, quant_scale, consumer_scale) + ): + return None + + dequant_zp = get_arg(dequantize, "zero_point", int) + quant_zp = get_arg(quantize, "zero_point", int) + consumer_zp = get_zero_point(exported_program, consumer_qparams) + replacement_scale = dequant_scale * consumer_scale / quant_scale + replacement_zp_float = ( + dequant_zp + (consumer_zp - quant_zp) * quant_scale / dequant_scale + ) + replacement_zp = round(replacement_zp_float) + if ( + not math.isfinite(replacement_scale) + or replacement_scale <= 0 + or not math.isclose( + replacement_zp_float, replacement_zp, rel_tol=0.0, abs_tol=1e-6 + ) + or not get_arg(dequantize, "quant_min", int) + <= replacement_zp + <= get_arg(dequantize, "quant_max", int) + ): + return None + return replacement_scale, replacement_zp + + +def _try_absorb_output( + ep: ExportedProgram, node: fx.Node, graph: fx.Graph, absorb_with_fork: bool +) -> bool: + """Try to absorb a downstream dequant→quant into node's out qparams. + + Pattern 1 (single user): + fused_quant Q(A) → [getitem 0] → [passthrough]* → DQ(D) → Q(B) → ... + Result: fused_quant Q(compose(A, D, B)) → [getitem 0] → [passthrough]* → ... + + Pattern 2 (forked dequant, requires absorb_with_fork=True): + fused_quant Q(A) → [getitem 0] → [passthrough]* → DQ(D) ┬→ Q(B) → ... + └→ other float users + + Result: fused_quant Q(R) → [getitem 0] → [passthrough]* ┬→ former Q(B) users + └→ DQ(B) → other float users + where R = compose(A, D, B) + + When A == D, compose(A, D, B) == B, which is the ordinary removal of an + identity quantize/dequantize boundary. A mismatch between A and D encodes + a real affine operation and must be retained in the composed qparams. + + In the forked case, R includes A, D, and B. The former quantized users + consume the composed integer output directly. DQ(B) on the float branch + undoes just the B part, leaving the original A-to-D effect intact. DQ(R) + would undo the entire composition and incorrectly erase that effect. When + A == D, R == B, so the distinction is invisible. + + ``out_node`` is the node carrying output 0: the fused op itself for a + single-output op, or ``getitem(node, 0)`` for a multi-output op (only output 0 + is quantized, and its qparams are the op's ``out_scale`` block). The walk to + the dequant, and all retyping, run from ``out_node`` -- the op node itself is + left alone for a multi-output op (its ``meta['val']`` is the output tuple). + + [passthrough]* is an optional chain of single-user data-movement ops + (e.g. permute_copy) that are transparent to quantization parameters. The + chain's structure is unaffected, but because absorbing the quant can change + the output dtype (e.g. int8 → uint8), each passthrough node's meta['val'] is + retyped to the absorbed dtype. + """ + out_node = _output_zero_node(node) + if out_node is None or len(out_node.users) != 1: + return False + + user = next(iter(out_node.users.keys())) + + # passthrough_chain is the (possibly empty) list of data-movement ops + # between the fused op and the dequant. Every node in it carries the fused + # op's quantized output, so its dtype must follow the new output qparams. + passthrough_chain: list[fx.Node] = [] + + # Direct case: fused_quant [→ getitem 0] → dequant + if user.target == _DEQUANT_PER_TENSOR: + dequant_node = user + # The node whose quantized output replaces the quant + quantized_source = out_node + elif user.target in _PASSTHROUGH_TARGETS: + # Passthrough case: walk through passthrough ops to find dequant + dequant_node, passthrough_chain = _find_dequant_through_passthrough(user) + if dequant_node is None: + return False + quantized_source = passthrough_chain[-1] + else: + return False + + quant_nodes = [u for u in dequant_node.users if u.target == _QUANT_PER_TENSOR] + if len(quant_nodes) != 1: + return False + quant_node = quant_nodes[0] + + has_other_users = len(dequant_node.users) > 1 + + if has_other_users and not absorb_with_fork: + return False + + replacement_qparams = _compose_output_qparams(ep, node, dequant_node, quant_node) + if replacement_qparams is None: + return False + replacement_scale, replacement_zp = replacement_qparams + + downstream_scale = get_arg(quant_node, "scale", float) + downstream_zp = get_arg(quant_node, "zero_point", int) + # The downstream quant may target a different dtype than the fused op's + # original output (e.g. int8 → uint8). Absorbing it means the fused op now + # produces that dtype, so it must propagate to the fused op's out qparams + # and to every node carrying its output (the passthrough chain, and the + # kept dequant's input in the forked case). + new_dtype = get_arg(quant_node, "dtype", torch.dtype) + new_quant_min = get_arg(quant_node, "quant_min", int) + new_quant_max = get_arg(quant_node, "quant_max", int) + + _set_qparams( + ep, + node, + "out", + scale=replacement_scale, + zero_point=replacement_zp, + dtype=new_dtype, + quant_min=new_quant_min, + quant_max=new_quant_max, + ) + + quant_node.replace_all_uses_with(quantized_source) + graph.erase_node(quant_node) + + if has_other_users: + # R includes A, D, and B. DQ(B) undoes just the B part for float users. + set_arg(dequant_node, "scale", downstream_scale) + set_arg(dequant_node, "zero_point", downstream_zp) + set_arg(dequant_node, "dtype", new_dtype) + set_arg(dequant_node, "quant_min", new_quant_min) + set_arg(dequant_node, "quant_max", new_quant_max) + else: + graph.erase_node(dequant_node) + + # out_node (the fused op for single-output, or its getitem-0 for + # multi-output) and every passthrough node now carry the requantized value + # flowing into the quant's former users, so retype their meta['val']. The + # multi-output op node itself keeps its output-tuple val untouched. + for retyped in (out_node, *passthrough_chain): + val = retyped.meta.get("val") + if isinstance(val, torch.Tensor): + retyped.meta["val"] = val.to(new_dtype) + + return True + + +def _try_absorb_quant_into_unquantized_out( + exported_program: ExportedProgram, node: fx.Node, graph: fx.Graph +) -> bool: + """Absorb a bare quantize that follows a fused op whose output is unquantized. + + Pattern: + fused_quant(float out) → quant(B) → ... + Result: fused_quant(out qparams = B) → ... + + Some fused ops (notably fused_quant.embedding) leave their output unquantized + at fusion time -- the produced value is float and a downstream consumer + quantizes it. When that consumer's quantize is the op's sole user, fold it + into the op's (previously null) output qparams so the op emits the quantized + value directly and the redundant quantize is removed. The quantize may itself + fan out to several consumers; they all read the same quantized value and are + simply repointed at the fused op. + + Requires the fused op to have a single user (the quantize): an op feeding both + the quantize and a separate float consumer would need to keep a float path, + which this case does not handle. + """ + if len(node.users) != 1: + return False + quant_node = next(iter(node.users.keys())) + if quant_node.target != _QUANT_PER_TENSOR: + return False + + new_scale = get_arg(quant_node, "scale", float) + new_zp = get_arg(quant_node, "zero_point", int) + new_dtype = get_arg(quant_node, "dtype", torch.dtype) + new_quant_min = get_arg(quant_node, "quant_min", int) + new_quant_max = get_arg(quant_node, "quant_max", int) + + _set_qparams( + exported_program, + node, + "out", + scale=new_scale, + zero_point=new_zp, + dtype=new_dtype, + quant_min=new_quant_min, + quant_max=new_quant_max, + ) + + quant_node.replace_all_uses_with(node) + graph.erase_node(quant_node) + + # The fused op now carries the quantized value, so its dtype follows the + # absorbed quantize (e.g. float32 → int8). + val = node.meta.get("val") + if isinstance(val, torch.Tensor): + node.meta["val"] = val.to(new_dtype) + + return True + + +def _set_qparams( + ep: ExportedProgram, + node: fx.Node, + prefix: str, + *, + scale: float, + zero_point: int, + dtype: torch.dtype, + quant_min: int, + quant_max: int, +) -> None: + """Replace a fused_quant op's qparams with new lifted constants. + + Lifted as ``CONSTANT_TENSOR`` to match how ``FusedQuantFusion`` lifts per-tensor + qparams: its ``_scale``/``_zero_point`` module attributes are plain + ``setattr`` tensors, which export classifies as lifted tensor constants (the + graph's ``c__scale*`` placeholders confirm this -- ``c_`` is the + CONSTANT_TENSOR prefix). Using a different kind here (e.g. BUFFER) would give + the new node a ``b__scale`` name whose fqn ``_scale`` still collides with the + fusion's ``ep.constants`` entry -- a cross-kind fqn clash node-name + uniquification can't catch, tripping add_constant's assert. + """ + new_scale_node = add_constant( + ep, + "_scale", + torch.tensor(scale, dtype=torch.float32), + node, + InputKind.CONSTANT_TENSOR, + ) + new_zp_node = add_constant( + ep, + "_zero_point", + torch.tensor(zero_point, dtype=torch.int64), + node, + InputKind.CONSTANT_TENSOR, + ) + + set_arg(node, f"{prefix}_scale", new_scale_node) + set_arg(node, f"{prefix}_zero_point", new_zp_node) + set_arg(node, f"{prefix}_dtype", dtype) + set_arg(node, f"{prefix}_quant_min", quant_min) + set_arg(node, f"{prefix}_quant_max", quant_max) + + +def _try_absorb_input( + exported_program: ExportedProgram, + consumer: fx.Node, + input_name: str, + graph: fx.Graph, +) -> bool: + """Absorb an upstream dequantize→quantize into fused input qparams.""" + consumer_qparams = get_qparams_from_node(consumer, input_name) + if consumer_qparams is None or not consumer_qparams.is_per_tensor(): + return False + + consumer_input = get_arg(consumer, input_name, Optional[fx.Node]) + if consumer_input is None: + return False + quantize, passthrough_chain = _find_quant_through_input_passthrough(consumer_input) + if quantize is None or len(quantize.users) != 1: + return False + + dequantize = get_arg(quantize, "input", Optional[fx.Node]) + if dequantize is None or dequantize.target != _DEQUANT_PER_TENSOR: + return False + + replacement_qparams = _compose_input_qparams( + exported_program, dequantize, quantize, consumer, input_name + ) + if replacement_qparams is None: + return False + replacement_scale, replacement_zp = replacement_qparams + + source = get_arg(dequantize, "input", Optional[fx.Node]) + if source is None: + return False + + if passthrough_chain: + passthrough_chain[-1].replace_input_with(quantize, source) + source_dtype = source.meta["val"].dtype + for passthrough in passthrough_chain: + passthrough.meta["val"] = passthrough.meta["val"].to(source_dtype) + else: + set_arg(consumer, input_name, source) + _set_qparams( + exported_program, + consumer, + input_name, + scale=replacement_scale, + zero_point=replacement_zp, + # Input dtype is the dequantized compute dtype; the source tensor carries + # its integer storage dtype. + dtype=dequantize.meta["val"].dtype, + quant_min=get_arg(dequantize, "quant_min", int), + quant_max=get_arg(dequantize, "quant_max", int), + ) + if not quantize.users: + graph.erase_node(quantize) + if not dequantize.users: + graph.erase_node(dequantize) + return True + + +def _try_absorb_inputs( + exported_program: ExportedProgram, + consumer: fx.Node, + graph: fx.Graph, +) -> bool: + assert isinstance(consumer.target, EdgeOpOverload) + input_names, _ = split_fused_arg_names(consumer.target) + modified = False + for input_name in input_names: + modified |= _try_absorb_input(exported_program, consumer, input_name, graph) + return modified + + +class QuantAbsorptionPass(ExportedProgramPassBase): + """Absorb adjacent quantization boundaries into fused_quant operators. + + On the output side, when a fused_quant op with per-tensor output quantization + feeds into a dequant→quant pair, this pass composes the producer, dequantize, + and downstream quantize qparams into the fused op's output qparams. When the + producer and dequantize qparams match, this reduces to the downstream target. + + On the input side, when a dequant→quant pair feeds any fused_quant input with + per-tensor qparams, this pass composes the dequantize, quantize, and consumer + input qparams, then bypasses the pair. This handles graph inputs, which have + no fused producer into which the quantization can be absorbed. + + It also absorbs a *bare* downstream quantize into a fused op whose output is + currently unquantized (null out qparams), turning that op into a directly + quantizing one. This is what folds the quantize that a consumer places on a + fused_quant.embedding's float gather into the embedding's output qparams. + + Args: + absorb_with_fork: When True, also absorb dequant→quant pairs where + the dequant has additional users (a "forked" dequant). The dequant + is kept for the other users but its qparams are updated to match + the new output qparams. This changes (not necessarily worsens) the + numerics of the float path through the dequant, but should improve + numerics on the int path since it removes a redundant dq-q pair. + """ + + def __init__(self, absorb_with_fork: bool = False) -> None: + super().__init__() + self._absorb_with_fork = absorb_with_fork + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph + modified = False + + for node in list(graph.nodes): + if not _is_fused_quant_op(node): + continue + if _has_per_tensor_out_qparams(node): + modified |= _try_absorb_output( + exported_program, node, graph, self._absorb_with_fork + ) + elif _has_unquantized_out(node): + modified |= _try_absorb_quant_into_unquantized_out( + exported_program, node, graph + ) + modified |= _try_absorb_inputs(exported_program, node, graph) + + if modified: + constant_prop_pass(exported_program) + + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/optimization_passes/replace_dequant_quant_with_requantize.py b/backends/fused_quant/optimization_passes/replace_dequant_quant_with_requantize.py new file mode 100644 index 00000000000..3b3d6169642 --- /dev/null +++ b/backends/fused_quant/optimization_passes/replace_dequant_quant_with_requantize.py @@ -0,0 +1,96 @@ +# 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. + + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import add_constant +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from torch import fx +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +def _lift_qparams( + exported_program: ExportedProgram, + node: fx.Node, + name: str, +) -> tuple[fx.Node, fx.Node, torch.dtype, int, int]: + scale = add_constant( + exported_program, + f"_{name}_scale", + torch.tensor(get_arg(node, "scale", float), dtype=torch.float32), + node, + InputKind.CONSTANT_TENSOR, + ) + zero_point = add_constant( + exported_program, + f"_{name}_zero_point", + torch.tensor(get_arg(node, "zero_point", int), dtype=torch.int64), + node, + InputKind.CONSTANT_TENSOR, + ) + return ( + scale, + zero_point, + get_arg(node, "dtype", torch.dtype), + get_arg(node, "quant_min", int), + get_arg(node, "quant_max", int), + ) + + +class ReplaceDequantQuantWithRequantize(ExportedProgramPassBase): + """Replace direct per-tensor dequantize-quantize branches with requantize.""" + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + for dequant in dequant_nodes: + quant_nodes = [ + user + for user in dequant.users + if user.target + == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + ] + if not quant_nodes: + continue + + inp_scale, inp_zero_point, _, inp_quant_min, inp_quant_max = _lift_qparams( + exported_program, dequant, "requantize_inp" + ) + inp_qparams = ( + inp_scale, + inp_zero_point, + dequant.meta["val"].dtype, + inp_quant_min, + inp_quant_max, + ) + + for quant in quant_nodes: + out_qparams = _lift_qparams(exported_program, quant, "requantize_out") + with graph.inserting_before(quant): + requantize = graph.call_function( + exir_ops.edge.fused_quant.requantize.default, + args=(dequant.args[0], *inp_qparams, *out_qparams), + ) + requantize.meta = quant.meta.copy() + + quant.replace_all_uses_with(requantize) + graph.erase_node(quant) + if not dequant.users: + graph.erase_node(dequant) + modified = True + + if modified: + exported_program.graph_module.recompile() + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/optimization_passes/sink_constant_cat.py b/backends/fused_quant/optimization_passes/sink_constant_cat.py new file mode 100644 index 00000000000..e77a0ff4855 --- /dev/null +++ b/backends/fused_quant/optimization_passes/sink_constant_cat.py @@ -0,0 +1,446 @@ +# 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. + + +from __future__ import annotations + +from dataclasses import dataclass + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch._dynamo.utils import detect_fake_mode +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind +from torch.fx.passes.fake_tensor_prop import FakeTensorProp + + +def _propagate_fake_tensors(graph_module: fx.GraphModule) -> None: + """Recompute meta["val"] for all nodes via FakeTensorProp.""" + inputs = [] + for node in graph_module.graph.find_nodes(op="placeholder"): + val = node.meta.get("val") + if val is not None: + inputs.append(val) + else: + inputs.append(node) + + fake_mode = detect_fake_mode(inputs) + FakeTensorProp(graph_module, mode=fake_mode).propagate_dont_convert_inputs(*inputs) + + +@dataclass +class CatDimTracker: + """Tracks the cat dimension through a chain of downstream ops. + + After a cat on dimension `dim`, this tracks where that dimension ends up + as ops like permute, view, unsqueeze, squeeze, and slice are applied. + If a reshape merges or splits the cat dimension, the tracker is invalidated. + """ + + dim: int + valid: bool = True + + def update_permute(self, dims: list[int]) -> None: + if not self.valid: + return + try: + self.dim = dims.index(self.dim) + except ValueError: + self.valid = False + + def update_view(self, old_shape: list[int], new_shape: list[int]) -> None: + if not self.valid: + return + + i = 0 # index into new_shape + j = 0 # index into old_shape + new_cat_dim = None + + while i < len(new_shape) and j < len(old_shape): + if new_shape[i] == old_shape[j]: + if j == self.dim: + new_cat_dim = i + break + i += 1 + j += 1 + elif new_shape[i] == 1: + # Inserted unary dim in new_shape — skip it. + i += 1 + else: + # Sizes don't match — accumulate product on the old side. + # If the cat dim is consumed in this merge, bail. + product = old_shape[j] + j += 1 + while product != new_shape[i]: + if j >= len(old_shape): + self.valid = False + return + if j == self.dim: + self.valid = False + return + product *= old_shape[j] + j += 1 + i += 1 + + if new_cat_dim is None: + self.valid = False + else: + self.dim = new_cat_dim + + def update_unsqueeze(self, unsqueeze_dim: int, ndim: int) -> None: + if not self.valid: + return + if unsqueeze_dim < 0: + unsqueeze_dim = ndim + 1 + unsqueeze_dim + if unsqueeze_dim <= self.dim: + self.dim += 1 + + def update_squeeze(self, squeeze_dims: list[int]) -> None: + if not self.valid: + return + if self.dim in squeeze_dims: + self.valid = False + return + shift = sum(1 for d in squeeze_dims if d < self.dim) + self.dim -= shift + + +_UNARY_ELEMENTWISE_OPS: frozenset[torch._ops.OpOverload] = frozenset( + { + exir_ops.edge.aten.relu.default, + exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.hardswish.default, + exir_ops.edge.aten.clamp.default, + exir_ops.edge.fused_quant.relu.default, + exir_ops.edge.fused_quant.hardswish.default, + } +) + + +def _track_dim_through_op( + node: fx.Node, + tracker: CatDimTracker, +) -> bool: + """Update the tracker for a single op. Returns True if the op is safe.""" + if not tracker.valid: + return False + + target = node.target + + if target in _UNARY_ELEMENTWISE_OPS: + return True + + if target == exir_ops.edge.aten.slice_copy.Tensor: + slice_dim = get_arg(node, "dim", int) + input_node = node.args[0] + assert isinstance(input_node, fx.Node) + ndim = input_node.meta["val"].ndim + if slice_dim < 0: + slice_dim = ndim + slice_dim + if slice_dim == tracker.dim: + return False + return True + + if target == exir_ops.edge.aten.permute_copy.default: + dims = get_arg(node, "dims", list[int]) + tracker.update_permute(dims) + return tracker.valid + + if target == exir_ops.edge.aten.unsqueeze_copy.default: + dim = get_arg(node, "dim", int) + input_node = get_arg(node, "input", fx.Node) + tracker.update_unsqueeze(dim, input_node.meta["val"].ndim) + return tracker.valid + + if target == exir_ops.edge.aten.squeeze_copy.dims: + dims = get_arg(node, "dims", list[int]) + tracker.update_squeeze(dims) + return tracker.valid + + if target == exir_ops.edge.aten.view_copy.default: + input_node = get_arg(node, "input", fx.Node) + old_shape = list(input_node.meta["val"].shape) + new_shape = list(node.meta["val"].shape) + tracker.update_view(old_shape, new_shape) + return tracker.valid + + # As we support more data movement patterns, we can add them here. + return False + + +@dataclass +class SinkPath: + """A single user path that the cat can be sunk through.""" + + chain: list[fx.Node] + cat_dim_per_node: dict[fx.Node, int] + new_cat_dim: int + + +def _find_sink_paths( + cat_node: fx.Node, + cat_dim: int, +) -> list[SinkPath]: + """Find sink paths for each user of the cat node. + + Walks forward from each user along a single-user chain, collecting + ops that the cat can be sunk through. Returns a SinkPath per user + that has at least one sinkable op. + """ + paths: list[SinkPath] = [] + for user in list(cat_node.users.keys()): + tracker = CatDimTracker(dim=cat_dim) + cursor = user + chain: list[fx.Node] = [] + cat_dim_per_node: dict[fx.Node, int] = {} + + while True: + if cursor.op != "call_function": + break + cat_dim_per_node[cursor] = tracker.dim + if not _track_dim_through_op(cursor, tracker): + break + chain.append(cursor) + + # While we support sinking cats if there are multiple users, we + # stop in the chain if we hit a fork. We can run this pass iteratively + # to address multiple forks if they appear. + if len(cursor.users) != 1: + break + cursor = next(iter(cursor.users.keys())) + + if chain: + paths.append( + SinkPath( + chain=chain, + cat_dim_per_node=cat_dim_per_node, + new_cat_dim=tracker.dim, + ) + ) + return paths + + +def _is_constant_node(node: fx.Node, exported_program: ExportedProgram) -> bool: + for spec in exported_program.graph_signature.input_specs: + if spec.arg.name == node.name: + if spec.kind in (InputKind.PARAMETER, InputKind.CONSTANT_TENSOR): + return True + if spec.kind == InputKind.BUFFER: + assert spec.target is not None + return ( + spec.target + not in exported_program.graph_signature.buffers_to_mutate.values() + ) + return False + return False + + +def _adjust_view_shape( + original_shape: list[int], + cat_dim: int, + new_input: fx.Node, +) -> list[int]: + """Compute the correct view shape for a replicated branch. + + The original view shape was computed for the full cat output. The + replicated branch has a different size on the cat dim, so we scale + the cat-dim component proportionally. + """ + original_input_shape = new_input.meta["val"].shape + new_cat_dim_size = original_input_shape[cat_dim] + + adjusted = list(original_shape) + adjusted[cat_dim] = new_cat_dim_size + return adjusted + + +def _replicate_chain( + chain: list[fx.Node], + node_map: dict[fx.Node, fx.Node], + cat_dim_per_node: dict[fx.Node, int], + graph: fx.Graph, + target_node: fx.Node, +) -> None: + """Replicate a chain of ops, fixing view shapes for the new input size.""" + for op_node in chain: + new_args = _replace_nodes_in_args(op_node.args, node_map) + + if op_node.target == exir_ops.edge.aten.view_copy.default: + cat_dim = cat_dim_per_node[op_node] + input_in_map = op_node.args[0] + assert isinstance(input_in_map, fx.Node) + new_input = node_map[input_in_map] + original_output_shape = list(op_node.meta["val"].shape) + adjusted_shape = _adjust_view_shape( + original_output_shape, cat_dim, new_input + ) + new_args = (new_args[0], adjusted_shape) + + with graph.inserting_before(target_node): + new_node = graph.call_function( + op_node.target, # pyre-ignore[6] + args=new_args, # pyre-ignore[6] + kwargs=op_node.kwargs, + ) + new_node.meta = op_node.meta.copy() + node_map[op_node] = new_node + + +def _replace_nodes_in_args( + args: tuple[object, ...], + node_map: dict[fx.Node, fx.Node], +) -> tuple[object, ...]: + result: list[object] = [] + for arg in args: + if isinstance(arg, fx.Node) and arg in node_map: + result.append(node_map[arg]) + elif isinstance(arg, (list, tuple)): + replaced = [ + node_map[a] if (isinstance(a, fx.Node) and a in node_map) else a + for a in arg + ] + result.append(type(arg)(replaced)) + else: + result.append(arg) + return tuple(result) + + +def _sink_cat_along_paths( + cat_node: fx.Node, + paths: list[SinkPath], + constant_inputs: list[fx.Node], + activation_inputs: list[fx.Node], + graph: fx.Graph, +) -> None: + """Push cat past downstream ops, splitting only constant inputs out. + + Only the constant inputs get their own replicated chain. Activation + inputs stay together in the original cat (or a reduced cat if some + constants were removed). At the sink point, a new cat merges the + activation path with each constant path. + + Before: + [act_a, act_b, const] -> cat -> op1 -> op2 -> consumers + + After (before constant folding): + [act_a, act_b] -> cat -> op1 -> op2 -> ─┐ + ├─> cat -> consumers + const -> op1 -> op2 -> ─────────────────┘ + + After constant folding: + [act_a, act_b] -> cat -> op1 -> op2 -> ─┐ + ├─> cat -> consumers + [folded_const] ────────────────────────┘ + """ + cat_dim = get_arg(cat_node, "dim", int) + all_chain_nodes: list[fx.Node] = [] + + for path in paths: + chain = path.chain + target_node = chain[-1] + all_chain_nodes.extend(chain) + + # The activation path: keep activations in a (possibly smaller) cat, + # then run the chain on that cat's output. + if len(activation_inputs) == 1: + act_head: fx.Node = activation_inputs[0] + else: + with graph.inserting_before(cat_node): + act_head = graph.call_function( + exir_ops.edge.aten.cat.default, + args=(activation_inputs, cat_dim), + ) + + act_node_map: dict[fx.Node, fx.Node] = {cat_node: act_head} + _replicate_chain(chain, act_node_map, path.cat_dim_per_node, graph, target_node) + + sink_outputs: list[fx.Node] = [act_node_map[chain[-1]]] + + # Each constant input gets its own replicated chain. + for const_inp in constant_inputs: + const_node_map: dict[fx.Node, fx.Node] = {cat_node: const_inp} + _replicate_chain( + chain, const_node_map, path.cat_dim_per_node, graph, target_node + ) + sink_outputs.append(const_node_map[chain[-1]]) + + with graph.inserting_before(target_node): + new_cat = graph.call_function( + exir_ops.edge.aten.cat.default, + args=(sink_outputs, path.new_cat_dim), + ) + target_node.replace_all_uses_with(new_cat) + + for op_node in reversed(all_chain_nodes): + if len(op_node.users) == 0: + graph.erase_node(op_node) + if len(cat_node.users) == 0: + graph.erase_node(cat_node) + + +class SinkConstantCat(ExportedProgramPassBase): + """Sink cat ops past downstream ops to separate constant from activation paths. + + When a cat combines constant and activation inputs, downstream ops + (view, permute, slice, etc.) operate on the full concatenated tensor, + preventing the constant subgraph from being folded. This pass pushes the + cat as far down as possible, applying downstream ops to each cat input + independently. The constant path then becomes a pure constant subgraph + that a subsequent constant propagation pass can fold away entirely, + eliminating redundant runtime compute. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph_module = exported_program.graph_module + graph = graph_module.graph + modified = False + + cat_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.cat.default + ) + + for cat_node in cat_nodes: + tensors = get_arg(cat_node, "tensors", list[fx.Node]) + constant_inputs = [ + t for t in tensors if _is_constant_node(t, exported_program) + ] + if not constant_inputs: + continue + activation_inputs = [ + t for t in tensors if not _is_constant_node(t, exported_program) + ] + + cat_dim = get_arg(cat_node, "dim", int) + paths = _find_sink_paths(cat_node, cat_dim) + if len(paths) != len(list(cat_node.users)): + # We require that for each user we should be able + # to sink the cat node, since if we cannot for + # every user, we still end up + continue + + _sink_cat_along_paths( + cat_node, paths, constant_inputs, activation_inputs, graph + ) + + # Make sure meta['val'] is correct after these types of modifications. + # Shouldn't be too bad in practice, since only updates after we handle + # all users of a cat node. + _propagate_fake_tensors(graph_module) + modified = True + + if modified: + graph_module.recompile() + exported_program = constant_prop_pass(exported_program) + + return ExportedProgramPassResult( + exported_program=exported_program, + modified=modified, + ) diff --git a/backends/fused_quant/optimization_passes/split_linear_at_slices.py b/backends/fused_quant/optimization_passes/split_linear_at_slices.py new file mode 100644 index 00000000000..c4263cd3053 --- /dev/null +++ b/backends/fused_quant/optimization_passes/split_linear_at_slices.py @@ -0,0 +1,346 @@ +# 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. + + +import logging +from typing import NamedTuple, Optional + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg +from executorch.backends.transforms.utils import ( + add_constant, + get_constant, + get_fqn, + get_input_kind, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +logger: logging.Logger = logging.getLogger(__name__) + + +def _is_output_dim_slice( + slice_node: fx.Node, linear_out_features: int +) -> Optional[tuple[int, int, int]]: + """Check if a slice_copy is on the output (last) dimension of a linear. + + Returns (start, end, step) if it's an output-dim slice, None otherwise. + """ + dim = get_arg(slice_node, "dim", int) + start = get_arg(slice_node, "start") + end = get_arg(slice_node, "end") + step = get_arg(slice_node, "step", int) + if start is None: + start = 0 + if end is None: + end = linear_out_features + + assert isinstance(start, int) and isinstance(end, int), ( + f"Symbolic slice bounds not supported: start={start!r}, end={end!r}" + ) + + input_node = slice_node.args[0] + assert isinstance(input_node, fx.Node) + input_shape = input_node.meta["val"].shape + + ndim = len(input_shape) + if dim < 0: + dim = ndim + dim + + if dim != ndim - 1: + return None + + if start < 0: + start = max(0, start + linear_out_features) + else: + start = min(start, linear_out_features) + + if end < 0: + end = max(0, end + linear_out_features) + else: + end = min(end, linear_out_features) + + return (start, end, step) + + +# Tensor args of fused_quant.linear that index the OUTPUT channels on dim 0, so +# each must be sliced to follow an output-dim split. The weight itself is handled +# separately (it is mandatory and always sliced). Note this includes the *bias* +# qparams (bias_scale/bias_zero_point): if a bias is per-output-channel quantized, +# its scale/zp ride on dim 0 just like the weight's and must be sliced too -- +# slicing the bias tensor while leaving full-rank bias qparams would be a silent +# miscompile. Per-tensor / broadcast instances of any of these are reused +# unchanged; see _plan_output_qparam. +# +# The output qparams (out_scale/out_zero_point) are deliberately NOT here: they +# index the output tensor's *last* dim, not dim 0, so they need a different +# slicing axis. They are per-tensor in the current flow; _try_split bails (rather +# than silently mis-slice) if one is ever per-channel -- see _OUTPUT_QPARAMS. +_PER_OUTPUT_DIM0_ARGS: tuple[str, ...] = ( + "bias", + "weight_scale", + "weight_zero_point", + "bias_scale", + "bias_zero_point", +) + +# Output requantization qparams. These index the output's last dim, not dim 0, so +# the dim-0 slicing above does not apply; per-channel output requant is not +# handled by the split, so _try_split bails when one of these is non-scalar. +_OUTPUT_QPARAMS: tuple[str, ...] = ("out_scale", "out_zero_point") + + +# A planned per-output tensor arg (bias / weight scale / zero_point / bias +# scale / zero_point): (do_slice, tensor, kind). +# do_slice False -> reuse the original node unchanged (it broadcasts over every +# output slice); tensor/kind are None. +# do_slice True -> slice `tensor` along dim 0 per output slice, re-adding the +# result as a constant of input kind `kind`. +_QParamPlan = tuple[bool, Optional[torch.Tensor], Optional[InputKind]] + + +def _plan_output_qparam( + ep: ExportedProgram, + qparam_node: Optional[fx.Node], + out_features: int, +) -> Optional[_QParamPlan]: + """Decide -- WITHOUT mutating the graph -- how a per-output tensor arg (bias / + weight scale / weight zero_point) is carried onto each split linear. + + Returns: + - ``(False, None, None)`` when there is nothing to slice: the arg is absent, + a scalar, or broadcasts over the output dim (its leading dim is not the + output channels). Every split linear then reuses the original node, which + broadcasts correctly over any output slice. + - ``(True, tensor, kind)`` when the arg varies along the output dim and its + constant value is available, so it can be sliced per output slice. + - ``None`` to signal the whole split must be ABORTED: the arg varies along + the output dim but its constant value is unavailable (e.g. a per-channel + scale still behind an unfolded ``view_copy``). Slicing it correctly is + impossible, and reusing the full-rank value on a smaller linear would be a + silent miscompile -- so the caller must bail before touching the graph + rather than proceed. + + Granularity is read from ``meta["val"]`` (always present), so the slice-vs- + reuse decision never depends on the constant value being resolvable; the value + is only needed when we actually slice. + """ + if qparam_node is None: + return (False, None, None) + val = qparam_node.meta["val"] + if val.ndim == 0 or val.shape[0] != out_features: + # Scalar / broadcast over the output dim -> reuse unchanged. + return (False, None, None) + tensor = get_constant(ep, qparam_node) + if tensor is None: + # Per-output but the constant value isn't available -> cannot slice. + return None + kind = get_input_kind(ep, qparam_node) + assert kind is not None + return (True, tensor, kind) + + +def _apply_output_qparam( + ep: ExportedProgram, + qparam_node: Optional[fx.Node], + plan: _QParamPlan, + slice_range: tuple[int, int, int], + linear_node: fx.Node, +) -> Optional[fx.Node]: + """Apply a plan from ``_plan_output_qparam`` to one output slice. + + Returns the new sliced constant node, or ``None`` to reuse the original node + (the plan determined the arg broadcasts / is absent). + """ + do_slice, tensor, kind = plan + if not do_slice: + return None + assert qparam_node is not None and tensor is not None and kind is not None + start, end, step = slice_range + fqn = get_fqn(ep, qparam_node) + assert fqn is not None + return add_constant( + ep, + f"{fqn.replace('.', '_')}_slice_{start}_{end}_{step}", + tensor[start:end:step], + linear_node, + kind=kind, + ) + + +class _SplitPlan(NamedTuple): + """A fully validated split -- everything needed to apply it, computed without + mutating the graph (see ``_plan_split`` / ``_apply_split``).""" + + users: list[fx.Node] + slice_ranges: list[tuple[int, int, int]] + weight_tensor: torch.Tensor + weight_kind: InputKind + # Per-output dim-0 arg nodes and their plans, keyed by _PER_OUTPUT_DIM0_ARGS. + arg_nodes: dict[str, Optional[fx.Node]] + arg_plans: dict[str, _QParamPlan] + + +class SplitLinearAtSlices(ExportedProgramPassBase): + """Split a fused_quant.linear into smaller linears when its output is + sliced on the output dimension. + + Matches: fused_quant.linear → (slice_copy[dim=-1], slice_copy[dim=-1], ...) + + When all users of the linear are output-dimension slices and the total + sliced output features does not exceed the original (no extra compute), + the linear is split into one smaller linear per slice. Each gets a + compile-time sliced copy of the weight (and bias/weight_scale/weight_zp + if present). + + The split is fully validated before the graph is modified: weight, bias, and + any per-output-channel weight scale / zero_point must all be resolvable to + constants that can be sliced. If any of them varies along the output dim but + can't be resolved, ``_try_split`` bails without touching the graph, so we + never leave a sub-linear with a sliced weight but a mismatched full-rank + qparam. (The ``ConstantFold`` pass run before the optimization passes bakes + the per-channel scale/zp ``view_copy`` fusion inserts, so these resolve in the + normal pipeline.) + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + + for linear_node in list( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + ): + if self._try_split(exported_program, linear_node): + modified = True + + if modified: + exported_program = constant_prop_pass(exported_program) + + return ExportedProgramPassResult(exported_program, modified) + + def _try_split(self, ep: ExportedProgram, linear_node: fx.Node) -> bool: + plan = self._plan_split(ep, linear_node) + if plan is None: + return False + self._apply_split(ep, linear_node, plan) + logger.info( + "Split linear %s into %d smaller linears at output slices", + linear_node.name, + len(plan.users), + ) + return True + + def _plan_split( + self, ep: ExportedProgram, linear_node: fx.Node + ) -> Optional[_SplitPlan]: + """Validate the split and compute everything needed to apply it, WITHOUT + mutating the graph. Returns ``None`` (the linear can't be safely split) if + any user is not an output-dim slice, the weight isn't a resolvable + constant, the slices over-cover the output, the output qparams are + per-channel, or any per-output dim-0 qparam varies along the output dim but + can't be sliced -- so the caller bails before touching the graph rather + than half-rewrite it.""" + users = list(linear_node.users.keys()) + if not all(u.target == exir_ops.edge.aten.slice_copy.Tensor for u in users): + return None + + weight_node = get_arg(linear_node, "weight", fx.Node) + weight_tensor = get_constant(ep, weight_node) + if weight_tensor is None: + return None + weight_kind = get_input_kind(ep, weight_node) + assert weight_kind is not None + out_features = weight_tensor.shape[0] + + # Every user must be an output-dim slice, together covering no more than + # out_features (no extra compute). + slice_ranges: list[tuple[int, int, int]] = [] + for user in users: + r = _is_output_dim_slice(user, out_features) + if r is None: + return None + slice_ranges.append(r) + if sum(len(range(s, e, st)) for s, e, st in slice_ranges) > out_features: + return None + + # Per-channel OUTPUT requant (out_scale/out_zero_point index the output's + # last dim -- a different axis than the dim-0 weight/bias qparams) isn't + # handled; bail rather than reuse a full-rank value on a smaller linear. + for name in _OUTPUT_QPARAMS: + node = get_arg(linear_node, name, Optional[fx.Node]) + if node is not None and node.meta["val"].numel() != 1: + return None + + # Plan every per-output dim-0 tensor arg. A None plan means the arg varies + # along the output dim but can't be sliced -> the whole split aborts. + arg_nodes: dict[str, Optional[fx.Node]] = { + name: get_arg(linear_node, name, Optional[fx.Node]) + for name in _PER_OUTPUT_DIM0_ARGS + } + arg_plans: dict[str, _QParamPlan] = {} + for name, node in arg_nodes.items(): + arg_plan = _plan_output_qparam(ep, node, out_features) + if arg_plan is None: + return None + arg_plans[name] = arg_plan + + return _SplitPlan( + users=users, + slice_ranges=slice_ranges, + weight_tensor=weight_tensor, + weight_kind=weight_kind, + arg_nodes=arg_nodes, + arg_plans=arg_plans, + ) + + def _apply_split( + self, ep: ExportedProgram, linear_node: fx.Node, plan: _SplitPlan + ) -> None: + """Apply a validated ``_SplitPlan``: replace each output slice with its own + smaller linear carrying sliced weight / bias / qparams.""" + graph = ep.graph_module.graph + weight_node = get_arg(linear_node, "weight", fx.Node) + weight_fqn = get_fqn(ep, weight_node) + assert weight_fqn is not None + for user, (start, end, step) in zip(plan.users, plan.slice_ranges): + new_weight_node = add_constant( + ep, + f"{weight_fqn.replace('.', '_')}_slice_{start}_{end}_{step}", + plan.weight_tensor[start:end:step], + linear_node, + kind=plan.weight_kind, + ) + new_arg_nodes = { + name: _apply_output_qparam( + ep, + plan.arg_nodes[name], + plan.arg_plans[name], + (start, end, step), + linear_node, + ) + for name in _PER_OUTPUT_DIM0_ARGS + } + + with graph.inserting_before(user): + new_linear = graph.call_function( + exir_ops.edge.fused_quant.linear.default, + args=linear_node.args, + ) + new_linear.meta["val"] = user.meta["val"] + + set_arg(new_linear, "weight", new_weight_node) + for name, new_node in new_arg_nodes.items(): + if new_node is not None: + set_arg(new_linear, name, new_node) + + user.replace_all_uses_with(new_linear) diff --git a/backends/fused_quant/optimization_passes/to_channels_first.py b/backends/fused_quant/optimization_passes/to_channels_first.py new file mode 100644 index 00000000000..fc63ed21239 --- /dev/null +++ b/backends/fused_quant/optimization_passes/to_channels_first.py @@ -0,0 +1,219 @@ +# 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. + +from __future__ import annotations + +import operator +from typing import cast + +import executorch.backends.fused_quant.ops # noqa: F401 +from executorch.backends.fused_quant.colorer import is_colored +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_flat, + get_qparams_from_node, +) +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import compute_meta_val +from executorch.exir.dialects._ops import ops as exir_ops +from torch import fx +from torch.fx.passes.infra.pass_base import PassBase, PassResult + + +def _copy_meta_and_compute(node: fx.Node, source: fx.Node) -> None: + node.meta = source.meta.copy() + node.meta["val"] = compute_meta_val(node) + + +def _permute( + graph: fx.Graph, tensor: fx.Node, dims: list[int], source: fx.Node +) -> fx.Node: + permute = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(tensor, dims), + ) + _copy_meta_and_compute(permute, source) + return permute + + +def _permute_qparams( + node: fx.Node, + prefix: str, + dims: list[int], +) -> list[fx.node.Argument]: + qparams = get_qparams_from_node(node, prefix) + if qparams is None or qparams.is_per_tensor(): + return get_qparams_flat(node, prefix) + + graph = node.graph + scale = _permute(graph, qparams.scale, dims, qparams.scale) + zero_point = _permute(graph, qparams.zero_point, dims, qparams.zero_point) + return [ + scale, + zero_point, + qparams.dtype, + qparams.quant_min, + qparams.quant_max, + ] + + +def _convert_convolution_to_channels_first(node: fx.Node) -> None: + graph = node.graph + inp = get_arg(node, "inp", fx.Node) + weight = get_arg(node, "weight", fx.Node) + bias = get_arg(node, "bias", fx.Node | None) + ndim = inp.meta["val"].ndim + to_channels_first = [0, ndim - 1, *range(1, ndim - 1)] + to_channels_last = [0, *range(2, ndim), 1] + + with graph.inserting_before(node): + inp_channels_first = _permute(graph, inp, to_channels_first, inp) + weight_channels_first = _permute(graph, weight, to_channels_first, weight) + args: tuple[fx.node.Argument, ...] = ( + inp_channels_first, + weight_channels_first, + bias, + *_permute_qparams(node, "inp", to_channels_first), + *_permute_qparams(node, "weight", to_channels_first), + *get_qparams_flat(node, "bias"), + *_permute_qparams(node, "out", to_channels_first), + get_arg(node, "stride"), + get_arg(node, "padding"), + get_arg(node, "dilation"), + get_arg(node, "transposed"), + get_arg(node, "output_padding"), + get_arg(node, "groups"), + ) + convolution = graph.call_function( + exir_ops.edge.fused_quant.convolution.default, + args=args, + ) + _copy_meta_and_compute(convolution, node) + output = _permute(graph, convolution, to_channels_last, node) + + node.replace_all_uses_with(output) + graph.erase_node(node) + + +def _convert_avg_pool2d_to_channels_first(node: fx.Node) -> None: + graph = node.graph + inp = get_arg(node, "inp", fx.Node) + to_channels_first = [0, 3, 1, 2] + to_channels_last = [0, 2, 3, 1] + + with graph.inserting_before(node): + inp_channels_first = _permute(graph, inp, to_channels_first, inp) + args: tuple[fx.node.Argument, ...] = ( + inp_channels_first, + *_permute_qparams(node, "inp", to_channels_first), + *_permute_qparams(node, "out", to_channels_first), + get_arg(node, "kernel_size"), + get_arg(node, "stride"), + get_arg(node, "padding"), + get_arg(node, "ceil_mode"), + get_arg(node, "count_include_pad"), + get_arg(node, "divisor_override"), + ) + pool = graph.call_function( + exir_ops.edge.fused_quant.avg_pool2d.default, + args=args, + ) + _copy_meta_and_compute(pool, node) + output = _permute(graph, pool, to_channels_last, node) + + node.replace_all_uses_with(output) + graph.erase_node(node) + + +def _convert_max_pool2d_with_indices_to_channels_first(node: fx.Node) -> None: + users = list(node.users) + if any( + user.target is not operator.getitem or user.args[1] not in (0, 1) + for user in users + ): + raise ValueError( + "fused_quant.max_pool2d_with_indices_channels_last must be consumed " + "through getitem(0) or getitem(1)" + ) + + graph = node.graph + inp = get_arg(node, "inp", fx.Node) + to_channels_first = [0, 3, 1, 2] + to_channels_last = [0, 2, 3, 1] + with graph.inserting_before(node): + inp_channels_first = _permute(graph, inp, to_channels_first, inp) + args: tuple[fx.node.Argument, ...] = ( + inp_channels_first, + *_permute_qparams(node, "inp", to_channels_first), + *_permute_qparams(node, "out", to_channels_first), + get_arg(node, "kernel_size"), + get_arg(node, "stride"), + get_arg(node, "padding"), + get_arg(node, "dilation"), + get_arg(node, "ceil_mode"), + ) + pool = graph.call_function( + exir_ops.edge.fused_quant.max_pool2d_with_indices.default, + args=args, + ) + _copy_meta_and_compute(pool, node) + + for user in users: + output_index = cast(int, user.args[1]) + output_channels_first = graph.call_function( + operator.getitem, + args=(pool, output_index), + ) + _copy_meta_and_compute(output_channels_first, user) + output = _permute( + graph, + output_channels_first, + to_channels_last, + user, + ) + user.replace_all_uses_with(output) + graph.erase_node(user) + + graph.erase_node(node) + + +class ToChannelsFirst(PassBase): + """Canonicalize channels-last fused ops to channels-first fused ops. + + Colored nodes are left alone: a backend that claimed a channels-last op + wants it in that layout, and relayouting would take it away. See + :class:`executorch.backends.fused_quant.colorer.ColorerBase`. + """ + + def call(self, graph_module: fx.GraphModule) -> PassResult: + graph = graph_module.graph + conversions = ( + ( + exir_ops.edge.fused_quant.convolution_channels_last.default, + _convert_convolution_to_channels_first, + ), + ( + exir_ops.edge.fused_quant.max_pool2d_with_indices_channels_last.default, + _convert_max_pool2d_with_indices_to_channels_first, + ), + ( + exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + _convert_avg_pool2d_to_channels_first, + ), + ) + + modified = False + for target, convert in conversions: + for node in list(graph.find_nodes(op="call_function", target=target)): + if is_colored(node): + continue + convert(node) + modified = True + + if modified: + graph.eliminate_dead_code() + graph.lint() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/fused_quant/optimization_passes/to_channels_last.py b/backends/fused_quant/optimization_passes/to_channels_last.py new file mode 100644 index 00000000000..2dd3ad05d1a --- /dev/null +++ b/backends/fused_quant/optimization_passes/to_channels_last.py @@ -0,0 +1,321 @@ +# 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. + + +import operator +from typing import cast, Optional + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.colorer import is_colored +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_flat, + get_qparams_from_node, +) +from executorch.backends.fused_quant.ops_utils import compute_conv_out_shape_nhwc +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch import fx +from torch._export.utils import _detect_fake_mode_from_gm + + +def _permute_qparams( + node: fx.Node, + prefix: str, + dims: list[int], +) -> list[fx.node.Argument]: + qparams = get_qparams_from_node(node, prefix) + if qparams is None or qparams.is_per_tensor(): + return get_qparams_flat(node, prefix) + + graph = node.graph + permuted: list[fx.Node] = [] + for value in (qparams.scale, qparams.zero_point): + value_permuted = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(value, dims), + ) + value_permuted.meta["val"] = value.meta["val"].permute(*dims) + permuted.append(value_permuted) + + return [ + *permuted, + qparams.dtype, + qparams.quant_min, + qparams.quant_max, + ] + + +def _convert_convolution_to_channels_last(node: fx.Node) -> None: + """Convert a fused_quant.convolution node to fused_quant.convolution_channels_last. + + This function: + 1. Inserts permute nodes to convert inputs from NCHW -> NHWC + 2. Inserts permute nodes to convert weights from OIHW -> OHWI + 3. Replaces the convolution with convolution_channels_last + 4. Inserts permute node to convert output from NHWC -> NCHW + """ + graph = node.graph + + inp = get_arg(node, "inp", fx.Node) + weight = get_arg(node, "weight", fx.Node) + + inp_shape = inp.meta["val"].shape + num_spatial_dims = len(inp_shape) - 2 # Subtract batch and channel dims + + # NCHW -> NHWC: (0, 2, 3, ..., n+1, 1) + nchw_to_nhwc = [0] + list(range(2, num_spatial_dims + 2)) + [1] + # NHWC -> NCHW: (0, n+1, 1, 2, ..., n) + nhwc_to_nchw = [0, num_spatial_dims + 1] + list(range(1, num_spatial_dims + 1)) + # OIHW -> OHWI: (0, 2, 3, ..., n+1, 1) + oihw_to_ohwi = [0] + list(range(2, num_spatial_dims + 2)) + [1] + + with graph.inserting_before(node): + # Transpose input from NCHW to NHWC + inp_nhwc = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(inp, nchw_to_nhwc), + ) + inp_nhwc.meta["val"] = inp.meta["val"].permute(*nchw_to_nhwc) + + # Transpose weight from OIHW to OHWI + weight_ohwi = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(weight, oihw_to_ohwi), + ) + weight_ohwi.meta["val"] = weight.meta["val"].permute(*oihw_to_ohwi) + + bias = get_arg(node, "bias", Optional[fx.Node]) + stride = get_arg(node, "stride", list[int]) + padding = get_arg(node, "padding", list[int]) + dilation = get_arg(node, "dilation", list[int]) + transposed = get_arg(node, "transposed", bool) + output_padding = get_arg(node, "output_padding", list[int]) + groups = get_arg(node, "groups", int) + + # pyre-ignore[60]: Pyre doesn't support multiple variadic tuple splats + new_args: tuple[fx.node.Argument, ...] = ( + inp_nhwc, + weight_ohwi, + bias, + *_permute_qparams(node, "inp", nchw_to_nhwc), + *_permute_qparams(node, "weight", oihw_to_ohwi), + *get_qparams_flat(node, "bias"), + *_permute_qparams(node, "out", nchw_to_nhwc), + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + conv_channels_last = graph.call_function( + exir_ops.edge.fused_quant.convolution_channels_last.default, + args=new_args, + ) + + out_shape_nhwc = compute_conv_out_shape_nhwc( + inp_nhwc.meta["val"], + weight_ohwi.meta["val"], + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + assert graph.owning_module is not None + fake_mode = _detect_fake_mode_from_gm(graph.owning_module) + assert fake_mode + with fake_mode: + conv_channels_last.meta["val"] = torch.empty( + out_shape_nhwc, dtype=node.meta["val"].dtype + ) + + # Transpose output from NHWC back to NCHW + out_nchw = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(conv_channels_last, nhwc_to_nchw), + ) + out_nchw.meta["val"] = conv_channels_last.meta["val"].permute(*nhwc_to_nchw) + + node.replace_all_uses_with(out_nchw) + graph.erase_node(node) + + +def _convert_max_pool2d_with_indices_to_channels_last(node: fx.Node) -> None: + """Convert both max_pool2d_with_indices tuple outputs to channels-last.""" + graph = node.graph + inp = get_arg(node, "inp", fx.Node) + nchw_to_nhwc = [0, 2, 3, 1] + nhwc_to_nchw = [0, 3, 1, 2] + users = list(node.users) + if any( + user.target != operator.getitem or user.args[1] not in (0, 1) for user in users + ): + raise ValueError( + "fused_quant.max_pool2d_with_indices must be consumed through " + "getitem(0) or getitem(1)" + ) + + with graph.inserting_before(node): + inp_nhwc = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(inp, nchw_to_nhwc), + ) + inp_nhwc.meta["val"] = inp.meta["val"].permute(*nchw_to_nhwc) + + # pyre-ignore[60]: Pyre doesn't support multiple variadic tuple splats + new_args: tuple[fx.node.Argument, ...] = ( + inp_nhwc, + *_permute_qparams(node, "inp", nchw_to_nhwc), + *_permute_qparams(node, "out", nchw_to_nhwc), + get_arg(node, "kernel_size", list[int]), + get_arg(node, "stride", list[int]), + get_arg(node, "padding", list[int]), + get_arg(node, "dilation", list[int]), + get_arg(node, "ceil_mode", bool), + ) + pool_channels_last = graph.call_function( + exir_ops.edge.fused_quant.max_pool2d_with_indices_channels_last.default, + args=new_args, + ) + node_values = cast(tuple[torch.Tensor, torch.Tensor], node.meta["val"]) + pool_channels_last.meta["val"] = tuple( + value.permute(*nchw_to_nhwc) for value in node_values + ) + + for user in users: + output_index = cast(int, user.args[1]) + output_nhwc = graph.call_function( + operator.getitem, + args=(pool_channels_last, output_index), + ) + output_nhwc.meta = user.meta.copy() + output_nhwc.meta["val"] = pool_channels_last.meta["val"][output_index] + + output_nchw = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(output_nhwc, nhwc_to_nchw), + ) + output_nchw.meta = user.meta.copy() + output_nchw.meta["val"] = node_values[output_index] + user.replace_all_uses_with(output_nchw) + graph.erase_node(user) + + graph.erase_node(node) + + +def _convert_avg_pool2d_to_channels_last(node: fx.Node) -> None: + graph = node.graph + inp = get_arg(node, "inp", fx.Node) + nchw_to_nhwc = [0, 2, 3, 1] + nhwc_to_nchw = [0, 3, 1, 2] + + with graph.inserting_before(node): + inp_nhwc = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(inp, nchw_to_nhwc), + ) + inp_nhwc.meta["val"] = inp.meta["val"].permute(*nchw_to_nhwc) + + # pyre-ignore[60]: Pyre doesn't support multiple variadic tuple splats + new_args: tuple[fx.node.Argument, ...] = ( + inp_nhwc, + *_permute_qparams(node, "inp", nchw_to_nhwc), + *_permute_qparams(node, "out", nchw_to_nhwc), + get_arg(node, "kernel_size", list[int]), + get_arg(node, "stride", list[int]), + get_arg(node, "padding", list[int]), + get_arg(node, "ceil_mode", bool), + get_arg(node, "count_include_pad", bool), + get_arg(node, "divisor_override", Optional[int]), + ) + pool_channels_last = graph.call_function( + exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + args=new_args, + ) + pool_channels_last.meta["val"] = node.meta["val"].permute(*nchw_to_nhwc) + + output_nchw = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(pool_channels_last, nhwc_to_nchw), + ) + output_nchw.meta = node.meta.copy() + + node.replace_all_uses_with(output_nchw) + graph.erase_node(node) + + +class ToChannelsLast(ExportPass): + """ + Convert fused quant convolution and pool ops to channels-last layout. + + This pass identifies: + fused_quant.convolution (NCHW layout) + fused_quant.max_pool2d_with_indices (NCHW layout) + fused_quant.avg_pool2d (NCHW layout) + + and replaces it with: + permute (NCHW->NHWC) -> fused_quant.convolution_channels_last -> permute (NHWC->NCHW) + permute (NCHW->NHWC) -> fused_quant.max_pool2d_with_indices_channels_last + -> getitem per live output -> permute (NHWC->NCHW) + permute (NCHW->NHWC) -> fused_quant.avg_pool2d_channels_last + -> permute (NHWC->NCHW) + + For most backends, NHWC layout is more efficient than NCHW layout, so this optimization + is applied in the fused_quant IR. + + Colored nodes are left alone: a backend that claimed an op is going to lower + it itself, and relayouting would take it away. See + :class:`executorch.backends.fused_quant.colorer.ColorerBase`. + """ + + def call(self, graph_module: fx.GraphModule) -> PassResult: + """ + Execute the channels-last optimization pass on the graph module. + + Args: + graph_module: The graph module to transform. + + Returns: + PassResult with the transformed graph module and modification status. + """ + modified = False + graph = graph_module.graph + + for node in graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution.default, + ): + if is_colored(node): + continue + _convert_convolution_to_channels_last(node) + modified = True + + for node in graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.max_pool2d_with_indices.default, + ): + if is_colored(node): + continue + _convert_max_pool2d_with_indices_to_channels_last(node) + modified = True + + for node in graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d.default, + ): + if is_colored(node): + continue + _convert_avg_pool2d_to_channels_last(node) + modified = True + + if modified: + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/fused_quant/optimization_passes/to_convolution.py b/backends/fused_quant/optimization_passes/to_convolution.py new file mode 100644 index 00000000000..88672996fbb --- /dev/null +++ b/backends/fused_quant/optimization_passes/to_convolution.py @@ -0,0 +1,49 @@ +# 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. + +from __future__ import annotations + +import executorch.backends.fused_quant.ops # noqa: F401 +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from torch import fx +from torch.fx.passes.infra.pass_base import PassBase, PassResult + + +class ToConvolution(PassBase): + """Canonicalize fused conv1d/2d/3d ops to fused_quant.convolution.""" + + def call(self, graph_module: fx.GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + + for target, spatial_dims in ( + (exir_ops.edge.fused_quant.conv1d.default, 1), + (exir_ops.edge.fused_quant.conv2d.default, 2), + (exir_ops.edge.fused_quant.conv3d.default, 3), + ): + for node in list(graph.find_nodes(op="call_function", target=target)): + values = [get_arg(node, arg.name) for arg in target._schema.arguments] + with graph.inserting_before(node): + convolution = graph.call_function( + exir_ops.edge.fused_quant.convolution.default, + args=( + *values[:-1], + False, + [0] * spatial_dims, + values[-1], + ), + ) + convolution.meta = node.meta.copy() + node.replace_all_uses_with(convolution) + graph.erase_node(node) + modified = True + + if modified: + graph.eliminate_dead_code() + graph.lint() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/fused_quant/pass_base.py b/backends/fused_quant/pass_base.py new file mode 100644 index 00000000000..5a47eb79fcb --- /dev/null +++ b/backends/fused_quant/pass_base.py @@ -0,0 +1,48 @@ +# 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. + + +import logging + +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.pass_manager import ExportedProgramPassManager, PassType +from torch.export import ExportedProgram + +logger: logging.Logger = logging.getLogger(__name__) + + +class IterativePassGroup(ExportedProgramPassBase): + """Runs a group of passes as a single ExportedProgramPassBase. + + Useful for composing pass pipelines — e.g., running a set of optimization + passes multiple times until convergence while other passes run once. + """ + + def __init__(self, passes: list[PassType], steps: int) -> None: + # Retain the children and step count so instrumentation can recursively + # rebuild this group with wrapped leaf passes. + self._passes = passes + self._steps = steps + self._pm = ExportedProgramPassManager(passes, steps=1) + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + modified = False + for _ in range(self._steps): + result = self._pm(exported_program) + exported_program = result.exported_program + if not result.modified: + return ExportedProgramPassResult(exported_program, modified) + modified = True + # A single-step group runs its passes once by design; only a group meant + # to iterate to a fixed point can fail to converge. + if self._steps > 1: + logger.warning( + "%s did not converge within %d steps; the graph was still changing " + "on the last step", + type(self).__name__, + self._steps, + ) + return ExportedProgramPassResult(exported_program, modified) diff --git a/backends/fused_quant/passes.py b/backends/fused_quant/passes.py new file mode 100644 index 00000000000..d81c1e6fd3d --- /dev/null +++ b/backends/fused_quant/passes.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. + +from typing import cast + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_from_node, + split_fused_arg_names, +) +from executorch.backends.fused_quant.optimization_passes.constant_fold import ( + ConstantFold, +) +from executorch.backends.fused_quant.optimization_passes.fold_scalar_mul import ( + FoldScalarMul, +) +from executorch.backends.fused_quant.optimization_passes.fuse_add_into_linear import ( + FuseAddIntoLinear, +) +from executorch.backends.fused_quant.optimization_passes.fuse_mul_into_linear import ( + FuseMulIntoLinear, +) +from executorch.backends.fused_quant.optimization_passes.hoist_activations import ( + HoistActivations, +) +from executorch.backends.fused_quant.optimization_passes.prequantize_embedding import ( + PrequantizeEmbedding, +) +from executorch.backends.fused_quant.optimization_passes.quant_absorption import ( + QuantAbsorptionPass, +) +from executorch.backends.fused_quant.optimization_passes.replace_dequant_quant_with_requantize import ( + ReplaceDequantQuantWithRequantize, +) +from executorch.backends.fused_quant.optimization_passes.sink_constant_cat import ( + SinkConstantCat, +) +from executorch.backends.fused_quant.optimization_passes.split_linear_at_slices import ( + SplitLinearAtSlices, +) +from executorch.backends.fused_quant.optimization_passes.to_channels_last import ( + ToChannelsLast, +) +from executorch.backends.fused_quant.optimization_passes.to_convolution import ( + ToConvolution, +) +from executorch.backends.fused_quant.pass_base import IterativePassGroup +from executorch.backends.transforms.advance_quantize_op_above_def_chain import ( + AdvanceQuantizeOpAboveDefChainPass, +) +from executorch.backends.transforms.advance_quantize_op_above_def_in_branch import ( + AdvanceQuantizeOpAboveDefInBranchPass, +) +from executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops import ( + FuseCascadedTransposeOrPermuteOps, +) +from executorch.backends.transforms.fuse_cascaded_view_ops import FuseCascadedViewOps +from executorch.backends.transforms.fuse_quant_dequant_to_requantize import ( + FuseQuantDequantToRequantizePass, +) +from executorch.backends.transforms.merge_split_concat_chain import ( + MergeSplitConcatChainPass, +) +from executorch.backends.transforms.move_permute_after_concat import ( + MovePermuteAfterConcat, +) +from executorch.backends.transforms.move_slice_before_permute import ( + MoveSliceBeforePermutePass, +) +from executorch.backends.transforms.move_slice_before_view import ( + MoveSliceBeforeViewPass, +) +from executorch.backends.transforms.postpone_dequantize_op_below_use_chain import ( + PostponeDequantizeOpBelowUseChainPass, +) +from executorch.backends.transforms.propagate_slice import PropagateSlice +from executorch.backends.transforms.remove_alias_copy_op import RemoveAliasCopyOpPass +from executorch.backends.transforms.remove_branched_quant_dequant import ( + RemoveBranchedQuantDequant, +) +from executorch.backends.transforms.remove_cat_from_slice_copy import ( + RemoveCatFromSliceCopyPass, +) +from executorch.backends.transforms.remove_clone_ops_transform_imported import ( + RemoveCloneOpsTransformImported, +) +from executorch.backends.transforms.remove_nop_as_strided_copy_op import ( + RemoveNopAsStridedCopyOpPass, +) +from executorch.backends.transforms.remove_nop_expand_op import RemoveNopExpandOpPass +from executorch.backends.transforms.remove_nop_slice_or_view_op import ( + RemoveNopSliceOrViewOpPass, +) +from executorch.backends.transforms.remove_permute_before_mean import ( + RemovePermuteBeforeMeanPass, +) +from executorch.backends.transforms.remove_permutes_around_elementwise_ops import ( + RemovePermutesAroundElementwiseOps, +) +from executorch.backends.transforms.remove_zero_sized_cat_args import ( + RemoveZeroSizedCatArgsPass, +) +from executorch.backends.transforms.replace_nop_transpose_or_permute_with_view import ( + ReplaceNopTransposeOrPermuteWithViewPass, +) +from executorch.backends.transforms.replace_select_with_view_op import ( + ReplaceSelectWithViewOpPass, +) +from executorch.backends.transforms.replace_split_with_slice import ( + ReplaceSplitWithSlicePass, +) +from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( + ReplaceSqueezeAndUnsqueezeWithViewPass, +) +from executorch.backends.transforms.split_dequantized_cat import SplitDequantizedCatPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_manager import PassType +from torch.fx.passes.dialect.common.cse_pass import CSEPass + +# fused_quant elementwise ops (add / mul in both their Tensor and Scalar +# overloads, and the unary activations) that RemovePermutesAroundElementwiseOps +# should treat as permutable, so a permutation from ToChannelsLast propagates +# through them and cancels at the boundary permutes -- running the whole region +# in NHWC. See the pass for how the lifted scale / zero_point operands are +# handled per qparam flavor. +FUSED_QUANT_ELEMENTWISE: set[EdgeOpOverload] = { + exir_ops.edge.fused_quant.add.default, + exir_ops.edge.fused_quant.add.Scalar, + exir_ops.edge.fused_quant.sub.default, + exir_ops.edge.fused_quant.sub.Scalar, + exir_ops.edge.fused_quant.mul.default, + exir_ops.edge.fused_quant.mul.Scalar, + exir_ops.edge.fused_quant.relu.default, + exir_ops.edge.fused_quant.hardswish.default, + exir_ops.edge.fused_quant.sigmoid.default, + exir_ops.edge.fused_quant.tanh.default, + exir_ops.edge.fused_quant.hard_tanh.default, + exir_ops.edge.fused_quant.silu.default, + exir_ops.edge.fused_quant.hardsigmoid.default, + exir_ops.edge.fused_quant.gelu.default, +} + +# Hardcoded max number of iterations for the optimization group. +_OPTIMIZATION_STEPS = 5 + + +def _has_only_per_tensor_qparams(node: torch.fx.Node) -> bool: + """Whether every qparam block the op carries is per-tensor. + + Which blocks exist depends on the op: unary fused ops have no ``other``. + """ + node_target = cast(EdgeOpOverload, node.target) + input_names, _ = split_fused_arg_names(node_target) + for prefix in (*input_names, "out"): + qparams = get_qparams_from_node(node, prefix) + if qparams is not None and not qparams.is_per_tensor(): + return False + return True + + +def get_optimization_passes( + *, + absorb_quant_with_fork: bool = False, + extra_quantizable_ops: dict[EdgeOpOverload, tuple[int, ...]] | None = None, +) -> IterativePassGroup: + """The rewrites applied repeatedly to the fused_quant edge graph. + + Returns the group rather than running it. The passes feed each other -- a + layout change exposes a fusion, which exposes a dead transpose -- so they are + re-run until no pass reports a change, capped at ``_OPTIMIZATION_STEPS`` + iterations. The cap matters: the group is not guaranteed to converge, and a + graph that is still changing on the last iteration is left as-is, with a + warning. + """ + return IterativePassGroup( + [ + ToConvolution(), + FoldScalarMul(), + QuantAbsorptionPass(absorb_with_fork=absorb_quant_with_fork), + ToChannelsLast(), + HoistActivations(), + CSEPass(), + FuseAddIntoLinear(), + FuseMulIntoLinear(), + RemoveAliasCopyOpPass(), + RemoveCloneOpsTransformImported(), + RemoveNopAsStridedCopyOpPass(), + RemoveNopExpandOpPass(), + RemoveNopSliceOrViewOpPass(), + FuseCascadedTransposeOrPermuteOps(), + FuseCascadedViewOps(), + # Canonicalize rank-changing shape ops to view_copy so the permute + # pass below sees a single representation. + ReplaceSqueezeAndUnsqueezeWithViewPass(), + MergeSplitConcatChainPass(), + ReplaceSplitWithSlicePass(), + ReplaceSelectWithViewOpPass(), + PostponeDequantizeOpBelowUseChainPass(), + RemovePermuteBeforeMeanPass(), + RemovePermutesAroundElementwiseOps( + extra_permutable_ops=FUSED_QUANT_ELEMENTWISE, + target_filters={ + target: _has_only_per_tensor_qparams + for target in FUSED_QUANT_ELEMENTWISE + }, + ), + MovePermuteAfterConcat(), + ReplaceNopTransposeOrPermuteWithViewPass(), + SplitDequantizedCatPass(), + AdvanceQuantizeOpAboveDefChainPass( + extra_quantizable_ops=extra_quantizable_ops + ), + PropagateSlice( + additional_binary_targets=[ + exir_ops.edge.fused_quant.add.default, + exir_ops.edge.fused_quant.mul.default, + ], + target_filters={ + exir_ops.edge.fused_quant.add.default: _has_only_per_tensor_qparams, + exir_ops.edge.fused_quant.mul.default: _has_only_per_tensor_qparams, + }, + ), + AdvanceQuantizeOpAboveDefInBranchPass(), + FuseQuantDequantToRequantizePass(allow_requantize=False), + RemoveBranchedQuantDequant(), + PrequantizeEmbedding(), + SinkConstantCat(), + SplitLinearAtSlices(), + RemoveCatFromSliceCopyPass(), + RemoveZeroSizedCatArgsPass(), + MoveSliceBeforePermutePass(), + MoveSliceBeforeViewPass(), + ], + _OPTIMIZATION_STEPS, + ) + + +def get_fused_quant_passes( + *, + absorb_quant_with_fork: bool = False, + extra_quantizable_ops: dict[EdgeOpOverload, tuple[int, ...]] | None = None, +) -> list[PassType]: + """The backend-neutral pipeline that produces the optimized fused_quant graph. + + Runs on a graph that already contains fused_quant ops in edge dialect (fusion + happens before ``to_edge``) and stops before coloring, so the result is + hardware-independent. A backend appends its own coloring, lowering, and + partitioning passes after these. + """ + return [ + # Fold pure-constant subgraphs before the optimization passes. Most + # importantly this bakes the aten.view_copy that fusion inserts to + # reshape a per-channel weight scale/zero_point to full rank: without + # it those qparams reach the optimization passes as a non-constant + # view_copy, so passes that read them via get_constant (e.g. + # SplitLinearAtSlices) see None and silently skip -- leaving a + # sub-linear with a mismatched full-rank scale. + ConstantFold(), + get_optimization_passes( + absorb_quant_with_fork=absorb_quant_with_fork, + extra_quantizable_ops=extra_quantizable_ops, + ), + ConstantFold(), + ReplaceDequantQuantWithRequantize(), + ] diff --git a/backends/fused_quant/pre_quantize_passes/__init__.py b/backends/fused_quant/pre_quantize_passes/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/fused_quant/pre_quantize_passes/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/fused_quant/pre_quantize_passes/fold_batch_norm.py b/backends/fused_quant/pre_quantize_passes/fold_batch_norm.py new file mode 100644 index 00000000000..dbf0f69a98d --- /dev/null +++ b/backends/fused_quant/pre_quantize_passes/fold_batch_norm.py @@ -0,0 +1,351 @@ +# 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. + +from __future__ import annotations + +import operator +from dataclasses import dataclass +from typing import cast + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import ( + add_constant, + get_constant, + get_input_kind, +) +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch._ops import OpOverload +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +_BATCH_NORM = torch.ops.aten._native_batch_norm_legit_no_training.default +_CONVOLUTION = torch.ops.aten.convolution.default +_CONV_TARGETS = ( + torch.ops.aten.conv1d.default, + torch.ops.aten.conv2d.default, + torch.ops.aten.conv3d.default, + _CONVOLUTION, +) +_SLICE_TARGETS = ( + torch.ops.aten.slice.Tensor, + torch.ops.aten.slice_copy.Tensor, +) +_CONTIGUOUS_OR_CLONE_TARGETS = ( + torch.ops.aten.contiguous.default, + torch.ops.aten.clone.default, +) + + +@dataclass(frozen=True) +class _FoldConstants: + conv_weight_node: fx.Node + conv_bias_node: fx.Node | None + conv_weight: torch.Tensor + conv_bias: torch.Tensor | None + bn_weight: torch.Tensor | None + bn_bias: torch.Tensor | None + running_mean: torch.Tensor + running_var: torch.Tensor + + +def _get_optional_constant( + exported_program: ExportedProgram, + node: fx.Node, + name: str, +) -> tuple[fx.Node | None, torch.Tensor | None]: + arg = get_arg(node, name) + if arg is None: + return None, None + if not isinstance(arg, fx.Node): + return None, None + return arg, get_constant(exported_program, arg) + + +def _get_fold_constants( + exported_program: ExportedProgram, + conv: fx.Node, + batch_norm: fx.Node, +) -> _FoldConstants | None: + conv_weight_node, conv_weight = _get_optional_constant( + exported_program, conv, "weight" + ) + conv_bias_node, conv_bias = _get_optional_constant(exported_program, conv, "bias") + bn_weight_node, bn_weight = _get_optional_constant( + exported_program, batch_norm, "weight" + ) + bn_bias_node, bn_bias = _get_optional_constant(exported_program, batch_norm, "bias") + mean_node, running_mean = _get_optional_constant( + exported_program, batch_norm, "running_mean" + ) + var_node, running_var = _get_optional_constant( + exported_program, batch_norm, "running_var" + ) + required = ( + conv_weight_node, + conv_weight, + mean_node, + running_mean, + var_node, + running_var, + ) + if any(value is None for value in required): + return None + optional = ( + (conv_bias_node, conv_bias), + (bn_weight_node, bn_weight), + (bn_bias_node, bn_bias), + ) + if any(node is not None and tensor is None for node, tensor in optional): + return None + assert conv_weight_node is not None + assert conv_weight is not None + assert running_mean is not None + assert running_var is not None + return _FoldConstants( + conv_weight_node, + conv_bias_node, + conv_weight, + conv_bias, + bn_weight, + bn_bias, + running_mean, + running_var, + ) + + +def _native_output_accessors(batch_norm: fx.Node) -> list[fx.Node] | None: + output_accessors: list[fx.Node] = [] + for user in batch_norm.users: + if user.target is not operator.getitem or not isinstance(user.args[1], int): + return None + if user.args[1] == 0: + output_accessors.append(user) + elif user.users: + return None + return output_accessors or None + + +def _is_supported_conv(conv: fx.Node) -> bool: + if conv.op != "call_function" or conv.target not in _CONV_TARGETS: + return False + if conv.target is _CONVOLUTION and get_arg(conv, "transposed", bool): + return False + return True + + +def _is_non_channel_slice(node: fx.Node) -> bool: + if node.op != "call_function" or node.target not in _SLICE_TARGETS: + return False + + dim = get_arg(node, "dim", int) + if dim >= 0: + return dim != 1 + + inp = node.args[0] + if not isinstance(inp, fx.Node): + return False + inp_val = inp.meta["val"] + if not isinstance(inp_val, torch.Tensor): + return False + return dim + inp_val.ndim != 1 + + +def _is_contiguous_or_clone(node: fx.Node) -> bool: + return node.op == "call_function" and node.target in _CONTIGUOUS_OR_CLONE_TARGETS + + +def _get_fusion_nodes(batch_norm: fx.Node) -> tuple[fx.Node, fx.Node] | None: + batch_norm_input = get_arg(batch_norm, "input") + if not isinstance(batch_norm_input, fx.Node): + return None + if _is_supported_conv(batch_norm_input): + if len(batch_norm_input.users) != 1: + return None + return batch_norm_input, batch_norm_input + + replacement = batch_norm_input + slice_node = batch_norm_input + if _is_contiguous_or_clone(slice_node): + if len(slice_node.users) != 1: + return None + slice_input = slice_node.args[0] + if not isinstance(slice_input, fx.Node): + return None + slice_node = slice_input + + if not _is_non_channel_slice(slice_node) or len(slice_node.users) != 1: + return None + conv = slice_node.args[0] + if ( + not isinstance(conv, fx.Node) + or not _is_supported_conv(conv) + or len(conv.users) != 1 + ): + return None + return conv, replacement + + +def _make_fused_parameters( + constants: _FoldConstants, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor] | None: + channels = constants.running_mean.numel() + channel_shape = (channels,) + if constants.conv_weight.shape[0] != channels: + return None + for tensor in ( + constants.running_mean, + constants.running_var, + constants.conv_bias, + constants.bn_weight, + constants.bn_bias, + ): + if tensor is not None and tensor.shape != channel_shape: + return None + + scale = torch.rsqrt(constants.running_var + eps) + if constants.bn_weight is not None: + scale = scale * constants.bn_weight + conv_bias = constants.conv_bias + if conv_bias is None: + conv_bias = torch.zeros_like(constants.running_mean) + fused_bias = (conv_bias - constants.running_mean) * scale + if constants.bn_bias is not None: + fused_bias = fused_bias + constants.bn_bias + scale_shape = (channels, *([1] * (constants.conv_weight.ndim - 1))) + fused_weight = constants.conv_weight * scale.reshape(scale_shape) + dtype = constants.conv_weight.dtype + return ( + fused_weight.to(dtype=dtype).detach().contiguous(), + fused_bias.to(dtype=dtype).detach().contiguous(), + ) + + +def _constant_kind( + exported_program: ExportedProgram, + node: fx.Node, +) -> InputKind: + return get_input_kind(exported_program, node) or InputKind.CONSTANT_TENSOR + + +def _insert_folded_conv( + exported_program: ExportedProgram, + conv: fx.Node, + batch_norm: fx.Node, + constants: _FoldConstants, + fused_weight: torch.Tensor, + fused_bias: torch.Tensor, +) -> fx.Node: + weight = add_constant( + exported_program, + f"{batch_norm.name}_folded_weight", + fused_weight, + batch_norm, + kind=_constant_kind(exported_program, constants.conv_weight_node), + ) + bias_source = constants.conv_bias_node or constants.conv_weight_node + bias = add_constant( + exported_program, + f"{batch_norm.name}_folded_bias", + fused_bias, + batch_norm, + kind=_constant_kind(exported_program, bias_source), + ) + owning_module = conv.graph.owning_module + if owning_module is None: + raise RuntimeError(f"Convolution {conv} does not belong to a graph module") + normalized_args = conv.normalized_arguments( + owning_module, + normalize_to_only_use_kwargs=True, + ) + if normalized_args is None: + raise RuntimeError(f"Could not normalize convolution arguments for {conv}") + conv_kwargs = dict(normalized_args.kwargs) + conv_input = conv_kwargs.pop("input") + conv_kwargs.pop("weight") + conv_kwargs.pop("bias") + with batch_norm.graph.inserting_before(conv): + folded_conv = batch_norm.graph.call_function( + cast(OpOverload, conv.target), + args=(conv_input, weight, bias), + kwargs=conv_kwargs, + ) + folded_conv.meta = conv.meta.copy() + return folded_conv + + +def _erase_folded_pattern( + conv: fx.Node, + batch_norm: fx.Node, + output_accessors: list[fx.Node], + replacement: fx.Node, + folded_conv: fx.Node, +) -> None: + graph = batch_norm.graph + conv.replace_all_uses_with(folded_conv) + if replacement is conv: + replacement = folded_conv + for accessor in output_accessors: + accessor.replace_all_uses_with(replacement) + for accessor in list(batch_norm.users): + graph.erase_node(accessor) + graph.erase_node(batch_norm) + graph.erase_node(conv) + + +class FoldBatchNorm(ExportedProgramPassBase): + """Fold inference BatchNorm parameters into an upstream convolution.""" + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + for batch_norm in graph.find_nodes(op="call_function", target=_BATCH_NORM): + modified |= self._fold(exported_program, batch_norm) + + if modified: + exported_program = constant_prop_pass(exported_program) + return ExportedProgramPassResult(exported_program, modified) + + def _fold( + self, + exported_program: ExportedProgram, + batch_norm: fx.Node, + ) -> bool: + output_accessors = _native_output_accessors(batch_norm) + fusion_nodes = _get_fusion_nodes(batch_norm) + if output_accessors is None or fusion_nodes is None: + return False + conv, replacement = fusion_nodes + + constants = _get_fold_constants(exported_program, conv, batch_norm) + if constants is None: + return False + fused_parameters = _make_fused_parameters( + constants, + get_arg(batch_norm, "eps", float), + ) + if fused_parameters is None: + return False + fused_weight, fused_bias = fused_parameters + folded_conv = _insert_folded_conv( + exported_program, + conv, + batch_norm, + constants, + fused_weight, + fused_bias, + ) + _erase_folded_pattern( + conv, + batch_norm, + output_accessors, + replacement, + folded_conv, + ) + return True diff --git a/backends/fused_quant/pre_quantize_passes/fuse_add_softmax_into_masked_softmax.py b/backends/fused_quant/pre_quantize_passes/fuse_add_softmax_into_masked_softmax.py new file mode 100644 index 00000000000..f1b966bf749 --- /dev/null +++ b/backends/fused_quant/pre_quantize_passes/fuse_add_softmax_into_masked_softmax.py @@ -0,0 +1,305 @@ +# 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. + + +from dataclasses import dataclass + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import ( + add_constant, + get_constant, + get_fqn, + get_input_kind, +) +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch._ops import OpOverload +from torch.export import ExportedProgram + +_ADD: OpOverload = torch.ops.aten.add.Tensor +_INDEX: OpOverload = torch.ops.aten.index.Tensor +_EXPAND: OpOverload = torch.ops.aten.expand.default +_MASKED_SOFTMAX: OpOverload = torch.ops.aten._masked_softmax.default +# _masked_softmax mask_type 2 is the generic "mask matches input shape" form. The +# 2D forms (0 = square LxL attention, 1 = BxL padding) don't cover an [L_q, L_k] +# mask with a KV cache (L_q != L_k), so we always broadcast to the full input. +_MASK_TYPE_FULL: int = 2 +_SOFTMAX_TARGETS: tuple[OpOverload, ...] = ( + torch.ops.aten._safe_softmax.default, + torch.ops.aten.softmax.int, + torch.ops.aten._softmax.default, +) + +# Sentinel mask values (-1e4, -inf, finfo.min, and their bf16 roundings) sit far +# below this bound; real attention logits never do. So a binary constant holding +# a value past this bound is an additive attention mask, not data. +_NEGATIVE_MASK_THRESHOLD: float = -1e3 + + +@dataclass +class _MaskSource: + """How the additive mask reaches the add. + + ``buffer_node`` / ``buffer_tensor`` are the constant holding the mask. + ``gather_node`` is the ``index.Tensor`` reading it at runtime (e.g. the MEC + mask buffer indexed by ``input_pos``), or None when the constant feeds the + add directly. + """ + + buffer_node: fx.Node + buffer_tensor: torch.Tensor + gather_node: fx.Node | None + + +def _is_additive_mask(t: torch.Tensor) -> bool: + """True for a canonical additive attention mask: a float tensor with exactly + two values -- 0 at keep positions and a very-negative sentinel at masked ones. + + The keep value must be exactly 0. A uniform nonzero keep value would in fact + cancel in the softmax, but demanding 0 keeps detection conservative: it avoids + matching arbitrary binary tensors that merely happen to hold a very-negative + entry. It also excludes a single-valued (all-masked) constant, which is *not* + equivalent to masking -- a uniform -1e4 row softmaxes to uniform weights, not + zeros. + """ + if not torch.is_floating_point(t): + return False + uniq = torch.unique(t) + if uniq.numel() != 2: + return False + return bool(uniq.max() == 0) and bool(uniq.min() <= _NEGATIVE_MASK_THRESHOLD) + + +def _has_fully_masked_row(mask: torch.Tensor, dim: int, scores_rank: int) -> bool: + """True if some softmax row (a slice along ``dim``) is entirely masked. + + Such a row has nothing to normalize over, so the softmax is undefined: eager + ``aten._masked_softmax`` writes NaN for it, which would poison quant + observers downstream during calibration. The rewrite bails in that case. + Canonical causal / sliding-window masks keep the diagonal unmasked and never + trip this; it guards pathological masks (e.g. an all-padding query row). + """ + # True = masked (binary {0, sentinel} per _is_additive_mask; min is the sentinel). + masked = mask == mask.min() + # The mask broadcasts to the scores from the right, so map the softmax dim to a + # right-counted (negative) axis that lines up regardless of the mask's rank. + neg_dim = dim if dim < 0 else dim - scores_rank + if mask.dim() < -neg_dim: + # Mask doesn't span the softmax dim (it broadcasts along it): every slice + # along that dim repeats a single value, so any masked element is a fully + # masked row. + return bool(masked.any()) + return bool(masked.all(dim=neg_dim).any()) + + +def _classify_mask(ep: ExportedProgram, operand: fx.Node) -> _MaskSource | None: + """Resolve an add operand to its additive-mask constant, else None. + + Handles a constant fed directly, or one gathered through a single-use + ``index.Tensor`` (the MEC mask-buffer-indexed-by-input_pos pattern). + """ + const = get_constant(ep, operand) + if const is not None and _is_additive_mask(const): + return _MaskSource(operand, const, None) + + if operand.target is _INDEX and len(operand.users) == 1: + buffer_node = operand.args[0] + if isinstance(buffer_node, fx.Node): + const = get_constant(ep, buffer_node) + if const is not None and _is_additive_mask(const): + return _MaskSource(buffer_node, const, operand) + return None + + +def _split_mask_scores( + ep: ExportedProgram, add: fx.Node +) -> tuple[fx.Node, _MaskSource] | None: + """Find which add operand is the additive mask. + + Returns (scores_node, mask_source) when exactly one operand resolves to an + additive mask; None if neither or both qualify (ambiguous). + """ + lhs, rhs = add.args[0], add.args[1] + matches: list[tuple[fx.Node, _MaskSource]] = [] + for mask_cand, scores_cand in ((lhs, rhs), (rhs, lhs)): + if not isinstance(mask_cand, fx.Node) or not isinstance(scores_cand, fx.Node): + continue + source = _classify_mask(ep, mask_cand) + if source is not None: + matches.append((scores_cand, source)) + if len(matches) != 1: + return None + return matches[0] + + +def _broadcastable_to(src: torch.Size, target: torch.Size) -> bool: + """True if a tensor of shape src can broadcast/expand to target.""" + try: + return tuple(torch.broadcast_shapes(src, target)) == tuple(target) + except RuntimeError: + return False + + +def _mask_operand(source: _MaskSource) -> fx.Node: + """The node feeding the add as the mask (the gather, or the constant itself).""" + return source.gather_node if source.gather_node is not None else source.buffer_node + + +def _materialize_bool_mask( + ep: ExportedProgram, + source: _MaskSource, + before: fx.Node, + cache: dict[str, fx.Node], +) -> fx.Node: + """Produce the bool mask node _masked_softmax consumes (True = masked). + + Mutates the graph: the additive buffer is converted to bool once per buffer + (cached). For a gathered mask the existing index op is rewired to read the + bool buffer so the runtime gather is preserved; otherwise the bool constant + feeds the op directly. Call only after mask_type has been validated. + """ + key = source.buffer_node.name + bool_const = cache.get(key) + if bool_const is None: + kind = get_input_kind(ep, source.buffer_node) + assert kind is not None + fqn = get_fqn(ep, source.buffer_node) + assert fqn is not None + # Binary buffer: the very-negative sentinel is the minimum (see + # _is_additive_mask). Equality recovers the bool mask exactly. + bool_buffer = source.buffer_tensor == source.buffer_tensor.min() + bool_const = add_constant( + ep, f"{fqn}_bool", bool_buffer.contiguous(), before, kind=kind + ) + cache[key] = bool_const + + if source.gather_node is None: + return bool_const + + source.gather_node.update_arg(0, bool_const) + source.gather_node.meta["val"] = source.gather_node.meta["val"].to(torch.bool) + return source.gather_node + + +def _is_plain_softmax(node: fx.Node) -> bool: + """False for softmaxes that upcast dtype -- _masked_softmax cannot express that.""" + if node.target is torch.ops.aten._softmax.default: + return not get_arg(node, "half_to_float") + return get_arg(node, "dtype") is None + + +class FuseAddSoftmaxIntoMaskedSoftmax(ExportedProgramPassBase): + """Fuse ``add(scores, additive_mask) -> softmax`` into ``aten._masked_softmax``. + + Runs post-trace / pre-quantize on the ExportedProgram, where the attention + mask is still an inspectable constant -- either fed directly or gathered + through an ``index.Tensor`` (the MEC mask-buffer-indexed-by-input_pos + pattern). The add + softmax pair is collapsed to a single ``_masked_softmax`` + that takes a bool mask (True = masked), so the quantizer can anchor it as one + op (skipping the mask operand) and the very-negative sentinel never enters the + quantized domain -- masking becomes exact rather than an additive -1e4 + approximation. + + The additive buffer is converted to a bool buffer once per buffer; a gathered + mask keeps its runtime gather (rewired to read the bool buffer). Because + ``_safe_softmax`` is still a single op here, replacing it also drops its + internal NaN guard. The original additive buffer goes dead and is pruned by + the subsequent export. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + # Bool mask constants keyed by source buffer name, so a buffer used by + # multiple layers/softmaxes is materialized once. + bool_mask_cache: dict[str, fx.Node] = {} + # find_nodes returns a snapshot, so mutation while iterating is safe. + for target in _SOFTMAX_TARGETS: + for softmax in graph.find_nodes(op="call_function", target=target): + modified |= self._try_fuse(exported_program, softmax, bool_mask_cache) + + if modified: + # The recompose replaced each additive mask with a bool mask, leaving + # the additive buffer a dead lifted constant. constant_prop_pass erases + # such dead constants from the graph, state_dict/constants, AND + # graph_signature (and folds the now-constant bool-mask subgraph), + # keeping the signature self-consistent without a manual remove_constant. + # It also DCEs and recompiles. + exported_program = constant_prop_pass(exported_program) + + return ExportedProgramPassResult(exported_program, modified) + + def _try_fuse( + self, + ep: ExportedProgram, + softmax: fx.Node, + bool_mask_cache: dict[str, fx.Node], + ) -> bool: + """Rewrite add(scores, mask) -> softmax into a single _masked_softmax.""" + if not _is_plain_softmax(softmax): + return False + add = softmax.args[0] + if ( + not isinstance(add, fx.Node) + or add.target is not _ADD + or len(add.users) != 1 + ): + return False + + # add.Tensor computes ``self + alpha * other``; a non-default alpha scales + # one operand, so scores and mask are not combined as-is and the recompose + # would be wrong. Only a plain (alpha=1) add is fusible. + if get_arg(add, "alpha") != 1: + return False + + split = _split_mask_scores(ep, add) + if split is None: + return False + scores, source = split + + # Validate before mutating: _materialize rewires the gather and adds a + # constant, so an unusable mask shape must bail first or it would leave a + # bool gather feeding the still-present add. The mask must broadcast to + # the scores shape (mask_type 2 requires a full-shape mask). + scores_shape = scores.meta["val"].shape + if not _broadcastable_to(_mask_operand(source).meta["val"].shape, scores_shape): + return False + + # A fully-masked softmax row is undefined and makes eager aten._masked_softmax + # emit NaN (poisoning quant calibration), so decline the rewrite if the mask + # has one. Checked on the buffer rather than the gathered operand so a runtime + # gather can't select a fully-masked row we never inspected. + dim = get_arg(softmax, "dim", int) + if _has_fully_masked_row(source.buffer_tensor, dim, len(scores_shape)): + return False + + graph = ep.graph_module.graph + mask_node = _materialize_bool_mask(ep, source, softmax, bool_mask_cache) + + # mask_type 2 wants the mask to match the input shape exactly; the gathered + # [L_q, L_k] mask broadcasts over batch/heads, so expand it to scores shape. + if tuple(mask_node.meta["val"].shape) != tuple(scores_shape): + with graph.inserting_before(softmax): + expanded = graph.call_function( + _EXPAND, args=(mask_node, list(scores_shape)) + ) + expanded.meta["val"] = mask_node.meta["val"].expand(scores_shape) + mask_node = expanded + + with graph.inserting_before(softmax): + fused = graph.call_function( + _MASKED_SOFTMAX, args=(scores, mask_node, dim, _MASK_TYPE_FULL) + ) + + fused.meta = softmax.meta + + softmax.replace_all_uses_with(fused) + graph.erase_node(softmax) + if not add.users: + graph.erase_node(add) + return True diff --git a/backends/fused_quant/pre_quantize_passes/replace_batch_norm_with_conv.py b/backends/fused_quant/pre_quantize_passes/replace_batch_norm_with_conv.py new file mode 100644 index 00000000000..c9e2940a398 --- /dev/null +++ b/backends/fused_quant/pre_quantize_passes/replace_batch_norm_with_conv.py @@ -0,0 +1,287 @@ +# 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. + +from __future__ import annotations + +import operator +from dataclasses import dataclass +from typing import Optional + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import ( + add_constant, + compute_meta_val, + get_constant, + get_input_kind, +) +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch import fx +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +_BATCH_NORM_TARGET = torch.ops.aten._native_batch_norm_legit_no_training.default + +_CONVOLUTION = torch.ops.aten.convolution.default + + +@dataclass(frozen=True) +class _BatchNormConstants: + weight_node: Optional[fx.Node] + bias_node: Optional[fx.Node] + mean_node: fx.Node + var_node: fx.Node + weight: Optional[torch.Tensor] + bias: Optional[torch.Tensor] + running_mean: torch.Tensor + running_var: torch.Tensor + + +def _constant_kind( + exported_program: ExportedProgram, + *nodes: Optional[fx.Node], +) -> InputKind: + for node in nodes: + if node is not None: + kind = get_input_kind(exported_program, node) + if kind is not None: + return kind + return InputKind.CONSTANT_TENSOR + + +def _native_output_accessors(batch_norm: fx.Node) -> Optional[list[fx.Node]]: + """Return output-0 accessors when the tuple's auxiliary outputs are unused.""" + output_accessors: list[fx.Node] = [] + for user in batch_norm.users: + if user.target is not operator.getitem or not isinstance(user.args[1], int): + return None + if user.args[1] == 0: + output_accessors.append(user) + elif user.users: + return None + return output_accessors or None + + +def _constant_arg( + exported_program: ExportedProgram, + batch_norm: fx.Node, + name: str, +) -> tuple[Optional[fx.Node], Optional[torch.Tensor]]: + arg = get_arg(batch_norm, name) + if arg is None: + return None, None + if not isinstance(arg, fx.Node): + return None, None + return arg, get_constant(exported_program, arg) + + +def _get_constants( + exported_program: ExportedProgram, + batch_norm: fx.Node, +) -> Optional[_BatchNormConstants]: + weight_node, weight = _constant_arg(exported_program, batch_norm, "weight") + bias_node, bias = _constant_arg(exported_program, batch_norm, "bias") + mean_node, running_mean = _constant_arg( + exported_program, batch_norm, "running_mean" + ) + var_node, running_var = _constant_arg(exported_program, batch_norm, "running_var") + if ( + mean_node is None + or var_node is None + or running_mean is None + or running_var is None + or (weight_node is not None and weight is None) + or (bias_node is not None and bias is None) + ): + return None + + channels = running_mean.numel() + expected_shape = (channels,) + tensors = (running_mean, running_var, weight, bias) + if any(tensor is not None and tensor.shape != expected_shape for tensor in tensors): + return None + return _BatchNormConstants( + weight_node, + bias_node, + mean_node, + var_node, + weight, + bias, + running_mean, + running_var, + ) + + +def _make_fused_parameters( + constants: _BatchNormConstants, + eps: float, + dtype: torch.dtype, + spatial_dims: int, +) -> tuple[torch.Tensor, torch.Tensor]: + scale = torch.rsqrt(constants.running_var + eps) + if constants.weight is not None: + scale = scale * constants.weight + fused_bias = -constants.running_mean * scale + if constants.bias is not None: + fused_bias = fused_bias + constants.bias + + channels = constants.running_mean.numel() + fused_weight = scale.reshape(channels, 1, *([1] * spatial_dims)) + return ( + fused_weight.to(dtype=dtype).detach().contiguous(), + fused_bias.to(dtype=dtype).detach().contiguous(), + ) + + +def _add_fused_parameters( + exported_program: ExportedProgram, + batch_norm: fx.Node, + constants: _BatchNormConstants, + fused_weight: torch.Tensor, + fused_bias: torch.Tensor, +) -> tuple[fx.Node, fx.Node]: + weight_kind = _constant_kind( + exported_program, constants.weight_node, constants.var_node + ) + bias_kind = _constant_kind( + exported_program, constants.bias_node, constants.mean_node + ) + conv_weight = add_constant( + exported_program, + f"{batch_norm.name}_conv_weight", + fused_weight, + batch_norm, + kind=weight_kind, + ) + conv_bias = add_constant( + exported_program, + f"{batch_norm.name}_conv_bias", + fused_bias, + batch_norm, + kind=bias_kind, + ) + return conv_weight, conv_bias + + +def _insert_convolution( + batch_norm: fx.Node, + inp: fx.Node, + conv_weight: fx.Node, + conv_bias: fx.Node, + output_meta: dict[str, object], +) -> fx.Node: + graph = batch_norm.graph + inp_val = inp.meta["val"] + assert isinstance(inp_val, torch.Tensor) + spatial_dims = max(inp_val.ndim - 2, 1) + conv_input = inp + with graph.inserting_before(batch_norm): + if inp_val.ndim == 2: + conv_input = graph.call_function( + torch.ops.aten.unsqueeze.default, + args=(inp, -1), + ) + conv_input.meta = inp.meta.copy() + conv_input.meta["val"] = inp_val.unsqueeze(-1) + + conv = graph.call_function( + _CONVOLUTION, + args=( + conv_input, + conv_weight, + conv_bias, + [1] * spatial_dims, + [0] * spatial_dims, + [1] * spatial_dims, + False, + [0] * spatial_dims, + conv_weight.meta["val"].shape[0], + ), + ) + conv.meta = batch_norm.meta.copy() + conv_val = compute_meta_val(conv) + assert isinstance(conv_val, torch.Tensor) + conv.meta["val"] = conv_val + + if inp_val.ndim != 2: + return conv + squeeze = graph.call_function( + torch.ops.aten.squeeze.dim, + args=(conv, -1), + ) + squeeze.meta = output_meta.copy() + squeeze.meta["val"] = conv_val.squeeze(-1) + return squeeze + + +def _replace_outputs( + batch_norm: fx.Node, + output_accessors: list[fx.Node], + replacement: fx.Node, +) -> None: + graph = batch_norm.graph + for accessor in output_accessors: + accessor.replace_all_uses_with(replacement) + for accessor in list(batch_norm.users): + graph.erase_node(accessor) + + graph.erase_node(batch_norm) + + +class ReplaceBatchNormWithConv(ExportedProgramPassBase): + """Replace inference batch normalization with a depthwise convolution.""" + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + modified = False + for bn_node in graph.find_nodes(op="call_function", target=_BATCH_NORM_TARGET): + modified |= self._replace(exported_program, bn_node) + + if modified: + exported_program = constant_prop_pass(exported_program) + + return ExportedProgramPassResult(exported_program, modified) + + def _replace( + self, + exported_program: ExportedProgram, + batch_norm: fx.Node, + ) -> bool: + output_accessors = _native_output_accessors(batch_norm) + if output_accessors is None: + return False + + inp = get_arg(batch_norm, "input", fx.Node) + inp_val = inp.meta["val"] + if not isinstance(inp_val, torch.Tensor) or not 2 <= inp_val.ndim <= 5: + return False + + constants = _get_constants(exported_program, batch_norm) + if constants is None: + return False + + spatial_dims = max(inp_val.ndim - 2, 1) + eps = get_arg(batch_norm, "eps", float) + fused_weight, fused_bias = _make_fused_parameters( + constants, eps, inp_val.dtype, spatial_dims + ) + conv_weight, conv_bias = _add_fused_parameters( + exported_program, + batch_norm, + constants, + fused_weight, + fused_bias, + ) + replacement = _insert_convolution( + batch_norm, + inp, + conv_weight, + conv_bias, + output_accessors[0].meta, + ) + _replace_outputs(batch_norm, output_accessors, replacement) + return True diff --git a/backends/fused_quant/pre_quantize_passes/replace_mm_with_bmm.py b/backends/fused_quant/pre_quantize_passes/replace_mm_with_bmm.py new file mode 100644 index 00000000000..1d74aa00203 --- /dev/null +++ b/backends/fused_quant/pre_quantize_passes/replace_mm_with_bmm.py @@ -0,0 +1,75 @@ +# 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. + + +from __future__ import annotations + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import compute_meta_val +from torch import fx +from torch.fx.passes.infra.pass_base import PassBase, PassResult + + +class ReplaceMmWithBmm(PassBase): + """Replace ``aten.mm(x, y)`` with a unit-batch ``aten.bmm`` between views. + + ``mm`` has no ``fused_quant`` equivalent, so this lets the ``fused_quant::bmm`` + quantizer handle the matmul instead of leaving it as an unfused float op. + Adding and dropping a leading size-1 dimension is a valid view + for any input strides, so the reshapes cost nothing at runtime. + """ + + def call(self, graph_module: fx.GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + + for mm_node in graph.find_nodes( + op="call_function", target=torch.ops.aten.mm.default + ): + self._replace(mm_node) + modified = True + + if modified: + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + @staticmethod + def _replace(mm_node: fx.Node) -> None: + graph = mm_node.graph + inp = get_arg(mm_node, "input", fx.Node) + mat2 = get_arg(mm_node, "mat2", fx.Node) + inp_val = inp.meta["val"] + mat2_val = mat2.meta["val"] + + with graph.inserting_before(mm_node): + inp_3d = graph.call_function( + torch.ops.aten.view.default, args=(inp, [1, *inp_val.shape]) + ) + inp_3d.meta = mm_node.meta.copy() + inp_3d.meta["val"] = inp_val.unsqueeze(0) + + mat2_3d = graph.call_function( + torch.ops.aten.view.default, args=(mat2, [1, *mat2_val.shape]) + ) + mat2_3d.meta = mm_node.meta.copy() + mat2_3d.meta["val"] = mat2_val.unsqueeze(0) + + bmm = graph.call_function( + torch.ops.aten.bmm.default, args=(inp_3d, mat2_3d) + ) + bmm.meta = mm_node.meta.copy() + bmm.meta["val"] = compute_meta_val(bmm) + + out_shape = list(mm_node.meta["val"].shape) + squeezed = graph.call_function( + torch.ops.aten.view.default, args=(bmm, out_shape) + ) + squeezed.meta = mm_node.meta.copy() + + mm_node.replace_all_uses_with(squeezed) + graph.erase_node(mm_node) diff --git a/backends/fused_quant/quantizer/__init__.py b/backends/fused_quant/quantizer/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/fused_quant/quantizer/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/fused_quant/quantizer/defaults.py b/backends/fused_quant/quantizer/defaults.py new file mode 100644 index 00000000000..6b03b804d07 --- /dev/null +++ b/backends/fused_quant/quantizer/defaults.py @@ -0,0 +1,298 @@ +# 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. + +import dataclasses +from typing import Optional + +import torch +from executorch.backends.fused_quant.quantizer.qspecs import ( + arg1_is_scalar, + make_qconfig, +) +from executorch.backends.fused_quant.quantizer.quantizer import ( + FusedQuantQuantizer, + MaxPoolQuantizer, + OpQuantizer, + QuantizerBase, +) +from executorch.backends.transforms.permute_pass_utils import get_arg +from torch import fx +from torchao.quantization.pt2e.quantizer import QuantizationConfig + + +# Why add/mul/sub are each registered twice -- once for tensor operands and once +# for a scalar operand: +# +# PyTorch does NOT give scalar arithmetic its own op overload. `tensor + 1e-5` +# (e.g. RMSNorm's `variance + eps`) traces to `aten.add.Tensor(tensor, 1e-5)` -- +# the `.Tensor` overload carrying a Python scalar literal in arg 1, NOT a +# separate `aten.add.Scalar`. So the op type (`aten.add.Tensor`) alone does not +# tell us whether arg 1 is a tensor or a scalar; both land on the same op. +# +# Annotating a scalar arg as a quant edge breaks PT2E prepare, so the two operand +# shapes are handled separately: +# - activation_names=("input", "other") quantizes both tensor operands; the +# float-tensor guard makes it skip a scalar arg 1. +# - activation_names=("input",) + node_filter=arg1_is_scalar quantizes only the +# tensor input and leaves the scalar unquantized. +# Between them every add/mul/sub is covered. The scalar variant folds into +# `fused_quant.{add,mul,sub}.Scalar`, which pass the scalar through as a trailing +# arg. +def _get_binary_quantizers(config: QuantizationConfig) -> list[QuantizerBase]: + quantizers: list[QuantizerBase] = [] + for aten_op, tensor_op, scalar_op in ( + ( + torch.ops.aten.add.Tensor, + torch.ops.fused_quant.add.default, + torch.ops.fused_quant.add.Scalar, + ), + ( + torch.ops.aten.mul.Tensor, + torch.ops.fused_quant.mul.default, + torch.ops.fused_quant.mul.Scalar, + ), + ( + torch.ops.aten.sub.Tensor, + torch.ops.fused_quant.sub.default, + torch.ops.fused_quant.sub.Scalar, + ), + ): + quantizers.append( + OpQuantizer(aten_op, tensor_op, config, activation_names=("input", "other")) + ) + quantizers.append( + OpQuantizer( + aten_op, + scalar_op, + config, + activation_names=("input",), + node_filter=arg1_is_scalar, + ) + ) + return quantizers + + +def get_core_quantizers(config: QuantizationConfig) -> list[QuantizerBase]: + """Quantizers for the ops every backend is expected to handle. + + Everything is quantized with the single ``config``, so weights follow whatever + granularity that config specifies. Use this for backends without per-channel + support or without kernels for the wider op set. + """ + return [ + *_get_binary_quantizers(config), + OpQuantizer( + torch.ops.aten.bmm.default, + torch.ops.fused_quant.bmm.default, + config, + activation_names=("input", "mat2"), + ), + OpQuantizer( + torch.ops.aten.relu.default, torch.ops.fused_quant.relu.default, config + ), + OpQuantizer( + torch.ops.aten.linear.default, + torch.ops.fused_quant.linear.default, + config, + activation_names=("input",), + weight_names=("weight",), + other_names=("bias",), + ), + OpQuantizer( + torch.ops.aten.conv2d.default, + torch.ops.fused_quant.conv2d.default, + config, + activation_names=("input",), + weight_names=("weight",), + other_names=("bias",), + ), + ] + + +def _transposed_convolution_quantizers( + default_config: QuantizationConfig, + weighted_config: QuantizationConfig, +) -> list[QuantizerBase]: + """Quantizers for ``aten.convolution`` with ``transposed=True``. + + A transposed convolution's weight is ``[C_in, C_out // groups, ...]``, so + per-output-channel weight scales sit on axis 1, not axis 0. With groups > 1 + axis 1 holds only ``C_out // groups`` channels, so those fall back to the + per-tensor weights of ``default_config``. Listed before the general + convolution quantizer so they claim transposed nodes first. + """ + weight = weighted_config.weight + if weight is None or weight.qscheme not in ( + torch.per_channel_affine, + torch.per_channel_symmetric, + ): + return [] + axis_1 = dataclasses.replace( + weighted_config, weight=dataclasses.replace(weight, ch_axis=1) + ) + + def transposed(single_group: bool): + def node_filter(node: fx.Node) -> bool: + return ( + bool(get_arg(node, "transposed")) + and (get_arg(node, "groups") == 1) == single_group + ) + + return node_filter + + return [ + OpQuantizer( + torch.ops.aten.convolution.default, + torch.ops.fused_quant.convolution.default, + config, + activation_names=("input",), + weight_names=("weight",), + other_names=("bias",), + node_filter=transposed(single_group), + ) + for config, single_group in ((axis_1, True), (default_config, False)) + ] + + +def get_default_quantizers( + default_config: QuantizationConfig, + weighted_config: QuantizationConfig, +) -> list[QuantizerBase]: + """Quantizers for every ATen op that has a ``fused_quant`` equivalent. + + Weighted ops (conv, linear) use ``weighted_config``; everything else uses + ``default_config``. Both are ordinary ``QuantizationConfig`` values -- the + split is by which ops they apply to, not by granularity, and nothing here + requires either one to be per-tensor or per-channel. + """ + weighted_ops = ( + (torch.ops.aten.conv1d.default, torch.ops.fused_quant.conv1d.default), + (torch.ops.aten.conv2d.default, torch.ops.fused_quant.conv2d.default), + (torch.ops.aten.conv3d.default, torch.ops.fused_quant.conv3d.default), + ( + torch.ops.aten.convolution.default, + torch.ops.fused_quant.convolution.default, + ), + (torch.ops.aten.linear.default, torch.ops.fused_quant.linear.default), + ) + unary_ops = ( + (torch.ops.aten.hardswish.default, torch.ops.fused_quant.hardswish.default), + (torch.ops.aten.sigmoid.default, torch.ops.fused_quant.sigmoid.default), + (torch.ops.aten.tanh.default, torch.ops.fused_quant.tanh.default), + (torch.ops.aten.hardtanh.default, torch.ops.fused_quant.hard_tanh.default), + (torch.ops.aten.silu.default, torch.ops.fused_quant.silu.default), + ( + torch.ops.aten.hardsigmoid.default, + torch.ops.fused_quant.hardsigmoid.default, + ), + (torch.ops.aten.gelu.default, torch.ops.fused_quant.gelu.default), + ( + torch.ops.aten.native_layer_norm.default, + torch.ops.fused_quant.native_layer_norm.default, + ), + (torch.ops.aten.rms_norm.default, torch.ops.fused_quant.rms_norm.default), + # Masked softmax: only the scores (input) and output are quantized; the + # bool mask is left unquantized (threaded through, never quantized). + ( + torch.ops.aten._masked_softmax.default, + torch.ops.fused_quant._masked_softmax.default, + ), + ) + return [ + *_get_binary_quantizers(default_config), + OpQuantizer( + torch.ops.aten.bmm.default, + torch.ops.fused_quant.bmm.default, + default_config, + activation_names=("input", "mat2"), + ), + OpQuantizer( + torch.ops.aten.relu.default, + torch.ops.fused_quant.relu.default, + default_config, + ), + MaxPoolQuantizer(default_config), + OpQuantizer( + torch.ops.aten.avg_pool2d.default, + torch.ops.fused_quant.avg_pool2d.default, + default_config, + ), + *_transposed_convolution_quantizers(default_config, weighted_config), + *( + OpQuantizer( + aten_op, + fused_op, + weighted_config, + activation_names=("input",), + weight_names=("weight",), + other_names=("bias",), + ) + for aten_op, fused_op in weighted_ops + ), + *( + OpQuantizer(aten_op, fused_op, default_config) + for aten_op, fused_op in unary_ops + ), + ] + + +def make_fused_quant_quantizer( + overrides: Optional[list[QuantizerBase]] = None, + *, + core_only: bool = False, + weight_per_channel: bool = True, + act_dtype: torch.dtype = torch.int8, + weight_dtype: torch.dtype = torch.int8, + act_symmetric: bool = False, + weight_symmetric: bool = True, + act_qmin: Optional[int] = None, + act_qmax: Optional[int] = None, + weight_qmin: Optional[int] = None, + weight_qmax: Optional[int] = None, +) -> FusedQuantQuantizer: + """Build a quantizer from override quantizers layered on top of the defaults. + + overrides are prepended to the defaults, so an overridden node is claimed + by its override first (annotation is first-match-wins on key presence) and the + default for that op skips it. An override is an ordinary OpQuantizer (or + NoopQuantizer) carrying a node_filter that selects the nodes it owns. + + Op-type-level changes do not need an override: pass the desired dtype / + symmetry / bounds knobs and the defaults handle them. Reach for an override + only when a specific module or FX node needs a different config than others of + the same op. Overrides are ordered; the first matching override claims the node. + + Args: + core_only: Restrict the defaults to :func:`get_core_quantizers`. + weight_per_channel: Quantize conv/linear weights per output channel. + act_dtype / weight_dtype: Integer container dtypes for activations and + weights (quantized independently). Default int8. + act_symmetric / weight_symmetric: Use a symmetric qscheme (zero-point + pinned to 0) for activations / weights. + act_qmin / act_qmax / weight_qmin / weight_qmax: Explicit quant bounds. + Activations and asymmetric weights default to the corresponding + dtype's full ``torch.iinfo`` range. Signed symmetric weights default + to a balanced range (e.g. -127/127 for int8). + """ + knobs = { + "act_dtype": act_dtype, + "weight_dtype": weight_dtype, + "act_symmetric": act_symmetric, + "weight_symmetric": weight_symmetric, + "act_qmin": act_qmin, + "act_qmax": act_qmax, + "weight_qmin": weight_qmin, + "weight_qmax": weight_qmax, + } + config = make_qconfig(weight_per_channel=weight_per_channel, **knobs) + defaults = ( + get_core_quantizers(config) + if core_only + else get_default_quantizers( + make_qconfig(weight_per_channel=False, **knobs), config + ) + ) + return FusedQuantQuantizer([*(overrides or []), *defaults]) diff --git a/backends/fused_quant/quantizer/qspecs.py b/backends/fused_quant/quantizer/qspecs.py new file mode 100644 index 00000000000..7fec83d77b8 --- /dev/null +++ b/backends/fused_quant/quantizer/qspecs.py @@ -0,0 +1,197 @@ +# 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. + +from typing import Optional + +import torch +from torch import fx +from torchao.quantization.granularity import PerBlock +from torchao.quantization.pt2e import ( + HistogramObserver, + MinMaxObserver, + PerChannelMinMaxObserver, +) +from torchao.quantization.pt2e._affine_quantization import AffineQuantizedMinMaxObserver +from torchao.quantization.pt2e.quantizer import QuantizationConfig, QuantizationSpec +from torchao.quantization.quant_primitives import MappingType + + +def make_per_tensor_qspec( + dtype: torch.dtype, + quant_min: int, + quant_max: int, + *, + histogram_observer: bool = False, + symmetric: bool = False, + eps: float | None = None, +) -> QuantizationSpec: + """Build a per-tensor ``QuantizationSpec`` (one scale/zero-point per tensor). + + Args: + symmetric: When True the qscheme is ``per_tensor_symmetric`` (zero-point + pinned to 0); otherwise ``per_tensor_affine`` (asymmetric). + histogram_observer: When True use ``HistogramObserver``; otherwise use ``MinMaxObserver``. + eps: optional floor on the computed scale (forwarded to the observer); + leave None to use the observer's default. Activation observers want an + explicit eps=2**-12: with a smaller floor a small-range activation can + produce a tiny scale, which makes affine-transform parameter extraction + infeasible on backends that approximate activations with a lookup table + (e.g. HardSigmoid). + """ + base_observer = HistogramObserver if histogram_observer else MinMaxObserver + observer_ctr = ( + base_observer.with_args(eps=eps) if eps is not None else base_observer + ) + return QuantizationSpec( + dtype=dtype, + quant_min=quant_min, + quant_max=quant_max, + qscheme=(torch.per_tensor_symmetric if symmetric else torch.per_tensor_affine), + is_dynamic=False, + observer_or_fake_quant_ctr=observer_ctr, + ) + + +def make_per_channel_qspec( + dtype: torch.dtype, + quant_min: int, + quant_max: int, + *, + ch_axis: int = 0, + eps: float = 2**-12, + symmetric: bool = False, +) -> QuantizationSpec: + """Build a per-channel ``QuantizationSpec`` (one scale/zero-point per slice + along ``ch_axis``; axis 0 is the output-channel axis for conv/linear weights). + + Args: + symmetric: When True the qscheme is ``per_channel_symmetric``; otherwise + ``per_channel_affine`` (asymmetric). + """ + return QuantizationSpec( + dtype=dtype, + quant_min=quant_min, + quant_max=quant_max, + qscheme=( + torch.per_channel_symmetric if symmetric else torch.per_channel_affine + ), + ch_axis=ch_axis, + is_dynamic=False, + observer_or_fake_quant_ctr=PerChannelMinMaxObserver.with_args(eps=eps), + ) + + +def make_blockwise_qspec( + dtype: torch.dtype, + quant_min: int, + quant_max: int, + *, + block_size: tuple[int, ...], + symmetric: bool = True, + eps: float = 2**-12, +) -> QuantizationSpec: + """Build a blockwise ``QuantizationSpec`` whose ``convert`` emits + ``torchao.quantize_affine`` / ``dequantize_affine`` nodes with an explicit + ``block_size``. + + ``block_size`` is right-aligned to the tensor's trailing dims (leading dims + pad with 1), so axis ``i`` gets ``shape[i] // block_size[i]`` scales -- i.e. + arbitrary rectangular blocks, not just a 1-D group. The common per-row group + case is ``block_size=(1, group_size)`` (one scale per ``group_size`` elements + along the last dim), giving a 2-D weight a ``[rows, cols // group_size]`` + scale; ``block_size=(2, 4)`` instead tiles a ``[R, C]`` weight into + ``[R//2, C//4]`` blocks. Each block dim must divide the tensor dim. + + Unlike the per-tensor / per-channel specs this carries no ``qscheme`` / + ``ch_axis`` (blockwise has no torch.qscheme): the upstream + ``AffineQuantizedMinMaxObserver`` owns node insertion via its ``convert`` hook + (which ``convert_pt2e`` calls when an observer defines one), driven by the + ``PerBlock`` granularity bound here. ``dtype`` is the integer container (e.g. + int8 with quant_min/max -8/7 for W4). + """ + mapping_type = MappingType.SYMMETRIC if symmetric else MappingType.ASYMMETRIC + return QuantizationSpec( + dtype=dtype, + quant_min=quant_min, + quant_max=quant_max, + qscheme=None, + ch_axis=None, + is_dynamic=False, + observer_or_fake_quant_ctr=AffineQuantizedMinMaxObserver.with_args( + mapping_type=mapping_type, + target_dtype=dtype, + granularity=PerBlock(block_size=tuple(block_size)), + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + ), + ) + + +def make_qconfig( + *, + weight_per_channel: bool, + act_dtype: torch.dtype = torch.int8, + weight_dtype: torch.dtype = torch.int8, + act_symmetric: bool = False, + weight_symmetric: bool = True, + act_qmin: Optional[int] = None, + act_qmax: Optional[int] = None, + weight_qmin: Optional[int] = None, + weight_qmax: Optional[int] = None, +) -> QuantizationConfig: + """Build a default quant config. + + Activations are per-tensor histogram (the same qspec instance is shared for + input and output); the weight qspec is per-channel along axis 0 when + ``weight_per_channel`` is set, otherwise per-tensor. Activation and weight + dtypes / symmetry are independent. Explicit quantization bounds are used + verbatim (e.g. -8/7 to pack W4 into int8). By default, signed symmetric + weights use a balanced range (e.g. -127/127 for int8); all other defaults + use the dtype's full range. + """ + act_info = torch.iinfo(act_dtype) + weight_info = torch.iinfo(weight_dtype) + act_qmin = act_qmin if act_qmin is not None else act_info.min + act_qmax = act_qmax if act_qmax is not None else act_info.max + weight_qmax = weight_qmax if weight_qmax is not None else weight_info.max + weight_qmin = ( + weight_qmin + if weight_qmin is not None + else ( + -weight_qmax + if weight_symmetric and weight_info.min < 0 + else weight_info.min + ) + ) + # eps=2**-12 floors the activation scale; without it a small-range activation + # gets a tiny scale that breaks affine-transform parameter extraction on + # backends that approximate activations with a lookup table (e.g. HardSigmoid + # in mobilenetv3). + act_qspec = make_per_tensor_qspec( + act_dtype, + act_qmin, + act_qmax, + histogram_observer=True, + symmetric=act_symmetric, + eps=2**-12, + ) + weight_qspec = ( + make_per_channel_qspec( + weight_dtype, weight_qmin, weight_qmax, symmetric=weight_symmetric + ) + if weight_per_channel + else make_per_tensor_qspec( + weight_dtype, weight_qmin, weight_qmax, symmetric=weight_symmetric + ) + ) + return QuantizationConfig(act_qspec, act_qspec, weight_qspec, None) + + +def arg1_is_scalar(node: fx.Node) -> bool: + """A node_filter for the .Scalar binary variants: True when the second + operand is a Python scalar (e.g. x + eps) rather than a tensor.""" + return not isinstance(node.args[1], fx.Node) diff --git a/backends/fused_quant/quantizer/quantizer.py b/backends/fused_quant/quantizer/quantizer.py new file mode 100644 index 00000000000..791d47ed7a1 --- /dev/null +++ b/backends/fused_quant/quantizer/quantizer.py @@ -0,0 +1,385 @@ +# 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-strict + +import re +from abc import ABC, abstractmethod +from typing import Callable, cast, Iterable, Optional + +import torch +from executorch.backends.fused_quant.fuse_aten import arg_names, fuse_aten, output_node +from executorch.backends.fused_quant.graph_utils import is_dequantize_node +from executorch.backends.fused_quant.pre_quantize_passes.replace_mm_with_bmm import ( + ReplaceMmWithBmm, +) +from executorch.backends.transforms.permute_pass_utils import get_arg +from torch import fx +from torch._export.utils import _detect_fake_mode_from_gm +from torch._ops import OpOverload +from torch.fx.passes.fake_tensor_prop import FakeTensorProp +from torch.fx.passes.tools_common import stable_topological_sort +from torchao.quantization.pt2e.quantizer import ( + ComposableQuantizer, + QuantizationConfig, + Quantizer, + SharedQuantizationSpec, +) +from torchao.quantization.pt2e.quantizer.quantizer import ( + Q_ANNOTATION_KEY, + QuantizationAnnotation, +) +from torchao.quantization.pt2e.quantizer.utils import ( + annotate_input_qspec_map, + annotate_output_qspec, +) + + +# Node meta flag set when a fused_quant quantizer claimed a node, so fuse() can +# tell its own claims from another PT2E quantizer's (which also set _annotated). +_OWNER_KEY = "fused_quant_owned" + + +class QuantizerBase(Quantizer, ABC): + """A torchao Quantizer that also fuses.""" + + def validate(self, model: fx.GraphModule) -> None: + """Unused; satisfies the torchao Quantizer contract.""" + + @abstractmethod + def annotate(self, model: fx.GraphModule) -> fx.GraphModule: + """Stamp QuantizationAnnotation meta on the ops this quantizer owns.""" + ... + + @abstractmethod + def fuse(self, gm: fx.GraphModule) -> bool: + """Fold each dequant -> op -> quant into a fused op; return whether the + graph changed.""" + ... + + @abstractmethod + def preserved_ops(self) -> Iterable[OpOverload]: + """Ops this quantizer matches, kept intact through decomposition.""" + ... + + +class OpQuantizer(QuantizerBase): + """A single-op quantizer implementation: annotates and fuses a single op type. + + Args: + op: The ATen op to match, annotate, and preserve from decomposition. + fused_op: The fused_quant op to lower matched instances into. + config: The per-instance quantization config. + activation_names: Schema names of quantized activation inputs (annotated + with config.input_activation). + weight_names: Schema names of quantized weight inputs (annotated with + config.weight). + other_names: Schema names of tensor passthroughs (e.g. "bias") that + occupy the fused op's tensor prefix with a null qparams block. + output_indices: Outputs to quantize with config.output_activation. Defaults + to just output 0; a multi-output op can pass e.g. (0, 1). All quantized + outputs share the single config.output_activation. + node_filter: Optional predicate restricting which op nodes this quantizer + owns. + """ + + def __init__( + self, + op: OpOverload, + fused_op: OpOverload, + config: QuantizationConfig, + *, + activation_names: tuple[str, ...] = ("input",), + weight_names: tuple[str, ...] = (), + other_names: tuple[str, ...] = (), + output_indices: tuple[int, ...] = (0,), + node_filter: Optional[Callable[[fx.Node], bool]] = None, + ) -> None: + # Validate arg names. + valid_names = arg_names(op) + for name in (*activation_names, *weight_names, *other_names): + if name not in valid_names: + raise ValueError( + f"{type(self).__name__}: {name!r} is not an argument of {op}; " + f"expected one of {valid_names}" + ) + + self.op = op + self.fused_op = fused_op + self.config = config + self.activation_names = activation_names + self.weight_names = weight_names + self.other_names = other_names + self.output_indices = output_indices + self.node_filter = node_filter + + def matches(self, node: fx.Node) -> bool: + """Whether this quantizer handles node, the shared gate for annotate and fuse. + + Requires node to pass node_filter and every quantized edge to be a float tensor. + """ + if self.node_filter is not None and not self.node_filter(node): + return False + + def is_float_tensor(x: fx.Node) -> bool: + # A dequantize node's output is float by definition; convert-inserted + # dq/q nodes may not carry meta["val"], so check the target directly. + if is_dequantize_node(x): + return True + return ( + isinstance(x.meta.get("val"), torch.Tensor) + and x.meta["val"].is_floating_point() + ) + + for name in (*self.activation_names, *self.weight_names): + # get_arg is not necessarily Node type, e.g. aten.add with scalar rhs. + arg = get_arg(node, name) + if not isinstance(arg, fx.Node) or not is_float_tensor(arg): + return False + if self.config.output_activation is not None: + for output_index in self.output_indices: + out = output_node(node, output_index) + if out is not None and not is_float_tensor(out): + return False + return True + + def annotate(self, model: fx.GraphModule) -> fx.GraphModule: + """Annotate every matching node in the graph. Skips already annotated nodes. + + Annotations are marked ``_annotated`` -- the flag other PT2E quantizers + check before claiming a node -- so composing with them is first-claim-wins + in either order. + """ + for node in model.graph.find_nodes(op="call_function", target=self.op): + if ( + node.meta.get("source_fn_stack") is not None + and Q_ANNOTATION_KEY not in node.meta + and self.matches(node) + ): + self.annotate_node(node) + self._mark_owned(node) + return model + + def _mark_owned(self, node: fx.Node) -> None: + node.meta[_OWNER_KEY] = True + outputs = [output_node(node, index) for index in self.output_indices] + for annotated in (node, *outputs): + if annotated is None: + continue + annotation = annotated.meta.get(Q_ANNOTATION_KEY) + if annotation is not None: + annotation._annotated = True + + def annotate_node(self, node: fx.Node) -> None: + """Annotate node. Subclasses may override this for custom annotation.""" + for name in self.activation_names: + annotate_input_qspec_map( + node, cast(fx.Node, get_arg(node, name)), self.config.input_activation + ) + for name in self.weight_names: + annotate_input_qspec_map( + node, cast(fx.Node, get_arg(node, name)), self.config.weight + ) + # output_node is None if a multi-output op's getitem(output_index) is unused. + if self.config.output_activation is not None: + for output_index in self.output_indices: + out = output_node(node, output_index) + if out is not None: + annotate_output_qspec(out, self.config.output_activation) + + def fuse(self, gm: fx.GraphModule) -> bool: + """Fold each annotated op node this quantizer matches into fused_op. + + A node is fused only when it matches and carries an annotation with a live + qspec, so a node claimed by a no-op annotation or one merely sandwiched + between its neighbors' dq/q is left alone. A node another PT2E quantizer + claimed (``_annotated`` set, but not by a fused_quant quantizer) is never + fused. + """ + modified = False + for node in list(gm.graph.find_nodes(op="call_function", target=self.op)): + if not self.matches(node): + continue + ann = node.meta.get(Q_ANNOTATION_KEY) + if ann is None or (ann._annotated and not node.meta.get(_OWNER_KEY)): + continue + if ( + not any(v is not None for v in ann.input_qspec_map.values()) + and ann.output_qspec is None + ): + continue + self.fuse_node(gm, node) + modified = True + return modified + + def fuse_node(self, gm: fx.GraphModule, node: fx.Node) -> fx.Node: + """Fuse a generic aten op to a matching fused_quant op. + + Subclasses may override for custom fusion logic. + """ + return fuse_aten( + node, + self.fused_op, + self.activation_names, + self.weight_names, + self.other_names, + self.output_indices, + ) + + def preserved_ops(self) -> Iterable[OpOverload]: + return (self.op,) + + +class NoopQuantizer(OpQuantizer): + """Annotates matching ops with a no-op config so that later quantizers do not + touch them; must be listed before the quantizer it suppresses. + + A skipped op stays an unfused float aten op, so the later to_edge decomposes + it into core-aten; the backend must have a kernel for that or compilation + fails at lowering. + """ + + def __init__(self, op: OpOverload, node_filter: Callable[[fx.Node], bool]) -> None: + super().__init__( + op, + op, # fused_op target, unused. + QuantizationConfig(None, None, None, None), + activation_names=(), + node_filter=node_filter, + ) + + def annotate_node(self, node: fx.Node) -> None: + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(input_qspec_map={}) + + def fuse(self, gm: fx.GraphModule) -> bool: + return False + + +class BmmQuantizer(OpQuantizer): + """Quantize ``aten.bmm``, first rewriting every ``aten.mm`` into a unit-batch bmm. + + ``mm`` has no ``fused_quant`` equivalent, so left alone it stays an unfused + float op. Rewriting it in the pre-annotation transform puts the matmul on the + fused bmm path -- the new bmm nodes are then annotated like any other. + """ + + def __init__(self, config: QuantizationConfig) -> None: + super().__init__( + torch.ops.aten.bmm.default, + torch.ops.fused_quant.bmm.default, + config, + activation_names=("input", "mat2"), + ) + + def transform_for_annotation(self, model: fx.GraphModule) -> fx.GraphModule: + ReplaceMmWithBmm().call(model) + return model + + +class MaxPoolQuantizer(OpQuantizer): + """max_pool2d_with_indices: output 0 (the pooled values, via getitem 0) shares + the input's observer; the indices output is left unquantized.""" + + def __init__(self, config: QuantizationConfig) -> None: + super().__init__( + torch.ops.aten.max_pool2d_with_indices.default, + torch.ops.fused_quant.max_pool2d_with_indices.default, + config, + ) + + def annotate_node(self, node: fx.Node) -> None: + """Annotate as usual, then re-point output 0 to share input 0's observer.""" + super().annotate_node(node) + out = output_node(node) + if out is not None: + annotate_output_qspec( + out, + SharedQuantizationSpec((cast(fx.Node, node.args[0]), node)), + ) + + +class FusedQuantQuantizer(ComposableQuantizer, QuantizerBase): + """Composes several QuantizerBase instances into a single Quantizer. + + Order matters: annotation runs in list order and the first quantizer to claim a + node wins (later ones skip an already-annotated node). A quantizer that narrows + or suppresses another, e.g. a node_filtered override or a NoopQuantizer, must + come before the general one it overrides. + """ + + def __init__(self, quantizers: list[QuantizerBase]) -> None: + for q in quantizers: + if not isinstance(q, QuantizerBase): + raise TypeError(f"{type(q).__name__} is not a QuantizerBase") + super().__init__(list(quantizers)) + + def fuse(self, gm: fx.GraphModule) -> bool: + """Fuse via every child quantizer; return whether anything fused. + + ``gm`` is the output of ``convert_pt2e``, which is first cleaned up so + the children can match on it; see :func:`_repair_convert_output`. + """ + _repair_convert_output(gm) + modified = False + for q in self.quantizers: + modified |= cast(QuantizerBase, q).fuse(gm) + if modified: + stable_topological_sort(gm) + gm.graph.eliminate_dead_code() + gm.recompile() + return modified + + def preserved_ops(self) -> list[OpOverload]: + """The union of every child quantizer's preserved ops.""" + return [ + op for q in self.quantizers for op in cast(QuantizerBase, q).preserved_ops() + ] + + +def _repair_convert_output(gm: fx.GraphModule) -> None: + """Work around two ``convert_pt2e`` (torchao) gaps that fusion trips on. + + ``AffineQuantizedObserverBase.convert`` inserts ``quantize_affine`` / + ``dequantize_affine`` as the ``OpOverloadPacket`` rather than ``.default``, + which fusion does not match. And ``convert_pt2e`` rewrites the graph without + refreshing ``meta["val"]``, which fusion reads for shapes and dtypes. + """ + replacements = { + torch.ops.torchao.quantize_affine: torch.ops.torchao.quantize_affine.default, + torch.ops.torchao.dequantize_affine: torch.ops.torchao.dequantize_affine.default, + } + for node in gm.graph.nodes: + if node.op == "call_function" and node.target in replacements: + node.target = replacements[node.target] + gm.recompile() + + # A hand-built graph may carry no fake values at all; there is nothing to refresh. + placeholders = gm.graph.find_nodes(op="placeholder") + fake_mode = _detect_fake_mode_from_gm(gm) + if fake_mode is None or any("val" not in node.meta for node in placeholders): + return + FakeTensorProp(gm, mode=fake_mode).propagate_dont_convert_inputs( + *(node.meta["val"] for node in placeholders) + ) + + +# ============================================================================= +# node_filter predicates +# ============================================================================= + + +def module_fqn_filter(pattern: str) -> Callable[[fx.Node], bool]: + """A node_filter matching nodes whose module FQN matches pattern (regex + search over nn_module_stack).""" + regex = re.compile(pattern) + + def _filter(node: fx.Node) -> bool: + return any( + regex.search(fqn) is not None + for fqn, _ in node.meta.get("nn_module_stack", {}).values() + ) + + return _filter diff --git a/backends/fused_quant/targets.bzl b/backends/fused_quant/targets.bzl new file mode 100644 index 00000000000..c9668237f13 --- /dev/null +++ b/backends/fused_quant/targets.bzl @@ -0,0 +1,471 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +# fused_quant is a backend-neutral quantization frontend, so it is consumed both +# from inside ExecuTorch and from downstream compilers built on top of it. +_VISIBILITY = ["PUBLIC"] + +def define_common_targets(): + runtime.python_library( + name = "ops_utils", + srcs = ["ops_utils.py"], + visibility = _VISIBILITY, + deps = [ + "//caffe2:torch", + ], + ) + + runtime.python_library( + name = "ops", + srcs = ["ops.py"], + visibility = _VISIBILITY, + deps = [ + ":ops_utils", + "//caffe2:torch", + "//executorch/exir:pass_base", + "//pytorch/ao:torchao", + ], + ) + + runtime.python_library( + name = "graph_utils", + srcs = ["graph_utils.py"], + visibility = _VISIBILITY, + deps = [ + ":ops", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir/dialects/edge:lib", + ], + ) + + runtime.python_library( + name = "pass_base", + srcs = ["pass_base.py"], + visibility = _VISIBILITY, + deps = [ + "//executorch/exir:pass_base", + "//executorch/exir:pass_manager", + ], + ) + + runtime.python_library( + name = "fuse_aten", + srcs = ["fuse_aten.py"], + visibility = _VISIBILITY, + deps = [ + ":graph_utils", + ":ops", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + ], + ) + + runtime.python_library( + name = "quantizer", + srcs = [ + "quantizer/__init__.py", + "quantizer/quantizer.py", + ], + visibility = _VISIBILITY, + deps = [ + ":replace_mm_with_bmm", + ":fuse_aten", + ":graph_utils", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//pytorch/ao:torchao", + ], + ) + + runtime.python_library( + name = "colorer", + srcs = ["colorer.py"], + visibility = _VISIBILITY, + deps = [ + ":pass_base", + "//caffe2:torch", + "//executorch/exir:pass_base", + "//executorch/exir/dialects/edge:lib", + ], + ) + + runtime.python_library( + name = "qspecs", + srcs = ["quantizer/qspecs.py"], + visibility = _VISIBILITY, + deps = [ + "//caffe2:torch", + "//pytorch/ao:torchao", + ], + ) + + runtime.python_library( + name = "quantizer_defaults", + srcs = ["quantizer/defaults.py"], + visibility = _VISIBILITY, + deps = [ + ":qspecs", + ":quantizer", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//pytorch/ao:torchao", + ], + ) + + runtime.python_library( + name = "frontend", + srcs = ["frontend.py"], + visibility = _VISIBILITY, + deps = [ + ":fold_batch_norm", + ":fuse_add_softmax_into_masked_softmax", + ":quantizer", + ":replace_batch_norm_with_conv", + "//caffe2:torch", + "//pytorch/ao:torchao", + ], + ) + + runtime.python_library( + name = "passes", + srcs = ["passes.py"], + visibility = _VISIBILITY, + deps = [ + ":constant_fold", + ":fold_scalar_mul", + ":fuse_add_into_linear", + ":fuse_mul_into_linear", + ":graph_utils", + ":hoist_activations", + ":ops", + ":pass_base", + ":prequantize_embedding", + ":quant_absorption", + ":replace_dequant_quant_with_requantize", + ":sink_constant_cat", + ":split_linear_at_slices", + ":to_channels_last", + ":to_convolution", + "//caffe2:torch", + "//executorch/backends/transforms:advance_quantize_op_above_def_chain", + "//executorch/backends/transforms:advance_quantize_op_above_def_in_branch", + "//executorch/backends/transforms:fuse_cascaded_transpose_or_permute_ops", + "//executorch/backends/transforms:fuse_cascaded_view_ops", + "//executorch/backends/transforms:fuse_quant_dequant_to_requantize", + "//executorch/backends/transforms:merge_split_concat_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:postpone_dequantize_op_below_use_chain", + "//executorch/backends/transforms:propagate_slice", + "//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_as_strided_copy_op", + "//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_permutes_around_elementwise_ops", + "//executorch/backends/transforms:remove_zero_sized_cat_args", + "//executorch/backends/transforms:replace_nop_transpose_or_permute_with_view", + "//executorch/backends/transforms:replace_select_with_view_op", + "//executorch/backends/transforms:replace_split_with_slice", + "//executorch/backends/transforms:replace_squeeze_unsqueeze_with_view", + "//executorch/backends/transforms:split_dequantized_cat", + "//executorch/exir:pass_manager", + "//executorch/exir/dialects:lib", + "//executorch/exir/dialects/edge:lib", + ], + ) + + runtime.python_library( + name = "compile", + srcs = ["compile.py"], + visibility = _VISIBILITY, + deps = [ + ":frontend", + ":passes", + ":quantizer", + ":quantizer_defaults", + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/exir:pass_manager", + "//pytorch/ao:torchao", + ], + ) + + runtime.python_library( + name = "decompose_fused_quant", + srcs = ["decompose_fused_quant.py"], + visibility = _VISIBILITY, + deps = [ + ":colorer", + ":graph_utils", + ":ops", + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/dialects/edge:lib", + ], + ) + + # ------------------------------------------------------------------ + # Pre-quantization passes: rewrites applied to the traced ExportedProgram + # before the quantizer annotates it. + # ------------------------------------------------------------------ + + runtime.python_library( + name = "replace_mm_with_bmm", + srcs = ["pre_quantize_passes/replace_mm_with_bmm.py"], + visibility = _VISIBILITY, + deps = [ + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + ], + ) + + runtime.python_library( + name = "fold_batch_norm", + srcs = ["pre_quantize_passes/fold_batch_norm.py"], + visibility = _VISIBILITY, + deps = [ + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "fuse_add_softmax_into_masked_softmax", + srcs = ["pre_quantize_passes/fuse_add_softmax_into_masked_softmax.py"], + visibility = _VISIBILITY, + deps = [ + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "replace_batch_norm_with_conv", + srcs = ["pre_quantize_passes/replace_batch_norm_with_conv.py"], + visibility = _VISIBILITY, + deps = [ + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + # ------------------------------------------------------------------ + # Optimization passes: rewrites applied to the fused_quant edge graph. + # ------------------------------------------------------------------ + + runtime.python_library( + name = "constant_fold", + srcs = ["optimization_passes/constant_fold.py"], + visibility = _VISIBILITY, + deps = [ + ":pass_base", + "//caffe2:torch", + "//executorch/exir:pass_base", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "hoist_activations", + srcs = ["optimization_passes/hoist_activations.py"], + visibility = _VISIBILITY, + deps = [ + ":ops", + "//caffe2:torch", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_library( + name = "to_channels_last", + srcs = ["optimization_passes/to_channels_last.py"], + visibility = _VISIBILITY, + deps = [ + ":colorer", + ":graph_utils", + ":ops", + ":ops_utils", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_library( + name = "to_channels_first", + srcs = ["optimization_passes/to_channels_first.py"], + visibility = _VISIBILITY, + deps = [ + ":colorer", + ":graph_utils", + ":ops", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_library( + name = "to_convolution", + srcs = ["optimization_passes/to_convolution.py"], + visibility = _VISIBILITY, + deps = [ + ":ops", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_library( + name = "fuse_add_into_linear", + srcs = ["optimization_passes/fuse_add_into_linear.py"], + visibility = _VISIBILITY, + deps = [ + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "fuse_mul_into_linear", + srcs = ["optimization_passes/fuse_mul_into_linear.py"], + visibility = _VISIBILITY, + deps = [ + ":graph_utils", + ":ops", + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "split_linear_at_slices", + srcs = ["optimization_passes/split_linear_at_slices.py"], + visibility = _VISIBILITY, + deps = [ + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "fold_scalar_mul", + srcs = ["optimization_passes/fold_scalar_mul.py"], + visibility = _VISIBILITY, + deps = [ + ":graph_utils", + ":ops", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "sink_constant_cat", + srcs = ["optimization_passes/sink_constant_cat.py"], + visibility = _VISIBILITY, + deps = [ + ":ops", + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "prequantize_embedding", + srcs = ["optimization_passes/prequantize_embedding.py"], + visibility = _VISIBILITY, + deps = [ + ":graph_utils", + ":ops", + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "quant_absorption", + srcs = ["optimization_passes/quant_absorption.py"], + visibility = _VISIBILITY, + deps = [ + ":graph_utils", + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + "//executorch/exir/dialects/edge:lib", + "//executorch/exir/passes:constant_prop_pass", + ], + ) + + runtime.python_library( + name = "replace_dequant_quant_with_requantize", + srcs = ["optimization_passes/replace_dequant_quant_with_requantize.py"], + visibility = _VISIBILITY, + deps = [ + ":ops", + ":pass_base", + "//caffe2:torch", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) diff --git a/backends/fused_quant/test/BUCK b/backends/fused_quant/test/BUCK new file mode 100644 index 00000000000..0a503173691 --- /dev/null +++ b/backends/fused_quant/test/BUCK @@ -0,0 +1,6 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +fbcode_target(_kind = define_common_targets) diff --git a/backends/fused_quant/test/__init__.py b/backends/fused_quant/test/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/fused_quant/test/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/fused_quant/test/helpers.py b/backends/fused_quant/test/helpers.py new file mode 100644 index 00000000000..93ac390494c --- /dev/null +++ b/backends/fused_quant/test/helpers.py @@ -0,0 +1,115 @@ +# 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. + + +from typing import Optional + +import torch +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir.pass_base import ProxyValue +from torch._inductor.decomposition import remove_decompositions +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +def create_per_tensor_qparams( + builder: ProgramBuilder, + scale: float = 1.0, + zero_point: int = 0, + dtype: torch.dtype = torch.int8, +) -> tuple[ProxyValue, ProxyValue, torch.dtype, int, int]: + """Create per-tensor QuantParams as lifted scalar constants.""" + idx = len(builder.input_specs) + scale_node = builder.placeholder( + f"qp_scale_{idx}", + torch.tensor(scale, dtype=torch.float32), + input_kind=InputKind.CONSTANT_TENSOR, + ) + zp_node = builder.placeholder( + f"qp_zero_point_{idx}", + torch.tensor(zero_point, dtype=torch.int64), + input_kind=InputKind.CONSTANT_TENSOR, + ) + return (scale_node, zp_node, dtype, -128, 127) + + +def create_per_channel_qparams( + builder: ProgramBuilder, + num_channels: int, + idx: int = 0, + scale: float = 1.0, + zero_point: int = 0, + dtype: torch.dtype = torch.int8, + weight_ndim: int = 2, +) -> tuple[ProxyValue, ProxyValue, torch.dtype, int, int]: + """Create per-channel QuantParams using BUFFER placeholders. + + Scales/zero_points are full-rank ``[num_channels, 1, ...]`` so their shape + encodes the affine block layout (channel axis 0). + """ + full_shape = (num_channels,) + (1,) * (weight_ndim - 1) + scale_node = builder.placeholder( + f"wt_scale_{idx}", + torch.full(full_shape, scale, dtype=torch.float32), + input_kind=InputKind.BUFFER, + ) + zp_node = builder.placeholder( + f"wt_zp_{idx}", + torch.full(full_shape, zero_point, dtype=torch.int64), + input_kind=InputKind.BUFFER, + ) + return (scale_node, zp_node, dtype, -128, 127) + + +def create_per_axis_qparams( + builder: ProgramBuilder, + num_channels: int, + ndim: int, + axis: int, + dtype: torch.dtype = torch.int8, + vary_values: bool = False, +) -> tuple[ProxyValue, ProxyValue, torch.dtype, int, int]: + """Create full-rank qparams whose non-unitary dimension is ``axis``.""" + idx = len(builder.input_specs) + shape = [1] * ndim + shape[axis] = num_channels + scales = ( + 0.5 + torch.arange(num_channels, dtype=torch.float32) / num_channels + if vary_values + else torch.ones(num_channels, dtype=torch.float32) + ) + zero_points = ( + torch.arange(num_channels, dtype=torch.int64) - num_channels // 2 + if vary_values + else torch.zeros(num_channels, dtype=torch.int64) + ) + scale_node = builder.placeholder( + f"axis_scale_{idx}", + scales.reshape(shape), + input_kind=InputKind.BUFFER, + ) + zp_node = builder.placeholder( + f"axis_zp_{idx}", + zero_points.reshape(shape), + input_kind=InputKind.BUFFER, + ) + return (scale_node, zp_node, dtype, -128, 127) + + +def trace( + model: torch.nn.Module, + inputs: tuple[object, ...], + ops_to_keep: Optional[list[torch._ops.OpOverload]] = None, +) -> ExportedProgram: + """Export ``model`` to core ATen, keeping ``ops_to_keep`` undecomposed.""" + model.eval() + decomp_table = torch.export.default_decompositions() + remove_decompositions( + decomp_table, [*(ops_to_keep or []), torch.ops.aten._safe_softmax.default] + ) + return torch.export.export(model, inputs, strict=True).run_decompositions( + decomp_table + ) diff --git a/backends/fused_quant/test/targets.bzl b/backends/fused_quant/test/targets.bzl new file mode 100644 index 00000000000..bf98644f2d4 --- /dev/null +++ b/backends/fused_quant/test/targets.bzl @@ -0,0 +1,323 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +_VISIBILITY = ["PUBLIC"] + +def define_common_targets(): + + runtime.python_library( + name = "test_helpers", + srcs = ["helpers.py"], + visibility = _VISIBILITY, + deps = [ + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir:pass_base", + ], + ) + + runtime.python_test( + name = "test_colorer", + srcs = ["test_colorer.py"], + deps = [ + "//executorch/backends/fused_quant:decompose_fused_quant", + "//executorch/backends/fused_quant:colorer", + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir/dialects:lib", + "//executorch/exir/dialects/edge:lib", + ], + ) + + runtime.python_test( + name = "test_compile", + srcs = ["test_compile.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/fused_quant:compile", + "//executorch/backends/fused_quant:decompose_fused_quant", + "//executorch/backends/fused_quant:frontend", + "//executorch/backends/fused_quant:quantizer_defaults", + "//executorch/backends/fused_quant:to_channels_first", + "//executorch/exir:lib", + "//executorch/exir:pass_manager", + "//executorch/kernels/quantized:custom_ops_generated_lib", + "//executorch/kernels/quantized:quantized_ops_lib", + ], + ) + + runtime.python_test( + name = "test_ops", + srcs = ["test_ops.py"], + deps = [ + "//executorch/backends/fused_quant:ops", + "//caffe2:torch", + ], + ) + + runtime.python_test( + name = "test_pass_base", + srcs = ["test_pass_base.py"], + deps = [ + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:pass_base", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_decompose_fused_quant", + srcs = ["test_decompose_fused_quant.py"], + deps = [ + "//executorch/backends/fused_quant:decompose_fused_quant", + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_fold_batch_norm", + srcs = ["test_fold_batch_norm.py"], + deps = [ + ":test_helpers", + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:fold_batch_norm", + "//executorch/backends/fused_quant:replace_batch_norm_with_conv", + "//caffe2:torch", + ], + ) + + runtime.python_test( + name = "test_fuse_add_softmax_into_masked_softmax", + srcs = ["test_fuse_add_softmax_into_masked_softmax.py"], + deps = [ + ":test_helpers", + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:fuse_add_softmax_into_masked_softmax", + "//caffe2:torch", + ], + ) + + runtime.python_test( + name = "test_replace_batch_norm_with_conv", + srcs = ["test_replace_batch_norm_with_conv.py"], + deps = [ + ":test_helpers", + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:replace_batch_norm_with_conv", + "//caffe2:torch", + "//executorch/backends/transforms:utils", + ], + ) + + runtime.python_test( + name = "test_hoist_activations", + srcs = ["test_hoist_activations.py"], + deps = [ + "//executorch/backends/fused_quant:hoist_activations", + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_to_channels_last", + srcs = ["test_to_channels_last.py"], + deps = [ + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//executorch/backends/fused_quant:to_channels_last", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_to_channels_first", + srcs = ["test_to_channels_first.py"], + deps = [ + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//executorch/backends/fused_quant:to_channels_first", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_to_convolution", + srcs = ["test_to_convolution.py"], + deps = [ + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//executorch/backends/fused_quant:to_convolution", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_fuse_add_into_linear", + srcs = ["test_fuse_add_into_linear.py"], + deps = [ + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:fuse_add_into_linear", + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_fuse_mul_into_linear", + srcs = ["test_fuse_mul_into_linear.py"], + deps = [ + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:fuse_mul_into_linear", + "//executorch/backends/fused_quant:ops", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_split_linear_at_slices", + srcs = ["test_split_linear_at_slices.py"], + deps = [ + "//executorch/backends/fused_quant:constant_fold", + "//executorch/backends/fused_quant:ops", + "//executorch/backends/fused_quant:split_linear_at_slices", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_replace_dequant_quant_with_requantize", + srcs = ["test_replace_dequant_quant_with_requantize.py"], + deps = [ + "//executorch/backends/fused_quant:graph_utils", + "//executorch/backends/fused_quant:ops", + "//executorch/backends/fused_quant:replace_dequant_quant_with_requantize", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_quant_absorption", + srcs = ["test_quant_absorption.py"], + deps = [ + "//executorch/backends/fused_quant:graph_utils", + "//executorch/backends/fused_quant:ops", + "//executorch/backends/fused_quant:quant_absorption", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/backends/transforms:utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_replace_mm_with_bmm", + srcs = ["test_replace_mm_with_bmm.py"], + deps = [ + ":test_helpers", + "//executorch/backends/fused_quant:replace_mm_with_bmm", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/backends/transforms:permute_pass_utils", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_fold_scalar_mul", + srcs = ["test_fold_scalar_mul.py"], + deps = [ + ":test_helpers", + "//executorch/backends/fused_quant:ops", + "//executorch/backends/fused_quant:fold_scalar_mul", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_sink_constant_cat", + srcs = ["test_sink_constant_cat.py"], + deps = [ + "//executorch/backends/fused_quant:ops", + "//executorch/backends/fused_quant:sink_constant_cat", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_prequantize_embedding", + srcs = ["test_prequantize_embedding.py"], + deps = [ + "//executorch/backends/fused_quant:ops", + "//executorch/backends/fused_quant:prequantize_embedding", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_op_numerics", + srcs = ["test_op_numerics.py"], + deps = [ + "fbsource//third-party/pypi/parameterized:parameterized", + "//executorch/backends/fused_quant:decompose_fused_quant", + "//executorch/backends/fused_quant:graph_utils", + "//executorch/backends/fused_quant:ops", + "//executorch/backends/fused_quant:to_channels_first", + ":test_helpers", + "//caffe2:torch", + "//executorch/backends/test:program_builder", + "//executorch/exir:pass_manager", + "//executorch/exir/dialects:lib", + ], + ) diff --git a/backends/fused_quant/test/test_colorer.py b/backends/fused_quant/test/test_colorer.py new file mode 100644 index 00000000000..525afed5313 --- /dev/null +++ b/backends/fused_quant/test/test_colorer.py @@ -0,0 +1,199 @@ +# 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. + + +import unittest + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.colorer import ColorerBase +from executorch.backends.fused_quant.decompose_fused_quant import DecomposeFusedQuant +from executorch.backends.fused_quant.test.helpers import create_per_tensor_qparams +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.export import ExportedProgram +from torch.fx import Node + + +def _build_relu_add_program() -> ExportedProgram: + """inp -> fused_quant.relu -> fused_quant.add(., inp) -> output.""" + builder = ProgramBuilder() + inp = builder.placeholder("inp", torch.randn(1, 4).to(torch.int8)) + + relu = builder.call_operator( + op=exir_ops.edge.fused_quant.relu.default, + args=( + inp, + *create_per_tensor_qparams(builder), + *create_per_tensor_qparams(builder), + ), + ) + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.default, + args=( + relu, + inp, + *create_per_tensor_qparams(builder), + *create_per_tensor_qparams(builder), + *create_per_tensor_qparams(builder), + ), + ) + builder.output([add]) + return builder.get_program() + + +def _tagged(program: ExportedProgram) -> dict[str, str]: + return { + node.name: node.meta["delegation_tag"] + for node in program.graph_module.graph.nodes + if "delegation_tag" in node.meta + } + + +def _build_two_relu_program() -> ExportedProgram: + """inp -> fused_quant.relu -> fused_quant.relu -> output.""" + builder = ProgramBuilder() + inp = builder.placeholder("inp", torch.randn(1, 4).to(torch.int8)) + + first = builder.call_operator( + op=exir_ops.edge.fused_quant.relu.default, + args=( + inp, + *create_per_tensor_qparams(builder), + *create_per_tensor_qparams(builder), + ), + ) + second = builder.call_operator( + op=exir_ops.edge.fused_quant.relu.default, + args=( + first, + *create_per_tensor_qparams(builder), + *create_per_tensor_qparams(builder), + ), + ) + builder.output([second]) + return builder.get_program() + + +def _fused_quant_targets(program: ExportedProgram) -> list[str]: + return [ + str(node.target) + for node in program.graph_module.graph.nodes + if node.op == "call_function" + and getattr(node.target, "namespace", None) == "fused_quant" + ] + + +class ReluOnlyColorer(ColorerBase): + def __init__(self) -> None: + super().__init__(delegation_tag="toy") + + def supported_ops(self) -> set[EdgeOpOverload]: + return { + exir_ops.edge.fused_quant.relu.default, + exir_ops.edge.fused_quant.add.default, + } + + def is_legal(self, exported_program: ExportedProgram, node: Node) -> bool: + return node.target == exir_ops.edge.fused_quant.relu.default + + +class SecondReluColorer(ColorerBase): + """A lower-priority colorer that also wants relu.""" + + def __init__(self) -> None: + super().__init__(delegation_tag="second") + + def supported_ops(self) -> set[EdgeOpOverload]: + return {exir_ops.edge.fused_quant.relu.default} + + +class EmptyColorer(ColorerBase): + """Claims no ops -- the shape a backend starts from.""" + + def __init__(self) -> None: + super().__init__(delegation_tag="empty") + + def supported_ops(self) -> set[EdgeOpOverload]: + return set() + + +class ColorerBaseTest(unittest.TestCase): + def test_a_colorer_with_no_supported_ops_tags_nothing(self) -> None: + program = _build_relu_add_program() + + result = EmptyColorer()(program) + + self.assertFalse(result.modified) + self.assertEqual(_tagged(result.exported_program), {}) + + def test_tags_only_supported_and_legal_nodes(self) -> None: + """add is a supported op but illegal, so only relu is claimed.""" + program = _build_relu_add_program() + + result = ReluOnlyColorer()(program) + + self.assertTrue(result.modified) + self.assertEqual( + list(_tagged(result.exported_program).values()), + ["toy"], + ) + (tagged_name,) = _tagged(result.exported_program) + self.assertIn("relu", tagged_name) + + def test_first_claim_wins(self) -> None: + program = _build_relu_add_program() + ReluOnlyColorer()(program) + + result = SecondReluColorer()(program) + + self.assertFalse(result.modified) + self.assertEqual( + list(_tagged(result.exported_program).values()), + ["toy"], + ) + + +class AllReluColorer(ColorerBase): + """Claims every fused_quant.relu.""" + + def __init__(self) -> None: + super().__init__(delegation_tag="toy") + + def supported_ops(self) -> set[EdgeOpOverload]: + return {exir_ops.edge.fused_quant.relu.default} + + +class ColoredNodesAreLeftAloneTest(unittest.TestCase): + """The ``ColorerBase`` contract: a claimed node is owned by its backend. + + Without this, claiming an op is meaningless -- the generic passes would + defuse or relayout it before the backend ever saw it. + + Every fused op in these graphs is claimed, so the pass has nothing left to + rewrite and must report no change. That keeps the assertion on the skip + itself: the synthetic ProgramBuilder graphs carry int8 metadata and are not + decomposable, so letting the pass touch one would fail for an unrelated + reason. End-to-end decomposition is covered by ``test_compile.py``. + """ + + def test_decompose_skips_colored_nodes(self) -> None: + program = _build_two_relu_program() + AllReluColorer()(program) + before = _fused_quant_targets(program) + + result = DecomposeFusedQuant()(program) + + self.assertFalse(result.modified) + self.assertEqual(_fused_quant_targets(result.exported_program), before) + + def test_decompose_rewrites_uncolored_nodes(self) -> None: + """Same graph, nothing claimed: the pass does try to rewrite.""" + program = _build_two_relu_program() + + with self.assertRaises(Exception): + DecomposeFusedQuant()(program) diff --git a/backends/fused_quant/test/test_compile.py b/backends/fused_quant/test/test_compile.py new file mode 100644 index 00000000000..a78c988ab85 --- /dev/null +++ b/backends/fused_quant/test/test_compile.py @@ -0,0 +1,202 @@ +# 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. + + +import unittest + +# Registers the out variants of the quantized_decomposed q/dq ops that +# to_executorch needs once DecomposeFusedQuant has emitted them. +import executorch.kernels.quantized # noqa: F401 +import torch +from executorch.backends.fused_quant.compile import compile_to_fused_quant +from executorch.backends.fused_quant.decompose_fused_quant import DecomposeFusedQuant +from executorch.backends.fused_quant.frontend import trace +from executorch.backends.fused_quant.optimization_passes.to_channels_first import ( + ToChannelsFirst, +) +from executorch.backends.fused_quant.quantizer.defaults import ( + make_fused_quant_quantizer, +) +from executorch.exir import EdgeProgramManager, ExecutorchProgramManager +from executorch.exir.pass_manager import ExportedProgramPassManager +from torch import nn + + +class ConvReluLinear(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu = nn.ReLU() + self.linear = nn.Linear(8 * 4 * 4, 5) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.relu(self.conv(x)) + return self.linear(x.flatten(1)) + + +def _calibration_inputs() -> list[tuple[torch.Tensor, ...]]: + torch.manual_seed(0) + return [(torch.randn(2, 3, 4, 4),) for _ in range(4)] + + +def _decompose_for_portable_kernels( + edge_program_manager: EdgeProgramManager, +) -> EdgeProgramManager: + """The tail a backend runs after :func:`compile_to_fused_quant`. + + This is the no-backend case written out: nothing is colored, so every + fused_quant op has to go back to quantize/dequantize plus ATen for the + portable kernels to run it. A real backend colors first, lowers what it + claimed, and runs this over the remainder. + + ``ToChannelsFirst`` comes first because no backend is consuming NHWC here, so + the layout ``ToChannelsLast`` introduced has to be undone before decomposing. + """ + return edge_program_manager.transform( + ExportedProgramPassManager([ToChannelsFirst(), DecomposeFusedQuant()]) + ) + + +def _op_targets(program: ExecutorchProgramManager) -> list[str]: + return [ + str(node.target) + for node in program.exported_program().graph_module.graph.nodes + if node.op == "call_function" + ] + + +class CompileTest(unittest.TestCase): + def test_compile_to_fused_quant_produces_fused_quant_ops(self) -> None: + """The pre-coloring seam hands back a graph in fused_quant form.""" + inputs = _calibration_inputs() + + edge_program_manager = compile_to_fused_quant(ConvReluLinear(), inputs) + + targets = [ + str(node.target) + for node in edge_program_manager.exported_program().graph_module.graph.nodes + if node.op == "call_function" + ] + self.assertTrue( + any("fused_quant" in target for target in targets), + f"expected fused_quant ops, got {targets}", + ) + + def test_passes_can_be_overridden(self) -> None: + """An explicit pass list replaces the default pipeline.""" + inputs = _calibration_inputs() + + unoptimized = compile_to_fused_quant(ConvReluLinear(), inputs, passes=[]) + + targets = [ + str(node.target) + for node in unoptimized.exported_program().graph_module.graph.nodes + if node.op == "call_function" + ] + # ToConvolution is part of the default pipeline and rewrites conv2d into + # the generic convolution op, so its absence shows the default did not run. + self.assertTrue( + any("fused_quant.conv2d" in target for target in targets), + f"expected the unoptimized conv2d form, got {targets}", + ) + + def test_decomposing_clears_every_fused_quant_op(self) -> None: + """With nothing colored, the tail must leave no fused_quant op behind.""" + inputs = _calibration_inputs() + + edge = _decompose_for_portable_kernels( + compile_to_fused_quant(ConvReluLinear(), inputs) + ) + + leftover = [ + str(node.target) + for node in edge.exported_program().graph_module.graph.nodes + if node.op == "call_function" and "fused_quant" in str(node.target) + ] + self.assertEqual(leftover, []) + + def test_decomposed_program_has_no_delegates(self) -> None: + """Nothing is claimed, so nothing should be delegated.""" + inputs = _calibration_inputs() + + edge = _decompose_for_portable_kernels( + compile_to_fused_quant(ConvReluLinear(), inputs) + ) + program = edge.to_executorch() + + self.assertNotIn("executorch_call_delegate", _op_targets(program)) + + def test_decomposed_program_matches_the_quantized_reference(self) -> None: + """Compiled output tracks the eager float model within quantization error.""" + model = ConvReluLinear().eval() + inputs = _calibration_inputs() + + edge = _decompose_for_portable_kernels(compile_to_fused_quant(model, inputs)) + + (example,) = inputs[0] + with torch.no_grad(): + expected = model(example) + actual = edge.exported_program().module()(example) + + self.assertEqual(actual.shape, expected.shape) + # int8 per-tensor activations over this range; compare against the + # magnitude of the reference rather than an absolute epsilon. + tolerance = 0.1 * expected.abs().max().item() + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=0.0) + + +class SoftmaxAttention(nn.Module): + """SDPA, which traces to aten._safe_softmax.""" + + def forward( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor + ) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention(q, k, v) + + +class CoreOnlyAttentionTest(unittest.TestCase): + """A quantizer that does not claim _masked_softmax on an attention model. + + `trace` holds `aten._safe_softmax` back from decomposition unconditionally, + but the recompose into `_masked_softmax` only runs for a quantizer that + claims it, so with the core set the op survives tracing. It is not core ATen + and it is not in the verifier exception list, which would be a problem if it + reached the verifier -- it does not, because `to_edge` decomposes it. This + pins that down, since the exception list only covers ops re-emitted *after* + `to_edge` by `DecomposeFusedQuant`. + """ + + def _inputs(self) -> list[tuple[torch.Tensor, ...]]: + torch.manual_seed(0) + shape = (1, 2, 8, 16) + return [ + (torch.randn(shape), torch.randn(shape), torch.randn(shape)) + for _ in range(2) + ] + + def test_safe_softmax_survives_tracing_but_not_to_edge(self) -> None: + quantizer = make_fused_quant_quantizer(core_only=True) + self.assertNotIn( + torch.ops.aten._masked_softmax.default, quantizer.preserved_ops() + ) + + program = trace(SoftmaxAttention().eval(), self._inputs()[0], quantizer) + traced = {n.target for n in program.graph.nodes if n.op == "call_function"} + self.assertIn(torch.ops.aten._safe_softmax.default, traced) + + edge = compile_to_fused_quant( + SoftmaxAttention().eval(), self._inputs(), quantizer + ) + + edge_ops = { + str(n.target) + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + } + self.assertFalse( + any("_safe_softmax" in op for op in edge_ops), + f"_safe_softmax reached the edge graph: {sorted(edge_ops)}", + ) diff --git a/backends/fused_quant/test/test_decompose_fused_quant.py b/backends/fused_quant/test/test_decompose_fused_quant.py new file mode 100644 index 00000000000..fbcad443be1 --- /dev/null +++ b/backends/fused_quant/test/test_decompose_fused_quant.py @@ -0,0 +1,313 @@ +# 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. + +from __future__ import annotations + +import operator +import unittest + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.decompose_fused_quant import DecomposeFusedQuant +from executorch.backends.fused_quant.test.helpers import ( + create_per_axis_qparams, + create_per_channel_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +class DecomposeFusedQuantTest(unittest.TestCase): + def assert_bit_exact(self, expected: object, actual: object) -> None: + torch.testing.assert_close(expected, actual, rtol=0, atol=0) + + def _assert_no_fused_quant_ops(self, program: ExportedProgram) -> None: + self.assertFalse( + any( + node.op == "call_function" + and getattr(node.target, "namespace", None) == "fused_quant" + for node in program.graph.nodes + ) + ) + + def test_decomposes_per_channel_convolution(self) -> None: + inp = torch.randint(-8, 8, (1, 2, 6, 6), dtype=torch.int8) + weight = torch.randint(-8, 8, (3, 2, 3, 3), dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + weight_node = builder.placeholder("weight", weight) + inp_qparams = create_per_axis_qparams( + builder, 2, 4, axis=1, dtype=torch.float32, vary_values=True + ) + weight_qparams = create_per_channel_qparams( + builder, 3, dtype=torch.float32, weight_ndim=4 + ) + out_qparams = create_per_axis_qparams(builder, 3, 4, axis=1, vary_values=True) + convolution = builder.call_operator( + exir_ops.edge.fused_quant.convolution.default, + args=( + inp_node, + weight_node, + None, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + -128, + 127, + *out_qparams, + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + 1, + ), + ) + builder.output([convolution]) + program = builder.get_program() + expected = program.module()(inp, weight) + + decomposed = DecomposeFusedQuant()(program) + + self.assertTrue(decomposed.modified) + self.assert_bit_exact( + expected, + decomposed.exported_program.module()(inp, weight), + ) + self._assert_no_fused_quant_ops(decomposed.exported_program) + self.assertEqual( + len( + decomposed.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.convolution.default, + ) + ), + 1, + ) + self.assertEqual( + len( + decomposed.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, + ) + ), + 2, + ) + self.assertEqual( + len( + decomposed.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_channel.default, + ) + ), + 1, + ) + + def test_preserves_transposed_convolution_arguments(self) -> None: + inp = torch.randn(1, 2, 4, 4) + weight = torch.randn(2, 3, 3, 3) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + weight_node = builder.placeholder("weight", weight) + no_qparams = (None, None, torch.float32, 0, 0) + convolution = builder.call_operator( + exir_ops.edge.fused_quant.convolution.default, + args=( + inp_node, + weight_node, + None, + *no_qparams, + *no_qparams, + *no_qparams, + *no_qparams, + [2, 2], + [1, 1], + [1, 1], + True, + [1, 1], + 1, + ), + ) + builder.output([convolution]) + program = builder.get_program() + expected = program.module()(inp, weight) + + result = DecomposeFusedQuant()(program) + + self.assertTrue(result.modified) + self.assert_bit_exact(expected, result.exported_program.module()(inp, weight)) + [aten_convolution] = result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.convolution.default, + ) + self.assertTrue(get_arg(aten_convolution, "transposed", bool)) + self.assertEqual(get_arg(aten_convolution, "output_padding", list[int]), [1, 1]) + + def test_decomposes_tensor_add(self) -> None: + inp = torch.tensor([[-4, 1, 7]], dtype=torch.int8) + other = torch.tensor([[2, -3, 5]], dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + other_node = builder.placeholder("other", other) + inp_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + other_qparams = create_per_tensor_qparams( + builder, scale=0.25, dtype=torch.float32 + ) + out_qparams = create_per_tensor_qparams(builder, scale=0.125, zero_point=3) + add = builder.call_operator( + exir_ops.edge.fused_quant.add.default, + args=( + inp_node, + other_node, + *inp_qparams, + *other_qparams, + *out_qparams, + 1.0, + ), + ) + builder.output([add]) + program = builder.get_program() + expected = program.module()(inp, other) + + result = DecomposeFusedQuant()(program) + + self.assertTrue(result.modified) + self.assert_bit_exact(expected, result.exported_program.module()(inp, other)) + self._assert_no_fused_quant_ops(result.exported_program) + graph = result.exported_program.graph + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ), + 2, + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.add.Tensor, + ) + ), + 1, + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + 1, + ) + + def test_decomposes_requantize_without_aten_op(self) -> None: + inp = torch.tensor([[-4, 1, 7]], dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + inp_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + out_qparams = create_per_tensor_qparams(builder, scale=0.25, zero_point=2) + requantize = builder.call_operator( + exir_ops.edge.fused_quant.requantize.default, + args=(inp_node, *inp_qparams, *out_qparams), + ) + builder.output([requantize]) + program = builder.get_program() + expected = program.module()(inp) + + result = DecomposeFusedQuant()(program) + + self.assertTrue(result.modified) + self.assert_bit_exact(expected, result.exported_program.module()(inp)) + self._assert_no_fused_quant_ops(result.exported_program) + call_targets = [ + node.target + for node in result.exported_program.graph.nodes + if node.op == "call_function" and node.target is not operator.getitem + ] + self.assertEqual( + call_targets.count( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + ), + 1, + ) + self.assertEqual( + call_targets.count( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + ), + 1, + ) + + def test_preserves_colored_fused_quant_op(self) -> None: + inp = torch.tensor([[-4, 1, 7]], dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + inp_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + out_qparams = create_per_tensor_qparams(builder, scale=0.25, zero_point=2) + requantize = builder.call_operator( + exir_ops.edge.fused_quant.requantize.default, + args=(inp_node, *inp_qparams, *out_qparams), + ) + builder.output([requantize]) + program = builder.get_program() + (requantize_node,) = program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.requantize.default, + ) + requantize_node.meta["delegation_tag"] = "test_backend" + + result = DecomposeFusedQuant()(program) + + self.assertFalse(result.modified) + self.assertEqual(requantize_node.meta["delegation_tag"], "test_backend") + self.assertIn(requantize_node, result.exported_program.graph.nodes) + + def test_rejects_group_quantization(self) -> None: + inp = torch.randint(-8, 8, (2, 4), dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + scale = builder.placeholder( + "group_scale", + torch.ones(1, 2, dtype=torch.float32), + input_kind=InputKind.BUFFER, + ) + zero_point = builder.placeholder( + "group_zero_point", + torch.zeros(1, 2, dtype=torch.int64), + input_kind=InputKind.BUFFER, + ) + relu = builder.call_operator( + exir_ops.edge.fused_quant.relu.default, + args=( + inp_node, + scale, + zero_point, + torch.float32, + -128, + 127, + None, + None, + torch.float32, + -128, + 127, + ), + ) + builder.output([relu]) + program = builder.get_program() + + with self.assertRaisesRegex( + ValueError, + "only per-tensor and per-channel quantization are supported", + ): + DecomposeFusedQuant()(program) diff --git a/backends/fused_quant/test/test_fold_batch_norm.py b/backends/fused_quant/test/test_fold_batch_norm.py new file mode 100644 index 00000000000..acf7cb209f1 --- /dev/null +++ b/backends/fused_quant/test/test_fold_batch_norm.py @@ -0,0 +1,163 @@ +# 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. + +from __future__ import annotations + +import unittest + +import torch +from executorch.backends.fused_quant.pre_quantize_passes.fold_batch_norm import ( + FoldBatchNorm, +) +from executorch.backends.fused_quant.pre_quantize_passes.replace_batch_norm_with_conv import ( + ReplaceBatchNormWithConv, +) +from executorch.backends.fused_quant.test.helpers import trace +from parameterized import parameterized +from torch import fx, nn +from torch._ops import OpOverload +from torch.export import ExportedProgram + +_BATCH_NORM = torch.ops.aten._native_batch_norm_legit_no_training.default +_CONVOLUTION = torch.ops.aten.convolution.default + + +def _conv_target(input_rank: int) -> OpOverload: + if input_rank == 3: + return torch.ops.aten.conv1d.default + if input_rank == 4: + return torch.ops.aten.conv2d.default + if input_rank == 5: + return torch.ops.aten.conv3d.default + raise ValueError(f"Unsupported input rank: {input_rank}") + + +def _make_conv(input_rank: int, channels: int, bias: bool) -> nn.Module: + if input_rank == 3: + return nn.Conv1d(channels, channels, kernel_size=3, padding=1, bias=bias) + if input_rank == 4: + return nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=bias) + if input_rank == 5: + return nn.Conv3d(channels, channels, kernel_size=3, padding=1, bias=bias) + raise ValueError(f"Unsupported input rank: {input_rank}") + + +def _make_batch_norm(input_rank: int, channels: int) -> nn.Module: + if input_rank == 3: + return nn.BatchNorm1d(channels) + if input_rank == 4: + return nn.BatchNorm2d(channels) + if input_rank == 5: + return nn.BatchNorm3d(channels) + raise ValueError(f"Unsupported input rank: {input_rank}") + + +def _randomize_batch_norm(batch_norm: nn.Module, channels: int) -> None: + batch_norm.running_mean = torch.randn(channels) + batch_norm.running_var = torch.rand(channels) + 0.5 + batch_norm.weight = nn.Parameter(torch.randn(channels)) + batch_norm.bias = nn.Parameter(torch.randn(channels)) + + +class _ConvBatchNormModel(nn.Module): + def __init__(self, input_rank: int, channels: int, bias: bool) -> None: + super().__init__() + self.conv = _make_conv(input_rank, channels, bias) + self.batch_norm = _make_batch_norm(input_rank, channels) + _randomize_batch_norm(self.batch_norm, channels) + + def forward(self, inp: torch.Tensor) -> torch.Tensor: + return self.batch_norm(self.conv(inp)) + + +class _FoldAndReplaceModel(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1) + self.fusable_batch_norm = nn.BatchNorm2d(channels) + self.standalone_batch_norm = nn.BatchNorm2d(channels) + _randomize_batch_norm(self.fusable_batch_norm, channels) + _randomize_batch_norm(self.standalone_batch_norm, channels) + + def forward(self, inp: torch.Tensor) -> torch.Tensor: + return self.fusable_batch_norm(self.conv(inp)) + self.standalone_batch_norm(inp) + + +def _trace( + model: nn.Module, inp: torch.Tensor, conv_target: OpOverload +) -> ExportedProgram: + return trace(model, (inp,), ops_to_keep=[conv_target, _CONVOLUTION]) + + +def _count(exported_program: ExportedProgram, target: object) -> int: + return len( + exported_program.graph_module.graph.find_nodes( + op="call_function", target=target + ) + ) + + +class FoldBatchNormTest(unittest.TestCase): + @parameterized.expand( + [ + ("conv1d_with_bias", (2, 3, 8), True), + ("conv1d_without_bias", (2, 3, 8), False), + ("conv2d_with_bias", (2, 3, 6, 8), True), + ("conv2d_without_bias", (2, 3, 6, 8), False), + ("conv3d_with_bias", (2, 3, 4, 6, 8), True), + ("conv3d_without_bias", (2, 3, 4, 6, 8), False), + ] + ) + def test_folds_batch_norm_into_preceding_convolution( + self, + _name: str, + input_shape: tuple[int, ...], + bias: bool, + ) -> None: + torch.manual_seed(0) + conv_target = _conv_target(len(input_shape)) + model = _ConvBatchNormModel(len(input_shape), input_shape[1], bias).eval() + inp = torch.randn(input_shape) + expected = model(inp) + exported_program = _trace(model, inp, conv_target) + + self.assertEqual(_count(exported_program, conv_target), 1) + self.assertEqual(_count(exported_program, _BATCH_NORM), 1) + + result = FoldBatchNorm()(exported_program) + + self.assertTrue(result.modified) + self.assertEqual(_count(result.exported_program, conv_target), 1) + self.assertEqual(_count(result.exported_program, _BATCH_NORM), 0) + folded_conv = result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=conv_target + )[0] + self.assertIsInstance(folded_conv.args[2], fx.Node) + self.assertFalse( + any(isinstance(value, fx.Node) for value in folded_conv.kwargs.values()) + ) + actual = result.exported_program.module()(inp) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_fold_then_replace_handles_only_remaining_batch_norm(self) -> None: + torch.manual_seed(0) + model = _FoldAndReplaceModel(3).eval() + inp = torch.randn(2, 3, 6, 8) + expected = model(inp) + exported_program = _trace(model, inp, torch.ops.aten.conv2d.default) + + fold_result = FoldBatchNorm()(exported_program) + + self.assertTrue(fold_result.modified) + self.assertEqual(_count(fold_result.exported_program, _BATCH_NORM), 1) + + replace_result = ReplaceBatchNormWithConv()(fold_result.exported_program) + + self.assertTrue(replace_result.modified) + self.assertEqual(_count(replace_result.exported_program, _BATCH_NORM), 0) + self.assertEqual(_count(replace_result.exported_program, _CONVOLUTION), 1) + actual = replace_result.exported_program.module()(inp) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) diff --git a/backends/fused_quant/test/test_fold_scalar_mul.py b/backends/fused_quant/test/test_fold_scalar_mul.py new file mode 100644 index 00000000000..89f2a8ec66f --- /dev/null +++ b/backends/fused_quant/test/test_fold_scalar_mul.py @@ -0,0 +1,216 @@ +# 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. + +from __future__ import annotations + +import unittest + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.optimization_passes.fold_scalar_mul import ( + FoldScalarMul, +) +from executorch.backends.fused_quant.test.helpers import ( + create_per_axis_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram + + +_NULL_QPARAMS = (None, None, torch.float32, 0, 0) + + +class FoldScalarMulTest(unittest.TestCase): + def _build_scalar_mul( + self, + *, + scalar: float = 1.5, + has_input_qparams: bool = True, + has_output_qparams: bool = True, + per_axis_input_qparams: bool = False, + per_axis_output_qparams: bool = False, + ) -> tuple[ExportedProgram, torch.Tensor]: + builder = ProgramBuilder() + inp_value = ( + torch.tensor([[-8, -1, 3, 8]], dtype=torch.int8) + if has_input_qparams + else torch.tensor([[-2.0, -0.25, 0.75, 2.0]]) + ) + inp = builder.placeholder("inp", inp_value) + if per_axis_input_qparams: + inp_qparams = create_per_axis_qparams( + builder, num_channels=4, ndim=2, axis=1, dtype=torch.float32 + ) + elif has_input_qparams: + inp_qparams = create_per_tensor_qparams( + builder, scale=0.25, zero_point=-3, dtype=torch.float32 + ) + else: + inp_qparams = _NULL_QPARAMS + if per_axis_output_qparams: + out_qparams = create_per_axis_qparams( + builder, num_channels=4, ndim=2, axis=1, dtype=torch.int8 + ) + elif has_output_qparams: + out_qparams = create_per_tensor_qparams( + builder, scale=0.2, zero_point=4, dtype=torch.int8 + ) + else: + out_qparams = _NULL_QPARAMS + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.Scalar, + args=(inp, *inp_qparams, *out_qparams, scalar), + ) + builder.output([mul]) + return builder.get_program(), inp_value + + def _assert_numerics( + self, + ep: ExportedProgram, + inp: torch.Tensor, + expected: tuple[torch.Tensor, ...], + ) -> None: + actual = ep.module()(inp) + self.assertEqual(len(expected), len(actual)) + for expected_output, actual_output in zip(expected, actual): + torch.testing.assert_close(actual_output, expected_output) + + # -- Case 1: both inp and out qparams (dq->q folding) -- + + def test_folds_fully_quantized_into_dequant_quant(self) -> None: + ep, inp = self._build_scalar_mul() + expected = ep.module()(inp) + + result = FoldScalarMul()(ep) + + self.assertTrue(result.modified) + self._assert_numerics(result.exported_program, inp, expected) + graph = result.exported_program.graph + self.assertEqual( + graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.mul.Scalar + ), + [], + ) + (dequant,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + (quant,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + self.assertLess( + list(graph.nodes).index(dequant), list(graph.nodes).index(quant) + ) + self.assertAlmostEqual(get_arg(dequant, "scale", float), 0.375) + self.assertEqual(get_arg(dequant, "zero_point", int), -3) + self.assertAlmostEqual(get_arg(quant, "scale", float), 0.2) + self.assertEqual(get_arg(quant, "zero_point", int), 4) + self.assertEqual(dequant.meta["val"].dtype, torch.float32) + self.assertEqual(quant.meta["val"].dtype, torch.int8) + result.exported_program.validate() + + # -- Case 2: only out qparams (fold scalar into quant) -- + + def test_folds_float_input_scalar_mul_into_quant(self) -> None: + ep, inp = self._build_scalar_mul(has_input_qparams=False) + expected = ep.module()(inp) + + result = FoldScalarMul()(ep) + + self.assertTrue(result.modified) + self._assert_numerics(result.exported_program, inp, expected) + graph = result.exported_program.graph + self.assertEqual( + graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.mul.Scalar + ), + [], + ) + dequants = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(dequants, []) + (quant,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + # scale = 0.2 / 1.5 = 0.1333... + self.assertAlmostEqual(get_arg(quant, "scale", float), 0.2 / 1.5, places=6) + self.assertEqual(get_arg(quant, "zero_point", int), 4) + self.assertEqual(quant.meta["val"].dtype, torch.int8) + result.exported_program.validate() + + def test_fold_into_quant_with_different_scalars(self) -> None: + for scalar in (0.5, 2.0, 3.14): + with self.subTest(scalar=scalar): + ep, inp = self._build_scalar_mul(scalar=scalar, has_input_qparams=False) + expected = ep.module()(inp) + + result = FoldScalarMul()(ep) + + self.assertTrue(result.modified) + self._assert_numerics(result.exported_program, inp, expected) + (quant,) = result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + self.assertAlmostEqual( + get_arg(quant, "scale", float), 0.2 / scalar, places=6 + ) + + # -- Negative cases -- + + def test_does_not_fold_without_output_qparams(self) -> None: + for has_input_qparams in (True, False): + with self.subTest(has_input_qparams=has_input_qparams): + ep, _ = self._build_scalar_mul( + has_input_qparams=has_input_qparams, + has_output_qparams=False, + ) + + result = FoldScalarMul()(ep) + + self.assertFalse(result.modified) + self.assertEqual( + len( + result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.mul.Scalar, + ) + ), + 1, + ) + + def test_does_not_fold_unsupported_scalars(self) -> None: + for scalar in (0.0, -1.5, float("inf"), float("nan")): + with self.subTest(scalar=scalar): + ep, _ = self._build_scalar_mul(scalar=scalar) + + result = FoldScalarMul()(ep) + + self.assertFalse(result.modified) + + def test_does_not_fold_per_axis_input_qparams(self) -> None: + ep, _ = self._build_scalar_mul(per_axis_input_qparams=True) + + result = FoldScalarMul()(ep) + + self.assertFalse(result.modified) + + def test_does_not_fold_per_axis_output_qparams(self) -> None: + ep, _ = self._build_scalar_mul( + has_input_qparams=False, per_axis_output_qparams=True + ) + + result = FoldScalarMul()(ep) + + self.assertFalse(result.modified) diff --git a/backends/fused_quant/test/test_fuse_add_into_linear.py b/backends/fused_quant/test/test_fuse_add_into_linear.py new file mode 100644 index 00000000000..35f2398ef5e --- /dev/null +++ b/backends/fused_quant/test/test_fuse_add_into_linear.py @@ -0,0 +1,499 @@ +# 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. + + +import unittest +from typing import Any + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.optimization_passes.fuse_add_into_linear import ( + FuseAddIntoLinear, +) +from executorch.backends.fused_quant.test.helpers import ( + create_per_axis_qparams, + create_per_channel_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import get_constant +from executorch.exir.dialects._ops import ops as exir_ops +from parameterized import parameterized +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +class FuseAddIntoLinearTest(unittest.TestCase): + def _build_affine_add_scalar( + self, target: Any, with_bias: bool + ) -> tuple[ExportedProgram, torch.Tensor]: + builder = ProgramBuilder() + out_channels = 4 + + if target == exir_ops.edge.fused_quant.linear.default: + inp_value = torch.randn(2, 8) + weight_value = torch.randint(-16, 16, (out_channels, 8), dtype=torch.int8) + elif target == exir_ops.edge.fused_quant.convolution.default: + inp_value = torch.randn(1, 3, 5, 5) + weight_value = torch.randint( + -16, 16, (out_channels, 3, 1, 1), dtype=torch.int8 + ) + else: + inp_value = torch.randn(1, 5, 5, 3) + weight_value = torch.randint( + -16, 16, (out_channels, 1, 1, 3), dtype=torch.int8 + ) + + bias_value = torch.arange(out_channels, dtype=torch.float32) + expected_bias = bias_value + 0.5 if with_bias else torch.full((4,), 0.5) + + inp = builder.placeholder("x", inp_value) + weight = builder.placeholder( + "weight", weight_value, input_kind=InputKind.BUFFER + ) + bias = ( + builder.placeholder("bias", bias_value, input_kind=InputKind.BUFFER) + if with_bias + else None + ) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_channel_qparams( + builder, out_channels, dtype=torch.float32 + ) + affine_out_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.int8 + ) + affine_args: tuple[Any, ...] = ( + inp, + weight, + bias, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + 0, + 0, + *affine_out_qparams, + ) + if target != exir_ops.edge.fused_quant.linear.default: + affine_args = ( + *affine_args, + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + 1, + ) + + affine = builder.call_operator(op=target, args=affine_args) + add_inp_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.float32 + ) + add_out_qparams = create_per_tensor_qparams( + builder, scale=0.125, dtype=torch.int8 + ) + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.Scalar, + args=(affine, *add_inp_qparams, *add_out_qparams, 2.0, 0.25), + ) + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(add, 0.125, 0, -128, 127, torch.int8), + ) + builder.output([dequantize]) + return builder.get_program(), expected_bias + + @parameterized.expand( + [ + ("linear_bias", exir_ops.edge.fused_quant.linear.default, True), + ("linear_no_bias", exir_ops.edge.fused_quant.linear.default, False), + ("conv_bias", exir_ops.edge.fused_quant.convolution.default, True), + ("conv_no_bias", exir_ops.edge.fused_quant.convolution.default, False), + ( + "conv_channels_last_bias", + exir_ops.edge.fused_quant.convolution_channels_last.default, + True, + ), + ] + ) + def test_scalar_add_folded_into_bias( + self, _name: str, target: Any, with_bias: bool + ) -> None: + ep, expected_bias = self._build_affine_add_scalar(target, with_bias) + + result = FuseAddIntoLinear().call(ep) + + self.assertTrue(result.modified) + graph = result.exported_program.graph_module.graph + self.assertEqual( + graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.add.Scalar + ), + [], + ) + affine_node = graph.find_nodes(op="call_function", target=target)[0] + bias_node = get_arg(affine_node, "bias", torch.fx.Node) + actual_bias = get_constant(result.exported_program, bias_node) + torch.testing.assert_close(actual_bias, expected_bias) + + def test_scalar_add_not_folded_when_qparams_mismatch(self) -> None: + builder = ProgramBuilder() + out_channels = 4 + inp_value = torch.randn(2, 8) + weight_value = torch.randint(-16, 16, (out_channels, 8), dtype=torch.int8) + bias_value = torch.arange(out_channels, dtype=torch.float32) + + inp = builder.placeholder("x", inp_value) + weight = builder.placeholder( + "weight", weight_value, input_kind=InputKind.BUFFER + ) + bias = builder.placeholder("bias", bias_value, input_kind=InputKind.BUFFER) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_channel_qparams( + builder, out_channels, dtype=torch.float32 + ) + affine_out_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.int8 + ) + affine = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + inp, + weight, + bias, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + 0, + 0, + *affine_out_qparams, + ), + ) + # Mismatched inp_scale (0.25) vs affine out_scale (0.5) + add_inp_qparams = create_per_tensor_qparams( + builder, scale=0.25, dtype=torch.float32 + ) + add_out_qparams = create_per_tensor_qparams( + builder, scale=0.125, dtype=torch.int8 + ) + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.Scalar, + args=(affine, *add_inp_qparams, *add_out_qparams, 2.0, 0.25), + ) + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(add, 0.125, 0, -128, 127, torch.int8), + ) + builder.output([dequantize]) + ep = builder.get_program() + + result = FuseAddIntoLinear().call(ep) + + self.assertFalse(result.modified) + + def test_scalar_add_folded_when_qparams_are_none(self) -> None: + builder = ProgramBuilder() + out_channels = 4 + inp_value = torch.randn(2, 8) + weight_value = torch.randint(-16, 16, (out_channels, 8), dtype=torch.int8) + bias_value = torch.arange(out_channels, dtype=torch.float32) + expected_bias = bias_value + 0.5 + + inp = builder.placeholder("x", inp_value) + weight = builder.placeholder( + "weight", weight_value, input_kind=InputKind.BUFFER + ) + bias = builder.placeholder("bias", bias_value, input_kind=InputKind.BUFFER) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_channel_qparams( + builder, out_channels, dtype=torch.float32 + ) + # Affine out_scale and out_zero_point are None + affine = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + inp, + weight, + bias, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.int8, + -128, + 127, + ), + ) + # Add inp_scale and inp_zero_point are also None + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.Scalar, + args=( + affine, + None, + None, + torch.float32, + -128, + 127, + None, + None, + torch.int8, + -128, + 127, + 2.0, + 0.25, + ), + ) + builder.output([add]) + ep = builder.get_program() + + result = FuseAddIntoLinear().call(ep) + + self.assertTrue(result.modified) + graph = result.exported_program.graph_module.graph + self.assertEqual( + graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.add.Scalar + ), + [], + ) + affine_node = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + )[0] + bias_node = get_arg(affine_node, "bias", torch.fx.Node) + actual_bias = get_constant(result.exported_program, bias_node) + torch.testing.assert_close(actual_bias, expected_bias) + + def test_scalar_add_not_folded_when_dtype_changes(self) -> None: + builder = ProgramBuilder() + out_channels = 4 + inp_value = torch.randn(2, 8) + weight_value = torch.randint(-16, 16, (out_channels, 8), dtype=torch.int8) + bias_value = torch.arange(out_channels, dtype=torch.float32) + + inp = builder.placeholder("x", inp_value) + weight = builder.placeholder( + "weight", weight_value, input_kind=InputKind.BUFFER + ) + bias = builder.placeholder("bias", bias_value, input_kind=InputKind.BUFFER) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_channel_qparams( + builder, out_channels, dtype=torch.float32 + ) + affine_out_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.int8 + ) + affine = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + inp, + weight, + bias, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + 0, + 0, + *affine_out_qparams, + ), + ) + add_inp_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.float32 + ) + # out_dtype=int16 differs from affine out_dtype=int8 + add_out_qparams = create_per_tensor_qparams( + builder, scale=0.125, dtype=torch.int16 + ) + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.Scalar, + args=(affine, *add_inp_qparams, *add_out_qparams, 2.0, 0.25), + ) + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(add, 0.125, 0, -32768, 32767, torch.int16), + ) + builder.output([dequantize]) + ep = builder.get_program() + + result = FuseAddIntoLinear().call(ep) + + self.assertFalse(result.modified) + + def test_scalar_add_folded_through_channels_last_permute(self) -> None: + builder = ProgramBuilder() + out_channels = 4 + # NHWC input for channels-last conv + inp_value = torch.randn(1, 5, 5, 3) + weight_value = torch.randint(-16, 16, (out_channels, 1, 1, 3), dtype=torch.int8) + bias_value = torch.arange(out_channels, dtype=torch.float32) + expected_bias = bias_value + 0.5 + + inp = builder.placeholder("x", inp_value) + weight = builder.placeholder( + "weight", weight_value, input_kind=InputKind.BUFFER + ) + bias = builder.placeholder("bias", bias_value, input_kind=InputKind.BUFFER) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_channel_qparams( + builder, out_channels, dtype=torch.float32, weight_ndim=4 + ) + affine_out_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.int8 + ) + affine = builder.call_operator( + op=exir_ops.edge.fused_quant.convolution_channels_last.default, + args=( + inp, + weight, + bias, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + 0, + 0, + *affine_out_qparams, + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + 1, + ), + ) + # NHWC [1,5,5,4] -> NCHW [1,4,5,5] + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(affine, [0, 3, 1, 2]), + ) + add_inp_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.float32 + ) + add_out_qparams = create_per_tensor_qparams( + builder, scale=0.125, dtype=torch.int8 + ) + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.Scalar, + args=(permute, *add_inp_qparams, *add_out_qparams, 2.0, 0.25), + ) + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(add, 0.125, 0, -128, 127, torch.int8), + ) + builder.output([dequantize]) + ep = builder.get_program() + + result = FuseAddIntoLinear().call(ep) + + self.assertTrue(result.modified) + graph = result.exported_program.graph_module.graph + self.assertEqual( + graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.add.Scalar + ), + [], + ) + affine_node = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution_channels_last.default, + )[0] + bias_node = get_arg(affine_node, "bias", torch.fx.Node) + actual_bias = get_constant(result.exported_program, bias_node) + torch.testing.assert_close(actual_bias, expected_bias) + + def test_scalar_add_not_folded_with_per_channel_qparams_through_passthrough( + self, + ) -> None: + builder = ProgramBuilder() + out_channels = 4 + inp_value = torch.randint(-16, 16, (2, 8), dtype=torch.int8) + weight_value = torch.randint(-16, 16, (out_channels, 8), dtype=torch.int8) + bias_value = torch.arange(out_channels, dtype=torch.float32) + + inp = builder.placeholder("x", inp_value) + weight = builder.placeholder( + "weight", weight_value, input_kind=InputKind.BUFFER + ) + bias = builder.placeholder("bias", bias_value, input_kind=InputKind.BUFFER) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_channel_qparams( + builder, out_channels, dtype=torch.float32 + ) + affine_out_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.int8 + ) + affine = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + inp, + weight, + bias, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + 0, + 0, + *affine_out_qparams, + ), + ) + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(affine, [0, 1]), + ) + add_inp_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.float32 + ) + add_out_qparams = create_per_axis_qparams( + builder, + out_channels, + ndim=2, + axis=1, + ) + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.Scalar, + args=(permute, *add_inp_qparams, *add_out_qparams, 0.0, 1.0), + ) + builder.output([add]) + ep = builder.get_program() + + result = FuseAddIntoLinear().call(ep) + + self.assertFalse(result.modified) + + def test_scalar_add_not_folded_when_affine_has_multiple_users(self) -> None: + ep, _ = self._build_affine_add_scalar( + exir_ops.edge.fused_quant.linear.default, True + ) + graph = ep.graph_module.graph + linear = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + )[0] + output = graph.find_nodes(op="output")[0] + with graph.inserting_before(output): + extra_user = graph.call_function( + exir_ops.edge.aten.clone.default, args=(linear,) + ) + extra_user.meta["val"] = linear.meta["val"] + output.args = ([output.args[0][0], extra_user],) + ep.graph_module.recompile() + + result = FuseAddIntoLinear().call(ep) + + self.assertFalse(result.modified) diff --git a/backends/fused_quant/test/test_fuse_add_softmax_into_masked_softmax.py b/backends/fused_quant/test/test_fuse_add_softmax_into_masked_softmax.py new file mode 100644 index 00000000000..8df938fac71 --- /dev/null +++ b/backends/fused_quant/test/test_fuse_add_softmax_into_masked_softmax.py @@ -0,0 +1,351 @@ +# 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. + + +import unittest +from typing import cast + +import torch +import torch.fx as fx +from executorch.backends.fused_quant.pre_quantize_passes.fuse_add_softmax_into_masked_softmax import ( + _has_fully_masked_row, + _is_additive_mask, + FuseAddSoftmaxIntoMaskedSoftmax, +) +from executorch.backends.fused_quant.test.helpers import trace +from parameterized import parameterized +from torch import nn +from torch.export import ExportedProgram + +_ADD = torch.ops.aten.add.Tensor +_SAFE_SOFTMAX = torch.ops.aten._safe_softmax.default +_MASKED_SOFTMAX = torch.ops.aten._masked_softmax.default +_INDEX = torch.ops.aten.index.Tensor +_MASK_TYPE_FULL = 2 + + +def _causal_mask(ctx: int) -> torch.Tensor: + """Binary additive causal mask: 0 on/below the diagonal, -inf above.""" + mask = torch.zeros(ctx, ctx) + mask.masked_fill_( + torch.triu(torch.ones(ctx, ctx), diagonal=1).bool(), float("-inf") + ) + return mask + + +class _DirectMaskModel(nn.Module): + """add(scores, mask) -> _safe_softmax, mask used as a direct constant buffer.""" + + def __init__(self, ctx: int) -> None: + super().__init__() + self.register_buffer("mask", _causal_mask(ctx).reshape(1, 1, ctx, ctx)) + + def forward(self, scores: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._safe_softmax.default(scores + self.mask, -1) + + +class _GatheredMaskModel(nn.Module): + """add(scores, mask[pos]) -> _safe_softmax: the MEC mask-gathered-by-pos form.""" + + def __init__(self, ctx: int) -> None: + super().__init__() + self.register_buffer("mask", _causal_mask(ctx)) + + def forward(self, scores: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._safe_softmax.default(scores + self.mask[pos], -1) + + +class _NoMaskModel(nn.Module): + """Plain softmax with no preceding mask add -- the pass must not fire.""" + + def forward(self, scores: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._safe_softmax.default(scores, -1) + + +class _NonZeroKeepMaskModel(nn.Module): + """add(scores, mask) -> _safe_softmax where the mask's keep value is 5.0, not 0. + + Not the canonical {0, -inf} additive mask, so the pass declines to fuse it.""" + + def __init__(self, ctx: int) -> None: + super().__init__() + mask = torch.full((ctx, ctx), 5.0) + mask.masked_fill_( + torch.triu(torch.ones(ctx, ctx), diagonal=1).bool(), float("-inf") + ) + self.register_buffer("mask", mask.reshape(1, 1, ctx, ctx)) + + def forward(self, scores: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._safe_softmax.default(scores + self.mask, -1) + + +class _NonConstAddModel(nn.Module): + """softmax over add(scores, bias) where bias is an activation, not a constant.""" + + def forward(self, scores: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._safe_softmax.default(scores + bias, -1) + + +class _AlphaMaskModel(nn.Module): + """add(scores, mask, alpha=2) -> _safe_softmax: a scaled add. + + add.Tensor computes self + alpha*other; only alpha=1 combines the operands + as-is, so the recompose must decline a non-default alpha.""" + + def __init__(self, ctx: int) -> None: + super().__init__() + self.register_buffer("mask", _causal_mask(ctx).reshape(1, 1, ctx, ctx)) + + def forward(self, scores: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._safe_softmax.default( + torch.ops.aten.add.Tensor(scores, self.mask, alpha=2), -1 + ) + + +def _count(ep: ExportedProgram, target: object) -> int: + return len(ep.graph_module.graph.find_nodes(op="call_function", target=target)) + + +def _only(ep: ExportedProgram, target: object) -> torch.fx.Node: + nodes = ep.graph_module.graph.find_nodes(op="call_function", target=target) + assert len(nodes) == 1 + return nodes[0] + + +def _lifted_fqns(ep: ExportedProgram) -> set[str]: + """Fully-qualified names of every lifted constant (params/buffers in state_dict + plus lifted tensor constants).""" + return set(ep.state_dict.keys()) | set(ep.constants.keys()) + + +class FuseAddSoftmaxIntoMaskedSoftmaxTest(unittest.TestCase): + _CTX = 4 + + def _build(self, gathered: bool) -> tuple[nn.Module, tuple[torch.Tensor, ...]]: + torch.manual_seed(0) + if gathered: + model = _GatheredMaskModel(self._CTX) + # scores: [B, H, seq, ctx]; pos selects `seq` rows of the mask. + inputs = (torch.randn(1, 2, 3, self._CTX), torch.tensor([0, 1, 2])) + else: + model = _DirectMaskModel(self._CTX) + inputs = (torch.randn(1, 2, self._CTX, self._CTX),) + return model, inputs + + @parameterized.expand([("direct", False), ("gathered", True)]) + def test_fuses_to_masked_softmax(self, _name: str, gathered: bool) -> None: + model, inputs = self._build(gathered) + ep = trace(model, inputs) + + result = FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + self.assertTrue(result.modified) + ep = result.exported_program + + # The add + safe_softmax pair is replaced by a single _masked_softmax. + self.assertEqual(_count(ep, _MASKED_SOFTMAX), 1) + self.assertEqual(_count(ep, _SAFE_SOFTMAX), 0) + self.assertEqual(_count(ep, _ADD), 0) + + masked_softmax = _only(ep, _MASKED_SOFTMAX) + _scores, mask, dim, mask_type = masked_softmax.args + mask = cast(fx.Node, mask) + # The mask operand is a bool tensor (never quantized), broadcast to the + # full scores shape for mask_type 2. + self.assertEqual(mask.meta["val"].dtype, torch.bool) + self.assertEqual( + tuple(mask.meta["val"].shape), + (1, 2, 3, self._CTX) if gathered else (1, 2, self._CTX, self._CTX), + ) + self.assertEqual(dim, -1) + self.assertEqual(mask_type, _MASK_TYPE_FULL) + + def test_gathered_mask_index_is_rewired_to_bool(self) -> None: + """The runtime index gather is preserved but now reads a bool buffer.""" + model, inputs = self._build(gathered=True) + ep = trace(model, inputs) + + FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + + index_nodes = ep.graph_module.graph.find_nodes( + op="call_function", target=_INDEX + ) + self.assertEqual(len(index_nodes), 1) + self.assertEqual(index_nodes[0].meta["val"].dtype, torch.bool) + + @parameterized.expand([("direct", False), ("gathered", True)]) + def test_numerically_equivalent(self, _name: str, gathered: bool) -> None: + model, inputs = self._build(gathered) + with torch.no_grad(): + ref = model(*inputs) + + ep = trace(model, inputs) + FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + with torch.no_grad(): + got = ep.module()(*inputs) + + # The mask uses -inf, so additive softmax and exact masking agree exactly + # (every causal row has at least one unmasked position -> no NaN guard). + torch.testing.assert_close(got, ref, atol=1e-6, rtol=1e-6) + + def test_source_fn_stack_is_copied(self) -> None: + """The quantizer's matcher skips nodes without source_fn_stack, so the + recompose must carry it from the softmax onto the new op.""" + model, inputs = self._build(gathered=False) + ep = trace(model, inputs) + sentinel = [("attn_softmax", torch.nn.functional.softmax)] + _only(ep, _SAFE_SOFTMAX).meta["source_fn_stack"] = sentinel + + FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + + self.assertEqual( + _only(ep, _MASKED_SOFTMAX).meta.get("source_fn_stack"), sentinel + ) + + def test_no_mask_is_noop(self) -> None: + torch.manual_seed(0) + ep = trace(_NoMaskModel(), (torch.randn(1, 2, self._CTX, self._CTX),)) + + result = FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + + self.assertFalse(result.modified) + self.assertEqual(_count(result.exported_program, _SAFE_SOFTMAX), 1) + self.assertEqual(_count(result.exported_program, _MASKED_SOFTMAX), 0) + + def test_non_constant_add_operand_is_noop(self) -> None: + """An add whose operands are both activations is not a mask add.""" + torch.manual_seed(0) + scores = torch.randn(1, 2, self._CTX, self._CTX) + bias = torch.randn(1, 2, self._CTX, self._CTX) + ep = trace(_NonConstAddModel(), (scores, bias)) + + result = FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + + self.assertFalse(result.modified) + self.assertEqual(_count(result.exported_program, _MASKED_SOFTMAX), 0) + self.assertEqual(_count(result.exported_program, _SAFE_SOFTMAX), 1) + + def test_nonzero_keep_value_is_noop(self) -> None: + """A binary mask whose keep positions are nonzero is not the canonical + {0, -inf} form, so the pass conservatively declines to fuse it.""" + torch.manual_seed(0) + ep = trace( + _NonZeroKeepMaskModel(self._CTX), + (torch.randn(1, 2, self._CTX, self._CTX),), + ) + + result = FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + + self.assertFalse(result.modified) + self.assertEqual(_count(result.exported_program, _MASKED_SOFTMAX), 0) + self.assertEqual(_count(result.exported_program, _SAFE_SOFTMAX), 1) + + def test_alpha_not_one_is_noop(self) -> None: + """add(scores, mask, alpha != 1) scales an operand, so it is not fusible.""" + torch.manual_seed(0) + ep = trace( + _AlphaMaskModel(self._CTX), + (torch.randn(1, 2, self._CTX, self._CTX),), + ) + + result = FuseAddSoftmaxIntoMaskedSoftmax().call(ep) + + self.assertFalse(result.modified) + self.assertEqual(_count(result.exported_program, _MASKED_SOFTMAX), 0) + self.assertEqual(_count(result.exported_program, _SAFE_SOFTMAX), 1) + + @parameterized.expand([("direct", False), ("gathered", True)]) + def test_additive_mask_constant_is_removed( + self, _name: str, gathered: bool + ) -> None: + """After fusion the additive mask buffer is gone from the program -- its + lifted constant, placeholder, and input spec are all removed (constant_prop + drops the now-dead constant) -- leaving the graph signature consistent.""" + model, inputs = self._build(gathered) + ep = trace(model, inputs) + # Sanity: the additive "mask" buffer is a lifted constant before fusion. + self.assertIn("mask", _lifted_fqns(ep)) + + ep = FuseAddSoftmaxIntoMaskedSoftmax().call(ep).exported_program + + # The additive mask is gone from state_dict/constants... + self.assertNotIn("mask", _lifted_fqns(ep)) + # ...and the signature stays consistent: one input spec per placeholder. + self.assertEqual( + len(ep.graph_signature.input_specs), + len(ep.graph_module.graph.find_nodes(op="placeholder")), + ) + + +class IsAdditiveMaskTest(unittest.TestCase): + """Direct coverage of the additive-mask predicate: only the canonical + {0, very-negative} two-value float tensor qualifies.""" + + def test_canonical_zero_and_neg_inf(self) -> None: + t = torch.tensor([[0.0, float("-inf")], [0.0, 0.0]]) + self.assertTrue(_is_additive_mask(t)) + + def test_canonical_zero_and_finite_sentinel(self) -> None: + t = torch.tensor([[0.0, -1e4], [0.0, 0.0]]) + self.assertTrue(_is_additive_mask(t)) + + def test_nonzero_keep_value_rejected(self) -> None: + t = torch.tensor([[5.0, float("-inf")], [5.0, 5.0]]) + self.assertFalse(_is_additive_mask(t)) + + def test_single_all_masked_value_rejected(self) -> None: + # A uniform very-negative tensor would softmax to uniform weights, not + # zeros -- it must not be treated as a mask. + t = torch.full((2, 2), -1e4) + self.assertFalse(_is_additive_mask(t)) + + def test_negative_value_above_threshold_rejected(self) -> None: + # The non-zero value is negative but not a masking sentinel. + t = torch.tensor([[0.0, -5.0], [0.0, 0.0]]) + self.assertFalse(_is_additive_mask(t)) + + def test_integer_tensor_rejected(self) -> None: + t = torch.tensor([[0, -10000], [0, 0]], dtype=torch.int64) + self.assertFalse(_is_additive_mask(t)) + + +class HasFullyMaskedRowTest(unittest.TestCase): + """Direct coverage of the fully-masked-row guard. Inputs are canonical additive + masks ({0, sentinel}), which is what _has_fully_masked_row assumes (it runs + only after _is_additive_mask).""" + + def test_causal_mask_has_no_fully_masked_row(self) -> None: + # 0 on/below the diagonal, -inf above: every row keeps its diagonal. + self.assertFalse(_has_fully_masked_row(_causal_mask(4), dim=-1, scores_rank=2)) + + def test_all_masked_row_is_detected(self) -> None: + mask = _causal_mask(4) + mask[2] = float("-inf") # query row 2 attends to nothing + self.assertTrue(_has_fully_masked_row(mask, dim=-1, scores_rank=2)) + + def test_finite_sentinel_row_is_detected(self) -> None: + # The sentinel need not be -inf; -1e4 is the min and marks masked positions. + mask = torch.zeros(3, 3) + mask[1] = -1e4 + self.assertTrue(_has_fully_masked_row(mask, dim=-1, scores_rank=2)) + + def test_lower_rank_mask_maps_dim_from_the_right(self) -> None: + # Gathered form: a [ctx, ctx] buffer checked against rank-4 scores -- the + # softmax dim -1 still lines up with the mask's last axis. + mask = _causal_mask(4) + self.assertFalse(_has_fully_masked_row(mask, dim=-1, scores_rank=4)) + mask[0] = float("-inf") + self.assertTrue(_has_fully_masked_row(mask, dim=-1, scores_rank=4)) + + def test_positive_dim_is_normalized(self) -> None: + # dim=3 on rank-4 scores is the same axis as dim=-1. + mask = _causal_mask(4).reshape(1, 1, 4, 4) + self.assertFalse(_has_fully_masked_row(mask, dim=3, scores_rank=4)) + + def test_mask_not_spanning_softmax_dim_is_conservative(self) -> None: + # A 1-D key mask softmaxed over a dim it doesn't span: the mask broadcasts + # along that dim, so any masked element implies a fully-masked row. + mask = torch.tensor([0.0, 0.0, -1e4]) + self.assertTrue(_has_fully_masked_row(mask, dim=-2, scores_rank=3)) diff --git a/backends/fused_quant/test/test_fuse_mul_into_linear.py b/backends/fused_quant/test/test_fuse_mul_into_linear.py new file mode 100644 index 00000000000..0fdd7c8690d --- /dev/null +++ b/backends/fused_quant/test/test_fuse_mul_into_linear.py @@ -0,0 +1,1125 @@ +# 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. + + +import unittest +from typing import Any, Optional + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.optimization_passes.fuse_mul_into_linear import ( + FuseMulIntoLinear, +) +from executorch.backends.fused_quant.test.helpers import ( + create_per_channel_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import get_constant +from executorch.exir.dialects._ops import ops as exir_ops +from parameterized import parameterized +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +class FuseMulIntoLinearTest(unittest.TestCase): + """Tests for the FuseMulIntoLinear optimization pass.""" + + IN_FEATURES = 8 + OUT_FEATURES = 16 + + def _build_linear_mul( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + mul_scale: torch.Tensor | float, + x_scale: float = 1.0, + w_scale_per_ch: Optional[torch.Tensor] = None, + w_zero_point_per_ch: Optional[torch.Tensor] = None, + per_tensor_weight_scale: bool = False, + linear_out_scale: float = 0.5, + mul_out_scale: float = 0.125, + mul_other_scale: float | None = None, + ) -> ExportedProgram: + out_f = weight.shape[0] + builder = ProgramBuilder() + + x_ph = builder.placeholder("x", x) + w_ph = builder.placeholder("weight", weight, input_kind=InputKind.BUFFER) + b_ph = ( + builder.placeholder("bias", bias, input_kind=InputKind.BUFFER) + if bias is not None + else None + ) + scale_ph = ( + builder.placeholder("scale", mul_scale, input_kind=InputKind.BUFFER) + if isinstance(mul_scale, torch.Tensor) + else None + ) + + inp_qparams = create_per_tensor_qparams( + builder, scale=x_scale, dtype=torch.float32 + ) + weight_qparams = ( + create_per_tensor_qparams(builder, dtype=torch.float32) + if per_tensor_weight_scale + else create_per_channel_qparams(builder, out_f, dtype=torch.float32) + ) + if w_scale_per_ch is not None: + builder.state_dict["wt_scale_0"] = w_scale_per_ch.clone().reshape(-1, 1) + if w_zero_point_per_ch is not None: + builder.state_dict["wt_zp_0"] = w_zero_point_per_ch.clone().reshape(-1, 1) + + out_qparams = create_per_tensor_qparams( + builder, scale=linear_out_scale, dtype=torch.int8 + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x_ph, + w_ph, + b_ph, + *inp_qparams, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + mul_inp_qp = create_per_tensor_qparams( + builder, scale=linear_out_scale, dtype=torch.float32 + ) + mul_out_qp = create_per_tensor_qparams( + builder, scale=mul_out_scale, dtype=torch.int8 + ) + + if scale_ph is not None: + if mul_other_scale is None: + mul_other_qp: tuple[Any, ...] = ( + None, + None, + torch.float32, + 0, + 0, + ) + else: + other_scale, other_zero_point, other_dtype, _, _ = ( + create_per_tensor_qparams( + builder, + scale=mul_other_scale, + dtype=torch.float32, + ) + ) + mul_other_qp = ( + other_scale, + other_zero_point, + other_dtype, + -32768, + 32767, + ) + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.default, + args=(linear, scale_ph, *mul_inp_qp, *mul_other_qp, *mul_out_qp), + ) + else: + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.Scalar, + args=(linear, *mul_inp_qp, *mul_out_qp, mul_scale), + ) + + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(mul, mul_out_scale, 0, -128, 127, torch.int8), + ) + builder.output([dequantize]) + + return builder.get_program() + + def _default_linear_mul(self, with_bias: bool = True) -> ExportedProgram: + x = torch.randn(1, self.IN_FEATURES) + weight = torch.randint( + -16, 16, (self.OUT_FEATURES, self.IN_FEATURES), dtype=torch.int8 + ) + bias = torch.randn(self.OUT_FEATURES) if with_bias else None + mul_scale = torch.randn(self.OUT_FEATURES).abs().clamp(min=0.5, max=2.0) + return self._build_linear_mul(x, weight, bias, mul_scale) + + def test_mul_folded_into_linear(self) -> None: + ep = self._default_linear_mul(with_bias=True) + + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + mul_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.mul.default, + ) + self.assertEqual(len(mul_nodes), 0, "Mul node should be removed") + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1, "Linear node should remain") + + def test_mul_folded_no_bias(self) -> None: + ep = self._default_linear_mul(with_bias=False) + + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + + mul_nodes = result.exported_program.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.mul.default, + ) + self.assertEqual(len(mul_nodes), 0, "Mul node should be removed") + + def test_scalar_mul_folded_with_per_tensor_weight_scale(self) -> None: + x = torch.randn(1, self.IN_FEATURES) + weight = torch.randint( + -16, + 16, + (self.OUT_FEATURES, self.IN_FEATURES), + dtype=torch.int8, + ) + ep = self._build_linear_mul( + x, + weight, + torch.randn(self.OUT_FEATURES), + torch.tensor(2.0), + per_tensor_weight_scale=True, + ) + + result = FuseMulIntoLinear().call(ep) + + self.assertTrue(result.modified) + linear = result.exported_program.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + )[0] + weight_scale_node = get_arg(linear, "weight_scale", torch.fx.Node) + weight_scale = get_constant(result.exported_program, weight_scale_node) + self.assertIsNotNone(weight_scale) + self.assertEqual(weight_scale.shape, torch.Size([])) + + def test_scalar_mul_folded_into_linear(self) -> None: + ep = self._build_linear_mul( + torch.randn(1, self.IN_FEATURES), + torch.randint( + -16, + 16, + (self.OUT_FEATURES, self.IN_FEATURES), + dtype=torch.int8, + ), + torch.randn(self.OUT_FEATURES), + 1.5, + ) + + result = FuseMulIntoLinear().call(ep) + + self.assertTrue(result.modified) + scalar_mul_nodes = result.exported_program.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.mul.Scalar, + ) + self.assertEqual(len(scalar_mul_nodes), 0, "Scalar mul should be removed") + + def test_negative_scalar_mul_preserves_numerics(self) -> None: + x = torch.tensor([[1, -1, 2, -2, 1, 0, -1, 2]], dtype=torch.int8) + weight = ( + torch.arange(self.OUT_FEATURES * self.IN_FEATURES) + .reshape(self.OUT_FEATURES, self.IN_FEATURES) + .remainder(5) + .sub(2) + .to(torch.int8) + ) + bias = torch.arange(self.OUT_FEATURES).remainder(3).sub(1).to(torch.float32) + weight_zero_point = ( + torch.arange(self.OUT_FEATURES).remainder(5).sub(2).to(torch.int64) + ) + multiplier = -2.0 + ep = self._build_linear_mul( + x, + weight, + bias, + multiplier, + w_scale_per_ch=torch.ones(self.OUT_FEATURES), + w_zero_point_per_ch=weight_zero_point, + linear_out_scale=1.0, + mul_out_scale=1.0, + ) + (unfused_output,) = ep.module()(x) + + result = FuseMulIntoLinear()(ep) + + self.assertTrue(result.modified) + fused_ep = result.exported_program + (fused_output,) = fused_ep.module()(x) + torch.testing.assert_close(fused_output, unfused_output, rtol=0, atol=0) + + linear = fused_ep.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + )[0] + weight_node = get_arg(linear, "weight", torch.fx.Node) + weight_scale_node = get_arg(linear, "weight_scale", torch.fx.Node) + weight_zero_point_node = get_arg(linear, "weight_zero_point", torch.fx.Node) + bias_node = get_arg(linear, "bias", torch.fx.Node) + actual_weight = get_constant(fused_ep, weight_node) + actual_weight_scale = get_constant(fused_ep, weight_scale_node) + actual_weight_zero_point = get_constant(fused_ep, weight_zero_point_node) + actual_bias = get_constant(fused_ep, bias_node) + self.assertIsNotNone(actual_weight) + self.assertIsNotNone(actual_weight_scale) + self.assertIsNotNone(actual_weight_zero_point) + self.assertIsNotNone(actual_bias) + torch.testing.assert_close( + actual_weight, + (-1 - weight.to(torch.int16)).to(torch.int8), + ) + torch.testing.assert_close( + actual_weight_scale, + torch.full((self.OUT_FEATURES, 1), abs(multiplier)), + ) + torch.testing.assert_close( + actual_weight_zero_point, + -1 - weight_zero_point.reshape(-1, 1), + ) + torch.testing.assert_close(actual_bias, bias * multiplier) + + def test_vector_mul_preserves_numerics(self) -> None: + x = torch.tensor([[1, -1, 2, -2, 1, 0, -1, 2]], dtype=torch.int8) + weight = ( + torch.arange(self.OUT_FEATURES * self.IN_FEATURES) + .reshape(self.OUT_FEATURES, self.IN_FEATURES) + .remainder(5) + .sub(2) + .to(torch.int8) + ) + bias = torch.arange(self.OUT_FEATURES).remainder(3).sub(1).to(torch.float32) + multiplier = torch.tensor([1.0, 2.0] * (self.OUT_FEATURES // 2)) + ep = self._build_linear_mul( + x, + weight, + bias, + multiplier, + w_scale_per_ch=torch.ones(self.OUT_FEATURES), + linear_out_scale=1.0, + mul_out_scale=1.0, + ) + (unfused_output,) = ep.module()(x) + + result = FuseMulIntoLinear()(ep) + + self.assertTrue(result.modified) + fused_ep = result.exported_program + self._assert_mul_removed(fused_ep) + (fused_output,) = fused_ep.module()(x) + torch.testing.assert_close(fused_output, unfused_output, rtol=0, atol=0) + + linear = fused_ep.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + )[0] + weight_scale_node = get_arg(linear, "weight_scale", torch.fx.Node) + bias_node = get_arg(linear, "bias", torch.fx.Node) + actual_weight_scale = get_constant(fused_ep, weight_scale_node) + actual_bias = get_constant(fused_ep, bias_node) + self.assertIsNotNone(actual_weight_scale) + self.assertIsNotNone(actual_bias) + torch.testing.assert_close(actual_weight_scale, multiplier.reshape(-1, 1)) + torch.testing.assert_close(actual_bias, bias * multiplier) + + def test_quantized_constant_mul_is_dequantized_before_fold(self) -> None: + represented_value = 0.5 + quantized_value = torch.tensor(32767, dtype=torch.int16) + ep = self._build_linear_mul( + torch.randn(1, self.IN_FEATURES), + torch.randint( + -16, + 16, + (self.OUT_FEATURES, self.IN_FEATURES), + dtype=torch.int8, + ), + torch.randn(self.OUT_FEATURES), + quantized_value, + mul_other_scale=represented_value / int(quantized_value), + ) + + result = FuseMulIntoLinear().call(ep) + + self.assertTrue(result.modified) + linear = result.exported_program.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + )[0] + weight_scale_node = get_arg(linear, "weight_scale", torch.fx.Node) + weight_scale = get_constant(result.exported_program, weight_scale_node) + self.assertIsNotNone(weight_scale) + torch.testing.assert_close( + weight_scale, + torch.full((self.OUT_FEATURES, 1), represented_value), + ) + + def test_graph_signature_updated(self) -> None: + """After folding, the scale placeholder used by mul becomes dead and + should be removed from the graph signature by constant_prop_pass.""" + ep = self._default_linear_mul(with_bias=True) + + input_names_before = {s.arg.name for s in ep.graph_signature.input_specs} + self.assertIn("scale", input_names_before) + + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + + fused_ep = result.exported_program + input_names_after = {s.arg.name for s in fused_ep.graph_signature.input_specs} + placeholder_names = { + n.name for n in fused_ep.graph_module.graph.find_nodes(op="placeholder") + } + self.assertEqual( + input_names_after, + placeholder_names, + "Graph signature input_specs should match actual placeholder nodes", + ) + self.assertNotIn( + "scale", + input_names_after, + "Dead scale placeholder should be removed from signature", + ) + + def test_numerics_improve(self) -> None: + """Fusing mul into linear eliminates an intermediate quantization + boundary, so the fused result should be closer to the true fp32 + result than the unfused quantized graph. + + We build an executable fused_quant graph (linear->mul->dequant), + run it before and after the optimization, and compare both outputs + against a pure fp32 reference. + """ + torch.manual_seed(42) + in_f, out_f = self.IN_FEATURES, self.OUT_FEATURES + batch = 4 + + x_float = torch.randn(batch, in_f) + weight_float = torch.randn(out_f, in_f) + bias_float = torch.randn(out_f) + mul_scale_val = torch.rand(out_f).clamp(min=0.5, max=2.0) + + # fp32 reference: linear(x, w, b) * scale + fp32_ref = ( + torch.nn.functional.linear(x_float, weight_float, bias_float) + * mul_scale_val + ) + + # Compute scales from data ranges: (max - min) / (qmax - qmin) + x_scale = (x_float.max() - x_float.min()).item() / 255.0 + w_scale_per_ch = ( + weight_float.max(dim=1).values - weight_float.min(dim=1).values + ) / 255.0 + + # Quantize tensors + x_int8 = torch.clamp(torch.round(x_float / x_scale), -128, 127).to(torch.int8) + weight_int8 = torch.clamp( + torch.round(weight_float / w_scale_per_ch.unsqueeze(1)), -128, 127 + ).to(torch.int8) + + # Compute output scales from approximate output ranges + linear_approx = torch.nn.functional.linear(x_float, weight_float, bias_float) + linear_out_scale = (linear_approx.max() - linear_approx.min()).item() / 255.0 + mul_approx = linear_approx * mul_scale_val + mul_out_scale = (mul_approx.max() - mul_approx.min()).item() / 255.0 + + ep = self._build_linear_mul( + x_int8, + weight_int8, + bias_float, + mul_scale_val, + x_scale=x_scale, + w_scale_per_ch=w_scale_per_ch, + linear_out_scale=linear_out_scale, + mul_out_scale=mul_out_scale, + ) + + (unfused_out,) = ep.module()(x_int8) + + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + (fused_out,) = result.exported_program.module()(x_int8) + + unfused_err = (unfused_out.float() - fp32_ref).abs().mean() + fused_err = (fused_out.float() - fp32_ref).abs().mean() + + self.assertLess( + fused_err, + unfused_err, + f"Fused error ({fused_err:.6f}) should be < unfused error ({unfused_err:.6f})", + ) + + def test_no_fold_when_linear_has_multiple_users(self) -> None: + """Mul should not be folded if the linear output is used elsewhere.""" + builder = ProgramBuilder() + + x = builder.placeholder("x", torch.randn(1, 8)) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (16, 8), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + scale = builder.placeholder( + "scale", + torch.randn(16).abs() + 0.1, + input_kind=InputKind.BUFFER, + ) + + inp_qparams = create_per_tensor_qparams(builder) + weight_qparams = create_per_channel_qparams(builder, 16) + out_qparams = create_per_tensor_qparams(builder) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + None, + *inp_qparams, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + mul_inp_qp = create_per_tensor_qparams(builder) + mul_other_qp = create_per_tensor_qparams(builder) + mul_out_qp = create_per_tensor_qparams(builder) + + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.default, + args=(linear, scale, *mul_inp_qp, *mul_other_qp, *mul_out_qp), + ) + + dequant_linear = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(linear, 0.5, 0, -128, 127, torch.int8), + ) + dequant_mul = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(mul, 0.125, 0, -128, 127, torch.int8), + ) + builder.output([dequant_linear, dequant_mul]) + + ep = builder.get_program() + result = FuseMulIntoLinear().call(ep) + + self.assertFalse( + result.modified, "Should not fold when linear has multiple users" + ) + + def test_no_fold_when_weight_scale_is_shared(self) -> None: + """Mul should not be folded if weight_scale is used by another op.""" + builder = ProgramBuilder() + + x = builder.placeholder("x", torch.randn(1, 8)) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (16, 8), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + scale = builder.placeholder( + "scale", + torch.randn(16).abs() + 0.1, + input_kind=InputKind.BUFFER, + ) + + inp_qparams = create_per_tensor_qparams(builder) + weight_qparams = create_per_channel_qparams(builder, 16) + out_qparams = create_per_tensor_qparams(builder) + + linear1 = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + None, + *inp_qparams, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + inp_qparams2 = create_per_tensor_qparams(builder) + out_qparams2 = create_per_tensor_qparams(builder) + + linear2 = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + None, + *inp_qparams2, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams2, + ), + ) + + mul_inp_qp = create_per_tensor_qparams(builder) + mul_other_qp = create_per_tensor_qparams(builder) + mul_out_qp = create_per_tensor_qparams(builder) + + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.default, + args=(linear1, scale, *mul_inp_qp, *mul_other_qp, *mul_out_qp), + ) + + dequant_mul = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(mul, 0.125, 0, -128, 127, torch.int8), + ) + dequant_linear2 = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(linear2, 0.5, 0, -128, 127, torch.int8), + ) + builder.output([dequant_mul, dequant_linear2]) + + ep = builder.get_program() + result = FuseMulIntoLinear().call(ep) + + self.assertFalse( + result.modified, + "Should not fold when weight_scale is shared with another linear", + ) + + def _build_linear_passthrough_mul( + self, + passthrough: str = "permute", + constant_shape: tuple[int, ...] = (1,), + linear_out_dtype: torch.dtype = torch.int8, + mul_out_dtype: torch.dtype = torch.int8, + constant_dtype: torch.dtype = torch.float32, + view_constant: bool = False, + scalar: Optional[float] = None, + ) -> tuple[ExportedProgram, torch.Tensor]: + """Build a linear → passthrough → mul(constant) graph. + + Args: + passthrough: "permute", "view", or "permute_view". + constant_shape: Shape of the mul constant. + linear_out_dtype: Output dtype of the linear's out qparams. + mul_out_dtype: Output dtype of the mul's out qparams. + constant_dtype: Dtype of the mul constant tensor. + view_constant: Reshape the constant through view_copy before the mul. + + Returns: + (program, input_tensor) + """ + + builder = ProgramBuilder() + in_features, out_features = 8, 16 + inp_tensor = torch.randn(1, in_features) + + x = builder.placeholder("x", inp_tensor) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + b = builder.placeholder( + "bias", torch.randn(out_features), input_kind=InputKind.BUFFER + ) + scale = ( + builder.placeholder( + "scale", + torch.randn(*constant_shape) + .abs() + .clamp(min=0.5, max=2.0) + .to(constant_dtype), + input_kind=InputKind.BUFFER, + ) + if scalar is None + else None + ) + + weight_qparams = create_per_channel_qparams( + builder, out_features, scale=0.05, dtype=torch.float32 + ) + out_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=linear_out_dtype + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + b, + None, + None, + torch.float32, + 0, + 0, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + node = linear + if "permute" in passthrough: + node = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(node, [0, 1]), + ) + if "view" in passthrough: + node = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(node, [1, out_features]), + ) + + mul_scale = scale + if scale is not None and view_constant: + mul_scale = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(scale, [1, out_features]), + ) + + mul_inp_qp = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + mul_out_qp = create_per_tensor_qparams( + builder, scale=0.125, dtype=mul_out_dtype + ) + + if mul_scale is not None: + mul_other_qp = create_per_tensor_qparams( + builder, scale=0.25, dtype=torch.float32 + ) + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.default, + args=(node, mul_scale, *mul_inp_qp, *mul_other_qp, *mul_out_qp), + ) + else: + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.Scalar, + args=(node, *mul_inp_qp, *mul_out_qp, scalar), + ) + + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=( + mul, + 0.125, + 0, + torch.iinfo(mul_out_dtype).min, + torch.iinfo(mul_out_dtype).max, + mul_out_dtype, + ), + ) + builder.output([dequantize]) + return builder.get_program(), inp_tensor + + def _run_ep(self, ep: ExportedProgram, inp: torch.Tensor) -> torch.Tensor: + full_inputs = [] + sig = ep.graph_signature + for spec in sig.input_specs: + if spec.kind == InputKind.USER_INPUT: + full_inputs.append(inp) + else: + assert spec.target is not None + full_inputs.append(ep.state_dict[spec.target]) + return ep.graph_module(*full_inputs) + + def _assert_mul_removed(self, ep: ExportedProgram) -> None: + mul_nodes = ep.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.mul.default + ) + self.assertEqual(len(mul_nodes), 0, "Mul should be removed") + + @parameterized.expand( + [ + ("permute_scalar", "permute", (1,)), + ("view_scalar", "view", (1,)), + ("permute_view_scalar", "permute_view", (1,)), + ] + ) + def test_mul_folded_through_passthrough( + self, _name: str, passthrough: str, constant_shape: tuple[int, ...] + ) -> None: + ep, _ = self._build_linear_passthrough_mul( + passthrough=passthrough, constant_shape=constant_shape + ) + + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + self._assert_mul_removed(result.exported_program) + + def test_scalar_mul_folded_through_passthrough(self) -> None: + ep, _ = self._build_linear_passthrough_mul( + passthrough="permute_view", scalar=1.5 + ) + + result = FuseMulIntoLinear().call(ep) + + self.assertTrue(result.modified) + scalar_mul_nodes = result.exported_program.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.mul.Scalar, + ) + self.assertEqual(len(scalar_mul_nodes), 0, "Scalar mul should be removed") + + def test_mul_not_folded_multi_dim_constant(self) -> None: + """Constant that varies in non-output dimensions should not fold.""" + + builder = ProgramBuilder() + in_features, out_features, batch = 8, 16, 4 + + x = builder.placeholder("x", torch.randn(batch, in_features)) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + scale = builder.placeholder( + "scale", + torch.randn(batch, 1).abs().clamp(min=0.5, max=2.0), + input_kind=InputKind.BUFFER, + ) + + weight_qparams = create_per_channel_qparams( + builder, out_features, scale=0.05, dtype=torch.float32 + ) + out_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + None, + None, + None, + torch.float32, + 0, + 0, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + mul_inp_qp = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + mul_other_qp = create_per_tensor_qparams( + builder, scale=0.25, dtype=torch.float32 + ) + mul_out_qp = create_per_tensor_qparams(builder, scale=0.125, dtype=torch.int8) + + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.default, + args=(linear, scale, *mul_inp_qp, *mul_other_qp, *mul_out_qp), + ) + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(mul, 0.125, 0, -128, 127, torch.int8), + ) + builder.output([dequantize]) + + ep = builder.get_program() + result = FuseMulIntoLinear().call(ep) + self.assertFalse( + result.modified, + "Should not fold constant that varies in non-output dimensions", + ) + + @parameterized.expand( + [ + ("permute", "permute", exir_ops.edge.aten.permute_copy.default), + ("view", "view", exir_ops.edge.aten.view_copy.default), + ] + ) + def test_mul_folded_through_passthrough_retypes_meta( + self, _name: str, passthrough: str, pt_target: Any + ) -> None: + """When the mul's output dtype differs from the linear's original + (int8 → uint8), folding a scalar through a passthrough must retype the + linear and every passthrough node's meta['val'] to the new dtype.""" + ep, _ = self._build_linear_passthrough_mul( + passthrough=passthrough, + constant_shape=(1,), + linear_out_dtype=torch.int8, + mul_out_dtype=torch.uint8, + ) + + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + self._assert_mul_removed(result.exported_program) + + graph = result.exported_program.graph_module.graph + linear_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + self.assertEqual(len(linear_nodes), 1) + self.assertEqual(linear_nodes[0].meta["val"].dtype, torch.uint8) + + pt_nodes = graph.find_nodes(op="call_function", target=pt_target) + self.assertEqual(len(pt_nodes), 1) + self.assertEqual( + pt_nodes[0].meta["val"].dtype, + torch.uint8, + "passthrough meta['val'] should take the mul's output dtype", + ) + + def test_mul_folded_fp16_constant(self) -> None: + """A non-float32 (fp16) constant still folds; multiplying it into the + float32 weight_scale promotes the result back to float32.""" + ep, _ = self._build_linear_passthrough_mul( + passthrough="permute", + constant_shape=(1,), + constant_dtype=torch.float16, + ) + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + self._assert_mul_removed(result.exported_program) + + fused_ep = result.exported_program + linear_nodes = fused_ep.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + self.assertEqual(len(linear_nodes), 1) + weight_scale_node = get_arg( + linear_nodes[0], "weight_scale", Optional[torch.fx.Node] + ) + assert isinstance(weight_scale_node, torch.fx.Node) + weight_scale = get_constant(fused_ep, weight_scale_node) + assert weight_scale is not None + self.assertEqual( + weight_scale.dtype, + torch.float32, + "weight_scale should stay float32 after folding an fp16 constant", + ) + + def test_mul_folded_quantized_constant(self) -> None: + """The constant may reach the mul through a quantize_per_tensor node; + _trace_to_constant recovers the float tensor behind it and folds.""" + + builder = ProgramBuilder() + in_features, out_features = self.IN_FEATURES, self.OUT_FEATURES + + x = builder.placeholder("x", torch.randn(1, in_features)) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + b = builder.placeholder( + "bias", torch.randn(out_features), input_kind=InputKind.BUFFER + ) + scale_fp = builder.placeholder( + "scale_fp", + torch.randn(1).abs().clamp(min=0.5, max=2.0), + input_kind=InputKind.BUFFER, + ) + + weight_qparams = create_per_channel_qparams( + builder, out_features, scale=0.05, dtype=torch.float32 + ) + out_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + b, + None, + None, + torch.float32, + 0, + 0, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + # The constant operand reaches the mul through a quantize_per_tensor. + quant_const = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(scale_fp, 0.01, 0, -128, 127, torch.int8), + ) + + mul_inp_qp = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + mul_other_qp = create_per_tensor_qparams( + builder, scale=0.01, dtype=torch.float32 + ) + mul_out_qp = create_per_tensor_qparams(builder, scale=0.125, dtype=torch.int8) + + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.default, + args=(linear, quant_const, *mul_inp_qp, *mul_other_qp, *mul_out_qp), + ) + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(mul, 0.125, 0, -128, 127, torch.int8), + ) + builder.output([dequantize]) + + ep = builder.get_program() + result = FuseMulIntoLinear().call(ep) + self.assertTrue(result.modified) + self._assert_mul_removed(result.exported_program) + + def test_mul_folded_with_viewed_quantized_constant(self) -> None: + ep, inp = self._build_linear_passthrough_mul( + passthrough="permute", + constant_shape=(self.OUT_FEATURES,), + constant_dtype=torch.int8, + view_constant=True, + ) + with torch.no_grad(): + (expected,) = ep.module()(inp) + + result = FuseMulIntoLinear().call(ep) + + self.assertTrue(result.modified) + self._assert_mul_removed(result.exported_program) + with torch.no_grad(): + (actual,) = result.exported_program.module()(inp) + # Folding drops the linear's output requantization, so values on a + # rounding boundary may move by one mul output step (scale 0.125). + torch.testing.assert_close(actual, expected, rtol=0, atol=0.125) + + def test_per_channel_mul_folded_through_passthrough(self) -> None: + ep, _ = self._build_linear_passthrough_mul( + passthrough="permute", + constant_shape=(1, 16), + ) + + result = FuseMulIntoLinear().call(ep) + + self.assertTrue(result.modified) + self._assert_mul_removed(result.exported_program) + + def test_mul_not_folded_when_constant_expands_rank(self) -> None: + """A scalar (numel 1) constant with more dims than the data broadcasts + the data up to a higher rank, so it must not fold (folding would drop + the added unary dimension).""" + x = torch.randn(1, self.IN_FEATURES) + weight = torch.randint( + -16, 16, (self.OUT_FEATURES, self.IN_FEATURES), dtype=torch.int8 + ) + bias = torch.randn(self.OUT_FEATURES) + # numel == 1 (scalar), but rank 3 vs the rank-2 (1, OUT_FEATURES) linear + # output: linear[1,16] * const[1,1,1] broadcasts to [1,1,16]. + mul_scale = torch.randn(1, 1, 1).abs().clamp(min=0.5, max=2.0) + ep = self._build_linear_mul(x, weight, bias, mul_scale) + + result = FuseMulIntoLinear().call(ep) + self.assertFalse( + result.modified, + "Should not fold a constant that broadcasts the data to a higher rank", + ) + + def test_mul_not_folded_wrong_axis_constant(self) -> None: + """A constant that is shape-preserving but varies along a non-output + dimension (here dim 0, with batch == out_features so numel matches) + must not fold.""" + + builder = ProgramBuilder() + # batch == out_features so a [out_features, 1] constant is both + # shape-preserving and numel-matching, yet varies on the wrong axis. + in_features = out_features = batch = 16 + + x = builder.placeholder("x", torch.randn(batch, in_features)) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + scale = builder.placeholder( + "scale", + torch.randn(out_features, 1).abs().clamp(min=0.5, max=2.0), + input_kind=InputKind.BUFFER, + ) + + weight_qparams = create_per_channel_qparams( + builder, out_features, scale=0.05, dtype=torch.float32 + ) + out_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + None, + None, + None, + torch.float32, + 0, + 0, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + mul_inp_qp = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + mul_other_qp = create_per_tensor_qparams( + builder, scale=0.25, dtype=torch.float32 + ) + mul_out_qp = create_per_tensor_qparams(builder, scale=0.125, dtype=torch.int8) + + mul = builder.call_operator( + op=exir_ops.edge.fused_quant.mul.default, + args=(linear, scale, *mul_inp_qp, *mul_other_qp, *mul_out_qp), + ) + dequantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(mul, 0.125, 0, -128, 127, torch.int8), + ) + builder.output([dequantize]) + + ep = builder.get_program() + result = FuseMulIntoLinear().call(ep) + self.assertFalse( + result.modified, + "Should not fold a constant that varies along a non-output axis", + ) diff --git a/backends/fused_quant/test/test_hoist_activations.py b/backends/fused_quant/test/test_hoist_activations.py new file mode 100644 index 00000000000..5fe8d11f01a --- /dev/null +++ b/backends/fused_quant/test/test_hoist_activations.py @@ -0,0 +1,206 @@ +# 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. + + +import unittest + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.optimization_passes.hoist_activations import ( + HoistActivations, +) +from executorch.backends.fused_quant.test.helpers import create_per_tensor_qparams +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir.dialects._ops import ops as exir_ops + + +def _build_relu_after_permutes_graph() -> torch.fx.GraphModule: + """Build a graph: inp -> permute -> permute -> relu -> output. + + Quant params are lifted front placeholders (matching production), so they + are graph inputs rather than inline body nodes fed to the relu. + """ + inp = torch.randn(1, 4, 8, 8).to(torch.int8) + + builder = ProgramBuilder() + node = builder.placeholder("inp", inp) + + # Quant params are lifted constants and thus front placeholders. + inp_qparams = create_per_tensor_qparams(builder) + out_qparams = create_per_tensor_qparams(builder) + + node = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(node, [0, 2, 3, 1]), + ) + node = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(node, [0, 3, 1, 2]), + ) + + relu = builder.call_operator( + op=exir_ops.edge.fused_quant.relu.default, + args=( + node, + *inp_qparams, + *out_qparams, + ), + ) + + builder.output([relu]) + graph_module = builder.get_program().graph_module + + return graph_module + + +class HoistActivationsTest(unittest.TestCase): + """Tests for the HoistActivations optimization pass.""" + + def test_relu_moves_above_permutes(self) -> None: + """Relu is hoisted above a chain of permute_copy ops. + + Before: inp -> permute -> permute -> relu -> output + After: inp -> relu -> permute -> permute -> output + + The relu never moves down; each single-user permute is moved forward to + sit after the relu. The relu's qparams are lifted front placeholders, + so they are trivially already before the relu and do not move. + """ + graph_module = _build_relu_after_permutes_graph() + + result = HoistActivations().call(graph_module) + + self.assertTrue(result.modified) + + graph = result.graph_module.graph + relu_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.relu.default + ) + permute_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + + self.assertEqual(len(relu_nodes), 1, "Expected 1 relu node") + self.assertEqual(len(permute_nodes), 2, "Expected 2 permute nodes") + + all_nodes = list(graph.nodes) + relu_idx = all_nodes.index(relu_nodes[0]) + + # Relu should come before both permutes + for permute_node in permute_nodes: + permute_idx = all_nodes.index(permute_node) + self.assertLess( + relu_idx, permute_idx, "Relu should come before permute nodes" + ) + + def test_convergence(self) -> None: + """First invocation modifies the graph; second invocation is a no-op.""" + graph_module = _build_relu_after_permutes_graph() + + opt_pass = HoistActivations() + + first_result = opt_pass.call(graph_module) + self.assertTrue(first_result.modified) + + second_result = opt_pass.call(first_result.graph_module) + self.assertFalse(second_result.modified) + + def test_no_hoist_when_data_movement_has_multiple_users(self) -> None: + """Should not hoist when the data-movement op feeding the activation + has more than one user. + + If permute feeds both relu and another consumer, hoisting would make + the other consumer see relu(permute(inp)) instead of permute(inp), + silently changing its semantics. + """ + inp = torch.randn(1, 4, 8, 8).to(torch.int8) + + builder = ProgramBuilder() + node = builder.placeholder("inp", inp) + inp_qparams = create_per_tensor_qparams(builder) + out_qparams = create_per_tensor_qparams(builder) + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(node, [0, 2, 3, 1]), + ) + # A second consumer of permute (besides relu below). + other = builder.call_operator( + op=exir_ops.edge.aten.abs.default, + args=(permute,), + ) + relu = builder.call_operator( + op=exir_ops.edge.fused_quant.relu.default, + args=(permute, *inp_qparams, *out_qparams), + ) + builder.output([relu, other]) + graph_module = builder.get_program().graph_module + + result = HoistActivations().call(graph_module) + + self.assertFalse( + result.modified, + "Should not hoist when the data-movement op has multiple users", + ) + + def test_requantizing_relu_keeps_output_dtype(self) -> None: + """A requantizing relu (int8 -> int16) keeps its output dtype on hoist. + + The relu's meta['val'] dtype is int16 (its output), not int8 (its + input). After hoisting, the relu must still report int16, and each + permute it now feeds must be retyped from int8 to int16 since it now + carries the relu's requantized output. + """ + inp = torch.randn(1, 4, 8, 8).to(torch.int8) + + builder = ProgramBuilder() + node = builder.placeholder("inp", inp) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.int8) + out_qparams = create_per_tensor_qparams(builder, dtype=torch.int16) + node = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(node, [0, 2, 3, 1]), + ) + node = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(node, [0, 3, 1, 2]), + ) + relu = builder.call_operator( + op=exir_ops.edge.fused_quant.relu.default, + args=(node, *inp_qparams, *out_qparams), + ) + builder.output([relu]) + graph_module = builder.get_program().graph_module + + relu_before = graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.relu.default + )[0] + self.assertEqual(relu_before.meta["val"].dtype, torch.int16) + + result = HoistActivations().call(graph_module) + self.assertTrue(result.modified) + + graph = result.graph_module.graph + relu_after = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.relu.default + )[0] + self.assertEqual( + relu_after.meta["val"].dtype, + torch.int16, + "Hoisted relu must keep its requantized output dtype (int16), not " + "the int8 dtype of its new input", + ) + + # The permutes now carry the relu's requantized output, so their + # meta['val'] dtype must follow (int16), not the stale int8 from before. + for permute_node in graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ): + self.assertEqual( + permute_node.meta["val"].dtype, + torch.int16, + "Permute carrying the hoisted relu's output must be retyped to " + "the relu's output dtype", + ) diff --git a/backends/fused_quant/test/test_op_numerics.py b/backends/fused_quant/test/test_op_numerics.py new file mode 100644 index 00000000000..fcd19de6218 --- /dev/null +++ b/backends/fused_quant/test/test_op_numerics.py @@ -0,0 +1,348 @@ +# 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. + +from __future__ import annotations + +import unittest +from collections.abc import Callable +from dataclasses import dataclass, field + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.decompose_fused_quant import DecomposeFusedQuant +from executorch.backends.fused_quant.graph_utils import split_fused_arg_names +from executorch.backends.fused_quant.optimization_passes.to_channels_first import ( + ToChannelsFirst, +) +from executorch.backends.fused_quant.test.helpers import create_per_tensor_qparams +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_manager import ExportedProgramPassManager +from parameterized import parameterized +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.utils import _pytree as pytree + +_QPARAM_FIELDS = ("scale", "zero_point", "dtype", "quant_min", "quant_max") +_INP_SCALE = 0.25 +_OUT_SCALE = 0.5 + + +def _int8(*shape: int) -> Callable[[], torch.Tensor]: + return lambda: torch.randint(-128, 128, shape, dtype=torch.int8) + + +def _float(*shape: int) -> Callable[[], torch.Tensor]: + return lambda: torch.randn(*shape) + + +@dataclass(frozen=True) +class _Case: + op: str + operands: dict[str, Callable[[], torch.Tensor]] + quantized: frozenset[str] + extras: dict[str, object] = field(default_factory=dict) + float_out: bool = True + + def overload(self) -> str: + return self.op.partition(".")[2] or "default" + + def packet_name(self) -> str: + return self.op.partition(".")[0] + + +_CONV_2D = { + "stride": [1, 1], + "padding": [0, 0], + "dilation": [1, 1], + "transposed": False, + "output_padding": [0, 0], + "groups": 1, +} +_POOL_2D = {"kernel_size": [2, 2], "stride": [2, 2], "padding": [0, 0]} + +_CASES: list[_Case] = [ + _Case("requantize", {"inp": _int8(1, 8)}, frozenset({"inp"}), float_out=False), + *( + _Case(name, {"inp": _int8(1, 8)}, frozenset({"inp"})) + for name in ( + "relu", + "hardswish", + "sigmoid", + "tanh", + "hard_tanh", + "silu", + "hardsigmoid", + ) + ), + _Case("gelu", {"inp": _int8(1, 8)}, frozenset({"inp"}), {"approximate": "none"}), + *( + _Case( + name, + {"inp": _int8(2, 4), "other": _int8(2, 4)}, + frozenset({"inp", "other"}), + ) + for name in ("add", "mul", "sub") + ), + _Case("add.Scalar", {"inp": _int8(2, 4)}, frozenset({"inp"}), {"other": 2.0}), + _Case("mul.Scalar", {"inp": _int8(2, 4)}, frozenset({"inp"}), {"other": 2.0}), + _Case("sub.Scalar", {"inp": _int8(2, 4)}, frozenset({"inp"}), {"other": 2.0}), + _Case( + "bmm", + {"inp": _int8(2, 3, 4), "other": _int8(2, 4, 5)}, + frozenset({"inp", "other"}), + ), + _Case( + "linear", + {"inp": _int8(2, 8), "weight": _int8(4, 8), "bias": _float(4)}, + frozenset({"inp", "weight"}), + ), + _Case( + "convolution", + {"inp": _int8(1, 2, 6, 6), "weight": _int8(3, 2, 3, 3), "bias": _float(3)}, + frozenset({"inp", "weight"}), + _CONV_2D, + ), + _Case( + "convolution_channels_last", + {"inp": _int8(1, 6, 6, 2), "weight": _int8(3, 3, 3, 2), "bias": _float(3)}, + frozenset({"inp", "weight"}), + _CONV_2D, + ), + _Case( + "conv1d", + {"inp": _int8(1, 2, 8), "weight": _int8(3, 2, 3), "bias": _float(3)}, + frozenset({"inp", "weight"}), + {"stride": [1], "padding": [0], "dilation": [1], "groups": 1}, + ), + _Case( + "conv2d", + {"inp": _int8(1, 2, 6, 6), "weight": _int8(3, 2, 3, 3), "bias": _float(3)}, + frozenset({"inp", "weight"}), + {"stride": [1, 1], "padding": [0, 0], "dilation": [1, 1], "groups": 1}, + ), + _Case( + "conv3d", + { + "inp": _int8(1, 2, 4, 4, 4), + "weight": _int8(3, 2, 3, 3, 3), + "bias": _float(3), + }, + frozenset({"inp", "weight"}), + {"stride": [1, 1, 1], "padding": [0, 0, 0], "dilation": [1, 1, 1], "groups": 1}, + ), + *( + _Case( + name, + {"inp": _int8(*shape)}, + frozenset({"inp"}), + {**_POOL_2D, "dilation": [1, 1], "ceil_mode": False}, + ) + for name, shape in ( + ("max_pool2d_with_indices", (1, 2, 6, 6)), + ("max_pool2d_with_indices_channels_last", (1, 6, 6, 2)), + ) + ), + *( + _Case( + name, + {"inp": _int8(*shape)}, + frozenset({"inp"}), + { + **_POOL_2D, + "ceil_mode": False, + "count_include_pad": True, + "divisor_override": None, + }, + ) + for name, shape in ( + ("avg_pool2d", (1, 2, 6, 6)), + ("avg_pool2d_channels_last", (1, 6, 6, 2)), + ) + ), + _Case( + "native_layer_norm", + {"inp": _int8(2, 8), "weight": _float(8), "bias": _float(8)}, + frozenset({"inp"}), + {"normalized_shape": [8], "eps": 1e-5}, + ), + _Case( + "rms_norm", + {"inp": _int8(2, 8), "weight": _float(8)}, + frozenset({"inp"}), + {"normalized_shape": [8], "eps": 1e-5}, + ), + _Case( + "_masked_softmax", + {"inp": _int8(2, 4), "mask": lambda: torch.rand(2, 4) > 0.5}, + frozenset({"inp"}), + {"dim": -1, "mask_type": 2}, + ), + _Case( + "embedding", + {"weight": _int8(10, 4), "indices": lambda: torch.tensor([0, 3, 9])}, + frozenset({"weight"}), + ), +] + +_VARIANTS: list[tuple[str, _Case, bool]] = [ + ( + f"{case.op}_{'float_out' if not quantize_out else 'quantized_out'}", + case, + quantize_out, + ) + for case in _CASES + for quantize_out in (True, False) + if quantize_out or case.float_out +] + + +def _eager_op(case: _Case) -> torch._ops.OpOverload: + return getattr(getattr(torch.ops.fused_quant, case.packet_name()), case.overload()) + + +def _edge_op(case: _Case) -> object: + return getattr( + getattr(exir_ops.edge.fused_quant, case.packet_name()), case.overload() + ) + + +def _qparam_values( + prefix: str, quantized: bool +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.dtype, int, int]: + if not quantized: + return (None, None, torch.float32, 0, 0) + is_out = prefix == "out" + return ( + torch.tensor(_OUT_SCALE if is_out else _INP_SCALE), + torch.tensor(0, dtype=torch.int64), + torch.int8 if is_out else torch.float32, + -128, + 127, + ) + + +def _build_args( + case: _Case, + quantize_out: bool, + operands: dict[str, object], + qparams: Callable[[str, bool], tuple[object, ...]], +) -> list[object]: + """Lay out ``case``'s arguments in schema order. + + Tensor operands come from ``operands``, each qparam block from ``qparams``, + and every remaining argument from ``case.extras`` or its schema default. + """ + op = _eager_op(case) + input_names, _ = split_fused_arg_names(op) + blocks = { + prefix: qparams( + prefix, quantize_out if prefix == "out" else prefix in case.quantized + ) + for prefix in (*input_names, "out") + } + extras = dict(case.extras) + args: list[object] = [] + for argument in op._schema.arguments: + name = argument.name + if name in operands: + args.append(operands[name]) + continue + block_field = next( + ( + (block, f) + for block in blocks + for f in _QPARAM_FIELDS + if name == f"{block}_{f}" + ), + None, + ) + if block_field is not None: + block, f = block_field + args.append(blocks[block][_QPARAM_FIELDS.index(f)]) + elif name in extras: + args.append(extras.pop(name)) + else: + assert argument.has_default_value(), f"{case.op}: no value for {name}" + args.append(argument.default_value) + assert not extras, f"{case.op}: unused extras {sorted(extras)}" + return args + + +def _first_output(result: object) -> torch.Tensor: + return result[0] if isinstance(result, (tuple, list)) else result + + +class FakeKernelDtypeTest(unittest.TestCase): + @parameterized.expand(_VARIANTS) + def test_fake_matches_eager( + self, _name: str, case: _Case, quantize_out: bool + ) -> None: + torch.manual_seed(0) + operands = {name: make() for name, make in case.operands.items()} + args = _build_args(case, quantize_out, operands, _qparam_values) + eager = _eager_op(case)(*args) + + with FakeTensorMode(allow_non_fake_inputs=True) as mode: + fake_args = pytree.tree_map_only(torch.Tensor, mode.from_tensor, args) + fake = _eager_op(case)(*fake_args) + + eager_flat, _ = pytree.tree_flatten(eager) + fake_flat, _ = pytree.tree_flatten(fake) + self.assertEqual(len(eager_flat), len(fake_flat)) + for e, f in zip(eager_flat, fake_flat): + self.assertEqual(f.dtype, e.dtype) + self.assertEqual(f.shape, e.shape) + + +class DecomposeNumericsTest(unittest.TestCase): + """The reference kernel and its DecomposeFusedQuant lowering agree exactly.""" + + @parameterized.expand(_VARIANTS) + def test_decomposed_matches_reference( + self, _name: str, case: _Case, quantize_out: bool + ) -> None: + torch.manual_seed(0) + inputs = {name: make() for name, make in case.operands.items()} + builder = ProgramBuilder() + nodes = { + name: builder.placeholder(name, value) for name, value in inputs.items() + } + + def lifted_qparams(prefix: str, quantized: bool) -> tuple[object, ...]: + if not quantized: + return _qparam_values(prefix, quantized) + scale = _OUT_SCALE if prefix == "out" else _INP_SCALE + dtype = torch.int8 if prefix == "out" else torch.float32 + return create_per_tensor_qparams(builder, scale=scale, dtype=dtype) + + args = _build_args(case, quantize_out, nodes, lifted_qparams) + result = builder.call_operator(_edge_op(case), tuple(args)) + if isinstance(result.data, (tuple, list)): + result = builder.call_getitem(result, 0) + builder.output([result]) + program = builder.get_program() + + user_inputs = [inputs[name] for name in case.operands] + reference = program.module()(*user_inputs) + # The documented fallback: channels-last ops go back to channels-first + # before decomposition, which only knows the ATen layouts. + decomposed = ExportedProgramPassManager( + [ToChannelsFirst(), DecomposeFusedQuant()] + )(program).exported_program + + self.assertFalse( + any( + getattr(node.target, "namespace", None) == "fused_quant" + for node in decomposed.graph.nodes + if node.op == "call_function" + ) + ) + torch.testing.assert_close( + _first_output(decomposed.module()(*user_inputs)), + _first_output(reference), + rtol=0, + atol=0, + ) diff --git a/backends/fused_quant/test/test_ops.py b/backends/fused_quant/test/test_ops.py new file mode 100644 index 00000000000..7ad29f9c4bf --- /dev/null +++ b/backends/fused_quant/test/test_ops.py @@ -0,0 +1,899 @@ +# 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. + + +import unittest +from contextlib import nullcontext + +import torch +from executorch.backends.fused_quant.ops import ( + _bmm_meta, + _embedding_meta, + _permute_contiguous, + _requantize_meta, + QuantParamsStruct, +) +from torch._subclasses.fake_tensor import FakeTensorMode + + +def _create_qparams( + dtype: torch.dtype = torch.int8, +) -> QuantParamsStruct[torch.Tensor]: + return QuantParamsStruct( + scale=torch.tensor([1.0]), + zero_point=torch.tensor([0], dtype=torch.int64), + dtype=dtype, + quant_min=torch.iinfo(dtype).min, + quant_max=torch.iinfo(dtype).max, + ) + + +def _create_per_channel_qparams( + num_channels: int, dtype: torch.dtype = torch.int8 +) -> QuantParamsStruct[torch.Tensor]: + # A 1D scale has a single non-unary dim, so it classifies as per-channel and + # its channel axis derives to 0. + return QuantParamsStruct( + scale=torch.ones(num_channels), + zero_point=torch.zeros(num_channels, dtype=torch.int64), + dtype=dtype, + quant_min=torch.iinfo(dtype).min, + quant_max=torch.iinfo(dtype).max, + ) + + +def _create_per_group_qparams( + scale: torch.Tensor, + zero_point: torch.Tensor | None = None, + dtype: torch.dtype = torch.int8, + quant_min: int = -8, + quant_max: int = 7, +) -> QuantParamsStruct[torch.Tensor]: + """Per-group qparams. ``scale`` is a full-rank [num_channels, num_groups] tensor. + + Defaults model the W4 case (4-bit range stored in an int8 container). + """ + if zero_point is None: + zero_point = torch.zeros_like(scale, dtype=torch.int64) + return QuantParamsStruct( + scale=scale, + zero_point=zero_point, + dtype=dtype, + quant_min=quant_min, + quant_max=quant_max, + ) + + +def _ref_quantize_per_group( + tensor: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + quant_min: int, + quant_max: int, + dtype: torch.dtype, +) -> torch.Tensor: + """Independent, loop-based per-group quantize reference (channel axis 0).""" + num_channels, grouped_dim = tensor.shape + num_groups = scale.shape[1] + group_size = grouped_dim // num_groups + out = torch.empty_like(tensor) + for k in range(num_channels): + for g in range(num_groups): + cols = slice(g * group_size, (g + 1) * group_size) + out[k, cols] = torch.clamp( + torch.round(tensor[k, cols] / scale[k, g]) + zero_point[k, g], + quant_min, + quant_max, + ) + return out.to(dtype) + + +def _ref_dequantize_per_group( + tensor: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, +) -> torch.Tensor: + """Independent, loop-based per-group dequantize reference (channel axis 0).""" + num_channels, grouped_dim = tensor.shape + num_groups = scale.shape[1] + group_size = grouped_dim // num_groups + out = torch.empty(tensor.shape, dtype=torch.float32) + for k in range(num_channels): + for g in range(num_groups): + cols = slice(g * group_size, (g + 1) * group_size) + out[k, cols] = ( + tensor[k, cols].to(torch.float32) - zero_point[k, g] + ) * scale[k, g] + return out + + +def _flat_qparams( + qp: QuantParamsStruct[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.dtype, int, int]: + return (qp.scale, qp.zero_point, qp.dtype, qp.quant_min, qp.quant_max) + + +class TestQuantParamsFunctions(unittest.TestCase): + def test_is_per_tensor(self) -> None: + qparams = _create_qparams() + self.assertTrue(qparams.is_per_tensor()) + self.assertFalse(qparams.is_per_channel()) + self.assertFalse(qparams.is_per_group()) + + def test_is_per_channel(self) -> None: + qparams = _create_per_channel_qparams(16) + self.assertFalse(qparams.is_per_tensor()) + self.assertTrue(qparams.is_per_channel()) + self.assertFalse(qparams.is_per_group()) + + def test_channel_axis_derivation(self) -> None: + # all-ones -> 0; single non-unary dim -> that dim; multiple -> None. + per_tensor = QuantParamsStruct( + torch.ones(1, 1), torch.zeros(1, 1, dtype=torch.int64), torch.int8, -8, 7 + ) + self.assertEqual(per_tensor.channel_axis(), 0) + per_channel_axis1 = QuantParamsStruct( + torch.ones(1, 4), torch.zeros(1, 4, dtype=torch.int64), torch.int8, -8, 7 + ) + self.assertEqual(per_channel_axis1.channel_axis(), 1) + blockwise = QuantParamsStruct( + torch.ones(4, 2), torch.zeros(4, 2, dtype=torch.int64), torch.int8, -8, 7 + ) + self.assertIsNone(blockwise.channel_axis()) + + def test_quant_min_max_in_qparams(self) -> None: + qparams_int8 = _create_qparams(torch.int8) + self.assertEqual(qparams_int8.quant_min, -128) + self.assertEqual(qparams_int8.quant_max, 127) + + qparams_uint8 = _create_qparams(torch.uint8) + self.assertEqual(qparams_uint8.quant_min, 0) + self.assertEqual(qparams_uint8.quant_max, 255) + + def test_validate_succeeds(self) -> None: + _create_qparams().validate() + _create_per_channel_qparams(16).validate() + _create_per_group_qparams(torch.ones(4, 4)).validate() + + def test_validate_shape_mismatch_raises(self) -> None: + qparams = QuantParamsStruct( + scale=torch.ones(4, 1), + zero_point=torch.zeros(4, 2, dtype=torch.int64), + dtype=torch.int8, + quant_min=-8, + quant_max=7, + ) + with self.assertRaisesRegex(ValueError, "same shape"): + qparams.validate() + + +class TestAffineQuant(unittest.TestCase): + def test_per_tensor_round_trip(self) -> None: + torch.manual_seed(0) + scale = torch.tensor([0.5]) + zero_point = torch.tensor([0], dtype=torch.int64) + quant_qp = QuantParamsStruct(scale, zero_point, torch.int8, -8, 7) + dequant_qp = QuantParamsStruct(scale, zero_point, torch.float32, -8, 7) + q = torch.randint(-8, 8, (3, 5), dtype=torch.int8) + weight = dequant_qp.dequantize(q) + self.assertEqual(weight.dtype, torch.float32) + self.assertTrue(torch.equal(quant_qp.quantize(weight), q)) + + def test_per_channel_round_trip(self) -> None: + torch.manual_seed(1) + num_channels, in_features = 4, 6 + scale = torch.rand(num_channels, 1) + 0.1 # full-rank [M, 1] -> per-channel + zero_point = torch.zeros(num_channels, 1, dtype=torch.int64) + quant_qp = QuantParamsStruct(scale, zero_point, torch.int8, -8, 7) + dequant_qp = QuantParamsStruct(scale, zero_point, torch.float32, -8, 7) + q = torch.randint(-8, 8, (num_channels, in_features), dtype=torch.int8) + weight = dequant_qp.dequantize(q) + self.assertTrue(torch.equal(quant_qp.quantize(weight), q)) + + +class TestRequantize(unittest.TestCase): + def test_requantize(self) -> None: + inp_qparams = QuantParamsStruct( + torch.tensor([0.5]), + torch.tensor([1], dtype=torch.int64), + torch.float32, + -128, + 127, + ) + out_qparams = QuantParamsStruct( + torch.tensor([0.25]), + torch.tensor([3], dtype=torch.int64), + torch.uint8, + 0, + 255, + ) + inp = torch.tensor([[-4, 0, 7]], dtype=torch.int8) + + expected = out_qparams.quantize(inp_qparams.dequantize(inp)) + actual = torch.ops.fused_quant.requantize( + inp, *_flat_qparams(inp_qparams), *_flat_qparams(out_qparams) + ) + meta = _requantize_meta( + inp, *_flat_qparams(inp_qparams), *_flat_qparams(out_qparams) + ) + + self.assertEqual(actual.dtype, torch.uint8) + self.assertEqual(meta.dtype, torch.uint8) + self.assertTrue(torch.equal(actual, expected)) + + +class TestConvNd(unittest.TestCase): + def test_rejects_wrong_spatial_dimensions(self) -> None: + no_qparams = (None, None, torch.float32, 0, 0) + for name, target, spatial_dims in ( + ("conv1d", torch.ops.fused_quant.conv1d.default, 1), + ("conv2d", torch.ops.fused_quant.conv2d.default, 2), + ("conv3d", torch.ops.fused_quant.conv3d.default, 3), + ): + wrong_spatial_dims = 2 if spatial_dims == 1 else 1 + for fake in (False, True): + with ( + self.subTest(name=name, fake=fake), + FakeTensorMode() if fake else nullcontext(), + ): + inp = torch.randn(1, 2, *([8] * wrong_spatial_dims)) + weight = torch.randn(3, 2, *([3] * wrong_spatial_dims)) + with self.assertRaisesRegex( + ValueError, + f"fused_quant::{name} expects input and weight rank " + f"{spatial_dims + 2}", + ): + target( + inp, + weight, + None, + *no_qparams, + *no_qparams, + *no_qparams, + *no_qparams, + [1] * wrong_spatial_dims, + [0] * wrong_spatial_dims, + [1] * wrong_spatial_dims, + 1, + ) + + def test_rejects_wrong_spatial_argument_lengths(self) -> None: + no_qparams = (None, None, torch.float32, 0, 0) + inp = torch.randn(1, 2, 8, 8) + weight = torch.randn(3, 2, 3, 3) + for arg_name, stride, padding, dilation in ( + ("stride", [1], [0, 0], [1, 1]), + ("padding", [1, 1], [0], [1, 1]), + ("dilation", [1, 1], [0, 0], [1]), + ): + with ( + self.subTest(arg_name=arg_name), + self.assertRaisesRegex( + ValueError, + f"fused_quant::conv2d expects {arg_name} to contain 2 values", + ), + ): + torch.ops.fused_quant.conv2d.default( + inp, + weight, + None, + *no_qparams, + *no_qparams, + *no_qparams, + *no_qparams, + stride, + padding, + dilation, + 1, + ) + + +class TestConvolutionChannelsLast(unittest.TestCase): + def test_qparams_use_channels_last_layout(self) -> None: + torch.manual_seed(0) + inp = torch.randint(-8, 8, (1, 4, 4, 2), dtype=torch.int8) + weight = torch.randint(-8, 8, (2, 3, 3, 2), dtype=torch.int8) + inp_qparams = QuantParamsStruct( + torch.tensor([[[[0.25, 0.5]]]]), + torch.zeros(1, 1, 1, 2, dtype=torch.int64), + torch.float32, + -8, + 7, + ) + weight_qparams = QuantParamsStruct( + torch.tensor( + [ + [[[0.25, 0.5]]], + [[[0.75, 1.0]]], + ] + ), + torch.zeros(2, 1, 1, 2, dtype=torch.int64), + torch.float32, + -8, + 7, + ) + out_qparams = QuantParamsStruct( + torch.tensor([[[[0.5, 0.25]]]]), + torch.zeros(1, 1, 1, 2, dtype=torch.int64), + torch.int8, + -128, + 127, + ) + + dq_inp = inp_qparams.dequantize(inp).permute(0, 3, 1, 2) + dq_weight = weight_qparams.dequantize(weight).permute(0, 3, 1, 2) + expected = out_qparams.quantize( + torch.nn.functional.conv2d(dq_inp, dq_weight).permute(0, 2, 3, 1) + ) + actual = torch.ops.fused_quant.convolution_channels_last.default( + inp, + weight, + None, + *_flat_qparams(inp_qparams), + *_flat_qparams(weight_qparams), + None, + None, + torch.uint8, + 0, + 0, + *_flat_qparams(out_qparams), + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + 1, + ) + self.assertTrue(torch.equal(actual, expected)) + + def test_transposed_convolution(self) -> None: + inp = torch.randn(1, 4, 4, 2) + weight = torch.randn(2, 3, 3, 3) + no_qparams = (None, None, torch.float32, 0, 0) + expected = torch.ops.aten.convolution.default( + _permute_contiguous(inp, [0, 3, 1, 2]), + _permute_contiguous(weight, [0, 3, 1, 2]), + None, + [2, 2], + [1, 1], + [1, 1], + True, + [1, 1], + 1, + ).permute(0, 2, 3, 1) + + actual = torch.ops.fused_quant.convolution_channels_last.default( + inp, + weight, + None, + *no_qparams, + *no_qparams, + *no_qparams, + *no_qparams, + [2, 2], + [1, 1], + [1, 1], + True, + [1, 1], + 1, + ) + + torch.testing.assert_close(expected, actual, rtol=0, atol=0) + + +class TestPermuteContiguous(unittest.TestCase): + def test_canonicalizes_singleton_dimension_strides(self) -> None: + tensor = torch.empty((16, 1, 1, 64)) + permuted = tensor.permute(0, 3, 1, 2) + + self.assertTrue(permuted.is_contiguous()) + self.assertEqual(permuted.stride(), (64, 1, 64, 64)) + self.assertEqual( + _permute_contiguous(tensor, [0, 3, 1, 2]).stride(), + (64, 1, 1, 1), + ) + + +class PoolingMetaTest(unittest.TestCase): + def test_integer_inputs(self) -> None: + inp_qparams = QuantParamsStruct( + torch.tensor([1.0]), + torch.tensor([0], dtype=torch.int64), + torch.float32, + -128, + 127, + ) + out_qparams = _create_qparams(torch.uint8) + + with FakeTensorMode(allow_non_fake_inputs=True) as mode: + inp_nchw = mode.from_tensor(torch.empty((1, 2, 4, 4), dtype=torch.int8)) + max_values, max_indices = ( + torch.ops.fused_quant.max_pool2d_with_indices.default( + inp_nchw, + *_flat_qparams(inp_qparams), + *_flat_qparams(out_qparams), + [2, 2], + [2, 2], + [0, 0], + [1, 1], + False, + ) + ) + avg_nchw = torch.ops.fused_quant.avg_pool2d.default( + inp_nchw, + *_flat_qparams(inp_qparams), + *_flat_qparams(out_qparams), + [2, 2], + [2, 2], + [0, 0], + False, + True, + None, + ) + + inp_nhwc = mode.from_tensor(torch.empty((1, 4, 4, 2), dtype=torch.int8)) + avg_nhwc = torch.ops.fused_quant.avg_pool2d_channels_last.default( + inp_nhwc, + *_flat_qparams(inp_qparams), + *_flat_qparams(out_qparams), + [2, 2], + [2, 2], + [0, 0], + False, + True, + None, + ) + + self.assertEqual(max_values.shape, torch.Size([1, 2, 2, 2])) + self.assertEqual(max_values.dtype, torch.uint8) + self.assertEqual(max_indices.shape, torch.Size([1, 2, 2, 2])) + self.assertEqual(max_indices.dtype, torch.int64) + self.assertEqual(avg_nchw.shape, torch.Size([1, 2, 2, 2])) + self.assertEqual(avg_nchw.dtype, torch.uint8) + self.assertEqual(avg_nhwc.shape, torch.Size([1, 2, 2, 2])) + self.assertEqual(avg_nhwc.dtype, torch.uint8) + + +class TestMaxPool2dWithIndicesChannelsLast(unittest.TestCase): + def test_qparams_use_channels_last_layout(self) -> None: + inp = torch.randint(-8, 8, (1, 4, 4, 2), dtype=torch.int8) + inp_qparams = QuantParamsStruct( + torch.tensor([[[[0.25, 0.5]]]]), + torch.zeros(1, 1, 1, 2, dtype=torch.int64), + torch.float32, + -8, + 7, + ) + out_qparams = QuantParamsStruct( + torch.tensor([[[[0.5, 0.25]]]]), + torch.zeros(1, 1, 1, 2, dtype=torch.int64), + torch.int8, + -128, + 127, + ) + + values_nchw, indices_nchw = torch.ops.aten.max_pool2d_with_indices.default( + inp_qparams.dequantize(inp).permute(0, 3, 1, 2), + [2, 2], + [2, 2], + ) + expected_values = out_qparams.quantize(values_nchw.permute(0, 2, 3, 1)) + expected_indices = indices_nchw.permute(0, 2, 3, 1) + actual_values, actual_indices = ( + torch.ops.fused_quant.max_pool2d_with_indices_channels_last.default( + inp, + *_flat_qparams(inp_qparams), + *_flat_qparams(out_qparams), + [2, 2], + [2, 2], + [0, 0], + [1, 1], + False, + ) + ) + self.assertTrue(torch.equal(actual_values, expected_values)) + self.assertTrue(torch.equal(actual_indices, expected_indices)) + + +class TestAvgPool2dChannelsLast(unittest.TestCase): + def test_output_has_canonical_strides_with_singleton_channel(self) -> None: + inp = torch.arange(16, dtype=torch.float32).reshape(1, 4, 4, 1) + no_qparams = (None, None, torch.float32, -2147483648, 2147483647) + + actual = torch.ops.fused_quant.avg_pool2d_channels_last.default( + inp, + *no_qparams, + *no_qparams, + [2, 2], + [2, 2], + [0, 0], + False, + False, + None, + ) + expected = torch.ops.aten.avg_pool2d.default( + inp.permute(0, 3, 1, 2), + [2, 2], + [2, 2], + ).permute(0, 2, 3, 1) + + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(actual.stride(), (4, 2, 1, 1)) + + def test_qparams_and_aten_arguments(self) -> None: + inp = torch.randint(-8, 8, (1, 4, 4, 2), dtype=torch.int8) + inp_qparams = QuantParamsStruct( + torch.tensor([[[[0.25, 0.5]]]]), + torch.zeros(1, 1, 1, 2, dtype=torch.int64), + torch.float32, + -8, + 7, + ) + out_qparams = QuantParamsStruct( + torch.tensor([[[[0.5, 0.25]]]]), + torch.zeros(1, 1, 1, 2, dtype=torch.int64), + torch.int8, + -128, + 127, + ) + + for ceil_mode, count_include_pad, divisor_override in ( + (False, False, None), + (True, True, 5), + ): + with self.subTest( + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + divisor_override=divisor_override, + ): + output_nchw = torch.ops.aten.avg_pool2d.default( + inp_qparams.dequantize(inp).permute(0, 3, 1, 2), + [2, 2], + [2, 2], + [1, 1], + ceil_mode, + count_include_pad, + divisor_override, + ) + expected = out_qparams.quantize(output_nchw.permute(0, 2, 3, 1)) + actual = torch.ops.fused_quant.avg_pool2d_channels_last.default( + inp, + *_flat_qparams(inp_qparams), + *_flat_qparams(out_qparams), + [2, 2], + [2, 2], + [1, 1], + ceil_mode, + count_include_pad, + divisor_override, + ) + self.assertTrue(torch.equal(actual, expected)) + + +class TestPerGroupQuant(unittest.TestCase): + def test_is_per_group_classification(self) -> None: + qparams = _create_per_group_qparams(torch.ones(4, 4)) + self.assertTrue(qparams.is_per_group()) + self.assertFalse(qparams.is_per_tensor()) + self.assertFalse(qparams.is_per_channel()) + + def test_quantize_matches_loop_reference(self) -> None: + torch.manual_seed(0) + weight = torch.randn(4, 16) + scale = torch.rand(4, 4) + 0.1 # strictly positive, num_groups=4 + qparams = _create_per_group_qparams(scale) + expected = _ref_quantize_per_group( + weight, scale, qparams.zero_point, -8, 7, torch.int8 + ) + actual = qparams.quantize(weight) + self.assertEqual(actual.dtype, torch.int8) + self.assertTrue(torch.equal(actual, expected)) + + def test_dequantize_matches_loop_reference(self) -> None: + torch.manual_seed(1) + scale = torch.rand(4, 4) + 0.1 + qweight = torch.randint(-8, 8, (4, 16), dtype=torch.int8) + # A dequant qparams carries the float out_dtype. + qparams = _create_per_group_qparams(scale, dtype=torch.float32) + expected = _ref_dequantize_per_group(qweight, scale, qparams.zero_point) + actual = qparams.dequantize(qweight) + self.assertEqual(actual.dtype, torch.float32) + self.assertTrue(torch.allclose(actual, expected)) + + def test_different_groups_use_different_scales(self) -> None: + # Two groups along the contraction axis with very different scales. + scale = torch.tensor([[1.0, 10.0]]) # 1 channel, 2 groups + qparams = _create_per_group_qparams(scale) + weight = torch.tensor([[4.0, 4.0, 40.0, 40.0]]) # group0 / 1, group1 / 10 + actual = qparams.quantize(weight) + # group0: round(4/1)=4 ; group1: round(40/10)=4 -> all 4s despite 10x range + self.assertTrue( + torch.equal(actual, torch.tensor([[4, 4, 4, 4]], dtype=torch.int8)) + ) + + def test_round_trip_on_grid_is_exact(self) -> None: + torch.manual_seed(2) + scale = torch.rand(4, 4) + 0.1 + zero_point = torch.zeros(4, 4, dtype=torch.int64) + quant_qp = _create_per_group_qparams( + scale, zero_point=zero_point, dtype=torch.int8 + ) + dequant_qp = _create_per_group_qparams( + scale, zero_point=zero_point, dtype=torch.float32 + ) + q = torch.randint(-8, 8, (4, 16), dtype=torch.int8) + # Build a float weight that lies exactly on the per-group grid, then + # re-quantizing must recover the original integer codes. + weight = dequant_qp.dequantize(q) + requantized = quant_qp.quantize(weight) + self.assertTrue(torch.equal(requantized, q)) + + def test_quantize_clamps_out_of_range(self) -> None: + scale = torch.tensor([[1.0, 1.0]]) + qparams = _create_per_group_qparams(scale) + weight = torch.tensor([[100.0, -100.0, 3.0, -3.0]]) + actual = qparams.quantize(weight) + self.assertTrue( + torch.equal(actual, torch.tensor([[7, -8, 3, -3]], dtype=torch.int8)) + ) + + def test_indivisible_grouped_dim_raises(self) -> None: + scale = torch.ones(4, 3) # 3 groups + qparams = _create_per_group_qparams(scale) + weight = torch.randn(4, 16) # 16 not divisible by 3 + with self.assertRaisesRegex(ValueError, "must be divisible by"): + qparams.quantize(weight) + + def test_rank_mismatch_raises(self) -> None: + # A non-singleton scale must be full-rank (match the tensor rank). + scale = torch.ones(4, 4) + qparams = _create_per_group_qparams(scale) + weight = torch.randn(4, 4, 4) + with self.assertRaisesRegex(ValueError, "match the tensor rank"): + qparams.quantize(weight) + + def test_nonzero_zero_point(self) -> None: + scale = torch.tensor([[2.0, 2.0]]) + zero_point = torch.tensor([[1, -1]], dtype=torch.int64) + quant_qp = _create_per_group_qparams( + scale, zero_point=zero_point, dtype=torch.int8 + ) + weight = torch.tensor([[4.0, 4.0, 4.0, 4.0]]) + # group0: round(4/2)+1=3 ; group1: round(4/2)-1=1 + actual = quant_qp.quantize(weight) + self.assertTrue( + torch.equal(actual, torch.tensor([[3, 3, 1, 1]], dtype=torch.int8)) + ) + # Dequantize inverts: group0 (3-1)*2=4 ; group1 (1+1)*2=4 + dequant_qp = _create_per_group_qparams( + scale, zero_point=zero_point, dtype=torch.float32 + ) + dq = dequant_qp.dequantize(actual) + self.assertTrue(torch.allclose(dq, torch.tensor([[4.0, 4.0, 4.0, 4.0]]))) + + +def _embedding_weight_qp( + scale: torch.Tensor, + zero_point: torch.Tensor | None = None, + quant_min: int = -8, + quant_max: int = 7, +) -> QuantParamsStruct[torch.Tensor]: + """The table (dequantize) qparams for fused_quant.embedding. + + dtype is float32 because the qparams describes dequantizing the quantized + table to float embeddings. Granularity is encoded by the scale shape relative + to the [num_embeddings, embedding_dim] table: a [num_embeddings, 1] scale is + per-row (per-channel), and a [num_embeddings, num_groups] scale is per-group + over the embedding dimension. + """ + if zero_point is None: + zero_point = torch.zeros_like(scale, dtype=torch.int64) + return QuantParamsStruct( + scale=scale, + zero_point=zero_point, + dtype=torch.float32, + quant_min=quant_min, + quant_max=quant_max, + ) + + +# A null output-qparams block: indices.shape + [emb_dim] embeddings stay float. +_NULL_OUT_QP_ARGS: tuple[None, None, torch.dtype, int, int] = ( + None, + None, + torch.float32, + 0, + 0, +) + + +def _call_embedding( + qtable: torch.Tensor, + indices: torch.Tensor, + weight_qp: QuantParamsStruct[torch.Tensor], + out_qp: QuantParamsStruct[torch.Tensor] | None = None, +) -> torch.Tensor: + """Invoke fused_quant.embedding: (table, weight_qp block, out_qp block, indices).""" + out_flat = _flat_qparams(out_qp) if out_qp is not None else _NULL_OUT_QP_ARGS + return torch.ops.fused_quant.embedding( + qtable, *_flat_qparams(weight_qp), *out_flat, indices + ) + + +class TestEmbedding(unittest.TestCase): + def test_embedding_per_channel(self) -> None: + torch.manual_seed(0) + num_embeddings, embedding_dim = 6, 8 + qtable = torch.randint(-8, 8, (num_embeddings, embedding_dim), dtype=torch.int8) + scale = torch.rand(num_embeddings, 1) + 0.1 # per-row + weight_qp = _embedding_weight_qp(scale) + indices = torch.tensor([0, 3, 5, 1]) + expected = weight_qp.dequantize(qtable)[indices] + actual = _call_embedding(qtable, indices, weight_qp) + self.assertEqual(tuple(actual.shape), (4, embedding_dim)) + self.assertEqual(actual.dtype, torch.float32) + self.assertTrue(torch.allclose(actual, expected)) + + def test_embedding_per_group(self) -> None: + torch.manual_seed(1) + num_embeddings, embedding_dim, num_groups = 6, 8, 4 + qtable = torch.randint(-8, 8, (num_embeddings, embedding_dim), dtype=torch.int8) + scale = torch.rand(num_embeddings, num_groups) + 0.1 # per-group over emb dim + weight_qp = _embedding_weight_qp(scale) + indices = torch.tensor([[0, 1], [5, 2]]) # 2D indices + expected = weight_qp.dequantize(qtable)[indices] + actual = _call_embedding(qtable, indices, weight_qp) + self.assertEqual(tuple(actual.shape), (2, 2, embedding_dim)) + self.assertTrue(torch.allclose(actual, expected)) + + def test_embedding_quantized_output(self) -> None: + # Exercises the optional output qparams block: the gathered embeddings are + # requantized to int8 (per-tensor). + torch.manual_seed(2) + num_embeddings, embedding_dim = 6, 8 + qtable = torch.randint(-8, 8, (num_embeddings, embedding_dim), dtype=torch.int8) + scale = torch.rand(num_embeddings, 1) + 0.1 + weight_qp = _embedding_weight_qp(scale) + out_qp = QuantParamsStruct( + scale=torch.tensor([0.05]), + zero_point=torch.tensor([0], dtype=torch.int64), + dtype=torch.int8, + quant_min=-128, + quant_max=127, + ) + indices = torch.tensor([0, 2, 4]) + expected = out_qp.quantize(weight_qp.dequantize(qtable)[indices]) + actual = _call_embedding(qtable, indices, weight_qp, out_qp) + self.assertEqual(actual.dtype, torch.int8) + self.assertTrue(torch.equal(actual, expected)) + + def test_embedding_requires_quantized_table(self) -> None: + qtable = torch.randint(-8, 8, (6, 8), dtype=torch.int8) + indices = torch.tensor([0, 1]) + with self.assertRaisesRegex(ValueError, "At least one"): + torch.ops.fused_quant.embedding( + qtable, *_NULL_OUT_QP_ARGS, *_NULL_OUT_QP_ARGS, indices + ) + + def test_embedding_non_2d_table_raises(self) -> None: + qtable = torch.randint(-8, 8, (6, 8, 2), dtype=torch.int8) + scale = torch.rand(6) + 0.1 + weight_qp = _embedding_weight_qp(scale) + indices = torch.tensor([0, 1]) + with self.assertRaisesRegex(ValueError, "must be 2D"): + _call_embedding(qtable, indices, weight_qp) + + def test_embedding_padding_idx_not_supported(self) -> None: + qtable = torch.randint(-8, 8, (6, 8), dtype=torch.int8) + weight_qp = _embedding_weight_qp(torch.rand(6, 1) + 0.1) + indices = torch.tensor([0, 1]) + with self.assertRaises(AssertionError): + torch.ops.fused_quant.embedding( + qtable, + *_flat_qparams(weight_qp), + *_NULL_OUT_QP_ARGS, + indices, + 0, # padding_idx (not the aten default -1) + False, # scale_grad_by_freq + False, # sparse + ) + + def test_embedding_scale_grad_by_freq_not_supported(self) -> None: + qtable = torch.randint(-8, 8, (6, 8), dtype=torch.int8) + weight_qp = _embedding_weight_qp(torch.rand(6, 1) + 0.1) + indices = torch.tensor([0, 1]) + with self.assertRaises(AssertionError): + torch.ops.fused_quant.embedding( + qtable, + *_flat_qparams(weight_qp), + *_NULL_OUT_QP_ARGS, + indices, + -1, # padding_idx + True, # scale_grad_by_freq + False, # sparse + ) + + def test_embedding_sparse_not_supported(self) -> None: + qtable = torch.randint(-8, 8, (6, 8), dtype=torch.int8) + weight_qp = _embedding_weight_qp(torch.rand(6, 1) + 0.1) + indices = torch.tensor([0, 1]) + with self.assertRaises(AssertionError): + torch.ops.fused_quant.embedding( + qtable, + *_flat_qparams(weight_qp), + *_NULL_OUT_QP_ARGS, + indices, + -1, # padding_idx + False, # scale_grad_by_freq + True, # sparse + ) + + def test_embedding_meta_shape(self) -> None: + scale = torch.rand(6, 1) + 0.1 + weight_qp = _embedding_weight_qp(scale) + qtable = torch.empty(6, 8, dtype=torch.int8) + indices = torch.zeros(3, 4, dtype=torch.int64) + out = _embedding_meta( + qtable, *_flat_qparams(weight_qp), *_NULL_OUT_QP_ARGS, indices + ) + self.assertEqual(tuple(out.shape), (3, 4, 8)) + self.assertEqual(out.dtype, torch.float32) + + def test_embedding_meta_dtype_follows_table_dequant_dtype(self) -> None: + # When the output is not quantized, the dtype is the table's dequantize + # output dtype (here bfloat16), not assumed float32. + weight_qp = QuantParamsStruct( + scale=torch.rand(6, 1) + 0.1, + zero_point=torch.zeros(6, 1, dtype=torch.int64), + dtype=torch.bfloat16, + quant_min=-8, + quant_max=7, + ) + qtable = torch.empty(6, 8, dtype=torch.int8) + indices = torch.zeros(3, dtype=torch.int64) + out = _embedding_meta( + qtable, *_flat_qparams(weight_qp), *_NULL_OUT_QP_ARGS, indices + ) + self.assertEqual(out.dtype, torch.bfloat16) + + +class TestOps(unittest.TestCase): + def test_bmm_incorrect_dimension(self) -> None: + qp = _flat_qparams(_create_qparams()) + with self.assertRaisesRegex(ValueError, "Input tensors must be 3D"): + _bmm_meta( + torch.randn(2, 3), + torch.randn(2, 3, 5), + *qp, + *qp, + *qp, + ) + + def test_bmm_incorrect_batch_dimension(self) -> None: + qp = _flat_qparams(_create_qparams()) + with self.assertRaisesRegex( + ValueError, "Input tensors must have the same batch dimension" + ): + _bmm_meta( + torch.randn(2, 3, 5), + torch.randn(3, 5, 3), + *qp, + *qp, + *qp, + ) + + def test_bmm_incorrect_inner_dimension(self) -> None: + qp = _flat_qparams(_create_qparams()) + with self.assertRaisesRegex( + ValueError, "Input tensors must have the same inner dimension" + ): + _bmm_meta( + torch.randn(2, 3, 5), + torch.randn(2, 4, 5), + *qp, + *qp, + *qp, + ) diff --git a/backends/fused_quant/test/test_pass_base.py b/backends/fused_quant/test/test_pass_base.py new file mode 100644 index 00000000000..6ca5ea57f2d --- /dev/null +++ b/backends/fused_quant/test/test_pass_base.py @@ -0,0 +1,139 @@ +# 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. + + +import unittest +from typing import cast + +import torch +from executorch.backends.fused_quant.pass_base import IterativePassGroup +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from parameterized import parameterized +from torch.export import ExportedProgram + + +class _NoopEPPass(ExportedProgramPassBase): + """EP pass that marks modified without changing the graph.""" + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + return ExportedProgramPassResult(exported_program, True) + + +class _FuseFirstMulPair(ExportedProgramPassBase): + """Fuses the first pair of consecutive muls it finds into one. + + Looks for mul(mul(a, b), c) and replaces it with mul(a, b * c), + folding the constant operand. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + graph = exported_program.graph_module.graph + for node in graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.mul.Tensor + ): + inp = node.args[0] + if not isinstance(inp, torch.fx.Node): + continue + if inp.target != exir_ops.edge.aten.mul.Tensor: + continue + if len(inp.users) != 1: + continue + node.replace_input_with(inp, inp.args[0]) + node.args = ( + node.args[0], + cast(float, inp.args[1]) * cast(float, node.args[1]), + ) + graph.erase_node(inp) + exported_program.graph_module.recompile() + return ExportedProgramPassResult(exported_program, True) + return ExportedProgramPassResult(exported_program, False) + + +def _build_mul_chain(n: int) -> ExportedProgram: + """Build a graph: x * 2 * 3 * 4 * ... * (n+1).""" + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(2, 4)) + cur = x + for i in range(n): + cur = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, + args=(cur, float(i + 2)), + ) + builder.output([cur]) + return builder.get_program() + + +class IterativePassGroupTest(unittest.TestCase): + @parameterized.expand([(1,), (2,), (3,)]) + def test_steps_reduce_mul_count(self, steps: int) -> None: + initial_muls = 4 + ep = _build_mul_chain(initial_muls) + + mul_count_before = len( + ep.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.mul.Tensor + ) + ) + self.assertEqual(mul_count_before, initial_muls) + + group = IterativePassGroup([_FuseFirstMulPair()], steps=steps) + result = group.call(ep) + self.assertTrue(result.modified) + + mul_count_after = len( + result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.mul.Tensor + ) + ) + self.assertEqual( + mul_count_before - mul_count_after, + steps, + f"Expected {steps} muls to be fused, but went from {mul_count_before} to {mul_count_after}", + ) + + def test_nested_error_message(self) -> None: + """A failing pass inside a nested IterativePassGroup should produce + a chained exception that names the exact pass that failed.""" + + class _FailingPass(ExportedProgramPassBase): + def call( + self, exported_program: ExportedProgram + ) -> ExportedProgramPassResult: + raise RuntimeError("intentional failure") + + ep = _build_mul_chain(2) + + inner_group = IterativePassGroup([_NoopEPPass(), _FailingPass()], steps=1) + outer_group = IterativePassGroup([_NoopEPPass(), inner_group], steps=1) + + with self.assertRaises(Exception) as ctx: + outer_group.call(ep) + + # Walk the exception chain to collect all messages + messages: list[str] = [] + exc: BaseException | None = ctx.exception + while exc is not None: + messages.append(str(exc)) + exc = exc.__cause__ + + full_trace = "\n".join(messages) + self.assertIn( + "_FailingPass", + full_trace, + f"Exception chain should name _FailingPass:\n{full_trace}", + ) + self.assertIn( + "IterativePassGroup", + full_trace, + f"Exception chain should name IterativePassGroup:\n{full_trace}", + ) + self.assertIn( + "intentional failure", + full_trace, + f"Exception chain should include the original error:\n{full_trace}", + ) diff --git a/backends/fused_quant/test/test_prequantize_embedding.py b/backends/fused_quant/test/test_prequantize_embedding.py new file mode 100644 index 00000000000..9e3d4a0775b --- /dev/null +++ b/backends/fused_quant/test/test_prequantize_embedding.py @@ -0,0 +1,210 @@ +# 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. + + +import unittest + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.optimization_passes.prequantize_embedding import ( + PrequantizeEmbedding, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +class PrequantizeEmbeddingTest(unittest.TestCase): + """Tests for the PrequantizeEmbedding pass.""" + + _NUM_EMBEDDINGS = 6 + _EMBEDDING_DIM = 8 + _NUM_GROUPS = 2 + + def _build_quantized_embedding_program( + self, out_scale: float = 0.05, out_zp: int = 3 + ) -> ExportedProgram: + """fused_quant.embedding with a group-quantized table and a per-tensor + quantized output -- the shape PrequantizeEmbedding folds into a plain int8 + gather. The gather feeds a nop view (not the graph output directly) so the + rewrite replaces an interior node and the output signature stays stable. + """ + torch.manual_seed(0) + builder = ProgramBuilder() + num_embeddings = self._NUM_EMBEDDINGS + embedding_dim = self._EMBEDDING_DIM + num_groups = self._NUM_GROUPS + + table = builder.placeholder( + "table", + torch.randint(-8, 7, (num_embeddings, embedding_dim), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + indices = builder.placeholder("indices", torch.tensor([0, 3, 5, 1])) + # Group-quant weight: full-rank [vocab, num_groups] scale/zp encodes block + # (1, dim/num_groups). Strictly-positive scales make dequant non-trivial. + scale = builder.placeholder( + "weight_scale", + torch.rand(num_embeddings, num_groups) + 0.1, + input_kind=InputKind.BUFFER, + ) + zp = builder.placeholder( + "weight_zero_point", + torch.zeros(num_embeddings, num_groups, dtype=torch.int64), + input_kind=InputKind.BUFFER, + ) + # Per-tensor output qparams are 0-dim lifted-constant placeholders (not + # aten.full ops). + out_scale_node = builder.placeholder( + "out_scale", + torch.tensor(out_scale, dtype=torch.float32), + input_kind=InputKind.BUFFER, + ) + out_zp_node = builder.placeholder( + "out_zero_point", + torch.tensor(out_zp, dtype=torch.int64), + input_kind=InputKind.BUFFER, + ) + embedding = builder.call_operator( + op=exir_ops.edge.fused_quant.embedding.default, + args=( + table, + scale, + zp, + torch.float32, + -8, + 7, + out_scale_node, + out_zp_node, + torch.int8, + -128, + 127, + indices, + ), + ) + out = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(embedding, [4, embedding_dim]), + ) + builder.output([out]) + return builder.get_program() + + @staticmethod + def _run(ep: ExportedProgram, indices: torch.Tensor) -> torch.Tensor: + """Run the program, feeding `indices` as the sole user input and pulling + params/buffers from the EP's backing stores.""" + full_inputs = [] + for spec in ep.graph_signature.input_specs: + if spec.kind == InputKind.USER_INPUT: + full_inputs.append(indices) + else: + assert spec.target is not None + if spec.target in ep.state_dict: + full_inputs.append(ep.state_dict[spec.target]) + else: + full_inputs.append(ep.constants[spec.target]) + return ep.graph_module(*full_inputs)[0] + + def test_prequantize_replaces_with_aten_embedding(self) -> None: + ep = self._build_quantized_embedding_program() + + result = PrequantizeEmbedding().call(ep) + self.assertTrue(result.modified) + ep = result.exported_program + graph = ep.graph_module.graph + + # fused_quant.embedding is replaced by a plain aten.embedding gather. + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.embedding.default, + ) + ), + 0, + ) + gather_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.embedding.default + ) + self.assertEqual(len(gather_nodes), 1) + # The gather emits int8 directly (no quant math downstream). + self.assertEqual(gather_nodes[0].meta["val"].dtype, torch.int8) + + # The scale/zero_point constants are dead and dropped from the signature. + buffer_targets = [ + spec.target + for spec in ep.graph_signature.input_specs + if spec.kind == InputKind.BUFFER + ] + self.assertNotIn("weight_scale", buffer_targets) + self.assertNotIn("weight_zero_point", buffer_targets) + + def test_prequantize_is_bit_exact(self) -> None: + ep = self._build_quantized_embedding_program() + indices = torch.tensor([0, 3, 5, 1]) + before = self._run(ep, indices) + + result = PrequantizeEmbedding().call(ep) + after = self._run(result.exported_program, indices) + + self.assertEqual(before.dtype, torch.int8) + self.assertTrue( + torch.equal(before, after), + f"prequantization changed the gather result:\n{before}\nvs\n{after}", + ) + + def test_prequantize_skips_float_output_embedding(self) -> None: + """An embedding whose output stays float (no out qparams) is left alone.""" + torch.manual_seed(0) + builder = ProgramBuilder() + table = builder.placeholder( + "table", + torch.randint(-8, 7, (6, 8), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + indices = builder.placeholder("indices", torch.tensor([0, 3, 5, 1])) + scale = builder.placeholder( + "weight_scale", + torch.rand(6, 2) + 0.1, + input_kind=InputKind.BUFFER, + ) + zp = builder.placeholder( + "weight_zero_point", + torch.zeros(6, 2, dtype=torch.int64), + input_kind=InputKind.BUFFER, + ) + embedding = builder.call_operator( + op=exir_ops.edge.fused_quant.embedding.default, + args=( + table, + scale, + zp, + torch.float32, + -8, + 7, + None, # out_scale -> float output + None, # out_zero_point + torch.int8, + -128, + 127, + indices, + ), + ) + builder.output([embedding]) + ep = builder.get_edge_program().exported_program() + + result = PrequantizeEmbedding().call(ep) + self.assertFalse(result.modified) + self.assertEqual( + len( + result.exported_program.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.embedding.default, + ) + ), + 1, + ) diff --git a/backends/fused_quant/test/test_quant_absorption.py b/backends/fused_quant/test/test_quant_absorption.py new file mode 100644 index 00000000000..ee9147a35d1 --- /dev/null +++ b/backends/fused_quant/test/test_quant_absorption.py @@ -0,0 +1,1774 @@ +# 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. + + +import operator +import unittest +from typing import Optional + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_from_node, + get_scale, + get_zero_point, +) +from executorch.backends.fused_quant.optimization_passes.quant_absorption import ( + QuantAbsorptionPass, +) +from executorch.backends.fused_quant.test.helpers import ( + create_per_axis_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import add_constant, get_input_kind +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensor +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + + +class QuantAbsorptionTest(unittest.TestCase): + """Tests for the QuantAbsorptionPass.""" + + def _build_fused_linear_with_dequant_quant( + self, + out_scale: float = 0.5, + out_zp: int = 0, + dequant_scale: Optional[float] = None, + dequant_zp: Optional[int] = None, + downstream_scale: float = 0.3, + downstream_zp: int = 2, + ) -> ExportedProgram: + """Build: fused_quant.linear(Q(A)) → DQ(D) → Q(B).""" + + dequant_scale = out_scale if dequant_scale is None else dequant_scale + dequant_zp = out_zp if dequant_zp is None else dequant_zp + + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(1, 8)) + + out_qparams = create_per_tensor_qparams( + builder, scale=out_scale, zero_point=out_zp, dtype=torch.int8 + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + x, + None, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(linear,), + kwargs={ + "scale": dequant_scale, + "zero_point": dequant_zp, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": downstream_scale, + "zero_point": downstream_zp, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + builder.output([quant]) + return builder.get_program() + + def test_absorb_mints_constant_tensor_qparams_no_cross_kind_clash(self) -> None: + """The out qparams QuantAbsorption mints must be the SAME kind as the + fusion's per-tensor _scale/_zero_point (lifted CONSTANT_TENSORs). Minting + them as a different kind (e.g. BUFFER) gives a b__scale node whose fqn + '_scale' still collides with the fusion's ep.constants entry -- a + cross-kind clash add_constant's node-name uniquification can't catch. Seed + the fusion-style constants and assert the absorb still succeeds and lands + its qparams in the same store.""" + ep = self._build_fused_linear_with_dequant_quant( + out_scale=0.5, out_zp=0, downstream_scale=0.3, downstream_zp=2 + ) + linear = next( + n + for n in ep.graph.nodes + if n.op == "call_function" + and n.target == exir_ops.edge.fused_quant.linear.default + ) + # Seed the exact fqns QuantAbsorption will mint, as CONSTANT_TENSOR -- + # exactly what FusedQuantFusion leaves for a per-tensor qparam. + add_constant(ep, "_scale", torch.tensor(1.0), linear, InputKind.CONSTANT_TENSOR) + add_constant( + ep, + "_zero_point", + torch.tensor(0, dtype=torch.int64), + linear, + InputKind.CONSTANT_TENSOR, + ) + + # A BUFFER mint would raise AssertionError (fqn '_scale' already backed). + result = QuantAbsorptionPass().call(ep) + assert result is not None + self.assertTrue(result.modified) + + fused = next( + n + for n in result.exported_program.graph.nodes + if n.op == "call_function" + and n.target == exir_ops.edge.fused_quant.linear.default + ) + out_qp = get_qparams_from_node(fused, "out") + assert out_qp is not None + self.assertEqual( + get_input_kind(result.exported_program, out_qp.scale), + InputKind.CONSTANT_TENSOR, + ) + self.assertEqual( + get_input_kind(result.exported_program, out_qp.zero_point), + InputKind.CONSTANT_TENSOR, + ) + + def test_absorb_multi_output_layer_norm_through_getitem(self) -> None: + """Absorb the output-0 requant of a MULTI-output fused op. + + ``fused_quant.native_layer_norm`` returns ``(out, mean, rstd)``; only + output 0 is quantized and its value flows through ``getitem(op, 0)``. The + pass must reach the ``getitem(0) → dequant → quant`` through the accessor, + fold the quant into the op's out qparams, and leave mean/rstd (getitem + 1/2) untouched.""" + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(2, 3, 4)) + weight = builder.placeholder( + "weight", torch.randn(4), input_kind=InputKind.BUFFER + ) + bias = builder.placeholder("bias", torch.randn(4), input_kind=InputKind.BUFFER) + out_qparams = create_per_tensor_qparams( + builder, scale=0.5, zero_point=0, dtype=torch.int8 + ) + ln = builder.call_operator( + op=exir_ops.edge.fused_quant.native_layer_norm.default, + args=( + x, + None, # inp_scale (input unquantized) + None, # inp_zero_point + torch.float32, + 0, + 0, + *out_qparams, # output-0 qparams + [4], # normalized_shape + weight, + bias, + 1e-5, # eps + ), + ) + g0 = builder.call_getitem(ln, 0) + g1 = builder.call_getitem(ln, 1) + g2 = builder.call_getitem(ln, 2) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(g0,), + kwargs={ + "scale": 0.5, + "zero_point": 0, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.3, + "zero_point": 2, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + builder.output([quant, g1, g2]) + ep = builder.get_program() + + result = QuantAbsorptionPass().call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ), + 0, + "output-0 dequant should be absorbed", + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + 0, + "output-0 quant should be absorbed", + ) + + ln_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.native_layer_norm.default, + ) + self.assertEqual(len(ln_nodes), 1) + out_qp = get_qparams_from_node(ln_nodes[0], "out") + assert out_qp is not None + self.assertAlmostEqual(get_scale(result.exported_program, out_qp), 0.3) + self.assertEqual(get_zero_point(result.exported_program, out_qp), 2) + + # out/mean/rstd getitems all survive; getitem(0) now yields int8. + getitems = [ + n + for n in graph.nodes + if n.op == "call_function" and n.target == operator.getitem + ] + self.assertEqual(len(getitems), 3) + g0_after = next(n for n in getitems if n.args[1] == 0) + self.assertEqual(g0_after.meta["val"].dtype, torch.int8) + + def test_absorb_dequant_quant(self) -> None: + """fused_quant.linear → dequant → quant should be absorbed.""" + + ep = self._build_fused_linear_with_dequant_quant( + out_scale=0.5, + out_zp=0, + downstream_scale=0.3, + downstream_zp=2, + ) + + result = QuantAbsorptionPass().call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(len(dequant_nodes), 0, "dequant should be removed") + + quant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + self.assertEqual(len(quant_nodes), 0, "quant should be removed") + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1) + + out_qp = get_qparams_from_node(linear_nodes[0], "out") + assert out_qp is not None + self.assertAlmostEqual(get_scale(result.exported_program, out_qp), 0.3) + + self.assertEqual(get_zero_point(result.exported_program, out_qp), 2) + + def test_absorb_composes_mismatched_dequant_affine_effect(self) -> None: + ep = self._build_fused_linear_with_dequant_quant( + out_scale=0.5, + out_zp=4, + dequant_scale=1.0, + dequant_zp=2, + downstream_scale=0.25, + downstream_zp=3, + ) + + result = QuantAbsorptionPass()(ep) + + self.assertTrue(result.modified) + graph = result.exported_program.graph + self.assertFalse( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ) + self.assertFalse( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ) + (linear,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + out_qparams = get_qparams_from_node(linear, "out") + assert out_qparams is not None + replacement_scale = get_scale(result.exported_program, out_qparams) + replacement_zp = get_zero_point(result.exported_program, out_qparams) + self.assertAlmostEqual(0.125, replacement_scale) + self.assertEqual(11, replacement_zp) + + source = torch.tensor([-2.0, 0.0, 3.0]) + original_zero_error_codes = ((source / 0.5 + 4) - 2) * 1.0 / 0.25 + 3 + replacement_zero_error_codes = source / replacement_scale + replacement_zp + torch.testing.assert_close( + original_zero_error_codes, replacement_zero_error_codes + ) + + def test_no_absorb_when_composed_zero_point_is_fractional(self) -> None: + ep = self._build_fused_linear_with_dequant_quant( + out_scale=0.5, + out_zp=1, + dequant_scale=0.3, + dequant_zp=0, + downstream_scale=0.2, + downstream_zp=0, + ) + + result = QuantAbsorptionPass()(ep) + + self.assertFalse(result.modified) + graph = result.exported_program.graph + self.assertEqual( + 1, + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ), + ) + self.assertEqual( + 1, + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + ) + + def _build_fused_embedding_with_bare_quant( + self, + downstream_scale: float = 0.3, + downstream_zp: int = 2, + ) -> ExportedProgram: + """Build: fused_quant.embedding(out=None, float gather) → quant(B). + + The embedding leaves its output unquantized and a consumer quantizes the + gathered rows -- the bare-quantize shape that QuantAbsorptionPass folds + into the embedding's (previously null) output qparams. + """ + builder = ProgramBuilder() + table = builder.placeholder( + "table", torch.randint(-8, 7, (6, 8), dtype=torch.int8) + ) + indices = builder.placeholder("indices", torch.tensor([0, 3, 5, 1])) + # Table quantized per-tensor; dtype is the dequantized (float) output + # dtype, so the gather is float. Output left unquantized. + weight_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.float32 + ) + embedding = builder.call_operator( + op=exir_ops.edge.fused_quant.embedding.default, + args=( + table, + *weight_qparams, + None, # out_scale + None, # out_zero_point + torch.int8, + -128, + 127, + indices, + ), + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(embedding,), + kwargs={ + "scale": downstream_scale, + "zero_point": downstream_zp, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + builder.output([quant]) + return builder.get_program() + + def test_absorb_bare_quant_into_unquantized_embedding_out(self) -> None: + """fused_quant.embedding(out=None) → quant folds the quantize into the + embedding's output qparams (the precondition for prequantization).""" + ep = self._build_fused_embedding_with_bare_quant( + downstream_scale=0.3, downstream_zp=2 + ) + result = QuantAbsorptionPass().call(ep) + self.assertTrue(result.modified) + graph = result.exported_program.graph + + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + 0, + "the bare quantize should be absorbed", + ) + embedding_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.embedding.default + ) + self.assertEqual(len(embedding_nodes), 1) + emb = embedding_nodes[0] + out_scale = get_arg(emb, "out_scale", Optional[torch.fx.Node]) + assert isinstance(out_scale, torch.fx.Node) + out_qp = get_qparams_from_node(emb, "out") + assert out_qp is not None + self.assertAlmostEqual(get_scale(result.exported_program, out_qp), 0.3) + self.assertEqual(get_zero_point(result.exported_program, out_qp), 2) + # The fused op now emits the quantized dtype. + self.assertEqual(emb.meta["val"].dtype, torch.int8) + + def test_no_absorb_bare_quant_when_embedding_output_forks(self) -> None: + """A fused op feeding both a quantize and another consumer keeps a float + path, so the bare-quantize fold must not fire (single-user only).""" + builder = ProgramBuilder() + table = builder.placeholder( + "table", torch.randint(-8, 7, (6, 8), dtype=torch.int8) + ) + indices = builder.placeholder("indices", torch.tensor([0, 3, 5, 1])) + weight_qparams = create_per_tensor_qparams( + builder, scale=0.5, dtype=torch.float32 + ) + embedding = builder.call_operator( + op=exir_ops.edge.fused_quant.embedding.default, + args=( + table, + *weight_qparams, + None, + None, + torch.int8, + -128, + 127, + indices, + ), + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(embedding,), + kwargs={ + "scale": 0.3, + "zero_point": 2, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + # A second (float) consumer of the embedding output. + relu = builder.call_operator( + op=exir_ops.edge.aten.relu.default, args=(embedding,) + ) + builder.output([quant, relu]) + ep = builder.get_program() + + result = QuantAbsorptionPass()(ep) + assert result is not None + self.assertFalse(result.modified) + emb = result.exported_program.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.embedding.default + )[0] + self.assertIsNone(get_arg(emb, "out_scale", Optional[torch.fx.Node])) + + def test_absorbed_qparams_vals_are_fake(self) -> None: + """The new out_scale/out_zero_point nodes must carry FakeTensor + meta['val']s, since the pass runs in fake tensor mode and downstream + passes rely on meta['val'] being fake.""" + + ep = self._build_fused_linear_with_dequant_quant( + out_scale=0.5, + out_zp=0, + downstream_scale=0.3, + downstream_zp=2, + ) + + result = QuantAbsorptionPass().call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1) + + for arg_name in ("out_scale", "out_zero_point"): + qparam_node = get_arg(linear_nodes[0], arg_name, Optional[torch.fx.Node]) + assert isinstance(qparam_node, torch.fx.Node) + val = qparam_node.meta.get("val") + self.assertIsInstance( + val, + FakeTensor, + f"{arg_name} node's meta['val'] must be a FakeTensor, got {type(val)}", + ) + + def test_no_absorb_when_no_mismatch(self) -> None: + """No absorption when there's no downstream requantize/dequant-quant.""" + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(1, 8)) + + out_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + x, + None, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + builder.output([linear]) + ep = builder.get_program() + + result = QuantAbsorptionPass().call(ep) + assert result is not None + self.assertFalse(result.modified) + + def _build_fused_linear_with_forked_dequant( + self, + out_scale: float = 0.5, + out_zp: int = 0, + dequant_scale: Optional[float] = None, + dequant_zp: Optional[int] = None, + downstream_scale: float = 0.3, + downstream_zp: int = 2, + ) -> ExportedProgram: + """Build: fused_quant.linear(Q(A)) → DQ(D) → Q(B) + └→ relu (float user) + """ + dequant_scale = out_scale if dequant_scale is None else dequant_scale + dequant_zp = out_zp if dequant_zp is None else dequant_zp + + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(1, 8)) + + out_qparams = create_per_tensor_qparams( + builder, scale=out_scale, zero_point=out_zp, dtype=torch.int8 + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + x, + None, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(linear,), + kwargs={ + "scale": dequant_scale, + "zero_point": dequant_zp, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": downstream_scale, + "zero_point": downstream_zp, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + float_user = builder.call_operator( + op=exir_ops.edge.aten.relu.default, + args=(dequant,), + ) + builder.output([quant, float_user]) + return builder.get_program() + + def test_no_absorb_forked_dequant_by_default(self) -> None: + """Forked dequant should NOT be absorbed when absorb_with_fork=False.""" + ep = self._build_fused_linear_with_forked_dequant() + + result = QuantAbsorptionPass(absorb_with_fork=False).call(ep) + assert result is not None + self.assertFalse(result.modified) + + def test_absorb_forked_dequant_when_enabled(self) -> None: + """Forked dequant should be absorbed when absorb_with_fork=True. + + Before: linear(out_scale=0.5) → dequant(0.5) → quant(0.3) + └→ relu + After: linear(out_scale=0.3) → dequant(0.3) → relu + └→ (direct int8, quant removed) + """ + ep = self._build_fused_linear_with_forked_dequant( + out_scale=0.5, + out_zp=0, + downstream_scale=0.3, + downstream_zp=2, + ) + + result = QuantAbsorptionPass(absorb_with_fork=True).call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + + quant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + self.assertEqual(len(quant_nodes), 0, "quant should be removed") + + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(len(dequant_nodes), 1, "dequant should remain for float user") + self.assertEqual( + get_arg(dequant_nodes[0], "scale", float), + 0.3, + "dequant scale should be updated to match new out_qparams", + ) + self.assertEqual( + get_arg(dequant_nodes[0], "zero_point", int), + 2, + "dequant zero_point should be updated to match new out_qparams", + ) + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1) + out_qp = get_qparams_from_node(linear_nodes[0], "out") + assert out_qp is not None + self.assertAlmostEqual(get_scale(result.exported_program, out_qp), 0.3) + + def test_absorb_forked_mismatched_dequant_preserves_float_effect(self) -> None: + ep = self._build_fused_linear_with_forked_dequant( + out_scale=0.5, + out_zp=4, + dequant_scale=1.0, + dequant_zp=2, + downstream_scale=0.25, + downstream_zp=3, + ) + + result = QuantAbsorptionPass(absorb_with_fork=True)(ep) + + self.assertTrue(result.modified) + graph = result.exported_program.graph + (linear,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + out_qparams = get_qparams_from_node(linear, "out") + assert out_qparams is not None + replacement_scale = get_scale(result.exported_program, out_qparams) + replacement_zp = get_zero_point(result.exported_program, out_qparams) + self.assertAlmostEqual(0.125, replacement_scale) + self.assertEqual(11, replacement_zp) + + (dequant,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + downstream_scale = get_arg(dequant, "scale", float) + downstream_zp = get_arg(dequant, "zero_point", int) + self.assertAlmostEqual(0.25, downstream_scale) + self.assertEqual(3, downstream_zp) + + source = torch.tensor([-2.0, 0.0, 3.0]) + original_float = (source / 0.5 + 4 - 2) * 1.0 + replacement_codes = source / replacement_scale + replacement_zp + replacement_float = (replacement_codes - downstream_zp) * downstream_scale + torch.testing.assert_close(original_float, replacement_float) + + def _build_fused_linear_with_permute_and_forked_dequant( + self, + out_scale: float = 0.5, + out_zp: int = 0, + downstream_scale: float = 0.3, + downstream_zp: int = 2, + downstream_dtype: torch.dtype = torch.int8, + ) -> ExportedProgram: + """Build: fused_quant.linear → permute → dequant → quant + └→ relu (float user) + """ + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(1, 8)) + + out_qparams = create_per_tensor_qparams( + builder, scale=out_scale, zero_point=out_zp, dtype=torch.int8 + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + x, + None, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(linear, [0, 1]), + ) + + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(permute,), + kwargs={ + "scale": out_scale, + "zero_point": out_zp, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": downstream_scale, + "zero_point": downstream_zp, + "quant_min": torch.iinfo(downstream_dtype).min, + "quant_max": torch.iinfo(downstream_dtype).max, + "dtype": downstream_dtype, + }, + ) + float_user = builder.call_operator( + op=exir_ops.edge.aten.relu.default, + args=(dequant,), + ) + builder.output([quant, float_user]) + return builder.get_program() + + def test_absorb_through_permute(self) -> None: + """Absorb dequant→quant through a permute passthrough op. + + Before: linear(out=0.5) → permute → dequant(0.5) → quant(0.3) + └→ relu + After: linear(out=0.3) → permute → dequant(0.3) → relu + └→ (direct int8, quant removed) + """ + ep = self._build_fused_linear_with_permute_and_forked_dequant( + out_scale=0.5, + out_zp=0, + downstream_scale=0.3, + downstream_zp=2, + ) + + result = QuantAbsorptionPass(absorb_with_fork=True).call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + + quant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + self.assertEqual(len(quant_nodes), 0, "quant should be removed") + + permute_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.permute_copy.default, + ) + self.assertEqual(len(permute_nodes), 1, "permute should remain") + + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(len(dequant_nodes), 1, "dequant should remain for float user") + self.assertEqual( + get_arg(dequant_nodes[0], "scale", float), + 0.3, + "dequant scale should be updated", + ) + self.assertEqual( + get_arg(dequant_nodes[0], "zero_point", int), + 2, + "dequant zero_point should be updated", + ) + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1) + out_qp = get_qparams_from_node(linear_nodes[0], "out") + assert out_qp is not None + self.assertAlmostEqual(get_scale(result.exported_program, out_qp), 0.3) + + def test_no_absorb_through_permute_without_flag(self) -> None: + """Passthrough absorption should not fire when dequant is forked + and absorb_with_fork is False.""" + ep = self._build_fused_linear_with_permute_and_forked_dequant() + + result = QuantAbsorptionPass(absorb_with_fork=False).call(ep) + assert result is not None + self.assertFalse(result.modified) + + def test_absorb_through_permute_single_user(self) -> None: + """Passthrough absorption works without fork flag when dequant + has a single user (the quant). + + Before: linear(out=0.5) → permute → dequant(0.5) → quant(0.3) + After: linear(out=0.3) → permute → ... + """ + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(1, 8)) + + out_qparams = create_per_tensor_qparams( + builder, scale=0.5, zero_point=0, dtype=torch.int8 + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + x, + None, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(linear, [0, 1]), + ) + + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(permute,), + kwargs={ + "scale": 0.5, + "zero_point": 0, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.3, + "zero_point": 2, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + builder.output([quant]) + ep = builder.get_program() + + result = QuantAbsorptionPass(absorb_with_fork=False).call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + + quant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + self.assertEqual(len(quant_nodes), 0, "quant should be removed") + + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(len(dequant_nodes), 0, "dequant should be removed") + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + out_qp = get_qparams_from_node(linear_nodes[0], "out") + assert out_qp is not None + self.assertAlmostEqual(get_scale(result.exported_program, out_qp), 0.3) + + def test_absorb_through_permute_changes_dtype_forked(self) -> None: + """When the downstream quant targets a different dtype than the fused + op's original output (int8 → uint8), the new dtype must propagate to the + fused op out qparams, the permute's meta['val'], and the kept dequant. + """ + ep = self._build_fused_linear_with_permute_and_forked_dequant( + out_scale=0.5, + out_zp=0, + downstream_scale=0.3, + downstream_zp=128, + downstream_dtype=torch.uint8, + ) + + result = QuantAbsorptionPass(absorb_with_fork=True).call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + + permute_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.permute_copy.default, + ) + self.assertEqual(len(permute_nodes), 1) + self.assertEqual( + permute_nodes[0].meta["val"].dtype, + torch.uint8, + "permute meta['val'] should take the downstream quant's out dtype", + ) + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1) + self.assertEqual( + get_arg(linear_nodes[0], "out_dtype", torch.dtype), torch.uint8 + ) + self.assertEqual(get_arg(linear_nodes[0], "out_quant_min", int), 0) + self.assertEqual(get_arg(linear_nodes[0], "out_quant_max", int), 255) + self.assertEqual(linear_nodes[0].meta["val"].dtype, torch.uint8) + + # The forked dequant is kept for the float (relu) user, but its input is + # now the requantized uint8 value, so its input qparams must follow. + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(len(dequant_nodes), 1) + self.assertEqual(get_arg(dequant_nodes[0], "dtype", torch.dtype), torch.uint8) + self.assertEqual(get_arg(dequant_nodes[0], "quant_min", int), 0) + self.assertEqual(get_arg(dequant_nodes[0], "quant_max", int), 255) + self.assertEqual(get_arg(dequant_nodes[0], "scale", float), 0.3) + + def test_absorb_through_permute_changes_dtype_single_user(self) -> None: + """Non-forked path: dequant is dropped, and the new dtype still + propagates to the fused op and the permute it now feeds.""" + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(1, 8)) + + out_qparams = create_per_tensor_qparams( + builder, scale=0.5, zero_point=0, dtype=torch.int8 + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + x, + None, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(linear, [0, 1]), + ) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(permute,), + kwargs={ + "scale": 0.5, + "zero_point": 0, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.3, + "zero_point": 128, + "quant_min": 0, + "quant_max": 255, + "dtype": torch.uint8, + }, + ) + builder.output([quant]) + ep = builder.get_program() + + result = QuantAbsorptionPass(absorb_with_fork=False).call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(len(dequant_nodes), 0, "dequant should be removed") + + permute_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.permute_copy.default, + ) + self.assertEqual(len(permute_nodes), 1) + self.assertEqual(permute_nodes[0].meta["val"].dtype, torch.uint8) + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual( + get_arg(linear_nodes[0], "out_dtype", torch.dtype), torch.uint8 + ) + self.assertEqual(linear_nodes[0].meta["val"].dtype, torch.uint8) + + def test_absorb_through_view_changes_dtype_single_user(self) -> None: + """A view_copy is a per-tensor-transparent passthrough, like permute. + + Before: linear(out=0.5, int8) -> view_copy -> dequant(0.5) -> quant(0.3, uint8) + After: linear(out=0.3, uint8) -> view_copy -> ... + + Covers both that view is a valid passthrough and that the dtype change + (int8 -> uint8) propagates to the fused op and the view's meta['val']. + """ + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(4, 8)) + + out_qparams = create_per_tensor_qparams( + builder, scale=0.5, zero_point=0, dtype=torch.int8 + ) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + x, + None, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.float32, + 0, + 0, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + # A genuine reshape (4, 4) -> (2, 8) to exercise the view passthrough. + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(linear, [2, 8]), + ) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(view,), + kwargs={ + "scale": 0.5, + "zero_point": 0, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.3, + "zero_point": 128, + "quant_min": 0, + "quant_max": 255, + "dtype": torch.uint8, + }, + ) + builder.output([quant]) + ep = builder.get_program() + + result = QuantAbsorptionPass(absorb_with_fork=False).call(ep) + assert result is not None + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + + quant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + self.assertEqual(len(quant_nodes), 0, "quant should be removed") + + dequant_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual(len(dequant_nodes), 0, "dequant should be removed") + + view_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.view_copy.default, + ) + self.assertEqual(len(view_nodes), 1, "view should remain") + self.assertEqual( + view_nodes[0].meta["val"].dtype, + torch.uint8, + "view carrying the absorbed output must be retyped to uint8", + ) + + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1) + out_qp = get_qparams_from_node(linear_nodes[0], "out") + assert out_qp is not None + self.assertAlmostEqual(get_scale(result.exported_program, out_qp), 0.3) + self.assertEqual( + get_arg(linear_nodes[0], "out_dtype", torch.dtype), torch.uint8 + ) + self.assertEqual(linear_nodes[0].meta["val"].dtype, torch.uint8) + + def _build_dequant_quant_pool( + self, + *, + dequant_scale: float = 0.25, + quant_scale: float = 0.5, + consumer_input_scale: float = 0.5, + consumer_input_zero_point: int = -5, + input_step: int = 1, + expose_quant: bool = False, + expose_dequant: bool = False, + input_passthrough: str | None = None, + expose_passthrough: bool = False, + ) -> tuple[ExportedProgram, torch.Tensor]: + builder = ProgramBuilder() + inp = (torch.arange(16, dtype=torch.uint8) * input_step).reshape(1, 4, 4, 1) + inp_node = builder.placeholder("inp", inp) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(inp_node,), + kwargs={ + "scale": dequant_scale, + "zero_point": 0, + "quant_min": 0, + "quant_max": 255, + "dtype": torch.uint8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": quant_scale, + "zero_point": -5, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + consumer_input = quant + if input_passthrough == "view": + consumer_input = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(quant, [1, 2, 8, 1]), + ) + elif input_passthrough == "permute": + consumer_input = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(quant, [0, 2, 1, 3]), + ) + elif input_passthrough is not None: + raise ValueError(f"Unsupported input passthrough: {input_passthrough}") + pool_input_qparams = create_per_tensor_qparams( + builder, + scale=consumer_input_scale, + zero_point=consumer_input_zero_point, + dtype=torch.float32, + ) + pool_output_qparams = create_per_tensor_qparams( + builder, + scale=0.1, + dtype=torch.uint8, + ) + pool = builder.call_operator( + op=exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + args=( + consumer_input, + *pool_input_qparams, + *pool_output_qparams[:3], + 0, + 255, + [2, 2], + [2, 2], + [0, 0], + False, + True, + None, + ), + ) + outputs = [pool] + if expose_quant: + outputs.append(quant) + if expose_dequant: + outputs.append(dequant) + if expose_passthrough: + outputs.append(consumer_input) + builder.output(outputs) + return builder.get_program(), inp + + @staticmethod + def _direct_pool(inp: torch.Tensor) -> torch.Tensor: + return torch.ops.fused_quant.avg_pool2d_channels_last.default( + inp, + torch.tensor(0.25, dtype=torch.float32), + torch.tensor(0, dtype=torch.int64), + torch.float32, + 0, + 255, + torch.tensor(0.1, dtype=torch.float32), + torch.tensor(0, dtype=torch.int64), + torch.uint8, + 0, + 255, + [2, 2], + [2, 2], + [0, 0], + False, + True, + None, + ) + + def _build_dequant_quant_add(self) -> tuple[ExportedProgram, torch.Tensor]: + builder = ProgramBuilder() + inp = torch.arange(8, dtype=torch.uint8).reshape(2, 4) + inp_node = builder.placeholder("inp", inp) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(inp_node,), + kwargs={ + "scale": 0.25, + "zero_point": 0, + "quant_min": 0, + "quant_max": 255, + "dtype": torch.uint8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.5, + "zero_point": -5, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + input_qparams = create_per_tensor_qparams( + builder, + scale=0.5, + zero_point=-5, + dtype=torch.float32, + ) + other_qparams = create_per_tensor_qparams( + builder, + scale=0.5, + zero_point=-5, + dtype=torch.float32, + ) + output_qparams = create_per_tensor_qparams(builder, scale=0.1) + add = builder.call_operator( + op=exir_ops.edge.fused_quant.add.default, + args=( + quant, + quant, + *input_qparams, + *other_qparams, + *output_qparams, + 1.0, + ), + ) + builder.output([add]) + return builder.get_program(), inp + + def test_absorbs_dequant_quant_into_fused_consumer_input(self) -> None: + ep, inp = self._build_dequant_quant_pool() + + result = QuantAbsorptionPass()(ep) + + self.assertTrue(result.modified) + (actual,) = result.exported_program.module()(inp) + self.assertTrue(torch.equal(self._direct_pool(inp), actual)) + + graph = result.exported_program.graph + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ), + ) + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ), + ) + (pool,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + ) + pool_input = get_arg(pool, "inp", torch.fx.Node) + self.assertEqual("placeholder", pool_input.op) + input_qparams = get_qparams_from_node(pool, "inp") + assert input_qparams is not None + self.assertAlmostEqual(0.25, get_scale(result.exported_program, input_qparams)) + self.assertEqual(0, get_zero_point(result.exported_program, input_qparams)) + self.assertEqual(torch.float32, input_qparams.dtype) + self.assertEqual(0, input_qparams.quant_min) + self.assertEqual(255, input_qparams.quant_max) + result.exported_program.validate() + + def test_absorbs_consumer_input_through_view_copy(self) -> None: + ep, inp = self._build_dequant_quant_pool(input_passthrough="view") + + result = QuantAbsorptionPass()(ep) + + self.assertTrue(result.modified) + (actual,) = result.exported_program.module()(inp) + expected = self._direct_pool( + torch.ops.aten.view_copy.default(inp, [1, 2, 8, 1]) + ) + self.assertTrue(torch.equal(expected, actual)) + + graph = result.exported_program.graph + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ), + ) + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ), + ) + (view,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.view_copy.default + ) + self.assertEqual("placeholder", view.args[0].op) + self.assertEqual(torch.uint8, view.meta["val"].dtype) + (pool,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + ) + self.assertIs(view, get_arg(pool, "inp", torch.fx.Node)) + input_qparams = get_qparams_from_node(pool, "inp") + assert input_qparams is not None + self.assertAlmostEqual(0.25, get_scale(result.exported_program, input_qparams)) + self.assertEqual(0, get_zero_point(result.exported_program, input_qparams)) + result.exported_program.validate() + + def test_absorbs_consumer_input_through_permute_copy(self) -> None: + ep, inp = self._build_dequant_quant_pool(input_passthrough="permute") + + result = QuantAbsorptionPass()(ep) + + self.assertTrue(result.modified) + (actual,) = result.exported_program.module()(inp) + expected = self._direct_pool( + torch.ops.aten.permute_copy.default(inp, [0, 2, 1, 3]) + ) + self.assertTrue(torch.equal(expected, actual)) + + graph = result.exported_program.graph + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ), + ) + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ), + ) + (permute,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + self.assertEqual("placeholder", permute.args[0].op) + self.assertEqual(torch.uint8, permute.meta["val"].dtype) + (pool,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + ) + self.assertIs(permute, get_arg(pool, "inp", torch.fx.Node)) + input_qparams = get_qparams_from_node(pool, "inp") + assert input_qparams is not None + self.assertAlmostEqual(0.25, get_scale(result.exported_program, input_qparams)) + self.assertEqual(0, get_zero_point(result.exported_program, input_qparams)) + result.exported_program.validate() + + def test_does_not_absorb_through_shared_input_passthrough(self) -> None: + ep, _ = self._build_dequant_quant_pool( + input_passthrough="view", expose_passthrough=True + ) + + result = QuantAbsorptionPass()(ep) + + self.assertFalse(result.modified) + self.assertEqual( + 1, + len( + result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + ) + + def test_absorbs_all_quantized_inputs_of_binary_op(self) -> None: + ep, inp = self._build_dequant_quant_add() + + result = QuantAbsorptionPass()(ep) + + self.assertTrue(result.modified) + (actual,) = result.exported_program.module()(inp) + expected = torch.ops.fused_quant.add.default( + inp, + inp, + torch.tensor(0.25, dtype=torch.float32), + torch.tensor(0, dtype=torch.int64), + torch.float32, + 0, + 255, + torch.tensor(0.25, dtype=torch.float32), + torch.tensor(0, dtype=torch.int64), + torch.float32, + 0, + 255, + torch.tensor(0.1, dtype=torch.float32), + torch.tensor(0, dtype=torch.int64), + torch.int8, + -128, + 127, + 1.0, + ) + self.assertTrue(torch.equal(expected, actual)) + + graph = result.exported_program.graph + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ), + ) + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ), + ) + (add,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.add.default, + ) + for input_name in ("inp", "other"): + self.assertEqual("placeholder", get_arg(add, input_name, torch.fx.Node).op) + input_qparams = get_qparams_from_node(add, input_name) + assert input_qparams is not None + self.assertAlmostEqual( + 0.25, get_scale(result.exported_program, input_qparams) + ) + self.assertEqual(0, get_zero_point(result.exported_program, input_qparams)) + self.assertEqual(torch.float32, input_qparams.dtype) + self.assertEqual(0, input_qparams.quant_min) + self.assertEqual(255, input_qparams.quant_max) + result.exported_program.validate() + + def test_composes_mismatched_consumer_input_qparams(self) -> None: + ep, inp = self._build_dequant_quant_pool( + consumer_input_scale=0.4, + consumer_input_zero_point=-3, + input_step=2, + ) + expected = ep.module()(inp) + + result = QuantAbsorptionPass()(ep) + + self.assertTrue(result.modified) + actual = result.exported_program.module()(inp) + self.assertEqual(len(expected), len(actual)) + for expected_output, actual_output in zip(expected, actual): + torch.testing.assert_close(actual_output, expected_output) + + graph = result.exported_program.graph + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ), + ) + self.assertEqual( + [], + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ), + ) + (pool,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + ) + input_qparams = get_qparams_from_node(pool, "inp") + assert input_qparams is not None + self.assertAlmostEqual(0.2, get_scale(result.exported_program, input_qparams)) + self.assertEqual(4, get_zero_point(result.exported_program, input_qparams)) + result.exported_program.validate() + + def test_does_not_absorb_unrepresentable_consumer_input_qparams(self) -> None: + ep, _ = self._build_dequant_quant_pool( + quant_scale=0.4, + consumer_input_scale=0.5, + consumer_input_zero_point=-4, + ) + + result = QuantAbsorptionPass()(ep) + + self.assertFalse(result.modified) + + def test_does_not_absorb_per_channel_consumer_input(self) -> None: + builder = ProgramBuilder() + inp = torch.arange(8, dtype=torch.uint8).reshape(1, 4, 2) + inp_node = builder.placeholder("inp", inp) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(inp_node,), + kwargs={ + "scale": 0.25, + "zero_point": 0, + "quant_min": 0, + "quant_max": 255, + "dtype": torch.uint8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.5, + "zero_point": -5, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + input_qparams = create_per_axis_qparams( + builder, + num_channels=2, + ndim=3, + axis=2, + dtype=torch.float32, + ) + output_qparams = create_per_tensor_qparams(builder) + relu = builder.call_operator( + op=exir_ops.edge.fused_quant.relu.default, + args=(quant, *input_qparams, *output_qparams), + ) + builder.output([relu]) + ep = builder.get_program() + + result = QuantAbsorptionPass()(ep) + + self.assertFalse(result.modified) + self.assertEqual( + 1, + len( + result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + ) + + def test_does_not_absorb_shared_quantize(self) -> None: + ep, _ = self._build_dequant_quant_pool(expose_quant=True) + + result = QuantAbsorptionPass()(ep) + + self.assertFalse(result.modified) + self.assertEqual( + 1, + len( + result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + ) + + def test_absorbs_quantize_and_preserves_shared_dequantize(self) -> None: + ep, inp = self._build_dequant_quant_pool(expose_dequant=True) + _, expected_dequant = ep.module()(inp) + + result = QuantAbsorptionPass()(ep) + + self.assertTrue(result.modified) + actual_pool, actual_dequant = result.exported_program.module()(inp) + self.assertTrue(torch.equal(self._direct_pool(inp), actual_pool)) + self.assertTrue(torch.equal(expected_dequant, actual_dequant)) + self.assertEqual( + [], + result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ), + ) + self.assertEqual( + 1, + len( + result.exported_program.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ), + ) + result.exported_program.validate() + + +# ── SinkConstantCat models ──────────────────────────────────────────── diff --git a/backends/fused_quant/test/test_replace_batch_norm_with_conv.py b/backends/fused_quant/test/test_replace_batch_norm_with_conv.py new file mode 100644 index 00000000000..838e78fcd74 --- /dev/null +++ b/backends/fused_quant/test/test_replace_batch_norm_with_conv.py @@ -0,0 +1,187 @@ +# 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. + +from __future__ import annotations + +import unittest +from typing import cast + +import torch +from executorch.backends.fused_quant.pre_quantize_passes.replace_batch_norm_with_conv import ( + ReplaceBatchNormWithConv, +) +from executorch.backends.fused_quant.test.helpers import trace +from executorch.backends.transforms.utils import get_constant +from parameterized import parameterized +from torch import fx, nn +from torch.export import ExportedProgram + +_NATIVE_BATCH_NORM = torch.ops.aten._native_batch_norm_legit_no_training.default +_CONVOLUTION = torch.ops.aten.convolution.default + + +class _BatchNormModel(nn.Module): + def __init__( + self, + channels: int, + input_rank: int, + affine: bool = True, + ) -> None: + super().__init__() + if input_rank in (2, 3): + self.batch_norm = nn.BatchNorm1d(channels, affine=affine) + elif input_rank == 4: + self.batch_norm = nn.BatchNorm2d(channels, affine=affine) + elif input_rank == 5: + self.batch_norm = nn.BatchNorm3d(channels, affine=affine) + else: + raise ValueError(f"Unsupported input rank: {input_rank}") + + self.batch_norm.running_mean = torch.randn(channels) + self.batch_norm.running_var = torch.rand(channels) + 0.5 + if affine: + self.batch_norm.weight = nn.Parameter(torch.randn(channels)) + self.batch_norm.bias = nn.Parameter(torch.randn(channels)) + + def forward(self, inp: torch.Tensor) -> torch.Tensor: + return self.batch_norm(inp) + + +class _NativeAuxOutputModel(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.randn(channels)) + self.bias = nn.Parameter(torch.randn(channels)) + self.register_buffer("running_mean", torch.randn(channels)) + self.register_buffer("running_var", torch.rand(channels) + 0.5) + + def forward(self, inp: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + output, saved_mean, _saved_invstd = _NATIVE_BATCH_NORM( + inp, + self.weight, + self.bias, + self.running_mean, + self.running_var, + 0.1, + 1e-5, + ) + return output, saved_mean + + +def _trace(model: nn.Module, inp: torch.Tensor) -> ExportedProgram: + return trace(model, (inp,)) + + +def _only(exported_program: ExportedProgram, target: object) -> fx.Node: + nodes = exported_program.graph_module.graph.find_nodes( + op="call_function", target=target + ) + assert len(nodes) == 1 + return nodes[0] + + +class ReplaceBatchNormWithConvTest(unittest.TestCase): + @parameterized.expand( + [ + ("batch_norm_1d_without_length", (2, 3)), + ("batch_norm_1d", (2, 3, 5)), + ("batch_norm_2d", (2, 3, 4, 5)), + ("batch_norm_3d", (2, 3, 2, 4, 5)), + ] + ) + def test_replaces_with_dimension_generic_depthwise_convolution( + self, + _name: str, + input_shape: tuple[int, ...], + ) -> None: + torch.manual_seed(0) + model = _BatchNormModel(input_shape[1], len(input_shape)).eval() + inp = torch.randn(input_shape) + expected = model(inp) + exported_program = _trace(model, inp) + + self.assertEqual( + len( + exported_program.graph_module.graph.find_nodes( + op="call_function", target=_NATIVE_BATCH_NORM + ) + ), + 1, + ) + + result = ReplaceBatchNormWithConv()(exported_program) + + self.assertTrue(result.modified) + exported_program = result.exported_program + self.assertEqual( + len( + exported_program.graph_module.graph.find_nodes( + op="call_function", target=_NATIVE_BATCH_NORM + ) + ), + 0, + ) + conv = _only(exported_program, _CONVOLUTION) + weight = get_constant(exported_program, cast(fx.Node, conv.args[1])) + self.assertIsNotNone(weight) + spatial_dims = max(len(input_shape) - 2, 1) + self.assertEqual( + tuple(cast(torch.Tensor, weight).shape), + (input_shape[1], 1, *([1] * spatial_dims)), + ) + self.assertEqual(conv.args[3], [1] * spatial_dims) + self.assertEqual(conv.args[4], [0] * spatial_dims) + self.assertEqual(conv.args[5], [1] * spatial_dims) + self.assertFalse(conv.args[6]) + self.assertEqual(conv.args[7], [0] * spatial_dims) + self.assertEqual(conv.args[8], input_shape[1]) + + actual = exported_program.module()(inp) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_supports_batch_norm_without_affine_parameters(self) -> None: + torch.manual_seed(0) + model = _BatchNormModel(3, input_rank=4, affine=False).eval() + inp = torch.randn(2, 3, 4, 5) + expected = model(inp) + exported_program = _trace(model, inp) + + result = ReplaceBatchNormWithConv()(exported_program) + + self.assertTrue(result.modified) + actual = result.exported_program.module()(inp) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_preserves_source_metadata_on_convolution(self) -> None: + model = _BatchNormModel(3, input_rank=4).eval() + inp = torch.randn(2, 3, 4, 5) + exported_program = _trace(model, inp) + sentinel = [("batch_norm", torch.nn.functional.batch_norm)] + _only(exported_program, _NATIVE_BATCH_NORM).meta["source_fn_stack"] = sentinel + + result = ReplaceBatchNormWithConv()(exported_program) + + self.assertEqual( + _only(result.exported_program, _CONVOLUTION).meta["source_fn_stack"], + sentinel, + ) + + def test_native_batch_norm_with_used_auxiliary_output_is_not_replaced(self) -> None: + model = _NativeAuxOutputModel(3).eval() + inp = torch.randn(2, 3, 4, 5) + exported_program = _trace(model, inp) + + result = ReplaceBatchNormWithConv()(exported_program) + + self.assertFalse(result.modified) + self.assertEqual( + len( + result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=_NATIVE_BATCH_NORM + ) + ), + 1, + ) diff --git a/backends/fused_quant/test/test_replace_dequant_quant_with_requantize.py b/backends/fused_quant/test/test_replace_dequant_quant_with_requantize.py new file mode 100644 index 00000000000..dd72cfa1035 --- /dev/null +++ b/backends/fused_quant/test/test_replace_dequant_quant_with_requantize.py @@ -0,0 +1,294 @@ +# 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. + + +import unittest + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.graph_utils import ( + get_qparams_from_node, + get_scale, + get_zero_point, +) +from executorch.backends.fused_quant.optimization_passes.replace_dequant_quant_with_requantize import ( + ReplaceDequantQuantWithRequantize, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram +from torch.export.graph_signature import TensorArgument + + +class ReplaceDequantQuantWithRequantizeTest(unittest.TestCase): + def _build_dequant_quant( + self, + *, + include_float_output: bool = False, + include_second_quant: bool = False, + ) -> tuple[ExportedProgram, torch.Tensor]: + builder = ProgramBuilder() + inp = torch.tensor([[-4, 0, 7]], dtype=torch.int8) + inp_node = builder.placeholder("inp", inp) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(inp_node,), + kwargs={ + "scale": 0.5, + "zero_point": 1, + "quant_min": -128, + "quant_max": 127, + "dtype": torch.int8, + }, + ) + quant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.25, + "zero_point": 3, + "quant_min": 0, + "quant_max": 255, + "dtype": torch.uint8, + }, + ) + outputs = [quant] + if include_second_quant: + outputs.append( + builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(dequant,), + kwargs={ + "scale": 0.125, + "zero_point": 5, + "quant_min": 0, + "quant_max": 255, + "dtype": torch.uint8, + }, + ) + ) + if include_float_output: + outputs.append(dequant) + builder.output(outputs) + return builder.get_program(), inp + + def test_replaces_dequant_quant(self) -> None: + ep, inp = self._build_dequant_quant() + expected = ep.module()(inp) + + result = ReplaceDequantQuantWithRequantize()(ep) + + self.assertTrue(result.modified) + actual = result.exported_program.module()(inp) + self.assertEqual(len(actual), len(expected)) + for actual_output, expected_output in zip(actual, expected): + self.assertTrue(torch.equal(actual_output, expected_output)) + + graph = result.exported_program.graph_module.graph + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ), + 0, + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + 0, + ) + (requantize,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.requantize.default + ) + inp_qparams = get_qparams_from_node(requantize, "inp") + out_qparams = get_qparams_from_node(requantize, "out") + assert inp_qparams is not None and out_qparams is not None + self.assertAlmostEqual(get_scale(result.exported_program, inp_qparams), 0.5) + self.assertEqual(get_zero_point(result.exported_program, inp_qparams), 1) + self.assertEqual(inp_qparams.dtype, torch.float32) + self.assertAlmostEqual(get_scale(result.exported_program, out_qparams), 0.25) + self.assertEqual(get_zero_point(result.exported_program, out_qparams), 3) + self.assertEqual(out_qparams.dtype, torch.uint8) + self.assertEqual(requantize.meta["val"].dtype, torch.uint8) + self.assertEqual( + result.exported_program.graph_signature.user_outputs, + (requantize.name,), + ) + result.exported_program.validate() + + def test_replaces_quant_branch_and_preserves_float_user(self) -> None: + ep, inp = self._build_dequant_quant(include_float_output=True) + expected = ep.module()(inp) + + result = ReplaceDequantQuantWithRequantize()(ep) + + self.assertTrue(result.modified) + actual = result.exported_program.module()(inp) + for actual_output, expected_output in zip(actual, expected): + self.assertTrue(torch.equal(actual_output, expected_output)) + graph = result.exported_program.graph_module.graph + (dequant,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + 0, + ) + (requantize,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.requantize.default + ) + self.assertEqual( + result.exported_program.graph_signature.user_outputs, + (requantize.name, dequant.name), + ) + result.exported_program.validate() + + def test_replaces_multiple_quant_branches(self) -> None: + ep, inp = self._build_dequant_quant(include_second_quant=True) + expected = ep.module()(inp) + + result = ReplaceDequantQuantWithRequantize()(ep) + + self.assertTrue(result.modified) + actual = result.exported_program.module()(inp) + for actual_output, expected_output in zip(actual, expected): + self.assertTrue(torch.equal(actual_output, expected_output)) + graph = result.exported_program.graph_module.graph + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ), + 0, + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + 0, + ) + requantize_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.requantize.default + ) + self.assertEqual(len(requantize_nodes), 2) + actual_out_qparams = [] + for requantize in requantize_nodes: + out_qparams = get_qparams_from_node(requantize, "out") + assert out_qparams is not None + actual_out_qparams.append( + ( + get_scale(result.exported_program, out_qparams), + get_zero_point(result.exported_program, out_qparams), + ) + ) + self.assertCountEqual(actual_out_qparams, [(0.25, 3), (0.125, 5)]) + first_inp_qparams = get_qparams_from_node(requantize_nodes[0], "inp") + second_inp_qparams = get_qparams_from_node(requantize_nodes[1], "inp") + assert first_inp_qparams is not None and second_inp_qparams is not None + self.assertIs(first_inp_qparams.scale, second_inp_qparams.scale) + self.assertIs(first_inp_qparams.zero_point, second_inp_qparams.zero_point) + self.assertEqual( + result.exported_program.graph_signature.user_outputs, + tuple(requantize.name for requantize in requantize_nodes), + ) + result.exported_program.validate() + + def test_replaces_multiple_quant_branches_and_preserves_float_user(self) -> None: + ep, inp = self._build_dequant_quant( + include_float_output=True, + include_second_quant=True, + ) + expected = ep.module()(inp) + + result = ReplaceDequantQuantWithRequantize()(ep) + + self.assertTrue(result.modified) + actual = result.exported_program.module()(inp) + for actual_output, expected_output in zip(actual, expected): + self.assertTrue(torch.equal(actual_output, expected_output)) + graph = result.exported_program.graph_module.graph + (dequant,) = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + ) + ), + 0, + ) + requantize_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.requantize.default + ) + self.assertEqual(len(requantize_nodes), 2) + self.assertEqual( + result.exported_program.graph_signature.user_outputs, + ( + requantize_nodes[0].name, + requantize_nodes[1].name, + dequant.name, + ), + ) + result.exported_program.validate() + + def test_dequant_without_quant_user_is_noop(self) -> None: + ep, _ = self._build_dequant_quant(include_float_output=True) + graph = ep.graph_module.graph + quant = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + )[0] + quant.replace_all_uses_with(quant.args[0]) + graph.erase_node(quant) + dequant = graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + )[0] + ep.graph_signature.output_specs[0].arg = TensorArgument(dequant.name) + ep.graph_module.recompile() + ep.validate() + + result = ReplaceDequantQuantWithRequantize()(ep) + + self.assertFalse(result.modified) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ) + ), + 1, + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.requantize.default, + ) + ), + 0, + ) diff --git a/backends/fused_quant/test/test_replace_mm_with_bmm.py b/backends/fused_quant/test/test_replace_mm_with_bmm.py new file mode 100644 index 00000000000..e0946da0126 --- /dev/null +++ b/backends/fused_quant/test/test_replace_mm_with_bmm.py @@ -0,0 +1,113 @@ +# 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. + + +from __future__ import annotations + +import unittest + +import torch +from executorch.backends.fused_quant.pre_quantize_passes.replace_mm_with_bmm import ( + ReplaceMmWithBmm, +) +from executorch.backends.fused_quant.test.helpers import trace +from torch import nn +from torch.export import ExportedProgram + +_MM = torch.ops.aten.mm.default +_BMM = torch.ops.aten.bmm.default + + +class _ConstantMmModel(nn.Module): + def __init__(self, in_features: int, out_features: int) -> None: + super().__init__() + self.proj = nn.Parameter(torch.randn(in_features, out_features)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.mm(x, self.proj) + + +class _DynamicMmModel(nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.mm(x, y) + + +class _MultiMmModel(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.proj_a = nn.Parameter(torch.randn(dim, dim)) + self.proj_b = nn.Parameter(torch.randn(dim, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + a = torch.mm(x, self.proj_a) + return torch.mm(a, self.proj_b) + + +def _count(ep: ExportedProgram, target: object) -> int: + return len(ep.graph_module.graph.find_nodes(op="call_function", target=target)) + + +class ReplaceMmWithBmmTest(unittest.TestCase): + def test_constant_weight_replaced(self) -> None: + torch.manual_seed(0) + model = _ConstantMmModel(8, 4).eval() + inp = torch.randn(2, 8) + expected = model(inp) + ep = trace(model, (inp,)) + + self.assertEqual(_count(ep, _MM), 1) + + result = ReplaceMmWithBmm().call(ep.graph_module) + + self.assertTrue(result.modified) + self.assertEqual(_count(ep, _MM), 0) + self.assertEqual(_count(ep, _BMM), 1) + actual = ep.module()(inp) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_dynamic_inputs_replaced(self) -> None: + torch.manual_seed(0) + model = _DynamicMmModel().eval() + x = torch.randn(3, 5) + y = torch.randn(5, 4) + expected = model(x, y) + ep = trace(model, (x, y)) + + self.assertEqual(_count(ep, _MM), 1) + + result = ReplaceMmWithBmm().call(ep.graph_module) + + self.assertTrue(result.modified) + self.assertEqual(_count(ep, _MM), 0) + self.assertEqual(_count(ep, _BMM), 1) + actual = ep.module()(x, y) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_multiple_mm_replaced(self) -> None: + torch.manual_seed(0) + model = _MultiMmModel(6).eval() + inp = torch.randn(2, 6) + expected = model(inp) + ep = trace(model, (inp,)) + + self.assertEqual(_count(ep, _MM), 2) + + result = ReplaceMmWithBmm().call(ep.graph_module) + + self.assertTrue(result.modified) + self.assertEqual(_count(ep, _MM), 0) + self.assertEqual(_count(ep, _BMM), 2) + actual = ep.module()(inp) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_no_mm_not_modified(self) -> None: + model = nn.Linear(4, 3).eval() + inp = torch.randn(2, 4) + ep = trace(model, (inp,)) + + result = ReplaceMmWithBmm().call(ep.graph_module) + + self.assertFalse(result.modified) diff --git a/backends/fused_quant/test/test_sink_constant_cat.py b/backends/fused_quant/test/test_sink_constant_cat.py new file mode 100644 index 00000000000..c4e4b95959e --- /dev/null +++ b/backends/fused_quant/test/test_sink_constant_cat.py @@ -0,0 +1,388 @@ +# 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. + + +import unittest + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.optimization_passes.sink_constant_cat import ( + SinkConstantCat, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export +from torch.export.graph_signature import InputKind, OutputKind + + +class _CatReluConstModel(torch.nn.Module): + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(torch.cat([x, self.const], dim=1)) + + +class _CatPermuteConstModel(torch.nn.Module): + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.cat([x, self.const], dim=1).permute(0, 2, 1) + + +class _CatReluPermuteConstModel(torch.nn.Module): + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + c = torch.cat([x, self.const], dim=1) + r = torch.relu(c) + return r.permute(0, 2, 1) + + +class _CatViewNegativeDimConstModel(torch.nn.Module): + """cat on dim=2, then view with -1 on a non-cat dimension. + + Input x is [2, 3, 4], const is [2, 3, 4], cat on dim=2 gives [2, 3, 8]. + The view reshapes to [2, -1] = [2, 24], merging dims 1 and 2. + Since cat dim (2) is merged with dim 1, the tracker invalidates and + the pass bails — but critically, the -1 must be resolved via + meta["val"].shape, not the raw arg (which would contain -1). + """ + + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 3, 4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + c = torch.cat([x, self.const], dim=2) # [2, 3, 8] + return c.view(c.shape[0], -1) # [2, 24] + + +class _CatViewNegativeDimPreservesCatModel(torch.nn.Module): + """cat on dim=1, then view with -1 that merges only non-cat dims. + + Input x is [2, 4, 2, 3], const is [2, 4, 2, 3], cat on dim=1 gives + [2, 8, 2, 3]. The view reshapes to [2, 8, -1] = [2, 8, 6], merging + dims 2 and 3 while cat dim (1) stays intact. + """ + + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 2, 3)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + c = torch.cat([x, self.const], dim=1) # [2, 8, 2, 3] + return c.view(c.shape[0], c.shape[1], -1) # [2, 8, 6] + + +class _CatUnsqueezeConstModel(torch.nn.Module): + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(4, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + c = torch.cat([x, self.const], dim=0) + return c.unsqueeze(0) + + +class _CatMultiUserConstModel(torch.nn.Module): + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + c = torch.cat([x, self.const], dim=1) + a = torch.relu(c) + b = c.permute(0, 2, 1) + return a, b + + +class _CatThreeInputsOneConstModel(torch.nn.Module): + """Two activations + one constant. The cat should shrink to two elements.""" + + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.relu(torch.cat([x, y, self.const], dim=1)) + + +class _CatAllActivationsModel(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.relu(torch.cat([x, y], dim=1)) + + +class _CatPartialSinkableModel(torch.nn.Module): + """Cat with two users: relu (sinkable) and matmul (not sinkable). + + Should bail because not all users are sinkable. + """ + + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + c = torch.cat([x, self.const], dim=1) + a = torch.relu(c) + b = torch.bmm(c, c.transpose(1, 2)) + return a, b + + +class _CatSliceOnNonCatDimModel(torch.nn.Module): + """cat on dim=1, then slice on dim=2 (not the cat dim). Safe to sink.""" + + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + c = torch.cat([x, self.const], dim=1) + return c[:, :, 2:6] + + +class _CatSliceOnCatDimModel(torch.nn.Module): + """cat on dim=1, then slice on the same dimension. + + Sinking would break because slicing on the cat dimension can cross + input boundaries — the result cannot be reproduced by slicing each + cat input independently. + """ + + const: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("const", torch.randn(2, 4, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + c = torch.cat([x, self.const], dim=1) + return c[:, 2:6, :] + + +class SinkConstantCatTest(unittest.TestCase): + def _run_pass( + self, + model: torch.nn.Module, + inputs: tuple[torch.Tensor, ...], + ) -> tuple[torch.fx.GraphModule, bool]: + exported = export(model, inputs) + edge = to_edge(exported, compile_config=EdgeCompileConfig(_skip_dim_order=True)) + ep = edge.exported_program() + + with torch.no_grad(): + ref_output = ep.module()(*inputs) + + result = SinkConstantCat()(ep) + + with torch.no_grad(): + new_output = result.exported_program.module()(*inputs) + + if isinstance(ref_output, torch.Tensor): + self.assertTrue( + (new_output == ref_output).all(), + "Pass should not change numerical output", + ) + else: + for i, (new, ref) in enumerate(zip(new_output, ref_output)): + self.assertTrue( + (new == ref).all(), + f"Pass changed numerical output at index {i}", + ) + return result.exported_program.graph_module, result.modified + + def test_cat_relu_pushed(self) -> None: + x = torch.randn(2, 4, 8) + _, modified = self._run_pass(_CatReluConstModel(), (x,)) + self.assertTrue(modified) + + def test_cat_permute_pushed(self) -> None: + x = torch.randn(2, 4, 8) + _, modified = self._run_pass(_CatPermuteConstModel(), (x,)) + self.assertTrue(modified) + + def test_cat_relu_permute_pushed(self) -> None: + x = torch.randn(2, 4, 8) + _, modified = self._run_pass(_CatReluPermuteConstModel(), (x,)) + self.assertTrue(modified) + + def test_cat_unsqueeze_pushed(self) -> None: + x = torch.randn(4, 8) + _, modified = self._run_pass(_CatUnsqueezeConstModel(), (x,)) + self.assertTrue(modified) + + def test_cat_multi_user_pushed(self) -> None: + """Each user of the cat gets its own sunk chain. + + Before: cat([x, const]) → relu, permute + After: relu(x), relu(const_folded) → cat → output_a + permute(x), permute(const_folded) → cat → output_b + """ + x = torch.randn(2, 4, 8) + gm, modified = self._run_pass(_CatMultiUserConstModel(), (x,)) + self.assertTrue(modified) + + graph = gm.graph + + cat_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.cat.default + ) + self.assertEqual(len(cat_nodes), 2, "Should have one sunk cat per user path") + + relu_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.relu.default + ) + self.assertEqual( + len(relu_nodes), 1, "Should have one relu for the activation path" + ) + + permute_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + self.assertEqual( + len(permute_nodes), 1, "Should have one permute for the activation path" + ) + + for cat_node in cat_nodes: + tensors = cat_node.args[0] + assert isinstance(tensors, (list, tuple)) + self.assertEqual( + len(tensors), 2, "Each sunk cat should merge activation + folded const" + ) + + def test_skips_activation_only_inputs(self) -> None: + x = torch.randn(2, 4, 8) + y = torch.randn(2, 4, 8) + _, modified = self._run_pass(_CatAllActivationsModel(), (x, y)) + self.assertFalse(modified) + + def test_skips_mutable_buffer(self) -> None: + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(2, 4, 8)) + state = builder.placeholder( + "state", torch.randn(2, 4, 8), input_kind=InputKind.BUFFER + ) + cat = builder.call_operator( + op=exir_ops.edge.aten.cat.default, args=([x, state], 1) + ) + relu = builder.call_operator(op=exir_ops.edge.aten.relu.default, args=(cat,)) + builder.output( + [state, relu], + output_kinds=[OutputKind.BUFFER_MUTATION, OutputKind.USER_OUTPUT], + output_targets=["state", None], + ) + program = builder.get_program() + self.assertTrue(program.graph_signature.buffers_to_mutate) + + result = SinkConstantCat()(program) + + self.assertFalse(result.modified) + graph = result.exported_program.graph_module.graph + (cat_node,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.cat.default + ) + (relu_node,) = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.relu.default + ) + state_node = next(node for node in graph.nodes if node.name == "state") + cat_inputs = cat_node.args[0] + assert isinstance(cat_inputs, (list, tuple)) + self.assertIn(state_node, cat_inputs) + self.assertIs(relu_node.args[0], cat_node) + + def test_skips_when_not_all_users_sinkable(self) -> None: + """Should bail when only some users of the cat are sinkable.""" + x = torch.randn(2, 4, 8) + _, modified = self._run_pass(_CatPartialSinkableModel(), (x,)) + self.assertFalse(modified) + + def test_three_inputs_one_const_keeps_activation_cat(self) -> None: + """With [act, act, const], the original cat should shrink to [act, act].""" + x = torch.randn(2, 4, 8) + y = torch.randn(2, 4, 8) + model = _CatThreeInputsOneConstModel() + + exported = export(model, (x, y)) + edge = to_edge(exported, compile_config=EdgeCompileConfig(_skip_dim_order=True)) + ep = edge.exported_program() + + with torch.no_grad(): + ref_output = ep.module()(x, y) + + result = SinkConstantCat()(ep) + self.assertTrue(result.modified) + + with torch.no_grad(): + new_output = result.exported_program.module()(x, y) + self.assertTrue( + (new_output == ref_output).all(), + "Pass should not change numerical output", + ) + + cat_nodes = result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.cat.default + ) + # Should have 2 cats: one with [act, act] feeding the chain, + # and one at the sink merging the activation path with the folded const + self.assertEqual(len(cat_nodes), 2) + for cat_node in cat_nodes: + tensors = cat_node.args[0] + assert isinstance(tensors, (list, tuple)) + self.assertEqual(len(tensors), 2) + + def test_cat_slice_on_non_cat_dim_pushed(self) -> None: + """Should sink when a slice is on a dimension other than the cat dim.""" + x = torch.randn(2, 4, 8) + _, modified = self._run_pass(_CatSliceOnNonCatDimModel(), (x,)) + self.assertTrue(modified) + + def test_skips_slice_on_cat_dim(self) -> None: + """Should not sink when a slice is applied on the cat dimension, + since the slice range may cross cat input boundaries.""" + x = torch.randn(2, 4, 8) + _, modified = self._run_pass(_CatSliceOnCatDimModel(), (x,)) + self.assertFalse(modified) + + def test_skips_view_that_merges_cat_dim(self) -> None: + """Should not sink when a view merges the cat dimension with another + dimension, since the cat can no longer be reconstructed after the reshape.""" + x = torch.randn(2, 3, 4) + _, modified = self._run_pass(_CatViewNegativeDimConstModel(), (x,)) + self.assertFalse(modified) + + def test_cat_view_negative_dim_preserves_cat_dim(self) -> None: + """view with -1 that merges only non-cat dims should allow sinking + and produce numerically correct results.""" + x = torch.randn(2, 4, 2, 3) + _, modified = self._run_pass(_CatViewNegativeDimPreservesCatModel(), (x,)) + self.assertTrue(modified) diff --git a/backends/fused_quant/test/test_split_linear_at_slices.py b/backends/fused_quant/test/test_split_linear_at_slices.py new file mode 100644 index 00000000000..dabf346d285 --- /dev/null +++ b/backends/fused_quant/test/test_split_linear_at_slices.py @@ -0,0 +1,698 @@ +# 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. + + +import unittest +from typing import Any, Optional, TypeAlias + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.optimization_passes.constant_fold import ( + ConstantFold, +) +from executorch.backends.fused_quant.optimization_passes.split_linear_at_slices import ( + SplitLinearAtSlices, +) +from executorch.backends.fused_quant.test.helpers import ( + create_per_channel_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.backends.transforms.utils import get_constant +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ProxyValue +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +_QBlock: TypeAlias = tuple[Any, Any, torch.dtype, int, int] + + +class SplitLinearAtSlicesTest(unittest.TestCase): + """Tests for the SplitLinearAtSlices optimization pass.""" + + def _build_linear_with_slices( + self, + in_features: int = 8, + out_features: int = 16, + slices: list[tuple[int, int, int]] | None = None, + with_bias: bool = True, + ) -> tuple[ExportedProgram, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Build a linear → (slice, slice, ...) graph. + + Returns (program, input_tensor, weight_tensor, bias_tensor). + """ + + if slices is None: + slices = [(0, out_features // 2, 1), (out_features // 2, out_features, 1)] + + builder = ProgramBuilder() + inp_tensor = torch.randn(1, in_features) + weight_tensor = torch.randint( + -16, 16, (out_features, in_features), dtype=torch.int8 + ) + bias_tensor = torch.randn(out_features) if with_bias else None + + x = builder.placeholder("x", inp_tensor) + w = builder.placeholder("weight", weight_tensor, input_kind=InputKind.BUFFER) + b = ( + builder.placeholder("bias", bias_tensor, input_kind=InputKind.BUFFER) + if bias_tensor is not None + else None + ) + + weight_qparams = create_per_channel_qparams( + builder, out_features, scale=0.05, dtype=torch.float32 + ) + out_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + b, + None, + None, + torch.float32, + 0, + 0, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + outputs = [] + for start, end, step in slices: + sliced = builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, + args=(linear, -1, start, end, step), + ) + outputs.append(sliced) + + builder.output(outputs) + return builder.get_program(), inp_tensor, weight_tensor, bias_tensor + + def _assert_signature_consistent(self, ep: ExportedProgram) -> None: + """Assert that the graph signature, graph placeholders, and state_dict are in sync.""" + sig = ep.graph_signature + placeholders = [n for n in ep.graph_module.graph.nodes if n.op == "placeholder"] + placeholder_names = {n.name for n in placeholders} + + spec_names = {s.arg.name for s in sig.input_specs} + self.assertEqual( + placeholder_names, + spec_names, + "Placeholder nodes and input_specs should match", + ) + + for spec in sig.input_specs: + if spec.kind == InputKind.BUFFER: + self.assertIn( + spec.target, + ep.state_dict, + f"Buffer {spec.target} missing from state_dict", + ) + + def _run_program( + self, ep: ExportedProgram, inp: torch.Tensor + ) -> list[torch.Tensor]: + """Run a program and return all outputs.""" + full_inputs = [] + sig = ep.graph_signature + for spec in sig.input_specs: + if spec.kind == InputKind.USER_INPUT: + full_inputs.append(inp) + else: + assert spec.target is not None + # Params/buffers live in state_dict; constant-folded tensors + # (e.g. from ConstantFold) live in constants. + if spec.target in ep.state_dict: + full_inputs.append(ep.state_dict[spec.target]) + else: + full_inputs.append(ep.constants[spec.target]) + out = ep.graph_module(*full_inputs) + if isinstance(out, torch.Tensor): + return [out] + return list(out) + + def test_split_two_contiguous_slices(self) -> None: + """linear → (slice[0:8], slice[8:16]) splits into two smaller linears.""" + ep, inp, _, _ = self._build_linear_with_slices(slices=[(0, 8, 1), (8, 16, 1)]) + before_outputs = self._run_program(ep, inp) + + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 2, "Should have 2 split linears") + + slice_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.slice_copy.Tensor, + ) + self.assertEqual(len(slice_nodes), 0, "Slice nodes should be removed") + + self._assert_signature_consistent(result.exported_program) + + after_outputs = self._run_program(result.exported_program, inp) + for before, after in zip(before_outputs, after_outputs): + self.assertTrue( + torch.equal(before, after), + f"Numerical mismatch: {before} vs {after}", + ) + + def test_split_strided_even_odd(self) -> None: + """linear → (slice[0::2], slice[1::2]) splits into even/odd rows.""" + ep, inp, _, _ = self._build_linear_with_slices(slices=[(0, 16, 2), (1, 16, 2)]) + before_outputs = self._run_program(ep, inp) + + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 2, "Should have 2 split linears") + + self._assert_signature_consistent(result.exported_program) + + after_outputs = self._run_program(result.exported_program, inp) + for before, after in zip(before_outputs, after_outputs): + self.assertTrue( + torch.equal(before, after), + f"Numerical mismatch: {before} vs {after}", + ) + + def test_split_no_bias(self) -> None: + """Split works correctly without bias.""" + ep, inp, _, _ = self._build_linear_with_slices( + slices=[(0, 8, 1), (8, 16, 1)], with_bias=False + ) + before_outputs = self._run_program(ep, inp) + + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + + self._assert_signature_consistent(result.exported_program) + + after_outputs = self._run_program(result.exported_program, inp) + for before, after in zip(before_outputs, after_outputs): + self.assertTrue( + torch.equal(before, after), + f"Numerical mismatch: {before} vs {after}", + ) + + def test_no_split_overlapping_slices(self) -> None: + """Overlapping slices exceed out_features — should not split.""" + ep, _, _, _ = self._build_linear_with_slices(slices=[(0, 12, 1), (4, 16, 1)]) + + result = SplitLinearAtSlices().call(ep) + self.assertFalse( + result.modified, + "Should not split when overlapping slices exceed out_features", + ) + + def test_split_single_slice(self) -> None: + """A single slice on the output dim replaces the linear with a smaller one.""" + ep, inp, _, _ = self._build_linear_with_slices( + slices=[(0, 8, 1)], + ) + before_outputs = self._run_program(ep, inp) + + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 1, "Should have 1 smaller linear") + + slice_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.slice_copy.Tensor, + ) + self.assertEqual(len(slice_nodes), 0, "Slice node should be removed") + + self._assert_signature_consistent(result.exported_program) + + after_outputs = self._run_program(result.exported_program, inp) + for before, after in zip(before_outputs, after_outputs): + self.assertTrue( + torch.equal(before, after), + f"Numerical mismatch: {before} vs {after}", + ) + + def test_split_negative_start_end(self) -> None: + """Negative start/end are normalized so range() and slicing agree.""" + # slice[-4:, :] is the last 4 rows, slice[:-4, :] is the first 12 rows + ep, inp, _, _ = self._build_linear_with_slices(slices=[(0, -4, 1), (-4, 16, 1)]) + before_outputs = self._run_program(ep, inp) + + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + + graph = result.exported_program.graph_module.graph + linear_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.linear.default, + ) + self.assertEqual(len(linear_nodes), 2, "Should have 2 split linears") + + slice_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.aten.slice_copy.Tensor, + ) + self.assertEqual(len(slice_nodes), 0, "Slice nodes should be removed") + + self._assert_signature_consistent(result.exported_program) + + after_outputs = self._run_program(result.exported_program, inp) + for before, after in zip(before_outputs, after_outputs): + self.assertTrue( + torch.equal(before, after), + f"Numerical mismatch: {before} vs {after}", + ) + + def test_no_split_non_output_dim_slice(self) -> None: + """Slices on the batch dimension (dim=0) should not trigger split.""" + + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(4, 8)) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (16, 8), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_channel_qparams(builder, 16, dtype=torch.float32) + out_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + + linear = builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + None, + *inp_qparams, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + ), + ) + + slice0 = builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, + args=(linear, 0, 0, 2, 1), + ) + slice1 = builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, + args=(linear, 0, 2, 4, 1), + ) + builder.output([slice0, slice1]) + + ep = builder.get_program() + result = SplitLinearAtSlices().call(ep) + self.assertFalse( + result.modified, "Should not split when slices are on batch dimension" + ) + + # ---- helpers for the qparam-handling tests below ---- + + def _make_linear( + self, + builder: ProgramBuilder, + x: ProxyValue, + w: ProxyValue, + b: Optional[ProxyValue], + weight_block: _QBlock, + out_block: _QBlock, + bias_block: Optional[_QBlock] = None, + ) -> ProxyValue: + """Assemble a fused_quant.linear with per-tensor (None) input qparams.""" + if bias_block is None: + bias_block = (None, None, torch.float32, 0, 0) + return builder.call_operator( + op=exir_ops.edge.fused_quant.linear.default, + args=( + x, + w, + b, + None, + None, + torch.float32, + 0, + 0, + *weight_block, + *bias_block, + *out_block, + ), + ) + + def _add_slices( + self, + builder: ProgramBuilder, + linear: ProxyValue, + slices: list[tuple[int, int, int]], + ) -> None: + outs = [] + for start, end, step in slices: + outs.append( + builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, + args=(linear, -1, start, end, step), + ) + ) + builder.output(outs) + + def test_split_slices_varying_per_channel_qparams(self) -> None: + """Per-channel weight scale/zp that VARY per output channel are sliced to + match each split linear. (The uniform-scale tests above would not catch a + mis-slice, since every channel shares the same value.)""" + out_features, in_features = 16, 8 + builder = ProgramBuilder() + inp = torch.randn(1, in_features) + x = builder.placeholder("x", inp) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + ws = builder.placeholder( + "wt_scale", + (torch.arange(out_features, dtype=torch.float32) * 0.01 + 0.02).reshape( + out_features, 1 + ), + input_kind=InputKind.BUFFER, + ) + wzp = builder.placeholder( + "wt_zp", + (torch.arange(out_features, dtype=torch.int64) % 5).reshape( + out_features, 1 + ), + input_kind=InputKind.BUFFER, + ) + out_block = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + # weight qparam "dtype" is the dequantize OUTPUT dtype (float); the int8 + # weight is stored in the tensor, its range is in quant_min/max. + linear = self._make_linear( + builder, x, w, None, (ws, wzp, torch.float32, -128, 127), out_block + ) + slices = [(0, 8, 1), (8, 16, 1)] + self._add_slices(builder, linear, slices) + ep = builder.get_program() + + before = self._run_program(ep, inp) + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + self._assert_signature_consistent(result.exported_program) + + after = self._run_program(result.exported_program, inp) + for b, a in zip(before, after): + self.assertTrue(torch.equal(b, a), f"Numerical mismatch: {b} vs {a}") + + # Each split linear must carry a sliced (not full-rank) scale/zp. + linears = result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + self.assertEqual(len(linears), 2) + for lin in linears: + ws_node = get_arg(lin, "weight_scale", torch.fx.Node) + sliced = get_constant(result.exported_program, ws_node) + assert sliced is not None + self.assertEqual(list(sliced.shape), [8, 1]) + + def test_split_unquantized_weight_no_qparams(self) -> None: + """A linear with no weight scale/zp (unquantized weight) still splits.""" + out_features, in_features = 16, 8 + builder = ProgramBuilder() + inp = torch.randn(1, in_features) + x = builder.placeholder("x", inp) + w = builder.placeholder( + "weight", + torch.randn(out_features, in_features), + input_kind=InputKind.BUFFER, + ) + out_block = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + linear = self._make_linear( + builder, x, w, None, (None, None, torch.float32, 0, 0), out_block + ) + self._add_slices(builder, linear, [(0, 8, 1), (8, 16, 1)]) + ep = builder.get_program() + + before = self._run_program(ep, inp) + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified, "no-qparam weight should still split") + self._assert_signature_consistent(result.exported_program) + + linears = result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + self.assertEqual(len(linears), 2) + + after = self._run_program(result.exported_program, inp) + for b, a in zip(before, after): + self.assertTrue(torch.equal(b, a), f"Numerical mismatch: {b} vs {a}") + + def test_split_slices_quantized_bias_qparams(self) -> None: + """A per-output-channel quantized bias has its bias_scale/zp sliced too, + not just the bias tensor.""" + out_features, in_features = 16, 8 + builder = ProgramBuilder() + inp = torch.randn(1, in_features) + x = builder.placeholder("x", inp) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + b = builder.placeholder( + "bias", + torch.randint(-8, 8, (out_features,), dtype=torch.int32), + input_kind=InputKind.BUFFER, + ) + weight_block = create_per_channel_qparams( + builder, out_features, scale=0.05, dtype=torch.float32 + ) + bias_scale = builder.placeholder( + "bias_scale", + torch.arange(out_features, dtype=torch.float32) * 0.01 + 0.05, + input_kind=InputKind.BUFFER, + ) + bias_zp = builder.placeholder( + "bias_zp", + torch.zeros(out_features, dtype=torch.int64), + input_kind=InputKind.BUFFER, + ) + out_block = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + linear = self._make_linear( + builder, + x, + w, + b, + weight_block, + out_block, + bias_block=(bias_scale, bias_zp, torch.float32, -(2**31), 2**31 - 1), + ) + slices = [(0, 8, 1), (8, 16, 1)] + self._add_slices(builder, linear, slices) + ep = builder.get_program() + + before = self._run_program(ep, inp) + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + self._assert_signature_consistent(result.exported_program) + + after = self._run_program(result.exported_program, inp) + for b_out, a_out in zip(before, after): + self.assertTrue( + torch.equal(b_out, a_out), f"Numerical mismatch: {b_out} vs {a_out}" + ) + + linears = result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + self.assertEqual(len(linears), 2) + for lin in linears: + bs_node = get_arg(lin, "bias_scale", torch.fx.Node) + sliced_bs = get_constant(result.exported_program, bs_node) + assert sliced_bs is not None + self.assertEqual(list(sliced_bs.shape), [8], "bias_scale must be sliced") + + def test_per_tensor_weight_scale_reused_not_sliced(self) -> None: + """A per-tensor (scalar) weight scale broadcasts, so it is reused/shared + across the split linears rather than sliced.""" + out_features, in_features = 16, 8 + builder = ProgramBuilder() + inp = torch.randn(1, in_features) + x = builder.placeholder("x", inp) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + # Per-tensor scalar weight qparams; dtype is the dequant output dtype + # (float), the int8 weight range is in quant_min/max. + weight_block = create_per_tensor_qparams( + builder, scale=0.05, dtype=torch.float32 + ) + out_block = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + linear = self._make_linear(builder, x, w, None, weight_block, out_block) + self._add_slices(builder, linear, [(0, 8, 1), (8, 16, 1)]) + ep = builder.get_program() + + before = self._run_program(ep, inp) + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + + linears = result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + self.assertEqual(len(linears), 2) + ws_nodes = {get_arg(lin, "weight_scale", torch.fx.Node) for lin in linears} + self.assertEqual( + len(ws_nodes), 1, "per-tensor weight scale should be shared, not sliced" + ) + + after = self._run_program(result.exported_program, inp) + for b, a in zip(before, after): + self.assertTrue(torch.equal(b, a), f"Numerical mismatch: {b} vs {a}") + + def _build_linear_with_scale_behind_view( + self, out_features: int = 16, in_features: int = 8 + ) -> tuple[ExportedProgram, torch.Tensor]: + """Build a split graph whose per-channel weight scale/zp ride behind an + ``aten.view_copy`` (the shape fusion inserts), so they are NOT directly + resolvable constants until folded.""" + builder = ProgramBuilder() + inp = torch.randn(1, in_features) + x = builder.placeholder("x", inp) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + ws_1d = builder.placeholder( + "wt_scale_1d", + torch.arange(out_features, dtype=torch.float32) * 0.01 + 0.02, + input_kind=InputKind.BUFFER, + ) + wzp_1d = builder.placeholder( + "wt_zp_1d", + torch.arange(out_features, dtype=torch.int64) % 5, + input_kind=InputKind.BUFFER, + ) + ws_view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(ws_1d, [out_features, 1]) + ) + wzp_view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(wzp_1d, [out_features, 1]) + ) + out_block = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.int8) + linear = self._make_linear( + builder, + x, + w, + None, + (ws_view, wzp_view, torch.float32, -128, 127), + out_block, + ) + self._add_slices(builder, linear, [(0, 8, 1), (8, 16, 1)]) + return builder.get_program(), inp + + def test_no_split_per_channel_out_qparams(self) -> None: + """Per-channel OUTPUT qparams (out_scale varying along the output dim) are + on a different axis than the weight/bias qparams and aren't handled, so the + split bails rather than reuse a full-rank out_scale on a smaller linear.""" + out_features, in_features = 16, 8 + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(1, in_features)) + w = builder.placeholder( + "weight", + torch.randint(-16, 16, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.BUFFER, + ) + weight_block = create_per_channel_qparams( + builder, out_features, scale=0.05, dtype=torch.float32 + ) + # Per-channel (non-scalar) output scale/zp. + out_scale = builder.placeholder( + "out_scale", + torch.full((out_features,), 0.5, dtype=torch.float32), + input_kind=InputKind.BUFFER, + ) + out_zp = builder.placeholder( + "out_zp", + torch.zeros(out_features, dtype=torch.int64), + input_kind=InputKind.BUFFER, + ) + linear = self._make_linear( + builder, + x, + w, + None, + weight_block, + (out_scale, out_zp, torch.int8, -128, 127), + ) + self._add_slices(builder, linear, [(0, 8, 1), (8, 16, 1)]) + ep = builder.get_program() + + result = SplitLinearAtSlices().call(ep) + self.assertFalse(result.modified, "Should bail on per-channel output qparams") + + def test_no_split_when_per_channel_qparam_behind_view(self) -> None: + """If a per-output weight scale is still behind a view_copy (not yet folded + to a constant), the split aborts -- it must not slice the weight while + leaving a mismatched full-rank scale.""" + ep, _ = self._build_linear_with_scale_behind_view() + result = SplitLinearAtSlices().call(ep) + self.assertFalse( + result.modified, + "Should abort when a per-output qparam is not a resolvable constant", + ) + + def test_split_after_constant_fold_resolves_view(self) -> None: + """ConstantFold bakes the weight scale/zp view_copy, after which the split + slices them correctly -- mirroring the real pipeline ordering.""" + ep, inp = self._build_linear_with_scale_behind_view() + before = self._run_program(ep, inp) + + ep = ConstantFold().call(ep).exported_program + result = SplitLinearAtSlices().call(ep) + self.assertTrue(result.modified) + self._assert_signature_consistent(result.exported_program) + + after = self._run_program(result.exported_program, inp) + for b, a in zip(before, after): + self.assertTrue(torch.equal(b, a), f"Numerical mismatch: {b} vs {a}") + + linears = result.exported_program.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.fused_quant.linear.default + ) + self.assertEqual(len(linears), 2) + for lin in linears: + ws_node = get_arg(lin, "weight_scale", torch.fx.Node) + sliced = get_constant(result.exported_program, ws_node) + assert sliced is not None + self.assertEqual(list(sliced.shape), [8, 1]) diff --git a/backends/fused_quant/test/test_to_channels_first.py b/backends/fused_quant/test/test_to_channels_first.py new file mode 100644 index 00000000000..e60804769af --- /dev/null +++ b/backends/fused_quant/test/test_to_channels_first.py @@ -0,0 +1,162 @@ +# 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. + +from __future__ import annotations + +import unittest + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.optimization_passes.to_channels_first import ( + ToChannelsFirst, +) +from executorch.backends.fused_quant.test.helpers import ( + create_per_axis_qparams, + create_per_channel_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops + + +class ToChannelsFirstTest(unittest.TestCase): + def assert_bit_exact(self, expected: object, actual: object) -> None: + torch.testing.assert_close(expected, actual, rtol=0, atol=0) + + def test_convolution(self) -> None: + inp = torch.randint(-8, 8, (1, 6, 6, 2), dtype=torch.int8) + weight = torch.randint(-8, 8, (3, 3, 3, 2), dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + weight_node = builder.placeholder("weight", weight) + inp_qparams = create_per_axis_qparams( + builder, 2, 4, axis=3, dtype=torch.float32, vary_values=True + ) + weight_qparams = create_per_channel_qparams( + builder, 3, dtype=torch.float32, weight_ndim=4 + ) + out_qparams = create_per_axis_qparams(builder, 3, 4, axis=3, vary_values=True) + convolution = builder.call_operator( + exir_ops.edge.fused_quant.convolution_channels_last.default, + args=( + inp_node, + weight_node, + None, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + -128, + 127, + *out_qparams, + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + 1, + ), + ) + builder.output([convolution]) + program = builder.get_program() + expected = program.module()(inp, weight) + + result = ToChannelsFirst().call(program.graph_module) + + self.assertTrue(result.modified) + self.assert_bit_exact(expected, program.module()(inp, weight)) + self.assertEqual( + result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution_channels_last.default, + ), + [], + ) + [channels_first] = result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution.default, + ) + self.assertFalse(get_arg(channels_first, "transposed", bool)) + self.assertEqual(get_arg(channels_first, "output_padding", list[int]), [0, 0]) + + def test_max_pool2d_with_indices(self) -> None: + inp = torch.randint(-8, 8, (1, 6, 6, 2), dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + inp_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + out_qparams = create_per_tensor_qparams(builder, scale=0.5) + pool = builder.call_operator( + exir_ops.edge.fused_quant.max_pool2d_with_indices_channels_last.default, + args=( + inp_node, + *inp_qparams, + *out_qparams, + [3, 3], + [2, 2], + [1, 1], + [1, 1], + False, + ), + ) + values = builder.call_getitem(pool, 0) + indices = builder.call_getitem(pool, 1) + builder.output([values, indices]) + program = builder.get_program() + expected = program.module()(inp) + + result = ToChannelsFirst().call(program.graph_module) + + self.assertTrue(result.modified) + self.assert_bit_exact(expected, program.module()(inp)) + self.assertEqual( + len( + result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.max_pool2d_with_indices.default, + ) + ), + 1, + ) + + def test_avg_pool2d(self) -> None: + inp = torch.randint(-8, 8, (1, 6, 6, 2), dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + inp_qparams = create_per_tensor_qparams(builder, scale=0.5, dtype=torch.float32) + out_qparams = create_per_tensor_qparams(builder, scale=0.25) + pool = builder.call_operator( + exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + args=( + inp_node, + *inp_qparams, + *out_qparams, + [3, 3], + [2, 2], + [1, 1], + False, + True, + None, + ), + ) + builder.output([pool]) + program = builder.get_program() + expected = program.module()(inp) + + result = ToChannelsFirst().call(program.graph_module) + + self.assertTrue(result.modified) + self.assert_bit_exact(expected, program.module()(inp)) + self.assertEqual( + len( + result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d.default, + ) + ), + 1, + ) diff --git a/backends/fused_quant/test/test_to_channels_last.py b/backends/fused_quant/test/test_to_channels_last.py new file mode 100644 index 00000000000..036d2d8635c --- /dev/null +++ b/backends/fused_quant/test/test_to_channels_last.py @@ -0,0 +1,599 @@ +# 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. + + +import operator +import unittest +from typing import Any + +import executorch.backends.fused_quant.ops # noqa +import torch +from executorch.backends.fused_quant.optimization_passes.to_channels_last import ( + ToChannelsLast, +) +from executorch.backends.fused_quant.test.helpers import ( + create_per_axis_qparams, + create_per_channel_qparams, + create_per_tensor_qparams, +) +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from parameterized import parameterized +from torch.export import ExportedProgram + +# These follow aten.convolution's argument order. +_CONVOLUTION1D_ARGS: tuple[Any, ...] = ([1], [0], [1], False, [0], 1) +_CONVOLUTION2D_ARGS: tuple[Any, ...] = ( + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + 1, +) +_CONVOLUTION3D_ARGS: tuple[Any, ...] = ( + [1, 1, 1], + [0, 0, 0], + [1, 1, 1], + False, + [0, 0, 0], + 1, +) + + +def _build_fused_quant_conv_graph( + input_shape: tuple[int, ...], + weight_shape: tuple[int, ...], + bias_shape: tuple[int, ...] | None, + conv_args: tuple[Any, ...], + per_channel_activations: bool = False, +) -> tuple[ExportedProgram, tuple[torch.Tensor, ...]]: + """Build a graph with a fused_quant.convolution edge op for testing. + + Quant params are lifted CONSTANT_TENSOR/BUFFER placeholders (matching how + production lowers them), so they sit at the front of the graph rather than + as inline body nodes. + """ + inp = torch.randn(*input_shape).to(torch.int8) + weight = torch.randn(*weight_shape).to(torch.int8) + bias = torch.randn(*bias_shape) if bias_shape else None + + builder = ProgramBuilder() + + # Create placeholders first + inp_placeholder = builder.placeholder("inp", inp) + weight_placeholder = builder.placeholder("weight", weight) + bias_placeholder = builder.placeholder("bias", bias) if bias is not None else None + + stride, padding, dilation, transposed, output_padding, groups = conv_args + + # Create qparams as lifted front placeholders + inp_qparams = ( + create_per_axis_qparams( + builder, + input_shape[1], + len(input_shape), + axis=1, + dtype=torch.float32, + vary_values=True, + ) + if per_channel_activations + else create_per_tensor_qparams(builder, dtype=torch.float32) + ) + weight_qparams = create_per_channel_qparams( + builder, + weight_shape[0], + dtype=torch.float32, + weight_ndim=len(weight_shape), + ) + out_qparams = ( + create_per_axis_qparams( + builder, + weight_shape[0], + len(input_shape), + axis=1, + vary_values=True, + ) + if per_channel_activations + else create_per_tensor_qparams(builder) + ) + + # Use edge op directly instead of ATen op + conv_result = builder.call_operator( + op=exir_ops.edge.fused_quant.convolution.default, + args=( + inp_placeholder, + weight_placeholder, + bias_placeholder, + *inp_qparams, + *weight_qparams, + None, + None, + torch.uint8, + 0, + 0, + *out_qparams, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ), + ) + + builder.output([conv_result]) + program = builder.get_program() + + if bias is not None: + example_inputs = (inp, weight, bias) + else: + example_inputs = (inp, weight) + + return program, example_inputs + + +def _build_fused_quant_max_pool2d_with_indices_graph( + input_shape: tuple[int, ...], + pool_args: tuple[Any, ...], + per_channel_qparams: bool = False, +) -> tuple[ExportedProgram, torch.Tensor]: + inp = torch.randn(*input_shape).to(torch.int8) + builder = ProgramBuilder() + inp_placeholder = builder.placeholder("inp", inp) + inp_qparams = ( + create_per_axis_qparams( + builder, + input_shape[1], + len(input_shape), + axis=1, + dtype=torch.float32, + vary_values=True, + ) + if per_channel_qparams + else create_per_tensor_qparams(builder, dtype=torch.float32) + ) + out_qparams = ( + create_per_axis_qparams( + builder, + input_shape[1], + len(input_shape), + axis=1, + vary_values=True, + ) + if per_channel_qparams + else create_per_tensor_qparams(builder) + ) + + pool = builder.call_operator( + op=exir_ops.edge.fused_quant.max_pool2d_with_indices.default, + args=( + inp_placeholder, + *inp_qparams, + *out_qparams, + *pool_args, + ), + ) + values = builder.call_getitem(pool, 0) + indices = builder.call_getitem(pool, 1) + builder.output([values, indices]) + return builder.get_program(), inp + + +def _build_fused_quant_avg_pool2d_graph( + input_shape: tuple[int, ...], + pool_args: tuple[Any, ...], + per_channel_qparams: bool = False, +) -> tuple[ExportedProgram, torch.Tensor]: + inp = torch.randn(*input_shape).to(torch.int8) + builder = ProgramBuilder() + inp_placeholder = builder.placeholder("inp", inp) + inp_qparams = ( + create_per_axis_qparams( + builder, + input_shape[1], + len(input_shape), + axis=1, + dtype=torch.float32, + vary_values=True, + ) + if per_channel_qparams + else create_per_tensor_qparams(builder, dtype=torch.float32) + ) + out_qparams = ( + create_per_axis_qparams( + builder, + input_shape[1], + len(input_shape), + axis=1, + vary_values=True, + ) + if per_channel_qparams + else create_per_tensor_qparams(builder) + ) + pool = builder.call_operator( + op=exir_ops.edge.fused_quant.avg_pool2d.default, + args=(inp_placeholder, *inp_qparams, *out_qparams, *pool_args), + ) + builder.output([pool]) + return builder.get_program(), inp + + +class ToChannelsLastTest(unittest.TestCase): + """Tests for the ToChannelsLast optimization pass.""" + + def assert_bit_exact(self, before: object, after: object) -> None: + torch.testing.assert_close(before, after, rtol=0, atol=0) + + def test_pointwise_conv_with_singleton_weight_dims_is_bit_exact(self) -> None: + indices = torch.arange(64) + inp_vector = ((indices * 37 + 17) % 256).to(torch.uint8) + weight_vector = ((indices * 53 + 29) % 256 - 128).to(torch.int8) + inp = inp_vector.reshape(1, 64, 1, 1).expand(1, 64, 7, 46).contiguous() + weight = weight_vector.reshape(1, 64, 1, 1).expand(16, 64, 1, 1).contiguous() + weight_ohwi = weight.permute(0, 2, 3, 1).clone( + memory_format=torch.contiguous_format + ) + roundtripped_weight = weight_ohwi.permute(0, 3, 1, 2).contiguous() + self.assertTrue(roundtripped_weight.is_contiguous()) + self.assertEqual(roundtripped_weight.stride(), (64, 1, 64, 64)) + self.assertEqual( + roundtripped_weight.clone(memory_format=torch.contiguous_format).stride(), + (64, 1, 1, 1), + ) + bias = torch.full((16,), -0.28) + inp_scale = torch.tensor(0.01256) + inp_zero_point = torch.tensor(124) + weight_scales = torch.full((16, 1, 1, 1), 0.00965) + weight_zero_points = torch.zeros((16, 1, 1, 1), dtype=torch.int64) + inp_qparams = (inp_scale, inp_zero_point, torch.float32, 0, 255) + weight_qparams = ( + weight_scales, + weight_zero_points, + torch.float32, + -128, + 127, + ) + bias_qparams = (None, None, torch.float32, -2147483648, 2147483647) + out_qparams = (None, None, torch.float32, -2147483648, 2147483647) + conv_args = ([1, 1], [0, 0], [1, 1], False, [0, 0], 1) + args = ( + inp, + weight, + bias, + *inp_qparams, + *weight_qparams, + *bias_qparams, + *out_qparams, + *conv_args, + ) + + builder = ProgramBuilder() + graph_args: list[object] = [] + runtime_inputs: list[torch.Tensor] = [] + for index, value in enumerate(args): + if isinstance(value, torch.Tensor): + graph_args.append(builder.placeholder(f"arg_{index}", value)) + runtime_inputs.append(value) + else: + graph_args.append(value) + conv = builder.call_operator( + exir_ops.edge.fused_quant.convolution.default, + args=tuple(graph_args), + ) + builder.output([conv]) + program = builder.get_program() + [baseline] = program.module()(*runtime_inputs) + ToChannelsLast().call(program.graph_module) + [converted] = program.module()(*runtime_inputs) + self.assert_bit_exact(baseline, converted) + + @parameterized.expand( + [ + # Conv1d with and without bias + ( + "conv1d_with_bias", + (1, 4, 8), + (2, 4, 3), + (2,), + _CONVOLUTION1D_ARGS, + ), + ( + "conv1d_no_bias", + (1, 4, 8), + (2, 4, 3), + None, + _CONVOLUTION1D_ARGS, + ), + # Conv2d with and without bias + ( + "conv2d_with_bias", + (1, 4, 8, 8), + (2, 4, 3, 3), + (2,), + _CONVOLUTION2D_ARGS, + ), + ( + "conv2d_no_bias", + (1, 4, 8, 8), + (2, 4, 3, 3), + None, + _CONVOLUTION2D_ARGS, + ), + # Conv3d with and without bias + ( + "conv3d_with_bias", + (1, 4, 8, 8, 8), + (2, 4, 3, 3, 3), + (2,), + _CONVOLUTION3D_ARGS, + ), + ( + "conv3d_no_bias", + (1, 4, 8, 8, 8), + (2, 4, 3, 3, 3), + None, + _CONVOLUTION3D_ARGS, + ), + ] + ) + def test_conv_to_channels_last( + self, + _name: str, + input_shape: tuple[int, ...], + weight_shape: tuple[int, ...], + bias_shape: tuple[int, ...] | None, + conv_args: tuple[Any, ...], + ) -> None: + """ + Test that ToChannelsLast correctly transforms fused_quant.convolution + to fused_quant.convolution_channels_last with appropriate permutes. + + The expected pattern after optimization is: + permute(NCHW->NHWC) -> permute(OIHW->OHWI) -> + fused_quant.convolution_channels_last -> permute(NHWC->NCHW) + """ + program, example_inputs = _build_fused_quant_conv_graph( + input_shape, weight_shape, bias_shape, conv_args + ) + before = program.module()(*example_inputs) + + # Apply the optimization pass + result = ToChannelsLast().call(program.graph_module) + after = program.module()(*example_inputs) + + self.assertTrue(result.modified) + self.assert_bit_exact(before, after) + + # Verify the graph structure + graph = result.graph_module.graph + + # We should have: + # 1. permute_copy (input NCHW -> NHWC) + # 2. permute_copy (weight OIHW -> OHWI) + # 3. permute_copy (weight scale OIHW -> OHWI) + # 4. permute_copy (weight zero point OIHW -> OHWI) + # 5. fused_quant.convolution_channels_last + # 6. permute_copy (output NHWC -> NCHW) + permute_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + conv_channels_last_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution_channels_last.default, + ) + + self.assertEqual( + len(permute_nodes), 5, "Expected 5 tensor/qparam permute nodes" + ) + self.assertEqual( + len(conv_channels_last_nodes), + 1, + "Expected 1 convolution_channels_last node", + ) + + # Verify no original convolution nodes remain + conv_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution.default, + ) + self.assertEqual( + len(conv_nodes), 0, "Original convolution node should be removed" + ) + + def test_conv_to_channels_last_permutes_activation_qparams(self) -> None: + program, example_inputs = _build_fused_quant_conv_graph( + (1, 4, 8, 8), + (6, 4, 3, 3), + None, + _CONVOLUTION2D_ARGS, + per_channel_activations=True, + ) + before = program.module()(*example_inputs) + + result = ToChannelsLast().call(program.graph_module) + self.assert_bit_exact(before, program.module()(*example_inputs)) + graph = result.graph_module.graph + [conv] = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution_channels_last.default, + ) + self.assertEqual( + tuple(get_arg(conv, "inp_scale", torch.fx.Node).meta["val"].shape), + (1, 1, 1, 4), + ) + self.assertEqual( + tuple(get_arg(conv, "weight_scale", torch.fx.Node).meta["val"].shape), + (6, 1, 1, 1), + ) + self.assertEqual( + tuple(get_arg(conv, "out_scale", torch.fx.Node).meta["val"].shape), + (1, 1, 1, 6), + ) + + def test_conv_to_channels_last_preserves_transposed_arguments(self) -> None: + program, example_inputs = _build_fused_quant_conv_graph( + (1, 2, 4, 4), + (2, 3, 3, 3), + None, + ([2, 2], [1, 1], [1, 1], True, [1, 1], 1), + ) + before = program.module()(*example_inputs) + + result = ToChannelsLast().call(program.graph_module) + + self.assertTrue(result.modified) + self.assert_bit_exact(before, program.module()(*example_inputs)) + [conv] = result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution_channels_last.default, + ) + self.assertTrue(get_arg(conv, "transposed", bool)) + self.assertEqual(get_arg(conv, "output_padding", list[int]), [1, 1]) + + def test_max_pool2d_with_indices_to_channels_last(self) -> None: + pool_args: tuple[Any, ...] = ( + [3, 3], + [2, 2], + [1, 1], + [1, 1], + True, + ) + program, inp = _build_fused_quant_max_pool2d_with_indices_graph( + (1, 4, 8, 8), pool_args + ) + before = program.module()(inp) + + result = ToChannelsLast().call(program.graph_module) + self.assertTrue(result.modified) + self.assert_bit_exact(before, program.module()(inp)) + graph = result.graph_module.graph + + permute_nodes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + self.assertEqual(len(permute_nodes), 3) + self.assertEqual(permute_nodes[0].args[1], [0, 2, 3, 1]) + self.assertEqual(permute_nodes[1].args[1], [0, 3, 1, 2]) + self.assertEqual(permute_nodes[2].args[1], [0, 3, 1, 2]) + + pool_nodes = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.max_pool2d_with_indices_channels_last.default, + ) + self.assertEqual(len(pool_nodes), 1) + pool = pool_nodes[0] + self.assertEqual(get_arg(pool, "kernel_size"), pool_args[0]) + self.assertEqual(get_arg(pool, "stride"), pool_args[1]) + self.assertEqual(get_arg(pool, "padding"), pool_args[2]) + self.assertEqual(get_arg(pool, "dilation"), pool_args[3]) + self.assertEqual(get_arg(pool, "ceil_mode"), pool_args[4]) + self.assertEqual( + [tuple(value.shape) for value in pool.meta["val"]], + [(1, 5, 5, 4), (1, 5, 5, 4)], + ) + getitems = graph.find_nodes(op="call_function", target=operator.getitem) + self.assertEqual({node.args[1] for node in getitems}, {0, 1}) + + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.max_pool2d_with_indices.default, + ) + ), + 0, + ) + + def test_max_pool2d_with_indices_permutes_qparams(self) -> None: + program, inp = _build_fused_quant_max_pool2d_with_indices_graph( + (1, 4, 8, 8), + ([3, 3], [2, 2], [1, 1], [1, 1], True), + per_channel_qparams=True, + ) + before = program.module()(inp) + + result = ToChannelsLast().call(program.graph_module) + self.assert_bit_exact(before, program.module()(inp)) + graph = result.graph_module.graph + [pool] = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.max_pool2d_with_indices_channels_last.default, + ) + self.assertEqual( + tuple(get_arg(pool, "inp_scale", torch.fx.Node).meta["val"].shape), + (1, 1, 1, 4), + ) + self.assertEqual( + tuple(get_arg(pool, "out_scale", torch.fx.Node).meta["val"].shape), + (1, 1, 1, 4), + ) + + def test_avg_pool2d_to_channels_last(self) -> None: + pool_args: tuple[Any, ...] = ( + [3, 3], + [], + [1, 1], + False, + False, + None, + ) + program, inp = _build_fused_quant_avg_pool2d_graph((1, 4, 8, 8), pool_args) + before = program.module()(inp) + result = ToChannelsLast().call(program.graph_module) + self.assert_bit_exact(before, program.module()(inp)) + graph = result.graph_module.graph + + permutes = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + self.assertEqual( + [node.args[1] for node in permutes], [[0, 2, 3, 1], [0, 3, 1, 2]] + ) + [pool] = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + ) + self.assertEqual(get_arg(pool, "kernel_size"), pool_args[0]) + self.assertEqual(get_arg(pool, "stride"), pool_args[1]) + self.assertEqual(get_arg(pool, "padding"), pool_args[2]) + self.assertEqual(get_arg(pool, "ceil_mode"), pool_args[3]) + self.assertEqual(get_arg(pool, "count_include_pad"), pool_args[4]) + self.assertEqual(get_arg(pool, "divisor_override"), pool_args[5]) + self.assertEqual(tuple(pool.meta["val"].shape), (1, 3, 3, 4)) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d.default, + ) + ), + 0, + ) + + def test_avg_pool2d_permutes_qparams(self) -> None: + program, inp = _build_fused_quant_avg_pool2d_graph( + (1, 4, 8, 8), + ([3, 3], [2, 2], [1, 1], False, True, None), + per_channel_qparams=True, + ) + before = program.module()(inp) + result = ToChannelsLast().call(program.graph_module) + self.assert_bit_exact(before, program.module()(inp)) + graph = result.graph_module.graph + [pool] = graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.avg_pool2d_channels_last.default, + ) + self.assertEqual( + tuple(get_arg(pool, "inp_scale", torch.fx.Node).meta["val"].shape), + (1, 1, 1, 4), + ) + self.assertEqual( + tuple(get_arg(pool, "out_scale", torch.fx.Node).meta["val"].shape), + (1, 1, 1, 4), + ) diff --git a/backends/fused_quant/test/test_to_convolution.py b/backends/fused_quant/test/test_to_convolution.py new file mode 100644 index 00000000000..30701c08771 --- /dev/null +++ b/backends/fused_quant/test/test_to_convolution.py @@ -0,0 +1,115 @@ +# 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. + +from __future__ import annotations + +import unittest +from typing import Any + +import executorch.backends.fused_quant.ops # noqa: F401 +import torch +from executorch.backends.fused_quant.optimization_passes.to_convolution import ( + ToConvolution, +) +from executorch.backends.fused_quant.test.helpers import create_per_tensor_qparams +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from parameterized import parameterized + + +class ToConvolutionTest(unittest.TestCase): + @parameterized.expand( + [ + ( + "conv1d", + exir_ops.edge.fused_quant.conv1d.default, + (1, 2, 8), + (3, 2, 3), + [2], + [1], + [1], + ), + ( + "conv2d", + exir_ops.edge.fused_quant.conv2d.default, + (1, 2, 8, 8), + (3, 2, 3, 3), + [2, 2], + [1, 1], + [1, 1], + ), + ( + "conv3d", + exir_ops.edge.fused_quant.conv3d.default, + (1, 2, 8, 8, 8), + (3, 2, 3, 3, 3), + [2, 2, 2], + [1, 1, 1], + [1, 1, 1], + ), + ] + ) + def test_converts_to_convolution( + self, + _name: str, + target: Any, + input_shape: tuple[int, ...], + weight_shape: tuple[int, ...], + stride: list[int], + padding: list[int], + dilation: list[int], + ) -> None: + inp = torch.randint(-8, 8, input_shape, dtype=torch.int8) + weight = torch.randint(-8, 8, weight_shape, dtype=torch.int8) + builder = ProgramBuilder() + inp_node = builder.placeholder("inp", inp) + weight_node = builder.placeholder("weight", weight) + inp_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + weight_qparams = create_per_tensor_qparams(builder, dtype=torch.float32) + out_qparams = create_per_tensor_qparams(builder) + conv = builder.call_operator( + target, + args=( + inp_node, + weight_node, + None, + *inp_qparams, + *weight_qparams, + None, + None, + torch.float32, + -128, + 127, + *out_qparams, + stride, + padding, + dilation, + 1, + ), + ) + builder.output([conv]) + program = builder.get_program() + expected = program.module()(inp, weight) + + result = ToConvolution().call(program.graph_module) + + self.assertTrue(result.modified) + torch.testing.assert_close( + expected, program.module()(inp, weight), rtol=0, atol=0 + ) + self.assertEqual( + result.graph_module.graph.find_nodes(op="call_function", target=target), + [], + ) + [convolution] = result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.fused_quant.convolution.default, + ) + self.assertFalse(get_arg(convolution, "transposed", bool)) + self.assertEqual( + get_arg(convolution, "output_padding", list[int]), [0] * len(stride) + ) 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..abdf51134e9 --- /dev/null +++ b/backends/transforms/advance_quantize_op_above_def_in_branch.py @@ -0,0 +1,195 @@ +# 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 math import prod +from typing import DefaultDict, List, Tuple + +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 the quant descendents. Cluster them into sets + # that have the same scale, zero_point, and dtype. We use quant_dict + # for the clustering + quant_dict: DefaultDict[Tuple, int] = defaultdict(int) + for quant_node in descendent_quant_ops: + quant_dict[quant_node.args[1:]] += 1 + rep_args = sorted(quant_dict.keys(), key=lambda x: x[1]).pop() + + # 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..8e286f7e886 --- /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 6cd86402f7e..5da602f0eaa 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 @@ -758,6 +759,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..ae0bf0d256e --- /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 46430621a12..d3baf583bbd 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -82,6 +82,18 @@ def define_common_targets(): ], ) + runtime.python_test( + name = "test_transforms_utils", + srcs = ["test/test_transforms_utils.py"], + deps = [ + ":utils", + "//caffe2:torch", + "//executorch/backends/test:graph_builder", + "//executorch/backends/test:program_builder", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_test( name = "test_fuse_rms_norm", srcs = ["test/test_fuse_rms_norm.py"], @@ -243,6 +255,7 @@ def define_common_targets(): srcs = ["utils.py"], visibility = [ "//executorch/backends/...", + "@EXECUTORCH_CLIENTS", ], deps = [ "//caffe2:torch", @@ -868,3 +881,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 c656e4e13a9..55d04e0bab8 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1678,6 +1678,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) diff --git a/backends/transforms/test/test_transforms_utils.py b/backends/transforms/test/test_transforms_utils.py new file mode 100644 index 00000000000..464ba06ccc2 --- /dev/null +++ b/backends/transforms/test/test_transforms_utils.py @@ -0,0 +1,245 @@ +# 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. + + +import unittest + +import torch +from executorch.backends.test.graph_builder import GraphBuilder +from executorch.backends.test.program_builder import ProgramBuilder +from executorch.backends.transforms.utils import ( + add_constant, + compute_meta_val, + get_constant, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensor +from torch.export import ExportedProgram +from torch.export.graph_signature import ( + ExportGraphSignature, + InputKind, + InputSpec, + TensorArgument, +) + + +def _build_ep_with_divergent_const( + node_name: str, fqn: str, value: torch.Tensor +) -> ExportedProgram: + """Build an ExportedProgram with a lifted constant whose node name != fqn. + + This mimics the state export produces for a fusion-created constant: the + placeholder node carries a ``c_``-prefixed name (e.g. ``c__scale_16``) while + its fqn into ``ep.constants`` is un-prefixed (``_scale_16``). Crucially the + fqn is NOT itself a graph node name, so a later ``graph.placeholder`` can + generate a node named exactly that fqn. + """ + builder = ProgramBuilder() + x = builder.placeholder("x", torch.randn(value.shape)) + c = builder.placeholder(node_name, value, input_kind=InputKind.CONSTANT_TENSOR) + out = builder.call_operator(op=exir_ops.edge.aten.add.Tensor, args=(x, c)) + builder.output([out]) + ep = builder.get_program() + + # Rewrite the constant's spec so its target (fqn) diverges from the node name, + # and move its backing value under the new fqn. + sig = ep.graph_signature + new_specs = [ + ( + InputSpec( + kind=s.kind, + arg=TensorArgument(name=node_name), + target=fqn, + persistent=True, + ) + if getattr(s.arg, "name", None) == node_name + else s + ) + for s in sig.input_specs + ] + ep._graph_signature = ExportGraphSignature( + input_specs=new_specs, output_specs=list(sig.output_specs) + ) + ep.constants[fqn] = ep.constants.pop(node_name) + return ep + + +class AddConstantTest(unittest.TestCase): + """Regression tests for add_constant's fqn uniquification. + + add_constant uniquifies the placeholder *node name* against the graph, but + the spec target (the fqn into ep.constants) lives in a separate namespace: + export gives lifted constants a ``c_``-prefixed node name over an un-prefixed + fqn. A node name generated here can therefore equal an existing constant's + fqn, so the derived target must be uniquified against the fqn namespace -- + otherwise two specs share one ep.constants entry and constant_prop, erasing + one dead placeholder, pops the value the other still needs (a KeyError in + get_lifted_tensor_constant, seen via QuantAbsorptionPass minting many + _scale/_zero_point constants until one collides with a fusion fqn). + """ + + def _find_node(self, ep: ExportedProgram, name: str) -> torch.fx.Node: + return next(n for n in ep.graph.nodes if n.name == name) + + def _spec_for(self, ep: ExportedProgram, node_name: str) -> InputSpec: + return next( + s + for s in ep.graph_signature.input_specs + if getattr(s.arg, "name", None) == node_name + ) + + def test_add_constant_uniquifies_target_against_existing_fqn(self) -> None: + orig_val = torch.tensor([1.0, 2.0, 3.0]) + ep = _build_ep_with_divergent_const("c_s", "s", orig_val) + # The fqn "s" already exists, so add_constant must not reuse it. + self.assertIn( + "s", ep.graph_signature.inputs_to_lifted_tensor_constants.values() + ) + + anchor = self._find_node(ep, "c_s") + new_val = torch.tensor([9.0, 9.0, 9.0]) + new_node = add_constant(ep, "s", new_val, anchor, InputKind.CONSTANT_TENSOR) + + # The node name is export-style prefixed for its kind (c_ for a constant), + # keeping the new placeholder in the same namespace as export's own. + self.assertTrue(new_node.name.startswith("c_")) + + # The fqn (spec target) was uniquified away from the existing "s". + new_spec = self._spec_for(ep, new_node.name) + self.assertNotEqual(new_spec.target, "s") + + # No two input specs share a target (fqn), so no shared ep.constants entry. + targets = [ + s.target for s in ep.graph_signature.input_specs if s.target is not None + ] + self.assertEqual(len(targets), len(set(targets))) + + # The original constant's value was not overwritten, and both round-trip. + orig_after = get_constant(ep, self._find_node(ep, "c_s")) + assert orig_after is not None + self.assertTrue(torch.equal(orig_after, orig_val)) + new_after = get_constant(ep, new_node) + assert new_after is not None + self.assertTrue(torch.equal(new_after, new_val)) + + def test_add_constant_prefixes_node_name_by_kind(self) -> None: + """add_constant owns the prefixing: node = , fqn bare.""" + ep = _build_ep_with_divergent_const("c_seed", "seed", torch.tensor([0.0])) + anchor = self._find_node(ep, "c_seed") + + # A bare logical name gets the kind's prefix on the node, bare fqn. + node = add_constant( + ep, "myscale", torch.tensor(0.5), anchor, InputKind.CONSTANT_TENSOR + ) + self.assertEqual(node.name, "c_myscale") + self.assertEqual(self._spec_for(ep, node.name).target, "myscale") + + # Passing a name that already carries the kind's prefix is a caller bug. + with self.assertRaises(AssertionError): + add_constant( + ep, "c_other", torch.tensor(0.25), anchor, InputKind.CONSTANT_TENSOR + ) + + def test_repeated_add_constant_never_collides(self) -> None: + """Minting many constants with the same base name keeps every fqn unique. + + Mirrors QuantAbsorptionPass adding a fresh _scale/_zero_point per absorbed + op: each must land on its own ep.constants entry, and none may clobber the + pre-existing constant that shares the base name. + """ + orig_val = torch.tensor([1.0, 2.0, 3.0]) + ep = _build_ep_with_divergent_const("c_s", "s", orig_val) + anchor = self._find_node(ep, "c_s") + + added = [] + for i in range(20): + val = torch.tensor([float(i)]) + node = add_constant(ep, "s", val, anchor, InputKind.CONSTANT_TENSOR) + added.append((node, val)) + + targets = [ + s.target for s in ep.graph_signature.input_specs if s.target is not None + ] + self.assertEqual(len(targets), len(set(targets)), "duplicate fqn targets") + + # The pre-existing constant survived every mint. + orig_after = get_constant(ep, self._find_node(ep, "c_s")) + assert orig_after is not None + self.assertTrue(torch.equal(orig_after, orig_val)) + + for node, val in added: + got = get_constant(ep, node) + assert got is not None + self.assertTrue(torch.equal(got, val)) + + +class ComputeMetaValTest(unittest.TestCase): + """compute_meta_val recomputes a node's meta['val'] by executing it under the + graph's fake mode. Verified against the val GraphBuilder itself computed, for + single-output, multi-output, and literal-only (aten.full) nodes.""" + + def _node(self, gm: torch.fx.GraphModule, target: object) -> torch.fx.Node: + return next( + n for n in gm.graph.nodes if n.op == "call_function" and n.target == target + ) + + def test_single_output_from_fake_inputs(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 3)) + y = builder.call_operator(op=torch.ops.aten.add.Tensor, args=(x, x)) + builder.output([y]) + gm = builder.get_graph_module() + + add_node = self._node(gm, torch.ops.aten.add.Tensor) + expected = add_node.meta["val"] + got = compute_meta_val(add_node) + + self.assertIsInstance(got, FakeTensor) + self.assertEqual(tuple(got.shape), tuple(expected.shape)) + self.assertEqual(got.dtype, expected.dtype) + + def test_multi_output_returns_tuple_of_fakes(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 4)) + w = builder.placeholder("w", torch.randn(4)) + b = builder.placeholder("b", torch.randn(4)) + ln = builder.call_operator( + op=torch.ops.aten.native_layer_norm.default, + args=(x, [4], w, b, 1e-5), + ) + builder.output([builder.call_getitem(ln, 0)]) + gm = builder.get_graph_module() + + ln_node = self._node(gm, torch.ops.aten.native_layer_norm.default) + got = compute_meta_val(ln_node) + + self.assertIsInstance(got, (tuple, list)) + expected = ln_node.meta["val"] + self.assertEqual(len(got), len(expected)) + for g, e in zip(got, expected): + self.assertIsInstance(g, FakeTensor) + self.assertEqual(tuple(g.shape), tuple(e.shape)) + + def test_literal_only_full_is_faked_under_mode(self) -> None: + """aten.full has no tensor inputs, so it yields a fake (not a real) tensor + only because compute_meta_val runs it under the graph's fake mode -- the + case the simpler implementation hinges on.""" + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2)) + full = builder.call_operator( + op=torch.ops.aten.full.default, + args=([4], 0), + kwargs={"dtype": torch.int32}, + ) + builder.output([x, full]) + gm = builder.get_graph_module() + + full_node = self._node(gm, torch.ops.aten.full.default) + got = compute_meta_val(full_node) + + self.assertIsInstance(got, FakeTensor) + self.assertEqual(tuple(got.shape), (4,)) + self.assertEqual(got.dtype, torch.int32) diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index d3c50c692b1..7bff5a946d5 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -10,6 +10,7 @@ import torch from executorch.exir import ExportedProgram from torch._export.utils import ( + _detect_fake_mode_from_gm, get_buffer, get_lifted_tensor_constant, get_param, @@ -17,7 +18,11 @@ is_lifted_tensor_constant, is_param, ) -from torch._subclasses.fake_tensor import FakeTensorConverter +from torch._subclasses.fake_tensor import ( + FakeTensor, + FakeTensorConverter, + FakeTensorMode, +) from torch.export.graph_signature import ( ExportGraphSignature, InputKind, @@ -26,6 +31,16 @@ OutputSpec, TensorArgument, ) +from torch.fx import map_arg + + +def _fake_mode_from_node(node: torch.fx.Node) -> FakeTensorMode: + """Pull the graph's FakeTensorMode off a node's meta['val'], which is a + FakeTensor for single-output nodes and a tuple/list of them otherwise.""" + val = node.meta["val"] + if isinstance(val, (tuple, list)): + val = next(v for v in val if isinstance(v, FakeTensor)) + return val.fake_mode def _get_fake_tensor_mode(graph: torch.fx.Graph, data: torch.Tensor) -> torch.Tensor: @@ -48,15 +63,9 @@ def _get_fake_tensor_mode(graph: torch.fx.Graph, data: torch.Tensor) -> torch.Te "Cannot create fake tensor: graph has no nodes to extract fake_mode from" ) - example_node = nodes[0] - if isinstance( - example_node.meta["val"], (tuple, torch.fx.immutable_collections.immutable_list) - ): - example_fake_tensor = example_node.meta["val"][0] - else: - example_fake_tensor = example_node.meta["val"] - - return FakeTensorConverter().from_real_tensor(example_fake_tensor.fake_mode, t=data) + return FakeTensorConverter().from_real_tensor( + _fake_mode_from_node(nodes[0]), t=data + ) def is_get_attr_node(node: torch.fx.Node) -> bool: @@ -95,6 +104,73 @@ def get_param_tensor( raise RuntimeError(f"unsupported param type, {node.op}.") +def get_constant(ep: ExportedProgram, node: torch.fx.Node) -> Optional[torch.Tensor]: + """Return the tensor backing a lifted constant placeholder, or None. + + Unlike get_param_tensor, non-constant nodes yield None instead of raising, and + parameters are unwrapped to plain tensors. + """ + if node.op != "placeholder" or not is_param_node(ep, node): + return None + tensor = get_param_tensor(ep, node) + return tensor.data if isinstance(tensor, torch.nn.Parameter) else tensor + + +def get_input_kind(ep: ExportedProgram, node: torch.fx.Node) -> Optional[InputKind]: + """Return the InputKind of a constant placeholder node, or None.""" + if node.op != "placeholder": + return None + + sig = ep.graph_signature + if node.name in sig.inputs_to_parameters: + return InputKind.PARAMETER + if node.name in sig.inputs_to_buffers: + return InputKind.BUFFER + if node.name in sig.inputs_to_lifted_tensor_constants: + return InputKind.CONSTANT_TENSOR + return None + + +def get_fqn(ep: ExportedProgram, node: torch.fx.Node) -> Optional[str]: + """Return the fqn backing a constant placeholder, or None. + + The fqn is the key into ``ep.state_dict`` / ``ep.constants`` and the + ``InputSpec.target`` -- the un-prefixed logical name (node ``p_a_b`` -> fqn + ``a.b``), i.e. the inverse of export's ``c_``/``b_``/``p_`` node-name prefixing. + """ + if node.op != "placeholder": + return None + + sig = ep.graph_signature + for mapping in ( + sig.inputs_to_parameters, + sig.inputs_to_buffers, + sig.inputs_to_lifted_tensor_constants, + ): + if node.name in mapping: + return mapping[node.name] + return None + + +def compute_meta_val(node: torch.fx.Node) -> object: + """Compute ``meta['val']`` for a new or mutated call_function node by running + its target on its inputs' fake vals under the graph's fake mode. + + Handles single-output, multi-output, and literal-only ops (e.g. ``aten.full``, + which is faked only because it runs under the fake mode). + """ + gm = node.graph.owning_module + assert gm is not None, "node's graph has no owning module" + fake_mode = _detect_fake_mode_from_gm(gm) + assert fake_mode is not None, "graph has no fake tensor mode" + target = node.target + assert callable(target), f"expected a callable target, got {target}" + args = map_arg(node.args, lambda n: n.meta["val"]) + kwargs = map_arg(node.kwargs, lambda n: n.meta["val"]) + with fake_mode: + return target(*args, **kwargs) + + def set_param_tensor( exp_prog: ExportedProgram, node: torch.fx.Node, data: torch.Tensor ) -> None: @@ -233,6 +309,106 @@ def create_constant_placeholder( return node +# The node-name prefix export gives each lifted input kind. Mirroring it keeps +# add_constant's placeholders indistinguishable from export's own and makes +# node-name uniqueness imply fqn uniqueness (fqn == node name minus prefix). +_INPUT_KIND_PREFIX: dict[InputKind, str] = { + InputKind.PARAMETER: "p_", + InputKind.BUFFER: "b_", + InputKind.CONSTANT_TENSOR: "c_", +} + + +def add_constant( + ep: ExportedProgram, + name: str, + tensor: torch.Tensor, + before_node: torch.fx.Node, + kind: InputKind, +) -> torch.fx.Node: + """Add a new constant placeholder, placed ahead of the user inputs. + + Unlike create_constant_placeholder, the caller does not pick the insertion + point and every call mints a fresh constant (no sharing by name). + + ``name`` must be a logical (unprefixed) base -- add_constant owns the + prefixing. The node name is the export-style ```` (``c_`` + constant, ``b_`` buffer, ``p_`` parameter) and the fqn is that node name with + the prefix stripped, so fx's node-name uniquification also uniquifies the fqn + within a kind. The assert guards a cross-kind fqn clash (the same fqn backing + both ep.constants and ep.state_dict), which would otherwise surface later as a + KeyError in constant_prop's get_lifted_tensor_constant. + + ``before_node`` supplies the fake mode for the new placeholder's meta['val']. + """ + graph = ep.graph_module.graph + placeholders = graph.find_nodes(op="placeholder") + + # Placeholder order and input_specs order must stay in lockstep, with + # constants ahead of user inputs (create_constant_placeholder enforces this). + user_inputs = { + spec.arg.name + for spec in ep.graph_signature.input_specs + if spec.kind == InputKind.USER_INPUT and hasattr(spec.arg, "name") + } + insert_at = len(ep.graph_signature.input_specs) + anchor = None + for index, candidate in enumerate(placeholders): + if candidate.name in user_inputs: + anchor = candidate + insert_at = index + break + + assert ( + kind in _INPUT_KIND_PREFIX + ), f"add_constant supports {set(_INPUT_KIND_PREFIX)}, got {kind}" + prefix = _INPUT_KIND_PREFIX[kind] + assert not name.startswith( + prefix + ), f"add_constant expects a logical (unprefixed) name for {kind}, got {name!r}" + + position = ( + graph.inserting_before(anchor) + if anchor is not None + else graph.inserting_after(placeholders[-1]) + ) + with position: + placeholder = graph.placeholder(prefix + name) + placeholder.target = placeholder.name + fake_mode = _fake_mode_from_node(before_node) + placeholder.meta["val"] = fake_mode.from_tensor(tensor, static_shapes=True) + + fqn = placeholder.name[len(prefix) :] + assert ( + fqn not in ep.constants and fqn not in ep.state_dict + ), f"add_constant: fqn {fqn!r} already backed; prefix invariant violated" + + input_specs = list(ep.graph_signature.input_specs) + input_specs.insert( + insert_at, + InputSpec( + kind=kind, + arg=TensorArgument(name=placeholder.name), + target=fqn, + persistent=True, + ), + ) + ep._graph_signature = ExportGraphSignature( + input_specs=input_specs, + output_specs=list(ep.graph_signature.output_specs), + ) + + match kind: + case InputKind.PARAMETER: + ep.state_dict[fqn] = torch.nn.Parameter(tensor, requires_grad=False) + case InputKind.BUFFER: + ep.state_dict[fqn] = tensor + case _: + ep.constants[fqn] = tensor + + return placeholder + + def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Node): """ Deletes a node of type parameter, buffer, or lifted constant tensor and its related diff --git a/pytest.ini b/pytest.ini index 576ed2b78ce..19ade57c55f 100644 --- a/pytest.ini +++ b/pytest.ini @@ -81,6 +81,7 @@ testpaths = # backends backends/apple/coreml/test + backends/fused_quant/test backends/native/test backends/test/harness/tests backends/test/suite/tests