1 /** 2 * Copyright 2023 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #ifndef MINDSPORE_CCSRC_PIPELINE_PYNATIVE_PREDICTOUTTYPEMAP_H_ 17 #define MINDSPORE_CCSRC_PIPELINE_PYNATIVE_PREDICTOUTTYPEMAP_H_ 18 19 #include <string> 20 #include <memory> 21 #include <vector> 22 23 #include "utils/hash_map.h" 24 #include "ops/base_operator.h" 25 #include "base/base.h" 26 #include "ir/dtype.h" 27 #include "pipeline/pynative/base.h" 28 29 namespace mindspore { 30 namespace pynative { 31 using PredictOutTypeMap = mindspore::HashMap<std::string, TypePtr>; 32 const TypePtr kTupleTensor1 = std::make_shared<Tuple>(TypePtrList{kTensorType}); 33 const TypePtr kTupleTensor2 = std::make_shared<Tuple>(TypePtrList{kTensorType, kTensorType}); 34 const TypePtr kTupleTensor3 = std::make_shared<Tuple>(TypePtrList{kTensorType, kTensorType, kTensorType}); 35 const TypePtr kTupleTensor4 = std::make_shared<Tuple>(TypePtrList{kTensorType, kTensorType, kTensorType, kTensorType}); 36 const TypePtr kTupleTensor5 = 37 std::make_shared<Tuple>(TypePtrList{kTensorType, kTensorType, kTensorType, kTensorType, kTensorType}); 38 const TypePtr kTupleTensor6 = 39 std::make_shared<Tuple>(TypePtrList{kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType}); 40 const TypePtr kTupleTensor7 = std::make_shared<Tuple>( 41 TypePtrList{kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType}); 42 const TypePtr kTupleTensor8 = std::make_shared<Tuple>( 43 TypePtrList{kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType}); 44 const TypePtr kTupleTensor9 = std::make_shared<Tuple>(TypePtrList{ 45 kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType, kTensorType}); 46 47 // If Abstract of the operator is constant, please add the type kTypeNone, such as "Size". 48 // If Abstract of the operator is Tensor and Tuple[Tensor], please add the type kTypeAny, such as "Eigh". 49 inline static PredictOutTypeMap out_type_prediction = {{"ActsULQ", kTupleTensor4}, 50 {"Adam", kTupleTensor3}, 51 {"AdamApplyOne", kTupleTensor3}, 52 {"AdamApplyOneAssign", kTupleTensor3}, 53 {"AdamApplyOneWithDecay", kTupleTensor3}, 54 {"AdamApplyOneWithDecayAssign", kTupleTensor3}, 55 {"AdamWeightDecay", kTupleTensor3}, 56 {"AdaptiveMaxPool2D", kTupleTensor2}, 57 {"AdaptiveMaxPool3D", kTupleTensor2}, 58 {"TopKRouter", kTypeAny}, 59 {"ApplyAdaMax", kTupleTensor3}, 60 {"ApplyAdadelta", kTupleTensor3}, 61 {"ApplyAdagrad", kTupleTensor2}, 62 {"ApplyAdagradV2", kTupleTensor2}, 63 {"ApplyAdamWithAmsgrad", kTupleTensor4}, 64 {"ApplyAdamWithAmsgradV2", kTupleTensor4}, 65 {"ApplyAddSign", kTupleTensor2}, 66 {"ApplyCamePart1", kTupleTensor3}, 67 {"ApplyCamePart2", kTupleTensor4}, 68 {"ApplyCamePart3", kTupleTensor4}, 69 {"ApplyCamePart4", kTupleTensor3}, 70 {"ApplyKerasMomentum", kTupleTensor2}, 71 {"ApplyPowerSign", kTupleTensor2}, 72 {"ApplyProximalAdagrad", kTupleTensor2}, 73 {"ArgMaxWithValue", kTupleTensor2}, 74 {"ArgMinWithValue", kTupleTensor2}, 75 {"BNTrainingReduce", kTupleTensor2}, 76 {"BNTrainingUpdate", kTupleTensor5}, 77 {"BNTrainingUpdateGrad", kTupleTensor2}, 78 {"BNTrainingUpdateV2", kTupleTensor3}, 79 {"BNTrainingUpdateV3", kTupleTensor5}, 80 {"BasicLSTMCell", kTupleTensor7}, 81 {"BasicLSTMCellCStateGrad", kTupleTensor2}, 82 {"BasicLSTMCellCStateGradV2", kTupleTensor2}, 83 {"BasicLSTMCellInputGrad", kTupleTensor2}, 84 {"BasicLSTMCellWeightGrad", kTupleTensor2}, 85 {"BatchNorm", kTupleTensor5}, 86 {"SyncBatchNorm", kTupleTensor5}, 87 {"BatchNormFold2GradD", kTupleTensor4}, 88 {"BatchNormFold2GradReduce", kTupleTensor2}, 89 {"BatchNormFoldD", kTupleTensor7}, 90 {"BatchNormGrad", kTupleTensor3}, 91 {"BatchNormGradWithActivation", kTupleTensor3}, 92 {"BatchNormGradWithAddAndActivation", kTupleTensor4}, 93 {"BatchNormGradGrad", kTupleTensor3}, 94 {"BiasDropoutAdd", kTupleTensor2}, 95 {"Broadcast", kTuple}, 96 {"CSRSparseMatrixToSparseTensor", kTupleTensor3}, 97 {"CTCGreedyDecoder", kTupleTensor4}, 98 {"CTCLoss", kTupleTensor2}, 99 {"CTCLossV2", kTupleTensor2}, 100 {"Coalesce", kTupleTensor3}, 101 {"ConcatOffset", kTypeNone}, 102 {"CombinedNonMaxSuppression", kTupleTensor4}, 103 {"ComputeAccidentalHits", kTupleTensor3}, 104 {"ConfusionMulGrad", kTupleTensor2}, 105 {"CorrectionMulGradReduce", kTupleTensor3}, 106 {"CudnnGRU", kTupleTensor4}, 107 {"Cummax", kTupleTensor2}, 108 {"Cummin", kTupleTensor2}, 109 {"Custom", kTypeAny}, 110 {"DSDGrad", kTupleTensor3}, 111 {"DeformableOffsetsGrad", kTupleTensor2}, 112 {"DenseGrad", kTupleTensor3}, 113 {"DenseToCSRSparseMatrix", kTupleTensor5}, 114 {"DenseToDenseSetOperation", kTupleTensor3}, 115 {"DenseToSparseSetOperation", kTupleTensor3}, 116 {"Dropout", kTupleTensor2}, 117 {"Dropout2D", kTupleTensor2}, 118 {"Dropout3D", kTupleTensor2}, 119 {"DynamicBroadcastGradientArgs", kTupleTensor2}, 120 {"DynamicGRUV2", kTupleTensor6}, 121 {"DynamicGRUV2Grad", kTupleTensor6}, 122 {"DynamicRNN", kTupleTensor8}, 123 {"Eig", kTupleTensor2}, 124 {"Eigh", kTypeAny}, 125 {"FakeLearnedScaleQuantPerChannelGradD", kTupleTensor2}, 126 {"FakeLearnedScaleQuantPerLayerGradD", kTupleTensor2}, 127 {"FakeQuantWithMinMaxVarsGradient", kTupleTensor3}, 128 {"FakeQuantWithMinMaxVarsPerChannelGradient", kTupleTensor3}, 129 {"FractionalAvgPool", kTupleTensor3}, 130 {"FractionalMaxPool", kTupleTensor3}, 131 {"FractionalMaxPool3DWithFixedKsize", kTupleTensor2}, 132 {"FractionalMaxPoolWithFixedKsize", kTupleTensor2}, 133 {"FusedDbnDw", kTupleTensor2}, 134 {"FusedMulAddNL2loss", kTupleTensor2}, 135 {"FusedMulApplyMomentum", kTupleTensor2}, 136 {"FusedMulApplyMomentumExtern", kTupleTensor3}, 137 {"FusedSparseAdam", kTupleTensor3}, 138 {"FusedSparseFtrl", kTupleTensor3}, 139 {"FusedSparseLazyAdam", kTupleTensor3}, 140 {"FusedSparseProximalAdagrad", kTupleTensor2}, 141 {"GRU", kTupleTensor4}, 142 {"GRUV2", kTupleTensor4}, 143 {"GRUV2Grad", kTupleTensor3}, 144 {"GRUV2HiddenGrad", kTupleTensor3}, 145 {"GRUV2HiddenGradCell", kTupleTensor3}, 146 {"Geqrf", kTupleTensor2}, 147 {"GetNext", kTuple}, 148 {"GridSampler2DGrad", kTupleTensor2}, 149 {"GridSampler3DGrad", kTupleTensor2}, 150 {"IFMR", kTupleTensor2}, 151 {"IdentityN", kTuple}, 152 {"InstanceNorm", kTupleTensor3}, 153 {"InstanceNormGrad", kTupleTensor3}, 154 {"InstanceNormV2", kTupleTensor3}, 155 {"InstanceNormV2Grad", kTupleTensor3}, 156 {"InvertPermutation", kTypeNone}, 157 {"LSTM", kTupleTensor5}, 158 {"LSTMGrad", kTupleTensor4}, 159 {"LSTMGradData", kTupleTensor3}, 160 {"LSTMInputGrad", kTupleTensor4}, 161 {"LSTMV2", kTupleTensor5}, 162 {"LU", kTupleTensor3}, 163 {"LambApplyOptimizerAssign", kTupleTensor3}, 164 {"LambNextMV", kTupleTensor4}, 165 {"LambNextMVWithDecay", kTupleTensor4}, 166 {"LambNextRight", kTupleTensor2}, 167 {"LayerNorm", kTupleTensor3}, 168 {"LayerNormV3", kTupleTensor3}, 169 {"LayerNormExt", kTupleTensor3}, 170 {"LayerNormBetaGammaBackprop", kTupleTensor2}, 171 {"LayerNormBetaGammaBackpropV2", kTupleTensor2}, 172 {"LayerNormGrad", kTupleTensor3}, 173 {"LayerNormGradV3", kTupleTensor3}, 174 {"LayerNormGradGrad", kTupleTensor3}, 175 {"LayerNormXBackpropV2", kTupleTensor2}, 176 {"LinearSumAssignment", kTupleTensor2}, 177 {"ListDiff", kTupleTensor2}, 178 {"LogMatrixDeterminant", kTupleTensor2}, 179 {"LogUniformCandidateSampler", kTupleTensor3}, 180 {"Lu", kTupleTensor2}, 181 {"LuUnpack", kTupleTensor3}, 182 {"LuUnpackGrad", kTupleTensor2}, 183 {"MakeTuple", kTypeNone}, 184 {"MapCacheIdx", kTupleTensor4}, 185 {"MapTensorGetData", kTupleTensor2}, 186 {"MatmulDDS", kTupleTensor2}, 187 {"MatmulDDSGrad", kTupleTensor2}, 188 {"MaxPool3DWithArgmax", kTupleTensor2}, 189 {"MaxPoolWithArgmax", kTupleTensor2}, 190 {"MaxPoolWithArgmaxV2", kTupleTensor2}, 191 {"MaximumGrad", kTupleTensor2}, 192 {"MaximumGradGrad", kTupleTensor3}, 193 {"Median", kTupleTensor2}, 194 {"MedianGrad", kTupleTensor2}, 195 {"Meshgrid", kTuple}, 196 {"MinMaxUpdatePerChannel", kTupleTensor2}, 197 {"MinMaxUpdatePerLayer", kTupleTensor2}, 198 {"MinimumGrad", kTupleTensor2}, 199 {"MinimumGradGrad", kTupleTensor3}, 200 {"MultilabelMarginLoss", kTupleTensor2}, 201 {"NLLLoss", kTupleTensor2}, 202 {"NMSWithMask", kTupleTensor3}, 203 {"PReLUGrad", kTupleTensor2}, 204 {"PSROIPooling", kTypeAny}, 205 {"PriorityReplayBufferSample", kTuple}, 206 {"Qr", kTupleTensor2}, 207 {"RNNTLoss", kTupleTensor2}, 208 {"RaggedRange", kTupleTensor2}, 209 {"RaggedTensorToSparse", kTupleTensor3}, 210 {"RandomChoiceWithMask", kTupleTensor2}, 211 {"ReduceStd", kTupleTensor2}, 212 {"ReservoirReplayBufferDestroy", kTupleTensor4}, 213 {"SampleDistortedBoundingBoxV2", kTupleTensor3}, 214 {"ScalarAdd", kTypeNone}, 215 {"ScalarBool", kTypeNone}, 216 {"ScalarDiv", kTypeNone}, 217 {"ScalarFloorDiv", kTypeNone}, 218 {"ScalarMod", kTypeNone}, 219 {"ScalarMul", kTypeNone}, 220 {"ScalarSub", kTypeNone}, 221 {"SelfAdjointEig", kTupleTensor2}, 222 {"SequenceAdd", kTypeNone}, 223 {"SequenceAddN", kTypeNone}, 224 {"SequenceConcat", kTypeNone}, 225 {"SequenceCount", kTypeNone}, 226 {"SequenceIndex", kTypeNone}, 227 {"SequenceMul", kTypeNone}, 228 {"SequenceMax", kTypeNone}, 229 {"SequenceMin", kTypeNone}, 230 {"SequenceSlice", kTypeNone}, 231 {"SequenceSliceGrad", kTypeNone}, 232 {"SequenceSliceSetItem", kTypeNone}, 233 {"SequenceStack", kTypeNone}, 234 {"SequenceZerosLike", kTypeNone}, 235 {"Size", kTypeNone}, 236 {"SolveTriangularGrad", kTupleTensor2}, 237 {"SoftmaxCrossEntropyWithLogits", kTupleTensor2}, 238 {"SoftmaxV2WithDropoutDoMaskV3", kTupleTensor2}, 239 {"Sort", kTupleTensor2}, 240 {"SparseAdd", kTupleTensor3}, 241 {"SparseAddGrad", kTupleTensor2}, 242 {"SparseApplyAdadelta", kTupleTensor3}, 243 {"SparseApplyAdagrad", kTupleTensor2}, 244 {"SparseApplyAdagradV2", kTupleTensor2}, 245 {"SparseApplyFtrl", kTupleTensor3}, 246 {"SparseApplyFtrlV2", kTupleTensor3}, 247 {"SparseApplyProximalAdagrad", kTupleTensor2}, 248 {"SparseApplyRMSProp", kTupleTensor3}, 249 {"SparseConcat", kTupleTensor3}, 250 {"SparseCountSparseOutput", kTupleTensor3}, 251 {"SparseCross", kTupleTensor3}, 252 {"SparseFillEmptyRows", kTupleTensor4}, 253 {"SparseFillEmptyRowsGrad", kTupleTensor2}, 254 {"SparseMatrixAdd", kTupleTensor5}, 255 {"SparseMatrixMul", kTupleTensor5}, 256 {"SparseMatrixSoftmax", kTupleTensor5}, 257 {"SparseMatrixSparseMatMul", kTupleTensor5}, 258 {"SparseMatrixTranspose", kTupleTensor5}, 259 {"SparseReorder", kTupleTensor2}, 260 {"SparseReshape", kTupleTensor2}, 261 {"SparseSlice", kTupleTensor3}, 262 {"SparseSoftmaxCrossEntropyWithLogitsV2", kTupleTensor2}, 263 {"SparseSparseMaximum", kTupleTensor2}, 264 {"SparseSparseMinimum", kTupleTensor2}, 265 {"SparseSplit", kTuple}, 266 {"SparseTensorToCSRSparseMatrix", kTupleTensor5}, 267 {"Split", kTuple}, 268 {"SplitTensor", kTuple}, 269 {"SplitWithSize", kTuple}, 270 {"SplitV", kTuple}, 271 {"Chunk", kTuple}, 272 {"SquareSumAll", kTupleTensor2}, 273 {"SquareSumV2", kTupleTensor2}, 274 {"Sspaddmm", kTupleTensor3}, 275 {"SubAndFilter", kTupleTensor2}, 276 {"Svd", kTupleTensor3}, 277 {"TensorToList", kTypeNone}, 278 {"TensorToScalar", kTypeNone}, 279 {"TensorToTuple", kTypeNone}, 280 {"TopK", kTupleTensor2}, 281 {"TupleGetItem", kTypeNone}, 282 {"UniformCandidateSampler", kTupleTensor3}, 283 {"Unique", kTupleTensor2}, 284 {"UniqueConsecutive", kTupleTensor3}, 285 {"UniqueWithPad", kTupleTensor2}, 286 {"Unpack", kTuple}, 287 {"Unstack", kTuple}, 288 {"bit_and", kTypeNone}, 289 {"bit_or", kTypeNone}, 290 {"make_range", kTypeNone}, 291 {"ScalarEq", kTypeNone}, 292 {"ScalarGe", kTypeNone}, 293 {"ScalarGt", kTypeNone}, 294 {"ScalarLe", kTypeNone}, 295 {"ScalarLt", kTypeNone}, 296 {"sequence_len", kTypeNone}, 297 {"tuple_setitem", kTypeNone}, 298 {"FlashAttentionScore", kTupleTensor4}, 299 {"FlashAttentionScoreGrad", kTupleTensor4}, 300 {"SplitTensor", kTypeAny}, 301 {"SplitWithSize", kTypeAny}, 302 {"SilentCheck", kTupleTensor5}, 303 {"BatchISendIRecv", kTuple}}; 304 305 TypePtr PredictOutTypeByName(const std::string &op_name); 306 307 TypePtr PredictOutType(const FrontendOpRunInfoPtr &op_run_info); 308 } // namespace pynative 309 } // namespace mindspore 310 #endif // MINDSPORE_CCSRC_PIPELINE_PYNATIVE_PREDICTOUTTYPEMAP_H_ 311