ai-edge-litert-nightly 2.0.2a1.dev20250829__cp312-cp312-manylinux_2_27_x86_64.whl → 2.0.2.dev20250916__cp312-cp312-manylinux_2_27_x86_64.whl

ai_edge_litert/__init__.py

@@ -1 +1 @@
-__version__ = "2.0.2a1.dev20250829"
+__version__ = "2.0.2.dev20250916"

ai_edge_litert/aot/aot_compile.py

@@ -50,8 +50,8 @@ def aot_compile(
       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.
+    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.
@@ -92,6 +92,7 @@ def aot_compile(
       output_dir.mkdir(parents=True, exist_ok=True)
       output_dir = str(output_dir)
   output_dir_path = pathlib.Path(output_dir)
+  output_dir_path.mkdir(parents=True, exist_ok=True)
 
   if isinstance(config, types.CompilationConfig) or not config:
     if config:

ai_edge_litert/aot/vendors/__init__.py

@@ -14,5 +14,6 @@
 # ==============================================================================
 """Vendor backends for LiteRt."""
 
+from ai_edge_litert.aot.vendors.google_tensor import google_tensor_backend as _
 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/google_tensor/google_tensor_backend.py

@@ -0,0 +1,153 @@
+# 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"
+)
+
+
+@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_TENSOR
+
+  @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("LITERT_GOOGLE_TENSOR_SDK", None)
+
+    extra_kwargs = {"libs": lib_dir, "sdk_libs_path": sdk_libs_dir}
+  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: 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, [])

ai_edge_litert/aot/vendors/google_tensor/target.py

@@ -0,0 +1,83 @@
+# 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_tensor"
+
+
+class SocModel(StrEnum):
+  """Google Tensor SOC model."""
+
+  ALL = "ALL"
+
+  G3 = "g3"
+  G4 = "g4"
+  G5 = "g5"
+
+
+class SocManufacturer(StrEnum):
+  """Google Tensor SOC manufacturer."""
+
+  GOOGLE_TENSOR = "GoogleTensor"
+
+
+@dataclasses.dataclass
+class Target(types.Target):
+  """Compilation target for Google Tensor SOCs."""
+
+  soc_model: SocModel
+  soc_manufacturer: SocManufacturer = SocManufacturer.GOOGLE_TENSOR
+
+  @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