ai-edge-torch-nightly 0.3.0.dev20240828__py3-none-any.whl → 0.3.0.dev20240830__py3-none-any.whl

ai_edge_torch/odml_torch/export.py

@@ -0,0 +1,320 @@
+# 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.
+# ==============================================================================
+"""APIs to convert and lower a PyTorch ExportedProgram to MLIR."""
+
+import dataclasses
+import enum
+import io
+import operator
+from typing import Any, Callable, Optional
+
+from jax.lib import xla_extension
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+import torch.utils._pytree as pytree
+
+from . import _torch_future
+from . import debuginfo
+from . import export_utils
+from . import lowerings
+
+LoweringContext = lowerings.context.LoweringContext
+
+
+def _build_flat_inputs(
+    ctx: ir.Context, exported_program: torch.export.ExportedProgram
+):
+  """Build flattened inputs and metadata from exported program's signature."""
+  placeholder_nodes = [
+      n for n in exported_program.graph.nodes if n.op == "placeholder"
+  ]
+  export_flat_args = _torch_future.graph_module_flat_inputs(
+      exported_program, *exported_program.example_inputs
+  )
+
+  ir_inputs = []
+  tensor_metas = []
+  for node, arg in zip(placeholder_nodes, export_flat_args):
+    tensor_meta = node.meta.get("tensor_meta")
+    if tensor_meta is None:
+      raise RuntimeError(f"{type(arg)} (for {node.name}) is not a tensor")
+
+    tensor_metas.append(tensor_meta)
+    # Assume all dynamic dimensions are unbounded.
+    # TODO: Add checks for ep.range_constraints in MLIR.
+    shape = tuple(
+        export_utils.IR_DYNAMIC if export_utils.is_torch_dynamic(s) else s
+        for s in tensor_meta.shape
+    )
+    ir_inputs.append(
+        ir.RankedTensorType.get(
+            shape,
+            export_utils.torch_dtype_to_ir_element_type(ctx, tensor_meta.dtype),
+        )
+    )
+  return tuple(ir_inputs), tuple(export_flat_args), tuple(tensor_metas)
+
+
+def _get_output_metas(exported_program: torch.export.ExportedProgram):
+  """Get the output node's tensor_meta from the exported program."""
+  outputs = [n for n in exported_program.graph.nodes if n.op == "output"]
+  assert len(outputs) == 1
+  outputs, _ = pytree.tree_flatten(outputs[0].args[0])
+  assert all(isinstance(output, torch.fx.Node) for output in outputs)
+  return tuple(output.meta["tensor_meta"] for output in outputs)
+
+
+class LoweringInterpreter(torch.fx.Interpreter):
+  """The FX interpreter to iterate and invoke corresponding lowering for each PyTorch op in the graph."""
+
+  def __init__(self, module: torch.fx.GraphModule, lctx: LoweringContext):
+    super().__init__(module)
+    self.lctx = lctx
+    self.outputs = None
+
+  def _build_loc(self, node: torch.fx.Node):
+
+    info = debuginfo.build_mlir_debuginfo(node)
+    if info is None:
+      return ir.Location.unknown()
+
+    return ir.Location.name(name=info)
+
+  def run_node(self, node: torch.fx.Node):
+    loc = self._build_loc(node)
+    with loc:
+      self.lctx = self.lctx.replace(ir_location=loc, node=node)
+      res = super().run_node(node)
+      self.lctx = self.lctx.replace(ir_location=None, node=None)
+    return res
+
+  def call_function(self, target, args, kwargs):
+    if target is operator.getitem:
+      return super().call_function(target, args, kwargs)
+
+    if hasattr(target, "_schema"):
+      new_args = []
+      for arg, spec in zip(args, target._schema.arguments):
+        if isinstance(spec.type, torch.TensorType):
+          if isinstance(arg, int):
+            arg = lowerings.utils.splat(arg, ir.IntegerType.get_signless(32))
+          elif isinstance(arg, float):
+            arg = lowerings.utils.splat(arg, ir.F32Type.get())
+
+        new_args.append(arg)
+      args = tuple(new_args)
+
+    lowering = lowerings.lookup(target)
+    if lowering is None:
+      raise RuntimeError(f"Lowering not found: {target}")
+    return lowering(self.lctx, *args, **kwargs)
+
+  def output(self, target, args, kwargs):
+    flat_outputs = pytree.tree_flatten(args[0])[0]
+    self.outputs = flat_outputs
+
+
+@dataclasses.dataclass
+class InputSpec:
+
+  class VariableType(enum.Enum):
+    USER_INPUT = "user_input"
+    PARAMETER = "parameter"
+
+  type_: VariableType
+  i: int = -1
+  name: str = ""
+
+  @classmethod
+  def parameter(cls, name: str):
+    return cls(type_=cls.VariableType.PARAMETER, name=name)
+
+  @classmethod
+  def user_input(cls, i: int):
+    return cls(type_=cls.VariableType.USER_INPUT, i=i)
+
+  @property
+  def is_parameter(self):
+    return self.type_ == self.VariableType.PARAMETER
+
+  @property
+  def is_user_input(self):
+    return self.type_ == self.VariableType.USER_INPUT
+
+
+@dataclasses.dataclass
+class VariableSignature:  # either argument or parameters
+  shape: list[int]
+  dtype: str
+  input_spec: InputSpec = None
+
+
+@dataclasses.dataclass
+class MlirLowered:
+  """The lowered MLIR module, metadata, and weight tensors bundle from exported program."""
+
+  ctx: ir.Context
+  module: ir.Module
+  state_dict: dict[str, torch.Tensor]
+  input_signature: list[VariableSignature]
+  output_signature: list[VariableSignature]
+
+  _tf_function: Optional[Callable[Any, Any]] = None
+
+  def __str__(self):
+    return str(self.get_text(enable_debug_info=False))
+
+  def __repr__(self):
+    return str(self.get_text(enable_debug_info=False))
+
+  def get_text(self, enable_debug_info=False):
+    return str(
+        self.module.operation.get_asm(enable_debug_info=enable_debug_info)
+    )
+
+  @property
+  def module_bytecode(self) -> bytes:
+    output = io.BytesIO()
+    self.module.operation.write_bytecode(file=output)
+    return output.getvalue()
+
+  @property
+  def module_bytecode_vhlo(self) -> bytes:
+    # HACK: In OSS, we use MLIR pybinding and StableHLO dialect from JAX's
+    # build, which may not have the same StableHLO version as what used in
+    # TFLite converter. Therefore we always serialize MLIR module in VHLO.
+    # TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
+    target_version = stablehlo.get_minimum_version()
+    module_bytecode = xla_extension.mlir.serialize_portable_artifact(
+        self.module_bytecode, target_version
+    )
+    return module_bytecode
+
+  @property
+  def tf_function(self):
+    # Lazy import
+    from . import tf_integration
+
+    if self._tf_function is None:
+      self._tf_function = tf_integration.mlir_to_tf_function(self)
+    return self._tf_function
+
+  def __call__(self, *args):
+    # Lazy importing TF when execution is needed.
+    return self.tf_function(*args)
+
+  def to_flatbuffer(self):
+    from . import tf_integration
+
+    return tf_integration.mlir_to_flatbuffer(self)
+
+
+def exported_program_to_mlir(
+    exported_program: torch.export.ExportedProgram,
+) -> MlirLowered:
+  """Lower the exported program to MLIR."""
+  if torch.__version__ >= "2.2":
+    # torch version 2.1 didn't expose this yet
+    exported_program = exported_program.run_decompositions()
+    exported_program = exported_program.run_decompositions(
+        lowerings.decompositions()
+    )
+
+  with export_utils.create_ir_context() as context, ir.Location.unknown():
+
+    module = ir.Module.create()
+    lctx = LoweringContext(context, module)
+    interpreter = LoweringInterpreter(exported_program.graph_module, lctx)
+    ir_flat_inputs, export_flat_args, tensor_metas = _build_flat_inputs(
+        context, exported_program
+    )
+
+    # HACK: OSS MLIR pybinding could mysteriously transform func.func under
+    # construction into a func.return op after calling ir.Module.parse(..)
+    # in the context, which happens in JAX bridge. This is a bug in MLIR
+    # pybinding.
+    # Workaround steps:
+    # 1. Create a temp func.func.
+    # 2. Create and insert ops to temp's entry block. During the process
+    #   the temp func.func would be broken, but the ops in the block are fine.
+    # 3. Create the main func.func and copy all the ops in temp's entry block
+    #   to main.
+    # 4. Erase the temp func.func.
+    temp_func = func.FuncOp(
+        "temp",
+        ir.FunctionType.get(ir_flat_inputs, []),
+        ip=ir.InsertionPoint.at_block_begin(module.body),
+    )
+    with ir.InsertionPoint(temp_func.add_entry_block()):
+      interpreter.run(*temp_func.arguments, enable_io_processing=False)
+      num_mutations = len(exported_program.graph_signature.buffers_to_mutate)
+      outputs = interpreter.outputs[num_mutations:]
+      func.ReturnOp(interpreter.outputs[num_mutations:])
+
+    main_func = func.FuncOp(
+        "main",
+        ir.FunctionType.get(ir_flat_inputs, [o.type for o in outputs]),
+        ip=ir.InsertionPoint.at_block_begin(module.body),
+    )
+    with ir.InsertionPoint(main_func.add_entry_block()):
+      outputs = export_utils.clone_func_body_ops(temp_func, main_func.arguments)
+      func.ReturnOp(outputs)
+
+    main_func.attributes["sym_visibility"] = ir.StringAttr.get("public")
+    temp_func.erase()
+
+  module.operation.verify()
+
+  input_signature = []
+  state_dict = {}
+
+  user_inputs_cnt = 0
+  for arg, tensor_meta, input_spec in zip(
+      export_flat_args,
+      tensor_metas,
+      exported_program.graph_signature.input_specs,
+  ):
+    # Assumption:
+    # All states comes first in the list of args, and user provided inputs
+    # comes later. Also there is no kwargs.
+    if input_spec.kind == torch.export.graph_signature.InputKind.USER_INPUT:
+      input_signature.append(
+          VariableSignature(
+              tensor_meta.shape,
+              tensor_meta.dtype,
+              input_spec=InputSpec.user_input(user_inputs_cnt),
+          )
+      )
+      user_inputs_cnt += 1
+    else:
+      # Parameter or constant
+      state_dict[input_spec.target] = arg
+      input_signature.append(
+          VariableSignature(
+              tensor_meta.shape,
+              tensor_meta.dtype,
+              input_spec=InputSpec.parameter(input_spec.target),
+          )
+      )
+
+  output_signature = [
+      VariableSignature(tensor_meta.shape, tensor_meta.dtype)
+      for tensor_meta in _get_output_metas(exported_program)
+  ]
+  return MlirLowered(
+      context, module, state_dict, input_signature, output_signature
+  )

