From 0afdc7f3dfb66becd0fc274cc488456dc406c2bf Mon Sep 17 00:00:00 2001 From: Charles Greenberg Date: Thu, 20 Aug 2026 16:45:57 -0700 Subject: [PATCH] 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 | 355 ++++++++++++++++ .../test_fold_dyt_alpha_into_lut_pass.py | 385 ++++++++++++++++++ 4 files changed, 766 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 b62e76aef72..f1ec243bf04 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..6b072771f3c --- /dev/null +++ b/backends/arm/_passes/fold_dyt_alpha_into_lut_pass.py @@ -0,0 +1,355 @@ +# 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 + + +@dataclass(frozen=True) +class _MulOperand: + index: int + source: Node + rescale: _RescaleParams + scalar_constant: Optional[torch.Tensor] + had_view: bool + + +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_mul_operands( + self, mul: Node + ) -> Optional[tuple[_MulOperand, _MulOperand]]: + if len(mul.args) < 2: + return None + + operands = [] + 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( + _MulOperand( + index=index, + source=source, + rescale=rescale, + scalar_constant=self._get_scalar_constant(rescale_node), + had_view=rescale_node is not arg, + ) + ) + return operands[0], operands[1] + + def _match(self, tanh: Node) -> Optional[_DyTMatch]: + 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.users) != 1: + return None + + mul_qparams = cast(dict[int, QuantArgs], mul.meta.get("input_qparams", {})) + if len(mul_qparams) != 2: + return None + + operands = self._match_mul_operands(mul) + if operands is None: + return None + scalar_operands = [ + operand for operand in operands if operand.scalar_constant 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 is not alpha_operand + ) + if activation_operand.had_view: + return None + + activation_int32_qargs = mul_qparams.get(activation_operand.index) + if ( + activation_int32_qargs is None + or activation_int32_qargs.dtype != torch.int32 + ): + return None + activation_qargs = self._source_qargs( + activation_operand.source, + activation_operand.rescale, + 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.scalar_constant + 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.source, + activation_qargs=activation_qargs, + activation_rescale=activation_operand.rescale, + alpha_code=alpha_code, + alpha_rescale=alpha_operand.rescale, + 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()