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/fold_dyt_alpha_into_lut_pass.py b/backends/arm/_passes/fold_dyt_alpha_into_lut_pass.py new file mode 100644 index 00000000000..679da3bfad4 --- /dev/null +++ b/backends/arm/_passes/fold_dyt_alpha_into_lut_pass.py @@ -0,0 +1,337 @@ +# 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 +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) 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()