ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -0,0 +1,252 @@
+# 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 functools
+import logging
+
+from ai_edge_torch.odml_torch import jax_bridge
+import torch
+import torch_xla2.ops.jaten  # Import to load torch_xla2 ops
+import torch_xla2.ops.ops_registry  # Import to load torch_xla2 ops
+
+from . import registry
+
+
+@functools.cache
+def _log_usage(op):
+  logging.warning("Use jax lowering: %s", str(op))
+
+
+def lower_by_jax(op, ir_input_names=None):
+  def inner(lowering):
+    bridged = jax_bridge.wrap(lowering, ir_input_names)
+
+    @registry.lower(op)
+    def _jax_lowering(lctx, *args, **kwargs):
+      _log_usage(op)
+      return bridged(lctx, *args, **kwargs)
+
+    return lowering
+
+  return inner
+
+
+_TORCH_XLA2_IMPLS = {
+    key: val.func
+    for key, val in torch_xla2.ops.ops_registry.all_aten_ops.items()
+    if val.is_jax_function
+}
+
+
+def lower_by_torch_xla2(op):
+  return lower_by_jax(op)(_TORCH_XLA2_IMPLS[op])
+
+
+lower_by_torch_xla2(torch.ops.aten._adaptive_avg_pool2d)
+lower_by_torch_xla2(torch.ops.aten._adaptive_avg_pool3d)
+lower_by_torch_xla2(torch.ops.aten._cdist_forward)
+lower_by_torch_xla2(torch.ops.aten._local_scalar_dense)
+lower_by_torch_xla2(torch.ops.aten._local_scalar_dense)
+lower_by_torch_xla2(torch.ops.aten._log_softmax)
+lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit)
+lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit_no_training)
+lower_by_torch_xla2(torch.ops.aten._pdist_forward)
+lower_by_torch_xla2(torch.ops.aten._softmax)
+lower_by_torch_xla2(torch.ops.aten._to_copy)
+lower_by_torch_xla2(torch.ops.aten._unsafe_index)
+lower_by_torch_xla2(torch.ops.aten._unsafe_view)
+lower_by_torch_xla2(torch.ops.aten.abs)
+lower_by_torch_xla2(torch.ops.aten.acos)
+lower_by_torch_xla2(torch.ops.aten.acosh)
+lower_by_torch_xla2(torch.ops.aten.add.Scalar)
+lower_by_torch_xla2(torch.ops.aten.add.Tensor)
+lower_by_torch_xla2(torch.ops.aten.addbmm.default)
+lower_by_torch_xla2(torch.ops.aten.addmm)
+lower_by_torch_xla2(torch.ops.aten.addmv)
+lower_by_torch_xla2(torch.ops.aten.alias)
+lower_by_torch_xla2(torch.ops.aten.allclose)
+lower_by_torch_xla2(torch.ops.aten.amax)
+lower_by_torch_xla2(torch.ops.aten.amin)
+lower_by_torch_xla2(torch.ops.aten.any)
+lower_by_torch_xla2(torch.ops.aten.arange.default)
+lower_by_torch_xla2(torch.ops.aten.arange.start)
+lower_by_torch_xla2(torch.ops.aten.arange.start_step)
+lower_by_torch_xla2(torch.ops.aten.argmax)
+lower_by_torch_xla2(torch.ops.aten.argmin)
+lower_by_torch_xla2(torch.ops.aten.as_strided)
+lower_by_torch_xla2(torch.ops.aten.as_strided_copy)
+lower_by_torch_xla2(torch.ops.aten.asin)
+lower_by_torch_xla2(torch.ops.aten.asinh)
+lower_by_torch_xla2(torch.ops.aten.atan)
+lower_by_torch_xla2(torch.ops.aten.atan2)
+lower_by_torch_xla2(torch.ops.aten.atanh)
+lower_by_torch_xla2(torch.ops.aten.avg_pool2d)
+lower_by_torch_xla2(torch.ops.aten.avg_pool3d)
+lower_by_torch_xla2(torch.ops.aten.bitwise_and)
+lower_by_torch_xla2(torch.ops.aten.bitwise_not)
+lower_by_torch_xla2(torch.ops.aten.bitwise_or)
+lower_by_torch_xla2(torch.ops.aten.bitwise_xor)
+lower_by_torch_xla2(torch.ops.aten.bmm)
+lower_by_torch_xla2(torch.ops.aten.cat)
+lower_by_torch_xla2(torch.ops.aten.ceil)
+lower_by_torch_xla2(torch.ops.aten.clamp.Tensor)
+lower_by_torch_xla2(torch.ops.aten.clamp.default)
+lower_by_torch_xla2(torch.ops.aten.clone)
+lower_by_torch_xla2(torch.ops.aten.clone.default)
+lower_by_torch_xla2(torch.ops.aten.constant_pad_nd)
+lower_by_torch_xla2(torch.ops.aten.cos)
+lower_by_torch_xla2(torch.ops.aten.cosh)
+lower_by_torch_xla2(torch.ops.aten.cumsum)
+lower_by_torch_xla2(torch.ops.aten.detach)
+lower_by_torch_xla2(torch.ops.aten.diagonal)
+lower_by_torch_xla2(torch.ops.aten.div)
+lower_by_torch_xla2(torch.ops.aten.dot)
+lower_by_torch_xla2(torch.ops.aten.embedding)
+lower_by_torch_xla2(torch.ops.aten.empty)
+lower_by_torch_xla2(torch.ops.aten.eq)
+lower_by_torch_xla2(torch.ops.aten.erf)
+lower_by_torch_xla2(torch.ops.aten.exp)
+lower_by_torch_xla2(torch.ops.aten.expand)
+lower_by_torch_xla2(torch.ops.aten.expand_copy)
+lower_by_torch_xla2(torch.ops.aten.expm1)
+lower_by_torch_xla2(torch.ops.aten.fill)
+lower_by_torch_xla2(torch.ops.aten.flip)
+lower_by_torch_xla2(torch.ops.aten.floor)
+lower_by_torch_xla2(torch.ops.aten.fmod)
+lower_by_torch_xla2(torch.ops.aten.full)
+lower_by_torch_xla2(torch.ops.aten.full_like)
+lower_by_torch_xla2(torch.ops.aten.gather)
+lower_by_torch_xla2(torch.ops.aten.ge)
+lower_by_torch_xla2(torch.ops.aten.gelu)
+lower_by_torch_xla2(torch.ops.aten.glu)
+lower_by_torch_xla2(torch.ops.aten.glu.default)
+lower_by_torch_xla2(torch.ops.aten.gt)
+lower_by_torch_xla2(torch.ops.aten.hardtanh)
+lower_by_torch_xla2(torch.ops.aten.index)
+lower_by_torch_xla2(torch.ops.aten.index.Tensor)
+lower_by_torch_xla2(torch.ops.aten.index_copy)
+lower_by_torch_xla2(torch.ops.aten.index_put)
+lower_by_torch_xla2(torch.ops.aten.index_select)
+lower_by_torch_xla2(torch.ops.aten.isinf)
+lower_by_torch_xla2(torch.ops.aten.isnan)
+lower_by_torch_xla2(torch.ops.aten.le)
+lower_by_torch_xla2(torch.ops.aten.leaky_relu)
+lower_by_torch_xla2(torch.ops.aten.lift_fresh_copy)
+lower_by_torch_xla2(torch.ops.aten.linalg_vector_norm)
+lower_by_torch_xla2(torch.ops.aten.log)
+lower_by_torch_xla2(torch.ops.aten.log10)
+lower_by_torch_xla2(torch.ops.aten.log1p)
+lower_by_torch_xla2(torch.ops.aten.log2)
+lower_by_torch_xla2(torch.ops.aten.logical_and)
+lower_by_torch_xla2(torch.ops.aten.logical_not)
+lower_by_torch_xla2(torch.ops.aten.logical_or)
+lower_by_torch_xla2(torch.ops.aten.logical_xor)
+lower_by_torch_xla2(torch.ops.aten.lt)
+lower_by_torch_xla2(torch.ops.aten.max)
+lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices)
+lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices_backward)
+lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices_backward)
+lower_by_torch_xla2(torch.ops.aten.max_pool3d_with_indices)
+lower_by_torch_xla2(torch.ops.aten.maximum)
+lower_by_torch_xla2(torch.ops.aten.mean)
+lower_by_torch_xla2(torch.ops.aten.min)
+lower_by_torch_xla2(torch.ops.aten.minimum)
+lower_by_torch_xla2(torch.ops.aten.mm)
+lower_by_torch_xla2(torch.ops.aten.mul.Scalar)
+lower_by_torch_xla2(torch.ops.aten.mul.Tensor)
+lower_by_torch_xla2(torch.ops.aten.native_batch_norm)
+lower_by_torch_xla2(torch.ops.aten.native_group_norm)
+lower_by_torch_xla2(torch.ops.aten.native_layer_norm_backward)
+lower_by_torch_xla2(torch.ops.aten.ne)
+lower_by_torch_xla2(torch.ops.aten.neg)
+lower_by_torch_xla2(torch.ops.aten.nonzero)
+lower_by_torch_xla2(torch.ops.aten.outer)
+lower_by_torch_xla2(torch.ops.aten.permute)
+lower_by_torch_xla2(torch.ops.aten.permute_copy)
+lower_by_torch_xla2(torch.ops.aten.pixel_shuffle)
+lower_by_torch_xla2(torch.ops.aten.pow)
+lower_by_torch_xla2(torch.ops.aten.prod)
+lower_by_torch_xla2(torch.ops.aten.rand)
+lower_by_torch_xla2(torch.ops.aten.randn)
+lower_by_torch_xla2(torch.ops.aten.reciprocal)
+lower_by_torch_xla2(torch.ops.aten.reflection_pad1d)
+lower_by_torch_xla2(torch.ops.aten.relu)
+lower_by_torch_xla2(torch.ops.aten.remainder)
+lower_by_torch_xla2(torch.ops.aten.repeat)
+lower_by_torch_xla2(torch.ops.aten.reshape)
+lower_by_torch_xla2(torch.ops.aten.roll)
+lower_by_torch_xla2(torch.ops.aten.round)
+lower_by_torch_xla2(torch.ops.aten.rsqrt)
+lower_by_torch_xla2(torch.ops.aten.scalar_tensor)
+lower_by_torch_xla2(torch.ops.aten.scatter.src)
+lower_by_torch_xla2(torch.ops.aten.scatter.value)
+lower_by_torch_xla2(torch.ops.aten.scatter_add)
+lower_by_torch_xla2(torch.ops.aten.scatter_reduce)
+lower_by_torch_xla2(torch.ops.aten.select)
+lower_by_torch_xla2(torch.ops.aten.select_copy)
+lower_by_torch_xla2(torch.ops.aten.select_scatter)
+lower_by_torch_xla2(torch.ops.aten.sigmoid)
+lower_by_torch_xla2(torch.ops.aten.sign)
+lower_by_torch_xla2(torch.ops.aten.silu)
+lower_by_torch_xla2(torch.ops.aten.sin)
+lower_by_torch_xla2(torch.ops.aten.sinh)
+lower_by_torch_xla2(torch.ops.aten.slice)
+lower_by_torch_xla2(torch.ops.aten.slice_copy)
+lower_by_torch_xla2(torch.ops.aten.sort)
+lower_by_torch_xla2(torch.ops.aten.split)
+lower_by_torch_xla2(torch.ops.aten.split_copy)
+lower_by_torch_xla2(torch.ops.aten.split_with_sizes)
+lower_by_torch_xla2(torch.ops.aten.sqrt)
+lower_by_torch_xla2(torch.ops.aten.squeeze)
+lower_by_torch_xla2(torch.ops.aten.squeeze_copy)
+lower_by_torch_xla2(torch.ops.aten.stack)
+lower_by_torch_xla2(torch.ops.aten.sub.Scalar)
+lower_by_torch_xla2(torch.ops.aten.sub.Tensor)
+lower_by_torch_xla2(torch.ops.aten.sum)
+lower_by_torch_xla2(torch.ops.aten.sym_size)
+lower_by_torch_xla2(torch.ops.aten.t)
+lower_by_torch_xla2(torch.ops.aten.tan)
+lower_by_torch_xla2(torch.ops.aten.tanh)
+lower_by_torch_xla2(torch.ops.aten.tensor_split.sections)
+lower_by_torch_xla2(torch.ops.aten.tensor_split.sections)
+lower_by_torch_xla2(torch.ops.aten.to.device)
+lower_by_torch_xla2(torch.ops.aten.to.device)
+lower_by_torch_xla2(torch.ops.aten.to.dtype)
+lower_by_torch_xla2(torch.ops.aten.topk)
+lower_by_torch_xla2(torch.ops.aten.transpose)
+lower_by_torch_xla2(torch.ops.aten.transpose_copy)
+lower_by_torch_xla2(torch.ops.aten.triu)
+lower_by_torch_xla2(torch.ops.aten.true_divide)
+lower_by_torch_xla2(torch.ops.aten.trunc)
+lower_by_torch_xla2(torch.ops.aten.unbind)
+lower_by_torch_xla2(torch.ops.aten.unbind_copy)
+lower_by_torch_xla2(torch.ops.aten.unsqueeze)
+lower_by_torch_xla2(torch.ops.aten.unsqueeze.default)
+lower_by_torch_xla2(torch.ops.aten.unsqueeze_copy)
+lower_by_torch_xla2(torch.ops.aten.var.correction)
+lower_by_torch_xla2(torch.ops.aten.var_mean.correction)
+lower_by_torch_xla2(torch.ops.aten.view)
+lower_by_torch_xla2(torch.ops.aten.view_as_complex)
+lower_by_torch_xla2(torch.ops.aten.view_as_real)
+lower_by_torch_xla2(torch.ops.aten.view_copy)
+lower_by_torch_xla2(torch.ops.aten.where.ScalarOther)
+lower_by_torch_xla2(torch.ops.aten.where.ScalarSelf)
+lower_by_torch_xla2(torch.ops.aten.where.self)
+lower_by_torch_xla2(torch.ops.prims.broadcast_in_dim)
+lower_by_torch_xla2(torch.ops.prims.var)
+
+
+@lower_by_jax(torch.ops.aten.copy, ir_input_names=["src"])
+def _aten_copy(self, src, **kwargs):
+  return _TORCH_XLA2_IMPLS[torch.ops.aten.copy](self, src)