ai_edge_torch/odml_torch/export_utils.py

@@ -0,0 +1,168 @@
+# 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.
+# ==============================================================================
+"""Utilities for ODML Torch export."""
+
+import functools
+import re
+from typing import Sequence, cast
+import jax._src.interpreters.mlir
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+import torch
+
+# std::numeric_limits<int64_t>::min()
+IR_DYNAMIC = -9223372036854775808
+
+
+def is_ir_dynamic(v):
+  return v == IR_DYNAMIC
+
+
+def is_torch_dynamic(v):
+  return isinstance(v, torch.SymInt)
+
+
+def is_iterable(v):
+  try:
+    iter(v)
+  except TypeError:
+    return False
+  return True
+
+
+def create_ir_context():
+  # HACK: Use ir context from JAX as base for better stability in OSS.
+  # TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
+  context = jax._src.interpreters.mlir.make_ir_context()
+  context.allow_unregistered_dialects = True
+
+  return context
+
+
+def inline(
+    symbol_table: ir.SymbolTable,
+    block: ir.Block,
+):
+  """Recursively inlines all func.call ops in the block.
+
+  The symbol_table must include all func.func called by func.call ops.
+  This inliner in Python is implemented because MLIR inline pass from JAX's
+  MLIR pybinding build in OSS cannot properly inline func.call ops.
+  """
+  while True:
+    is_changed = False
+    for op in block.operations:
+      if op.OPERATION_NAME != func.CallOp.OPERATION_NAME:
+        continue
+
+      call_op = cast(func.CallOp, op)
+      func_op = cast(func.FuncOp, symbol_table[call_op.callee.value])
+      with ir.InsertionPoint(op):
+        new_results = clone_func_body_ops(func_op, call_op.operands)
+
+      for old_result, new_result in zip(call_op.results, new_results):
+        old_result = cast(ir.Value, old_result)
+        old_result.replace_all_uses_with(new_result)
+      call_op.erase()
+      is_changed = True
+
+    if not is_changed:
+      break
+
+  for op in block.operations:
+    for region in op.regions:
+      for block in region.blocks:
+        inline(symbol_table, block)
+
+
+def clone_func_body_ops(func_op: func.FuncOp, ir_inputs: Sequence[ir.Value]):
+  """Clone operations in the func_op's body by one into the current context."""
+  func_args = list(func_op.arguments)
+  ir_inputs = list(ir_inputs)
+  assert len(func_args) == len(ir_inputs)
+
+  value_mapping = {arg: ir_input for arg, ir_input in zip(func_args, ir_inputs)}
+
+  for op in list(func_op.entry_block.operations):
+    cloned_operands = [value_mapping[val] for val in op.operands]
+    if op.OPERATION_NAME == func.ReturnOp.OPERATION_NAME:
+      return cloned_operands
+
+    cloned = cast(ir.Operation, op.operation.clone())
+
+    for i in range(len(op.operands)):
+      cloned.operands[i] = cloned_operands[i]
+
+    for i in range(len(op.results)):
+      value_mapping[op.results[i]] = cloned.results[i]
+
+  return []
+
+
+def sanitize_aten_op_name(op, chars=":."):
+  return re.sub("[{}]".format(chars), "_", str(op))
+
+
+def build_ir_attr(val):
+  if val is None:
+    return ir.StringAttr.get("py_None")
+  if isinstance(val, bool):
+    return ir.BoolAttr.get(val)
+  if isinstance(val, int):
+    return ir.IntegerAttr.get(ir.IntegerType.get_signless(64), val)
+  if isinstance(val, float):
+    return ir.BoolAttr.get(val)
+  if isinstance(val, str):
+    return ir.StringAttr.get(val)
+  if isinstance(val, dict):
+    return ir.DictAttr.get({k: build_ir_attr(v) for k, v in val.items()})
+  if isinstance(val, (list, tuple)):
+    return ir.ArrayAttr.get([build_ir_attr(v) for v in val])
+
+  # Stringify the value to a StringAttr by default
+  return ir.StringAttr.get(str(val))
+
+
+def torch_dtype_to_ir_element_type(ctx, dtype):
+  ty_get = {
+      torch.double: ir.F64Type.get,
+      torch.float32: ir.F32Type.get,
+      torch.half: ir.F16Type.get,
+      torch.long: functools.partial(ir.IntegerType.get_signless, 64),
+      torch.int32: functools.partial(ir.IntegerType.get_signless, 32),
+      torch.int16: functools.partial(ir.IntegerType.get_signless, 16),
+      torch.bool: functools.partial(ir.IntegerType.get_signless, 1),
+  }.get(dtype)
+  return ty_get(ctx)
+
+
+def ir_element_type_to_torch_dtype(ty):
+  if isinstance(ty, ir.F32Type):
+    return torch.float32
+  if isinstance(ty, ir.F64Type):
+    return torch.float64
+  if isinstance(ty, ir.F16Type):
+    return torch.half
+  if isinstance(ty, ir.IntegerType):
+    if ty.is_signless:
+      if ty.width == 64:
+        return torch.long
+      if ty.width == 32:
+        return torch.int32
+      if ty.width == 16:
+        return torch.int16
+      if ty.width == 1:
+        return torch.bool
+  raise RuntimeError(f"Unsupported ir element type: {ty}")

