Arm backend: Add LSTM test for int16x8

Signed-off-by: Saoirse Stewart <[email protected]>

Author: SaoirseARM

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)
@@ -123,3 +131,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()