ai_edge_torch/odml_torch/lowerings/_layer_norm.py

@@ -0,0 +1,78 @@
+# 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.
+# ==============================================================================
+"""Provides lowering for coreaten to stablehlo for LayerNorm."""
+
+import math
+from typing import Optional
+from ai_edge_torch.odml_torch.lowerings import registry
+from ai_edge_torch.odml_torch.lowerings import utils
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+
+# native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight,
+# Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)
+@registry.lower(torch.ops.aten.native_layer_norm)
+def _aten_native_layer_norm(
+    lctx,
+    data: ir.Value,
+    normalized_shape: list[int],
+    weight: Optional[ir.Value],
+    bias: Optional[ir.Value],
+    eps: float,
+):
+  data_type: ir.RankedTensorType = data.type
+  unnormalized_count = math.prod(data_type.shape) // math.prod(normalized_shape)
+  dest_shape = [
+      1,
+      unnormalized_count,
+      math.prod(normalized_shape),
+  ]
+  dest_type = ir.RankedTensorType.get(dest_shape, data_type.element_type)
+
+  reshaped_data = stablehlo.reshape(dest_type, data)
+
+  one = utils.splat(1, data_type.element_type, [unnormalized_count])
+  zero = utils.splat(0, data_type.element_type, [unnormalized_count])
+  output, mean, var = stablehlo.batch_norm_training(
+      reshaped_data, one, zero, eps, 1
+  )
+  eps_splat = utils.splat(eps, var.type.element_type, var.type.shape)
+  rstd = stablehlo.rsqrt(stablehlo.add(var, eps_splat))
+
+  stats_shape = data_type.shape[: -1 * len(normalized_shape)] + [1] * len(
+      normalized_shape
+  )
+  stats_type = ir.RankedTensorType.get(stats_shape, data_type.element_type)
+  mean = stablehlo.reshape(stats_type, mean)
+  rstd = stablehlo.reshape(stats_type, rstd)
+
+  output = stablehlo.reshape(data_type, output)
+
+  data_rank = len(data_type.shape)
+  normalized_rank = len(normalized_shape)
+  if weight is not None:
+    weight = stablehlo.broadcast_in_dim(
+        data_type, weight, list(range(data_rank - normalized_rank, data_rank))
+    )
+    output = stablehlo.multiply(weight, output)
+  if bias is not None:
+    bias = stablehlo.broadcast_in_dim(
+        data_type, bias, list(range(data_rank - normalized_rank, data_rank))
+    )
+    output = stablehlo.add(bias, output)
+
+  return output, mean, rstd

