| /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. |
| |
| 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. |
| ==============================================================================*/ |
| |
| // This is the operation definition file for MHLO ops. |
| |
| #ifndef HLO_OPS |
| #define HLO_OPS |
| |
| include "mlir/Dialect/Shape/IR/ShapeBase.td" |
| include "mlir/Interfaces/InferTypeOpInterface.td" |
| include "mlir/Interfaces/SideEffectInterfaces.td" |
| include "mlir/IR/OpAsmInterface.td" |
| include "mlir/IR/OpBase.td" |
| include "mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td" |
| include "mlir-hlo/Dialect/mhlo/IR/hlo_utils.td" |
| |
| class HLO_Op<string mnemonic, list<Trait> traits> : |
| Op<HLO_Dialect, mnemonic, traits> { |
| // Whether this operation has a custom conversion to HLO or not. |
| bit hasCustomHLOConverter = 0b0; |
| } |
| |
| class HLO_ShapedInterfaceOp<string mnemonic, list<Trait> traits> : |
| HLO_Op<mnemonic, traits # [DeclareOpInterfaceMethods<InferShapedTypeOpInterface, |
| ["reifyReturnTypeShapes"]>]> { |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO nullary op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_ConstantOp : HLO_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; |
| |
| // Constant has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| |
| let hasFolder = 1; |
| |
| let extraClassDeclaration = [{ |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r); |
| }]; |
| } |
| |
| def HLO_IotaOp : HLO_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); |
| |
| // TODO(b/130357376): Iota has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| let hasCanonicalizer = 1; |
| let hasFolder = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_DynamicIotaOp: HLO_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); |
| |
| let hasCanonicalizer = 1; |
| // Cannot be exported to legacy formats. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| |
| def HLO_CreateTokenOp : HLO_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 = mhlo.create_token : !mhlo.token |
| ``` |
| }]; |
| |
| let results = (outs HLO_Token:$output); |
| |
| let assemblyFormat = "attr-dict `:` type(results)"; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO unary elementwise op definitions. |
| //===----------------------------------------------------------------------===// |
| // See https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions |
| |
| class HLO_UnaryElementwiseOp<string mnemonic, list<Trait> traits, |
| Type OperandType, Type ResultType = OperandType> : HLO_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::mhlo::deriveShapeFromOperand(&builder, getOperation(), |
| operands.front(), |
| &reifiedReturnShapes); |
| } |
| // Relax the strict default implementation with one that allows |
| // for MHLO-specific differences. |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) { |
| if (l.size() != r.size()) return false; |
| for (auto it : llvm::zip(l, r)) |
| if (!isCompatibleForMhloTypeInference(std::get<0>(it), std::get<1>(it))) |
| return false; |
| return true; |
| } |
| }]; |
| let extraClassDefinition = [{ |
| ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) { |
| return ::mlir::mhlo::parseUnaryOp(parser, result); |
| } |
| void $cppClass::print(OpAsmPrinter &p) { |
| ::mlir::mhlo::printUnaryOp(getOperation(), p); |
| } |
| }]; |
| let hasCustomAssemblyFormat = 1; |
| } |
| |
| // Abs supports complex to real, so element type is not guaranteed to match. |
| def HLO_AbsOp: HLO_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 HLO_CbrtOp: HLO_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 HLO_CeilOp: HLO_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 HLO_ConvertOp : HLO_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)>]; |
| |
| let hasFolder = 1; |
| |
| let hasCanonicalizer = 1; |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_ClzOp: HLO_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 HLO_CosineOp: HLO_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. |
| }]; |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_ExpOp: HLO_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 HLO_Expm1Op: HLO_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 HLO_FloorOp: HLO_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 HLO_ImagOp: HLO_UnaryElementwiseOp<"imag", |
| [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>], |
| HLO_FpOrComplexTensor> { |
| let summary = "Imag operator"; |
| let description = [{ |
| Returns `Imag(operand)` element-wise. |
| }]; |
| let results = (outs HLO_FpTensor); |
| let hasFolder = 1; |
| } |
| |
| def HLO_IsFiniteOp: HLO_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 HLO_LogOp: HLO_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 HLO_Log1pOp: HLO_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 HLO_LogisticOp: HLO_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 HLO_NotOp: HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_NegOp: HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_PopulationCountOp: HLO_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 HLO_RealOp: HLO_UnaryElementwiseOp<"real", |
| [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>], |
| HLO_FpOrComplexTensor> { |
| let summary = "Real operator"; |
| let description = [{ |
| Returns `Real(operand)` element-wise. |
| }]; |
| let results = (outs HLO_FpTensor); |
| let hasFolder = 1; |
| } |
| |
| def HLO_RoundOp: HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_RoundNearestEvenOp: HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_RsqrtOp: HLO_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 HLO_SignOp: HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_SineOp: HLO_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. |
| }]; |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_SqrtOp: HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_TanhOp: HLO_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. |
| }]; |
| } |
| //===----------------------------------------------------------------------===// |
| // MHLO binary elementwise op definitions. |
| //===----------------------------------------------------------------------===// |
| // See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations |
| |
| class HLO_BinaryElementwiseOpNoAssembly<string mnemonic, list<Trait> traits> : |
| HLO_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::mhlo::deriveShapeFromOperand(&builder, getOperation(), |
| operands.front(), |
| &reifiedReturnShapes); |
| } |
| // Relax the strict default implementation with one that allows |
| // for MHLO-specific differences. |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) { |
| if (l.size() != r.size()) return false; |
| for (auto it : llvm::zip(l, r)) |
| if (!isCompatibleForMhloTypeInference(std::get<0>(it), std::get<1>(it))) |
| return false; |
| return true; |
| } |
| }]; |
| |
| let results = (outs HLO_Tensor:$result); |
| } |
| |
| class HLO_BinaryElementwiseOp<string mnemonic, list<Trait> traits> : |
| HLO_BinaryElementwiseOpNoAssembly<mnemonic, traits> { |
| let extraClassDefinition = [{ |
| ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) { |
| return ::mlir::mhlo::parseBinaryOp(parser, result); |
| } |
| void $cppClass::print(OpAsmPrinter &p) { |
| ::mlir::mhlo::printBinaryOp(getOperation(), p); |
| } |
| }]; |
| let hasCustomAssemblyFormat = 1; |
| } |
| |
| def HLO_AddOp : HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_Atan2Op : HLO_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 HLO_ComplexOp: HLO_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); |
| |
| let hasFolder = 1; |
| |
| // TODO(b/241767457): Remove parens when cleaning up BinaryOps. |
| let assemblyFormat = "`(`operands`)` attr-dict `:` `(`type(operands)`)` `->` type($result)"; |
| } |
| |
| def HLO_DivOp : HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_MaxOp : HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_MinOp : HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_MulOp : HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_PowOp : HLO_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 HLO_RemOp : HLO_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. |
| }]; |
| let hasFolder = 1; |
| } |
| |
| def HLO_ShiftLeftOp : HLO_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 HLO_ShiftRightArithmeticOp : HLO_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 HLO_ShiftRightLogicalOp : HLO_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 HLO_SubtractOp : HLO_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. |
| }]; |
| let hasFolder = 1; |
| let hasCustomHLOConverter = 1; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO binary logical elementwise op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| // See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations |
| class HLO_BinaryBiwiseOrLogicalElementwiseOp<string mnemonic> : |
| HLO_BinaryElementwiseOp<mnemonic, |
| [Commutative, NoSideEffect, HLO_CompatibleOperandsAndResultType]> { |
| let arguments = (ins |
| HLO_PredOrIntTensor:$lhs, |
| HLO_PredOrIntTensor:$rhs |
| ); |
| |
| let hasFolder = 1; |
| } |
| |
| def HLO_AndOp: HLO_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 HLO_OrOp: HLO_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 HLO_XorOp : HLO_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. |
| }]; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO communication op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| // InfeedOp corresponds to 'InfeedWithToken' xla client API and not 'Infeed'. |
| // InfeedWithToken allows ordering of infeed HLO instructions using tokens. |
| def HLO_InfeedOp : HLO_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 hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| } |
| |
| // OutfeedOp corresponds to 'OutfeedWithToken' xla client API and not 'Outfeed'. |
| // OutfeedWithToken allows ordering of outfeed HLO instructions using tokens. |
| def HLO_OutfeedOp : HLO_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); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_SendOp : HLO_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, |
| ChannelHandle:$channel_handle, |
| DefaultValuedAttr<BoolAttr, "false">:$is_host_transfer |
| ); |
| |
| let results = (outs HLO_Token); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_RecvOp : HLO_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, |
| ChannelHandle:$channel_handle, |
| DefaultValuedAttr<BoolAttr, "false">:$is_host_transfer |
| ); |
| |
| let results = (outs Variadic<HLO_TensorOrToken>); |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO parallelism related op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_ReplicaIdOp : HLO_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 = mhlo.replica_id : tensor<ui32> |
| ``` |
| }]; |
| let results = (outs TensorOf<[UI32]>); |
| |
| let assemblyFormat = "attr-dict `:` type(results)"; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO control flow op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_AddDependencyOp : HLO_Op<"add_dependency", [NoSideEffect, |
| DeclareOpInterfaceMethods<InferTypeOpInterface>]> { |
| let summary = "AddDependency operator"; |
| let description = [{ |
| AddDependency takes two operands: a data operand and a token. The output of |
| the operation is the data operand. When used with AfterAll this operation |
| enables ordering non-side-effecting operations (those that do not produce |
| token values). |
| |
| Example: |
| |
| ```mlir |
| %1 = mhlo.add_dependency %arg0, %0 : (tensor<3x4xf32>, !mhlo.token) -> tensor<3x4xf32> |
| ``` |
| }]; |
| |
| let arguments = (ins HLO_TensorOrToken:$operand, HLO_Token:$token); |
| let results = (outs HLO_TensorOrToken:$output); |
| let hasCustomHLOConverter = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| |
| def HLO_AfterAllOp : HLO_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 = mhlo.after_all %arg0, %arg1 : (!mhlo.token, !mhlo.token) -> !mhlo.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 HLO_IfOp: HLO_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>); |
| |
| // TODO(b/129422361): ConditionalOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| |
| let hasCanonicalizer = 1; |
| 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 HLO_CaseOp: HLO_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 hasCustomHLOConverter = 1; |
| |
| let hasCanonicalizer = 1; |
| let hasVerifier = 1; |
| } |
| |
| |
| def HLO_WhileOp: HLO_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 ®ion, OpAsmSetValueNameFn setNameFn) { |
| for (BlockArgument arg : region.getArguments()) |
| setNameFn(arg, "iterArg"); |
| } |
| }]; |
| // TODO(b/129422361): WhileOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| let hasCanonicalizer = 1; |
| let hasCustomAssemblyFormat = 1; |
| let hasFolder = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_AllGatherOp : HLO_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<ChannelHandle>:$channel_handle |
| ); |
| let results = (outs HLO_Tensor); |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_AllReduceOp : HLO_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<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::mhlo::ChannelHandleAttr":$channel_handle)>]; |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_ReduceScatterOp : HLO_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<ChannelHandle>:$channel_handle |
| ); |
| let regions = (region SizedRegion<1>:$computation); |
| let results = (outs HLO_Tensor); |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_AllToAllOp : HLO_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 HLO_ReduceOp: HLO_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 hasCanonicalizer = 1; |
| let hasCustomAssemblyFormat = 1; |
| let hasFolder = 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); |
| |
| // TODO(b/129422361): ReduceOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO tuple op definitions. |
| //===----------------------------------------------------------------------===// |
| def HLO_GetTupleElementOp: HLO_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 hasFolder = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_TupleOp : HLO_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 hasCanonicalizer = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_CompareOp: HLO_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 = mhlo.compare LT, %arg0, %arg1 : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> |
| %1 = mhlo.compare LT, %arg0, %arg1, TOTALORDER : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> |
| ``` |
| }]; |
| let arguments = (ins |
| HLO_Tensor:$lhs, |
| HLO_Tensor:$rhs, |
| HLO_ComparisonDirectionAttr:$comparison_direction, |
| OptionalAttr<HLO_ComparisonTypeAttr>:$compare_type |
| ); |
| let results = (outs HLO_PredTensor); |
| |
| let hasFolder = 1; |
| |
| let builders = [ |
| OpBuilder<(ins "Value":$lhs, "Value":$rhs, |
| "::mlir::mhlo::ComparisonDirection":$comparison_direction, |
| CArg<"::mlir::mhlo::ComparisonType", |
| "::mlir::mhlo::ComparisonType::NOTYPE">:$compare_type)>, |
| ]; |
| |
| let hasCustomHLOConverter = 1; |
| |
| 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) |
| }]; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO Slice definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_SliceOp: HLO_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); |
| |
| let hasCanonicalizer = 1; |
| let hasFolder = 1; |
| } |
| |
| def HLO_DynamicSliceOp: HLO_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 hasCanonicalizer = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_DynamicUpdateSliceOp: HLO_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 = mhlo.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 hasFolder = 1; |
| let hasVerifier = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO Other op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_DomainOp : HLO_Op<"domain", [HLO_CompatibleOperandsAndResultType, InferTypeOpInterface, NoSideEffect]> { |
| let summary = "Marks groups of instructions (domains) with a property"; |
| let description = [{ |
| Domain instructions are used to group instructions with the same |
| DomainMetadata property. ShardingMetadata is the main use case today to |
| group instructions on the same device. Domain instructions provide two |
| major benefits: |
| - Prevent unintentionally optimizing instructions across domains. |
| - Automatically assign the metadata of the instructions created in the domain. |
| Without domain instructions, each HLO optimization pass would have to check |
| and propagate the metadata, which would be easy to miss and also adds |
| complexity to the compiler. Since domain instructions connect two different |
| domains, each domain instruction is associated with two DomainMetadata -- |
| one on the operand side and one on the user side of the domain. |
| }]; |
| let arguments = (ins |
| HLO_TensorOrToken:$operand, |
| HLO_DomainKindAttr:$kind, |
| StrAttr:$entry_metadata, |
| StrAttr:$exit_metadata |
| ); |
| let results = (outs HLO_TensorOrToken:$result); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_BatchNormGradOp : HLO_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 hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_BatchNormInferenceOp : HLO_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 HLO_BatchNormTrainingOp : HLO_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; |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_BitcastConvertOp : HLO_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 = mhlo.bitcast_convert %arg0 : (tensor<2xi32>) -> tensor<2xf32> |
| ``` |
| }]; |
| |
| let arguments = (ins HLO_Tensor:$operand); |
| let results = (outs HLO_Tensor); |
| let hasVerifier = 1; |
| let hasCustomHLOConverter = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_BroadcastOp : HLO_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 hasFolder = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_BroadcastInDimOp : HLO_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 hasFolder = 1; |
| let hasCanonicalizer = 1; |
| let hasVerifier = 1; |
| // Only handles a static subset of the legacy format. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_DynamicBroadcastInDimOp : HLO_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 hasCanonicalizer = 1; |
| let hasVerifier = 1; |
| // Cannot be exported to legacy formats. |
| let hasCustomHLOConverter = 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 HLO_CholeskyOp : HLO_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 HLO_ClampOp : HLO_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 = mhlo.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 it : llvm::zip(l, r)) |
| if (!isCompatibleForMhloTypeInference(std::get<0>(it), std::get<1>(it))) |
| return false; |
| return true; |
| } |
| }]; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_ConcatenateOp : HLO_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 hasCanonicalizer = 1; |
| let hasFolder = 1; |
| let hasVerifier = 1; |
| |
| let extraClassDeclaration = [{ |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) { |
| return succeeded(mlir::verifyCompatibleShapes(l, r)); |
| } |
| }]; |
| } |
| |
| def HLO_CollectivePermuteOp: HLO_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 HLO_ConvolutionOp : HLO_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), |
| ConvolutionAttributes.attributes); |
| |
| let results = (outs HLO_Tensor); |
| let hasCanonicalizer = 1; |
| let hasCustomHLOConverter = 1; |
| 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 HLO_CopyOp: HLO_Op<"copy", [NoSideEffect, HLO_CompatibleOperandsAndResultType]> { |
| let summary = "Copy operator"; |
| let description = [{ |
| Returns a copy of `operand`. |
| |
| Example: |
| |
| ```mlir |
| %0 = mhlo.copy %arg0 : (tensor<f32>) -> tensor<f32> |
| ``` |
| }]; |
| let arguments = (ins HLO_Tensor); |
| let results = (outs HLO_Tensor); |
| let hasFolder = 1; |
| |
| // TODO (b/241767462): Condense type printing if input/output type match. |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_CrossReplicaSumOp : HLO_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 HLO_CustomCallOp: HLO_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< |
| HLO_CustomCallApiVersionAttr, |
| "::mlir::mhlo::CustomCallApiVersion::API_VERSION_ORIGINAL">: |
| $api_version, |
| DefaultValuedAttr<HLO_FlatSymbolRefArrayAttr, "{}">:$called_computations, |
| OptionalAttr<HLO_ArrayOfLayoutAttr>:$operand_layouts, |
| OptionalAttr<HLO_ArrayOfLayoutAttr>:$result_layouts |
| ); |
| let results = (outs Variadic<HLO_TensorOrTokenOrTuple>); |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_DotOp: HLO_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, |
| HLO_PrecisionConfigAttr:$precision_config |
| ); |
| let results = (outs HLO_Tensor); |
| // Dot op required custom exporter to pass the preferred element type |
| // to Xla builder. |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| |
| let extraClassDeclaration = [{ |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) { |
| return succeeded(mlir::verifyCompatibleShapes(l, r)); |
| } |
| }]; |
| } |
| |
| def HLO_DotGeneralOp: HLO_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, |
| DotDimensionNumbers:$dot_dimension_numbers, |
| HLO_PrecisionConfigAttr:$precision_config |
| ); |
| |
| let results = (outs HLO_Tensor); |
| let hasCanonicalizer = 1; |
| // DotGeneral op required custom exporter to pass the preferred element type |
| // to Xla builder. |
| let hasCustomHLOConverter = 1; |
| 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 HLO_EinsumOp: HLO_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 HLO_UnaryEinsumOp: HLO_Op<"unary_einsum", [NoSideEffect]>, BASE_EinsumOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| StrAttr:$einsum_config |
| ); |
| |
| let results = (outs HLO_Tensor); |
| |
| let hasCanonicalizer = 1; |
| |
| // UnaryEinsumOp is unconditionally canonicalized to the binary EinsumOp so |
| // the HLO converter shouldn't be invoked. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_FftOp: HLO_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, |
| HLO_FftTypeAttr: $fft_type, |
| I64ElementsAttr:$fft_length |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_GatherOp: HLO_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, |
| GatherDimensionNumbers:$dimension_numbers, |
| I64ElementsAttr:$slice_sizes, |
| DefaultValuedAttr<BoolAttr, "false">:$indices_are_sorted |
| ); |
| |
| let results = (outs HLO_Tensor); |
| |
| let hasCanonicalizer = 1; |
| |
| let extraClassDeclaration = [{ |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) { |
| return succeeded(mlir::verifyCompatibleShapes(l, r)); |
| } |
| }]; |
| } |
| |
| def HLO_GetDimensionSizeOp: HLO_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 hasFolder = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_MapOp: HLO_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 hasFolder = 1; |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_ReshapeOp: HLO_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 = mhlo.reshape %arg0 : (tensor<2xf32>) -> tensor<1x2xf32> |
| ``` |
| }]; |
| |
| let arguments = (ins HLO_Tensor:$operand); |
| |
| let results = (outs HLO_StaticShapeTensor); |
| let hasFolder = 1; |
| let hasCanonicalizer = 1; |
| let hasVerifier = 1; |
| |
| let hasCustomHLOConverter = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_DynamicReshapeOp: HLO_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 = mhlo.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 hasCanonicalizer = 1; |
| // Cannot be exported to legacy formats. |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_ScatterOp: HLO_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, |
| 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 hasCustomHLOConverter = 1; |
| |
| let hasFolder = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_SelectOp: HLO_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 hasFolder = 1; |
| let hasVerifier = 1; |
| let hasCanonicalizer = 1; |
| |
| let extraClassDeclaration = [{ |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) { |
| return succeeded(mlir::verifyCompatibleShapes(l, r)); |
| } |
| }]; |
| } |
| |
| def HLO_SelectAndScatterOp: HLO_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; |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_SetDimensionSizeOp: HLO_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 it : llvm::zip(l, r)) |
| if (!isCompatibleForMhloTypeInference(std::get<0>(it), std::get<1>(it))) |
| return false; |
| return true; |
| } |
| }]; |
| |
| |
| let hasFolder = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_SortOp : HLO_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)>]; |
| |
| // TODO(b/129422361): SortOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| |
| let hasCanonicalizer = 1; |
| let hasVerifier = 1; |
| } |
| |
| def HLO_ReverseOp: HLO_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); |
| |
| let hasFolder = 1; |
| } |
| |
| def HLO_PartitionIdOp : HLO_Op<"partition_id", []> { |
| let summary = "PartitionId operator"; |
| let description = [{ |
| Returns the value of the partition id of the currently executing device. |
| XLA supports two mechanisms for parallel execution: replication and |
| partition. A module can be replicated to run on multiple devices |
| (replicas) and a module can also be partitioned to split the work |
| between devices. replica-id and partition-id returns the id values of the |
| current device. |
| |
| Example: |
| |
| ```mlir |
| %1 = mhlo.partition_id : tensor<ui32> |
| ``` |
| }]; |
| let results = (outs TensorOf<[UI32]>); |
| let hasCustomHLOConverter = 1; |
| |
| let assemblyFormat = "attr-dict `:` type(results)"; |
| } |
| |
| def HLO_PadOp: HLO_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); |
| |
| // TODO(b/129422361): PadOp has a custom constructor for HLO. |
| let hasCustomHLOConverter = 1; |
| |
| let hasCanonicalizer = 1; |
| let hasFolder = 1; |
| } |
| |
| def HLO_TraceOp: HLO_Op<"trace", []> { |
| let summary = "Trace operator"; |
| let description = [{ |
| Emits a logging message `tag` with the `operand`. |
| |
| Example: |
| |
| ```mlir |
| mhlo.trace %arg0, "In test code." : tensor<5x1x5xi32> |
| ``` |
| }]; |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| StrAttr:$tag |
| ); |
| let hasCustomHLOConverter = 1; |
| let assemblyFormat = "$operand `,` $tag attr-dict `:` type($operand)"; |
| } |
| |
| def HLO_TransposeOp: HLO_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)); |
| } |
| }]; |
| |
| let hasFolder = 1; |
| let hasCanonicalizer = 1; |
| } |
| |
| def HLO_TriangularSolveOp: HLO_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, |
| HLO_TransposeAttr:$transpose_a |
| ); |
| let results = (outs HLO_FpOrComplexTensor); |
| |
| let hasVerifier = 1; |
| } |
| |
| def HLO_ReduceWindowOp: HLO_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 hasCustomHLOConverter = 1; |
| 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 HLO_ReturnOp : HLO_Op<"return", [NoSideEffect, Terminator]> { |
| let summary = [{ |
| The `hlo.return` operation terminates a region and returns values. |
| |
| Example: |
| |
| ```mlir |
| %0 = mhlo.reduce %arg0, %arg1 { |
| ... |
| mhlo.return %1 : tensor<f32> |
| } |
| ``` |
| }]; |
| |
| let arguments = (ins |
| Variadic<HLO_TensorOrTokenOrTuple >:$results |
| ); |
| |
| // Disable conversion operator for return op as the op is not an actual XLA |
| // instruction and is only used as a terminator for regions. |
| let hasCustomHLOConverter = 1; |
| |
| let assemblyFormat = "$results attr-dict (`:` type($results)^)?"; |
| } |
| |
| def HLO_TorchIndexSelectOp : HLO_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 HLO_OptimizationBarrierOp : HLO_Op<"optimization_barrier", |
| [NoSideEffect, HLO_PairwiseSameOperandAndResultType]> { |
| let summary = [{ |
| The `mhlo.optimization_barrier` op blocks optimizations. |
| |
| Example: |
| |
| ```mlir |
| %0:2 = mhlo.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>); |
| |
| let hasCustomHLOConverter = 1; |
| |
| // TODO(b/241767462): Enhance type printing to condense pairwise ops. |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO RNG Operators. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_RngOp : HLO_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, |
| HLO_RngDistributionAttr:$rng_distribution |
| ); |
| |
| let results = (outs HLO_PredIntOrFpTensor:$result); |
| |
| let hasCustomHLOConverter = 1; |
| 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 HLO_RngBitGeneratorOp : HLO_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 |
| HLO_RngAlgorithmAttr:$rng_algorithm, |
| HLO_IntOrFpTensor:$initial_state |
| ); |
| |
| let results = (outs |
| HLO_IntOrFpTensor:$output_state, |
| HLO_IntOrFpTensor:$output |
| ); |
| |
| let hasVerifier = 1; |
| // TODO(jpienaar): This should not be needed. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_XlaRngGetAndUpdateStateOp: HLO_Op<"xla.rng_get_and_update_state", [DeclareOpInterfaceMethods<InferTypeOpInterface>]> { |
| let summary = "RNG state change"; |
| let description = [{ |
| This instruction represents the change of the global random number generator |
| state for rng instructions. The global state is incremented by delta and |
| the old state is returned. |
| |
| The output is currently defined for a single output type. If this changes in |
| the future to support multiple types, lowering to use of a global memref |
| must ensure that a single memref is still used and updated appropriately. |
| }]; |
| let arguments = (ins I64Attr:$delta); |
| let results = (outs StaticShapeTensorOf<[UI64]>); |
| |
| let hasVerifier = 1; |
| let assemblyFormat = "attr-dict"; |
| |
| // Doesn't have an XLA builder equivalent. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO Quantize Operator. |
| //===----------------------------------------------------------------------===// |
| |
| // TODO(b/230662142): Implement unknown scales/zero_point cases. |
| def HLO_UniformQuantizeOp : HLO_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 = mhlo.uniform_quantize %arg0 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>> |
| ``` |
| }]; |
| |
| // Currently, it doesn't have an XLA builder equivalent. |
| // TODO(b/230671877): Implement XLA import/export for quantized MHLO ops. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_UniformDequantizeOp : HLO_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 = mhlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<i8:f32, 34.0:16>>) -> tensor<16x16xf32> |
| ``` |
| }]; |
| |
| // Currently, it doesn't have an XLA builder equivalent. |
| // TODO(b/230671877): Implement XLA import/export for quantized MHLO ops. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_FusionOp : HLO_Op<"fusion", []> { |
| let summary = "Fusion operator"; |
| let description = [{ |
| Models the fusion instruction. |
| |
| A fusion op is consists of a group of basic ops (represented as a region |
| attached to it). It serves as a hint to the backend that it is beneficial |
| to emit the contained ops into a single loop nest or kernel. |
| }]; |
| let regions = (region SizedRegion<1>:$fused_computation); |
| |
| let arguments = (ins |
| Variadic<HLO_TensorOrToken>:$operands, |
| OptionalAttr<HLO_FusionKindAttr>:$fusion_kind |
| ); |
| |
| let results = (outs |
| Variadic<HLO_TensorOrTuple>:$results |
| ); |
| |
| // FusionOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| // This is an op for purposes internal to XLA/GPU. |
| def HLO_BitcastOp : HLO_Op<"bitcast", [NoSideEffect]> { |
| let summary = "Bitcast operator"; |
| let description = [{ |
| This op changes the shape of the input in the way that the physical |
| arrangement of elements are unchanged. |
| |
| However, the op needs layout information to make sense of "physical |
| arrangement of elements". Layout support in MHLO is currently under |
| exploration. |
| |
| Example: |
| |
| ```mlir |
| %0 = mhlo.bitcast %arg0 : (tensor<3x4xf32>) -> tensor<3x4x1xf32> |
| ``` |
| }]; |
| |
| let arguments = (ins HLO_Tensor:$operand); |
| let results = (outs HLO_Tensor); |
| let hasCustomHLOConverter = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_ReducePrecisionOp : |
| HLO_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 HLO_RealDynamicSliceOp: HLO_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 = mhlo.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 hasCanonicalizer = 1; |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_DynamicPadOp: HLO_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 = mhlo.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 hasCanonicalizer = 1; |
| let hasCustomHLOConverter = 1; |
| let hasVerifier = 1; |
| |
| let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; |
| } |
| |
| def HLO_DynamicGatherOp: HLO_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, |
| GatherDimensionNumbers:$dimension_numbers, |
| DefaultValuedAttr<BoolAttr, "false">:$indices_are_sorted |
| ); |
| let results = (outs HLO_Tensor); |
| |
| let hasCustomHLOConverter = 1; |
| |
| let extraClassDeclaration = [{ |
| static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) { |
| return succeeded(mlir::verifyCompatibleShapes(l, r)); |
| } |
| }]; |
| } |
| |
| def HLO_DynamicConvOp : HLO_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), |
| ConvolutionAttributes.attributes); |
| let results = (outs HLO_Tensor); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_ComputeReshapeShapeOp : |
| HLO_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)"; |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_CstrReshapableOp : |
| HLO_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 = mhlo.cstr_reshapable 12, [2, -1] -> success |
| %1 = mhlo.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)"; |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| #endif // HLO_OPS |