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

ai_edge_litert/aot/prepare_for_npu.py

@@ -0,0 +1,152 @@
+# 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.
+# ==============================================================================
+
+"""Implementations for the main public API functionalities."""
+
+import pathlib
+from typing import cast
+
+# pylint: disable=g-import-not-at-top
+# pytype: disable=import-error
+try:
+  from tqdm import auto as autotqdm
+except ImportError:
+  from tqdm.tqdm import auto as autotqdm
+# pytype: enable=import-error
+
+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
+
+# pylint: enable=g-import-not-at-top
+
+
+def resolve_backend(config: types.Config) -> types.BackendT:
+  # Import the backend based on the ID.
+  backend_id = config.get("backend_id", None)
+  if backend_id is None:
+    raise ValueError("Backend ID is required.")
+  return import_vendor.import_vendor(backend_id)
+
+
+def prepare_for_npu_multiple_configs(
+    flatbuffer: types.Model,
+    output_dir: pathlib.Path,
+    configs: list[tuple[types.BackendT, types.Config]],
+    plugin: components.ApplyPluginT,
+    transforms: components.MlirTransformsT | None = None,
+    quantizer: components.AieQuantizerT | None = None,
+    keep_going: bool = False,
+) -> types.CompilationResult:
+  """Prepares a TFLite model for NPU execution."""
+  backends = []
+  for backend_class, config in configs:
+    backend = backend_class.create(config)
+    backends += list(backend.specialize())
+
+  pipeline: list[types.Component] = [
+      c for c in [transforms, quantizer, plugin] if c is not None
+  ]
+  return compile_model(flatbuffer, output_dir, backends, pipeline, keep_going)
+
+
+def prepare_for_npu(
+    flatbuffer: types.Model,
+    output_dir: pathlib.Path,
+    backend_class: types.BackendT,
+    config: types.Config,
+    plugin: components.ApplyPluginT,
+    transforms: components.MlirTransformsT | None = None,
+    quantizer: components.AieQuantizerT | None = None,
+    keep_going: bool = False,
+) -> types.CompilationResult:
+  """Prepares a TFLite model for NPU execution.
+
+  High level command that erforms various backend specific pre-processing steps
+  and then applies an NPU compiler to the given model.
+
+  Args:
+      flatbuffer: Path to the input flatbuffer file.
+      output_dir: Directory to write the output flatbuffer file.
+      backend_class: The backend to prepare the model for.
+      config: The configuration for the backend.
+      plugin: The plugin to apply to the model.
+      transforms: The transforms to apply to the model.
+      quantizer: The quantizer to apply to the model.
+      keep_going: Whether to keep going if some backends fail.
+
+  Returns:
+      List of the paths to the output flatbuffer file.
+
+  Raises:
+      ValueError: If the given path is not a valid flatbuffer file.
+  """
+
+  backend = backend_class.create(config)
+
+  pipeline: list[types.Component] = [
+      c for c in [transforms, quantizer, plugin] if c is not None
+  ]
+  backends = list(backend.specialize())
+  return compile_model(flatbuffer, output_dir, backends, pipeline, keep_going)
+
+
+def compile_model(
+    flatbuffer: types.Model,
+    output_dir: pathlib.Path,
+    backends: list[types.Backend],
+    pipeline: list[types.Component],
+    keep_going: bool = False,
+) -> types.CompilationResult:
+  """Compiles a TFLite model for NPU execution."""
+  if flatbuffer.in_memory:
+    base_name = "model"
+  else:
+    base_name = flatbuffer.path.name.removesuffix(common.DOT_TFLITE)
+  compile_models = types.CompilationResult()
+  with autotqdm.tqdm(backends, desc="Backend") as t_backends:
+    for backend in t_backends:
+      component_input = flatbuffer
+      backend = cast(types.Backend, backend)
+      input_name_pref = base_name + backend.target_id_suffix
+      t_backends.set_description(f"Compiling {backend.target_id}")
+      try:
+        for component in pipeline:
+          component = cast(types.Component, component)
+          t_backends.set_description(
+              f"Compiling {backend.target_id}: {component.component_name}"
+          )
+          component_output = types.Model.create_from_path(
+              output_dir
+              / f"{input_name_pref}_{component.component_name}{common.DOT_TFLITE}"
+          )
+          backend.call_component(component_input, component_output, component)
+          if not component_output.in_memory and not common.is_tflite(
+              component_output.path
+          ):
+            raise ValueError(
+                f"{component.component_name} failed to produce a TFLite model."
+            )
+          component_input = component_output
+        compile_models.models_with_backend.append((backend, component_input))
+      except ValueError as e:
+        if keep_going:
+          print(f"Skipping failed compilation for {backend.target}. Error: {e}")
+          compile_models.failed_backends.append((backend, str(e)))
+        else:
+          raise
+
+  return compile_models

