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/aot_compile.py

@@ -1,152 +0,0 @@
-# Copyright 2025 The LiteRT 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.
-# ==============================================================================
-
-"""AOT Compilation for LiteRT model."""
-import pathlib
-import tempfile
-
-from ai_edge_litert.aot import prepare_for_npu as core
-from ai_edge_litert.aot.core import apply_plugin
-from ai_edge_litert.aot.core import components
-from ai_edge_litert.aot.core import mlir_transforms
-from ai_edge_litert.aot.core import types
-from ai_edge_litert.aot.vendors import import_vendor
-
-
-def aot_compile(
-    input_model: types.Model | str,
-    output_dir: str | pathlib.Path | None = None,
-    target: types.Target | list[types.Target] | None = None,
-    config: (
-        types.CompilationConfig | list[types.CompilationConfig] | None
-    ) = None,
-    quantizer: components.AieQuantizerT | None = None,
-    keep_going: bool = True,
-    subgraphs_to_compile: list[int] | None = None,
-    **kwargs,
-) -> 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:
-    input_model: The input model to compile.
-    output_dir: Directory to write the output files to. If not specified, the
-      output files will be written to the same directory as the input file.
-    target: The target to compile for. If not specified, will compile to all
-      registered targets.
-    config: The compilation config(s). Cannot be specified with target.
-    quantizer: The quantizer to use for quantization.
-    keep_going: Whether to keep going if some backends fail. If False, fail
-      fast on the first error and raise an exception.
-    subgraphs_to_compile: The subgraph index list to compile to NPU. If None,
-      compile all subgraphs.
-    **kwargs: Additional arguments to pass to the backend.
-
-  Returns:
-    Compiled models.
-  """
-  # Only one of target or config is needed.
-  if target and config:
-    raise ValueError("Cannot specify both target and config.")
-
-  if config is None:
-    if target is None:
-      target = import_vendor.AllRegisteredTarget()
-    if isinstance(target, types.Target):
-      config = types.CompilationConfig(target=target)
-    elif isinstance(target, list):
-      config = [types.CompilationConfig(target=t) for t in target]
-    else:
-      raise ValueError("Unsupported target type.")
-
-  if isinstance(input_model, str):
-    input_path = pathlib.Path(input_model)
-    input_model = types.Model.create_from_path(input_path)
-
-  # Resolve output paths.
-  temp_dir = None
-  if not output_dir:
-    if input_model.in_memory:
-      # Use a temp dir for in-memory models.
-      # The temp dir will be cleaned up after the models are compiled and loaded
-      # back to memory (i.e. function returns).
-      temp_dir = tempfile.TemporaryDirectory()
-      output_dir = temp_dir.name
-    else:
-      input_path = input_model.path
-      output_dir = input_path.parent / "_compiled_models"
-      output_dir.mkdir(parents=True, exist_ok=True)
-      output_dir = str(output_dir)
-  output_dir_path = pathlib.Path(output_dir)
-
-  if isinstance(config, types.CompilationConfig) or not config:
-    if config:
-      # Make pytype happy.
-      assert isinstance(config, types.CompilationConfig)
-      kw_config = config.to_dict() | kwargs
-    else:
-      kw_config = kwargs
-
-    backend_class = core.resolve_backend(kw_config)
-
-    quant_recipe = kw_config.get("quantize_recipe", None)
-    if quant_recipe:
-      assert quantizer is not None, "Quantizer is required for quantization."
-
-    results = core.prepare_for_npu(
-        input_model,
-        output_dir_path,
-        backend_class,
-        kw_config,
-        transforms=mlir_transforms.MlirTransforms(),
-        quantizer=quantizer,
-        plugin=apply_plugin.ApplyPlugin(
-            experimental_capture_stderr=True,
-            subgraphs_to_compile=subgraphs_to_compile,
-        ),
-        keep_going=keep_going,
-    )
-  elif isinstance(config, list):
-    kw_configs = [c.to_dict() | kwargs for c in config]
-
-    configs_with_backend = [(core.resolve_backend(c), c) for c in kw_configs]
-    requires_quantizer = any("quantize_recipe" in c for c in kw_configs)
-    if requires_quantizer and quantizer is None:
-      raise ValueError("Quantizer is required for quantization.")
-
-    results = core.prepare_for_npu_multiple_configs(
-        input_model,
-        output_dir_path,
-        configs_with_backend,
-        transforms=mlir_transforms.MlirTransforms(),
-        quantizer=quantizer,
-        plugin=apply_plugin.ApplyPlugin(
-            experimental_capture_stderr=True,
-            subgraphs_to_compile=subgraphs_to_compile,
-        ),
-        keep_going=keep_going,
-    )
-  else:
-    # Should not reach here.
-    raise ValueError("Unsupported config type.")
-
-  if temp_dir:
-    # Load the models to memory before cleaning up the temp dir.
-    results.load()
-    temp_dir.cleanup()
-
-  return results

ai_edge_litert/aot/core/apply_plugin.py

@@ -1,146 +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.
-# ==============================================================================
-
-"""Wrapper for calling the apply plugin tooling."""
-
-
-import os
-import pathlib
-import re
-import subprocess
-import tempfile
-
-from ai_edge_litert.aot.core import common
-from ai_edge_litert.aot.core import components
-from ai_edge_litert.aot.core import types
-
-_BINARY = pathlib.Path("tools/apply_plugin_main")
-
-_RE_PARTITION_STATS = re.compile(
-    r"PartitionSubgraph: (\d+), selected num ops: (\d+), from totoal ops:"
-    r" (\d+), num partitions: (\d+)"
-)
-
-
-class ApplyPlugin(components.ApplyPluginT):
-  """Wrapper for calling the apply plugin tooling."""
-
-  def __init__(
-      self,
-      experimental_capture_stderr: bool = False,
-      subgraphs_to_compile: list[int] | None = None,
-  ):
-    self._experimental_capture_stderr = experimental_capture_stderr
-    self._subgraphs_to_compile = subgraphs_to_compile
-
-  @property
-  def default_err(self) -> str:
-    # NOTE: Capture stderr from underlying binary.
-    return "--"
-
-  @property
-  def component_name(self) -> str:
-    return "apply_plugin"
-
-  def __call__(
-      self,
-      input_model: types.Model,
-      output_model: types.Model,
-      soc_manufacturer: str,
-      soc_model: str,
-      sdk_libs_path: str | None = None,
-      **kwargs,
-  ):
-    """Applies a plugin to the input model.
-
-    Args:
-      input_model: The path to the input model.
-      output_model: The path to the output model.
-      soc_manufacturer: The SOC manufacturer of the plugin.
-      soc_model: The SOC model of the plugin.
-      sdk_libs_path: The path to the SDK libs. If not provided,
-        the default SDK path will be used.
-      **kwargs: Additional arguments to pass to the underlying binary.
-
-    Returns:
-      The output model.
-
-    Raises:
-      ValueError: If no tflite model was created by the underying binary.
-    """
-    if input_model.in_memory:
-      tmp_file = tempfile.NamedTemporaryFile(mode="wb")
-      input_model.save(tmp_file.name)
-    else:
-      tmp_file = None
-
-    binary = common.get_resource(_BINARY)
-    args = [
-        str(binary),
-        "--cmd=apply",
-        f"--model={str(input_model.path)}",
-        f"--o={str(output_model.path)}",
-        f"--soc_manufacturer={soc_manufacturer}",
-        f"--soc_model={soc_model}",
-        f"--err={self.default_err}",
-    ]
-    extra_args = [f"--{key}={value}" for key, value in kwargs.items()]
-    args.extend(extra_args)
-    if self._subgraphs_to_compile:
-      subgraphs_to_compile = ",".join(
-          str(s) for s in self._subgraphs_to_compile
-      )
-      args.append(f"--subgraphs={subgraphs_to_compile}")
-    env = os.environ.copy()
-    ld_library_path = common.construct_ld_library_path()
-    if ld_library_path:
-      if sdk_libs_path:
-        ld_library_path = f"{sdk_libs_path}{os.pathsep}{ld_library_path}"
-      env["LD_LIBRARY_PATH"] = ld_library_path
-
-    result = subprocess.run(
-        args,
-        check=False,
-        text=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        env=env,
-    )
-    if result.returncode:
-      log_file = tempfile.NamedTemporaryFile(mode="w", delete=False)
-      log_file.write(result.stdout)
-      log_file.close()
-      raise ValueError(
-          f"{self.component_name} failed to apply plugin. See"
-          f" {log_file.name} for details."
-      )
-
-    if not common.is_tflite(output_model.path):
-      raise ValueError(f"{output_model.path} is not a TFLite model.")
-
-    partition_stats = _RE_PARTITION_STATS.findall(result.stdout)
-    output_model.partition_stats = types.PartitionStats(
-        subgraph_stats=[
-            types.SubgraphPartitionStats(
-                subgraph_index=int(s[0]),
-                num_ops_offloaded=int(s[1]),
-                num_total_ops=int(s[2]),
-                num_partitions_offloaded=int(s[3]),
-            )
-            for s in partition_stats
-        ]
-    )
-    if tmp_file is not None:
-      tmp_file.close()

ai_edge_litert/aot/core/common.py

@@ -1,95 +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.
-# ==============================================================================
-
-"""Constants and other small generic utilities."""
-
-from importlib import resources
-import os
-import pathlib
-
-TFLITE = "tflite"
-DOT_TFLITE = f".{TFLITE}"
-NPU = "npu"
-
-
-_WORKSPACE_PREFIX = "litert"
-_AI_EDGE_LITERT_PREFIX = "ai_edge_litert"
-_LITERT_ROOT = ""
-_PYTHON_ROOT = "python/aot"
-
-MODULE_ROOT = ".".join([
-    _WORKSPACE_PREFIX,
-    _LITERT_ROOT.replace("/", "."),
-    _PYTHON_ROOT.replace("/", "."),
-])
-
-
-def get_resource(litert_relative_path: pathlib.Path) -> pathlib.Path:
-  """Returns the path to a resource in the Litert workspace."""
-  try:
-    resource_root = resources.files(_WORKSPACE_PREFIX)
-  except ModuleNotFoundError:
-    resource_root = resources.files(_AI_EDGE_LITERT_PREFIX)
-  litert_resource = resource_root.joinpath(
-      _LITERT_ROOT, str(litert_relative_path)
-  )
-  if not litert_resource.is_file():
-    raise FileNotFoundError(f"Resource {litert_resource} does not exist.")
-  return pathlib.Path(str(litert_resource))
-
-
-def is_tflite(path: pathlib.Path) -> bool:
-  return path.exists() and path.is_file() and path.suffix == f".{TFLITE}"
-
-
-def construct_ld_library_path() -> str:
-  """Constructs a string suitable for the LD_LIBRARY_PATH environment variable.
-
-  This function is used in ai_edge_litert python package, when the shared
-  libraries are not in a static location. This function will construct the
-  LD_LIBRARY_PATH environment variable using the ai_edge_litert directory, and
-  all subdirectories.
-
-  If the module is built from source, this function will return an empty string.
-
-  Returns:
-    A string suitable for the LD_LIBRARY_PATH environment variable.
-  """
-  try:
-    resource_root = resources.files(_AI_EDGE_LITERT_PREFIX)
-  except ModuleNotFoundError:
-    # Bulit from source case.
-    return ""
-  root_package_path = str(resource_root)
-
-  library_paths = set()
-
-  library_paths.add(os.path.abspath(root_package_path))
-
-  for dirpath, _, _ in os.walk(root_package_path):
-    library_paths.add(os.path.abspath(dirpath))
-
-  sorted_paths = sorted(list(library_paths))
-  new_ld_library_path = os.pathsep.join(sorted_paths)
-  current_ld_library_path = os.environ.get("LD_LIBRARY_PATH")
-
-  if current_ld_library_path:
-    if current_ld_library_path not in new_ld_library_path:
-      lib_paths = f"{new_ld_library_path}{os.pathsep}{current_ld_library_path}"
-    else:
-      lib_paths = new_ld_library_path
-  else:
-    lib_paths = new_ld_library_path
-  return lib_paths

ai_edge_litert/aot/core/components.py

@@ -1,93 +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.
-# ==============================================================================
-
-"""Interfaces for specific components used in the LiteRt AOT flow."""
-
-import abc
-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, 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
-
-QuantRecipe: TypeAlias = list[dict[str, Any]] | str
-
-
-class AieQuantizerT(metaclass=abc.ABCMeta):
-  """Interface for AIE quantizer components."""
-
-  @property
-  def component_name(self) -> str:
-    return "aie_quantizer"
-
-  @abc.abstractmethod
-  def __call__(
-      self,
-      input_model: types.Model,
-      output_model: types.Model,
-      quantization_recipe: QuantRecipe | None = None,
-      *args,
-      **kwargs,
-  ):
-    pass
-
-
-class ApplyPluginT(metaclass=abc.ABCMeta):
-  """Interface for apply plugin components."""
-
-  @property
-  def default_err(self) -> str:
-    # NOTE: Capture stderr from underlying binary.
-    return "none"
-
-  @property
-  def component_name(self) -> str:
-    return "apply_plugin"
-
-  @abc.abstractmethod
-  def __call__(
-      self,
-      input_model: types.Model,
-      output_model: types.Model,
-      soc_manufacturer: str,
-      soc_model: str,
-      *args,
-      **kwargs,
-  ):
-    pass
-
-
-class MlirTransformsT(metaclass=abc.ABCMeta):
-  """Interface for MLIR transforms components."""
-
-  @property
-  def component_name(self) -> str:
-    return "mlir_transforms"
-
-  @abc.abstractmethod
-  def __call__(
-      self, input_model: types.Model, output_model: types.Model, *args, **kwargs
-  ):
-    pass

ai_edge_litert/aot/core/mlir_transforms.py

@@ -1,36 +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.
-# ==============================================================================
-
-"""Wrapper for suite of MLIR passes."""
-
-from ai_edge_litert.aot.core import components
-from ai_edge_litert.aot.core import tflxx_util
-from ai_edge_litert.aot.core import types
-
-
-class MlirTransforms(components.MlirTransformsT):
-  """Wrapper for suite of MLIR passes."""
-
-  def __call__(
-      self,
-      input_model: types.Model,
-      output_model: types.Model,
-      pass_name: str,
-  ):
-    if not input_model.in_memory:
-      input_model.load()
-    input_bytes = input_model.model_bytes
-    output_bytes = tflxx_util.call_tflxx(input_bytes, pass_name)
-    output_model.set_bytes(output_bytes)

ai_edge_litert/aot/core/tflxx_util.py

@@ -1,30 +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.
-# ==============================================================================
-
-# pylint: disable=g-import-not-at-top
-# pytype: disable=import-error
-# pytype: disable=not-callable
-
-"""Shim layer for TFLXX while it is in experimental."""
-
-
-import importlib.util
-from typing import Callable
-
-call_tflxx: Callable[[bytes, str], bytes] = lambda input, pass_name: input
-
-
-def tflxx_enabled() -> bool:
-  return False