diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 5a4dde37f17..3273923b1e1 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -116,6 +116,7 @@ from .deduplicate_get_attr_pass import DeduplicateGetAttrPass # noqa from .ensure_unique_output_nodes_pass import EnsureUniqueOutputNodesPass # noqa from .exir_to_tosa_pass import ExirToTosaPass # noqa +from .fold_dyt_affine_into_conv_pass import FoldDyTAffineIntoConvPass # noqa from .fold_dyt_alpha_into_lut_pass import FoldDyTAlphaIntoLUTPass # noqa from .fold_qdq_with_annotated_qparams_pass import ( # noqa FoldAndAnnotateQParamsPass, diff --git a/backends/arm/_passes/fold_dyt_affine_into_conv_pass.py b/backends/arm/_passes/fold_dyt_affine_into_conv_pass.py new file mode 100644 index 00000000000..8eccdae32ed --- /dev/null +++ b/backends/arm/_passes/fold_dyt_affine_into_conv_pass.py @@ -0,0 +1,890 @@ +# 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 +"""Fold exact quantized DyT affine maps into following convolutions.""" + +from __future__ import annotations + +from copy import copy +from dataclasses import dataclass +from typing import cast, Set, Type + +import torch +from executorch.backends.arm._passes.arm_pass import ArmPass +from executorch.backends.arm._passes.arm_pass_utils import ( + get_constant_placeholder_kind, + get_param_tensor, + is_persistent_buffer, +) +from executorch.backends.arm._passes.fold_dyt_alpha_into_lut_pass import ( + _apply_tosa_rescale, + _RescaleParams, +) +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.transforms.utils import ( + create_constant_placeholder, + delete_constant_placeholder, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import Graph, GraphModule, Node + + +@dataclass(frozen=True) +class _Operand: + source: Node + rescale: _RescaleParams + constant: torch.Tensor | None + view_shape: tuple[int, ...] | None + + +@dataclass(frozen=True) +class _DyTAffineMatch: + table: Node + table_values: torch.Tensor + table_qargs: QuantArgs + gamma_activation_rescale: _RescaleParams + gamma: Node + gamma_codes: torch.Tensor + gamma_rescale: _RescaleParams + gamma_output: Node + gamma_output_rescale: _RescaleParams + add_activation_rescale: _RescaleParams + beta: Node + beta_codes: torch.Tensor + beta_rescale: _RescaleParams + add_output: Node + add_output_rescale: _RescaleParams + conv: Node + layout_chain: tuple[Node, ...] + + +@dataclass(frozen=True) +class _AddChainMatch: + output: Node + output_rescale: _RescaleParams + activation_operand: _Operand + beta_operand: _Operand + layout_chain: tuple[Node, ...] + + +@dataclass(frozen=True) +class _GammaChainMatch: + output: Node + output_rescale: _RescaleParams + table_operand: _Operand + gamma_operand: _Operand + + +@dataclass(frozen=True) +class _AffineCodes: + """Validated INT8 constants backing one DyT affine map.""" + + table_values: torch.Tensor + table_qargs: QuantArgs + gamma_codes: torch.Tensor + beta_codes: torch.Tensor + + +@dataclass(frozen=True) +class _ConvOperands: + """Weight/bias placeholders and quantization args of a foldable conv.""" + + weight_node: Node + bias_node: Node + weight: torch.Tensor + bias: torch.Tensor + input_qparams: dict[int, QuantArgs] + activation_qargs: QuantArgs + weight_qargs: QuantArgs + + +@dataclass(frozen=True) +class _ConvConstants: + """A conv whose constants are shape-compatible with the affine map.""" + + operands: _ConvOperands + groups: int + out_channels: int + in_channels_per_group: int + weight_zero_points: torch.Tensor + + +class FoldDyTAffineIntoConvPass(ArmPass): + """Fold exact integer DyT gamma/beta maps into a following convolution. + + INT8 requantization makes generic floating-point affine folding inexact. + + The pass evaluates the 256-entry TOSA TABLE and rewrites only exact affine + maps. + + Position-dependent padding keeps beta; gamma folds only for an exact + identity. + + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + + _VIEW_TARGETS: Set[object] = { + exir_ops.edge.aten.view_copy.default, + } + _LAYOUT_TARGETS: Set[object] = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.slice_copy.Tensor, + } + + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__() + self.exported_program = exported_program + + @staticmethod + def _single_qargs(node: Node, key: str) -> QuantArgs | None: + qparams = cast(dict[int, QuantArgs], node.meta.get(key, {})) + if len(qparams) != 1: + return None + qargs = next(iter(qparams.values())) + if qargs.per_channel: + return None + return qargs + + @staticmethod + def _tensor_shape(node: Node) -> tuple[int, ...] | None: + value = node.meta.get("val") + if not isinstance(value, torch.Tensor) or not all( + type(dim) is int for dim in value.shape + ): + return None + return cast(tuple[int, ...], tuple(value.shape)) + + def _constant(self, node: Node) -> torch.Tensor | None: + try: + return get_param_tensor(self.exported_program, node) + except RuntimeError: + return None + + def _unwrap_views(self, node: Node) -> tuple[Node, tuple[int, ...] | None] | None: + view_shape = None + while node.target in self._VIEW_TARGETS: + if ( + len(node.args) < 2 + or not isinstance(node.args[0], Node) + or len(node.users) != 1 + ): + return None + shape = node.args[1] + if ( + view_shape is not None + or not isinstance(shape, (list, tuple)) + or not all(type(value) is int for value in shape) + ): + return None + view_shape = tuple(shape) + node = node.args[0] + return node, view_shape + + def _operand(self, node: Node) -> _Operand | None: + unwrapped = self._unwrap_views(node) + if unwrapped is None: + return None + rescale_node, view_shape = unwrapped + rescale = _RescaleParams.from_node(rescale_node) + if ( + rescale is None + or rescale.output_dtype != torch.int32 + or not isinstance(rescale_node.args[0], Node) + ): + return None + source = rescale_node.args[0] + return _Operand( + source=source, + rescale=rescale, + constant=self._constant(source), + view_shape=view_shape, + ) + + def _binary_operands(self, node: Node) -> tuple[_Operand, _Operand] | None: + if len(node.args) < 2: + return None + lhs, rhs = node.args[:2] + if not isinstance(lhs, Node) or not isinstance(rhs, Node): + return None + lhs_operand = self._operand(lhs) + rhs_operand = self._operand(rhs) + if lhs_operand is None or rhs_operand is None: + return None + return lhs_operand, rhs_operand + + def _trace_layout_source(self, node: Node) -> tuple[Node, tuple[Node, ...]] | None: + """Walk back through layout ops to the affine site. + + Exactly one ``permute_copy`` is required and pinned to the NHWC->NCHW + dim order this pass is written against. ``slice_copy`` is deliberately not inspected: a slice + that changes which channels the conv consumes leaves the site's + per-channel slope/offset count disagreeing with the conv weight's input + channels, and ``_fold_conv_constants`` refuses the fold on that + mismatch. Slices on the batch or spatial dims cannot invalidate a + per-channel affine map. Both paths are pinned by + ``test_channel_narrowing_slice_is_rejected`` and + ``test_identity_affine_behind_channel_slice_leaves_conv_constants``. + + """ + chain = [] + permute_count = 0 + while node.target in self._LAYOUT_TARGETS: + if len(node.args) == 0 or not isinstance(node.args[0], Node): + return None + if node.target == exir_ops.edge.aten.permute_copy.default: + dims = node.args[1] if len(node.args) > 1 else None + if ( + permute_count != 0 + or not isinstance(dims, (list, tuple)) + or tuple(dims) != (0, 3, 1, 2) + ): + return None + permute_count += 1 + chain.append(node) + node = node.args[0] + if permute_count != 1: + return None + return node, tuple(chain) + + def _match_add_chain(self, conv: Node) -> _AddChainMatch | None: + if ( + conv.op != "call_function" + or conv.target != exir_ops.edge.aten.convolution.default + or len(conv.args) < 9 + or bool(conv.args[6]) + or not isinstance(conv.args[0], Node) + ): + return None + + traced = self._trace_layout_source(conv.args[0]) + if traced is None: + return None + add_output, layout_chain = traced + add_output_rescale = _RescaleParams.from_node(add_output) + if ( + add_output_rescale is None + or add_output_rescale.output_dtype != torch.int8 + or not isinstance(add_output.args[0], Node) + ): + return None + + add = add_output.args[0] + if add.target != exir_ops.edge.aten.add.Tensor or len(add.users) != 1: + return None + add_operands = self._binary_operands(add) + if add_operands is None: + return None + beta_operands = [ + operand for operand in add_operands if operand.constant is not None + ] + if len(beta_operands) != 1: + return None + beta_operand = beta_operands[0] + activation_operand = next( + operand for operand in add_operands if operand is not beta_operand + ) + return _AddChainMatch( + output=add_output, + output_rescale=add_output_rescale, + activation_operand=activation_operand, + beta_operand=beta_operand, + layout_chain=layout_chain, + ) + + def _match_gamma_chain(self, gamma_output: Node) -> _GammaChainMatch | None: + gamma_output_rescale = _RescaleParams.from_node(gamma_output) + if ( + gamma_output_rescale is None + or gamma_output_rescale.output_dtype != torch.int8 + or len(gamma_output.users) != 1 + or not isinstance(gamma_output.args[0], Node) + ): + return None + + mul = gamma_output.args[0] + if mul.target != exir_ops.edge.aten.mul.Tensor or len(mul.users) != 1: + return None + mul_operands = self._binary_operands(mul) + if mul_operands is None: + return None + gamma_operands = [ + operand for operand in mul_operands if operand.constant is not None + ] + if len(gamma_operands) != 1: + return None + gamma_operand = gamma_operands[0] + table_operand = next( + operand for operand in mul_operands if operand is not gamma_operand + ) + return _GammaChainMatch( + output=gamma_output, + output_rescale=gamma_output_rescale, + table_operand=table_operand, + gamma_operand=gamma_operand, + ) + + @staticmethod + def _table_values_node(table: Node) -> Node | None: + """Return the node holding a TOSA TABLE's lookup values. + + Returns ``None`` when ``table`` is not a TABLE or does not carry its + values as a node, which also narrows the operand for the caller. + + """ + if ( + table.target != exir_ops.backend.tosa.TABLE.default + or len(table.args) < 2 + or not isinstance(table.args[1], Node) + ): + return None + return table.args[1] + + def _affine_codes( + self, + table: Node, + table_values_node: Node, + gamma_match: _GammaChainMatch, + add_match: _AddChainMatch, + ) -> _AffineCodes | None: + """Collect the INT8 table, gamma and beta constants of an affine map. + + Returns ``None`` unless every constant is present, INT8, and sized as + the fold requires: a 256-entry table and matching gamma/beta lengths. + + """ + table_values = self._constant(table_values_node) + table_qargs = self._single_qargs(table, "output_qparams") + gamma_codes = gamma_match.gamma_operand.constant + beta_codes = add_match.beta_operand.constant + if ( + table_values is None + or table_values.dtype != torch.int8 + or table_values.numel() != 256 + or table_qargs is None + or table_qargs.dtype != torch.int8 + or gamma_codes is None + or gamma_codes.dtype != torch.int8 + or beta_codes is None + or beta_codes.dtype != torch.int8 + or gamma_codes.numel() != beta_codes.numel() + ): + return None + return _AffineCodes( + table_values=table_values, + table_qargs=table_qargs, + gamma_codes=gamma_codes, + beta_codes=beta_codes, + ) + + @staticmethod + def _views_are_channel_broadcasts( + gamma_match: _GammaChainMatch, + add_match: _AddChainMatch, + channels: int, + ) -> bool: + """Return True when only gamma/beta carry the channel broadcast view.""" + channel_view_shape = (1, 1, 1, channels) + return ( + gamma_match.table_operand.view_shape is None + and add_match.activation_operand.view_shape is None + and gamma_match.gamma_operand.view_shape == channel_view_shape + and add_match.beta_operand.view_shape == channel_view_shape + ) + + @staticmethod + def _table_shape_matches( + table_shape: tuple[int, ...] | None, + add_output_shape: tuple[int, ...] | None, + channels: int, + ) -> bool: + """Return True when the TABLE is NHWC and channel-aligned with gamma.""" + return ( + table_shape is not None + and table_shape == add_output_shape + and len(table_shape) == 4 + and table_shape[-1] == channels + ) + + def _match(self, conv: Node) -> _DyTAffineMatch | None: + add_match = self._match_add_chain(conv) + if add_match is None: + return None + gamma_match = self._match_gamma_chain(add_match.activation_operand.source) + if gamma_match is None: + return None + + table = gamma_match.table_operand.source + table_values_node = self._table_values_node(table) + if table_values_node is None: + return None + + codes = self._affine_codes(table, table_values_node, gamma_match, add_match) + if codes is None: + return None + channels = codes.gamma_codes.numel() + if not self._views_are_channel_broadcasts(gamma_match, add_match, channels): + return None + if not self._table_shape_matches( + self._tensor_shape(table), + self._tensor_shape(add_match.output), + channels, + ): + return None + + return _DyTAffineMatch( + table=table, + table_values=codes.table_values, + table_qargs=codes.table_qargs, + gamma_activation_rescale=gamma_match.table_operand.rescale, + gamma=gamma_match.gamma_operand.source, + gamma_codes=codes.gamma_codes, + gamma_rescale=gamma_match.gamma_operand.rescale, + gamma_output=gamma_match.output, + gamma_output_rescale=gamma_match.output_rescale, + add_activation_rescale=add_match.activation_operand.rescale, + beta=add_match.beta_operand.source, + beta_codes=codes.beta_codes, + beta_rescale=add_match.beta_operand.rescale, + add_output=add_match.output, + add_output_rescale=add_match.output_rescale, + conv=conv, + layout_chain=add_match.layout_chain, + ) + + @staticmethod + def _checked_int32(values: torch.Tensor) -> torch.Tensor | None: + limits = torch.iinfo(torch.int32) + if values.numel() and ( + int(values.min()) < limits.min or int(values.max()) > limits.max + ): + return None + return values.to(torch.int32) + + def _gamma_outputs(self, match: _DyTAffineMatch) -> torch.Tensor | None: + table_codes = match.table_values.reshape(-1, 1) + activation_i32 = _apply_tosa_rescale( + table_codes, + match.gamma_activation_rescale, + ) + gamma_i32 = _apply_tosa_rescale( + match.gamma_codes.reshape(1, -1), + match.gamma_rescale, + ) + product = self._checked_int32( + activation_i32.to(torch.int64) * gamma_i32.to(torch.int64) + ) + if product is None: + return None + return _apply_tosa_rescale(product, match.gamma_output_rescale) + + def _affine_outputs( + self, + match: _DyTAffineMatch, + gamma_outputs: torch.Tensor, + ) -> torch.Tensor | None: + activation_i32 = _apply_tosa_rescale( + gamma_outputs, + match.add_activation_rescale, + ) + beta_i32 = _apply_tosa_rescale( + match.beta_codes.reshape(1, -1), + match.beta_rescale, + ) + summed = self._checked_int32( + activation_i32.to(torch.int64) + beta_i32.to(torch.int64) + ) + if summed is None: + return None + return _apply_tosa_rescale(summed, match.add_output_rescale) + + @staticmethod + def _fit_integer_affine( + input_codes: torch.Tensor, + output_codes: torch.Tensor, + *, + input_zp: int, + output_zp: int, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + centered_inputs = input_codes.to(torch.int64).reshape(-1) - input_zp + centered_outputs = output_codes.to(torch.int64) - output_zp + slopes = [] + offsets = [] + + for channel in range(centered_outputs.shape[1]): + mapping: dict[int, int] = {} + for row in range(centered_inputs.numel()): + x = int(centered_inputs[row].item()) + y = int(centered_outputs[row, channel].item()) + previous = mapping.get(x) + if previous is not None and previous != y: + return None + mapping[x] = y + + points = sorted(mapping.items()) + if len(points) == 1: + slope = 0 + offset = points[0][1] + else: + x0, y0 = points[0] + x1, y1 = points[1] + dx = x1 - x0 + dy = y1 - y0 + if dy % dx != 0: + return None + slope = dy // dx + offset = y0 - slope * x0 + + if any(y != slope * x + offset for x, y in points): + return None + slopes.append(slope) + offsets.append(offset) + + return ( + torch.tensor(slopes, dtype=torch.int64), + torch.tensor(offsets, dtype=torch.int64), + ) + + @staticmethod + def _has_padding(conv: Node) -> bool: + padding = conv.args[4] + if not isinstance(padding, (list, tuple)): + return True + for value in padding: + if not isinstance(value, int) or value != 0: + return True + return False + + @staticmethod + def _exclusive_conv_input(match: _DyTAffineMatch) -> bool: + expected_user = match.conv + for node in match.layout_chain: + if set(node.users) != {expected_user}: + return False + expected_user = node + return set(match.add_output.users) == {expected_user} + + @staticmethod + def _weight_zero_points( + weight_qargs: QuantArgs, + out_channels: int, + weight_dim: int, + ) -> torch.Tensor | None: + if weight_qargs.per_channel: + if weight_qargs.axis != 0: + return None + zero_points = weight_qargs.get_zp_per_channel() + if len(zero_points) != out_channels: + return None + return torch.tensor(zero_points, dtype=torch.int64).reshape( + (out_channels,) + (1,) * (weight_dim - 1) + ) + return torch.tensor( + weight_qargs.get_zp_per_tensor(), + dtype=torch.int64, + ) + + def _create_constant( + self, + graph: Graph, + original: Node, + *, + name: str, + data: torch.Tensor, + ) -> Node: + kind = get_constant_placeholder_kind(self.exported_program, original) + persistent_buffer = is_persistent_buffer(self.exported_program, original) + with graph.inserting_before(original): + return create_constant_placeholder( + self.exported_program, + graph=graph, + name=name, + kind=kind, + data=data, + persistent_buffer=persistent_buffer, + ) + + def _conv_operands(self, match: _DyTAffineMatch) -> _ConvOperands | None: + """Read the conv's weight/bias constants and quantization arguments. + + Returns ``None`` unless the conv carries INT8 weights, an INT32 bias, + and per-tensor INT8 activation qparams whose zero point already agrees + with the affine output rescale. + + """ + conv = match.conv + if not isinstance(conv.args[1], Node) or not isinstance(conv.args[2], Node): + return None + weight_node = conv.args[1] + bias_node = conv.args[2] + weight = self._constant(weight_node) + bias = self._constant(bias_node) + input_qparams = cast(dict[int, QuantArgs], conv.meta.get("input_qparams", {})) + activation_qargs = input_qparams.get(0) + weight_qargs = input_qparams.get(1) + if ( + weight is None + or weight.dtype != torch.int8 + or weight.dim() != 4 + or bias is None + or bias.dtype != torch.int32 + or bias.dim() != 1 + ): + return None + if ( + activation_qargs is None + or activation_qargs.per_channel + or activation_qargs.dtype != torch.int8 + or activation_qargs.get_zp_per_tensor() + != match.add_output_rescale.output_zp + or weight_qargs is None + or weight_qargs.dtype != torch.int8 + ): + return None + return _ConvOperands( + weight_node=weight_node, + bias_node=bias_node, + weight=weight, + bias=bias, + input_qparams=input_qparams, + activation_qargs=activation_qargs, + weight_qargs=weight_qargs, + ) + + def _validated_conv_constants( + self, + match: _DyTAffineMatch, + slopes: torch.Tensor, + offsets: torch.Tensor, + ) -> _ConvConstants | None: + """Check the conv grouping and channel counts against the affine map. + + Returns ``None`` when the group layout is unusable or when the per + channel slopes/offsets do not line up with the conv's input channels. + + """ + operands = self._conv_operands(match) + if operands is None: + return None + groups = match.conv.args[8] + if not isinstance(groups, int) or groups <= 0: + return None + weight = operands.weight + out_channels = weight.shape[0] + in_channels_per_group = weight.shape[1] + in_channels = in_channels_per_group * groups + if ( + slopes.numel() != in_channels + or offsets.numel() != in_channels + or out_channels % groups != 0 + or operands.bias.numel() != out_channels + ): + return None + weight_zero_points = self._weight_zero_points( + operands.weight_qargs, + out_channels, + weight.dim(), + ) + if weight_zero_points is None: + return None + return _ConvConstants( + operands=operands, + groups=groups, + out_channels=out_channels, + in_channels_per_group=in_channels_per_group, + weight_zero_points=weight_zero_points, + ) + + def _fold_conv_constants( + self, + graph: Graph, + match: _DyTAffineMatch, + slopes: torch.Tensor, + offsets: torch.Tensor, + ) -> bool: + conv = match.conv + constants = self._validated_conv_constants(match, slopes, offsets) + if constants is None: + return False + operands = constants.operands + weight_node = operands.weight_node + bias_node = operands.bias_node + weight = operands.weight + bias = operands.bias + input_qparams = operands.input_qparams + activation_qargs = operands.activation_qargs + weight_qargs = operands.weight_qargs + groups = constants.groups + out_channels = constants.out_channels + in_channels_per_group = constants.in_channels_per_group + weight_zero_points = constants.weight_zero_points + + out_channels_per_group = out_channels // groups + output_groups = torch.arange(out_channels, dtype=torch.int64).div( + out_channels_per_group, + rounding_mode="floor", + ) + local_inputs = torch.arange(in_channels_per_group, dtype=torch.int64) + global_inputs = output_groups.reshape( + -1, 1 + ) * in_channels_per_group + local_inputs.reshape(1, -1) + broadcast_shape = ( + out_channels, + in_channels_per_group, + *([1] * (weight.dim() - 2)), + ) + channel_slopes = slopes[global_inputs].reshape(broadcast_shape) + channel_offsets = offsets[global_inputs].reshape(broadcast_shape) + + centered_weight = weight.to(torch.int64) - weight_zero_points + folded_centered_weight = centered_weight * channel_slopes + folded_weight_i64 = folded_centered_weight + weight_zero_points + if folded_weight_i64.numel() and ( + int(folded_weight_i64.min()) < weight_qargs.qmin + or int(folded_weight_i64.max()) > weight_qargs.qmax + ): + return False + + correction_dims = tuple(range(1, centered_weight.dim())) + bias_correction = (centered_weight * channel_offsets).sum(dim=correction_dims) + folded_bias_i64 = bias.to(torch.int64) + bias_correction + int32_limits = torch.iinfo(torch.int32) + if folded_bias_i64.numel() and ( + int(folded_bias_i64.min()) < int32_limits.min + or int(folded_bias_i64.max()) > int32_limits.max + ): + return False + + folded_weight = folded_weight_i64.to(torch.int8) + folded_bias = folded_bias_i64.to(torch.int32) + new_weight_node = weight_node + new_bias_node = bias_node + if not torch.equal(folded_weight, weight): + new_weight_node = self._create_constant( + graph, + weight_node, + name=f"{weight_node.name}_{conv.name}_dyt_affine_folded", + data=folded_weight, + ) + if not torch.equal(folded_bias, bias): + new_bias_node = self._create_constant( + graph, + bias_node, + name=f"{bias_node.name}_{conv.name}_dyt_affine_folded", + data=folded_bias, + ) + + conv.args = ( + conv.args[0], + new_weight_node, + new_bias_node, + *conv.args[3:], + ) + for original, replacement in ( + (weight_node, new_weight_node), + (bias_node, new_bias_node), + ): + if original is not replacement and len(original.users) == 0: + delete_constant_placeholder(self.exported_program, original) + updated_qparams = copy(input_qparams) + # Keep the original convolution activation scale: the folded integer + # weights and bias preserve that accumulator domain. Only the source + # zero point changes when the input is rewired to the TABLE. + updated_qparams[0] = QuantArgs( + scale=activation_qargs.scale, + zp=match.table_qargs.get_zp_per_tensor(), + qmin=activation_qargs.qmin, + qmax=activation_qargs.qmax, + dtype=activation_qargs.dtype, + axis=activation_qargs.axis, + per_channel=False, + ) + conv.meta["input_qparams"] = updated_qparams + return True + + def _fold_unpadded( + self, + graph: Graph, + match: _DyTAffineMatch, + affine_outputs: torch.Tensor, + ) -> bool: + if not self._exclusive_conv_input(match): + return False + fitted = self._fit_integer_affine( + match.table_values, + affine_outputs, + input_zp=match.table_qargs.get_zp_per_tensor(), + output_zp=match.add_output_rescale.output_zp, + ) + if fitted is None: + return False + slopes, offsets = fitted + if not self._fold_conv_constants(graph, match, slopes, offsets): + return False + match.add_output.replace_all_uses_with(match.table) + return True + + @staticmethod + def _affine_is_identity( + match: _DyTAffineMatch, + affine_outputs: torch.Tensor, + ) -> bool: + expected = match.table_values.reshape(-1, 1).expand_as(affine_outputs) + return torch.equal(affine_outputs, expected) + + @staticmethod + def _gamma_is_identity( + match: _DyTAffineMatch, + gamma_outputs: torch.Tensor, + ) -> bool: + expected = match.table_values.reshape(-1, 1).expand_as(gamma_outputs) + return torch.equal(gamma_outputs, expected) + + def call(self, graph_module: GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + constants_to_delete: set[Node] = set() + for node in list(graph.nodes): + match = self._match(node) + if match is None: + continue + gamma_outputs = self._gamma_outputs(match) + if gamma_outputs is None: + continue + + affine_outputs = self._affine_outputs(match, gamma_outputs) + if affine_outputs is None: + continue + + folded = False + if not self._has_padding(match.conv): + folded = self._fold_unpadded(graph, match, affine_outputs) + if folded: + constants_to_delete.update((match.gamma, match.beta)) + # Identity here means equality of the emitted INT8 codes. Keep the + # consumers' original zero points: changing them to the TABLE zero + # point would reinterpret the same bytes and alter convolution + # padding or downstream RESCALE arithmetic. + elif self._affine_is_identity(match, affine_outputs): + match.add_output.replace_all_uses_with(match.table) + constants_to_delete.update((match.gamma, match.beta)) + folded = True + elif self._gamma_is_identity(match, gamma_outputs): + match.gamma_output.replace_all_uses_with(match.table) + constants_to_delete.add(match.gamma) + folded = True + modified = modified or folded + + if modified: + graph.eliminate_dead_code() + for constant in constants_to_delete: + if constant.op == "placeholder" and len(constant.users) == 0: + delete_constant_placeholder(self.exported_program, constant) + graph.lint() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/arm/test/passes/test_fold_dyt_affine_into_conv_pass.py b/backends/arm/test/passes/test_fold_dyt_affine_into_conv_pass.py new file mode 100644 index 00000000000..41b4c2d7601 --- /dev/null +++ b/backends/arm/test/passes/test_fold_dyt_affine_into_conv_pass.py @@ -0,0 +1,1038 @@ +# 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 +"""Tests for folding exact quantized DyT affine maps into following convs.""" + +from __future__ import annotations + +import importlib +from types import ModuleType +from typing import cast, ClassVar, Dict, Tuple + +import executorch.backends.arm.tosa.dialect # noqa: F401 +import torch + +from executorch.backends.arm._passes import ( + FoldAndAnnotateQParamsPass, + InsertRescaleInt32Pass, + MatchArgRanksPass, +) +from executorch.backends.arm._passes.arm_pass_utils import get_param_tensor +from executorch.backends.arm._passes.fold_dyt_affine_into_conv_pass import ( + FoldDyTAffineIntoConvPass, +) +from executorch.backends.arm._passes.fold_dyt_alpha_into_lut_pass import ( + FoldDyTAlphaIntoLUTPass, +) +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from torch.export import export +from torch.fx import Node + + +_CHANNELS: int = 2 + + +class _PostRescaleAffineFixture(torch.nn.Module): + # Declared so the checker sees the registered buffers as Tensors rather than + # the ``Tensor | Module`` that ``nn.Module.__getattr__`` is annotated to give. + table: torch.Tensor + gamma: torch.Tensor + beta: torch.Tensor + weight: torch.Tensor + bias: torch.Tensor + + def __init__( + self, + *, + table: torch.Tensor, + gamma: torch.Tensor, + beta: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + ) -> None: + super().__init__() + self.register_buffer("table", table) + self.register_buffer("gamma", gamma) + self.register_buffer("beta", beta) + self.register_buffer("weight", weight) + self.register_buffer("bias", bias) + + def forward(self, x_code: torch.Tensor) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + return x_code, self.table, self.gamma, self.beta, self.weight, self.bias + + +def _qargs(scale: float, zp: int, dtype: torch.dtype = torch.int8) -> QuantArgs: + dtype_range = torch.iinfo(dtype) + return QuantArgs( + scale=scale, + zp=zp, + qmin=dtype_range.min, + qmax=dtype_range.max, + dtype=dtype, + ) + + +# ``QuantArgs.scale``/``zp`` are typed to also cover the per-channel case, where +# they are lists. Every fixture in this file is per-tensor, so narrow them once +# here instead of casting at each arithmetic site. +def _scale_of(qargs: QuantArgs) -> float: + return cast(float, qargs.scale) + + +def _zp_of(qargs: QuantArgs) -> int: + return cast(int, qargs.zp) + + +def _channel_codes(value: int | tuple[int, ...]) -> torch.Tensor: + values = (value,) * _CHANNELS if isinstance(value, int) else value + return torch.tensor(values, dtype=torch.int8) + + +def _buffer_nodes(exported_program: ExportedProgram) -> dict[str, Node]: + graph = exported_program.graph_module.graph + nodes_by_name = {node.name: node for node in graph.nodes} + return { + buffer_name: nodes_by_name[placeholder_name] + for placeholder_name, buffer_name in exported_program.graph_signature.inputs_to_buffers.items() + } + + +def _pass_module() -> ModuleType: + return importlib.import_module( + "executorch.backends.arm._passes.fold_dyt_affine_into_conv_pass" + ) + + +def _fixture_weight_and_bias( + *, depthwise: bool, channel_slice: bool +) -> tuple[torch.Tensor, torch.Tensor]: + """Pick the conv constants matching the graph shape under test.""" + if depthwise: + weight = torch.tensor( + [ + [[[1, -2, 1]]], + [[[2, 1, -1]]], + ], + dtype=torch.int8, + ) + elif channel_slice: + # A single input channel keeps the graph well formed behind the + # narrowing slice: the conv really does consume 1 of the 2 affine + # channels, which is what makes this a genuine mismatch rather than an + # impossible graph. + weight = torch.tensor([[[[1]]], [[[3]]]], dtype=torch.int8) + else: + weight = torch.tensor( + [ + [[[1]], [[2]]], + [[[3]], [[-2]]], + ], + dtype=torch.int8, + ) + return weight, torch.tensor([5, -7], dtype=torch.int32) + + +def _fixture_rescale_params( + actual_dyt_identity_qparams: bool, +) -> tuple[QuantArgs, int, float, int, int, float, float, float]: + """Return the TABLE qparams and the rescale scales/zero points to build. + + The identity variant replays the scales a real DyT block produces, so the + fixture exercises the same integer arithmetic the pass sees in a model. + + """ + if actual_dyt_identity_qparams: + table_qargs = _qargs(scale=0.00588326808065176, zp=-6) + beta_scale = 1.52587890625e-05 + common_scale = (2.0 * _scale_of(table_qargs)) / (1 << 20) + return ( + table_qargs, + -128, + 1.0 / 255.0, + _zp_of(table_qargs), + -128, + _scale_of(table_qargs) / common_scale, + beta_scale / common_scale, + common_scale / _scale_of(table_qargs), + ) + table_qargs = _qargs(scale=0.1, zp=0) + return (table_qargs, 0, 1.0, 0, 0, 1.0, 1.0, 1.0) + + +def _fixture_conv_input( + graph: torch.fx.Graph, + layout_source: Node, + *, + layout_permute_count: int, + slice_passthrough: bool, + channel_slice: bool, +) -> tuple[Node, Node]: + """Build the NHWC->NCHW layout chain feeding the conv. + + Returns the layout output (shared users hang off it) and the node the conv + actually consumes, which may sit behind a slice. + + """ + layout_output = layout_source + for _ in range(layout_permute_count): + layout_output = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + (layout_output, [0, 3, 1, 2]), + ) + conv_input = layout_output + if slice_passthrough: + conv_input = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + (layout_output, 0, 0, 1, 1), + ) + if channel_slice: + conv_input = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + (layout_output, 1, 0, 1, 1), + ) + return layout_output, conv_input + + +def _build_post_rescale_fixture( + *, + table: torch.Tensor, + gamma_code: int | tuple[int, ...], + beta_code: int | tuple[int, ...], + actual_dyt_identity_qparams: bool, + padded_depthwise: bool, + unpadded_depthwise: bool = False, + slice_passthrough: bool = False, + channel_slice: bool = False, + shared_layout_user: bool = False, + input_width: int = 4, + input_height: int = 1, + input_channels: int = _CHANNELS, + affine_view_shape: tuple[int, ...] | None = None, + activation_view_shape: tuple[int, ...] | None = None, + layout_permute_count: int = 1, + conv_input_zp: int | None = None, + gamma_output_zp: int | None = None, + add_output_zp: int | None = None, +) -> tuple[ExportedProgram, torch.Tensor, torch.Tensor]: + weight, bias = _fixture_weight_and_bias( + depthwise=padded_depthwise or unpadded_depthwise, + channel_slice=channel_slice, + ) + + test_input = ( + torch.arange( + input_height * input_width * input_channels, dtype=torch.int8 + ).reshape(1, input_height, input_width, input_channels), + ) + exported_program = export( + _PostRescaleAffineFixture( + table=table, + gamma=_channel_codes(gamma_code), + beta=_channel_codes(beta_code), + weight=weight, + bias=bias, + ), + test_input, + strict=True, + ) + graph = exported_program.graph_module.graph + buffers = _buffer_nodes(exported_program) + activation = next( + node + for node in graph.nodes + if node.op == "placeholder" + and node.name not in exported_program.graph_signature.inputs_to_buffers + ) + output = next(node for node in graph.nodes if node.op == "output") + view_shape = list(affine_view_shape or (1, 1, 1, _CHANNELS)) + + ( + table_qargs, + gamma_input_zp, + gamma_output_scale, + fixture_gamma_output_zp, + beta_input_zp, + add_activation_scale, + add_beta_scale, + add_output_scale, + ) = _fixture_rescale_params(actual_dyt_identity_qparams) + resolved_gamma_output_zp = ( + fixture_gamma_output_zp if gamma_output_zp is None else gamma_output_zp + ) + resolved_add_output_zp = ( + _zp_of(table_qargs) if add_output_zp is None else add_output_zp + ) + + with graph.inserting_before(output): + table_node = graph.call_function( + exir_ops.backend.tosa.TABLE.default, + (activation, buffers["table"]), + ) + activation_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (table_node, torch.int32, [1.0], table_qargs.zp, 0), + ) + gamma_activation = activation_rescale + if activation_view_shape is not None: + gamma_activation = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (activation_rescale, list(activation_view_shape)), + ) + gamma_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (buffers["gamma"], torch.int32, [1.0], gamma_input_zp, 0), + ) + gamma_view = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (gamma_rescale, view_shape), + ) + mul = graph.call_function( + exir_ops.edge.aten.mul.Tensor, + (gamma_activation, gamma_view), + ) + mul_output_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (mul, torch.int8, [gamma_output_scale], 0, resolved_gamma_output_zp), + ) + add_activation_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + ( + mul_output_rescale, + torch.int32, + [add_activation_scale], + resolved_gamma_output_zp, + 0, + ), + ) + beta_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (buffers["beta"], torch.int32, [add_beta_scale], beta_input_zp, 0), + ) + beta_view = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (beta_rescale, view_shape), + ) + add = graph.call_function( + exir_ops.edge.aten.add.Tensor, + (add_activation_rescale, beta_view), + ) + add_output_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (add, torch.int8, [add_output_scale], 0, resolved_add_output_zp), + ) + layout_output, conv_input = _fixture_conv_input( + graph, + add_output_rescale, + layout_permute_count=layout_permute_count, + slice_passthrough=slice_passthrough, + channel_slice=channel_slice, + ) + conv = graph.call_function( + exir_ops.edge.aten.convolution.default, + ( + conv_input, + buffers["weight"], + buffers["bias"], + [1, 1], + [0, 1] if padded_depthwise else [0, 0], + [1, 1], + False, + [0, 0], + _CHANNELS if padded_depthwise or unpadded_depthwise else 1, + ), + ) + shared_output = None + if shared_layout_user: + shared_output = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (layout_output, [1, _CHANNELS, 1, input_width]), + ) + + table_node.meta["output_qparams"] = {0: table_qargs} + table_node.meta["val"] = torch.empty( + (1, input_height, input_width, input_channels), + dtype=torch.int8, + device="meta", + ) + add_output_rescale.meta["val"] = torch.empty( + (1, input_height, input_width, _CHANNELS), + dtype=torch.int8, + device="meta", + ) + conv.meta["input_qparams"] = { + 0: _qargs( + scale=_scale_of(table_qargs), + zp=(resolved_add_output_zp if conv_input_zp is None else conv_input_zp), + ), + 1: _qargs(scale=0.02, zp=0), + } + output.args = ((conv, shared_output) if shared_output is not None else (conv,),) + graph.eliminate_dead_code() + graph.lint() + exported_program.graph_module.recompile() + return exported_program, weight, bias + + +def _add_second_shared_weight_branch( + exported_program: ExportedProgram, +) -> None: + graph = exported_program.graph_module.graph + buffers = _buffer_nodes(exported_program) + output = next(node for node in graph.nodes if node.op == "output") + table = next( + node + for node in graph.nodes + if node.target == exir_ops.backend.tosa.TABLE.default + ) + conv = next( + node + for node in graph.nodes + if node.target == exir_ops.edge.aten.convolution.default + ) + table_qargs = cast(dict[int, QuantArgs], table.meta["output_qparams"])[0] + + with graph.inserting_before(output): + activation_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (table, torch.int32, [1.0], table_qargs.zp, 0), + ) + gamma_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (buffers["gamma"], torch.int32, [1.0], -1, 0), + ) + gamma_view = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (gamma_rescale, [1, 1, 1, _CHANNELS]), + ) + mul = graph.call_function( + exir_ops.edge.aten.mul.Tensor, + (activation_rescale, gamma_view), + ) + gamma_output = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (mul, torch.int8, [1.0], 0, table_qargs.zp), + ) + add_activation = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (gamma_output, torch.int32, [1.0], table_qargs.zp, 0), + ) + beta_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (buffers["beta"], torch.int32, [1.0], 1, 0), + ) + beta_view = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (beta_rescale, [1, 1, 1, _CHANNELS]), + ) + add = graph.call_function( + exir_ops.edge.aten.add.Tensor, + (add_activation, beta_view), + ) + add_output = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (add, torch.int8, [1.0], 0, table_qargs.zp), + ) + nchw = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + (add_output, [0, 3, 1, 2]), + ) + second_conv = graph.call_function( + exir_ops.edge.aten.convolution.default, + (nchw, *conv.args[1:]), + ) + + add_output.meta["val"] = table.meta["val"] + second_conv.meta["input_qparams"] = dict(conv.meta["input_qparams"]) + output.args = ((conv, second_conv),) + graph.eliminate_dead_code() + graph.lint() + exported_program.graph_module.recompile() + + +def _call_pass(exported_program: ExportedProgram) -> PassResult: + pass_class = _pass_module().FoldDyTAffineIntoConvPass + return pass_class(exported_program).call(exported_program.graph_module) + + +def _call_targets(exported_program: ExportedProgram) -> list[str]: + return [ + str(node.target) + for node in exported_program.graph_module.graph.nodes + if node.op == "call_function" + ] + + +def _conv_constants( + exported_program: ExportedProgram, +) -> tuple[torch.Tensor, torch.Tensor]: + conv = next( + node + for node in exported_program.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.convolution.default + ) + weight_node = cast(Node, conv.args[1]) + bias_node = cast(Node, conv.args[2]) + weight = get_param_tensor(exported_program, weight_node) + bias = get_param_tensor(exported_program, bias_node) + assert weight is not None + assert bias is not None + return weight, bias + + +def test_unpadded_conv_folds_exact_integer_affine_into_weight_and_bias() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + ) + original_constants = _buffer_nodes(exported_program) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + expected_weight = weight.to(torch.int16).mul(2).to(torch.int8) + expected_bias = bias + weight.to(torch.int32).sum(dim=(1, 2, 3)).mul(3) + placeholder_names = { + node.name + for node in exported_program.graph_module.graph.nodes + if node.op == "placeholder" + } + + assert result.modified + assert torch.equal(folded_weight, expected_weight) + assert torch.equal(folded_bias, expected_bias) + assert not any("aten.mul" in target for target in _call_targets(exported_program)) + assert not any("aten.add" in target for target in _call_targets(exported_program)) + assert original_constants["weight"].name not in placeholder_names + assert original_constants["bias"].name not in placeholder_names + assert original_constants["gamma"].name not in placeholder_names + assert original_constants["beta"].name not in placeholder_names + + +def test_unpadded_depthwise_folds_nonuniform_gamma_and_beta() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=(2, 3), + beta_code=(3, -2), + actual_dyt_identity_qparams=False, + padded_depthwise=False, + unpadded_depthwise=True, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + expected_weight = ( + weight.to(torch.int16) + .mul(torch.tensor([2, 3], dtype=torch.int16).reshape(2, 1, 1, 1)) + .to(torch.int8) + ) + expected_bias = bias + weight.to(torch.int32).sum(dim=(1, 2, 3)).mul( + torch.tensor([3, -2], dtype=torch.int32) + ) + + assert result.modified + assert torch.equal(folded_weight, expected_weight) + assert torch.equal(folded_bias, expected_bias) + assert not any("aten.mul" in target for target in _call_targets(exported_program)) + assert not any("aten.add" in target for target in _call_targets(exported_program)) + + +def test_mismatched_conv_input_zero_point_is_rejected() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + conv_input_zp=1, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert any("aten.mul" in target for target in targets) + assert any("aten.add" in target for target in targets) + + +def test_shared_conv_constants_get_distinct_folded_values() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + ) + _add_second_shared_weight_branch(exported_program) + + result = _call_pass(exported_program) + convs = [ + node + for node in exported_program.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.convolution.default + ] + folded_constants = [] + for conv in convs: + folded_weight = get_param_tensor(exported_program, cast(Node, conv.args[1])) + folded_bias = get_param_tensor(exported_program, cast(Node, conv.args[2])) + assert folded_weight is not None + assert folded_bias is not None + folded_constants.append((folded_weight, folded_bias)) + + weight_sum = weight.to(torch.int32).sum(dim=(1, 2, 3)) + assert result.modified + assert len(folded_constants) == 2 + assert torch.equal( + folded_constants[0][0], weight.to(torch.int16).mul(2).to(torch.int8) + ) + assert torch.equal(folded_constants[0][1], bias + weight_sum.mul(3)) + assert torch.equal( + folded_constants[1][0], weight.to(torch.int16).mul(3).to(torch.int8) + ) + assert torch.equal(folded_constants[1][1], bias + weight_sum.mul(2)) + + +def test_unpadded_identity_affine_removes_ops_without_changing_constants() -> None: + table = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=127, + beta_code=-128, + actual_dyt_identity_qparams=True, + padded_depthwise=False, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + + assert result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert not any("aten.mul" in target for target in _call_targets(exported_program)) + assert not any("aten.add" in target for target in _call_targets(exported_program)) + + +def test_padded_depthwise_removes_identity_gamma_but_keeps_beta_add() -> None: + table = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=127, + beta_code=127, + actual_dyt_identity_qparams=True, + padded_depthwise=True, + shared_layout_user=True, + ) + original_constants = _buffer_nodes(exported_program) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert not any("aten.mul" in target for target in targets) + assert sum("aten.add" in target for target in targets) == 1 + placeholder_names = { + node.name + for node in exported_program.graph_module.graph.nodes + if node.op == "placeholder" + } + assert original_constants["gamma"].name not in placeholder_names + assert original_constants["beta"].name in placeholder_names + + +def test_identity_affine_preserves_convolution_input_zero_point() -> None: + table = torch.zeros(256, dtype=torch.int8) + exported_program, _, _ = _build_post_rescale_fixture( + table=table, + gamma_code=1, + beta_code=-7, + actual_dyt_identity_qparams=False, + padded_depthwise=True, + add_output_zp=7, + ) + + result = _call_pass(exported_program) + conv = next( + node + for node in exported_program.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.convolution.default + ) + table_node = next( + node + for node in exported_program.graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.TABLE.default + ) + input_qparams = cast(dict[int, QuantArgs], conv.meta["input_qparams"]) + table_qparams = cast(dict[int, QuantArgs], table_node.meta["output_qparams"]) + + assert result.modified + assert table_qparams[0].get_zp_per_tensor() == 0 + assert input_qparams[0].get_zp_per_tensor() == 7 + assert not any("aten.mul" in target for target in _call_targets(exported_program)) + assert not any("aten.add" in target for target in _call_targets(exported_program)) + + +def test_identity_gamma_preserves_downstream_rescale_input_zero_point() -> None: + table = torch.ones(256, dtype=torch.int8) + exported_program, _, _ = _build_post_rescale_fixture( + table=table, + gamma_code=-6, + beta_code=0, + actual_dyt_identity_qparams=False, + padded_depthwise=True, + gamma_output_zp=7, + ) + + result = _call_pass(exported_program) + add = next( + node + for node in exported_program.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.add.Tensor + ) + activation_rescale = cast(Node, add.args[0]) + table_node = cast(Node, activation_rescale.args[0]) + table_qparams = cast(dict[int, QuantArgs], table_node.meta["output_qparams"]) + + assert result.modified + assert table_node.target == exir_ops.backend.tosa.TABLE.default + assert table_qparams[0].get_zp_per_tensor() == 0 + assert activation_rescale.args[3] == 7 + assert not any("aten.mul" in target for target in _call_targets(exported_program)) + + +def test_zero_layout_permute_is_rejected() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + input_height=_CHANNELS, + layout_permute_count=0, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + + +def test_repeated_layout_permute_is_rejected() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + input_width=_CHANNELS, + layout_permute_count=2, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + + +def test_wrong_axis_affine_views_are_rejected() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + input_width=_CHANNELS, + affine_view_shape=(1, 1, _CHANNELS, 1), + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert any("aten.mul" in target for target in targets) + assert any("aten.add" in target for target in targets) + + +def test_activation_side_view_is_rejected() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + activation_view_shape=(1, 1, 4, _CHANNELS), + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert any("aten.mul" in target for target in targets) + assert any("aten.add" in target for target in targets) + + +def test_singleton_table_channel_broadcast_is_rejected() -> None: + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + input_channels=1, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert any("aten.mul" in target for target in targets) + assert any("aten.add" in target for target in targets) + + +def test_unpadded_shared_layout_removes_only_identity_gamma() -> None: + table = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=127, + beta_code=127, + actual_dyt_identity_qparams=True, + padded_depthwise=False, + shared_layout_user=True, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert not any("aten.mul" in target for target in targets) + assert sum("aten.add" in target for target in targets) == 1 + + +def test_unpadded_conv_folds_through_slice_passthrough() -> None: + table = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=127, + beta_code=-128, + actual_dyt_identity_qparams=True, + padded_depthwise=False, + slice_passthrough=True, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert not any("aten.mul" in target for target in targets) + assert not any("aten.add" in target for target in targets) + assert any("aten.slice_copy" in target for target in targets) + + +def test_channel_narrowing_slice_is_rejected() -> None: + """A slice that changes which channels the conv consumes must not fold. + + ``_trace_layout_source`` deliberately does not inspect ``slice_copy`` + arguments; safety comes from the channel-count guard in + ``_fold_conv_constants``, which compares the affine site's per-channel + slope/offset count against the conv weight's input channels. Here the + affine site produces two channels but the slice leaves the conv consuming + one, so the counts disagree and the fold must decline. This is the + fail-closed path that keeps the unvalidated ``slice_copy`` passthrough + sound, so it is pinned here rather than left implicit. + + """ + table = (torch.arange(256, dtype=torch.int16).remainder(21) - 10).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + channel_slice=True, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert any("aten.mul" in target for target in targets) + assert any("aten.add" in target for target in targets) + assert any("aten.slice_copy" in target for target in targets) + + +def test_identity_affine_behind_channel_slice_leaves_conv_constants() -> None: + """An exact-identity gamma/beta may be dropped even behind a channel slice. + + Identity is established per channel over the whole affine site, so removing + the Mul/Add is a no-op on every channel and stays sound no matter which + channels the conv goes on to consume. The conv constants must be left + untouched: nothing is folded into them, the redundant ops are just deleted. + Contrast ``test_channel_narrowing_slice_is_rejected``, where a real + (non-identity) affine behind the same slice is refused outright. + + """ + table = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=127, + beta_code=-128, + actual_dyt_identity_qparams=True, + padded_depthwise=False, + channel_slice=True, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + + assert result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + targets = _call_targets(exported_program) + assert not any("aten.mul" in target for target in targets) + assert not any("aten.add" in target for target in targets) + assert any("aten.slice_copy" in target for target in targets) + + +def test_non_affine_integer_mapping_is_rejected() -> None: + table = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + exported_program, weight, bias = _build_post_rescale_fixture( + table=table, + gamma_code=2, + beta_code=3, + actual_dyt_identity_qparams=False, + padded_depthwise=False, + ) + + result = _call_pass(exported_program) + folded_weight, folded_bias = _conv_constants(exported_program) + targets = _call_targets(exported_program) + + assert not result.modified + assert torch.equal(folded_weight, weight) + assert torch.equal(folded_bias, bias) + assert any("aten.mul" in target for target in targets) + assert any("aten.add" in target for target in targets) + + +class DyTAffineModule(torch.nn.Module): + """A full DyT site between two convs, mirroring the real module. + + conv -> NHWC permute -> tanh(alpha * x) -> x * gamma + beta -> back to NCHW + -> conv. The trailing conv is what the affine folds into. + + """ + + test_data: ClassVar[Dict[str, Tuple[torch.Tensor]]] = { + "rand": (torch.rand(1, 3, 8, 8),), + } + + def __init__(self, channels: int = 3, alpha: float = 0.5) -> None: + super().__init__() + self.conv_in = torch.nn.Conv2d(channels, channels, kernel_size=1) + self.alpha = torch.nn.Parameter(torch.tensor([alpha])) + self.gamma = torch.nn.Parameter(torch.ones(channels)) + self.beta = torch.nn.Parameter(torch.zeros(channels)) + self.conv_out = torch.nn.Conv2d(channels, channels, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = torch.permute(self.conv_in(x), (0, 2, 3, 1)) + y = torch.tanh(self.alpha * y) + y = y * self.gamma + self.beta + y = torch.permute(y, (0, 3, 1, 2)) + return self.conv_out(y) + + +@common.parametrize("test_data", DyTAffineModule.test_data) +def test_fold_dyt_affine_into_conv_tosa_INT(test_data: Tuple[torch.Tensor]) -> None: + """Pipeline-level counterpart to the IR-level regressions above. + + ``MatchArgRanksPass`` is required, not incidental: this pass only matches + gamma/beta operands that carry an explicit ``(1, 1, 1, C)`` view. A bare + ``(C,)`` constant broadcasts against the NHWC activation without one, and the + pass then declines to fold. ``MatchArgRanksPass`` is what materialises that + view, and it sits between ``InsertRescaleInt32Pass`` and + ``InsertTableOpsPass`` in ``ArmPassManager`` for exactly this reason. + + """ + pipeline = PassPipeline[Tuple[torch.Tensor]]( + DyTAffineModule(), + test_data, + quantize=True, + ops_after_pass={ + "executorch_exir_dialects_backend__ops_tosa_TABLE_default": 1, + "executorch_exir_dialects_edge__ops_aten_convolution_default": 2, + }, + ops_not_after_pass=[ + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + "executorch_exir_dialects_edge__ops_aten_tanh_default", + ], + pass_list=[FoldAndAnnotateQParamsPass, InsertRescaleInt32Pass], + passes_with_exported_program=[ + MatchArgRanksPass, + FoldDyTAlphaIntoLUTPass, + FoldDyTAffineIntoConvPass, + ], + ) + # The partial ``pass_list`` above stops short of a full TOSA lowering, so no + # runnable program is left for the comparison stage to execute. Dropped for + # the same reason as in ``test_insert_rescale_i32_pass.py``, which drives + # the same two passes. Skipping it does not leave the rewritten weights and + # biases unchecked: the IR-level regressions above assert the folded + # constants exactly, and the fold is only ever applied when the per-channel + # mapping is provably integer-affine, so it is exact by construction rather + # than approximate. + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run()