Add a TOSA specification class

per · per · commit 2cd27d494b81 · 2024-11-08T08:28:40.000+01:00
Signed-off-by: Per Åstrand &lt;per.astrand@arm.com&gt;
Change-Id: I45374812daaa831f8084f6534b4903d14d6634b8
diff --git a/backends/arm/test/misc/test_tosa_spec.py b/backends/arm/test/misc/test_tosa_spec.py
@@ -0,0 +1,105 @@
+# Copyright 2024 Arm Limited and/or its 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.
+
+import unittest
+
+from executorch.backends.arm.tosa_specification import (
+    Tosa_0_80,
+    Tosa_1_00,
+    TosaSpecification,
+)
+
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from parameterized import parameterized
+
+test_valid_0_80_strings = [
+    "TOSA-0.80.0+BI",
+    "TOSA-0.80.0+MI+8k",
+    "TOSA-0.80.0+BI+u55",
+]
+test_valid_1_00_strings = [
+    "TOSA-1.00.0+INT+FP+fft",
+    "TOSA-1.00.0+FP+bf16+fft",
+    "TOSA-1.00.0+INT+int4+cf",
+    "TOSA-1.00.0+FP+cf+bf16+8k",
+    "TOSA-1.00.0+FP+INT+bf16+fft+int4+cf",
+    "TOSA-1.00.0+FP+INT+fft+int4+cf+8k",
+]
+
+test_valid_1_00_extensions = {
+    "INT": ["int16", "int4", "var", "cf"],
+    "FP": ["bf16", "fp8e4m3", "fp8e5m2", "fft", "var", "cf"],
+}
+
+test_invalid_strings = [
+    "TOSA-0.80.0+bi",
+    "TOSA-0.80.0",
+    "TOSA-0.80.0+8k",
+    "TOSA-0.80.0+BI+MI",
+    "TOSA-0.80.0+BI+U55",
+    "TOSA-1.00.0+fft",
+    "TOSA-1.00.0+fp+bf16+fft",
+    "TOSA-1.00.0+INT+INT4+cf",
+    "TOSA-1.00.0+BI",
+    "TOSA-1.00.0+FP+FP+INT",
+    "TOSA-1.00.0+FP+CF+bf16",
+    "TOSA-1.00.0+BF16+fft+int4+cf+INT",
+]
+
+test_compile_specs = [
+    ([CompileSpec("tosa_version", "TOSA-0.80.0+BI".encode())],),
+    ([CompileSpec("tosa_version", "TOSA-0.80.0+BI+u55".encode())],),
+    ([CompileSpec("tosa_version", "TOSA-1.00.0+INT".encode())],),
+]
+
+test_compile_specs_no_version = [
+    ([CompileSpec("other_key", "TOSA-0.80.0+BI".encode())],),
+    ([CompileSpec("other_key", "some_value".encode())],),
+]
+
+
+class TestTosaSpecification(unittest.TestCase):
+    """Tests the TOSA specification class"""
+
+    @parameterized.expand(test_valid_0_80_strings)
+    def test_version_string_0_80(self, version_string: str):
+        tosa_spec = TosaSpecification.create_from_string(version_string)
+        assert isinstance(tosa_spec, Tosa_0_80)
+        assert tosa_spec.profile in ["BI", "MI"]
+
+    @parameterized.expand(test_valid_1_00_strings)
+    def test_version_string_1_00(self, version_string: str):
+        tosa_spec = TosaSpecification.create_from_string(version_string)
+        assert isinstance(tosa_spec, Tosa_1_00)
+        assert [profile in ["INT", "FP"] for profile in tosa_spec.profiles].count(
+            True
+        ) > 0
+
+        for profile in tosa_spec.profiles:
+            assert [
+                e in test_valid_1_00_extensions[profile] for e in tosa_spec.extensions
+            ]
+
+    @parameterized.expand(test_invalid_strings)
+    def test_invalid_version_strings(self, version_string: str):
+        tosa_spec = None
+        with self.assertRaises(ValueError):
+            tosa_spec = TosaSpecification.create_from_string(version_string)
+
+        assert tosa_spec is None
+
+    @parameterized.expand(test_compile_specs)
+    def test_create_from_compilespec(self, compile_specs: list[CompileSpec]):
+        tosa_spec = TosaSpecification.create_from_compilespecs(compile_specs)
+        assert isinstance(tosa_spec, TosaSpecification)
+
+    @parameterized.expand(test_compile_specs_no_version)
+    def test_create_from_invalid_compilespec(self, compile_specs: list[CompileSpec]):
+        tosa_spec = None
+        with self.assertRaises(ValueError):
+            tosa_spec = TosaSpecification.create_from_compilespecs(compile_specs)
+
+        assert tosa_spec is None
diff --git a/backends/arm/tosa_specification.py b/backends/arm/tosa_specification.py
@@ -0,0 +1,226 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+#
+# Main implementation of AoT flow to partition and preprocess for Arm target
+# backends. Converts via TOSA as an intermediate form supported by AoT and
+# JIT compiler flows.
+#
+
+import re
+from typing import List
+
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from packaging.version import Version
+
+
+class TosaSpecification:
+    """
+    This class implements a representation of TOSA specification
+    (https://www.mlplatform.org/tosa/tosa_spec.html) with a version, a profile
+    (with extension) and a level (8k).
+    For 0.80 releases the profile is BI or MI, with u55 handled as an inofficial extension
+    For 1.00 releases the profile is INT or FP, and the extensions are for
+        INT: int16, int4, var, cf
+        FP: bf16, fp8e4m3, fp8e5m2, fft, var, cf
+
+    The TOSA specification is encoded in the string represenatation
+        TOSA-major.minor.patch+profile[+level][+extensions]
+
+    For 0.80 MI implies BI, while for 1.0 the profiles has to explicitely be specified.
+
+    Profiles are uppercase letters and extensions and level is lowercase.
+    """
+
+    version: Version
+
+    def support_integer(self) -> bool:
+        """
+        Returns true if any integer operations are supported for the specification.
+        """
+        raise NotImplementedError
+
+    def support_float(self) -> bool:
+        """
+        Returns true if any float operations are supported for the specification.
+        """
+        raise NotImplementedError
+
+    def __init__(self, version: Version):
+        self.version = version
+
+    @staticmethod
+    def create_from_compilespecs(
+        compile_specs: List[CompileSpec],
+    ) -> "TosaSpecification":
+        """
+        Search the CompileSpec list for 'tosa_version' and instantiate a
+        class from the found value or return None on failure.
+        """
+        for spec in compile_specs:
+            if spec.key == "tosa_version":
+                return TosaSpecification.create_from_string(spec.value.decode())
+        raise ValueError(
+            "No TOSA version key found in any of the supplied CompileSpecs"
+        )
+
+    @staticmethod
+    def create_from_string(repr: str) -> "TosaSpecification":
+        """
+        Creates a TOSA specification class from a string representation:
+        TOSA-0.80.0+MI
+        TOSA-0.80.0+BI+8k
+        TOSA-0.80.0+BI+u55   # Ethos-U55 extension to handle TOSA subset
+        TOSA-0.90.0+MI
+        TOSA-1.00.0+INT+FP+int4+cf
+        """
+
+        pattern = r"^(TOSA)-([\d.]+)\+(.+)$"
+        match = re.match(pattern, repr)
+        if match:
+            name = match.group(1)
+            version = Version(match.group(2))
+            extras = match.group(3).split("+")
+            if name != "TOSA":
+                raise ValueError(f"Malformed TOSA specification representation: {repr}")
+            match version:
+                case _ if version.major == 0 and version.minor == 80:
+                    return Tosa_0_80(version, extras)
+                case _ if version.major == 1 and version.minor == 0:
+                    return Tosa_1_00(version, extras)
+                case _:
+                    raise ValueError(f"Wrong TOSA version: {version} from {repr}")
+
+        raise ValueError(f"Failed to parse TOSA specification representation: {repr}")
+
+
+class Tosa_0_80(TosaSpecification):
+    profile: str
+    level_8k: bool
+    is_U55_subset: bool
+    available_profiles = ["BI", "MI"]  # MT is not defined
+
+    def __init__(self, version: Version, extras: List[str]):
+        super().__init__(version)
+        assert version >= Version("0.80") and version < Version("0.90")
+
+        # Check that we only have one profile in the extensions list
+        if [e in Tosa_0_80.available_profiles for e in extras].count(True) != 1:
+            raise ValueError(
+                f"Bad combination of extras: {extras}, more than one of {Tosa_0_80.available_profiles} found."
+            )
+
+        # The list contains one profile at most, so pick it
+        self.profile = [e for e in extras if e in Tosa_0_80.available_profiles][0]
+        extras.remove(self.profile)
+
+        self.level_8k = "8k" in extras
+        if self.level_8k:
+            extras.remove("8k")
+        self.is_U55_subset = "u55" in extras
+        if self.is_U55_subset:
+            extras.remove("u55")
+
+        if len(extras) > 0:
+            raise ValueError(f"Unhandled extras found: {extras}")
+
+    def __repr__(self):
+        extensions = ""
+        if self.level_8k:
+            extensions += "+8K"
+        if self.is_U55_subset:
+            extensions += "+u55"
+        return f"TOSA-{str(self.version)}+{self.profile}{extensions}"
+
+    def __hash__(self) -> int:
+        return hash(str(self.version) + self.profile)
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Tosa_0_80):
+            return (self.version == other.version) and (self.profile == other.profile)
+        return False
+
+    def support_integer(self):
+        return True
+
+    def support_float(self):
+        return self.profile == "MI"
+
+
+class Tosa_1_00(TosaSpecification):
+    profiles: List[str]
+    level_8k: bool
+    extensions: List[str]
+
+    available_profiles = ["INT", "FP"]
+    valid_extensions = {
+        "INT": ["int16", "int4", "var", "cf"],
+        "FP": ["bf16", "fp8e4m3", "fp8e5m2", "fft", "var", "cf"],
+    }
+
+    def __init__(self, version: Version, extras: List[str]):
+        super().__init__(version)
+
+        # Check that we have at least one profile in the extensions list
+        if [e in Tosa_1_00.available_profiles for e in extras].count(True) == 0:
+            raise ValueError(
+                f"No profile ({Tosa_1_00.available_profiles}) found in: {extras}."
+            )
+
+        # and not more than number of available profiles
+        if [e in Tosa_1_00.available_profiles for e in extras].count(True) > len(
+            Tosa_1_00.available_profiles
+        ):
+            raise ValueError(
+                f"Too many profiles ({Tosa_1_00.available_profiles}) found in: {extras}."
+            )
+
+        # The list contains one profile at least, so pick them
+        self.profiles = [e for e in extras if e in Tosa_1_00.available_profiles]
+        for p in self.profiles:
+            extras.remove(p)
+
+        self.level_8k = "8k" in extras
+        if self.level_8k:
+            extras.remove("8k")
+
+        combined_extensions = []
+        for p in self.profiles:
+            combined_extensions += Tosa_1_00.valid_extensions[p]
+
+        if not all(e in combined_extensions for e in extras):
+            raise ValueError(
+                f"Bad extensions for TOSA-{version}{self._get_profiles_string()}: {extras}"
+            )
+
+        # all the rest of the extras are handled extensions
+        self.extensions = extras
+
+    def _get_profiles_string(self) -> str:
+        return "".join(["+" + p for p in self.profiles])
+
+    def _get_extensions_string(self) -> str:
+        return "".join(["+" + e for e in self.extensions])
+
+    def __repr__(self):
+        return f"TOSA-{self.version}{self._get_profiles_string()}{self._get_profiles_string()}"
+
+    def __hash__(self) -> int:
+        return hash(str(self.version) + self._get_profiles_string())
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Tosa_1_00):
+            return (self.version == other.version) and (
+                self._get_profiles_string() == other._get_profiles_string()
+            )
+        return False
+
+    def support_integer(self):
+        return "INT" in self.profiles
+
+    def support_float(self):
+        return "FP" in self.profiles
