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/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,25 @@
+# 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 _layer_norm
+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,258 @@
+# 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)
+
+
+# Schema:
+#   - aten::slice_scatter(Tensor self, Tensor src, int dim=0, SymInt?
+#       start=None, SymInt? end=None, SymInt step=1) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/stable/generated/torch.slice_scatter.html
+#   - https://github.com/pytorch/pytorch/blob/18f9331e5deb4c02ae5c206e133a9b4add49bd97/aten/src/ATen/native/TensorShape.cpp#L4002
+@lower(torch.ops.aten.slice_scatter)
+def _aten_slice_scatter(lctx, self, src, dim=0, start=None, end=None, step=1):
+  start = start if start is not None else 0
+  end = end if end is not None else self.type.shape[dim]
+
+  start, end = np.clip(
+      [start, end], -self.type.shape[dim], self.type.shape[dim]
+  )
+
+  if start < 0:
+    start = self.type.shape[dim] + start
+  if end < 0:
+    end = self.type.shape[dim] + end
+
+  if end <= start or np.prod(src.type.shape) == 0:
+    return self
+
+  end = start + step * math.ceil((end - start) / step) - (step - 1)
+  padding_low = start
+  padding_high = self.type.shape[dim] - end
+  interior_padding = step - 1
+
+  rank = len(self.type.shape)
+  src = stablehlo.pad(
+      src,
+      utils.splat(0, src.type.element_type, []),
+      edge_padding_low=[padding_low if i == dim else 0 for i in range(rank)],
+      edge_padding_high=[padding_high if i == dim else 0 for i in range(rank)],
+      interior_padding=[
+          interior_padding if i == dim else 0 for i in range(rank)
+      ],
+  )
+  pred = np.ones(self.type.shape, dtype=np.bool_)
+  pred[*[
+      slice(start, end, step) if i == dim else slice(None, None, None)
+      for i in range(rank)
+  ]] = False
+  pred = stablehlo.constant(
+      ir.DenseElementsAttr.get(
+          np.packbits(pred, bitorder="little"),
+          type=ir.IntegerType.get_signless(1),
+          shape=pred.shape,
+      )
+  )
+  out = stablehlo.select(pred, self, src)
+  return out

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,241 @@
+# 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 Convolution."""
+
+import math
+from typing import Optional
+
+from ai_edge_torch.odml_torch.lowerings import registry
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+
+def make_padding(padding):
+  """Change the padding from pytorch to stablehlo style.
+
+  Stablehlo allows start and end padding for each dimension while aten only
+  allows symmetric padding and so only has one number per dimension.
+
+  Args:
+    padding: The padding of the convolution
+
+  Returns:
+    The padding in stablehlo style
+  """
+  return tuple((p, p) for p in padding)
+
+
+def create_conv_dimension_numbers(lhs, transposed: bool = False):
+  """Create the dimension numbers for the convolution.
+
+  Args:
+    lhs: The input tensor
+    transposed: Whether the convolution is transposed
+
+  Returns:
+    The dimension numbers for the convolution
+  """
+  num_spatial_dims = len(lhs.type.shape) - 2
+  spatial_dimensions = []
+  for i in range(0, num_spatial_dims):
+    spatial_dimensions.append(i + 2)
+
+  # Regular kernels are OIHW
+  # TransposedConv kernels are IOHW
+  dimension_numbers = stablehlo.ConvDimensionNumbers.get(
+      input_batch_dimension=0,
+      input_feature_dimension=1,
+      input_spatial_dimensions=spatial_dimensions,
+      kernel_input_feature_dimension=0 if transposed else 1,
+      kernel_output_feature_dimension=1 if transposed else 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,
+    rhs,
+    stride,
+    dilation,
+    padding,
+    transposed: bool = False,
+    output_padding: list[int] = 0,
+):
+  """Infer the output shape of the convolution.
+
+  Args:
+    lhs: The input tensor
+    rhs: The kernel tensor
+    stride: The stride of the convolution (dilation of input in transposed conv)
+    dilation: The kernel dilation of the convolution
+    padding: The padding of the convolution
+    transposed: Whether the convolution is transposed
+    output_padding: The output padding of the convolution
+
+  Returns:
+    The output shape of the convolution
+  """
+  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 for regular conv
+  # Input layout is: (N)CHW and Kernel layout is: I(O)HW for transposed conv
+  output_shape = (
+      [lhs_shape[0], rhs_shape[1]]
+      if transposed
+      else [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 = spatial_dim + 2
+    dim_size = lhs_shape[dim]
+    kernel_dim_size = rhs_shape[dim]
+
+    if transposed:
+      output_dim_size = (
+          (dim_size - 1) * stride[spatial_dim]
+          - 2 * padding[spatial_dim]
+          + dilation[spatial_dim] * (kernel_dim_size - 1)
+          + output_padding[spatial_dim]
+          + 1
+      )
+    else:
+      output_dim_size = math.floor(
+          (
+              (
+                  dim_size
+                  + 2 * padding[spatial_dim]
+                  - dilation[spatial_dim] * (kernel_dim_size - 1)
+                  - 1
+              )
+              / stride[spatial_dim]
+          )
+          + 1
+      )
+
+    output_shape.append(output_dim_size)
+
+  return output_shape
+
+
+def build_transpose_conv(
+    lctx,
+    output_type: ir.RankedTensorType,
+    lhs: ir.Value,
+    rhs: ir.Value,
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    output_padding: list[int],
+    groups: int,
+):
+  lhs_type: ir.RankedTensorType = lhs.type
+  num_spatial_dims = len(lhs_type.shape) - 2
+  rhs = stablehlo.reverse(rhs, list(range(2, 2 + num_spatial_dims)))
+
+  kernel_size = rhs.type.shape
+  # We need to additional padding on the input to get the right output size.
+  adjusted_padding = [
+      dilation[dim] * (kernel_size[dim + 2] - 1) - padding[dim]
+      for dim in range(num_spatial_dims)
+  ]
+  return stablehlo.convolution(
+      result=output_type,
+      lhs=lhs,
+      rhs=rhs,
+      dimension_numbers=create_conv_dimension_numbers(lhs, True),
+      feature_group_count=groups,
+      batch_group_count=1,
+      padding=make_padding(adjusted_padding),
+      lhs_dilation=stride,
+      rhs_dilation=dilation,
+  )
+
+
+# convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride,
+#   SymInt[] padding, SymInt[] dilation, bool transposed,
+#   SymInt[] output_padding, SymInt groups) -> Tensor
+@registry.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,
+):
+
+  # TODO(b/365559296) Add support for output_padding
+  if any(output_padding):
+    raise NotImplementedError(
+        "Output padding on convolution is not implemented."
+    )
+
+  lhs_type: ir.RankedTensorType = lhs.type
+  output_shape = infer_output_shape(
+      lhs, rhs, stride, dilation, padding, transposed, output_padding
+  )
+  output_type = ir.RankedTensorType.get(
+      output_shape,
+      lhs_type.element_type,
+  )
+
+  if transposed:
+    res = build_transpose_conv(
+        lctx,
+        output_type,
+        lhs,
+        rhs,
+        stride,
+        padding,
+        dilation,
+        output_padding,
+        groups,
+    )
+  else:
+    res = stablehlo.convolution(
+        result=output_type,
+        lhs=lhs,
+        rhs=rhs,
+        dimension_numbers=create_conv_dimension_numbers(lhs),
+        feature_group_count=groups,
+        batch_group_count=1,
+        window_strides=stride,
+        padding=make_padding(padding),
+        rhs_dilation=dilation,
+    )
+
+  if bias is not None:
+    # broadcast [C] to [NCHW]
+    broadcasted_bias = stablehlo.broadcast_in_dim(output_type, bias, [1])
+    res = stablehlo.add(
+        lhs=res,
+        rhs=broadcasted_bias,
+    )
+
+  return res