| /* 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/IR/OpBase.td" |
| include "mlir/Interfaces/InferTypeOpInterface.td" |
| include "mlir/Interfaces/SideEffectInterfaces.td" |
| include "mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td" |
| include "mlir-hlo/Dialect/mhlo/IR/hlo_utils.td" |
| include "mlir-hlo/Dialect/mhlo/IR/infer_fusibility_op_interface.td" |
| |
| def HLO_Dialect : Dialect { |
| let name = "mhlo"; |
| let cppNamespace = "mhlo"; |
| } |
| |
| class HLO_Op<string mnemonic, list<OpTrait> traits> : |
| Op<HLO_Dialect, mnemonic, traits> { |
| // Whether this operation has a custom conversion to HLO or not. |
| bit hasCustomHLOConverter = 0b0; |
| |
| // TODO(b/129012527) Much of this custom verification should be expressed as |
| // type constraints. |
| let verifier = [{ return Verify(*this); }]; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO nullary op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_ConstOp : HLO_Op<"constant", |
| [ConstantLike, NoSideEffect, AllTypesMatch<["value", "output"]>]>, |
| BASE_HLO_ConstOp { |
| let arguments = (ins |
| ElementsAttr:$value |
| ); |
| |
| let results = (outs |
| HLO_StaticShapeTensor:$output |
| ); |
| |
| let builders = [OpBuilder< |
| "OpBuilder &builder, OperationState &result, Attribute value" |
| >]; |
| |
| let printer = [{ return Print(*this, &p); }]; |
| let parser = [{ return ParseConstOp(&parser, &result); }]; |
| |
| let hasFolder = 1; |
| |
| // Constant has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_IotaOp : HLO_Op<"iota", [NoSideEffect]>, BASE_HLO_IotaOp { |
| 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; |
| } |
| |
| def HLO_DynamicIotaOp: HLO_Op<"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]> { |
| string summary = "Create Token operator"; |
| |
| string description = [{ |
| Produces a HLO token. Tokens are used for ordering side-effecting perations. |
| This is exported to HLO as an AfterAll operation with no operands to |
| generate a token. |
| }]; |
| |
| let results = (outs HLO_Token:$output); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO unary elementwise op definitions. |
| //===----------------------------------------------------------------------===// |
| // See https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions |
| |
| class HLO_UnaryElementwiseOp<string mnemonic, list<OpTrait> traits, |
| Type TensorType>: HLO_Op<mnemonic, |
| !listconcat(traits, [InferShapedTypeOpInterface, InferFusibilityOpInterface])> { |
| let arguments = (ins TensorType:$operand); |
| let results = (outs TensorType); |
| let extraClassDeclaration = [{ |
| static LogicalResult inferReturnTypeComponents( |
| MLIRContext* context, Optional<Location> location, |
| ValueRange operands, DictionaryAttr attributes, RegionRange regions, |
| SmallVectorImpl<ShapedTypeComponents>& inferedReturnShapes) { |
| return failure(); |
| } |
| LogicalResult reifyReturnTypeShapes( |
| OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) { |
| return deriveShapeFromFirstOperand(&builder, getOperation(), |
| &reifiedReturnShapes); |
| } |
| bool inferInputOutputShapeEquality(int input, int output) { |
| return true; |
| } |
| llvm::Optional<Value> inferEffectiveWorkloadShape() { |
| return getOperation()->getResult(0); |
| } |
| }]; |
| } |
| |
| // Abs supports complex to real, so element type is not guaranteed to match. |
| def HLO_AbsOp: HLO_UnaryElementwiseOp<"abs", |
| [NoSideEffect, SameOperandsAndResultShape], |
| TensorOf<[HLO_SInt, AnyFloat, HLO_Complex]>>, BASE_HLO_AbsOp { |
| let builders = [OpBuilder< |
| "OpBuilder &builder, OperationState &result, Value operand" |
| >]; |
| } |
| |
| def HLO_CeilOp: HLO_UnaryElementwiseOp<"ceil", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpTensor>, BASE_HLO_CeilOp; |
| |
| def HLO_ConvertOp : HLO_UnaryElementwiseOp< |
| "convert", [NoSideEffect, SameOperandsAndResultShape], HLO_Tensor>, |
| BASE_HLO_ConvertOp { |
| |
| let builders = [OpBuilder< |
| "OpBuilder &, OperationState &tblgen_state, Value operand, " |
| "Type result_element_ty" |
| >]; |
| |
| let hasFolder = 1; |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_ClzOp: HLO_UnaryElementwiseOp<"count_leading_zeros", |
| [NoSideEffect, SameOperandsAndResultType], HLO_IntTensor>, |
| BASE_HLO_ClzOp; |
| |
| def HLO_CosOp: HLO_UnaryElementwiseOp<"cosine", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_CosOp; |
| |
| def HLO_ExpOp: HLO_UnaryElementwiseOp<"exponential", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_ExpOp; |
| |
| def HLO_Expm1Op: HLO_UnaryElementwiseOp<"exponential_minus_one", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_Expm1Op; |
| |
| def HLO_FloorOp: HLO_UnaryElementwiseOp<"floor", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpTensor>, BASE_HLO_FloorOp; |
| |
| def HLO_ImagOp: HLO_Op< |
| "imag", [NoSideEffect, SameOperandsAndResultShape]>, BASE_HLO_ImagOp { |
| let builders = [OpBuilder< |
| "OpBuilder &, OperationState &tblgen_state, Value val">]; |
| |
| let arguments = (ins HLO_ComplexTensor); |
| let results = (outs HLO_FpTensor); |
| let hasFolder = 1; |
| } |
| |
| def HLO_IsFiniteOp: HLO_UnaryElementwiseOp<"is_finite", |
| [NoSideEffect, SameOperandsAndResultShape], HLO_Tensor>, |
| BASE_HLO_IsFiniteOp { |
| let arguments = (ins HLO_FpTensor:$x); |
| let results = (outs HLO_PredTensor:$y); |
| } |
| |
| def HLO_LogOp: HLO_UnaryElementwiseOp<"log", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_LogOp; |
| |
| def HLO_Log1pOp: HLO_UnaryElementwiseOp<"log_plus_one", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_Log1pOp; |
| |
| def HLO_LogisticOp: HLO_UnaryElementwiseOp<"logistic", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_LogisticOp; |
| |
| def HLO_NotOp: HLO_UnaryElementwiseOp<"not", |
| [NoSideEffect, SameOperandsAndResultType], HLO_PredOrIntTensor>, |
| BASE_HLO_NotOp; |
| |
| def HLO_NegOp: HLO_UnaryElementwiseOp<"negate", |
| [NoSideEffect, SameOperandsAndResultType], HLO_IntFpOrComplexTensor>, |
| BASE_HLO_NegOp; |
| |
| def HLO_PopulationCountOp: HLO_UnaryElementwiseOp<"popcnt", |
| [NoSideEffect, SameOperandsAndResultType], HLO_IntTensor>, |
| BASE_HLO_PopulationCountOp; |
| |
| def HLO_RealOp: HLO_Op< |
| "real", [NoSideEffect, SameOperandsAndResultShape]>, BASE_HLO_RealOp { |
| let builders = [OpBuilder< |
| "OpBuilder &, OperationState &tblgen_state, Value val">]; |
| |
| let arguments = (ins HLO_ComplexTensor); |
| let results = (outs HLO_FpTensor); |
| let hasFolder = 1; |
| } |
| |
| def HLO_RoundOp: HLO_UnaryElementwiseOp<"round_nearest_afz", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpTensor>, BASE_HLO_RoundOp; |
| |
| def HLO_RsqrtOp: HLO_UnaryElementwiseOp<"rsqrt", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_RsqrtOp; |
| |
| def HLO_SignOp: HLO_UnaryElementwiseOp<"sign", |
| [NoSideEffect, SameOperandsAndResultType], |
| TensorOf<[HLO_SInt, AnyFloat, HLO_Complex]>>, |
| BASE_HLO_SignOp; |
| |
| def HLO_SinOp: HLO_UnaryElementwiseOp<"sine", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_SinOp; |
| |
| def HLO_SqrtOp: HLO_UnaryElementwiseOp<"sqrt", |
| [NoSideEffect, SameOperandsAndResultType], HLO_FpOrComplexTensor>, |
| BASE_HLO_SqrtOp; |
| |
| def HLO_TanhOp: HLO_UnaryElementwiseOp<"tanh", |
| [NoSideEffect, SameOperandsAndResultType], |
| HLO_FpOrComplexTensor>, BASE_HLO_TanhOp; |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO binary elementwise op definitions. |
| //===----------------------------------------------------------------------===// |
| // See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations |
| |
| class HLO_BinaryElementwiseOp<string mnemonic, list<OpTrait> traits> : |
| HLO_Op<mnemonic, !listconcat(traits, [InferShapedTypeOpInterface, InferFusibilityOpInterface])> { |
| let arguments = (ins |
| HLO_Tensor:$lhs, |
| HLO_Tensor:$rhs |
| ); |
| |
| let extraClassDeclaration = [{ |
| static LogicalResult inferReturnTypeComponents( |
| MLIRContext* context, Optional<Location> location, ValueRange operands, |
| DictionaryAttr attributes, RegionRange regions, |
| SmallVectorImpl<ShapedTypeComponents>& inferedReturnShapes) { |
| return failure(); |
| } |
| LogicalResult reifyReturnTypeShapes( |
| OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) { |
| return deriveShapeFromFirstOperand(&builder, getOperation(), |
| &reifiedReturnShapes); |
| } |
| bool inferInputsShapeEquality(int lhs, int rhs) { |
| return true; |
| } |
| bool inferInputOutputShapeEquality(int input, int output) { |
| return true; |
| } |
| llvm::Optional<Value> inferEffectiveWorkloadShape() { |
| return getOperation()->getResult(0); |
| } |
| }]; |
| |
| let results = (outs HLO_Tensor); |
| let parser = [{ return mlir::impl::parseOneResultSameOperandTypeOp(parser, result); }]; |
| let printer = [{ return mlir::impl::printOneResultOp(getOperation(), p); }]; |
| } |
| |
| def HLO_AddOp : HLO_BinaryElementwiseOp<"add", |
| [Commutative, NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_AddOp { |
| let hasFolder = 1; |
| } |
| |
| def HLO_Atan2Op : HLO_BinaryElementwiseOp<"atan2", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_Atan2Op; |
| |
| def HLO_ComplexOp: HLO_Op<"complex", |
| [NoSideEffect, SameOperandsAndResultShape]>, |
| BASE_HLO_ComplexOp { |
| let builders = [OpBuilder< |
| "OpBuilder &, OperationState &tblgen_state, Value lhs, Value rhs">]; |
| |
| let arguments = (ins HLO_FpTensor:$lhs, HLO_FpTensor:$rhs); |
| let results = (outs HLO_ComplexTensor); |
| let hasFolder = 1; |
| } |
| |
| def HLO_DivOp : HLO_BinaryElementwiseOp<"divide", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_DivOp { |
| let hasFolder = 1; |
| } |
| |
| def HLO_MaxOp : HLO_BinaryElementwiseOp<"maximum", |
| [Commutative, NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_MaxOp { |
| let hasFolder = 1; |
| } |
| |
| def HLO_MinOp : HLO_BinaryElementwiseOp<"minimum", |
| [Commutative, NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_MinOp { |
| let hasFolder = 1; |
| } |
| |
| def HLO_MulOp : HLO_BinaryElementwiseOp<"multiply", |
| [Commutative, NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_MulOp { |
| let hasFolder = 1; |
| } |
| |
| def HLO_PowOp : HLO_BinaryElementwiseOp<"power", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_PowOp; |
| |
| def HLO_RemOp : HLO_BinaryElementwiseOp<"remainder", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_RemOp; |
| |
| def HLO_ShiftLeftOp : HLO_BinaryElementwiseOp<"shift_left", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_ShiftLeftOp; |
| |
| def HLO_ShiftRightArithmeticOp : HLO_BinaryElementwiseOp<"shift_right_arithmetic", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_ShiftRightArithmeticOp; |
| |
| def HLO_ShiftRightLogicalOp : HLO_BinaryElementwiseOp<"shift_right_logical", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_ShiftRightLogicalOp; |
| |
| def HLO_SubOp : HLO_BinaryElementwiseOp<"subtract", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_SubOp { |
| let hasFolder = 1; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO binary logical elementwise op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| // See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations |
| class HLO_BinaryLogicalElementwiseOp<string mnemonic> : |
| HLO_BinaryElementwiseOp< |
| mnemonic, [Commutative, NoSideEffect, SameOperandsAndResultType]> { |
| let arguments = (ins |
| HLO_PredOrIntTensor:$lhs, |
| HLO_PredOrIntTensor:$rhs |
| ); |
| } |
| |
| def HLO_AndOp: HLO_BinaryLogicalElementwiseOp<"and">, BASE_HLO_AndOp; |
| def HLO_OrOp: HLO_BinaryLogicalElementwiseOp<"or">, BASE_HLO_OrOp; |
| def HLO_XorOp : HLO_BinaryLogicalElementwiseOp<"xor">, BASE_HLO_XorOp; |
| |
| //===----------------------------------------------------------------------===// |
| // 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", []> { |
| |
| string summary = "Infeed operator"; |
| |
| string 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. |
| |
| See https://www.tensorflow.org/xla/operation_semantics#infeed. |
| }]; |
| |
| let arguments = (ins |
| HLO_Token:$token, |
| DefaultValuedAttr<StrAttr, "">:$infeed_config |
| ); |
| let results = (outs HLO_Tuple); |
| let hasCustomHLOConverter = 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", []> { |
| |
| string summary = "Outfeed operator"; |
| |
| string 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 |
| HLO_TensorOrTuple:$operand, |
| HLO_Token:$token, |
| DefaultValuedAttr<StrAttr, "">:$outfeed_config |
| ); |
| let results = (outs HLO_Token); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_SendOp : HLO_Op<"send", []> { |
| |
| string summary = "Send operator"; |
| |
| string 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 |
| HLO_TensorOrTuple:$operand, |
| HLO_Token:$token, |
| ChannelHandle<HLO_Dialect>:$channel_id, |
| DefaultValuedAttr<BoolAttr, "false">:$is_host_transfer |
| ); |
| |
| let results = (outs HLO_Token); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_RecvOp : HLO_Op<"recv", []> { |
| |
| string summary = "Recv operator"; |
| |
| string 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<HLO_Dialect>:$channel_id, |
| DefaultValuedAttr<BoolAttr, "false">:$is_host_transfer |
| ); |
| |
| let results = (outs HLO_Tuple); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO parallelism related op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_ReplicaIdOp : HLO_Op<"replica_id", [NoSideEffect]>, |
| BASE_HLO_ReplicaIdOp { |
| // TODO(prakalps): The output should unsigned 32-bit integer but mlir does |
| // not differentiate between signed and unsigned int. |
| let results = (outs I32Tensor); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO control flow op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_AfterAllOp : HLO_Op<"after_all", [NoSideEffect]> { |
| |
| string summary = "AfterAll operator"; |
| |
| string 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. |
| }]; |
| |
| let arguments = (ins Variadic<HLO_Token>:$operands); |
| let results = (outs HLO_Token); |
| } |
| |
| // 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]> { |
| string summary = "If operator"; |
| |
| string description = [{ |
| Returns the result of executing either a true or false function depending on |
| the result of a condition function. |
| |
| See https://www.tensorflow.org/xla/operation_semantics#conditional. |
| }]; |
| |
| let arguments = (ins |
| HLO_PredTensor:$pred, |
| HLO_TensorOrTuple:$true_arg, |
| HLO_TensorOrTuple:$false_arg |
| ); |
| |
| let regions = (region AnyRegion:$true_branch, |
| AnyRegion:$false_branch); |
| |
| let results = (outs HLO_TensorOrTuple); |
| |
| // TODO(b/129422361): ConditionalOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 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]>, |
| BASE_HLO_CaseOp { |
| |
| let arguments = (ins |
| I32Tensor:$index, |
| Variadic<HLO_TensorOrTuple>:$branch_operands |
| ); |
| |
| let regions = (region VariadicRegion<AnyRegion>:$branches); |
| |
| let results = (outs Variadic<HLO_TensorOrTuple>); |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| |
| def HLO_WhileOp: HLO_Op<"while", [RecursiveSideEffects, |
| SameOperandsAndResultType]>, |
| BASE_HLO_WhileOp { |
| let arguments = (ins HLO_TensorOrTuple:$val); |
| |
| let regions = (region AnyRegion:$cond, AnyRegion:$body); |
| |
| let results = (outs HLO_TensorOrTuple); |
| |
| // TODO(b/129422361): WhileOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_AllReduceOp : HLO_Op<"all_reduce", |
| [SameOperandsAndResultType]>, BASE_HLO_AllReduceOp { |
| |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I64ElementsAttr:$replica_groups, |
| OptionalAttr<ChannelHandle<HLO_Dialect>>:$channel_id |
| ); |
| let regions = (region SizedRegion<1>:$computation); |
| let results = (outs HLO_Tensor); |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_AllToAllOp : HLO_Op<"all_to_all", |
| [NoSideEffect, SameOperandsElementType, SameOperandsShape]>, BASE_HLO_AllToAllOp { |
| |
| 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_Op<"reduce", [ |
| RecursiveSideEffects, |
| SameVariadicOperandSize, |
| SingleBlockImplicitTerminator<"ReturnOp">, |
| InferFusibilityOpInterface |
| ]>, BASE_HLO_ReduceOp { |
| let arguments = (ins |
| Variadic<HLO_TensorOrTuple>:$operands, |
| Variadic<HLO_TensorOrTuple>:$init_values, |
| I64ElementsAttr:$dimensions |
| ); |
| |
| let results = (outs Variadic<HLO_TensorOrTuple>); |
| |
| let builders = [OpBuilder< |
| "OpBuilder &, OperationState &state, ValueRange operands, " |
| "ValueRange init_values, DenseIntElementsAttr dimensions" |
| >]; |
| |
| let extraClassDeclaration = [{ |
| bool isFusibleWithConsumer() { |
| return false; |
| } |
| llvm::Optional<Value> inferEffectiveWorkloadShape() { |
| return getOperation()->getOperand(0); |
| } |
| }]; |
| |
| let hasFolder = 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(hinsu): Implement custom printer and parser. |
| |
| // 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]>, BASE_HLO_GetTupleElementOp { |
| let arguments = (ins |
| HLO_Tuple, |
| I32Attr:$index |
| ); |
| |
| let results = (outs HLO_TensorOrTokenOrTuple); |
| |
| let hasFolder = 1; |
| |
| let builders = [OpBuilder< |
| "OpBuilder &builder, OperationState &results, " |
| "Value value, int32_t index">]; |
| } |
| |
| def HLO_TupleOp : HLO_Op<"tuple", [NoSideEffect]>, BASE_HLO_TupleOp { |
| let arguments = (ins Variadic<HLO_TensorOrTuple>:$val); |
| let results = (outs HLO_Tuple); |
| |
| let builders = [OpBuilder< |
| "OpBuilder &builder, OperationState &results, " |
| "ValueRange values">]; |
| |
| } |
| |
| def HLO_CompareOp: HLO_Op<"compare", |
| [NoSideEffect, SameTypeOperands, SameOperandsAndResultShape]>, |
| BASE_HLO_CompareOp { |
| let arguments = (ins |
| HLO_Tensor:$lhs, |
| HLO_Tensor:$rhs, |
| HLO_ComparisonDirectionAttr:$comparison_direction |
| ); |
| let results = (outs HLO_PredTensor); |
| |
| let builders = [OpBuilder< |
| "OpBuilder &builder, OperationState &result, Value lhs, Value rhs, " |
| "StringAttr comparison_direction" |
| >]; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO Slice definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_SliceOp: HLO_Op< |
| "slice", |
| [NoSideEffect, SameOperandsAndResultElementType, |
| AllTypesMatch<["start_indices", "limit_indices", "strides"]>]> { |
| 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; |
| |
| let builders = [OpBuilder< |
| "OpBuilder &builder, OperationState &result, Value operand, " |
| "DenseIntElementsAttr start_indices, DenseIntElementsAttr limit_indices, " |
| "DenseIntElementsAttr strides" |
| >]; |
| |
| let extraClassDeclaration = [{ |
| // Infers output type for given operand and attributes. Result type is |
| // unranked if any of the attributes is illegal. |
| static Type InferOutputTypes(Builder *builder, Value operand, |
| DenseIntElementsAttr start_indices, |
| DenseIntElementsAttr limit_indices, |
| DenseIntElementsAttr strides); |
| }]; |
| } |
| |
| def HLO_DynamicSliceOp: HLO_Op<"dynamic-slice", |
| [NoSideEffect, AllElementTypesMatch<["operand", "result"]>]> { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| Variadic<HLO_ScalarIntTensor>:$start_indices, |
| I64ElementsAttr:$slice_sizes |
| ); |
| |
| let results = (outs HLO_Tensor:$result); |
| let hasCanonicalizer = 1; |
| } |
| |
| def HLO_DynamicUpdateSliceOp: HLO_Op<"dynamic-update-slice", |
| [NoSideEffect, AllElementTypesMatch<["operand", "update", "result"]>, |
| AllShapesMatch<["operand", "result"]>]> { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_Tensor:$update, |
| Variadic<HLO_ScalarIntTensor>:$start_indices |
| ); |
| |
| let results = (outs HLO_Tensor:$result); |
| } |
| |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO Other op definitions. |
| //===----------------------------------------------------------------------===// |
| |
| def HLO_BatchNormGradOp : HLO_Op<"batch_norm_grad", [NoSideEffect]>, |
| BASE_HLO_BatchNormGradOp { |
| |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_Tensor:$scale, |
| HLO_Tensor:$mean, |
| HLO_Tensor:$variance, |
| HLO_Tensor:$grad_output, |
| F32Attr:$epsilon, |
| I64Attr:$feature_index |
| ); |
| |
| let results = (outs HLO_Tuple); |
| } |
| |
| def HLO_BatchNormInferenceOp : HLO_Op<"batch_norm_inference", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_BatchNormInferenceOp { |
| |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_Tensor:$scale, |
| HLO_Tensor:$offset, |
| HLO_Tensor:$mean, |
| HLO_Tensor:$variance, |
| F32Attr:$epsilon, |
| I64Attr:$feature_index |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_BatchNormTrainingOp : HLO_Op<"batch_norm_training", [NoSideEffect]>, |
| BASE_HLO_BatchNormTrainingOp { |
| |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_Tensor:$scale, |
| HLO_Tensor:$offset, |
| F32Attr:$epsilon, |
| I64Attr:$feature_index |
| ); |
| |
| let results = (outs HLO_Tuple); |
| } |
| |
| def HLO_BitcastConvertOp : HLO_Op<"bitcast_convert", |
| [NoSideEffect, SameOperandsAndResultShape]>, BASE_HLO_BitcastConvertOp { |
| |
| let arguments = (ins HLO_Tensor:$operand); |
| let results = (outs HLO_Tensor); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_BroadcastOp : HLO_Op<"broadcast", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_BroadcastOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I64ElementsAttr:$broadcast_sizes |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_BroadcastInDimOp : HLO_Op<"broadcast_in_dim", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_BroadcastInDimOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| BroadcastDimAttr:$broadcast_dimensions |
| ); |
| |
| let results = (outs HLO_StaticShapeTensor); |
| |
| let hasFolder = 1; |
| // Only handles a static subset of the legacy format. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_DynamicBroadcastInDimOp : HLO_Op<"dynamic_broadcast_in_dim", |
| [NoSideEffect]> { |
| string summary = "Broadcast a tensor into the given dynamic shape by adding dimensions."; |
| string 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. |
| }]; |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_DimensionTensor:$output_dimensions, |
| BroadcastDimAttr:$broadcast_dimensions |
| ); |
| |
| let results = (outs HLO_Tensor); |
| |
| let hasCanonicalizer = 1; |
| // Cannot be exported to legacy formats. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| // Note: There is no HLO_CallOp because the standard call operation mlir::CallOp |
| // is used instead. A mlir::CallOp is exported to a HLO call instruction |
| // directly. |
| |
| def HLO_CholeskyOp : HLO_Op<"cholesky", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_CholeskyOp { |
| let arguments = (ins |
| HLO_FpOrComplexTensor:$a, |
| DefaultValuedAttr<BoolAttr, "false">:$lower |
| ); |
| |
| let results = (outs HLO_FpOrComplexTensor); |
| } |
| |
| def HLO_ClampOp : HLO_Op<"clamp", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_ClampOp { |
| let arguments = (ins |
| HLO_Tensor:$min, |
| HLO_Tensor:$operand, |
| HLO_Tensor:$max |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_ConcatenateOp : HLO_Op<"concatenate", |
| [NoSideEffect, SameOperandsAndResultElementType, DeclareOpInterfaceMethods<InferTypeOpInterface>]>, BASE_HLO_ConcatenateOp { |
| |
| let arguments = (ins |
| Variadic<HLO_Tensor>:$val, |
| I64Attr: $dimension |
| ); |
| |
| let results = (outs HLO_Tensor); |
| |
| let hasCanonicalizer = 1; |
| let hasFolder = 1; |
| |
| } |
| |
| def HLO_CollectivePermuteOp: HLO_Op<"collective_permute", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_CollectivePermuteOp { |
| |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I64ElementsAttr:$source_target_pairs |
| ); |
| let results = (outs HLO_Tensor); |
| } |
| |
| // TODO(hinsu): Make this struct dialect independent so that it can be shared |
| // between HLO and LHLO dialect. |
| def ConvDimensionNumbers : StructAttr<"ConvDimensionNumbers", HLO_Dialect, [ |
| StructFieldAttr<"input_batch_dimension",I64Attr>, |
| StructFieldAttr<"input_feature_dimension", I64Attr>, |
| StructFieldAttr<"input_spatial_dimensions", I64ElementsAttr>, |
| StructFieldAttr<"kernel_input_feature_dimension", I64Attr>, |
| StructFieldAttr<"kernel_output_feature_dimension", I64Attr>, |
| StructFieldAttr<"kernel_spatial_dimensions", I64ElementsAttr>, |
| StructFieldAttr<"output_batch_dimension", I64Attr>, |
| StructFieldAttr<"output_feature_dimension", I64Attr>, |
| StructFieldAttr<"output_spatial_dimensions", I64ElementsAttr>] > { |
| |
| let description = "Structure of dimension information for conv op"; |
| } |
| |
| def HLO_ConvOp : HLO_Op<"convolution", [NoSideEffect]>, BASE_HLO_ConvOp { |
| let arguments = (ins |
| HLO_Tensor:$lhs, |
| HLO_Tensor:$rhs, |
| // Default value: one for each of the spatial dimension. |
| OptionalAttr<I64ElementsAttr>:$window_strides, |
| // Default value: zero for each of the spatial dimension. |
| OptionalAttr<I64ElementsAttr>:$padding, |
| // Default value: one for each of the spatial dimension. |
| OptionalAttr<I64ElementsAttr>:$lhs_dilation, |
| // Default value: one for each of the spatial dimension. |
| OptionalAttr<I64ElementsAttr>:$rhs_dilation, |
| ConvDimensionNumbers:$dimension_numbers, |
| I64Attr:$feature_group_count, |
| I64Attr:$batch_group_count, |
| HLO_PrecisionConfigAttr:$precision_config |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_CopyOp: HLO_Op<"copy", [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_CopyOp { |
| let arguments = (ins HLO_Tensor); |
| let results = (outs HLO_Tensor); |
| let hasFolder = 1; |
| } |
| |
| def HLO_CrossReplicaSumOp : HLO_Op<"cross-replica-sum", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_CrossReplicaSumOp { |
| |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I64ElementsAttr:$replica_groups |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_CustomCallOp: HLO_Op<"custom_call", []>, BASE_HLO_CustomCallOp { |
| let arguments = (ins |
| Variadic<HLO_Tensor>:$args, |
| StrAttr:$call_target_name, |
| DefaultValuedAttr<BoolAttr, "false">:$has_side_effect, |
| DefaultValuedAttr<StrAttr, "">:$backend_config |
| ); |
| let results = (outs HLO_Tensor); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_DotOp: HLO_Op<"dot", [NoSideEffect]>, BASE_HLO_DotOp { |
| let arguments = ( |
| ins HLO_Tensor:$lhs, |
| HLO_Tensor:$rhs, |
| HLO_PrecisionConfigAttr:$precision_config |
| ); |
| let results = (outs HLO_Tensor); |
| } |
| |
| def DotDimensionNumbers : StructAttr<"DotDimensionNumbers", HLO_Dialect, [ |
| StructFieldAttr<"lhs_batching_dimensions", I64ElementsAttr>, |
| StructFieldAttr<"rhs_batching_dimensions", I64ElementsAttr>, |
| StructFieldAttr<"lhs_contracting_dimensions", I64ElementsAttr>, |
| StructFieldAttr<"rhs_contracting_dimensions", I64ElementsAttr> |
| ]> { |
| let description = "Structure of dimension information for dot product"; |
| } |
| |
| def HLO_DotGeneralOp: HLO_Op<"dot_general", [NoSideEffect]>, BASE_HLO_DotGeneralOp { |
| let arguments = (ins |
| HLO_Tensor:$lhs, |
| HLO_Tensor:$rhs, |
| DotDimensionNumbers:$dot_dimension_numbers, |
| HLO_PrecisionConfigAttr:$precision_config |
| ); |
| |
| let results = (outs HLO_Tensor); |
| let verifier = [{ return Verify(*this); }]; |
| } |
| |
| // 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", [NoSideEffect]>, BASE_HLO_FftOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_FftTypeAttr: $fft_type, |
| I64ElementsAttr:$fft_length |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def GatherDimensionNumbers : StructAttr<"GatherDimensionNumbers", HLO_Dialect, |
| [StructFieldAttr<"offset_dims", I64ElementsAttr>, |
| StructFieldAttr<"collapsed_slice_dims", I64ElementsAttr>, |
| StructFieldAttr<"start_index_map", I64ElementsAttr>, |
| StructFieldAttr<"index_vector_dim", I64Attr>]> { |
| let description = "Structure of dimension information for gather"; |
| } |
| |
| def HLO_GatherOp: HLO_Op<"gather", [NoSideEffect]>, BASE_HLO_GatherOp { |
| 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); |
| } |
| |
| def HLO_GetDimensionSizeOp: HLO_Op<"get_dimension_size", [NoSideEffect]>, |
| BASE_HLO_GetDimensionSizeOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I32Attr:$dimension |
| ); |
| let results = (outs HLO_IntTensor); |
| } |
| |
| def HLO_MapOp: HLO_Op<"map", |
| [RecursiveSideEffects, SameOperandsElementType, |
| SameOperandsAndResultShape, SingleBlockImplicitTerminator<"ReturnOp">]>, |
| BASE_HLO_MapOp { |
| let arguments = (ins |
| Variadic<HLO_Tensor>:$operands, |
| I64ElementsAttr:$dimensions |
| ); |
| let regions = (region SizedRegion<1>:$computation); |
| let results = (outs HLO_Tensor); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_ReshapeOp: HLO_Op<"reshape", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_ReshapeOp { |
| let arguments = (ins HLO_Tensor:$operand); |
| |
| let results = (outs HLO_StaticShapeTensor); |
| let hasFolder = 1; |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_DynamicReshapeOp: HLO_Op<"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`). |
| }]; |
| |
| 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; |
| } |
| |
| def HLO_ScatterOp: HLO_Op<"scatter", [RecursiveSideEffects]>, |
| BASE_HLO_ScatterOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_Tensor:$scatter_indices, |
| HLO_Tensor:$updates, |
| ScatterDimensionNumbers<HLO_Dialect>:$scatter_dimension_numbers, |
| DefaultValuedAttr<BoolAttr, "false">:$indices_are_sorted, |
| DefaultValuedAttr<BoolAttr, "false">:$unique_indices |
| ); |
| |
| let regions = (region SizedRegion<1>:$update_computation); |
| |
| let results = (outs HLO_Tensor); |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| // TODO(jpienaar): Add broadcastable trait. |
| def HLO_SelectOp: HLO_Op<"select", [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]>, BASE_HLO_SelectOp { |
| let arguments = (ins |
| HLO_PredTensor:$pred, |
| HLO_Tensor:$on_true, |
| HLO_Tensor:$on_false |
| ); |
| |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_SelectAndScatterOp: HLO_Op<"select_and_scatter", |
| [RecursiveSideEffects]>, BASE_HLO_SelectAndScatterOp { |
| 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 hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_SetDimensionSizeOp: HLO_Op<"set_dimension_size", [NoSideEffect]>, |
| BASE_HLO_SetDimensionSizeOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I32Tensor:$size, |
| I32Attr:$dimension |
| ); |
| let results = (outs HLO_Tensor); |
| } |
| |
| def HLO_SortOp : HLO_Op<"sort", [RecursiveSideEffects]>, BASE_HLO_SortOp { |
| let arguments = (ins |
| Variadic<HLO_Tensor>:$operands, |
| DefaultValuedAttr<I64Attr, "-1">:$dimension, |
| DefaultValuedAttr<BoolAttr, "false">:$is_stable |
| ); |
| |
| let results = (outs HLO_TensorOrTuple); |
| |
| let regions = (region SizedRegion<1>:$comparator); |
| |
| let builders = [OpBuilder< |
| "OpBuilder &builder, OperationState &state, ValueRange operands, " |
| "int64_t dimension = -1, bool is_stable = false" |
| >]; |
| |
| // TODO(b/129422361): SortOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_ReverseOp: HLO_Op<"reverse", |
| [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_ReverseOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I64ElementsAttr:$dimensions |
| ); |
| |
| let results = (outs HLO_Tensor); |
| |
| let hasFolder = 1; |
| } |
| |
| def HLO_PadOp: HLO_Op<"pad", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_PadOp { |
| 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); |
| |
| let description = [{ |
| Pads the `operand` according to TBD. |
| }]; |
| |
| // TODO(b/129422361): PadOp has a custom constructor for HLO. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_TraceOp: HLO_Op<"trace", []>, BASE_HLO_TraceOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| StrAttr:$tag |
| ); |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_TransposeOp: HLO_Op<"transpose", |
| [NoSideEffect, SameOperandsAndResultElementType]>, BASE_HLO_TransposeOp { |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| I64ElementsAttr:$permutation |
| ); |
| let results = (outs HLO_Tensor); |
| |
| let hasFolder = 1; |
| } |
| |
| def HLO_TriangularSolveOp: HLO_Op<"triangular_solve", |
| [NoSideEffect, SameOperandsAndResultElementType]>, |
| BASE_HLO_TriangularSolveOp { |
| 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); |
| } |
| |
| def HLO_ReduceWindowOp: HLO_Op<"reduce_window", [ |
| RecursiveSideEffects, |
| SingleBlockImplicitTerminator<"ReturnOp"> |
| ]>, BASE_HLO_ReduceWindowOp { |
| |
| // TODO(hinsu): Verify that padding attribute is 2-d and the remaining |
| // attributes are 1-d. Attributes' leading dimension should match rank of the |
| // inputs. |
| let arguments = (ins |
| HLO_Tensor:$operand, |
| HLO_Tensor:$init_value, |
| I64ElementsAttr:$window_dimensions, |
| // If strides or dilations attributes are missing then the default value is |
| // one for each of the input 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 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); |
| |
| let hasCustomHLOConverter = 1; |
| |
| // TODO(hinsu): Implement custom printer and parser. |
| } |
| |
| def HLO_ReturnOp : HLO_Op<"return", [NoSideEffect, Terminator]> { |
| let summary = [{ |
| The `hlo.return` operation terminates a region and returns values. |
| }]; |
| |
| let arguments = (ins |
| Variadic<HLO_TensorOrTuple>:$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; |
| |
| // TODO(hinsu): Implement custom printer and parser. |
| } |
| |
| def HLO_TorchIndexSelectOp : HLO_Op<"torch_index_select", [NoSideEffect]> { |
| let arguments = (ins |
| HLO_Tensor:$input, |
| 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. |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO RngUniform Operator. |
| //===----------------------------------------------------------------------===// |
| def HLO_RngUniformOp : HLO_Op<"rng_uniform", []>, BASE_HLO_RngUniformOp { |
| let arguments = (ins |
| HLO_PredIntOrFpTensor:$a, |
| HLO_PredIntOrFpTensor:$b, |
| HLO_DimensionTensor:$shape |
| ); |
| |
| let results = (outs HLO_PredIntOrFpTensor); |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| def HLO_RngNormalOp : HLO_Op<"rng_normal", []>, BASE_HLO_RngNormalOp { |
| let arguments = (ins |
| HLO_FpTensor:$mu, |
| HLO_FpTensor:$sigma, |
| HLO_DimensionTensor:$shape |
| ); |
| |
| let results = (outs HLO_FpTensor); |
| |
| let hasCustomHLOConverter = 1; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // MHLO Quantize Operator. |
| //===----------------------------------------------------------------------===// |
| def HLO_DequantizeOp : HLO_Op<"dequantize", [NoSideEffect]>, |
| BASE_HLO_DequantizeOp { |
| let arguments = (ins |
| TensorOf<[I32]>:$input, |
| F32Attr:$min_range, |
| F32Attr:$max_range, |
| HLO_DequantizeModeAttr:$mode, |
| BoolAttr:$transpose_output, |
| DefaultValuedAttr<BoolAttr, "false">:$is_16bits |
| ); |
| |
| let results = (outs TensorOf<[BF16]>:$output); |
| |
| 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_TensorOrTuple>:$operands |
| ); |
| |
| let results = (outs |
| Variadic<HLO_TensorOrTuple>:$results |
| ); |
| |
| // FusionOp has special conversion logic to HLO. |
| let hasCustomHLOConverter = 1; |
| } |
| |
| #endif // HLO_OPS |