Arm backend: Dynamic shape lowering for NSS & NFRU (pytorch#23307)

Change-Id: I06477ff9d1641d4ce0dbd6fbb62fd86258866c30

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Michiel Olieslagers <michiel.olieslagers@arm.com>

Author: Michiel-Olieslagers

backends/arm/test/models/test_nfru.py

@@ -38,6 +38,7 @@
 from ng_model_gym.usecases.nfru.model.nfru_v1_nn import (  # type: ignore[import-not-found,import-untyped]
     NFRUAutoEncoder,
 )
+from torch.export import Dim
 from torchao.quantization.pt2e import (
     allow_exported_model_train_eval,
     move_exported_model_to_eval,
@@ -47,6 +48,9 @@
 input_t = Tuple[torch.Tensor]  # Input x
 
 pytestmark = skip_if_frozen_release("NFRU")
+_NFRU_HEIGHT = 2 * Dim("_nfru_height", min=16, max=135)
+_NFRU_WIDTH = 2 * Dim("_nfru_width", min=16, max=240)
+_NFRU_DYNAMIC_SHAPES = ({2: _NFRU_HEIGHT, 3: _NFRU_WIDTH},)
 
 
 def nfru() -> NFRUAutoEncoder:
@@ -73,7 +77,7 @@ def nfru() -> NFRUAutoEncoder:
     return model
 
 
-def prequantized_nfru(inputs: input_t) -> torch.fx.GraphModule:
+def prequantized_nfru(inputs: input_t, dynamic_shapes=None) -> torch.fx.GraphModule:
     weights = hf_hub_download(  # nosec B615
         repo_id="Arm/neural-frame-rate-upscaling",
         filename="nfru_v1_int8.pt",
@@ -88,7 +92,9 @@ def prequantized_nfru(inputs: input_t) -> torch.fx.GraphModule:
     assert all(key.startswith(prefix) for key in checkpoint)
     state_dict = {key.removeprefix(prefix): value for key, value in checkpoint.items()}
 
-    exported = torch.export.export(nfru().eval(), inputs, strict=True).module()
+    exported = torch.export.export(
+        nfru().eval(), inputs, dynamic_shapes=dynamic_shapes, strict=True
+    ).module()
     quantizer = TOSAQuantizer(TosaSpecification.create_from_string("TOSA-1.0+INT"))
     quantizer.set_global(
         get_symmetric_quantization_config(is_per_channel=False, is_qat=True)
@@ -141,7 +147,10 @@ def random_inputs(memory_format: torch.memory_format = torch.channels_last):
 
 
 def _set_nfru_calibration_samples(pipeline):
-    return pipeline.set_quantization_calibration(iter_nfru_test_calibration_samples())
+    return pipeline.set_quantization_calibration(
+        iter_nfru_test_calibration_samples(),
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+    )
 
 
 @common.parametrize("use_real_data", input_test_data)
@@ -306,3 +315,182 @@ def test_nfru_vgf_quant_a16w8(use_real_data):
     if use_real_data:
         _set_nfru_calibration_samples(pipeline)
     pipeline.run()
+
+
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_tosa_FP_dynamic_shapes(use_real_data):
+    pipeline = TosaPipelineFP[input_t](
+        nfru().eval(),
+        example_inputs() if use_real_data else random_inputs(),
+        aten_op=[],
+        exir_op=[],
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_tosa_ref_model=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("is_qat", is_qat_test_data)
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_tosa_INT_dynamic_shapes(use_real_data, is_qat):
+    pipeline_kwargs = (
+        {
+            "per_channel_quantization": False,
+            "use_to_edge_transform_and_lower": True,
+            "frobenius_threshold": None,
+            "cosine_threshold": None,
+        }
+        if is_qat
+        else {}
+    )
+    pipeline = TosaPipelineINT[input_t](
+        nfru().eval(),
+        example_inputs() if use_real_data else random_inputs(),
+        aten_op=[],
+        exir_op=[],
+        atol=0.2,
+        qtol=2 if use_real_data else 1,
+        is_qat=is_qat,
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_tosa_ref_model=False,
+        **pipeline_kwargs,
+    )
+    if use_real_data:
+        _set_nfru_calibration_samples(pipeline)
+    pipeline.run()
+
+
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_prequantized_tosa_INT_dynamic_shapes(use_real_data):
+    inputs = example_inputs() if use_real_data else random_inputs()
+    pipeline = TosaPipelineINT[input_t](
+        prequantized_nfru(inputs, _NFRU_DYNAMIC_SHAPES),
+        inputs,
+        aten_op=[],
+        exir_op=[],
+        atol=0.2,
+        qtol=2 if use_real_data else 1,
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_tosa_ref_model=False,
+    )
+    pipeline.pop_stage("quantize")
+    pipeline.pop_stage("check.quant_nodes")
+    pipeline.add_stage_after(
+        "export",
+        pipeline.tester.check,
+        ["torch.ops.quantized_decomposed.dequantize_per_tensor.default"],
+        suffix="prequant_nodes",
+    )
+    pipeline.run()
+
+
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_tosa_INT_a16w8_dynamic_shapes(use_real_data):
+    pipeline = TosaPipelineINT[input_t](
+        nfru().eval(),
+        example_inputs() if use_real_data else random_inputs(),
+        aten_op=[],
+        exir_op=[],
+        tosa_extensions=["int16"],
+        atol=0.1,
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_tosa_ref_model=False,
+    )
+    if use_real_data:
+        _set_nfru_calibration_samples(pipeline)
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_vgf_no_quant_dynamic_shapes(use_real_data):
+    pipeline = VgfPipeline[input_t](
+        nfru().eval(),
+        example_inputs() if use_real_data else random_inputs(),
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+FP",
+        quantize=False,
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_vulkan_runtime=False,
+    )
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+@common.parametrize("is_qat", is_qat_test_data)
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_vgf_quant_dynamic_shapes(use_real_data, is_qat):
+    pipeline_kwargs = (
+        {
+            "per_channel_quantization": False,
+            "use_to_edge_transform_and_lower": True,
+        }
+        if is_qat
+        else {}
+    )
+    pipeline = VgfPipeline[input_t](
+        nfru().eval(),
+        example_inputs() if use_real_data else random_inputs(),
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+INT",
+        symmetric_io_quantization=True,
+        quantize=True,
+        is_qat=is_qat,
+        atol=0.2,
+        qtol=2 if use_real_data else 1,
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_vulkan_runtime=False,
+        **pipeline_kwargs,
+    )
+    if use_real_data:
+        _set_nfru_calibration_samples(pipeline)
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_prequantized_vgf_INT_dynamic_shapes(use_real_data):
+    inputs = example_inputs() if use_real_data else random_inputs()
+    pipeline = VgfPipeline[input_t](
+        prequantized_nfru(inputs, _NFRU_DYNAMIC_SHAPES),
+        inputs,
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+INT",
+        quantize=True,
+        atol=0.2,
+        qtol=2 if use_real_data else 1,
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_vulkan_runtime=False,
+    )
+    pipeline.pop_stage("quantize")
+    pipeline.pop_stage("check.quant_nodes")
+    pipeline.add_stage_after(
+        "export",
+        pipeline.tester.check,
+        ["torch.ops.quantized_decomposed.dequantize_per_tensor.default"],
+        suffix="prequant_nodes",
+    )
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+@common.parametrize("use_real_data", input_test_data)
+def test_nfru_vgf_quant_a16w8_dynamic_shapes(use_real_data):
+    pipeline = VgfPipeline[input_t](
+        nfru().eval(),
+        example_inputs() if use_real_data else random_inputs(),
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+INT",
+        tosa_extensions=["int16"],
+        symmetric_io_quantization=True,
+        atol=0.2,
+        dynamic_shapes=_NFRU_DYNAMIC_SHAPES,
+        run_on_vulkan_runtime=False,
+    )
+    if use_real_data:
+        _set_nfru_calibration_samples(pipeline)
+    pipeline.run()