ai-edge-litert-nightly 2.2.0.dev20260102__cp312-cp312-manylinux_2_27_x86_64.whl

ai_edge_litert/aot/vendors/google_tensor/target.py

@@ -0,0 +1,84 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Compilation target for Google Tensor SOCs."""
+
+import dataclasses
+import sys
+from typing import Any
+
+from ai_edge_litert.aot.core import types
+
+# pylint: disable=g-importing-member
+# pylint: disable=g-import-not-at-top
+# pylint: disable=g-bad-import-order
+if sys.version_info >= (3, 11):
+  from enum import StrEnum  # pylint: disable=g-importing-member
+else:
+  from backports.strenum import StrEnum  # pylint: disable=g-importing-member
+# pylint: enable=g-bad-import-order
+# pylint: enable=g-import-not-at-top
+# pylint: enable=g-importing-member
+
+
+_GOOGLE_TENSOR_BACKEND_ID = "GOOGLE"
+
+
+class SocModel(StrEnum):
+  """Google Tensor SOC model."""
+
+  ALL = "ALL"
+
+  TENSOR_G3 = "Tensor_G3"
+  TENSOR_G4 = "Tensor_G4"
+  TENSOR_G5 = "Tensor_G5"
+  TENSOR_G6 = "Tensor_G6"
+
+
+class SocManufacturer(StrEnum):
+  """Google Tensor SOC manufacturer."""
+
+  GOOGLE = "Google"
+
+
+@dataclasses.dataclass
+class Target(types.Target):
+  """Compilation target for Google Tensor SOCs."""
+
+  soc_model: SocModel
+  soc_manufacturer: SocManufacturer = SocManufacturer.GOOGLE
+
+  @classmethod
+  def backend_id(cls) -> str:
+    return _GOOGLE_TENSOR_BACKEND_ID
+
+  def __hash__(self) -> int:
+    return hash((self.soc_manufacturer, self.soc_model))
+
+  def __eq__(self, other: "Target") -> bool:
+    return (
+        self.soc_manufacturer == other.soc_manufacturer
+        and self.soc_model == other.soc_model
+    )
+
+  def __repr__(self) -> str:
+    return f"{self.soc_manufacturer.value}_{self.soc_model.value}"
+
+  def flatten(self) -> dict[str, Any]:
+    flattend_target = super().flatten()
+    flattend_target.update({
+        "soc_manufacturer": self.soc_manufacturer.value,
+        "soc_model": self.soc_model.value,
+    })
+    return flattend_target

ai_edge_litert/aot/vendors/import_vendor.py

@@ -0,0 +1,132 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utility for dynamically importing vendor backends based on ID."""
+
+import copy
+import dataclasses
+from typing import Any, Iterable
+
+from ai_edge_litert.aot.core import types
+from ai_edge_litert.aot.vendors import fallback_backend
+
+
+@dataclasses.dataclass
+class VendorRegistry:
+  """A vendor registry."""
+
+  backend_class: types.BackendT
+
+
+_VENDOR_REGISTRY: dict[str, VendorRegistry] = {}
+
+
+def register_backend(
+    backend_class: types.BackendT,
+):
+  backend_id = backend_class.id()
+  _VENDOR_REGISTRY[backend_id] = VendorRegistry(backend_class)
+
+  return backend_class
+
+register_backend(fallback_backend.FallbackBackend)
+
+
+def import_vendor(backend_id: str) -> types.BackendT:
+  """Imports a vendor backend class based on its ID.
+
+  Args:
+    backend_id: The ID of the backend to import.
+
+  Returns:
+    The imported backend class.
+
+  Raises:
+    ValueError: If the backend ID is not supported.
+  """
+  vendor_module = _VENDOR_REGISTRY.get(backend_id, None)
+  if vendor_module is None:
+    raise ValueError(f'Unsupported backend id: {backend_id}')
+
+  return vendor_module.backend_class
+
+
+class AllRegisteredTarget(types.Target):
+  """A virtual Compilation target."""
+
+  def __hash__(self) -> int:
+    raise NotImplementedError()
+
+  def __eq__(self, other) -> bool:
+    raise NotImplementedError()
+
+  def __repr__(self) -> str:
+    raise NotImplementedError()
+
+  @classmethod
+  def backend_id(cls) -> str:
+    return 'all'
+
+  def flatten(self) -> dict[str, Any]:
+    return {'backend_id': self.backend_id()}
+
+
+@register_backend
+class AllRegisteredBackend(types.Backend):
+  """A virtual backend that represents all registered backends."""
+
+  # NOTE: Only initialize through "create".
+  def __init__(self, config: types.Config):
+    self._config = config
+
+  @classmethod
+  def create(cls, config: types.Config) -> 'AllRegisteredBackend':
+    return AllRegisteredBackend(config)
+
+  @classmethod
+  def id(cls) -> str:
+    return 'all'
+
+  @property
+  def target(self) -> AllRegisteredTarget:
+    return AllRegisteredTarget()
+
+  @property
+  def target_id(self) -> str:
+    return ''
+
+  @property
+  def config(self) -> types.Config:
+    return self._config
+
+  def call_component(
+      self,
+      input_model: types.Model,
+      output_model: types.Model,
+      component: types.Component,
+  ):
+    del input_model, output_model, component
+    raise NotImplementedError(
+        'AllRegisteredBackend does not support any component.'
+    )
+
+  def specialize(self) -> Iterable[types.Backend]:
+    for backend_id, vendor_module in _VENDOR_REGISTRY.items():
+      if backend_id == 'all':
+        continue
+      config = copy.deepcopy(self.config)
+      config['backend_id'] = backend_id
+      backend = vendor_module.backend_class.create(config)
+      yield from backend.specialize()

