ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/lowertools/torch_xla_utils.py

@@ -0,0 +1,301 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+import copy
+import dataclasses
+from dataclasses import dataclass
+import gc
+import itertools
+import logging
+import os
+import tempfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+if "PJRT_DEVICE" not in os.environ:
+  # https://github.com/google-ai-edge/ai-edge-torch/issues/326
+  os.environ["PJRT_DEVICE"] = "CPU"
+
+os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
+
+
+from ai_edge_torch import model
+from ai_edge_torch._convert import conversion_utils
+from ai_edge_torch._convert import signature as signature_module
+from ai_edge_torch.lowertools import common_utils
+from ai_edge_torch.lowertools import translate_recipe
+from ai_edge_torch.quantize import quant_config as qcfg
+import torch
+from torch_xla import stablehlo
+
+try:
+  import tensorflow as tf
+
+  from tensorflow.compiler.tf2xla.python import xla as tfxla
+
+  from tensorflow.lite.python import conversion_metadata_schema_py_generated as conversion_metadata_fb  # isort:skip
+except ImportError:
+  logging.error(
+      "This module needs tensorflow with xla support.\n"
+      "Please install tensorflow with `pip install tf-nightly`.\n"
+  )
+  raise
+
+MlirBundle = stablehlo.StableHLOModelBundle
+
+
+@dataclasses.dataclass
+class MergedBundle:
+
+  bundle: stablehlo.StableHLOModelBundle
+  exported_programs: list[torch.export.ExportedProgram]
+  deduped_tf_vars: list[tf.Variable]
+
+
+def exported_program_to_mlir(
+    exported_program: torch.export.ExportedProgram,
+    sample_args: tuple[torch.Tensor],
+) -> stablehlo.StableHLOModelBundle:
+  # Setting export_weights to False here so that pytorch/xla avoids copying the
+  # weights to a numpy array which would lead to memory bloat. This means that
+  # the state_dict in the returned bundle is going to be empty.
+  return stablehlo.exported_program_to_stablehlo(
+      exported_program,
+      stablehlo.StableHLOExportOptions(
+          override_tracing_arguments=sample_args, export_weights=False
+      ),
+  )._bundle
+
+
+def merge_mlir_bundles(
+    bundles: list[stablehlo.StableHLOModelBundle],
+    signatures: list[signature_module.Signature],
+    exported_programs: list[torch.export.ExportedProgram],
+) -> stablehlo.StableHLOGraphModule:
+  state_dict, deduped_tf_vars = common_utils.gather_state_dict(
+      exported_programs, signatures
+  )
+
+  new_shlo_model_bundle = stablehlo.StableHLOModelBundle(
+      state_dict=state_dict, additional_constants=[], stablehlo_funcs=[]
+  )
+
+  for bundle, signature in zip(bundles, signatures):
+    const_offset = len(new_shlo_model_bundle.additional_constants)
+    for func in bundle.stablehlo_funcs:
+      func.meta.name = signature.name + "_" + func.meta.name
+      for loc in func.meta.input_locations:
+        if loc.type_ == stablehlo.VariableType.CONSTANT:
+          loc.position += const_offset
+        elif loc.type_ == stablehlo.VariableType.PARAMETER:
+          loc.name = signature.name + "_" + loc.name
+      new_shlo_model_bundle.stablehlo_funcs.append(func)
+    new_shlo_model_bundle.additional_constants.extend(
+        bundle.additional_constants
+    )
+  return MergedBundle(
+      bundle=new_shlo_model_bundle,
+      exported_programs=exported_programs,
+      deduped_tf_vars=deduped_tf_vars,
+  )
+
+
+def _get_shape_with_dynamic(signature: stablehlo.VariableSignature):
+  shape = copy.copy(signature.shape)
+  for i in signature.dynamic_dims:
+    shape[i] = None
+  return shape
+
+
+def _wrap_as_tf_func(
+    func: stablehlo.StableHLOFunc,
+    bundle: stablehlo.StableHLOModelBundle,
+    exported_program: torch.export.ExportedProgram,
+):
+  def inner(*args):
+    type_info = [sig.dtype for sig in func.meta.output_signature]
+    shape_info = [
+        _get_shape_with_dynamic(sig) for sig in func.meta.output_signature
+    ]
+    call_args = stablehlo._extract_call_parameters(args, func.meta, bundle)
+    call_module_return = tfxla.call_module(
+        tuple(call_args),
+        version=5,
+        Tout=type_info,
+        Sout=shape_info,
+        function_list=[],
+        module=func.bytecode,
+    )
+    spec = exported_program.call_spec.out_spec
+
+    # The module returning a flat array.
+    if not spec.context:
+      return call_module_return
+
+    flat_names = common_utils.flat_dict_names(spec.children_specs, spec.context)
+    return {name: value for name, value in zip(flat_names, call_module_return)}
+
+  return inner
+
+
+def _make_tf_signature(
+    meta: stablehlo.StableHLOFunctionMeta,
+    signature: signature_module.Signature,
+) -> list[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
+  input_pos_to_spec = {
+      loc.position: spec
+      for loc, spec in itertools.chain(
+          zip(meta.input_locations, meta.input_signature), meta.unused_inputs
+      )
+      if loc.type_ == stablehlo.VariableType.INPUT_ARG
+  }
+  assert len(input_pos_to_spec) == len(input_names)
+
+  primitive_type_to_tf_type = {"int": "int32", "float": "float32"}
+  ret: list[tf.TensorSpec] = []
+  for i, name in enumerate(input_names):
+    spec = input_pos_to_spec[i]
+    shape = _get_shape_with_dynamic(spec)
+    ret.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=primitive_type_to_tf_type[spec.dtype]
+            if spec.dtype in primitive_type_to_tf_type
+            else spec.dtype,
+            name=name,
+        )
+    )
+  return ret
+
+
+def exported_program_to_mlir_text(
+    exported_program: torch.export.ExportedProgram,
+) -> str:
+  """Converts a ExportedProgram to a MLIR text."""
+  return stablehlo.exported_program_to_stablehlo(
+      exported_program
+  ).get_stablehlo_text()
+
+
+def merged_bundle_to_tfl_model(
+    merged_bundle: MergedBundle,
+    signatures: list[signature_module.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: dict = {},
+    _saved_model_dir: Optional[str] = None,
+) -> None:
+  """Converts a StableHLOGraphModule to a tflite model.
+
+  Args: shlo_bundle - model to export and save
+
+    signatures: List of signatures from which names of the signatures is
+    extracted.
+    quant_config: User-defined quantization method and scheme of the model.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
+    _tfl_converter_flags: A nested dictionary allowing setting flags for the
+    underlying tflite converter.
+  """
+
+  tf_module = tf.Module()
+
+  shlo_bundle = merged_bundle.bundle
+
+  shlo_bundle.additional_constants = [
+      tf.Variable(v, trainable=False) for v in shlo_bundle.additional_constants
+  ]
+  tf_signatures: list[list[tf.TensorSpec]] = list(
+      _make_tf_signature(func.meta, sig)
+      for func, sig in zip(shlo_bundle.stablehlo_funcs, signatures)
+  )
+
+  tf_functions = [
+      _wrap_as_tf_func(func, shlo_bundle, ep)
+      for func, ep in zip(
+          shlo_bundle.stablehlo_funcs, merged_bundle.exported_programs
+      )
+  ]
+
+  tf_module.f = []
+  for tf_sig, func in zip(tf_signatures, tf_functions):
+    tf_module.f.append(
+        tf.function(
+            func,
+            input_signature=tf_sig,
+        )
+    )
+
+  tf_module._variables = (
+      merged_bundle.deduped_tf_vars + shlo_bundle.additional_constants
+  )
+  del shlo_bundle
+  gc.collect()
+
+  tf_concrete_funcs = [
+      func.get_concrete_function(*tf_sig)
+      for func, tf_sig in zip(tf_module.f, tf_signatures)
+  ]
+
+  # We need to temporarily save since TFLite's from_concrete_functions does not
+  # allow providing names for each of the concrete functions.
+  with tempfile.TemporaryDirectory() as temp_dir_path:
+    if _saved_model_dir is not None:
+      temp_dir_path = _saved_model_dir
+
+    tf.saved_model.save(
+        tf_module,
+        temp_dir_path,
+        signatures={
+            sig.name: tf_concrete_funcs[idx]
+            for idx, sig in enumerate(signatures)
+        },
+    )
+    # Clean up intermediate memory early.
+    del tf_functions
+    del tf_module
+    del tf_concrete_funcs
+    gc.collect()
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    converter._set_original_model_type(conversion_metadata_fb.ModelType.PYTORCH)
+    converter._experimental_enable_composite_direct_lowering = True
+
+    conversion_utils.set_tfl_converter_quant_flags(converter, quant_config)
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      translated_recipe = translate_recipe.translate_to_ai_edge_recipe(
+          quant_config.generative_recipe
+      )
+
+    conversion_utils.apply_tfl_converter_flags(converter, _tfl_converter_flags)
+
+    tflite_model = converter.convert()
+    del converter
+    gc.collect()
+
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      tflite_model = translate_recipe.quantize_model(
+          tflite_model, translated_recipe
+      )
+
+  return tflite_model

ai_edge_torch/lowertools/translate_recipe.py

@@ -0,0 +1,163 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+from ai_edge_quantizer import quantizer
+from ai_edge_torch.generative.quantize import quant_attrs
+from ai_edge_torch.generative.quantize import quant_recipe
+
+_ComputePrecision = quantizer.qtyping.ComputePrecision
+_QuantGranularity = quantizer.qtyping.QuantGranularity
+_OpName = quantizer.qtyping.TFLOperationName
+_TensorQuantConfig = quantizer.qtyping.TensorQuantizationConfig
+_OpQuantConfig = quantizer.qtyping.OpQuantizationConfig
+
+_DEFAULT_REGEX_STR = '.*'
+_SINGULAR_TRANSFORMER_BLOCK_REGEX_STR = 'transformer_block'
+_IDX_TRANSFORMER_BLOCKS_REGEX_STR = 'transformer_blocks\[{}\]'
+_ATTENTION_REGEX_STR = 'ai_edge_torch.generative.layers.attention'
+_FEEDFORWARD_REGEX_STR = 'ai_edge_torch.generative.layers.feed_forward'
+_EMBEDDING_REGEX_STR = 'Embedding_tok_embedding'
+_ANY_TWO_DIGITS_REGEX_STR = '\d{1,2}'
+
+
+def _get_nbits_from_dtype(dtype: quant_attrs.Dtype) -> int:
+  if dtype == quant_attrs.Dtype.FP32:
+    return 32
+  elif dtype == quant_attrs.Dtype.FP16:
+    return 16
+  elif dtype == quant_attrs.Dtype.INT8:
+    return 8
+  raise ValueError('Unimplemented number of bits')
+
+
+def _get_dtype_from_dtype(
+    dtype: quant_attrs.Dtype,
+) -> quantizer.qtyping.TensorDataType:
+  if dtype == quant_attrs.Dtype.FP32 or dtype == quant_attrs.Dtype.FP16:
+    return quantizer.qtyping.TensorDataType.FLOAT
+  else:
+    return quantizer.qtyping.TensorDataType.INT
+
+
+def _get_compute_precision_from_mode(
+    mode: quant_attrs.Mode,
+) -> _ComputePrecision:
+  if mode == quant_attrs.Mode.DYNAMIC_RANGE:
+    return _ComputePrecision.INTEGER
+  elif mode == quant_attrs.Mode.WEIGHT_ONLY:
+    return _ComputePrecision.FLOAT
+  raise ValueError('Unimplemented execution mode')
+
+
+def _get_explicit_dequant_from_mode(mode: quant_attrs.Mode) -> bool:
+  if mode == quant_attrs.Mode.DYNAMIC_RANGE:
+    return False
+  elif mode == quant_attrs.Mode.WEIGHT_ONLY:
+    return True
+  raise ValueError('Unimplemented execution mode')
+
+
+def _get_granularity(
+    granularity: quant_attrs.Granularity,
+) -> bool:
+  if granularity == quant_attrs.Granularity.CHANNELWISE:
+    return _QuantGranularity.CHANNELWISE
+  if granularity == quant_attrs.Granularity.NONE:
+    return _QuantGranularity.TENSORWISE
+  raise ValueError('Unimplemented granularity')
+
+
+def _get_algorithm_key_from_algorithm(algo: quant_attrs.Algorithm) -> str:
+  if algo == quant_attrs.Algorithm.MIN_MAX:
+    return quantizer.algorithm_manager.AlgorithmName.MIN_MAX_UNIFORM_QUANT
+  elif algo == quant_attrs.Algorithm.FLOAT_CAST:
+    return quantizer.algorithm_manager.AlgorithmName.FLOAT_CASTING
+  raise ValueError('Unimplemented algorithm')
+
+
+def _set_quant_config(
+    rm: quantizer.recipe_manager.RecipeManager,
+    layer_recipe: quant_recipe.LayerQuantRecipe,
+    regex: str,
+):
+  rm.add_quantization_config(
+      regex=regex,
+      operation_name=_OpName.ALL_SUPPORTED,
+      op_config=_OpQuantConfig(
+          weight_tensor_config=_TensorQuantConfig(
+              num_bits=_get_nbits_from_dtype(layer_recipe.weight_dtype),
+              symmetric=True,
+              granularity=_get_granularity(layer_recipe.granularity),
+              dtype=_get_dtype_from_dtype(layer_recipe.weight_dtype),
+          ),
+          compute_precision=_get_compute_precision_from_mode(layer_recipe.mode),
+          explicit_dequantize=_get_explicit_dequant_from_mode(
+              layer_recipe.mode
+          ),
+      ),
+      algorithm_key=_get_algorithm_key_from_algorithm(layer_recipe.algorithm),
+  )
+
+
+def translate_to_ai_edge_recipe(
+    recipe: quant_recipe.GenerativeQuantRecipe,
+) -> quantizer.recipe_manager.ModelQuantizationRecipe:
+  rm = quantizer.recipe_manager.RecipeManager()
+
+  if recipe.default is not None:
+    _set_quant_config(rm, recipe.default, _DEFAULT_REGEX_STR)
+
+  if recipe.embedding is not None:
+    _set_quant_config(rm, recipe.embedding, _EMBEDDING_REGEX_STR)
+
+  if recipe.attention is not None:
+    if isinstance(recipe.attention, dict):
+      for idx, layer in recipe.attention.items():
+        _set_quant_config(
+            rm,
+            layer,
+            f'{_IDX_TRANSFORMER_BLOCKS_REGEX_STR.format(idx)}/{_ATTENTION_REGEX_STR}',
+        )
+    else:
+      _set_quant_config(
+          rm,
+          recipe.attention,
+          f'{_SINGULAR_TRANSFORMER_BLOCK_REGEX_STR}/{_ATTENTION_REGEX_STR}',
+      )
+
+  if recipe.feedforward is not None:
+    if isinstance(recipe.feedforward, dict):
+      for idx, layer in recipe.feedforward.items():
+        _set_quant_config(
+            rm,
+            layer,
+            f'{_IDX_TRANSFORMER_BLOCKS_REGEX_STR.format(idx)}/{_FEEDFORWARD_REGEX_STR}',
+        )
+    else:
+      _set_quant_config(
+          rm,
+          recipe.feedforward,
+          f'{_SINGULAR_TRANSFORMER_BLOCK_REGEX_STR}/{_FEEDFORWARD_REGEX_STR}',
+      )
+
+  return rm.get_quantization_recipe()
+
+
+def quantize_model(
+    model: bytes, recipe: quantizer.recipe_manager.ModelQuantizationRecipe
+) -> bytearray:
+  qt = quantizer.Quantizer(model, recipe)
+  result = qt.quantize()
+  return result.quantized_model

ai_edge_torch/model.py

@@ -0,0 +1,177 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Represents an ai_edge_torch model.
+
+PyTorch models can be converted to this representation through
+`ai_edge_torch.convert`.
+"""
+from __future__ import annotations
+
+import abc
+import re
+from typing import Callable
+
+import numpy.typing as npt
+import tensorflow as tf
+
+from ai_edge_litert import interpreter as tfl_interpreter  # pylint: disable=g-direct-tensorflow-import
+
+DEFAULT_SIGNATURE_NAME = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+
+
+class Model(abc.ABC):
+  """Represents and edge model."""
+
+  @abc.abstractmethod
+  def __call__(
+      self,
+      *args: npt.ArrayLike,
+      signature_name: str = DEFAULT_SIGNATURE_NAME,
+      **kwargs,
+  ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
+    raise NotImplementedError()
+
+  @abc.abstractmethod
+  def export(self, path: str):
+    raise NotImplementedError()
+
+  @staticmethod
+  def load(path: str) -> TfLiteModel:
+    tflite_model = TfLiteModel.load(path)
+    if tflite_model:
+      return tflite_model
+
+    raise ValueError(f'File format in {path} cannot be deserialized.')
+
+
+class TfLiteModel(Model):
+  """An edge model which uses tflite under-the-hood."""
+
+  def __init__(self, tflite_model):
+    """Initializes the TfLiteModel instance using a TFLite serialized object.
+
+    Args:
+      tflite_model: A TFlite serialized object.
+    """
+    self._tflite_model = tflite_model
+    self._interpreter_builder = lambda: tfl_interpreter.Interpreter(
+        model_content=self._tflite_model,
+        experimental_default_delegate_latest_features=True,
+    )
+
+  def tflite_model(self) -> bytes:
+    """Returns the wrapped tflite model."""
+    return self._tflite_model
+
+  def set_interpreter_builder(
+      self, builder: Callable[[], tfl_interpreter.Interpreter]
+  ) -> None:
+    """Sets a custom interpreter builder.
+
+    Args:
+      builder: A function that returns a `tfl_interpreter.Interpreter` or its
+        subclass.
+    """
+    self._interpreter_builder = builder
+
+  def __call__(
+      self,
+      *args: npt.ArrayLike,
+      signature_name: str = DEFAULT_SIGNATURE_NAME,
+      **kwargs,
+  ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
+    """Runs inference on the edge model using the provided arguments.
+
+    Args:
+      *args: The arguments to be passed to the model for inference.
+      **kwargs: The arguments with specific names to be passed to the model for
+        inference.
+      signature_name: The name of the signature to be used for inference. The
+        default signature is used if not provided.
+    """
+    interpreter = self._interpreter_builder()
+    interpreter.allocate_tensors()
+
+    signature_list = interpreter.get_signature_list()
+    if signature_name not in signature_list:
+      raise ValueError(
+          'Invalid signature name provided. Available signatures:'
+          f' {", ".join(signature_list.keys())}'
+      )
+
+    try:
+      runner = interpreter.get_signature_runner(signature_name)
+    except ValueError as exception:
+      if 'Invalid signature_key provided.' in str(exception):
+        raise ValueError(
+            'Invalid signature key provided. Available signatures:'
+            f' {list(signature_list.keys())}'
+        )
+      else:
+        raise exception
+
+    if len(signature_list[signature_name]['inputs']) != len(args) + len(kwargs):
+      raise ValueError(
+          'The model requires'
+          f' {len(signature_list[signature_name]["inputs"])} arguments but'
+          f' {len(args)} was provided.'
+      )
+
+    # Gather the input dictionary based on the signature.
+    inputs = {f'args_{idx}': args[idx] for idx in range(len(args))}
+    inputs = {**inputs, **kwargs}
+    outputs = runner(**inputs)
+
+    # When attempting to run a model, check if all the output tensors are named
+    # output_<number>. If so, assume the pytorch model returned a tuple and not
+    # a dictionary.
+    output_heuristic = lambda key: bool(re.search(r'output_\d+', key))
+    if all(output_heuristic(key) for key in outputs.keys()):
+      return (
+          outputs['output_0']
+          if len(outputs) == 1
+          else [outputs[f'output_{idx}'] for idx in range(len(outputs))]
+      )
+
+    return outputs
+
+  def export(self, path: str) -> None:
+    """Serializes the edge model to disk.
+
+    Args:
+      path: The path to file to which the model is serialized.
+    """
+    with open(path, 'wb') as file_handle:
+      file_handle.write(self._tflite_model)
+
+  @staticmethod
+  def load(path: str) -> TfLiteModel | None:
+    """Returns an edge (tflite) model by reading it from the disk.
+
+    Args:
+      str: The path to the model.
+    """
+    with open(path, 'rb') as file_handle:
+      model_content = file_handle.read()
+
+    # Check if this is indeed a tflite model:
+    try:
+      interpreter = tfl_interpreter.Interpreter(model_content=model_content)
+      interpreter.get_signature_list()
+    except:
+      return None
+
+    return TfLiteModel(model_content)

ai_edge_torch/odml_torch/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+from . import composite
+from . import debuginfo
+from . import export
+from . import export_utils
+from . import lowerings
+from . import passes