diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 71df953a65b..3273923b1e1 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -1,3 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. # Copyright 2025-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the @@ -114,6 +116,8 @@ 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, QuantizeClampArgumentsPass, diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 1131fd26ae0..42e29bc8860 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -208,6 +208,11 @@ class PassInsertions: _registered_pass_insertions: dict[type, PassInsertions] = {} +_registered_pass_factories_before: dict[ + type, list[Callable[[ExportedProgram], ExportPass]] +] = {} + + def _graph_pass_name(graph_pass: Callable[[GraphModule], PassResult | None]) -> str: if isinstance(graph_pass, ExportPass): return ArmPass.get_name(graph_pass) @@ -271,9 +276,21 @@ def register_pass_insertions_after( _registered_pass_insertions[target_pass_type].after_passes.extend(passes) +def register_pass_factories_before( + target_pass_type: type, + factories: list[Callable[[ExportedProgram], ExportPass]], +) -> None: + """Register factories instantiated with the backend ExportedProgram.""" + registered = _registered_pass_factories_before.setdefault(target_pass_type, []) + for factory in factories: + if factory not in registered: + registered.append(factory) + + def clear_registered_pass_insertions() -> None: - """Clear all globally registered pass insertions.""" + """Clear all globally registered pass insertions and pass factories.""" _registered_pass_insertions.clear() + _registered_pass_factories_before.clear() class ArmPassManager(ExportedProgramPassManager): @@ -430,6 +447,11 @@ def _configure_pass_insertions(self, exported_program: ExportedProgram) -> None: self.insert_passes_before(pass_type, list(insertions.before_passes)) if insertions.after_passes: self.insert_passes_after(pass_type, list(insertions.after_passes)) + for pass_type, factories in _registered_pass_factories_before.items(): + self.insert_passes_before( + pass_type, + [factory(exported_program) for factory in factories], + ) def add_passes(self, passes: Sequence[ExportPass | None]): for p in passes: 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..346b05b366c --- /dev/null +++ b/backends/arm/_passes/fold_dyt_affine_into_conv_pass.py @@ -0,0 +1,675 @@ +# 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 exactly representable quantized DyT affine maps into following convs.""" + +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.insert_table_ops import InsertTableOpsPass +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_codes: torch.Tensor + gamma_rescale: _RescaleParams + gamma_output: Node + gamma_output_rescale: _RescaleParams + add_activation_rescale: _RescaleParams + beta_codes: torch.Tensor + beta_rescale: _RescaleParams + add_output: Node + add_output_rescale: _RescaleParams + conv: Node + layout_chain: tuple[Node, ...] + + +class FoldDyTAffineIntoConvPass(ArmPass): + """Fold exact integer DyT gamma/beta maps into a following convolution. + + Intermediate INT8 requantization makes a generic float affine fold inexact. + This pass evaluates the real TOSA integer path over the site's materialized + 256-entry TABLE and rewrites only when the resulting per-channel map is + exactly integer-affine. Padded convolutions retain beta because constant + padding makes its boundary contribution position dependent; their gamma path + is removed only when it is 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. + + ``permute_copy`` is 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 = [] + 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 not isinstance(dims, (list, tuple)) or tuple(dims) != (0, 3, 1, 2): + return None + chain.append(node) + node = node.args[0] + return node, tuple(chain) + + def _match(self, conv: Node) -> _DyTAffineMatch | None: # noqa: C901 + 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] + add_activation_operand = next( + operand for operand in add_operands if operand is not beta_operand + ) + + gamma_output = add_activation_operand.source + 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 + ) + table = table_operand.source + if ( + table.target != exir_ops.backend.tosa.TABLE.default + or len(table.args) < 2 + or not isinstance(table.args[1], Node) + ): + return None + + table_values = self._constant(table.args[1]) + table_qargs = self._single_qargs(table, "output_qparams") + gamma_codes = gamma_operand.constant + beta_codes = 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 + channel_view_shape = (1, 1, 1, gamma_codes.numel()) + if ( + table_operand.view_shape is not None + or add_activation_operand.view_shape is not None + or gamma_operand.view_shape != channel_view_shape + or beta_operand.view_shape != channel_view_shape + ): + return None + table_shape = self._tensor_shape(table) + if ( + table_shape is None + or table_shape != self._tensor_shape(add_output) + or len(table_shape) != 4 + or table_shape[-1] != gamma_codes.numel() + ): + return None + + return _DyTAffineMatch( + table=table, + table_values=table_values, + table_qargs=table_qargs, + gamma_activation_rescale=table_operand.rescale, + gamma_codes=gamma_codes, + gamma_rescale=gamma_operand.rescale, + gamma_output=gamma_output, + gamma_output_rescale=gamma_output_rescale, + add_activation_rescale=add_activation_operand.rescale, + beta_codes=beta_codes, + beta_rescale=beta_operand.rescale, + add_output=add_output, + add_output_rescale=add_output_rescale, + conv=conv, + layout_chain=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 _fold_conv_constants( + self, + graph: Graph, + match: _DyTAffineMatch, + slopes: torch.Tensor, + offsets: torch.Tensor, + ) -> bool: # noqa: C901 + conv = match.conv + if not isinstance(conv.args[1], Node) or not isinstance(conv.args[2], Node): + return False + + 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 + or activation_qargs is None + or activation_qargs.per_channel + or activation_qargs.dtype != torch.int8 + or weight_qargs is None + or weight_qargs.dtype != torch.int8 + ): + return False + + groups = conv.args[8] + if not isinstance(groups, int) or groups <= 0: + return False + 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 bias.numel() != out_channels + ): + return False + + weight_zero_points = self._weight_zero_points( + weight_qargs, + out_channels, + weight.dim(), + ) + if weight_zero_points is None: + return False + + 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) + 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, + gamma_outputs: torch.Tensor, + ) -> bool: + if not self._exclusive_conv_input(match): + return False + affine_outputs = self._affine_outputs(match, gamma_outputs) + if affine_outputs is None: + 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 _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 + 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 + + padded = self._has_padding(match.conv) + gamma_identity = self._gamma_is_identity(match, gamma_outputs) + folded = False + if not padded: + folded = self._fold_unpadded(graph, match, gamma_outputs) + if not folded and gamma_identity: + match.gamma_output.replace_all_uses_with(match.table) + folded = True + modified = modified or folded + + if modified: + graph.eliminate_dead_code() + graph.lint() + graph_module.recompile() + return PassResult(graph_module, modified) + + +def register_fold_dyt_affine_into_conv_pass() -> None: + """Enable exact DyT affine folding after the alpha TABLE rewrite.""" + # Imported here rather than at module scope: _passes/__init__.py imports + # arm_pass_manager last, and arm_pass_manager imports back from the package, + # so a module-level import from a globbed pass module is circular. + from executorch.backends.arm._passes.arm_pass_manager import ( + register_pass_factories_before, + ) + + register_pass_factories_before( + InsertTableOpsPass, + [FoldDyTAffineIntoConvPass], + ) diff --git a/backends/arm/_passes/fold_dyt_alpha_into_lut_pass.py b/backends/arm/_passes/fold_dyt_alpha_into_lut_pass.py new file mode 100644 index 00000000000..c081465e436 --- /dev/null +++ b/backends/arm/_passes/fold_dyt_alpha_into_lut_pass.py @@ -0,0 +1,341 @@ +# 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 quantized tanh(alpha * x) into one exact INT8 TOSA table.""" + +import math +from dataclasses import dataclass +from typing import cast, Optional, Set, Type + +import torch +from executorch.backends.arm._passes.arm_pass import ArmPass +from executorch.backends.arm._passes.arm_pass_utils import create_node, get_param_tensor +from executorch.backends.arm._passes.insert_table_ops import ( + create_constant_placeholder, + InsertTableOpsPass, +) +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm.operators.op_tosa_rescale import ( + _compute_multiplier_and_shift, +) +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.export.graph_signature import InputKind +from torch.fx import GraphModule, Node + + +@dataclass(frozen=True) +class _RescaleParams: + scale: float + input_zp: int + output_zp: int + output_dtype: torch.dtype + + @classmethod + def from_node(cls, node: Node) -> Optional["_RescaleParams"]: + if node.target != exir_ops.backend.tosa.RESCALE.default: + return None + scales = cast(list[float], node.args[2]) + if len(scales) != 1: + return None + if node.kwargs.get("input_unsigned", False) or node.kwargs.get( + "output_unsigned", False + ): + return None + return cls( + scale=float(scales[0]), + input_zp=cast(int, node.args[3]), + output_zp=cast(int, node.args[4]), + output_dtype=cast(torch.dtype, node.args[1]), + ) + + +def _apply_tosa_rescale( + values: torch.Tensor, + params: _RescaleParams, +) -> torch.Tensor: + """Apply TOSA-1.0 RESCALE SINGLE_ROUND exactly.""" + multipliers, shifts = _compute_multiplier_and_shift([params.scale]) + multiplier = multipliers[0] + shift = shifts[0] + # _compute_multiplier_and_shift asserts shift is in [2, 62], so shift - 1 is + # never negative here. Do not "guard" this by zeroing rounding on small + # shifts: that would silently change SINGLE_ROUND's round-half behaviour and + # break the bit-exactness this whole pass depends on. + rounding = 1 << (shift - 1) + centered = values.to(torch.int64) - params.input_zp + scaled = (centered * multiplier + rounding) >> shift + shifted = scaled + params.output_zp + dtype_range = torch.iinfo(params.output_dtype) + return shifted.clamp(dtype_range.min, dtype_range.max).to(params.output_dtype) + + +def _generate_dyt_lut( + *, + activation_qargs: QuantArgs, + alpha_code: torch.Tensor, + activation_rescale: _RescaleParams, + alpha_rescale: _RescaleParams, + mul_output_rescale: _RescaleParams, + tanh_input_qargs: QuantArgs, + tanh_output_qargs: QuantArgs, +) -> torch.Tensor: + """Compose the integer Mul/RESCALE path with the existing tanh mapping.""" + if alpha_code.numel() != 1: + raise ValueError(f"Expected scalar DyT alpha, got shape {alpha_code.shape}") + + int8_info = torch.iinfo(torch.int8) + domain = torch.arange( + int8_info.min, + int8_info.max + 1, + dtype=torch.int16, + ).clamp(activation_qargs.qmin, activation_qargs.qmax) + activation_i32 = _apply_tosa_rescale(domain.to(torch.int8), activation_rescale) + alpha_i32 = _apply_tosa_rescale(alpha_code, alpha_rescale) + product = (activation_i32.to(torch.int64) * alpha_i32.to(torch.int64)).to( + torch.int32 + ) + mul_codes = _apply_tosa_rescale(product, mul_output_rescale) + mul_codes = mul_codes.clamp(tanh_input_qargs.qmin, tanh_input_qargs.qmax) + tanh_values = torch.tanh(tanh_input_qargs.dequantize_value(mul_codes)) + return tanh_output_qargs.quantize_value(tanh_values).to(torch.int8) + + +@dataclass(frozen=True) +class _DyTMatch: + tanh: Node + activation: Node + activation_qargs: QuantArgs + activation_rescale: _RescaleParams + alpha_code: torch.Tensor + alpha_rescale: _RescaleParams + mul_output_rescale: _RescaleParams + tanh_input_qargs: QuantArgs + tanh_output_qargs: QuantArgs + + +class FoldDyTAlphaIntoLUTPass(ArmPass): + """Replace an INT8 DyT alpha Mul, requantize, and tanh with one TABLE.""" + + _passes_required_after: Set[Type[ExportPass]] = set() + + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__() + self.exported_program = exported_program + + @staticmethod + def _unwrap_view(node: Node) -> Node: + while ( + node.target == exir_ops.edge.aten.view_copy.default + and len(node.args) > 0 + and isinstance(node.args[0], Node) + ): + node = node.args[0] + return node + + def _get_scalar_constant(self, rescale_node: Node) -> Optional[torch.Tensor]: + source = rescale_node.args[0] + if not isinstance(source, Node): + return None + try: + value = get_param_tensor(self.exported_program, source) + except RuntimeError: + return None + if value is None or value.numel() != 1: + return None + return value + + @staticmethod + def _single_qargs(node: Node, key: str) -> Optional[QuantArgs]: + 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 _source_qargs( + source: Node, + rescale: _RescaleParams, + int32_qargs: QuantArgs, + ) -> Optional[QuantArgs]: + candidates = cast( + dict[int, QuantArgs], + source.meta.get("output_qparams", {}), + ) + for qargs in candidates.values(): + if qargs.per_channel or qargs.dtype != torch.int8: + continue + if qargs.get_zp_per_tensor() != rescale.input_zp: + continue + expected_scale = ( + qargs.get_scale_per_tensor() / int32_qargs.get_scale_per_tensor() + ) + if math.isclose(expected_scale, rescale.scale, rel_tol=1e-6, abs_tol=0.0): + return qargs + return None + + def _match(self, tanh: Node) -> Optional[_DyTMatch]: # noqa: C901 + if tanh.target != exir_ops.edge.aten.tanh.default: + return None + if len(tanh.args) != 1 or not isinstance(tanh.args[0], Node): + return None + + output_rescale_node = tanh.args[0] + output_rescale = _RescaleParams.from_node(output_rescale_node) + if ( + output_rescale is None + or output_rescale.output_dtype != torch.int8 + or len(output_rescale_node.users) != 1 + or not isinstance(output_rescale_node.args[0], Node) + ): + return None + + mul = output_rescale_node.args[0] + if ( + mul.target != exir_ops.edge.aten.mul.Tensor + or len(mul.args) < 2 + or len(mul.users) != 1 + ): + return None + + mul_qparams = cast(dict[int, QuantArgs], mul.meta.get("input_qparams", {})) + if len(mul_qparams) != 2: + return None + + operands: list[ + tuple[int, Node, Node, _RescaleParams, Optional[torch.Tensor], bool] + ] = [] + for index, arg in enumerate(mul.args[:2]): + if not isinstance(arg, Node): + return None + rescale_node = self._unwrap_view(arg) + rescale = _RescaleParams.from_node(rescale_node) + if rescale is None or rescale.output_dtype != torch.int32: + return None + source = rescale_node.args[0] + if not isinstance(source, Node): + return None + operands.append( + ( + index, + source, + rescale_node, + rescale, + self._get_scalar_constant(rescale_node), + rescale_node is not arg, + ) + ) + + scalar_operands = [operand for operand in operands if operand[4] is not None] + if len(scalar_operands) != 1: + return None + alpha_operand = scalar_operands[0] + activation_operand = next( + operand for operand in operands if operand != alpha_operand + ) + if activation_operand[5]: + return None + + activation_int32_qargs = mul_qparams.get(activation_operand[0]) + if ( + activation_int32_qargs is None + or activation_int32_qargs.dtype != torch.int32 + ): + return None + activation_qargs = self._source_qargs( + activation_operand[1], + activation_operand[3], + activation_int32_qargs, + ) + tanh_input_qargs = self._single_qargs(tanh, "input_qparams") + tanh_output_qargs = self._single_qargs(tanh, "output_qparams") + alpha_code = alpha_operand[4] + if ( + activation_qargs is None + or tanh_input_qargs is None + or tanh_output_qargs is None + or alpha_code is None + or tanh_input_qargs.dtype != torch.int8 + or tanh_output_qargs.dtype != torch.int8 + or output_rescale.output_zp != tanh_input_qargs.get_zp_per_tensor() + ): + return None + + return _DyTMatch( + tanh=tanh, + activation=activation_operand[1], + activation_qargs=activation_qargs, + activation_rescale=activation_operand[3], + alpha_code=alpha_code, + alpha_rescale=alpha_operand[3], + mul_output_rescale=output_rescale, + tanh_input_qargs=tanh_input_qargs, + tanh_output_qargs=tanh_output_qargs, + ) + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + for node in list(graph_module.graph.nodes): + match = self._match(node) + if match is None: + continue + + table = _generate_dyt_lut( + activation_qargs=match.activation_qargs, + alpha_code=match.alpha_code, + activation_rescale=match.activation_rescale, + alpha_rescale=match.alpha_rescale, + mul_output_rescale=match.mul_output_rescale, + tanh_input_qargs=match.tanh_input_qargs, + tanh_output_qargs=match.tanh_output_qargs, + ) + insert_pos = next(iter(graph_module.graph.nodes)) + with graph_module.graph.inserting_before(insert_pos): + table_constant = create_constant_placeholder( + exp_program=self.exported_program, + graph=graph_module.graph, + kind=InputKind.BUFFER, + name=f"b_{match.tanh.name}_dyt_table_constant", + data=table, + persistent_buffer=True, + ) + with graph_module.graph.inserting_before(match.tanh): + table_node = create_node( + graph=graph_module.graph, + op_target=exir_ops.backend.tosa.TABLE.default, + args=(match.activation, table_constant), + from_node=match.tanh, + ) + table_node.meta["input_qparams"] = {0: match.activation_qargs} + table_node.meta["output_qparams"] = {0: match.tanh_output_qargs} + match.tanh.replace_all_uses_with(table_node) + graph_module.graph.erase_node(match.tanh) + modified = True + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + return PassResult(graph_module, modified) + + +def register_fold_dyt_alpha_into_lut_pass() -> None: + """Enable exact DyT folding for subsequently lowered Arm programs.""" + # Imported here rather than at module scope: _passes/__init__.py imports + # arm_pass_manager last, and arm_pass_manager imports back from the package, + # so a module-level import from a globbed pass module is circular. + from executorch.backends.arm._passes.arm_pass_manager import ( + register_pass_factories_before, + ) + + register_pass_factories_before( + InsertTableOpsPass, + [FoldDyTAlphaIntoLUTPass], + ) 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..d7fdfa9f783 --- /dev/null +++ b/backends/arm/test/passes/test_fold_dyt_affine_into_conv_pass.py @@ -0,0 +1,796 @@ +# 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.""" + +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 _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 _build_post_rescale_fixture( + *, + table: torch.Tensor, + gamma_code: int, + beta_code: int, + actual_dyt_identity_qparams: bool, + padded_depthwise: bool, + slice_passthrough: bool = False, + channel_slice: bool = False, + shared_layout_user: bool = False, + input_width: int = 4, + input_channels: int = _CHANNELS, + affine_view_shape: tuple[int, ...] | None = None, + activation_view_shape: tuple[int, ...] | None = None, +) -> tuple[ExportedProgram, torch.Tensor, torch.Tensor]: + if padded_depthwise: + weight = torch.tensor( + [ + [[[1, -2, 1]]], + [[[2, 1, -1]]], + ], + dtype=torch.int8, + ) + bias = torch.tensor([5, -7], dtype=torch.int32) + 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) + bias = torch.tensor([5, -7], dtype=torch.int32) + else: + weight = torch.tensor( + [ + [[[1]], [[2]]], + [[[3]], [[-2]]], + ], + dtype=torch.int8, + ) + bias = torch.tensor([5, -7], dtype=torch.int32) + + test_input = ( + torch.arange(input_width * input_channels, dtype=torch.int8).reshape( + 1, 1, input_width, input_channels + ), + ) + exported_program = export( + _PostRescaleAffineFixture( + table=table, + gamma=torch.full((_CHANNELS,), gamma_code, dtype=torch.int8), + beta=torch.full((_CHANNELS,), beta_code, dtype=torch.int8), + 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)) + + if actual_dyt_identity_qparams: + table_qargs = _qargs(scale=0.00588326808065176, zp=-6) + gamma_input_zp = -128 + gamma_output_scale = 1.0 / 255.0 + gamma_output_zp = _zp_of(table_qargs) + beta_input_zp = -128 + beta_scale = 1.52587890625e-05 + common_scale = (2.0 * _scale_of(table_qargs)) / (1 << 20) + add_activation_scale = _scale_of(table_qargs) / common_scale + add_beta_scale = beta_scale / common_scale + add_output_scale = common_scale / _scale_of(table_qargs) + else: + table_qargs = _qargs(scale=0.1, zp=0) + gamma_input_zp = 0 + gamma_output_scale = 1.0 + gamma_output_zp = 0 + beta_input_zp = 0 + add_activation_scale = 1.0 + add_beta_scale = 1.0 + add_output_scale = 1.0 + + 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, gamma_output_zp), + ) + add_activation_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + ( + mul_output_rescale, + torch.int32, + [add_activation_scale], + 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, table_qargs.zp), + ) + nchw = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + (add_output_rescale, [0, 3, 1, 2]), + ) + conv_input = nchw + if slice_passthrough: + conv_input = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + (nchw, 0, 0, 1, 1), + ) + if channel_slice: + conv_input = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + (nchw, 1, 0, 1, 1), + ) + 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 else 1, + ), + ) + shared_output = None + if shared_layout_user: + shared_output = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (nchw, [1, _CHANNELS, 1, input_width]), + ) + + table_node.meta["output_qparams"] = {0: table_qargs} + table_node.meta["val"] = torch.empty( + (1, 1, input_width, input_channels), dtype=torch.int8, device="meta" + ) + add_output_rescale.meta["val"] = torch.empty( + (1, 1, input_width, _CHANNELS), dtype=torch.int8, device="meta" + ) + conv.meta["input_qparams"] = { + 0: table_qargs, + 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 + + +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=-128, + actual_dyt_identity_qparams=True, + padded_depthwise=True, + 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_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=-128, + 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) + + +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() diff --git a/backends/arm/test/passes/test_fold_dyt_alpha_into_lut_pass.py b/backends/arm/test/passes/test_fold_dyt_alpha_into_lut_pass.py new file mode 100644 index 00000000000..195b5d12ede --- /dev/null +++ b/backends/arm/test/passes/test_fold_dyt_alpha_into_lut_pass.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 +"""Tests for folding quantized DyT alpha multiplication into a tanh LUT. + +The tests prove the rewrite is byte-exact: the generated 256-entry TABLE is +compared against the real TOSA integer RESCALE/Mul/tanh path over every int8 +input code, so the fold cannot change quantized output on any model. + +""" + +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, +) +from executorch.backends.arm._passes.fold_dyt_alpha_into_lut_pass import ( + _generate_dyt_lut, + _RescaleParams, + 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.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export + + +class _PostRescaleFixture(torch.nn.Module): + # Declared so the checker sees the registered buffer as a Tensor rather than + # the ``Tensor | Module`` that ``nn.Module.__getattr__`` is annotated to give. + alpha_code: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "alpha_code", + torch.tensor([127], dtype=torch.int8), + ) + + def forward(self, x_code: torch.Tensor) -> torch.Tensor: + # Export a buffer and user input; the test replaces this placeholder op + # with the exact post-InsertRescaleInt32Pass topology consumed by the pass. + return x_code + self.alpha_code + + +_TEST_DATA = ( + torch.arange(-128, 128, dtype=torch.int16).to(torch.int8).reshape(1, 1, 16, 16), +) + + +def _qargs( + scale: float, + zp: int, + qmin: int | None = None, + qmax: int | None = None, + dtype: torch.dtype = torch.int8, +) -> QuantArgs: + dtype_range = torch.iinfo(dtype) + return QuantArgs( + scale=scale, + zp=zp, + qmin=dtype_range.min if qmin is None else qmin, + qmax=dtype_range.max if qmax is None else qmax, + 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) + + +_ACTIVATION_QARGS = _qargs(scale=0.015, zp=3) +_ALPHA_QARGS = _qargs(scale=0.0019607844296842813, zp=-128) +_TANH_INPUT_QARGS = _qargs(scale=0.0077, zp=2) +_TANH_OUTPUT_QARGS = _qargs(scale=0.0078, zp=0) +_ACTIVATION_RESCALE = _RescaleParams(1.0, 3, 0, torch.int32) +_ALPHA_RESCALE = _RescaleParams(1.0, -128, 0, torch.int32) +_MUL_OUTPUT_RESCALE = _RescaleParams( + (_scale_of(_ACTIVATION_QARGS) * _scale_of(_ALPHA_QARGS)) + / _scale_of(_TANH_INPUT_QARGS), + 0, + 2, + torch.int8, +) + + +def _build_post_rescale_fixture( + *, + activation_rank_view: bool = False, +) -> ExportedProgram: + exported_program = export(_PostRescaleFixture(), _TEST_DATA, strict=True) + graph = exported_program.graph_module.graph + alpha_name = next(iter(exported_program.graph_signature.inputs_to_buffers)) + alpha = next(node for node in graph.nodes if node.name == alpha_name) + activation = next( + node for node in graph.nodes if node.op == "placeholder" and node is not alpha + ) + original_add = next(node for node in graph.nodes if node.op == "call_function") + output = next(node for node in graph.nodes if node.op == "output") + + with graph.inserting_before(output): + activation_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (activation, torch.int32, [1.0], 3, 0), + ) + alpha_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + (alpha, torch.int32, [1.0], -128, 0), + ) + activation_mul_arg = activation_rescale + if activation_rank_view: + activation_mul_arg = graph.call_function( + exir_ops.edge.aten.view_copy.default, + (activation_rescale, [1, 1, 1, 16, 16]), + ) + mul = graph.call_function( + exir_ops.edge.aten.mul.Tensor, + (activation_mul_arg, alpha_rescale), + ) + mul_output_rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + ( + mul, + torch.int8, + [_MUL_OUTPUT_RESCALE.scale], + 0, + 2, + ), + ) + tanh = graph.call_function( + exir_ops.edge.aten.tanh.default, + (mul_output_rescale,), + ) + + activation.meta["output_qparams"] = {0: _ACTIVATION_QARGS} + mul.meta["input_qparams"] = { + 0: _qargs(_scale_of(_ACTIVATION_QARGS), 0, dtype=torch.int32), + 1: _qargs(_scale_of(_ALPHA_QARGS), 0, dtype=torch.int32), + } + tanh.meta["input_qparams"] = {0: _TANH_INPUT_QARGS} + tanh.meta["output_qparams"] = {0: _TANH_OUTPUT_QARGS} + output.replace_input_with(original_add, tanh) + graph.erase_node(original_add) + graph.lint() + exported_program.graph_module.recompile() + return exported_program + + +def _tosa_reference_outputs(domain: torch.Tensor) -> torch.Tensor: + spec = TosaSpecification.create_from_string("TOSA-1.0+INT") + with TosaLoweringContext(spec): + activation_i32 = exir_ops.backend.tosa.RESCALE.default( + domain, + torch.int32, + [_ACTIVATION_RESCALE.scale], + _ACTIVATION_RESCALE.input_zp, + _ACTIVATION_RESCALE.output_zp, + ) + alpha_i32 = exir_ops.backend.tosa.RESCALE.default( + torch.tensor([127], dtype=torch.int8), + torch.int32, + [_ALPHA_RESCALE.scale], + _ALPHA_RESCALE.input_zp, + _ALPHA_RESCALE.output_zp, + ) + product = activation_i32 * alpha_i32 + mul_codes = exir_ops.backend.tosa.RESCALE.default( + product, + torch.int8, + [_MUL_OUTPUT_RESCALE.scale], + _MUL_OUTPUT_RESCALE.input_zp, + _MUL_OUTPUT_RESCALE.output_zp, + ) + return _TANH_OUTPUT_QARGS.quantize_value( + torch.tanh(_TANH_INPUT_QARGS.dequantize_value(mul_codes)) + ).to(torch.int8) + + +def test_lut_preserves_intermediate_integer_rounding() -> None: + """The composed table matches Mul+RESCALE+tanh, not float alpha folding.""" + input_qargs = _qargs(scale=0.1, zp=3) + alpha_qargs = _qargs(scale=0.05, zp=0) + tanh_input_qargs = _qargs(scale=0.07, zp=-2) + tanh_output_qargs = _qargs(scale=0.006, zp=1) + alpha_code = torch.tensor([7], dtype=torch.int8) + + lut = _generate_dyt_lut( + activation_qargs=input_qargs, + alpha_code=alpha_code, + activation_rescale=_RescaleParams( + scale=1.0, + input_zp=_zp_of(input_qargs), + output_zp=0, + output_dtype=torch.int32, + ), + alpha_rescale=_RescaleParams( + scale=1.0, + input_zp=_zp_of(alpha_qargs), + output_zp=0, + output_dtype=torch.int32, + ), + mul_output_rescale=_RescaleParams( + scale=(_scale_of(input_qargs) * _scale_of(alpha_qargs)) + / _scale_of(tanh_input_qargs), + input_zp=0, + output_zp=_zp_of(tanh_input_qargs), + output_dtype=torch.int8, + ), + tanh_input_qargs=tanh_input_qargs, + tanh_output_qargs=tanh_output_qargs, + ) + + domain = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + alpha = alpha_qargs.dequantize_value(alpha_code) + naive_float_fold = tanh_output_qargs.quantize_value( + torch.tanh(input_qargs.dequantize_value(domain) * alpha) + ).to(torch.int8) + + assert lut.shape == (256,) + assert lut.dtype == torch.int8 + assert not torch.equal(lut, naive_float_fold) + + +def test_lut_clamps_narrowed_tanh_input_range() -> None: + activation_qargs = _qargs(scale=0.01, zp=0) + tanh_input_qargs = _qargs(scale=0.01, zp=0, qmin=-127, qmax=127) + tanh_output_qargs = _qargs(scale=0.01, zp=0) + identity_rescale = _RescaleParams(1.0, 0, 0, torch.int32) + + lut = _generate_dyt_lut( + activation_qargs=activation_qargs, + alpha_code=torch.tensor([1], dtype=torch.int8), + activation_rescale=identity_rescale, + alpha_rescale=identity_rescale, + mul_output_rescale=_RescaleParams(1.0, 0, 0, torch.int8), + tanh_input_qargs=tanh_input_qargs, + tanh_output_qargs=tanh_output_qargs, + ) + + domain = torch.arange(-128, 128, dtype=torch.int16).to(torch.int8) + effective_codes = domain.clamp( + tanh_input_qargs.qmin, + tanh_input_qargs.qmax, + ) + expected = tanh_output_qargs.quantize_value( + torch.tanh(tanh_input_qargs.dequantize_value(effective_codes)) + ).to(torch.int8) + + assert lut[0] == expected[0] + assert torch.equal(lut, expected) + + +def test_pass_removes_alpha_mul_and_materializes_one_table() -> None: + exported_program = _build_post_rescale_fixture() + result = FoldDyTAlphaIntoLUTPass(exported_program).call( + exported_program.graph_module + ) + + targets = [ + str(node.target) + for node in result.graph_module.graph.nodes + if node.op == "call_function" + ] + assert result.modified + assert sum("tosa.TABLE" in target for target in targets) == 1 + assert not any("aten.mul" in target for target in targets) + assert not any("aten.tanh" in target for target in targets) + assert not any("tosa.RESCALE" in target for target in targets) + + +def test_pass_rejects_activation_side_rank_view() -> None: + exported_program = _build_post_rescale_fixture(activation_rank_view=True) + result = FoldDyTAlphaIntoLUTPass(exported_program).call( + exported_program.graph_module + ) + + targets = [ + str(node.target) + for node in result.graph_module.graph.nodes + if node.op == "call_function" + ] + assert not result.modified + assert not any("tosa.TABLE" in target for target in targets) + assert any("aten.view_copy" in target for target in targets) + assert any("aten.mul" in target for target in targets) + assert any("aten.tanh" in target for target in targets) + + +def test_tosa_output_is_bit_exact_after_fold() -> None: + exported_program = _build_post_rescale_fixture() + result = FoldDyTAlphaIntoLUTPass(exported_program).call( + exported_program.graph_module + ) + table = next( + value + for name, value in exported_program.state_dict.items() + if "dyt_table_constant" in name + ) + domain = _TEST_DATA[0].flatten() + expected = _tosa_reference_outputs(domain) + table_outputs = table[(domain.to(torch.int16) + 128).to(torch.int64)] + + assert result.modified + assert torch.equal(table, expected) + assert torch.equal(table_outputs, expected) + + +class DyTModule(torch.nn.Module): + """A DyT site fed by a conv, i.e. the shape DyT takes in a real encoder. + + The conv matters: the pass recovers the activation's int8 quantization from + its producer's ``output_qparams``. Feeding the DyT straight off a graph input + leaves a bare ``quantize_per_tensor`` as the producer, which carries no + ``output_qparams``, and the pass then correctly declines to fold. + + """ + + test_data: ClassVar[Dict[str, Tuple[torch.Tensor]]] = { + "rand": (torch.rand(1, 3, 8, 8),), + } + + def __init__(self, alpha: float = 0.5) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 3, kernel_size=1) + self.alpha = torch.nn.Parameter(torch.tensor([alpha])) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # The permute mirrors the real DyT, which applies its affine in NHWC. + # It also keeps the alpha Mul adjacent to something other than the conv, + # so FoldScalarMulIntoConvPass does not absorb it before this pass runs. + y = torch.permute(self.conv(x), (0, 2, 3, 1)) + return torch.tanh(self.alpha * y) + + +@common.parametrize("test_data", DyTModule.test_data) +def test_fold_dyt_alpha_into_lut_tosa_INT(test_data: Tuple[torch.Tensor]) -> None: + """Pipeline-level counterpart to the exhaustive parity tests above. + + Those pin the numerics on hand-built post-InsertRescale IR. This one starts + from an nn.Module, runs the real quantization and rescale passes, and checks + the alpha Mul and the tanh are both replaced by a single TABLE. + + """ + pipeline = PassPipeline[Tuple[torch.Tensor]]( + DyTModule(), + test_data, + quantize=True, + ops_after_pass={ + "executorch_exir_dialects_backend__ops_tosa_TABLE_default": 1, + }, + ops_not_after_pass=[ + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_tanh_default", + ], + pass_list=[FoldAndAnnotateQParamsPass, InsertRescaleInt32Pass], + passes_with_exported_program=[FoldDyTAlphaIntoLUTPass], + ) + # 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. Output equivalence is not lost here: it is pinned + # exhaustively by ``test_tosa_output_is_bit_exact_after_fold`` above, which + # checks the generated TABLE against the reference RESCALE/Mul/tanh path + # over every int8 input code. + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run()