ai_edge_litert/aot/vendors/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+# ==============================================================================
+"""Vendor backends for LiteRt."""
+import os
+
+from ai_edge_litert.aot.vendors.mediatek import mediatek_backend as _
+from ai_edge_litert.aot.vendors.qualcomm import qualcomm_backend as _
+
+if os.environ.get("GOOGLE_TENSOR_COMPILER_LIB") is not None:
+  from ai_edge_litert.aot.vendors.google_tensor import google_tensor_backend as _  # pylint: disable=g-import-not-at-top

ai_edge_litert/aot/vendors/example/example_backend.py

@@ -0,0 +1,157 @@
+# 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 functools
+from typing import Any
+
+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
+
+
+class ExampleTarget(types.Target):
+  """Compilation target for the example backend."""
+
+  def __init__(self, soc_manufacturer: str, soc_model: str):
+    self.soc_manufacturer = soc_manufacturer
+    self.soc_model = soc_model
+
+  def __hash__(self) -> int:
+    return hash((self.soc_manufacturer, self.soc_model))
+
+  def __eq__(self, other) -> bool:
+    return (
+        self.soc_manufacturer == other.soc_manufacturer
+        and self.soc_model == other.soc_model
+    )
+
+  def __repr__(self) -> str:
+    return f"{self.soc_manufacturer}_{self.soc_model}"
+
+  def flatten(self) -> dict[str, Any]:
+    return {
+        "soc_manufacturer": self.soc_manufacturer,
+        "soc_model": self.soc_model,
+    }
+
+  @classmethod
+  def backend_id(cls) -> str:
+    return "example"
+
+
+# Note this is not a real target so not auto-registered unless the module is
+# imported.
+@import_vendor.register_backend
+class ExampleBackend(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)
+
+  @classmethod
+  def target_(cls) -> ExampleTarget:
+    return ExampleTarget("ExampleSocManufacturer", "ExampleSocModel")
+
+  @property
+  def target(self) -> ExampleTarget:
+    return self.target_()
+
+  @classmethod
+  def soc_manufacturer(cls) -> str:
+    return cls.target_().soc_manufacturer
+
+  @classmethod
+  def soc_model(cls) -> str:
+    return cls.target_().soc_model
+
+  @classmethod
+  def id(cls) -> str:
+    return "example"
+
+  @property
+  def target_id(self) -> str:
+    return ""
+
+  @property
+  def shared_pass_names(self) -> list[str]:
+    return ["example-pass"]
+
+  @classmethod
+  def create(cls, config: types.Config) -> "ExampleBackend":
+    if config.get("backend_id", "") != cls.id():
+      raise ValueError("Invalid backend id")
+    return cls(config)
+
+  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: ExampleBackend,
+    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: ExampleBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(
+      input_model,
+      output_model,
+      backend.soc_manufacturer,
+      backend.soc_model,
+  )
+
+
+@_call_component.register
+def _aie_quantizer(
+    component: components.AieQuantizerT,
+    unused_backend: ExampleBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(
+      input_model,
+      output_model,
+  )
+
+
+@_call_component.register
+def _mlir_transforms(
+    component: components.MlirTransformsT,
+    backend: ExampleBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(input_model, output_model, backend.shared_pass_names)

ai_edge_litert/aot/vendors/fallback_backend.py

@@ -0,0 +1,128 @@
+# 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.
+# ==============================================================================
+
+"""A Fallback backend for LITERT."""
+
+import functools
+from typing import Any
+
+from ai_edge_litert.aot.core import components
+from ai_edge_litert.aot.core import types
+
+
+class FallbackTarget(types.Target):
+  """A virtual Compilation target."""
+
+  def __hash__(self) -> int:
+    return hash(self.backend_id())
+
+  def __eq__(self, other: types.Target) -> bool:
+    return self.backend_id() == other.backend_id()
+
+  def __repr__(self) -> str:
+    return f"{self.backend_id()}"
+
+  @classmethod
+  def backend_id(cls) -> str:
+    return "fallback"
+
+  def flatten(self) -> dict[str, Any]:
+    return {"backend_id": self.backend_id()}
+
+
+class FallbackBackend(types.Backend):
+  """Fallback backend for LITERT."""
+
+  @property
+  def target(self) -> FallbackTarget:
+    return FallbackTarget()
+
+  @property
+  def target_id(self) -> str:
+    return repr(self.target)
+
+  @classmethod
+  def id(cls) -> str:
+    return "fallback"
+
+  @classmethod
+  def create(cls, config: types.Config) -> "FallbackBackend":
+    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: FallbackBackend,
+    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: FallbackBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  """A no-op component that just copies the input model to the output model."""
+  del component, backend
+  if input_model.in_memory:
+    output_model.set_bytes(input_model.model_bytes)
+  else:
+    output_model.set_path(input_model.path)
+
+
+@_call_component.register
+def _aie_quantizer(
+    component: components.AieQuantizerT,
+    backend: FallbackBackend,
+    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: FallbackBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(input_model, output_model, [])

ai_edge_litert/aot/vendors/google_tensor/google_tensor_backend.py

@@ -0,0 +1,168 @@
+# 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.
+
+"""Backend implementation for the Google Tensor 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.google_tensor import target as target_lib
+
+COMPILER_PLUGIN_LIB_PATH = pathlib.Path(
+    "vendors/google_tensor/compiler/libLiteRtCompilerPlugin_google_tensor.so"
+)
+
+
+# Returns true if the flag is a google_tensor flag.
+def _is_google_tensor_flag(flag: str) -> bool:
+  return flag.startswith("google_tensor_")
+
+
+@import_vendor.register_backend
+class GoogleTensorBackend(types.Backend):
+  """Backend implementation for the Google Tensor 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.GOOGLE
+
+  @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["GoogleTensorBackend"]:
+    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._GOOGLE_TENSOR_BACKEND_ID  # pylint: disable=protected-access
+
+  @classmethod
+  def create(cls, config: types.Config) -> "GoogleTensorBackend":
+    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: GoogleTensorBackend,
+    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: GoogleTensorBackend,
+    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)
+    sdk_libs_dir = os.environ.get("GOOGLE_TENSOR_COMPILER_LIB", None)
+
+    extra_kwargs = {"libs": lib_dir, "sdk_libs_path": sdk_libs_dir}
+  except FileNotFoundError:
+    extra_kwargs = {}
+
+  # Add google_tensor specific flags from the backend config.
+  for flag, value in backend.config.items():
+    if _is_google_tensor_flag(flag):
+      extra_kwargs[flag] = value
+
+  for flag, value in backend.config.get("compilation_config", {}).items():
+    if _is_google_tensor_flag(flag):
+      extra_kwargs[flag] = value
+
+  return component(
+      input_model,
+      output_model,
+      backend.soc_manufacturer,
+      backend.soc_model,
+      **extra_kwargs,
+  )
+
+
+@_call_component.register
+def _aie_quantizer(
+    component: components.AieQuantizerT,
+    backend: GoogleTensorBackend,
+    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: GoogleTensorBackend,
+    input_model: types.Model,
+    output_model: types.Model,
+):
+  return component(input_model, output_model, [])