From 0cec5b554d850fe05ad66e305aadf4b619ae3226 Mon Sep 17 00:00:00 2001 From: Charles Greenberg Date: Thu, 20 Aug 2026 09:32:16 -0700 Subject: [PATCH 1/2] Fold DyT alpha scalar into tanh LUT Summary: Dynamic Tanh (DyT) normalization computes `tanh(alpha * x)` with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. On Ethos-U the TABLE is free but the Mul is a real per-element cost, so the Mul is pure overhead. This adds `FoldDyTAlphaIntoLUTPass` in `backends/arm/_passes/`, which folds the alpha multiply into the tanh lookup table and deletes the Mul and its surrounding rescales. Folding alpha in floating point before quantization would change rounding, so the pass does it in the integer domain instead: it replays the exact TOSA SINGLE_ROUND RESCALE and Mul arithmetic over all 256 int8 input codes, feeds the results through the existing tanh quantization mapping, and materializes the result as one 256-entry TABLE. The rewrite is therefore byte-exact rather than approximate. It fails closed on anything it cannot prove: non-scalar alpha, activation-side rank views, and narrowed tanh ranges are all handled explicitly. Also adds `register_pass_factories_before()` to the Arm pass manager, a small hook for inserting a pass that needs access to the `ExportedProgram` ahead of a named target pass. `FoldDyTAlphaIntoLUTPass` needs it to read constant tensors. Note on one import: the pass imports `register_pass_factories_before` inside `register_fold_dyt_alpha_into_lut_pass()` 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-scope import from a pass module that `__init__.py` re-exports is circular. Only the register helper needs it. The pass is inert until a model registers it, so this diff changes no behaviour on its own. Differential Revision: D116560649 --- backends/arm/_passes/__init__.py | 3 + backends/arm/_passes/arm_pass_manager.py | 24 +- .../_passes/fold_dyt_alpha_into_lut_pass.py | 341 ++++++++++++++++ .../test_fold_dyt_alpha_into_lut_pass.py | 385 ++++++++++++++++++ 4 files changed, 752 insertions(+), 1 deletion(-) create mode 100644 backends/arm/_passes/fold_dyt_alpha_into_lut_pass.py create mode 100644 backends/arm/test/passes/test_fold_dyt_alpha_into_lut_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 71df953a65b..5a4dde37f17 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,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_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_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_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() From bef6d1d31a832979644a33fd2417108813c3fb44 Mon Sep 17 00:00:00 2001 From: Charles Greenberg Date: Thu, 20 Aug 2026 09:32:16 -0700 Subject: [PATCH 2/2] Fold DyT affine maps into following convolutions (#21953) Summary: Second half of the Dynamic Tanh (DyT) lowering cost, and a companion to `FoldDyTAlphaIntoLUTPass`. Once the alpha multiply has been folded into the tanh TABLE, a DyT site still emits a per-channel gamma Mul and a per-channel beta Add between the TABLE and the convolution that consumes it. Both are full-tensor elementwise ops on Ethos-U. This adds `FoldDyTAffineIntoConvPass` in `backends/arm/_passes/`, which folds that affine into the weights and bias of the following convolution, the same algebra as BatchNorm folding: `conv(gamma * x + beta) == conv_with_scaled_weights(x) + conv(beta)`. Doing this in floating point is not safe here, because the intermediate INT8 requantization between the affine and the convolution is nonlinear, so a float fold can change rounding. The pass instead evaluates the site's real TOSA integer path over the materialized 256-entry TABLE and only rewrites when the resulting per-channel map is provably exactly integer-affine. Everything else fails closed: saturating or nonlinear maps, unsupported constant layouts, non-exclusive passthrough edges, rank or shape mismatches, and singleton-channel broadcast are all rejected rather than approximated. Padded convolutions are a special case. Gamma still folds exactly, but beta does not: with constant padding its contribution becomes position dependent at the boundary and cannot be represented by a single conv bias. For those sites the pass folds gamma only when gamma is an exact identity and leaves the beta Add in place. Also exposes a small hook in `insert_table_ops` that the fold needs to locate the materialized TABLE. As with the alpha fold, `register_pass_factories_before` is imported inside `register_fold_dyt_affine_into_conv_pass()` rather than at module scope, to avoid a circular import through `_passes/__init__.py`. The pass is inert until a model registers it, so this diff changes no behaviour on its own. Differential Revision: D116573000 --- backends/arm/_passes/__init__.py | 1 + .../_passes/fold_dyt_affine_into_conv_pass.py | 675 +++++++++++++++ .../test_fold_dyt_affine_into_conv_pass.py | 796 ++++++++++++++++++ 3 files changed, 1472 insertions(+) create mode 100644 backends/arm/_passes/fold_dyt_affine_into_conv_pass.py create mode 100644 backends/arm/test/passes/test_fold_dyt_affine_into_conv_pass.py 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..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/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()