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/jax_bridge/utils.py

@@ -0,0 +1,75 @@
+# 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 Jax bridge."""
+
+from ai_edge_torch import odml_torch
+import jax
+import jax.numpy as jnp
+from jax._src.lib.mlir import ir
+import torch
+
+
+def t2j_dtype(dtype):
+  return {
+      torch.bfloat16: jnp.bfloat16,
+      torch.half: jnp.float16,
+      torch.float32: jnp.float32,
+      torch.double: jnp.double,
+      torch.long: jnp.int64,
+      torch.int64: jnp.int64,
+      torch.int32: jnp.int32,
+      torch.int16: jnp.int16,
+      torch.int8: jnp.int8,
+      torch.uint8: jnp.uint8,
+      torch.bool: jnp.bool_,
+      torch.complex64: jnp.complex64,
+      torch.complex128: jnp.complex128,
+  }.get(dtype)
+
+
+def is_ir_variable(value):
+  if isinstance(value, ir.Value):
+    return True
+  if isinstance(value, (list, tuple)):
+    return any(is_ir_variable(x) for x in value)
+  return False
+
+
+def ir_variable_to_jax(value):
+  if isinstance(value, (list, tuple)):
+    return tuple([ir_variable_to_jax(x) for x in value])
+  elif not isinstance(value, ir.Value):
+    return value
+  elif not isinstance(value.type, ir.RankedTensorType):
+    raise ValueError(
+        f"ir.Value to JAX must be in ir.RankedTensorType, got {value}"
+    )
+
+  return jax.ShapeDtypeStruct(
+      value.type.shape,
+      t2j_dtype(
+          odml_torch.export_utils.ir_element_type_to_torch_dtype(
+              value.type.element_type
+          )
+      ),
+  )
+
+
+def tree_map_list_to_tuple(value):
+  if isinstance(value, dict):
+    return {k: tree_map_list_to_tuple(v) for k, v in value.items()}
+  if isinstance(value, (list, tuple)):
+    return tuple([tree_map_list_to_tuple(v) for v in value])
+  return value

ai_edge_torch/odml_torch/lowerings/__init__.py

