xref: /external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/dialect/StablehloOps.td

/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
   Copyright 2022 The StableHLO 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.
==============================================================================*/

#ifndef STABLEHLO_DIALECT_STABLEHLO_OPS
#define STABLEHLO_DIALECT_STABLEHLO_OPS

include "dialect/Base.td"
include "mlir/Dialect/Shape/IR/ShapeBase.td"
include "mlir/IR/OpBase.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/IR/OpAsmInterface.td"

def StableHLO_Dialect : Dialect {
  let name = "stablehlo";
  let cppNamespace = "::mlir::stablehlo";

  let description = [{
    StableHLO is an operation set that expresses ML computations. It has been
    originally bootstrapped from the MHLO dialect and enhances it with additional
    functionality, including serialization and versioning, to be used as
    a portability layer between ML frameworks and ML compilers.
  }];

  let emitAccessorPrefix = kEmitAccessorPrefix_Raw;
  let useDefaultAttributePrinterParser = 0;
  let useDefaultTypePrinterParser = 0;
}

class StableHLO_Op<string mnemonic, list<Trait> traits> :
    Op<StableHLO_Dialect, mnemonic, traits> {
}

include "dialect/StablehloEnums.td"
include "dialect/StablehloAttrs.td"

class StableHLO_ShapedInterfaceOp<string mnemonic, list<Trait> traits> :
    StableHLO_Op<mnemonic, traits # [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
    ["reifyReturnTypeShapes"]>]> {
}

//===----------------------------------------------------------------------===//
// StableHLO nullary op definitions.
//===----------------------------------------------------------------------===//

def StableHLO_ConstantOp : StableHLO_Op<"constant",
    [ConstantLike, NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Constant operator";
  let description = [{
    Represents a constant value.
  }];
  let arguments = (ins
    ElementsAttr:$value
  );

  let results = (outs
    HLO_StaticShapeTensor:$output
  );

  let builders = [
    OpBuilder<(ins "Attribute":$value)>];

  let hasCustomAssemblyFormat = 1;
  let hasFolder = 1;

  let extraClassDeclaration = [{
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def StableHLO_IotaOp : StableHLO_Op<"iota", [NoSideEffect]> {
  let summary = "Iota operator";
  let description = [{
    Creates a rank 1 array of values starting at zero and incrementing by one.
  }];
  let arguments = (ins I64Attr:$iota_dimension);

  let results = (outs HLO_IntFpOrComplexTensor:$output);

  let hasVerifier = 1;
}

def StableHLO_DynamicIotaOp: StableHLO_ShapedInterfaceOp<"dynamic_iota", [NoSideEffect]> {
  let summary = "Create linear increasing values from 0 to length -1.";
  let description = [{
    Produces an HLO Tensor of the specified shape, with an incremental set of
    values along the specified dimension starting at 0.

    Requires:
    - The output length of the tensor result.
  }];

  let arguments = (ins HLO_DimensionTensor:$output_shape, I64Attr:$iota_dimension);
  let results = (outs HLO_Tensor:$result);
}

def StableHLO_CreateTokenOp : StableHLO_Op<"create_token", [NoSideEffect]> {
  let summary = "Create Token operator";

  let description = [{
    Produces a HLO token. Tokens are used for ordering side-effecting operations.
    This is exported to HLO as an AfterAll operation with no operands to
    generate a token.

    Example:

    ```mlir
    %1 = stablehlo.create_token : !stablehlo.token
    ```
  }];

  let results = (outs HLO_Token:$output);

  let assemblyFormat = "attr-dict `:` type(results)";
}

//===----------------------------------------------------------------------===//
// StableHLO unary elementwise op definitions.
//===----------------------------------------------------------------------===//
// See https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions

class StableHLO_UnaryElementwiseOp<string mnemonic, list<Trait> traits,
    Type OperandType, Type ResultType = OperandType> : StableHLO_Op<mnemonic, traits # [Elementwise,
    InferShapedTypeOpInterface, SameOperandsAndResultShape]> {
  let arguments = (ins OperandType:$operand);
  let results = (outs ResultType:$result);
  let extraClassDeclaration = [{
    LogicalResult reifyReturnTypeShapes(
        OpBuilder& builder, ValueRange operands,
        SmallVectorImpl<Value>& reifiedReturnShapes) {
      return ::mlir::hlo::deriveShapeFromOperand(&builder, getOperation(),
                                                operands.front(),
                                                &reifiedReturnShapes);
    }
    // Relax the strict default implementation with one that allows
    // for StableHLO-specific differences.
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      if (l.size() != r.size()) return false;
      for (auto [lt, rt] : llvm::zip(l, r))
        if (!mlir::hlo::isCompatibleForHloTypeInference(lt, rt))
          return false;
      return true;
    }
  }];
  let extraClassDefinition = [{
    ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) {
      return ::mlir::stablehlo::parseUnaryOp(parser, result);
    }
    void $cppClass::print(OpAsmPrinter &p) {
      ::mlir::stablehlo::printUnaryOp(getOperation(), p);
    }
  }];
  let hasCustomAssemblyFormat = 1;
}

