ai-edge-litert-nightly 1.4.0.dev20250813__cp310-cp310-manylinux_2_27_aarch64.whl → 1.4.0.dev20250815__cp310-cp310-manylinux_2_27_aarch64.whl

ai_edge_litert/aot/prepare_for_npu.py

@@ -1,152 +0,0 @@
-# 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

@@ -1,18 +0,0 @@
-# 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."""
-
-from ai_edge_litert.aot.vendors.mediatek import mediatek_backend as _
-from ai_edge_litert.aot.vendors.qualcomm import qualcomm_backend as _

ai_edge_litert/aot/vendors/example/example_backend.py

@@ -1,157 +0,0 @@
-# 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

@@ -1,128 +0,0 @@
-# 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/import_vendor.py

@@ -1,132 +0,0 @@
-# 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()