@@ -0,0 +1,24 @@
+# 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 _basic
+from . import _batch_norm
+from . import _convolution
+from . import _jax_lowerings
+from . import context
+from . import registry
+from . import utils
+from .registry import decompositions
+from .registry import lookup
+from .registry import lower

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -0,0 +1,204 @@
+# 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 math
+from typing import Optional, Union
+
+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 numpy as np
+import torch
+
+from .registry import lower
+
+
+# add(Tensor self, Tensor other) -> Tensor
+# @lower(torch.ops.aten.add)
+def _aten_add(lctx, x: ir.Value, y: ir.Value, alpha=1):
+  x, y = utils.upcast_to_same_type(x, y)
+  x, y = utils.broadcast_args_if_needed(x, y)
+  if alpha == 1:
+    return stablehlo.add(x, y)
+
+  alpha_splat = utils.splat(alpha, y.type.element_type, y.type.shape)
+  return stablehlo.add(x, stablehlo.multiply(y, alpha_splat))
+
+
+# mul.Tensor(Tensor self, Tensor other) -> Tensor
+# @lower(torch.ops.aten.mul.Tensor)
+def _aten_mul_tensor(lctx, self: ir.Value, other: ir.Value):
+  self, other = utils.upcast_to_same_type(self, other)
+  self, other = utils.broadcast_args_if_needed(self, other)
+
+  return stablehlo.multiply(self, other)
+
+
+# cat(Tensor[] tensors, int dim=0) -> Tensor
+# @lower(torch.ops.aten.cat)
+def _aten_cat(lctx, tensors: list[ir.Value], dim: int = 1):
+  return stablehlo.ConcatenateOp(tensors, dim).result
+
+
+# view(Tensor(a) self, SymInt[] size) -> Tensor(a)
+# @lower(torch.ops.aten.view)
+def _aten_view(lctx, self: ir.Value, size: list[int]):
+  return stablehlo.ReshapeOp(
+      ir.RankedTensorType.get(size, self.type.element_type), self
+  ).result
+
+
+# hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> Tensor
+@lower(torch.ops.aten.hardtanh)
+def _aten_hardtanh(
+    lctx,
+    self: ir.Value,
+    min_val: Union[int, float] = -1.0,
+    max_val: Union[int, float] = 1.0,
+):
+  elty = self.type.element_type
+  min_val = utils.splat(min_val, elty)
+  max_val = utils.splat(max_val, elty)
+
+  return stablehlo.clamp(min_val, self, max_val)
+
+
+# mean(Tensor self, *, ScalarType? dtype=None) -> Tensor
+# mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *,
+#   ScalarType? dtype=None) -> Tensor
+@lower(torch.ops.aten.mean)
+@lower(torch.ops.aten.mean.dim)
+def _aten_mean_dim(
+    lctx,
+    self: ir.Value,
+    dim: Optional[list[int]] = None,
+    keepdim: bool = False,
+    *,
+    dtype=None,
+):
+  self_shape = self.type.shape
+  self_elty = self.type.element_type
+  if dim is None:
+    dim = list(range(len(self_shape)))
+  dim = [len(self_shape) + d if d < 0 else d for d in dim]
+  dim_ = ir.DenseI64ArrayAttr.get(np.asarray(dim, np.int64))
+  dim_to_keep = [d for d in range(len(self_shape)) if d not in dim]
+  dim_to_keep_ = ir.DenseI64ArrayAttr.get(np.asarray(dim_to_keep, np.int64))
+
+  zero_ = utils.splat(0.0, self_elty)
+
+  reduce_result_shape = [
+      s for d, s in enumerate(self_shape) if d in dim_to_keep
+  ]
+  reduce_result_ty = ir.RankedTensorType.get(reduce_result_shape, self_elty)
+  reduce_op = stablehlo.ReduceOp([reduce_result_ty], [self], [zero_], dim_)
+
+  reducer_arg_ty = ir.RankedTensorType.get(tuple(), self_elty)
+  reducer = reduce_op.regions[0].blocks.append(reducer_arg_ty, reducer_arg_ty)
+  with ir.InsertionPoint(reducer):
+    stablehlo.return_(
+        [stablehlo.add(reducer.arguments[0], reducer.arguments[1])]
+    )
+
+  sum_ = reduce_op.result
+  if keepdim:
+    sum_ = stablehlo.broadcast_in_dim(
+        ir.RankedTensorType.get(
+            [s if d in dim_to_keep else 1 for d, s in enumerate(self_shape)],
+            self_elty,
+        ),
+        sum_,
+        dim_to_keep_,
+    )
+
+  dim_els = math.prod([s for d, s in enumerate(self_shape) if d in dim])
+  dim_els_ = utils.splat(dim_els, self_elty)
+  div_ = stablehlo.broadcast_in_dim(
+      sum_.type, dim_els_, ir.DenseI64ArrayAttr.get([])
+  )
+  mean_ = stablehlo.divide(sum_, div_)
+
+  return mean_
+
+
+# https://pytorch.org/docs/stable/generated/torch.clone.html
+# https://github.com/pytorch/pytorch/blob/a95ceb51a23ae33c00b3a99224143c609b1b3eb3/aten/src/ATen/native/TensorFactories.cpp#L1730
+@lower(torch.ops.aten.clone)
+def _aten_clone(lctx, x: ir.Value, *, memory_format=None):
+  return x
+
+
+# https://pytorch.org/docs/stable/generated/torch.permute.html
+# https://github.com/pytorch/pytorch/blob/519151a062a9bd4f0d32a9c7c7eae47d7ed847b2/aten/src/ATen/native/TensorShape.cpp#L1448
+# https://github.com/openxla/stablehlo/blob/main/docs/spec.md#transpose
+@lower(torch.ops.aten.permute)
+def _aten_permute(lctx, x: ir.Value, dims: list[int]):
+  dim = len(x.type.shape)
+  return stablehlo.transpose(x, ir.DenseI64ArrayAttr.get(dims))
+
+
+# https://pytorch.org/docs/stable/generated/torch.mm.html
+# https://github.com/pytorch/pytorch/blob/ffabb25c489df1dc631a577c12a0c843c8b202f3/aten/src/ATen/native/LinearAlgebra.cpp#L193
+# https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dot_general
+@lower(torch.ops.aten.mm)
+def _aten_mm(mod, mat1: ir.Value, mat2: ir.Value) -> ir.Value:
+  mat1_shape = mat1.type.shape
+  mat2_shape = mat2.type.shape
+  mat1_dims = len(mat1_shape)
+  mat2_dims = len(mat2_shape)
+
+  if mat1_dims != 2 or mat1_dims != 2:
+    raise ValueError(
+        "Both arguments must be 2D matrices, received dimensions %d and %d"
+        % (mat1_dims, mat2_dims)
+    )
+
+  if mat1_shape[1] != mat2_shape[0]:
+    raise ValueError(
+        "mat1 and mat2 shapes cannot be multiplied, received shapes %s and %s"
+        % (mat1_shape, mat2_shape)
+    )
+
+  dot_dnums = stablehlo.DotDimensionNumbers.get(
+      lhs_batching_dimensions=[],
+      rhs_batching_dimensions=[],
+      lhs_contracting_dimensions=(1,),
+      rhs_contracting_dimensions=(0,),
+  )
+  return stablehlo.dot_general(
+      ir.RankedTensorType.get(
+          (mat1.type.shape[0], mat2.type.shape[1]), mat1.type.element_type
+      ),
+      mat1,
+      mat2,
+      dot_dnums,
+  )
+
+
+# https://pytorch.org/docs/stable/generated/torch.div.html
+# https://openxla.org/stablehlo/spec#divide
+# TODO: support rounding mode and type promotion (see torch.div spec).
+# @lower(torch.ops.aten.div)
+def _aten_div(mod, x, y, *, rounding_mode=None, out=None) -> ir.Value:
+  # By default, PyTorch performs a "true" division like Python 3. This requires
+  # casting integer input types to float to achieve the same semantics using
+  # stablehlo.divide.
+  if isinstance(x.type.element_type, ir.IntegerType):
+    x = utils.convert_int_to_float(x)
+  if isinstance(y.type.element_type, ir.IntegerType):
+    y = utils.convert_int_to_float(y)
+
+  x, y = utils.broadcast_args_if_needed(x, y)
+
+  return stablehlo.divide(x, y)