ai_edge_litert/aot/vendors/mediatek/mediatek_backend.py

@@ -0,0 +1,196 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Backend implementation for the example compiler plugin.."""
+
+import copy
+import functools
+import itertools
+import os
+import pathlib
+from typing import Iterable
+
+from ai_edge_litert.aot.core import common
+from ai_edge_litert.aot.core import components
+from ai_edge_litert.aot.core import types
+from ai_edge_litert.aot.vendors import import_vendor
+from ai_edge_litert.aot.vendors.mediatek import target as target_lib
+
+COMPILER_PLUGIN_LIB_PATH = pathlib.Path(
+    "vendors/mediatek/compiler/libLiteRtCompilerPlugin_MediaTek.so"
+)
+
+
+@import_vendor.register_backend
+class MediaTekBackend(types.Backend):
+  """Backend implementation for the example compiler plugin."""
+
+  def __init__(self, config: types.Config):
+    super().__init__(config)
+    self._compilation_config = config.get("compilation_config", None)
+
+  @property
+  def soc_manufacturer(self) -> target_lib.SocManufacturer:
+    return target_lib.SocManufacturer.MEDIATEK
+
+  @property
+  def soc_model(self) -> target_lib.SocModel:
+    return target_lib.SocModel(self.config.get("soc_model", "ALL"))
+
+  @property
+  def android_os_version(self) -> target_lib.AndroidOsVersion:
+    return target_lib.AndroidOsVersion(
+        self.config.get("android_os_version", "ALL")
+    )
+
+  @property
+  def target(self) -> target_lib.Target:
+    return target_lib.Target(
+        self.soc_model, self.soc_manufacturer, self.android_os_version
+    )
+
+  @property
+  def target_id(self) -> str:
+    return repr(self.target)
+
+  def specialize(self) -> Iterable["MediaTekBackend"]:
+    if (
+        self.soc_model != target_lib.SocModel.ALL
+        and self.android_os_version != target_lib.AndroidOsVersion.ALL
+    ):
+      yield self
+    else:
+      if self.soc_model == target_lib.SocModel.ALL:
+        soc_models = filter(
+            lambda x: x != target_lib.SocModel.ALL, target_lib.SocModel
+        )
+      else:
+        soc_models = [self.soc_model]
+      if self.android_os_version == target_lib.AndroidOsVersion.ALL:
+        android_os_versions = [
+            x
+            for x in target_lib.AndroidOsVersion
+            if x != target_lib.AndroidOsVersion.ALL
+        ]
+      else:
+        android_os_versions = [self.android_os_version]
+      for soc_model, android_os_version in itertools.product(
+          soc_models, android_os_versions
+      ):
+        new_config = copy.deepcopy(self.config)
+        new_config["soc_model"] = soc_model.value
+        new_config["android_os_version"] = android_os_version.value
+        yield self.create(new_config)
+
+  @classmethod
+  def id(cls) -> str:
+    return target_lib._MEDIATEK_BACKEND_ID  # pylint: disable=protected-access
+
+  @classmethod
+  def create(cls, config: types.Config) -> "MediaTekBackend":
+    if config.get("backend_id", "") != cls.id():
+      raise ValueError("Invalid backend id")
+    return cls(config)
+
+  @property
+  def quantize_recipe(self) -> str | None:
+    return self.config.get("quantize_recipe", None)
+
+  def call_component(
+      self,
+      input_model: types.Model,
+      output_model: types.Model,
+      component: types.Component,
+  ):
+    return _call_component(component, self, input_model, output_model)
+
+
+@functools.singledispatch
+def _call_component(
+    component: types.Component,
+    backend: MediaTekBackend,
+    unused_input_model: types.Model,
+    unused_output_model: types.Model,
+):
+  raise NotImplementedError(
+      f"{backend.id()} backend does not support"
+      f" {component.component_name} component."
+  )
+
+
+# TODO(toribiosteven): Translate SOC | OS version to the corresponding
+# MediaTek SDK version and pass to the plugin.
+@_call_component.register
+def _apply_plugin(
+    component: components.ApplyPluginT,
+    backend: MediaTekBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  """Calls the apply plugin component."""
+  try:
+    # If the plugin is not built from source (i.e. using ai_edge_litert wheel),
+    # we find the plugin library directory from the package path.
+    # Otherwise we use the default library path.
+    plugin_path = common.get_resource(COMPILER_PLUGIN_LIB_PATH)
+    lib_dir = os.path.dirname(plugin_path)
+
+    try:
+      # pytype: disable=import-error
+      import ai_edge_litert_sdk_mediatek  # pylint: disable=g-import-not-at-top
+      # pytype: enable=import-error
+
+      # TODO(weiyiw): Translate SOC | OS version to the corresponding
+      # MediaTek SDK version and pass to the plugin.
+      sdk_version = "v8"
+      sdk_libs_path = str(
+          ai_edge_litert_sdk_mediatek.path_to_sdk_libs(sdk_version)
+      )
+    except ImportError:
+      sdk_libs_path = None
+    extra_kwargs = {"libs": lib_dir, "sdk_libs_path": sdk_libs_path}
+  except FileNotFoundError:
+    extra_kwargs = {}
+  return component(
+      input_model,
+      output_model,
+      backend.soc_manufacturer,
+      backend.soc_model.lower(),
+      **extra_kwargs,
+  )
+
+
+@_call_component.register
+def _aie_quantizer(
+    component: components.AieQuantizerT,
+    backend: MediaTekBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(
+      input_model,
+      output_model,
+      quantization_recipe=backend.quantize_recipe,
+  )
+
+
+@_call_component.register
+def _mlir_transforms(
+    component: components.MlirTransformsT,
+    unused_backend: MediaTekBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(input_model, output_model, [])

ai_edge_litert/aot/vendors/mediatek/target.py

@@ -0,0 +1,94 @@
+"""Compilation target for MediaTek SOCs."""
+
+import dataclasses
+import sys
+from typing import Any
+
+from ai_edge_litert.aot.core import types
+
+# pylint: disable=g-importing-member
+# pylint: disable=g-import-not-at-top
+# pylint: disable=g-bad-import-order
+if sys.version_info >= (3, 11):
+  from enum import StrEnum  # pylint: disable=g-importing-member
+else:
+  from backports.strenum import StrEnum  # pylint: disable=g-importing-member
+# pylint: enable=g-bad-import-order
+# pylint: enable=g-import-not-at-top
+# pylint: enable=g-importing-member
+
+
+_MEDIATEK_BACKEND_ID = "mediatek"
+
+
+class SocModel(StrEnum):
+  """MediaTek SOC model."""
+
+  ALL = "ALL"
+
+  MT6853 = "MT6853"
+  MT6877 = "MT6877"
+  MT6878 = "MT6878"
+  MT6879 = "MT6879"
+  MT6886 = "MT6886"
+  MT6893 = "MT6893"
+  MT6895 = "MT6895"
+  MT6897 = "MT6897"
+  MT6983 = "MT6983"
+  MT6985 = "MT6985"
+  MT6989 = "MT6989"
+  MT6991 = "MT6991"
+  MT8171 = "MT8171"
+  MT8188 = "MT8188"
+  MT8189 = "MT8189"
+
+
+class SocManufacturer(StrEnum):
+  """MediaTek SOC manufacturer."""
+
+  MEDIATEK = "MediaTek"
+
+
+class AndroidOsVersion(StrEnum):
+  """Android OS version."""
+
+  ALL = "ALL"
+
+  ANDROID_15 = "ANDROID_15"
+
+
+@dataclasses.dataclass
+class Target(types.Target):
+  """Compilation target for MediaTek SOCs."""
+
+  soc_model: SocModel
+  soc_manufacturer: SocManufacturer = SocManufacturer.MEDIATEK
+  android_os_version: AndroidOsVersion = AndroidOsVersion.ANDROID_15
+
+  @classmethod
+  def backend_id(cls) -> str:
+    return _MEDIATEK_BACKEND_ID
+
+  def __hash__(self) -> int:
+    return hash(
+        (self.soc_manufacturer, self.soc_model, self.android_os_version)
+    )
+
+  def __eq__(self, other: "Target") -> bool:
+    return (
+        self.soc_manufacturer == other.soc_manufacturer
+        and self.soc_model == other.soc_model
+        and self.android_os_version == other.android_os_version
+    )
+
+  def __repr__(self) -> str:
+    return f"{self.soc_manufacturer.value}_{self.soc_model.value}_{self.android_os_version.value}"
+
+  def flatten(self) -> dict[str, Any]:
+    flattend_target = super().flatten()
+    flattend_target.update({
+        "soc_manufacturer": self.soc_manufacturer.value,
+        "soc_model": self.soc_model.value,
+        "android_os_version": self.android_os_version.value,
+    })
+    return flattend_target

ai_edge_litert/aot/vendors/qualcomm/qualcomm_backend.py

@@ -0,0 +1,161 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Backend implementation for the example compiler plugin.."""
+
+import copy
+import functools
+import os
+import pathlib
+from typing import Iterable
+
+from ai_edge_litert.aot.core import common
+from ai_edge_litert.aot.core import components
+from ai_edge_litert.aot.core import types
+from ai_edge_litert.aot.vendors import import_vendor
+from ai_edge_litert.aot.vendors.qualcomm import target as target_lib
+
+COMPILER_PLUGIN_LIB_PATH = pathlib.Path(
+    "vendors/qualcomm/compiler/libLiteRtCompilerPlugin_Qualcomm.so"
+)
+
+
+@import_vendor.register_backend
+class QualcommBackend(types.Backend):
+  """Backend implementation for the example compiler plugin."""
+
+  def __init__(self, config: types.Config):
+    super().__init__(config)
+    self._compilation_config = config.get("compilation_config", None)
+
+  @property
+  def soc_manufacturer(self) -> target_lib.SocManufacturer:
+    return target_lib.SocManufacturer.QUALCOMM
+
+  @property
+  def soc_model(self) -> target_lib.SocModel:
+    return target_lib.SocModel(self.config.get("soc_model", "ALL"))
+
+  @property
+  def target(self) -> target_lib.Target:
+    return target_lib.Target(self.soc_model, self.soc_manufacturer)
+
+  @property
+  def target_id(self) -> str:
+    return repr(self.target)
+
+  def specialize(self) -> Iterable["QualcommBackend"]:
+    if self.soc_model != target_lib.SocModel.ALL:
+      yield self
+    else:
+      for soc_model in target_lib.SocModel:
+        if soc_model != target_lib.SocModel.ALL:
+          new_config = copy.deepcopy(self.config)
+          new_config["soc_model"] = soc_model.value
+          yield self.create(new_config)
+
+  @classmethod
+  def id(cls) -> str:
+    return target_lib._QUALCOMM_BACKEND_ID  # pylint: disable=protected-access
+
+  @classmethod
+  def create(cls, config: types.Config) -> "QualcommBackend":
+    if config.get("backend_id", "") != cls.id():
+      raise ValueError("Invalid backend id")
+    return cls(config)
+
+  @property
+  def quantize_recipe(self) -> str | None:
+    return self.config.get("quantize_recipe", None)
+
+  def call_component(
+      self,
+      input_model: types.Model,
+      output_model: types.Model,
+      component: types.Component,
+  ):
+    return _call_component(component, self, input_model, output_model)
+
+
+@functools.singledispatch
+def _call_component(
+    component: types.Component,
+    backend: QualcommBackend,
+    unused_input_model: types.Model,
+    unused_output_model: types.Model,
+):
+  raise NotImplementedError(
+      f"{backend.id()} backend does not support"
+      f" {component.component_name} component."
+  )
+
+
+@_call_component.register
+def _apply_plugin(
+    component: components.ApplyPluginT,
+    backend: QualcommBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  """Calls the apply plugin component."""
+  try:
+    # If the plugin is not built from source (i.e. using ai_edge_litert wheel),
+    # we find the plugin library directory from the package path.
+    # Otherwise we use the default library path.
+    plugin_path = common.get_resource(COMPILER_PLUGIN_LIB_PATH)
+    lib_dir = os.path.dirname(plugin_path)
+
+    try:
+      # pytype: disable=import-error
+      import ai_edge_litert_sdk_qualcomm  # pylint: disable=g-import-not-at-top
+      # pytype: enable=import-error
+
+      sdk_libs_path = str(ai_edge_litert_sdk_qualcomm.path_to_sdk_libs())
+    except ImportError:
+      sdk_libs_path = None
+    extra_kwargs = {"libs": lib_dir, "sdk_libs_path": sdk_libs_path}
+  except FileNotFoundError:
+    extra_kwargs = {}
+  return component(
+      input_model,
+      output_model,
+      backend.soc_manufacturer,
+      backend.soc_model,
+      **extra_kwargs,
+  )
+
+
+@_call_component.register
+def _aie_quantizer(
+    component: components.AieQuantizerT,
+    backend: QualcommBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(
+      input_model,
+      output_model,
+      quantization_recipe=backend.quantize_recipe,
+  )
+
+
+@_call_component.register
+def _mlir_transforms(
+    component: components.MlirTransformsT,
+    unused_backend: QualcommBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(input_model, output_model, [])