ai_edge_torch/odml_torch/lowerings/context.py

@@ -0,0 +1,42 @@
+# 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.
+# ==============================================================================
+"""Define context object for export and MLIR lowerings."""
+
+import dataclasses
+from jax._src.lib.mlir import ir
+import torch
+
+
+@dataclasses.dataclass
+class LoweringContext:
+  """The context object used in export interpreter and MLIR lowerings."""
+
+  ir_context: ir.Context
+  ir_module: ir.Module
+  ir_location: ir.Location = None
+  node: torch.fx.Node = None
+
+  @property
+  def ctx(self):
+    """Shortcut for ir_context."""
+    return self.ir_context
+
+  @property
+  def loc(self):
+    """Shortcut for ir_location."""
+    return self.ir_location
+
+  def replace(self, **kwargs):
+    return dataclasses.replace(self, **kwargs)

ai_edge_torch/odml_torch/lowerings/registry.py

@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+"""Torch op decompositions and MLIR lowerings registry."""
+
+from typing import Any, Callable
+
+import torch
+
+from . import context
+
+
+class LoweringRegistry:
+  """Registry object for torch op decompositions and to-MLIR lowerings."""
+
+  def __init__(self):
+    self.registered_ops = {}
+    self.decompositions = {}
+
+  def lookup(self, op_or_name):
+    candidate = self._get_lowering(op_or_name)
+    if candidate is None:
+      if isinstance(op_or_name, torch._ops.OpOverloadPacket):
+        candidate = self._get_lowering(op_or_name.default)
+      if isinstance(op_or_name, torch._ops.OpOverload):
+        candidate = self._get_lowering(op_or_name.overloadpacket)
+    return candidate
+
+  def _get_lowering(self, op):
+    candidate = self.registered_ops.get(op)
+    return candidate
+
+  def register(self, op, lowering):
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      ops = [getattr(op, overload) for overload in op.overloads()]
+    else:
+      ops = [op]
+
+    for op in ops:
+      self.registered_ops[op] = lowering
+
+
+global_registry = LoweringRegistry()
+global_registry.decompositions.update(torch._decomp.core_aten_decompositions())
+global_registry.decompositions.update(
+    torch._decomp.get_decompositions([
+        torch.ops.aten.upsample_nearest2d,
+        torch.ops.aten._native_batch_norm_legit.no_stats,
+        torch.ops.aten._native_batch_norm_legit_functional,
+        torch.ops.aten._adaptive_avg_pool2d,
+        torch.ops.aten._adaptive_avg_pool3d,
+        torch.ops.aten.grid_sampler_2d,
+        torch.ops.aten.native_dropout,
+        torch.ops.aten.reflection_pad1d,
+        torch.ops.aten.reflection_pad2d,
+        torch.ops.aten.reflection_pad3d,
+        torch.ops.aten.replication_pad1d,
+        torch.ops.aten.replication_pad2d,
+        torch.ops.aten.replication_pad3d,
+        torch.ops.aten.addmm,
+    ])
+)
+
+torch._decomp.remove_decompositions(
+    global_registry.decompositions,
+    [
+        torch.ops.aten.roll,
+    ],
+)
+
+
+def lookup(op):
+  return global_registry.lookup(op)
+
+
+def lower(op):
+  def inner(lowering: Callable[[context.LoweringContext, ...], Any]):
+    global_registry.register(op, lowering)
+    return lowering
+
+  return inner
+
+
+def decompositions():
+  return global_registry.decompositions

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -0,0 +1,185 @@
+# 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 building MLIR lowerings."""
+
+import numbers
+from typing import Any
+from typing import Optional
+
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import numpy as np
+
+
+def splat(val, ty, shape=tuple(), *, loc: Optional[Any] = None):
+  if isinstance(ty, ir.IntegerType):
+    if ty.width == 1:
+      attr = ir.BoolAttr.get(bool(val))
+    else:
+      attr = ir.IntegerAttr.get(ty, int(val))
+  elif isinstance(ty, ir.FloatType):
+    attr = ir.FloatAttr.get(ty, val)
+  else:
+    raise ValueError("Unsupported type: %s" % str(ty))
+
+  return stablehlo.constant(
+      ir.DenseElementsAttr.get_splat(
+          ir.RankedTensorType.get(shape, ty),
+          attr,
+      ),
+      loc=loc,
+  )
+
+
+def get_common_broadcast_shape(
+    shape_1: list[int], shape_2: list[int]
+) -> Optional[list[int]]:
+  if not shape_1 and not shape_2:
+    return None
+
+  shape_1 = shape_1 if shape_1 else [1]
+  shape_2 = shape_2 if shape_2 else [1]
+
+  length_diff = abs(len(shape_1) - len(shape_2))
+  if len(shape_1) < len(shape_2):
+    shape_1 = [1] * length_diff + shape_1
+  elif len(shape_1) > len(shape_2):
+    shape_2 = [1] * length_diff + shape_2
+
+  common_broadcast_shape = []
+  for idx in reversed(range(len(shape_1))):
+    dim_size1 = shape_1[idx]
+    dim_size2 = shape_2[idx]
+
+    if dim_size1 == dim_size2:
+      common_broadcast_shape.insert(0, dim_size1)
+    elif dim_size1 == 1 or dim_size2 == 1:
+      common_broadcast_shape.insert(0, max(dim_size1, dim_size2))
+    else:
+      return None
+
+  return common_broadcast_shape
+
+
+def get_broadcast_dimensions(
+    shape_from: list[int], shape_to: list[int]
+) -> list[int]:
+  assert get_common_broadcast_shape(shape_from, shape_to) == shape_to
+
+  ret = []
+  for val in range(len(shape_to) - len(shape_from), len(shape_to)):
+    ret.append(val)
+
+  return ir.DenseI64ArrayAttr.get(np.asarray(ret, np.int64))
+
+
+def broadcast_args_if_needed(
+    val_1: ir.Value, val_2: ir.Value
+) -> tuple[Optional[ir.Value], Optional[ir.Value]]:
+  broadcast_shape = get_common_broadcast_shape(
+      val_1.type.shape, val_2.type.shape
+  )
+  if broadcast_shape is None:
+    return None, None
+
+  new_val_1, new_val_2 = val_1, val_2
+
+  if val_1.type.shape != broadcast_shape:
+    new_val_1 = stablehlo.broadcast_in_dim(
+        result=ir.RankedTensorType.get(
+            broadcast_shape, val_1.type.element_type
+        ),
+        operand=val_1,
+        broadcast_dimensions=get_broadcast_dimensions(
+            val_1.type.shape, broadcast_shape
+        ),
+    )
+  if val_2.type.shape != broadcast_shape:
+    new_val_2 = stablehlo.broadcast_in_dim(
+        result=ir.RankedTensorType.get(
+            broadcast_shape, val_2.type.element_type
+        ),
+        operand=val_2,
+        broadcast_dimensions=get_broadcast_dimensions(
+            val_2.type.shape, broadcast_shape
+        ),
+    )
+  return new_val_1, new_val_2
+
+
+def upcast_to_same_type(*vals: ir.Value):
+  if not vals:
+    return None
+  if len(vals) == 1:
+    return vals[0]
+
+  def get_priority(ty: ir.Type):
+    priorities = [
+        ir.IntegerType.get_signless(1),
+        ir.IntegerType.get_signless(16),
+        ir.IntegerType.get_signless(32),
+        ir.IntegerType.get_signless(64),
+        ir.F16Type,
+        ir.F32Type,
+        ir.F64Type,
+    ]
+    for i, tycls in enumerate(priorities):
+      if tycls.isinstance(ty):
+        return i
+    raise ValueError("Unsupported type: %s" % str(ty))
+
+  cast_tycls = type(max([v.type.element_type for v in vals], key=get_priority))
+  new_vals = []
+  for val in vals:
+    if not cast_tycls.isinstance(val.type.element_type):
+      val = stablehlo.convert(
+          ir.RankedTensorType.get(val.type.shape, cast_tycls.get()), val
+      )
+    new_vals.append(val)
+  return tuple(new_vals)
+
+
+def minmax(ty: ir.Type) -> tuple[numbers.Number, numbers.Number]:
+  if isinstance(ty, ir.IntegerType):
+    if ty.is_unsigned:
+      return (0, 1 << ty.width)
+    else:
+      return (-(1 << (ty.width - 1)), (1 << (ty.width - 1)) - 1)
+  elif isinstance(ty, ir.F16Type):
+    return (np.finfo(np.float16).min, np.finfo(np.float16).max)
+  elif isinstance(ty, ir.F32Type):
+    return (np.finfo(np.float32).min, np.finfo(np.float32).max)
+  elif isinstance(ty, ir.F64Type):
+    return (np.finfo(np.float64).min, np.finfo(np.float64).max)
+  else:
+    raise ValueError("Unsupported type: %s" % ty)
+
+
+def convert_int_to_float(t: ir.Value) -> ir.Value:
+  """Converts an input with type ir.IntegerType to an ir.FloatType of equivalent width."""
+  elty = t.type.element_type
+  if not isinstance(elty, ir.IntegerType):
+    raise ValueError(
+        "Expected input with integer type, received %s" % type(elty)
+    )
+
+  if elty.width == 32:
+    return stablehlo.convert(
+        ir.RankedTensorType.get(t.type.shape, ir.F32Type.get()), t
+    )
+  elif elty.width == 64:
+    return stablehlo.convert(
+        ir.RankedTensorType.get(t.type.shape, ir.F64Type.get()), t
+    )

ai_edge_torch/odml_torch/passes/__init__.py

@@ -0,0 +1,38 @@
+# 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 jax._src.lib.mlir import ir
+from jax._src.lib.mlir import passmanager
+
+
+def run_pass(pipeline, module: ir.Module):
+  pm = passmanager.PassManager.parse(pipeline)
+  pm.run(module.operation)
+  return module
+
+
+def canonicalize(module: ir.Module):
+  return run_pass("builtin.module(canonicalize)", module)
+
+
+def cse(module: ir.Module):
+  return run_pass("builtin.module(cse)", module)
+
+
+def inline(module: ir.Module):
+  return run_pass("builtin.module(inline)", module)
+
+
+def strip_debuginfo(module: ir.Module):
+  return run_pass("builtin.module(strip-debuginfo)", module)