// Abs supports complex to real, so element type is not guaranteed to match.
def StableHLO_AbsOp: StableHLO_UnaryElementwiseOp<"abs",
    [NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>],
     TensorOf<[HLO_SInt, HLO_Float, HLO_Complex]>> {
  let summary = "Absolute value operator";
  let description = [{
    Returns `abs(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_CbrtOp: StableHLO_UnaryElementwiseOp<"cbrt",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpTensor> {
  let summary = "Cubic root operator";
  let description = [{
    Returns element-wise cubic root of the operand.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_CeilOp: StableHLO_UnaryElementwiseOp<"ceil",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpTensor> {
  let summary = "Ceil operator";
  let description = [{
    Returns `Ceil(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_ConvertOp : StableHLO_UnaryElementwiseOp<"convert",
    [NoSideEffect, SameOperandsAndResultShape], HLO_Tensor> {
  let summary = "Convert operator";
  let description = [{
    Performs element-wise conversion of values from one type to another, e.g.
    float to int.

    See https://www.tensorflow.org/xla/operation_semantics#convertelementtype.
  }];
  let builders = [
    OpBuilder<(ins "Value":$operand, "Type":$result_element_ty)>];
}

def StableHLO_ClzOp: StableHLO_UnaryElementwiseOp<"count_leading_zeros",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_IntTensor> {
  let summary = "Count-leading-zeros (Clz) operator";
  let description = [{
    Returns the number of leading zeros in each operand element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_CosineOp: StableHLO_UnaryElementwiseOp<"cosine",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Cos operator";
  let description = [{
    Returns `Cos(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_ExpOp: StableHLO_UnaryElementwiseOp<"exponential",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Exponential operator";
  let description = [{
    Returns `e^(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_Expm1Op: StableHLO_UnaryElementwiseOp<"exponential_minus_one",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Exponential minus one operator";
  let description = [{
    Returns `e^(operand) - 1` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_FloorOp: StableHLO_UnaryElementwiseOp<"floor",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpTensor> {
  let summary = "Floor operator";
  let description = [{
    Returns `Floor(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_ImagOp: StableHLO_UnaryElementwiseOp<"imag",
    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>],
    HLO_FpOrComplexTensor> {
  let summary = "Imag operator";
  let description = [{
    Returns `Imag(operand)` element-wise.
  }];
  let results = (outs HLO_FpTensor);
}

def StableHLO_IsFiniteOp: StableHLO_UnaryElementwiseOp<"is_finite", [NoSideEffect,
    DeclareOpInterfaceMethods<InferTypeOpInterface>], HLO_Tensor> {
  let summary = "IsFinite operator";
  let description = [{
    Tests whether each element of operand is finite, i.e., is not positive or
    negative infinity, and is not NaN. Returns a tensor of 1-bit integers with
    the same shape as the input, where each element is nonzero (i.e. true) if
    and only if the corresponding input element is finite.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
  let arguments = (ins HLO_FpTensor:$x);
  let results = (outs HLO_PredTensor:$y);
}

def StableHLO_LogOp: StableHLO_UnaryElementwiseOp<"log",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Logarithm operator";
  let description = [{
    Returns `log(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_Log1pOp: StableHLO_UnaryElementwiseOp<"log_plus_one",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Log1p operator";
  let description = [{
    Returns `log(operand+1)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_LogisticOp: StableHLO_UnaryElementwiseOp<"logistic",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Logistic operator";
  let description = [{
    Returns `logistic(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_NotOp: StableHLO_UnaryElementwiseOp<"not",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_PredOrIntTensor> {
  let summary = "Not operator";
  let description = [{
    Returns biwise-NOT of `operand` element-wise. The input tensor must be
    of type integer `HLO_Int` or boolean `HLO_Pred`.

    Note: For boolean tensor, the bitwise-NOT is equivalent to logical-NOT.
  }];
}

def StableHLO_NegOp: StableHLO_UnaryElementwiseOp<"negate",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_IntFpOrComplexTensor> {
  let summary = "Negation operator";
  let description = [{
    Returns `-operand` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_PopulationCountOp: StableHLO_UnaryElementwiseOp<"popcnt",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_IntTensor> {
  let summary = "PopulationCount operator";
  let description = [{
    Returns the number of bits set in each operand element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_RealOp: StableHLO_UnaryElementwiseOp<"real",
    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>],
    HLO_FpOrComplexTensor> {
  let summary = "Real operator";
  let description = [{
    Returns `Real(operand)` element-wise.
  }];
  let results = (outs HLO_FpTensor);
}

def StableHLO_RoundOp: StableHLO_UnaryElementwiseOp<"round_nearest_afz",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpTensor> {
  let summary = "Round operator, ties away from zero";
  let description = [{
    Returns `Round(operand)` element-wise, rounding to nearest integer with
    half-way cases rounding away from zero.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_RoundNearestEvenOp: StableHLO_UnaryElementwiseOp<"round_nearest_even",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpTensor> {
  let summary = "Round operator, ties to even";
  let description = [{
    Returns `Round(operand)` element-wise, rounding to nearest integer with
    half-way cases rounding towards even numbers.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_RsqrtOp: StableHLO_UnaryElementwiseOp<"rsqrt",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Reciprocal Square-root operator";
  let description = [{
    Returns `1.0 / sqrt(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
def StableHLO_SignOp: StableHLO_UnaryElementwiseOp<"sign",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType],
    TensorOf<[HLO_SInt, HLO_Float, HLO_Complex]>> {
  let summary = "Sign operator";
  let description = [{
    Returns `sign(operand)` element-wise, where

    ```
    sign(x) = -1  : x < 0
            = -0  : x = -0
            = NaN : x = NaN
            = +0  : x = +0
            = 1   : x > 0
    ```

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_SineOp: StableHLO_UnaryElementwiseOp<"sine",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Sin operator";
  let description = [{
    Returns `Sin(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_SqrtOp: StableHLO_UnaryElementwiseOp<"sqrt",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType], HLO_FpOrComplexTensor> {
  let summary = "Square-root operator";
  let description = [{
    Returns `sqrt(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}

def StableHLO_TanhOp: StableHLO_UnaryElementwiseOp<"tanh",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType],
    HLO_FpOrComplexTensor> {
  let summary = "Tanh operator";
  let description = [{
    Returns `tanh(operand)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions.
  }];
}
//===----------------------------------------------------------------------===//
// StableHLO binary elementwise op definitions.
//===----------------------------------------------------------------------===//
// See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations

class StableHLO_BinaryElementwiseOpNoAssembly<string mnemonic, list<Trait> traits> :
    StableHLO_Op<mnemonic, traits # [InferShapedTypeOpInterface,
    SameOperandsAndResultShape, Elementwise]> {
  let arguments = (ins
    HLO_Tensor:$lhs,
    HLO_Tensor:$rhs
  );

  let extraClassDeclaration = [{
    LogicalResult reifyReturnTypeShapes(
        OpBuilder& builder, ValueRange operands,
        SmallVectorImpl<Value>& reifiedReturnShapes) {
      return ::mlir::hlo::deriveShapeFromOperand(&builder, getOperation(),
                                                 operands.front(),
                                                 &reifiedReturnShapes);
    }
    // Relax the strict default implementation with one that allows
    // for StableHLO-specific differences.
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      if (l.size() != r.size()) return false;
      for (auto [lt, rt] : llvm::zip(l, r))
        if (!mlir::hlo::isCompatibleForHloTypeInference(lt, rt))
          return false;
      return true;
    }
  }];

  let results = (outs HLO_Tensor:$result);
}

class StableHLO_BinaryElementwiseOp<string mnemonic, list<Trait> traits> :
    StableHLO_BinaryElementwiseOpNoAssembly<mnemonic, traits> {
  let extraClassDefinition = [{
    ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) {
      return ::mlir::stablehlo::parseBinaryOp(parser, result);
    }
    void $cppClass::print(OpAsmPrinter &p) {
      ::mlir::stablehlo::printBinaryOp(getOperation(), p);
    }
  }];
  let hasCustomAssemblyFormat = 1;
}

def StableHLO_AddOp : StableHLO_BinaryElementwiseOp<"add",
      [Commutative, NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Addition operator";
  let description = [{
    Returns `lhs + rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_Atan2Op : StableHLO_BinaryElementwiseOp<"atan2",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Atan2 operator";
  let description = [{
    Returns `atan2(lhs/rhs)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_ComplexOp: StableHLO_BinaryElementwiseOpNoAssembly<"complex", [NoSideEffect,
    SameOperandsElementType, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Complex operator";
  let description = [{
    Performs element-wise conversion of a pair of real and imaginary values to
    a complex value.
  }];
  let arguments = (ins HLO_Fp32Or64Tensor:$lhs, HLO_Fp32Or64Tensor:$rhs);
  let results = (outs HLO_ComplexTensor:$result);

  // TODO(b/241767457): Remove parens when cleaning up BinaryOps.
  let assemblyFormat = "`(`operands`)` attr-dict `:` `(`type(operands)`)` `->` type($result)";
}

def StableHLO_DivOp : StableHLO_BinaryElementwiseOp<"divide",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Division operator";
  let description = [{
    Returns `lhs / rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_MaxOp : StableHLO_BinaryElementwiseOp<"maximum",
      [Commutative, NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Maximum operator";
  let description = [{
    Returns `max(lhs, rhs)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_MinOp : StableHLO_BinaryElementwiseOp<"minimum",
      [Commutative, NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Minimum operator";
  let description = [{
    Returns `min(lhs, rhs)` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_MulOp : StableHLO_BinaryElementwiseOp<"multiply",
      [Commutative, NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Multiplication operator";
  let description = [{
    Returns `lhs * rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_PowOp : StableHLO_BinaryElementwiseOp<"power",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Power operator";
  let description = [{
    Returns `lhs ^ rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}
def StableHLO_RemOp : StableHLO_BinaryElementwiseOp<"remainder",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Remainder operator";
  let description = [{
    Returns `lhs % rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_ShiftLeftOp : StableHLO_BinaryElementwiseOp<"shift_left",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Shift Left operator";
  let description = [{
    Returns `lhs << rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_ShiftRightArithmeticOp : StableHLO_BinaryElementwiseOp<"shift_right_arithmetic",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Shift right arithmetic operator";
  let description = [{
    Returns arithmetic `lhs >> rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_ShiftRightLogicalOp : StableHLO_BinaryElementwiseOp<"shift_right_logical",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Shift right logical operator";
  let description = [{
    Returns logical `lhs >> rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

def StableHLO_SubtractOp : StableHLO_BinaryElementwiseOp<"subtract",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Subtraction operator";
  let description = [{
    Returns `lhs - rhs` element-wise.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations.
  }];
}

//===----------------------------------------------------------------------===//
// StableHLO binary logical elementwise op definitions.
//===----------------------------------------------------------------------===//

// See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations
class StableHLO_BinaryBiwiseOrLogicalElementwiseOp<string mnemonic> :
        StableHLO_BinaryElementwiseOp<mnemonic,
          [Commutative, NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let arguments = (ins
    HLO_PredOrIntTensor:$lhs,
    HLO_PredOrIntTensor:$rhs
  );
}

def StableHLO_AndOp: StableHLO_BinaryBiwiseOrLogicalElementwiseOp<"and"> {
  let summary = "And operator";
  let description = [{
    Returns biwise-AND of `lhs` and `rhs` element-wise. The input tensors must
    be of type integer `HLO_Int` or boolean `HLO_Pred`.

    Note: For boolean tensor, the bitwise-AND is equivalent to logical-AND.
  }];
}

def StableHLO_OrOp: StableHLO_BinaryBiwiseOrLogicalElementwiseOp<"or"> {
  let summary = "Or operator";
  let description = [{
    Returns biwise-OR of `lhs` and `rhs` element-wise. The input tensors must
    be of type integer `HLO_Int` or boolean `HLO_Pred`.

    Note: For boolean tensor, the bitwise-OR is equivalent to logical-OR.
  }];
}

def StableHLO_XorOp : StableHLO_BinaryBiwiseOrLogicalElementwiseOp<"xor"> {
  let summary = "Xor operator";
  let description = [{
    Returns biwise-XOR of `lhs` and `rhs` element-wise. The input tensors must
    be of type integer `HLO_Int` or boolean `HLO_Pred`.

    Note: For boolean tensor, the bitwise-XOR is equivalent to logical-XOR.
  }];
}

//===----------------------------------------------------------------------===//
// StableHLO communication op definitions.
//===----------------------------------------------------------------------===//

// InfeedOp corresponds to 'InfeedWithToken' xla client API and not 'Infeed'.
// InfeedWithToken allows ordering of infeed HLO instructions using tokens.
def StableHLO_InfeedOp : StableHLO_Op<"infeed", []> {

  let summary = "Infeed operator";

  let description = [{
    Reads a single data item from the implicit Infeed streaming interface of
    the device, interpreting the data as the given shape, and returns a XlaOp
    of the data. Multiple Infeed operations are allowed in a computation, but
    there must be a total order among the Infeed operations.

    Attributes:
      layout:  Array attribute. Each element of the array is a minor_to_major
               array corresponding to the shape of the data read from the infeed
               interface.

    See https://www.tensorflow.org/xla/operation_semantics#infeed.
  }];

  let arguments = (ins
    HLO_Token:$token,
    DefaultValuedStrAttr<StrAttr, "">:$infeed_config,
    OptionalAttr<ArrayAttr>:$layout
  );
  let results = (outs Variadic<HLO_TensorOrToken>);
  let hasVerifier = 1;
}

// OutfeedOp corresponds to 'OutfeedWithToken' xla client API and not 'Outfeed'.
// OutfeedWithToken allows ordering of outfeed HLO instructions using tokens.
def StableHLO_OutfeedOp : StableHLO_Op<"outfeed", []> {

  let summary = "Outfeed operator";

  let description = [{
    Generates outgoing data transfers for the given data. It takes data and a
    token type operand and produces a token type value. Tokens are used for
    ordering side-effecting operations.

    See https://www.tensorflow.org/xla/operation_semantics#outfeed.
  }];

  let arguments = (ins
    Variadic<HLO_Tensor>:$operands,
    HLO_Token:$token,
    DefaultValuedStrAttr<StrAttr, "">:$outfeed_config
  );
  let results = (outs HLO_Token);
}

def StableHLO_SendOp : StableHLO_Op<"send", []> {

  let summary = "Send operator";

  let description = [{
    Sends the given operand data to a Recv instruction in another computation
    that shares the same channel handle. Does not return any data. Similar to
    the Recv operation, Send operation represents synchronous communication,
    and is internally decomposed into 2 HLO instructions (Send and SendDone) to
    enable asynchronous data transfers.

    See https://www.tensorflow.org/xla/operation_semantics#send.
  }];

  let arguments = (ins
    Variadic<HLO_Tensor>:$operands,
    HLO_Token:$token,
    StableHLO_ChannelHandle:$channel_handle,
    DefaultValuedAttr<BoolAttr, "false">:$is_host_transfer
  );

  let results = (outs HLO_Token);
}

def StableHLO_RecvOp : StableHLO_Op<"recv", []> {

  let summary = "Recv operator";

  let description = [{
    Receives data of the given shape from a Send instruction in another
    computation that shares the same channel handle. Returns a tuple containing
    value for the received data and a token. Recv operation represents
    synchronous communication. However, the instruction is internally decomposed
    into 2 HLO instructions (Recv and RecvDone) to enable asynchronous data
    transfers.

    See https://www.tensorflow.org/xla/operation_semantics#recv.
  }];

  let arguments = (ins
    HLO_Token:$token,
    StableHLO_ChannelHandle:$channel_handle,
    DefaultValuedAttr<BoolAttr, "false">:$is_host_transfer
  );

  let results = (outs Variadic<HLO_TensorOrToken>);
  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// StableHLO parallelism related op definitions.
//===----------------------------------------------------------------------===//

def StableHLO_ReplicaIdOp : StableHLO_Op<"replica_id", [NoSideEffect,
    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "ReplicaId operator";
  let description = [{
    Returns the unique ID (int32 scalar) of the replica.

    The unique ID of each replica is an unsigned integer in the interval [0, N),
    where N is the number of replicas. Since all the replicas are running the
    same program, a ReplicaId() call in the program will return a different
    value on each replica.

    See https://www.tensorflow.org/xla/operation_semantics#replicaid.

    Example:

    ```mlir
    %0 = stablehlo.replica_id : tensor<ui32>
    ```
  }];
  let results = (outs TensorOf<[UI32]>);

  let assemblyFormat = "attr-dict `:` type(results)";
}

//===----------------------------------------------------------------------===//
// StableHLO control flow op definitions.
//===----------------------------------------------------------------------===//

def StableHLO_AfterAllOp : StableHLO_Op<"after_all", [NoSideEffect]> {

  let summary = "AfterAll operator";

  let description = [{
    AfterAll takes a variadic number of tokens and produces a single token.
    Tokens are primitive types which can be threaded between side-effecting
    operations to enforce ordering. AfterAll can be used as a join of tokens
    for ordering a operation after a set operations.

    See https://www.tensorflow.org/xla/operation_semantics#afterall.

    Example:

    ```mlir
    %0 = stablehlo.after_all %arg0, %arg1 : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token
    ```
  }];

  let arguments = (ins Variadic<HLO_Token>:$operands);
  let results = (outs HLO_Token);

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

// Xla Client API has two separate calls for indexed and predicated conditional,
// although both eventually map to kConditional HLO. IfOp maps to predicated
// conditional use of kConditional HLO.
def StableHLO_IfOp: StableHLO_Op<"if", [
    RecursiveSideEffects,
    SingleBlockImplicitTerminator<"ReturnOp">]> {
  let summary = "If operator";

  let description = [{
    Executes the function `true_branch` if `pred` is true or `false_branch` if
    pred is false, and returns the result.

    The type of the returned values of `true_branch` and `false_branch`
    functions must be the same and equal to the types of the values returned by
    the operation.

    Note that only one of two functions will be executed depending on the value
    of `pred`.
  }];

  let arguments = (ins
    HLO_PredTensor:$pred
  );

  let regions = (region SizedRegion<1>:$true_branch,
                        SizedRegion<1>:$false_branch);

  let results = (outs Variadic<HLO_TensorOrToken>);

  let hasVerifier = 1;
}

// Xla Client API has two separate calls for indexed and predicated conditional,
// although both eventually map to kConditional HLO. CaseOp maps to indexed
// conditional use of kConditional HLO.
def StableHLO_CaseOp: StableHLO_Op<"case", [
      RecursiveSideEffects,
      SingleBlockImplicitTerminator<"ReturnOp">
    ]> {
  let summary = "Switch-Case operator";
  let description = [{
    Returns the result of executing `branches[index]`. If `index` is < 0 or >=
    N, then `branches[N-1]` is executed as the default branch.

    The type of the returned values of each branch must be the same and equal
    to the types of the values returned by the operation.

    Note that only one of the branches will be executed depending on the value
    of index.
  }];

  let arguments = (ins
    I32Tensor:$index
  );

  let regions = (region VariadicRegion<SizedRegion<1>>:$branches);

  let results = (outs Variadic<HLO_TensorOrToken>);

  let hasVerifier = 1;
}


def StableHLO_WhileOp: StableHLO_Op<"while", [
      RecursiveSideEffects,
      HLO_PairwiseSameOperandAndResultType,
      SingleBlockImplicitTerminator<"ReturnOp">,
      OpAsmOpInterface
    ]> {
  let summary = "While operator";
  let description = [{
    Returns the result of executing a body function until the cond body returns
    true.

    See https://www.tensorflow.org/xla/operation_semantics#while.
  }];
  let arguments = (ins Variadic<HLO_TensorOrToken>:$operand);

  let regions = (region SizedRegion<1>:$cond, SizedRegion<1>:$body);

  let results = (outs Variadic<HLO_TensorOrToken>);

  let extraClassDeclaration = [{
    // Method of OpAsmOpInterface used during custom printing to name the block
    // arguments in the nested regions. We name both the condition and the body
    // regions entry arguments the same way, with a `iterArg` prefix. Since the
    // two regions are side-by-side they will have the same name, which allows
    // us to print them once and share it for the two regions, and still be able
    // to parse them back.
    void getAsmBlockArgumentNames(Region &region, OpAsmSetValueNameFn setNameFn) {
      for (BlockArgument arg : region.getArguments())
        setNameFn(arg, "iterArg");
    }
  }];
  let hasCustomAssemblyFormat = 1;
  let hasVerifier = 1;
}

def StableHLO_AllGatherOp : StableHLO_Op<"all_gather", [SameOperandsAndResultElementType]> {

  string summary = "AllGather operator";

  string description = [{
    Performs concatenation across replicas.

    See https://www.tensorflow.org/xla/operation_semantics#allgather
  }];

  let arguments = (ins
    HLO_Tensor:$operand,
    I64Attr:$all_gather_dim,
    I64ElementsAttr:$replica_groups,
    OptionalAttr<StableHLO_ChannelHandle>:$channel_handle
  );
  let results = (outs HLO_Tensor);
  let hasVerifier = 1;
}

def StableHLO_AllReduceOp : StableHLO_Op<"all_reduce",
    [HLO_CompatibleOperandsAndResultType]> {
  let summary = "AllReduce operator";
  let description = [{
    Performs a custom reduction across replicas.

    See https://www.tensorflow.org/xla/operation_semantics#allreduce.
  }];

  let arguments = (ins
    HLO_Tensor:$operand,
    I64ElementsAttr:$replica_groups,
    OptionalAttr<StableHLO_ChannelHandle>:$channel_handle,
    UnitAttr:$use_global_device_ids
  );
  let regions = (region SizedRegion<1>:$computation);
  let results = (outs HLO_Tensor);
  // use_global_device_ids is rarely used, so we add a simplified
  // builder method for convenience.
  let builders = [
    OpBuilder<(ins
      "::mlir::Type":$result_type, "::mlir::Value":$operand,
      "::mlir::DenseIntElementsAttr":$replica_groups,
      "::mlir::stablehlo::ChannelHandleAttr":$channel_handle)>];
}

def StableHLO_ReduceScatterOp : StableHLO_Op<"reduce_scatter",
    [SameOperandsAndResultElementType]> {
  let summary = "ReduceScatter operator";
  let description = [{
     Performs all_reduce followed by a scatter.

     See https://www.tensorflow.org/xla/operation_semantics#reducescatter
  }];

  let arguments = (ins
    HLO_Tensor:$operand,
    I64Attr:$scatter_dimension,
    I64ElementsAttr:$replica_groups,
    OptionalAttr<StableHLO_ChannelHandle>:$channel_handle
  );
  let regions = (region SizedRegion<1>:$computation);
  let results = (outs HLO_Tensor);
  let hasVerifier = 1;
}

def StableHLO_AllToAllOp : StableHLO_Op<"all_to_all",
    [NoSideEffect, SameOperandsElementType, SameOperandsShape,
     InferTensorType]> {

  let arguments = (ins
    HLO_Tensor:$operand,
    I64Attr:$split_dimension,
    I64Attr:$concat_dimension,
    I64Attr:$split_count,
    I64ElementsAttr:$replica_groups
  );
  let results = (outs HLO_Tensor);
}

def StableHLO_ReduceOp: StableHLO_ShapedInterfaceOp<"reduce", [
      RecursiveSideEffects,
      SameVariadicOperandSize,
      SingleBlockImplicitTerminator<"ReturnOp">
    ]> {
  let summary = "Reduce operator";
  let description = [{
    Returns the result of executing a reduction function on one or more arrays
    in parallel.

    See https://www.tensorflow.org/xla/operation_semantics#reduce.
  }];
  let arguments = (ins
    Variadic<HLO_Tensor>:$operands,
    Variadic<HLO_Tensor>:$init_values,
    I64ElementsAttr:$dimensions
  );

  let results = (outs Variadic<HLO_Tensor>);

  let builders = [
    OpBuilder<(ins "ValueRange":$operands, "ValueRange":$init_values,
      "DenseIntElementsAttr":$dimensions)>];

  let hasCustomAssemblyFormat = 1;
  let hasVerifier = 1;

  // TODO(hinsu): Verify that the attached body arguments and results are
  // compatible with reduce op's operands.
  let regions = (region SizedRegion<1>:$body);
}

//===----------------------------------------------------------------------===//
// StableHLO tuple op definitions.
//===----------------------------------------------------------------------===//
def StableHLO_GetTupleElementOp: StableHLO_Op<"get_tuple_element", [NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "GetTupleElement operator";
  let description = [{
    Returns a member of a tuple specified by an index.

    See https://www.tensorflow.org/xla/operation_semantics#gettupleelement.
  }];
  let arguments = (ins
    HLO_Tuple,
    I32Attr:$index
  );

  let results = (outs HLO_TensorOrTokenOrTuple);

  let hasVerifier = 1;
}

def StableHLO_TupleOp : StableHLO_Op<"tuple", [NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "XLA's tuple op";
  let description = [{
     Groups a set of tensor inputs into a single tuple object.

     See https://www.tensorflow.org/xla/operation_semantics#tuple.
   }];
  let arguments = (ins Variadic<HLO_TensorOrTokenOrTuple>:$val);
  let results = (outs HLO_Tuple);

  let hasVerifier = 1;
}

def StableHLO_CompareOp: StableHLO_Op<"compare", [NoSideEffect, SameOperandsElementType,
    SameOperandsAndResultShape, Elementwise, InferTensorTypeWithReify]> {
  let summary = "Comparison operator";
  let description = [{
    Compares `lhs` and `rhs` elementwise according to `comparison_direction`
    and `compare_type`. If unspecified, `compare_type` is FLOAT for float element
    types, SIGNED for signed element types and UNSIGNED for unsigned element
    types.

    See
    https://www.tensorflow.org/xla/operation_semantics#element-wise_comparison_operations.

    Example:

    ```mlir
    %0 = stablehlo.compare LT, %arg0, %arg1 : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
    %1 = stablehlo.compare LT, %arg0, %arg1, TOTALORDER : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
    ```
  }];
  let arguments = (ins
    HLO_Tensor:$lhs,
    HLO_Tensor:$rhs,
    StableHLO_ComparisonDirectionAttr:$comparison_direction,
    OptionalAttr<StableHLO_ComparisonTypeAttr>:$compare_type
  );
  let results = (outs HLO_PredTensor);

  let builders = [
    OpBuilder<(ins "Value":$lhs, "Value":$rhs,
      "::mlir::stablehlo::ComparisonDirection":$comparison_direction,
      CArg<"::mlir::stablehlo::ComparisonType",
      "::mlir::stablehlo::ComparisonType::NOTYPE">:$compare_type)>,
  ];

  let extraClassDeclaration = [{
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(mlir::verifyCompatibleShapes(l, r));
    }
  }];

  let assemblyFormat = [{
    $comparison_direction `,` $lhs `,` $rhs (`,` $compare_type^)?
      attr-dict `:` functional-type(operands, results)
  }];
}

//===----------------------------------------------------------------------===//
// StableHLO Slice definitions.
//===----------------------------------------------------------------------===//

def StableHLO_SliceOp: StableHLO_Op<
      "slice",
      [NoSideEffect, SameOperandsAndResultElementType,
       AllTypesMatch<["start_indices", "limit_indices", "strides"]>,
       DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let arguments = (ins
    HLO_Tensor:$operand,
    I64ElementsAttr:$start_indices,
    I64ElementsAttr:$limit_indices,
    I64ElementsAttr:$strides
  );

  let results = (outs HLO_Tensor);
}

def StableHLO_DynamicSliceOp: StableHLO_Op<"dynamic_slice",
      [NoSideEffect, AllElementTypesMatch<["operand", "result"]>,
       InferTensorType]> {
  let summary = "Dynamic Slice operator";
  let description = [{
    Extracts a sub-array from the input array at dynamic start_indices.

    See https://www.tensorflow.org/xla/operation_semantics#dynamicslice.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    Variadic<HLO_ScalarIntTensor>:$start_indices,
    I64ElementsAttr:$slice_sizes
  );

  let results = (outs HLO_Tensor:$result);
  let hasVerifier = 1;
}

def StableHLO_DynamicUpdateSliceOp: StableHLO_Op<"dynamic_update_slice",
      [NoSideEffect, AllElementTypesMatch<["operand", "update", "result"]>,
       AllShapesMatch<["operand", "result"]>]> {
  let summary = "Dynamic Update Slice operator";
  let description = [{
    DynamicUpdateSlice generates a result which is the value of the input array
    operand, with a slice update overwritten at start_indices.

    See https://www.tensorflow.org/xla/operation_semantics#dynamicupdateslice.

    Example:

    ```mlir
    %0 = stablehlo.dynamic_update_slice %arg0, %arg1, %arg2
           : (tensor<4xf32>, tensor<2xf32>, tensor<i32>) -> tensor<4xf32>
    ```
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_Tensor:$update,
    Variadic<HLO_ScalarIntTensor>:$start_indices
  );
  let results = (outs HLO_Tensor:$result);
  let hasVerifier = 1;

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}


//===----------------------------------------------------------------------===//
// StableHLO Other op definitions.
//===----------------------------------------------------------------------===//

def StableHLO_BatchNormGradOp : StableHLO_Op<"batch_norm_grad", [NoSideEffect,
    AllShapesMatch<["scale", "mean", "variance", "grad_scale",
        "grad_offset"]>,
    AllShapesMatch<["operand", "grad_output"]>,
    AllElementTypesMatch<["operand", "grad_scale", "grad_offset"]>,
    AllTypesMatch<["operand", "grad_operand"]>]> {
  let summary = "Batch Normalization Gradient";
  let description = [{
    Calculates gradients of batch norm.

    See https://www.tensorflow.org/xla/operation_semantics#batchnormgrad
  }];

  let arguments = (ins
    RankedTensorOf<[HLO_Float]>:$operand,
    1DTensorOf<[HLO_Float]>:$scale,
    1DTensorOf<[HLO_Float]>:$mean,
    1DTensorOf<[HLO_Float]>:$variance,
    RankedTensorOf<[HLO_Float]>:$grad_output,
    F32Attr:$epsilon,
    I64Attr:$feature_index
  );

  let results = (outs Variadic<HLO_TensorOrToken>);
  let results = (outs
      RankedTensorOf<[HLO_Float]>:$grad_operand,
      1DTensorOf<[HLO_Float]>:$grad_scale,
      1DTensorOf<[HLO_Float]>:$grad_offset);

  let hasVerifier = 1;
}

def StableHLO_BatchNormInferenceOp : StableHLO_Op<"batch_norm_inference",
    [NoSideEffect, AllTypesMatch<["operand", "result"]>,
    AllShapesMatch<["scale", "offset", "mean", "variance"]>]> {
  let summary = "Batch Normalization for Inference";
  let description = [{
    Normalizes an array across batch and spatial dimensions.

    See https://www.tensorflow.org/xla/operation_semantics#batchnorminference
  }];

  let arguments = (ins
    RankedTensorOf<[HLO_Float]>:$operand,
    1DTensorOf<[HLO_Float]>:$scale,
    1DTensorOf<[HLO_Float]>:$offset,
    1DTensorOf<[HLO_Float]>:$mean,
    1DTensorOf<[HLO_Float]>:$variance,
    F32Attr:$epsilon,
    I64Attr:$feature_index
  );

  let results = (outs RankedTensorOf<[HLO_Float]>:$result);

  let hasVerifier = 1;
}

def StableHLO_BatchNormTrainingOp : StableHLO_Op<"batch_norm_training",
    [NoSideEffect, AllTypesMatch<["operand", "output"]>,
    AllElementTypesMatch<["operand", "batch_mean", "batch_var"]>,
    AllShapesMatch<["scale", "offset", "batch_mean", "batch_var"]>]> {
  let summary = "Batch Normalization for Training";
  let description = [{
    Normalizes an array across batch and spatial dimensions.

    See https://www.tensorflow.org/xla/operation_semantics#batchnormtraining
  }];

  let arguments = (ins
    RankedTensorOf<[HLO_Float]>:$operand,
    1DTensorOf<[HLO_Float]>:$scale,
    1DTensorOf<[HLO_Float]>:$offset,
    F32Attr:$epsilon,
    I64Attr:$feature_index
  );

  let results = (outs
      RankedTensorOf<[HLO_Float]>:$output,
      1DTensorOf<[HLO_Float]>:$batch_mean,
      1DTensorOf<[HLO_Float]>:$batch_var);

  let hasVerifier = 1;
}

def StableHLO_BitcastConvertOp : StableHLO_ShapedInterfaceOp<"bitcast_convert",
    [NoSideEffect]> {
  let summary = "BitcastConvert operator";
  let description = [{
    Similar to a 'tf.bitcast' in TensorFlow, performs an element-wise bitcast
    operation from a data shape to a target shape. The dimensions must match,
    and the conversion is an element-wise one. Bitcast is implemented as a
    low-level cast, so machines with different floating-point representations
    will give different results.

    See https://www.tensorflow.org/xla/operation_semantics#bitcastconverttype.

    Example:

    ```mlir
    %0 = stablehlo.bitcast_convert %arg0 : (tensor<2xi32>) -> tensor<2xf32>
    ```
  }];

  let arguments = (ins HLO_Tensor:$operand);
  let results = (outs HLO_Tensor);
  let hasVerifier = 1;

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

def StableHLO_BroadcastOp : StableHLO_ShapedInterfaceOp<"broadcast",
    [NoSideEffect, SameOperandsAndResultElementType, InferTensorType]> {
  let summary = "Broadcast a tensor to a higher rank by prepending dimensions";
  let description = [{
    Broadcasts the operand tensor to a higher rank by prepending
    `broadcast_sizes` to the dimensions. The current values of the operand are
    copied into the other dimensions.

    This is a more limited form of broadcasting, that corresponds to the XLA
    client Broadcast method. For a more general form of broadcasting, see the
    BroadcastInDimOp.

    See https://www.tensorflow.org/xla/operation_semantics#broadcast.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    I64ElementsAttr:$broadcast_sizes
  );

  let results = (outs HLO_Tensor);

  let hasVerifier = 1;
}

def StableHLO_BroadcastInDimOp : StableHLO_Op<"broadcast_in_dim",
      [NoSideEffect, SameOperandsAndResultElementType]> {
  let summary = "Broadcast a tensor into the given shape by adding dimensions.";
  let description = [{
    Broadcasts the `operand` tensor to a higher rank. This is not the limited
    form of broadcasting exposed as the XLA client broadcast op, but rather the
    more powerful "InDim" broadcasting, which is closer to the HLO broadcast op
    and exposed in the XLA client BroadcastInDim method.

    `broadcast_dimensions` maps the operand dimension number to the target shape
    dimension number. It must have the same size as the rank of the operand. The
    mapped dimensions must either be the same size or the dimension being
    broadcast from must be size 1 (degenerate broadcasting).

    For a scalar (0D tensor) operand, `broadcast_dimensions` must be empty. The
    The scalar value will be broadcast to every element in the target shape.

    See https://www.tensorflow.org/xla/broadcasting.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    BroadcastDimAttr:$broadcast_dimensions
  );

  let results = (outs HLO_StaticShapeTensor);

  let hasVerifier = 1;
}

def StableHLO_DynamicBroadcastInDimOp : StableHLO_ShapedInterfaceOp<
    "dynamic_broadcast_in_dim", [NoSideEffect]> {
  let summary = "Broadcast a tensor into the given dynamic shape by adding dimensions.";
  let description = [{
    This is a generalization of the BroadcastInDimOp which accepts its output
    dimensions as an argument. It should eventually supercede the statically
    shaped original, but is being phased as a separate op in order to support
    compatibility with lowerings and translations that precede dynamic shapes.

    The op accepts optional attributes to express static knowledge about the
    expanding behavior of dimensions. If not specified, all dimensions are
    assumed to be possibly expanding. The sets of dimensions that are known to
    be expanding and the set of dimensions that are known to be non-expanding
    must be disjoint and they must be a subset of the operand's dimensions.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_DimensionTensor:$output_dimensions,
    BroadcastDimAttr:$broadcast_dimensions,
    OptionalAttr<BroadcastDimAttr>:$known_expanding_dimensions,
    OptionalAttr<BroadcastDimAttr>:$known_nonexpanding_dimensions
  );

  let results = (outs HLO_Tensor);

  let builders = [
    OpBuilder<(ins
        "Type":$result_type, "Value":$operand, "Value":$output_dimensions,
        "DenseIntElementsAttr":$broadcast_dimensions), [{
      build($_builder, $_state, result_type, operand, output_dimensions,
          broadcast_dimensions, /*known_expanding_dimensions=*/{},
          /*known_nonexpanding_dimensions=*/{});
    }]>
  ];

  let hasVerifier = 1;
}

// Note: There is no HLO_CallOp because the standard call operation mlir::func::CallOp
// is used instead. A mlir::func::CallOp is exported to a HLO call instruction
// directly.

def StableHLO_CholeskyOp : StableHLO_Op<"cholesky",
      [NoSideEffect, SameOperandsAndResultElementType, InferTensorType]> {
  let summary = "Cholesky operator";
  let description = [{
  Computes the Cholesky decomposition of a batch of symmetric (Hermitian)
  positive definite matrices.

  If lower is true, computes lower-triangular matrices l such that
  `a=l.Transpose(l)`. If lower is false, computes upper-triangular matrices u such
  that `a=Transpose(u).u`.

  Input data is read only from the lower/upper triangle of a, depending on the
  value of lower. Values from the other triangle are ignored. Output data is
  returned in the same triangle; the values in the other triangle are
  implementation-defined and may be anything.

  If the rank of a is greater than 2, a is treated as a batch of matrices, where
  all except the minor 2 dimensions are batch dimensions.

  If a is not symmetric (Hermitian) positive definite, the result is
  implementation-defined.

    See https://www.tensorflow.org/xla/operation_semantics#cholesky.
  }];
  let arguments = (ins
    HLO_FpOrComplexTensor:$a,
    DefaultValuedAttr<BoolAttr, "false">:$lower
  );

  let results = (outs HLO_FpOrComplexTensor);
}

def StableHLO_ClampOp : StableHLO_ShapedInterfaceOp<"clamp", [NoSideEffect,
  SameOperandsAndResultElementType, HLO_BroadcastingElementwise,
  InferTensorType]> {
  let summary = "Clamp operator";
  let description = [{
    Clamps an operand to within the range between a minimum and maximum value.

    Note: All three arrays must be the same shape. Alternatively, as a
          restricted form of broadcasting, min and/or max can be a scalar (0D
          tensor) of the element type of the tensor operand.

    See https://www.tensorflow.org/xla/operation_semantics#clamp.

    Example:

    ```mlir
    %0 = stablehlo.clamp %arg0, %arg1, %arg2 : (tensor<f32>, tensor<4xf32>, tensor<f32>) -> tensor<4xf32>
    ```
  }];

  let arguments = (ins
    HLO_Tensor:$min,
    HLO_Tensor:$operand,
    HLO_Tensor:$max
  );
  let results = (outs HLO_Tensor);

  let hasVerifier = 1;

  let extraClassDeclaration = [{
    // Method from InferTypeOpInterface interface.
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      if (l.size() != r.size()) return false;
      for (auto [lt, rt] : llvm::zip(l, r))
        if (!mlir::hlo::isCompatibleForHloTypeInference(lt, rt))
          return false;
      return true;
    }
  }];

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

def StableHLO_ConcatenateOp : StableHLO_ShapedInterfaceOp<"concatenate",
    [NoSideEffect, SameOperandsAndResultElementType,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "XLA's concatenate op";
  let description = [{
     Concatenates a set of tensors along the specified dimension.

     See https://www.tensorflow.org/xla/operation_semantics#concatenate.
   }];

  let arguments = (ins
    Variadic<HLO_Tensor>:$val,
    I64Attr: $dimension
  );

  let results = (outs HLO_Tensor);

  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(mlir::verifyCompatibleShapes(l, r));
    }
  }];
}

def StableHLO_CollectivePermuteOp: StableHLO_Op<"collective_permute",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "CollectivePermute operator";
  let description = [{
    CollectivePermute is a collective operation that sends and receives data
    cross replicas.
    Note that there are the following restrictions on the source_target_pair:
    - Any two pairs should not have the same target replica id, and they should
    not have the same source replica id.
    - If a replica id is not a target in any pair, then the output on that
    replica is a tensor consists of 0(s) with the same shape as the input.

    See https://www.tensorflow.org/xla/operation_semantics#collectivepermute.

  }];

  let arguments = (ins
    HLO_Tensor:$operand,
    I64ElementsAttr:$source_target_pairs
  );
  let results = (outs HLO_Tensor);
  let hasVerifier = 1;
}

def StableHLO_ConvolutionOp : StableHLO_Op<"convolution", [NoSideEffect]> {
  let summary = "Convolution operator";
  let description = [{
    Computes a convolution of the kind used in neural networks.

    See https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
  }];
  let arguments = !con(
    (ins
       HLO_Tensor:$lhs,
       HLO_Tensor:$rhs),
    StableHLO_ConvolutionAttributes.attributes);

  let results = (outs HLO_Tensor);
  let hasVerifier = 1;

  code extraClassDeclaration = [{
    bool hasWindowReversal() {
      auto reversal = window_reversalAttr();
      return reversal && llvm::any_of(reversal.getValues<bool>(),
                                      [](bool v) { return v; });
    }
  }];

 let assemblyFormat = [{
    `(`operands`)`
       `dim_numbers` `=` custom<ConvolutionDimensions>($dimension_numbers) `,`
       `window` `=` `{` custom<WindowAttributes>($window_strides, $padding,
                                                 $lhs_dilation, $rhs_dilation,
                                                 $window_reversal) `}`
       attr-dict `:` functional-type(operands, results)
  }];
}

def StableHLO_CrossReplicaSumOp : StableHLO_Op<"cross-replica-sum",
    [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Sums input across replicated instances.";
  let description = [{
     For each of the replica groups, operands of the group devices are summed
     so that each device has the sum.

     For example, suppose there are 8 TPU devices: `[A, B, C, D, E, F, G, H]`.
     Passing group_assignment=`[[0,2,4,6],[1,3,5,7]]` sets `A, C, E, G` as group 0,
     and `B, D, F, H` as group 1. Thus we get the outputs:
     `[A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H]`.

     See https://www.tensorflow.org/xla/operation_semantics#crossreplicasum.
   }];

  let arguments = (ins
    HLO_Tensor:$operand,
    I64ElementsAttr:$replica_groups
  );

  let results = (outs HLO_Tensor);
}

def StableHLO_CustomCallOp: StableHLO_Op<"custom_call",
    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
  let summary = "CustomCall operator";
  let description = [{
    A custom call invokes code external to XLA. The `args` are passed to the
    external code, and the external code is expected to produce a result of the
    given type. The exact mechanism is backend-specific. For example, in the CPU
    backend, a call instruction is emitted which targets a symbol with the name
    `call_target_name`.

    `call_target_name` and `backend_config` can be arbitrary strings, but
    `call_target_name` should be short as it may be used in labels.
    `backend_config` can encode arbitrarily large amounts of information.

    `has_side_effect` must be true if the custom call has side-effects.
    `api_version` specifies the version of the API used by the custom call
    function.

    A custom call may apply functions within the scope of the parent module.
    They can be referenced using `called_computations` attribute.

    A custom call can also have layout constraints on operands and results which
    can be specified as optional `operand_layouts` and `result_layouts`
    attributes. The layout attribute is an array of rank-1 index tensors and the
    i-th layout attribute specifies the layout for i-th operand/result.

    The `operand_layouts` & `result_layouts` attributes can be specified under
    the following constraints:
    1) Either both `operand_layouts` and `result_layouts` are specified or none.
    2) None of the operands are of tuple type.
    3) None of the results are of tuple type except the common case of single
       tuple result packing non-tuple values is allowed. In this case the i-th
       `result_layouts` attribute specifies the layout of i-th element in the
       result tuple.

    See https://www.tensorflow.org/xla/operation_semantics#customcall.
  }];
  let arguments = (ins
    Variadic<HLO_TensorOrTokenOrTuple>:$operands,
    StrAttr:$call_target_name,
    DefaultValuedAttr<BoolAttr, "false">:$has_side_effect,
    DefaultValuedStrAttr<StrAttr, "">:$backend_config,
    // TODO(b/189822916): Remove this field when all clients are migrated to
    // the status-returning API.
    DefaultValuedAttr<
        StableHLO_CustomCallApiVersionAttr,
        "::mlir::stablehlo::CustomCallApiVersion::API_VERSION_ORIGINAL">:
        $api_version,
    DefaultValuedAttr<StableHLO_FlatSymbolRefArrayAttr, "{}">:$called_computations,
    OptionalAttr<StableHLO_ArrayOfLayoutAttr>:$operand_layouts,
    OptionalAttr<StableHLO_ArrayOfLayoutAttr>:$result_layouts
  );
  let results = (outs Variadic<HLO_TensorOrTokenOrTuple>);
  let hasVerifier = 1;
}

def StableHLO_DotOp: StableHLO_Op<"dot",
    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Dot operator";
  let description = [{
    Performs dot products between vectors, vector/matrix and matrix/matrix
    multiplication.

    See https://www.tensorflow.org/xla/operation_semantics#dot.
  }];
  let arguments = (
    ins HLO_Tensor:$lhs,
    HLO_Tensor:$rhs,
    StableHLO_PrecisionConfigAttr:$precision_config
  );
  let results = (outs HLO_Tensor);
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(mlir::verifyCompatibleShapes(l, r));
    }
  }];
}

def StableHLO_DotGeneralOp: StableHLO_ShapedInterfaceOp<"dot_general", [NoSideEffect]> {
  let summary = "General Dot operator";
  let description = [{
    Performs general dot products between vectors, vector/matrix and
    matrix/matrix multiplication.

    See https://www.tensorflow.org/xla/operation_semantics#dotgeneral.
  }];
  let arguments = (ins
    HLO_Tensor:$lhs,
    HLO_Tensor:$rhs,
    StableHLO_DotDimensionNumbers:$dot_dimension_numbers,
    StableHLO_PrecisionConfigAttr:$precision_config
  );

  let results = (outs HLO_Tensor);
  let hasVerifier = 1;
}

// Define Base Einsum op within the HLO dialect as these are client ops and
// therefore this class is not common between HLO and LHLO ops.
class BASE_EinsumOp {
  string summary = "Einsum operator";

  string description = [{
    Returns a tensor whose elements are defined by equation, which is written
    in a shorthand form inspired by the Einstein summation convention.
  }];
}

def StableHLO_EinsumOp: StableHLO_Op<"einsum", [NoSideEffect]>, BASE_EinsumOp {
  let arguments = (ins
    HLO_Tensor:$lhs,
    HLO_Tensor:$rhs,
    StrAttr:$einsum_config
  );

  let results = (outs HLO_Tensor);

  // TODO(hinsu): Canonicalize to lower this client side HLO op to server
  // side HLO ops.
}

def StableHLO_UnaryEinsumOp: StableHLO_Op<"unary_einsum", [NoSideEffect]>, BASE_EinsumOp {
  let arguments = (ins
    HLO_Tensor:$operand,
    StrAttr:$einsum_config
  );

  let results = (outs HLO_Tensor);
}

def StableHLO_FftOp: StableHLO_Op<"fft", [InferTensorType, NoSideEffect]> {
  let summary = "Fast fourier transform operator";
  let description = [{
    Returns the fast-fourier-transform of the input array.

    See
    https://www.tensorflow.org/xla/operation_semantics#fft.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    StableHLO_FftTypeAttr: $fft_type,
    I64ElementsAttr:$fft_length
  );

  let results = (outs HLO_Tensor);
}

def StableHLO_GatherOp: StableHLO_Op<"gather", [InferTensorTypeWithReify, NoSideEffect]> {
  let summary = "Gather operator";
  let description = [{
    Stitches together several slices of `operand` from offsets specified in
    `start_indices` (each slice at a potentially different runtime offset).

    See https://www.tensorflow.org/xla/operation_semantics#gather.
  }];

  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_IntTensor:$start_indices,
    StableHLO_GatherDimensionNumbers:$dimension_numbers,
    I64ElementsAttr:$slice_sizes,
    DefaultValuedAttr<BoolAttr, "false">:$indices_are_sorted
  );

  let results = (outs HLO_Tensor);

  let extraClassDeclaration = [{
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(mlir::verifyCompatibleShapes(l, r));
    }
  }];
}

def StableHLO_GetDimensionSizeOp: StableHLO_Op<"get_dimension_size", [NoSideEffect]> {
  let summary = "GetDimensionSize operator";
  let description = [{
    Returns the size of the given dimension of the operand.

    See
    https://www.tensorflow.org/xla/operation_semantics#getdimensionsize.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    I64Attr:$dimension
  );
  // TODO(hinsu): Allow 64-bit result types once XLA HLO dialect based on the
  // XLA semantics is available. This limitation is because of the current XLA
  // implementation.
  let results = (outs I32Tensor);

  let hasVerifier = 1;
}

def StableHLO_MapOp: StableHLO_ShapedInterfaceOp<"map",
      [RecursiveSideEffects, SameOperandsAndResultShape,
       SingleBlockImplicitTerminator<"ReturnOp">]> {
  let summary = "Map operator";
  let description = [{
  Applies a scalar function over the given operands arrays, producing an array
  of the same dimensions where each element is the result of the mapped function
  applied to the corresponding elements in the input arrays.

  The mapped function is an arbitrary computation with the restriction that it
  has N inputs of scalar type T and a single output with type S. The output has
  the same dimensions as the operands except that the element type T is replaced
  with S.

  See https://www.tensorflow.org/xla/operation_semantics#map.
  }];
  let arguments = (ins
    Variadic<HLO_Tensor>:$operands,
    I64ElementsAttr:$dimensions
  );
  let regions = (region SizedRegion<1>:$computation);
  let results = (outs HLO_Tensor);
  let hasVerifier = 1;
}

def StableHLO_ReshapeOp: StableHLO_Op<"reshape",
      [NoSideEffect, SameOperandsAndResultElementType]> {
  let summary = "Reshape operator";
  let description = [{
    Reshapes the dimensions of `operand` into a new configuration.

    See https://www.tensorflow.org/xla/operation_semantics#reshape.

    Example:

    ```mlir
    %0 = stablehlo.reshape %arg0 : (tensor<2xf32>) -> tensor<1x2xf32>
    ```
  }];

  let arguments = (ins HLO_Tensor:$operand);

  let results = (outs HLO_StaticShapeTensor);
  let hasVerifier = 1;

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

def StableHLO_DynamicReshapeOp: StableHLO_ShapedInterfaceOp<"dynamic_reshape", [NoSideEffect]> {
  let summary = "Reshape a tensor to a given, possibly dynamic, shape.";
  let description = [{
    Reshapes `operand` to `output_shape`.

    Requires:
    - The length of `output_shape` is equal to the rank of `result`.
    - The number of elements in `operand` (that is, the product of extents of
      its shape) is equal to the number of elements in `output_shape` (that is,
      the product of values in `output_shape`).

    Example:

    ```mlir
    %0 = stablehlo.dynamic_reshape %arg0, %shape : (tensor<?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
    ```
  }];

  let arguments = (ins HLO_Tensor:$operand, HLO_DimensionTensor:$output_shape);
  let results = (outs HLO_Tensor:$result);

  let hasVerifier = 1;

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

def StableHLO_ScatterOp: StableHLO_Op<"scatter", [SameVariadicOperandSize, RecursiveSideEffects]> {
  let summary = "Scatter operator";
  let description = [{
    Generates a result which is the value of the input array `operand`,
    with several slices (at indices specified by `scatter_indices`)
    updated with the values in `updates` using `update_computation`.

    See https://www.tensorflow.org/xla/operation_semantics#scatter.
  }];
  let arguments = (ins
    Variadic<HLO_Tensor>:$operands,
    TensorOf<[AnyInteger, Index]>:$scatter_indices,
    Variadic<HLO_Tensor>:$updates,
    StableHLO_ScatterDimensionNumbers:$scatter_dimension_numbers,
    DefaultValuedAttr<BoolAttr, "false">:$indices_are_sorted,
    DefaultValuedAttr<BoolAttr, "false">:$unique_indices
  );

  let regions = (region SizedRegion<1>:$update_computation);

  let results = (outs Variadic<HLO_Tensor>);

  let hasVerifier = 1;
}

def StableHLO_SelectOp: StableHLO_Op<"select", [NoSideEffect, HLO_BroadcastingElementwise,
    InferTensorTypeWithReify]> {
  let summary = "Select operator";
  let description = [{
    Constructs an output tensor from the elements of `on_true` and `on_false`
    based on the values of `pred`. All three operands must be of the same shape
    with the exception of `pred`, which may also be a scalar in which case it is
    broadcasted.

    See https://www.tensorflow.org/xla/operation_semantics#select.
  }];
  let arguments = (ins
    HLO_PredTensor:$pred,
    HLO_Tensor:$on_true,
    HLO_Tensor:$on_false
  );

  let results = (outs HLO_Tensor);

  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(mlir::verifyCompatibleShapes(l, r));
    }
  }];
}

def StableHLO_SelectAndScatterOp: StableHLO_Op<"select_and_scatter",
      [RecursiveSideEffects]> {
  let summary = "SelectAndScatter operator";
  let description = [{
    Runs a windowed selection `select` function over `operand` with shape
    `window_dimensions` and stride `window_strides`. This will produce an amount
    of selected locations whose shape matches `source`. These are then scattered
    to the output which is initialized with `init_value`.
    Multiple scattered elements which land in the same output location are
    combined using the `scatter` function.

    See https://www.tensorflow.org/xla/operation_semantics#selectandscatter.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_Tensor:$source,
    HLO_Tensor:$init_value,
    OptionalAttr<I64ElementsAttr>:$window_dimensions,
    OptionalAttr<I64ElementsAttr>:$window_strides,
    OptionalAttr<I64ElementsAttr>:$padding
  );

  let regions = (region SizedRegion<1>:$select, SizedRegion<1>:$scatter);

  let results = (outs HLO_Tensor);

  let hasVerifier = 1;
}

def StableHLO_SetDimensionSizeOp: StableHLO_Op<"set_dimension_size", [NoSideEffect,
    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "SetDimensionSize operator";
  let description = [{
    Sets the dynamic size of operand's given dimension. Pass through the operand
    as result, with dynamic dimension tracked by the compiler. Padded values
    will be ignored by downstream reduction ops.

    See https://www.tensorflow.org/xla/operation_semantics#setdimensionsize.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    I32Tensor:$size,
    I64Attr:$dimension
  );
  let results = (outs HLO_Tensor);

  let extraClassDeclaration = [{
    // Method from InferTypeOpInterface interface.
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      if (l.size() != r.size()) return false;
      for (auto [lt, rt] : llvm::zip(l, r))
        if (!mlir::hlo::isCompatibleForHloTypeInference(lt, rt))
          return false;
      return true;
    }
  }];

  let hasVerifier = 1;
}

def StableHLO_SortOp : StableHLO_Op<"sort", [RecursiveSideEffects,
                                 SameOperandsAndResultShape]> {
  let summary = "Sort operator";
  let description = [{
    Sorts the given `operands` at the given `dimension` with the given
    `comparator`.

    See https://www.tensorflow.org/xla/operation_semantics#sort.
  }];
  let arguments = (ins
    Variadic<HLO_Tensor>:$operands,
    DefaultValuedAttr<I64Attr, "-1">:$dimension,
    DefaultValuedAttr<BoolAttr, "false">:$is_stable
  );

  let results = (outs Variadic<HLO_Tensor>);

  let regions = (region SizedRegion<1>:$comparator);

  let builders = [
    OpBuilder<(ins "ValueRange":$operands, CArg<"int64_t", "-1">:$dimension,
      CArg<"bool", "false">:$is_stable)>];

  let hasVerifier = 1;
}

def StableHLO_ReverseOp: StableHLO_Op<"reverse",
      [NoSideEffect, HLO_CompatibleOperandsAndResultType]> {
  let summary = "Reverse operator";
  let description = [{
    Reverses the specified dimensions of `operand` according to the given
    `dimensions`.

    See https://www.tensorflow.org/xla/operation_semantics#rev_reverse.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    I64ElementsAttr:$dimensions
  );

  let results = (outs HLO_Tensor);
}

def StableHLO_PadOp: StableHLO_ShapedInterfaceOp<"pad",
      [NoSideEffect, SameOperandsAndResultElementType, InferTensorType]> {
  let summary = "Pad operator";
  let description = [{
    Pads edges and between the elements of `operand` with the `padding_value`
    according to the configuration parameters described below.

    `edge_padding_low` and `edge_padding_high` specify the amount of padding
    added at the low-end (next to index 0) and the high-end (next to the
    highest index) of each dimension respectively. The amount of edge
    padding can be negative -- the absolute value of negative padding indicates
    the number of elements to remove from the specified dimension.

    `interior_padding` specifies the amount of padding (non-negative) added
    between any two elements in each dimension. Interior padding occurs
    logically before edge padding, so in the case of negative edge padding,
    elements are removed from the interior-padded operand.

    This operation is a no-op if, for all dimensions, the edge padding pairs are
    all (0, 0) and the interior padding values are all 0. The figure below shows
    examples of different `edge_padding` and `interior_padding` values for a
    two-dimensional array.

    ![Examples](https://www.tensorflow.org/xla/images/ops_pad.png)

  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_Tensor:$padding_value,
    I64ElementsAttr: $edge_padding_low,
    I64ElementsAttr: $edge_padding_high,
    I64ElementsAttr: $interior_padding
  );

  let results = (outs HLO_Tensor);
}

def StableHLO_TraceOp: StableHLO_Op<"trace", []> {
  let summary = "Trace operator";
  let description = [{
    Emits a logging message `tag` with the `operand`.

    Example:

    ```mlir
    stablehlo.trace %arg0, "In test code." : tensor<5x1x5xi32>
    ```
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    StrAttr:$tag
  );
  let assemblyFormat = "$operand `,` $tag attr-dict `:` type($operand)";
}

def StableHLO_TransposeOp: StableHLO_ShapedInterfaceOp<"transpose",
      [NoSideEffect, SameOperandsAndResultElementType,
      DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Transpose operator";
  let description = [{
    Permutes the dimensions of `operand` according to the given `permutation`.

    `res_dimensions[i] = operand_dimensions[permutation[i]]`

    See https://www.tensorflow.org/xla/operation_semantics#transpose.
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    I64ElementsAttr:$permutation
  );
  let results = (outs HLO_Tensor);

  let extraClassDeclaration = [{
    // Method from  InferTypeOpInterface interface.
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(mlir::verifyCompatibleShapes(l, r));
    }
  }];
}

def StableHLO_TriangularSolveOp: StableHLO_Op<"triangular_solve",
    [NoSideEffect, SameOperandsAndResultElementType]> {
  let summary = "TriangularSolve operator";
  let description = [{
    Solves systems of linear equations with lower or upper triangular
    coefficient matrices by forward- or back-substitution. Broadcasting along
    leading dimensions, this routine solves one of the matrix systems
    op(a) * x = b, or x * op(a) = b, for the variable x, given a and b, where
    op(a) is either op(a) = a, or op(a) = Transpose(a), or
    op(a) = Conj(Transpose(a)).

    Input data is read only from the lower/upper triangle of a, depending on the
    value of lower. Values from the other triangle are ignored. Output data is
    returned in the same triangle; the values in the other triangle are
    implementation-defined and may be anything.

    If the rank of a and b are greater than 2, they are treated as batches of
    matrices, where all except the minor 2 dimensions are batch dimensions. a
    and b must have equal batch dimensions.

    See https://www.tensorflow.org/xla/operation_semantics#triangularsolve.
  }];
  let arguments = (ins
    HLO_FpOrComplexTensor:$a,
    HLO_FpOrComplexTensor:$b,
    BoolAttr:$left_side,
    BoolAttr:$lower,
    BoolAttr:$unit_diagonal,
    StableHLO_TransposeAttr:$transpose_a
  );
  let results = (outs HLO_FpOrComplexTensor);

  let hasVerifier = 1;
}

def StableHLO_ReduceWindowOp: StableHLO_Op<"reduce_window", [
      RecursiveSideEffects,
      SameVariadicOperandSize,
      SingleBlockImplicitTerminator<"ReturnOp">
    ]> {
  let summary = "ReduceWindow operator";
  let description = [{
    Returns the result of executing a reduction function over all elements in
    each window of one or more arrays in parallel.

    See https://www.tensorflow.org/xla/operation_semantics#reducewindow.
  }];

  // TODO(hinsu): Verify that padding attribute is 2-d and the remaining
  // attributes are 1-d. Attributes' leading dimension should match rank of the
  // operands.
  let arguments = (ins
    Variadic<HLO_Tensor>:$operands,
    Variadic<HLO_Tensor>:$init_values,
    I64ElementsAttr:$window_dimensions,
    // If strides or dilations attributes are missing then the default value is
    // one for each of the operand dimensions. Similarly, padding values are zero
    // for both low and high in each of the dimensions, if not specified.
    OptionalAttr<I64ElementsAttr>:$window_strides,
    OptionalAttr<I64ElementsAttr>:$base_dilations,
    OptionalAttr<I64ElementsAttr>:$window_dilations,
    OptionalAttr<I64ElementsAttr>:$padding
  );

  let results = (outs Variadic<HLO_Tensor>);

  // TODO(hinsu): Verify that the attached body arguments and results are
  // compatible with reduce op's operands.
  let regions = (region SizedRegion<1>:$body);

  // Builder for non-variadic version of the operation.
  let builders = [
    OpBuilder<(ins "Type":$result_type, "Value":$operand,
      "Value":$init_value,
      "DenseIntElementsAttr":$window_dimensions,
      "DenseIntElementsAttr":$window_strides,
      "DenseIntElementsAttr":$base_dilations,
      "DenseIntElementsAttr":$window_dilations,
      "DenseIntElementsAttr":$padding),
    [{
      build($_builder, $_state, TypeRange(result_type), ValueRange(operand),
            ValueRange(init_value), window_dimensions, window_strides,
            base_dilations, window_dilations, padding);
    }]>
  ];

  let hasVerifier = 1;
  // TODO(hinsu): Implement custom printer and parser.

  let extraClassDeclaration = [{
     // Get the operation used for reduction applied to `result_index`th result.
     Operation *getReductionOp(int result_index);
  }];
}

def StableHLO_ReturnOp : StableHLO_Op<"return", [NoSideEffect, Terminator]> {
  let summary = [{
    The `hlo.return` operation terminates a region and returns values.

    Example:

    ```mlir
    %0 = stablehlo.reduce %arg0, %arg1 {
      ...
      stablehlo.return %1 : tensor<f32>
    }
    ```
  }];

  let arguments = (ins
    Variadic<HLO_TensorOrTokenOrTuple >:$results
  );

  let assemblyFormat = "$results attr-dict (`:` type($results)^)?";
}

def StableHLO_TorchIndexSelectOp : StableHLO_Op<"torch_index_select", [NoSideEffect]> {
  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_Tensor:$index,
    I64Attr:$dim,
    I64Attr:$batch_dims
  );

  let results = (outs HLO_Tensor);

  // TODO(hinsu): Canonicalize to lower this client side HLO op to server
  // side HLO ops.
}

def StableHLO_OptimizationBarrierOp : StableHLO_Op<"optimization_barrier",
      [NoSideEffect, HLO_PairwiseSameOperandAndResultType]> {
  let summary = [{
    The `stablehlo.optimization_barrier` op blocks optimizations.

    Example:

    ```mlir
    %0:2 = stablehlo.optimization_barrier %arg0, %arg1 : (tensor<4x4xf32>, tensor<3x4xf32>) -> (tensor<4x4xf32>, tensor<3x4xf32>)
    ```
  }];

  let description = [{
    Blocks any optimization pass from moving computations across the barrier.

    Ensures that all inputs are evaluated before any operators that depend on the barrier's outputs.
    See
    https://www.tensorflow.org/xla/operation_semantics#optimizationbarrier
  }];

  let arguments = (ins Variadic<HLO_TensorOrToken>:$operand);

  let results = (outs Variadic<HLO_TensorOrToken>);

  // TODO(b/241767462): Enhance type printing to condense pairwise ops.
  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

//===----------------------------------------------------------------------===//
// StableHLO RNG Operators.
//===----------------------------------------------------------------------===//

def StableHLO_RngOp : StableHLO_Op<"rng", [InferTensorTypeWithReify, AllElementTypesMatch<["a", "b", "result"]>]> {
  let summary = "RNG with uniform distribution.";
  let description = [{
    Constructs an output of a given shape with random numbers generated
    following the given `rng_distribution` with two parameters:
      `UNIFORM`: the uniform distribution over the interval `[a,b)`. The parameters
                 and output element type have to be a boolean type, an integral type or a
                 floating point types, and the types have to be consistent.

                 See https://www.tensorflow.org/xla/operation_semantics#rnguniform.

      `NORMAL`: the normal distribution with parameters `mu` (=`a`) and
                `sigma` (=`b`). The parameters and output shape have to have a
                floating point elemental type. The parameters furthermore have
                to be scalar valued.

                See https://www.tensorflow.org/xla/operation_semantics#rngnormal.
  }];
  let arguments = (ins
    0DTensorOf<[HLO_Pred, HLO_Int, HLO_Float]>:$a,
    0DTensorOf<[HLO_Pred, HLO_Int, HLO_Float]>:$b,
    HLO_DimensionTensor:$shape,
    StableHLO_RngDistributionAttr:$rng_distribution
  );

  let results = (outs HLO_PredIntOrFpTensor:$result);

  let hasVerifier = 1;

  let extraClassDeclaration = [{
    // Returns whether the return types are compatible.
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(::mlir::verifyCompatibleShapes(l, r));
    }
  }];
}

def StableHLO_RngBitGeneratorOp : StableHLO_Op<"rng_bit_generator", [NoSideEffect]> {
  let summary = "Uniform random number generator operator";
  let description = [{
    Returns an output with a given shape filled with uniform random bits using
    the specified algorithm (or backend default) and returns an updated state
    (with the same shape as initial state) and the generated random data.

    See https://www.tensorflow.org/xla/operation_semantics#rngbitgenerator.
  }];
  let arguments = (ins
    StableHLO_RngAlgorithmAttr:$rng_algorithm,
    HLO_IntOrFpTensor:$initial_state
  );

  let results = (outs
      HLO_IntOrFpTensor:$output_state,
      HLO_IntOrFpTensor:$output
      );

  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// StableHLO Quantize Operator.
//===----------------------------------------------------------------------===//

// TODO(b/230662142): Implement unknown scales/zero_point cases.
def StableHLO_UniformQuantizeOp : StableHLO_UnaryElementwiseOp<"uniform_quantize",
      [NoSideEffect], TensorOf<[F32, BF16, HLO_QuantizedInt]>,
      HLO_QuantizedIntTensor> {
  let summary = "Uniform quantize operator";
  let description = [{
    Converts floating point tensors or uniform quantized integer tensors to
    uniform quantized integer tensors according to the quantization parameters
    defined by the output type.

    Example:

    ```mlir
    %0 = stablehlo.uniform_quantize %arg0 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
    ```
  }];
}

def StableHLO_UniformDequantizeOp : StableHLO_UnaryElementwiseOp<"uniform_dequantize",
      [InferTensorType, NoSideEffect], HLO_QuantizedIntTensor, TensorOf<[F32, BF16]>> {
  let summary = "Uniform dequantize operator";
  let description = [{
    Converts quantized array of integers to floating-points according to the
    quantization parameters defined by the input type.

    Example:

    ```mlir
    %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<i8:f32, 34.0:16>>) -> tensor<16x16xf32>
    ```
  }];
}

def StableHLO_ReducePrecisionOp :
    StableHLO_Op<"reduce_precision", [HLO_CompatibleOperandsAndResultType]> {
  let summary = "Reduce precision operator";
  let description = [{
    Models the effect of converting floating - point values to a lower -
    precision format(such as IEEE - FP16) and back to the original
    format. The number of exponent and mantissa bits in the lower -
    precision format can be specified arbitrarily,
    although all bit sizes may not be supported on all hardware
    implementations.

    See https://www.tensorflow.org/xla/operation_semantics#reduceprecision.
  }];
  let arguments = (ins
    HLO_FpTensor:$operand,
    I32Attr:$exponent_bits,
    I32Attr:$mantissa_bits
  );
  let hasVerifier = 1;
  let results = (outs HLO_FpTensor:$output);
}

def StableHLO_RealDynamicSliceOp: StableHLO_ShapedInterfaceOp<
      "real_dynamic_slice",
      [NoSideEffect, AllElementTypesMatch<["operand", "result"]>,
       AllTypesMatch<["start_indices", "limit_indices", "strides"]>]> {
  let summary = "Real Dynamic Slice operator";
  let description = [{
    The dynamic shape version of SliceOp. Extracts a sub-array from the input
    array according to start_indices, limit_indices and strides. Expect
    start_indices/limit_indices/strides to be statically shaped and matching
    the rank of the input.

    Example:

    ```mlir
    %0 = stablehlo.real_dynamic_slice %input, %start, %limit, %strides
           : (tensor<256x?xf32>, tensor<2xindex>, tensor<2xindex>, tensor<2xindex>) -> tensor<256x?xf32>
    ```
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_DimensionTensor:$start_indices,
    HLO_DimensionTensor:$limit_indices,
    HLO_DimensionTensor:$strides
  );
  let results = (outs HLO_Tensor:$result);
  let hasVerifier = 1;

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

def StableHLO_DynamicPadOp: StableHLO_ShapedInterfaceOp<"dynamic_pad",
      [NoSideEffect, AllElementTypesMatch<["operand", "padding_value", "result"]>,
      AllTypesMatch<["edge_padding_low", "edge_padding_high", "interior_padding"]>]> {
  let summary = "Dynamic Pad operator";
  let description = [{
    The dynamic shape version of PadOp. Pads the edges of `operand` with the
    `padding_value` and according to the passed configuration. Expect
    edge_padding_low/edge_padding_high/interior_padding to be statically shaped
    and matching the rank of the input.

    See https://www.tensorflow.org/xla/operation_semantics#pad.

    Example:

    ```mlir
    %0 = stablehlo.dynamic_pad %arg0, %arg1, %arg2, %arg3, %arg4
           : (tensor<?x?xf32>, tensor<f32>, tensor<2xindex>, tensor<2xindex>, tensor<2xindex>) -> tensor<?x?xf32>
    ```
  }];
  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_Tensor:$padding_value,
    HLO_DimensionTensor:$edge_padding_low,
    HLO_DimensionTensor:$edge_padding_high,
    HLO_DimensionTensor:$interior_padding
  );
  let results = (outs HLO_Tensor:$result);
  let description = [{
    Dynamically Pads the `operand`, with amount of padding added at
    low-end/high-end/interior is passed through input tensors.
  }];
  let hasVerifier = 1;

  let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
}

def StableHLO_DynamicGatherOp: StableHLO_Op<"dynamic_gather",
                                [InferTensorTypeWithReify, NoSideEffect]> {
  string summary = "Dynamic Gather operator";
  string description = [{
    The dynamic shape version of GatherOp. Stitches together several slices of
    an input array.
  }];

  let arguments = (ins
    HLO_Tensor:$operand,
    HLO_IntTensor:$start_indices,
    HLO_IntTensor:$slice_sizes,
    StableHLO_GatherDimensionNumbers:$dimension_numbers,
    DefaultValuedAttr<BoolAttr, "false">:$indices_are_sorted
  );
  let results = (outs HLO_Tensor);

  let extraClassDeclaration = [{
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return succeeded(mlir::verifyCompatibleShapes(l, r));
    }
  }];
}

def StableHLO_DynamicConvOp : StableHLO_Op<"dynamic_conv", [NoSideEffect]> {
  let summary = "Dynamic Convolution operator";
  let description = [{
    The dynamic shape version of ConvOp. Computes a convolution with dynamic padding.
  }];

  let arguments = !con(
    (ins
       HLO_Tensor:$lhs,
       HLO_Tensor:$rhs,
       HLO_Tensor:$d_padding),
    StableHLO_ConvolutionAttributes.attributes);
  let results = (outs HLO_Tensor);
}

def StableHLO_ComputeReshapeShapeOp :
    StableHLO_Op<"compute_reshape_shape", [NoSideEffect]> {
  string summary = "Compute input for reshape with any dynamic dim resolved";

  string description = [{
    This operation handles the dynamic aspect of a TF/NumPy/CHLO reshape. The
    dynamic aspect is that a single extent can be -1 and that dimension will
    instead be computed. This handles the computation and can then be passed to
    an HLO DynamicReshapeOp to replicate the TF/NumPy reshape behavior.

    This op has undefined behavior if the dimensions do not evenly divide the
    number of elements, or if there are multiple -1 values. It is an identity op
    if no dimensions are -1.

    ```
    %0 = hlo.compute_reshape_shape 12, [2, -1] -> [2, 6]
    ```
  }];

  let arguments = (ins Index:$num_elements, 1DTensorOf<[AnyInteger, Index]>:$dynamic_shape);
  let results = (outs 1DTensorOf<[AnyInteger, Index]>:$result);

  // TODO (b/241767462): Use functional-type for type printing for consistency.
  let assemblyFormat = "$num_elements `,` $dynamic_shape attr-dict `:` type($num_elements) `,` type($dynamic_shape) `->` type($result)";
}

def StableHLO_CstrReshapableOp :
    StableHLO_Op<"cstr_reshapable", [NoSideEffect]> {
  string summary = "Compute input for reshape with any dynamic dim resolved";

  string description = [{
    This operation creates a witness on the constraint that a given shape would
    be a valid reshape for the given number of elements.

    ```
    %0 = stablehlo.cstr_reshapable 12, [2, -1] -> success
    %1 = stablehlo.cstr_reshapable 13, [2, -1] -> failure
    ```
  }];

  let arguments = (ins Index:$num_elements, 1DTensorOf<[AnyInteger, Index]>:$dynamic_shape);
  let results = (outs Shape_WitnessType:$result);

  // TODO (b/241767462): Use functional-type for type printing for consistency.
  let assemblyFormat = "$num_elements `,` $dynamic_shape attr-dict `:` type($num_elements) `,` type($dynamic_shape)";
}

#endif // STABLEHLO_DIALECT_STABLEHLO_OPS