ai_edge_torch/odml_torch/lowerings/_batch_norm.py

@@ -0,0 +1,65 @@
+# 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 mlir stablehlo op: Convolution"""
+
+from typing import Optional
+
+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
+
+from .registry import lower
+
+
+# _native_batch_norm_legit_no_training(
+#     Tensor input,
+#     Tensor? weight,
+#     Tensor? bias,
+#     Tensor running_mean,
+#     Tensor running_var,
+#     float momentum,
+#     float eps) -> (Tensor, Tensor, Tensor)
+@lower(torch.ops.aten._native_batch_norm_legit_no_training)
+def _native_batch_norm_legit_no_training(
+    lctx,
+    input_tensor: ir.Value,
+    weight: Optional[ir.Value],
+    bias: Optional[ir.Value],
+    running_mean: ir.Value,
+    running_var: ir.Value,
+    momentum: float,
+    eps: float,
+):
+  if weight is None:
+    weight = utils.splat(
+        1, running_mean.type.element_type, running_mean.type.shape
+    )
+  if bias is None:
+    bias = utils.splat(
+        0, running_mean.type.element_type, running_mean.type.shape
+    )
+
+  return [
+      stablehlo.batch_norm_inference(
+          input_tensor, weight, bias, running_mean, running_var, eps, 1
+      ),
+      utils.splat(
+          0, input_tensor.type.element_type
+      ),  # TODO: return empty array instead
+      utils.splat(
+          0, input_tensor.type.element_type
+      ),  # TODO: return empty array instead
+  ]

