Arm backend: Add extension support for test pipelines

Adds the possibility to add extension to the TOSA specification
for the test pipeline.

Signed-off-by: Per Åstrand <[email protected]>
Change-Id: I64d8dcce3ea1bf6aa76230396690fd007974f792

Author: per

backends/arm/test/tester/test_pipeline.py

@@ -306,9 +306,14 @@ def __init__(
         rtol: float = 1e-03,
         qtol: int = 1,
         dynamic_shapes: Optional[Tuple[Any]] = None,
+        tosa_extensions: Optional[List[str]] = None,
     ):
+        if tosa_extensions is None:
+            tosa_extensions = []
         tosa_profiles = {
-            "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT"),
+            "1.0": TosaSpecification.create_from_string(
+                "TOSA-1.0+INT" + "".join([f"+{ext}" for ext in tosa_extensions])
+            ),
         }
         tosa_version = conftest.get_option("tosa_version")
 
@@ -406,9 +411,14 @@ def __init__(
         transform_passes: Optional[
             Union[Sequence[PassType], Dict[str, Sequence[PassType]]]
         ] = None,
+        tosa_extensions: Optional[List[str]] = None,
     ):
+        if tosa_extensions is None:
+            tosa_extensions = []
         tosa_profiles = {
-            "1.0": TosaSpecification.create_from_string("TOSA-1.0+FP"),
+            "1.0": TosaSpecification.create_from_string(
+                "TOSA-1.0+FP" + "".join([f"+{ext}" for ext in tosa_extensions])
+            ),
         }
         tosa_version = conftest.get_option("tosa_version")
 
@@ -655,10 +665,15 @@ def __init__(
         pass_functions: Optional[List[Callable]] = None,
         passes_with_exported_program: Optional[List[Type[ExportPass]]] = None,
         custom_path: str = None,
+        tosa_extensions: Optional[List[str]] = None,
     ):
+        if tosa_extensions is None:
+            tosa_extensions = []
         tosa_profiles = {
             "1.0": TosaSpecification.create_from_string(
-                "TOSA-1.0+" + ("INT" if quantize else "FP")
+                "TOSA-1.0+"
+                + ("INT" if quantize else "FP")
+                + "".join([f"+{ext}" for ext in tosa_extensions]),
             ),
         }
         tosa_version = conftest.get_option("tosa_version")
@@ -721,9 +736,14 @@ def __init__(
         module: torch.nn.Module,
         test_data: T,
         custom_path: str = None,
+        tosa_extensions: Optional[List[str]] = None,
     ):
+        if tosa_extensions is None:
+            tosa_extensions = []
         tosa_profiles = {
-            "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT"),
+            "1.0": TosaSpecification.create_from_string(
+                "TOSA-1.0+INT" + "".join([f"+{ext}" for ext in tosa_extensions]),
+            ),
         }
         tosa_version = conftest.get_option("tosa_version")
 
@@ -779,18 +799,23 @@ def __init__(
         custom_path: str = None,
         quantize: Optional[bool] = False,
         u55_subset: Optional[bool] = False,
+        tosa_extensions: Optional[List[str]] = None,
     ):
+        if tosa_extensions is None:
+            tosa_extensions = []
         tosa_profiles = {
-            "1.0": "TOSA-1.0+" + ("INT" if quantize else "FP"),
+            "1.0": TosaSpecification.create_from_string(
+                "TOSA-1.0+"
+                + ("INT" if quantize else "FP")
+                + ("+u55" if u55_subset and quantize else "")
+                + "".join([f"+{ext}" for ext in tosa_extensions]),
+            ),
         }
-        tosa_version = tosa_profiles[conftest.get_option("tosa_version")]
+        tosa_version = conftest.get_option("tosa_version")
 
-        if u55_subset and quantize:
-            tosa_version = f"{tosa_version}+u55"
+        tosa_spec = tosa_profiles[tosa_version]
 
-        compile_spec = common.get_tosa_compile_spec(
-            tosa_version, custom_path=custom_path
-        )
+        compile_spec = common.get_tosa_compile_spec(tosa_spec, custom_path=custom_path)
         super().__init__(
             module,
             test_data,
@@ -799,7 +824,7 @@ def __init__(
             [],
         )
 
-        if "INT" in tosa_version:
+        if tosa_spec.support_integer():
             self.add_stage(self.tester.quantize, pos=0)
 
         self.change_args("check_not.exir", [])
@@ -855,11 +880,16 @@ def __init__(
         transform_passes: Optional[
             Union[Sequence[PassType], Dict[str, Sequence[PassType]]]
         ] = None,
+        tosa_extensions: Optional[List[str]] = None,
     ):
 
-        tosa_profile = TosaSpecification.create_from_string(tosa_version)
+        if tosa_extensions is None:
+            tosa_extensions = []
+        tosa_spec = TosaSpecification.create_from_string(
+            tosa_version + "".join([f"+{ext}" for ext in tosa_extensions])
+        )
         compile_spec = common.get_vgf_compile_spec(
-            tosa_profile, compiler_flags=vgf_compiler_flags, custom_path=custom_path
+            tosa_spec, compiler_flags=vgf_compiler_flags, custom_path=custom_path
         )
 
         super().__init__(
@@ -873,7 +903,7 @@ def __init__(
             transform_passes=transform_passes,
         )
 
-        if "INT" in tosa_version:
+        if tosa_spec.support_integer():
             quantizer = VgfQuantizer(compile_spec)
             quantization_config = get_symmetric_quantization_config(
                 is_per_channel=per_channel_quantization