1/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 Copyright 2022 The StableHLO Authors. 3 4Licensed under the Apache License, Version 2.0 (the "License"); 5you may not use this file except in compliance with the License. 6You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10Unless required by applicable law or agreed to in writing, software 11distributed under the License is distributed on an "AS IS" BASIS, 12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13See the License for the specific language governing permissions and 14limitations under the License. 15==============================================================================*/ 16 17#ifndef STABLEHLO_DIALECT_STABLEHLO_ENUMS 18#define STABLEHLO_DIALECT_STABLEHLO_ENUMS 19 20include "mlir/IR/EnumAttr.td" 21include "mlir/IR/PatternBase.td" 22 23//===----------------------------------------------------------------------===// 24// Precision Config enum definitions. 25//===----------------------------------------------------------------------===// 26 27// These mirror the XLA PrecisionConfig proto enum. 28def STABLEHLO_PRECISION_DEFAULT : I32EnumAttrCase<"DEFAULT", 0>; 29def STABLEHLO_PRECISION_HIGH : I32EnumAttrCase<"HIGH", 1>; 30def STABLEHLO_PRECISION_HIGHEST : I32EnumAttrCase<"HIGHEST", 2>; 31 32def StableHLO_Precision : I32EnumAttr<"Precision", 33 "XLA precision for an operand. Has backend specific meaning.", 34 [ 35 STABLEHLO_PRECISION_DEFAULT, 36 STABLEHLO_PRECISION_HIGH, 37 STABLEHLO_PRECISION_HIGHEST 38 ]> { 39 let genSpecializedAttr = 0; 40 let cppNamespace = "::mlir::stablehlo"; 41} 42 43def StableHLO_PrecisionAttr : EnumAttr<StableHLO_Dialect, StableHLO_Precision, "precision">; 44 45// TODO(b/129153247) See if it's possible to also validate the size. 46def StableHLO_PrecisionConfigAttr: 47 OptionalAttr< 48 TypedArrayAttrBase<StableHLO_PrecisionAttr, "Precision Config attribute">>; 49 50//===----------------------------------------------------------------------===// 51// Fast Fourier Transform Type enum definitions. 52//===----------------------------------------------------------------------===// 53 54// These mirror the XLA FftType proto enum. 55def STABLEHLO_FFT_TYPE_FFT : I32EnumAttrCase<"FFT", 0>; 56def STABLEHLO_FFT_TYPE_IFFT : I32EnumAttrCase<"IFFT", 1>; 57def STABLEHLO_FFT_TYPE_RFFT : I32EnumAttrCase<"RFFT", 2>; 58def STABLEHLO_FFT_TYPE_IRFFT : I32EnumAttrCase<"IRFFT", 3>; 59 60def StableHLO_FftType : I32EnumAttr<"FftType", 61 "XLA fast fourier transform type.", 62 [ 63 STABLEHLO_FFT_TYPE_FFT, 64 STABLEHLO_FFT_TYPE_IFFT, 65 STABLEHLO_FFT_TYPE_RFFT, 66 STABLEHLO_FFT_TYPE_IRFFT 67 ]> { 68 let genSpecializedAttr = 0; 69 let cppNamespace = "::mlir::stablehlo"; 70} 71 72def StableHLO_FftTypeAttr : EnumAttr<StableHLO_Dialect, StableHLO_FftType, "fft_type">; 73 74//===----------------------------------------------------------------------===// 75// Custom call enum definitions. 76//===----------------------------------------------------------------------===// 77 78// TODO(b/189822916): Remove this enum when all clients are migrated to the 79// status-returning API. 80def STABLEHLO_CUSTOM_CALL_API_VERISON_UNSPECIFIED : 81 I32EnumAttrCase<"API_VERSION_UNSPECIFIED", 0>; 82def STABLEHLO_CUSTOM_CALL_API_VERSION_ORIGINAL : 83 I32EnumAttrCase<"API_VERSION_ORIGINAL", 1>; 84def STABLEHLO_CUSTOM_CALL_API_VERSION_STATUS_RETURNING : 85 I32EnumAttrCase<"API_VERSION_STATUS_RETURNING", 2>; 86def STABLEHLO_CUSTOM_CALL_API_VERSION_STATUS_RETURNING_UNIFIED : 87 I32EnumAttrCase<"API_VERSION_STATUS_RETURNING_UNIFIED", 3>; 88def StableHLO_CustomCallApiVersionAttr : 89 I32EnumAttr<"CustomCallApiVersion", "Custom call API version", [ 90 STABLEHLO_CUSTOM_CALL_API_VERISON_UNSPECIFIED, 91 STABLEHLO_CUSTOM_CALL_API_VERSION_ORIGINAL, 92 STABLEHLO_CUSTOM_CALL_API_VERSION_STATUS_RETURNING, 93 STABLEHLO_CUSTOM_CALL_API_VERSION_STATUS_RETURNING_UNIFIED 94 ]> { 95 let cppNamespace = "::mlir::stablehlo"; 96} 97 98//===----------------------------------------------------------------------===// 99// Comparison op definitions. 100//===----------------------------------------------------------------------===// 101 102// These mirror the XLA ComparisonDirection enum. 103def STABLEHLO_COMPARISON_DIRECTION_EQ : I32EnumAttrCase<"EQ", 0>; 104def STABLEHLO_COMPARISON_DIRECTION_NE : I32EnumAttrCase<"NE", 1>; 105def STABLEHLO_COMPARISON_DIRECTION_GE : I32EnumAttrCase<"GE", 2>; 106def STABLEHLO_COMPARISON_DIRECTION_GT : I32EnumAttrCase<"GT", 3>; 107def STABLEHLO_COMPARISON_DIRECTION_LE : I32EnumAttrCase<"LE", 4>; 108def STABLEHLO_COMPARISON_DIRECTION_LT : I32EnumAttrCase<"LT", 5>; 109 110def StableHLO_ComparisonDirection : I32EnumAttr<"ComparisonDirection", 111 "Which comparison operation to perform.", 112 [ 113 STABLEHLO_COMPARISON_DIRECTION_EQ, 114 STABLEHLO_COMPARISON_DIRECTION_NE, 115 STABLEHLO_COMPARISON_DIRECTION_GE, 116 STABLEHLO_COMPARISON_DIRECTION_GT, 117 STABLEHLO_COMPARISON_DIRECTION_LE, 118 STABLEHLO_COMPARISON_DIRECTION_LT 119 ]> { 120 let genSpecializedAttr = 0; 121 let cppNamespace = "::mlir::stablehlo"; 122} 123 124def StableHLO_ComparisonDirectionAttr : EnumAttr<StableHLO_Dialect, StableHLO_ComparisonDirection, "comparison_direction">; 125 126def STABLEHLO_DEFAULT_COMPARISON_TYPE : NativeCodeCall<"::mlir::stablehlo::ComparisonTypeAttr()">; 127def STABLEHLO_COMPARISON_TYPE_NOTYPE : I32EnumAttrCase<"NOTYPE", 0>; 128def STABLEHLO_COMPARISON_TYPE_FLOAT : I32EnumAttrCase<"FLOAT", 1>; 129def STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>; 130def STABLEHLO_COMPARISON_TYPE_SIGNED : I32EnumAttrCase<"SIGNED", 3>; 131def STABLEHLO_COMPARISON_TYPE_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>; 132 133def StableHLO_ComparisonType : I32EnumAttr<"ComparisonType", 134 "Which comparison type to use.", 135 [ 136 STABLEHLO_COMPARISON_TYPE_NOTYPE, 137 STABLEHLO_COMPARISON_TYPE_FLOAT, 138 STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER, 139 STABLEHLO_COMPARISON_TYPE_SIGNED, 140 STABLEHLO_COMPARISON_TYPE_UNSIGNED 141 ]> { 142 let genSpecializedAttr = 0; 143 let cppNamespace = "::mlir::stablehlo"; 144} 145 146def StableHLO_ComparisonTypeAttr : EnumAttr<StableHLO_Dialect, StableHLO_ComparisonType, "comparison_type">; 147 148// These mirror the XLA Transpose enum in Triangular Solve options. 149def STABLEHLO_TRANSPOSE_INVALID : I32EnumAttrCase<"TRANSPOSE_INVALID", 0>; 150def STABLEHLO_NO_TRANSPOSE : I32EnumAttrCase<"NO_TRANSPOSE", 1>; 151def STABLEHLO_TRANSPOSE : I32EnumAttrCase<"TRANSPOSE", 2>; 152def STABLEHLO_ADJOINT : I32EnumAttrCase<"ADJOINT", 3>; 153 154def StableHLO_Transpose : I32EnumAttr<"Transpose", 155 "Transpose options", 156 [ 157 STABLEHLO_TRANSPOSE_INVALID, 158 STABLEHLO_NO_TRANSPOSE, 159 STABLEHLO_TRANSPOSE, 160 STABLEHLO_ADJOINT 161 ]> { 162 let genSpecializedAttr = 0; 163 let cppNamespace = "::mlir::stablehlo"; 164} 165 166def StableHLO_TransposeAttr : EnumAttr<StableHLO_Dialect, StableHLO_Transpose, "transpose">; 167 168def STABLEHLO_RNG_DISTRIBUTION_UNIFORM : I32EnumAttrCase<"UNIFORM", 1>; 169def STABLEHLO_RNG_DISTRIBUTION_NORMAL : I32EnumAttrCase<"NORMAL", 2>; 170 171def StableHLO_RngDistribution : I32EnumAttr<"RngDistribution", 172 "XLA PRNG distribution to be used.", 173 [ 174 STABLEHLO_RNG_DISTRIBUTION_UNIFORM, 175 STABLEHLO_RNG_DISTRIBUTION_NORMAL 176 ]> { 177 let genSpecializedAttr = 0; 178 let cppNamespace = "::mlir::stablehlo"; 179} 180 181def StableHLO_RngDistributionAttr : EnumAttr<StableHLO_Dialect, StableHLO_RngDistribution, "rng_distribution"> { 182 let assemblyFormat = "`<` $value `>`"; 183} 184 185def STABLEHLO_RNG_ALGORITHM_DEFAULT : I32EnumAttrCase<"DEFAULT", 0>; 186def STABLEHLO_RNG_ALGORITHM_THREE_FRY : I32EnumAttrCase<"THREE_FRY", 1>; 187def STABLEHLO_RNG_ALGORITHM_PHILOX : I32EnumAttrCase<"PHILOX", 2>; 188 189def StableHLO_RngAlgorithm : I32EnumAttr<"RngAlgorithm", 190 "XLA PRNG algorithm to be used.", 191 [ 192 STABLEHLO_RNG_ALGORITHM_DEFAULT, 193 STABLEHLO_RNG_ALGORITHM_THREE_FRY, 194 STABLEHLO_RNG_ALGORITHM_PHILOX 195 ]> { 196 let genSpecializedAttr = 0; 197 let cppNamespace = "::mlir::stablehlo"; 198} 199 200def StableHLO_RngAlgorithmAttr : EnumAttr<StableHLO_Dialect, StableHLO_RngAlgorithm, "rng_algorithm"> { 201 let assemblyFormat = "`<` $value `>`"; 202} 203 204#endif // STABLEHLO_DIALECT_STABLEHLO_ENUMS 205