ai_edge_torch/odml_torch/lowerings/_convolution.py

@@ -0,0 +1,119 @@
+# 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 mlir stablehlo op: Convolution"""
+
+import math
+from typing import Optional
+
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+from .registry import lower
+
+
+# convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride,
+#   SymInt[] padding, SymInt[] dilation, bool transposed,
+#   SymInt[] output_padding, SymInt groups) -> Tensor
+# @lower(torch.ops.aten.convolution)
+def _aten_convolution(
+    lctx,
+    lhs: ir.Value,
+    rhs: ir.Value,
+    bias: Optional[ir.Value],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    transposed: bool,
+    output_padding: list[int],
+    groups: int,
+):
+  if transposed:
+    raise NotImplementedError("Transposed convolution is not implemented.")
+
+  if bias is not None:
+    raise NotImplementedError("Bias on convolution is not implemented.")
+
+  # Stablehlo allows start and end padding for each dimension while aten only
+  # allows symmetric padding and so only has one number per dimension.
+  def make_padding(padding):
+    return tuple((p, p) for p in padding)
+
+  def create_conv_dimension_numbers():
+    num_spatial_dims = len(lhs.type.shape) - 2
+    spatial_dimensions = []
+    for i in range(0, num_spatial_dims):
+      spatial_dimensions.append(i + 2)
+
+    dimension_numbers = stablehlo.ConvDimensionNumbers.get(
+        input_batch_dimension=0,
+        input_feature_dimension=1,
+        input_spatial_dimensions=spatial_dimensions,
+        kernel_input_feature_dimension=1,
+        kernel_output_feature_dimension=0,
+        kernel_spatial_dimensions=spatial_dimensions,
+        output_batch_dimension=0,
+        output_feature_dimension=1,
+        output_spatial_dimensions=spatial_dimensions,
+    )
+    return dimension_numbers
+
+  def infer_output_shape():
+    lhs_type: ir.RankedTensorType = lhs.type
+    lhs_shape: list[int] = lhs_type.shape
+    rhs_shape: list[int] = rhs.type.shape
+
+    # Input layout is: (N)CHW and Kernel layout is: (O)IHW
+    output_shape = [lhs_shape[0], rhs_shape[0]]
+    num_spatial_dims = len(lhs.type.shape) - 2
+
+    # looping over the spatial dims (skipping the first 2 dims which are
+    # batch and features)
+    for spatial_dim in range(0, num_spatial_dims):
+      dim_size = lhs_shape[spatial_dim + 2]
+      kernel_dim_size = rhs_shape[spatial_dim + 2]
+
+      # for example, a dilation of 2 increases the dimension size by 2
+      dim_size *= dilation[spatial_dim]
+
+      # padding added to both sides
+      dim_size += 2 * padding[spatial_dim]
+
+      output_dim_size = math.ceil(
+          (dim_size - kernel_dim_size + 1) / stride[spatial_dim]
+      )
+
+      output_shape.append(output_dim_size)
+
+    return output_shape
+
+  lhs_type: ir.RankedTensorType = lhs.type
+
+  op = stablehlo.ConvolutionOp(
+      result=ir.RankedTensorType.get(
+          infer_output_shape(), lhs_type.element_type
+      ),
+      lhs=lhs,
+      rhs=rhs,
+      dimension_numbers=create_conv_dimension_numbers(),
+      feature_group_count=groups,
+      batch_group_count=1,
+      window_strides=stride,
+      padding=make_padding(padding),
+      lhs_dilation=(1,) * len(stride),
+      rhs_dilation=dilation,
+  )
+
+  return op.result