ai_edge_torch/odml_torch/jax_bridge/__init__.py

@@ -0,0 +1,15 @@
+# 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_torch.odml_torch.jax_bridge._wrap import wrap

ai_edge_torch/odml_torch/jax_bridge/_wrap.py

@@ -0,0 +1,152 @@
+# 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.
+# ==============================================================================
+"""APIs to wrap JAX functions for using in ODML Torch lowerings."""
+
+import functools
+import inspect
+from typing import Any, Callable, cast
+import uuid
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.odml_torch import passes
+from ai_edge_torch.odml_torch.jax_bridge import utils
+import jax
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+import torch.utils._pytree as pytree
+
+# Jax double (64bit) precision is required to generate StableHLO mlir with
+# i64/f64 tensors from Jax bridged lowerings. If not set properly, all the
+# 64bit tensors would be truncated to 32bit dtype and potentially break the
+# lowering.
+jax.config.update("jax_enable_x64", True)
+
+
+def _lower_to_ir_text(
+    jaxfn, args, kwargs, ir_input_names: list[str] = None
+) -> str:
+  args = utils.tree_map_list_to_tuple(args)
+  kwargs = utils.tree_map_list_to_tuple(kwargs)
+
+  names_args = [
+      *zip(inspect.signature(jaxfn).parameters.keys(), args),
+      *kwargs.items(),
+  ]
+
+  static_argnames = []
+  jax_lower_static_kwargs = {}
+  jax_lower_args = []
+  jax_lower_argnames = []
+  ir_inputs = []
+
+  for i, (name, arg) in enumerate(names_args):
+    is_positional = i < len(args)
+    if not utils.is_ir_variable(arg):
+      static_argnames.append(name)
+      jax_lower_static_kwargs[name] = arg
+    else:
+      # Enforce the arg order in the mlir is the same as the lowering func
+      jax_lower_args.append(utils.ir_variable_to_jax(arg))
+
+      if is_positional and len(jax_lower_args) == i + 1:
+        # The first N continuous tensor args are passed to the lowering func
+        # as positional args, when they passed to the bridged func as
+        # positional args also.
+        jax_lower_argnames.append(None)
+      else:
+        # Otherwise pass the arg to the lowering func as keyword arg.
+        jax_lower_argnames.append(name)
+
+      if ir_input_names is None or name in ir_input_names:
+        # ir variable can be a nested tuple, while mlir args should be flat.
+        ir_inputs += [
+            x for x in pytree.tree_flatten(arg)[0] if isinstance(x, ir.Value)
+        ]
+
+  def new_lowering(*args, **jax_lower_static_kwargs):
+    jaxfn_args = []
+    jaxfn_kwargs = jax_lower_static_kwargs.copy()
+    for name, arg in zip(jax_lower_argnames, args):
+      if name is None:
+        jaxfn_args.append(arg)
+      else:
+        jaxfn_kwargs[name] = arg
+
+    return jaxfn(*jaxfn_args, **jaxfn_kwargs)
+
+  return (
+      jax.jit(new_lowering, static_argnames=static_argnames)
+      .lower(*jax_lower_args, **jax_lower_static_kwargs)
+      .as_text()
+  ), ir_inputs
+
+
+def wrap(jaxfn: Callable[Any, Any], ir_input_names: list[str] = None):
+  """Return the wrapped JAX function to be used in ODMLTorch lowerings.
+
+  If the given jaxfn has signature `jaxfn(*args, **kwargs) -> return`, the
+  wrapped function would:
+  - Have signature `wrapped(lctx: odml_torch.export.LoweringContext, *args,
+  **kwargs) -> return`.
+  - Accept mlir.ir.Value for all params expecting jax.Array as inputs.
+  - Return mlir.ir.Value for all jax.Array outputs from jaxfn.
+
+  Args:
+    jaxfn: The JAX function to be wrapped.
+    ir_input_names: The input (param) names of the JAX function to be used in
+      the MLIR lowering. This is useful when the JAX impl only depends on
+      specific inputs to the function. If not specified, all ir.Value passed to
+      the wrapped function are assumed to be used in the lowering.
+  """
+
+  @functools.wraps(jaxfn)
+  def wrapped(lctx, *args, **kwargs):
+
+    ir_text, ir_inputs = _lower_to_ir_text(
+        jaxfn,
+        args,
+        kwargs,
+        ir_input_names=ir_input_names,
+    )
+
+    module = ir.Module.parse(ir_text)
+    passes.strip_debuginfo(module)
+
+    symbol_table = ir.SymbolTable(module.operation)
+    main_func = symbol_table["main"]
+
+    with ir.InsertionPoint.at_block_begin(lctx.ir_module.body):
+      cloned_func = cast(func.FuncOp, main_func.clone())
+      cloned_func_name = f"{jaxfn.__name__}_{uuid.uuid4().hex[:8]}"
+      cloned_func.attributes["sym_name"] = ir.StringAttr.get(cloned_func_name)
+      cloned_func.attributes["sym_visibility"] = ir.StringAttr.get("private")
+
+    # HACK: Use the custom inliner implemented in Python because MLIR inline
+    # pass from JAX's MLIR pybinding build in OSS cannot properly inline
+    # func.call ops.
+    # This should be switched to `passes.inline(module)` when we have our own
+    # MLIR pybinding build.
+    export_utils.inline(symbol_table, cloned_func.entry_block)
+
+    if not cloned_func.arguments:
+      # Known edge case: when the lowering does not depend on input but
+      # just the meta of input like shape or dtype.
+      ir_inputs = []
+
+    results = func.CallOp(cloned_func, ir_inputs).results
+    if len(results) == 1:
+      return results[0]
+    return results
+
+  return wrapped