diff --git a/backends/arm/test/models/test_lstm_arm.py b/backends/arm/test/models/test_lstm_arm.py index d9691efab25..42744c151fc 100644 --- a/backends/arm/test/models/test_lstm_arm.py +++ b/backends/arm/test/models/test_lstm_arm.py @@ -5,9 +5,14 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) -from executorch.backends.arm.test import common +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, @@ -16,6 +21,9 @@ VgfPipeline, ) +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize + from torch.nn.quantizable.modules import rnn input_t = Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] # (h0, c0) @@ -134,3 +142,71 @@ def test_lstm_vgf_FP(): use_to_edge_transform_and_lower=True, ) pipeline.run() + + +def get_symmetric_a16w8_lstm_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization, epsilon=2**-16 + ) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization, epsilon=2**-16 + ), + ) + + +def test_lstm_16a8w_tosa_INT(): + """Test LSTM model with 16A8W quantization (16-bit activations, 8-bit weights)""" + + pipeline = TosaPipelineINT[input_t]( + TestLSTM.lstm, + TestLSTM.model_example_inputs, + aten_op=[], + exir_op=[], + per_channel_quantization=False, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args("quantize", get_symmetric_a16w8_lstm_quantizer()) + pipeline.run() + + +@pytest.mark.xfail( + reason="MLETORCH-1452: AssertionError: Output 0 does not match reference output." +) +@common.XfailIfNoCorstone300 +def test_lstm_16a8w_u55_INT(): + pipeline = EthosU55PipelineINT[input_t]( + TestLSTM.lstm, + TestLSTM.model_example_inputs, + aten_ops=[], + exir_ops=[], + use_to_edge_transform_and_lower=True, + ) + + pipeline.change_args("quantize", get_symmetric_a16w8_lstm_quantizer()) + pipeline.run() + + +@common.XfailIfNoCorstone320 +def test_lstm_16a8w_u85_INT(): + pipeline = EthosU85PipelineINT[input_t]( + TestLSTM.lstm, + TestLSTM.model_example_inputs, + aten_ops=[], + exir_ops=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.change_args("quantize", get_symmetric_a16w8_lstm_quantizer()) + pipeline.run()