ai-edge-litert-nightly 1.4.0.dev20250729__cp311-cp311-macosx_12_0_arm64.whl → 1.4.0.dev20250814__cp311-cp311-macosx_12_0_arm64.whl

ai_edge_litert/aot/core/types.py

@@ -1,374 +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.
-# ==============================================================================
-
-"""Basic types used in the LiteRt AOT flow."""
-
-import abc
-from collections.abc import Iterable
-import dataclasses
-import pathlib
-import sys
-from typing import Any, MutableMapping, Protocol, Type
-
-# pylint: disable=g-importing-member
-# pylint: disable=g-import-not-at-top
-# pylint: disable=g-bad-import-order
-if sys.version_info < (3, 10):
-  from typing_extensions import TypeAlias
-else:
-  from typing import TypeAlias
-# pylint: enable=g-bad-import-order
-# pylint: enable=g-import-not-at-top
-# pylint: enable=g-importing-member
-
-
-@dataclasses.dataclass(frozen=True)
-class SubgraphPartitionStats:
-  """Subgraph partition stats."""
-
-  subgraph_index: int
-  num_ops_offloaded: int
-  num_total_ops: int
-  num_partitions_offloaded: int
-
-  def __str__(self) -> str:
-    is_full_offload = self.num_ops_offloaded == self.num_total_ops
-    return (
-        'Subgraph'
-        f' {self.subgraph_index} {"fully" if is_full_offload else "partially"}'
-        f' compiled:\t{self.num_ops_offloaded} /'
-        f' {self.num_total_ops} ops offloaded to'
-        f' {self.num_partitions_offloaded} partitions.'
-    )
-
-
-@dataclasses.dataclass(frozen=True)
-class PartitionStats:
-  """Model partition stats."""
-
-  subgraph_stats: list[SubgraphPartitionStats]
-
-  def __str__(self) -> str:
-    return '\n'.join(str(s) for s in self.subgraph_stats)
-
-
-class Model:
-  """A model.
-
-  Note: If the model is not in memory, data_ will be a path to a file on disk.
-  If the model is in memory, data_ will be the model bytes.
-
-  However, there's no guarantee that the path will be a valid path to a file
-  on disk, and/or that the file are a valid TFLite model.
-  """
-
-  data_: pathlib.Path | bytes
-  partition_stats: PartitionStats | None = None
-
-  def __init__(
-      self,
-      path: pathlib.Path | str | None = None,
-      model_bytes: bytes | None = None,
-  ):
-    if path is not None:
-      if isinstance(path, str):
-        path = pathlib.Path(path)
-      if model_bytes:
-        raise ValueError('Cannot specify both path and model_bytes.')
-      self.data_ = path
-    else:
-      if model_bytes is None:
-        raise ValueError('Cannot specify neither path nor model_bytes.')
-      self.data_ = model_bytes
-
-  @property
-  def in_memory(self) -> bool:
-    return isinstance(self.data_, bytes)
-
-  @property
-  def path(self) -> pathlib.Path:
-    if not isinstance(self.data_, pathlib.Path):
-      raise ValueError('Model is not on disk.')
-    return self.data_
-
-  @property
-  def model_bytes(self) -> bytes:
-    if not isinstance(self.data_, bytes):
-      raise ValueError('Model is not in memory.')
-    return self.data_
-
-  @classmethod
-  def create_from_path(cls, path: pathlib.Path) -> 'Model':
-    return Model(path=path, model_bytes=None)
-
-  @classmethod
-  def create_from_bytes(cls, model_bytes: bytes) -> 'Model':
-    return Model(path=None, model_bytes=model_bytes)
-
-  def set_path(self, path: pathlib.Path | str):
-    if isinstance(path, str):
-      path = pathlib.Path(path)
-    self.data_ = path
-
-  def set_bytes(self, model_bytes: bytes):
-    self.data_ = model_bytes
-
-  def load(self):
-    """Loads the model from the given path.
-
-    Raises:
-      ValueError: If the model is already in memory.
-    """
-    if not isinstance(self.data_, pathlib.Path):
-      raise ValueError('Cannot load a model that is already in memory.')
-    self.data_ = self.data_.read_bytes()
-
-  def save(self, path: pathlib.Path | str, export_only: bool = False):
-    """Saves the model to the given path from the in-memory model content.
-
-    If export_only is True, the model will be copied to the given path without
-    modifying the internal state, regardless of whether the model is already on
-    disk or in memory.
-
-    Args:
-      path: The path to save the model to.
-      export_only: Whether to only export the model without modifying the
-        internal stat (i.e. transfer the in-memory model to disk).
-
-    Raises:
-      ValueError: If export_only is False and the model is not in memory.
-    """
-    if isinstance(path, str):
-      path = pathlib.Path(path)
-    if isinstance(self.data_, pathlib.Path):
-      if not export_only:
-        raise ValueError(
-            'Cannot save a model that is not in memory. Use export_only=True'
-            ' for copying the model to a new path.'
-        )
-      with open(self.data_, 'rb') as f:
-        model_content = f.read()
-    else:
-      model_content = self.data_
-    path.write_bytes(model_content)
-    if not export_only:
-      self.data_ = path
-
-
-@dataclasses.dataclass()
-class CompilationResult:
-  """Compilation result, as a collection of compiled models."""
-
-  models_with_backend: list[tuple['Backend', Model]] = dataclasses.field(
-      default_factory=list
-  )
-  failed_backends: list[tuple['Backend', str]] = dataclasses.field(
-      default_factory=list
-  )
-
-  @property
-  def models(self) -> list[Model]:
-    return [model for _, model in self.models_with_backend]
-
-  def load(self):
-    for _, model in self.models_with_backend:
-      if not model.in_memory:
-        model.load()
-
-  def export(self, output_dir: pathlib.Path | str, model_name: str = 'model'):
-    if isinstance(output_dir, str):
-      output_dir = pathlib.Path(output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    for backend, model in self.models_with_backend:
-      model.save(
-          output_dir / (model_name + backend.target_id_suffix + '.tflite'),
-          export_only=True,
-      )
-
-  def compilation_report(self) -> str:
-    """Returns a human readable compilation report."""
-    report = []
-    for backend, model in self.models_with_backend:
-      report.append(f'{backend.target_id}')
-      report.append('==========================')
-      report.append(f'Partition Stats:\n{model.partition_stats}\n')
-    report = '\n'.join(report)
-
-    failed_report = []
-    if self.failed_backends:
-      failed_report.append('==========================')
-      failed_report.append('COMPILATION FAILURES:')
-      failed_report.append('==========================')
-      for backend, error in self.failed_backends:
-        failed_report.append(f'{backend.target_id}\t{error}')
-    failed_report = '\n'.join(failed_report)
-    return '\n'.join([report, failed_report])
-
-
-class Component(Protocol):
-  """An arbitrary module in the AOT flow that inputs and outputs a Model.
-
-  For example quantizer, graph rewriter, compiler plugin etc.
-  """
-
-  @property
-  def component_name(self) -> str:
-    ...
-
-  def __call__(self, input_model: Model, output_model: Model, *args, **kwargs):
-    ...
-
-
-# A user provided configuration. This will contain all the information needed
-# to select the proper backend and run components (e.g. quant recipe,
-# backend id etc). Backends will validate and resolve configurations and are
-# ultimately responsible deciding how to configure the components.
-# NOTE: Consider a typed config approach (proto, data class, etc.)
-Config: TypeAlias = MutableMapping[str, Any]
-
-
-# Backend specific compilation configuration.
-BackendCompilationConfig: TypeAlias = MutableMapping[str, Any]
-
-
-# The following is experimental and for protyping only.
-class CompilationConfig:
-  """A typed configuration."""
-
-  target: 'Target'
-  compilation_config: BackendCompilationConfig = dataclasses.field(
-      default_factory=dict
-  )
-  quant_recipe: str | None = None
-
-  def __init__(self, target: 'Target', **kwargs: Any):
-    self.target = target
-    self.quant_recipe = kwargs.pop('quantize_recipe', None)
-    self.compilation_config = kwargs
-
-  def to_dict(self) -> dict[str, Any]:
-    ret = self.target.flatten()
-    ret['compilation_config'] = self.compilation_config
-    if self.quant_recipe is not None:
-      ret['quantize_recipe'] = self.quant_recipe
-    return ret
-
-
-class Backend(metaclass=abc.ABCMeta):
-  """A backend pertaining to a particular SoC vendor.
-
-  Mainly responsible for resolving configurations and managing vendor specific
-  resources (e.g. .so etc).
-  """
-
-  # NOTE: Only initialize through "create".
-  def __init__(self, config: Config):
-    self._config = config
-
-  @classmethod
-  @abc.abstractmethod
-  def create(cls, config: Config) -> 'Backend':
-    """Creates a backend instance.
-
-    If no target is specified, the backend will represent all targets.
-
-    Args:
-      config: The compilation configuration.
-
-    Returns:
-      The backend instance.
-    """
-
-  @classmethod
-  @abc.abstractmethod
-  def id(cls) -> str:
-    pass
-
-  @property
-  @abc.abstractmethod
-  def target(self) -> 'Target':
-    pass
-
-  @property
-  @abc.abstractmethod
-  def target_id(self) -> str:
-    pass
-
-  @property
-  def target_id_suffix(self) -> str:
-    if self.target_id:
-      return '_' + self.target_id
-    return ''
-
-  @property
-  def config(self) -> Config:
-    return self._config
-
-  @property
-  def soc_manufacturer(self) -> str:
-    """Manufacturer name or enum."""
-    raise NotImplementedError()
-
-  @property
-  def soc_model(self) -> str:
-    """Model name or enum."""
-    raise NotImplementedError()
-
-  @property
-  def shared_pass_names(self) -> list[str]:
-    """Names of shared passes."""
-    raise NotImplementedError()
-
-  @property
-  def quantize_recipe(self) -> str | None:
-    """Optional quantization recipe."""
-    return None
-
-  @abc.abstractmethod
-  def call_component(
-      self, input_model: Model, output_model: Model, component: Component
-  ):
-    pass
-
-  def specialize(self) -> Iterable['Backend']:
-    yield self
-
-
-BackendT: TypeAlias = Type[Backend]
-
-
-class Target(metaclass=abc.ABCMeta):
-  """Compilation target."""
-
-  @abc.abstractmethod
-  def __hash__(self) -> int:
-    pass
-
-  @abc.abstractmethod
-  def __eq__(self, other) -> bool:
-    pass
-
-  @abc.abstractmethod
-  def __repr__(self) -> str:
-    pass
-
-  @classmethod
-  @abc.abstractmethod
-  def backend_id(cls) -> str:
-    pass
-
-  @abc.abstractmethod
-  def flatten(self) -> dict[str, Any]:
-    return {'backend_id': self.backend_id()}

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)