1 /** 2 * This is the C++ adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/). 3 * 4 * Copyright 2019 Huawei Technologies Co., Ltd 5 * 6 * Licensed under the Apache License, Version 2.0 (the "License"); 7 * you may not use this file except in compliance with the License. 8 * You may obtain a copy of the License at 9 * 10 * http://www.apache.org/licenses/LICENSE-2.0 11 * 12 * Unless required by applicable law or agreed to in writing, software 13 * distributed under the License is distributed on an "AS IS" BASIS, 14 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 * See the License for the specific language governing permissions and 16 * limitations under the License. 17 */ 18 19 #ifndef MINDSPORE_CCSRC_PIPELINE_JIT_STATIC_ANALYSIS_PRIM_H_ 20 #define MINDSPORE_CCSRC_PIPELINE_JIT_STATIC_ANALYSIS_PRIM_H_ 21 22 #include <algorithm> 23 #include <memory> 24 #include <string> 25 #include <unordered_map> 26 #include <vector> 27 28 #include "pipeline/jit/static_analysis/evaluator.h" 29 #include "abstract/primitive_infer_map.h" 30 31 namespace mindspore { 32 namespace abstract { 33 class StandardPrimEvaluator : public TrivialPrimEvaluator { 34 public: StandardPrimEvaluator(const PrimitivePtr & primitive,const StandardPrimitiveImplReg & eval_impl)35 StandardPrimEvaluator(const PrimitivePtr &primitive, const StandardPrimitiveImplReg &eval_impl) 36 : TrivialPrimEvaluator("StandardPrimEvaluator"), prim_(primitive), eval_impl_(eval_impl) {} 37 ~StandardPrimEvaluator() override = default; 38 MS_DECLARE_PARENT(StandardPrimEvaluator, TrivialPrimEvaluator); 39 EvalResultPtr EvalPrim(const AnalysisEnginePtr &engine, const AbstractBasePtrList &args) override; prim()40 PrimitivePtr prim() { return prim_; } 41 ToString()42 std::string ToString() const override { return identifier_ + prim_->name(); } 43 44 private: 45 EvalResultPtr EvalPyCheckPrim(const AnalysisEnginePtr &engine, const AbstractBasePtrList &args); 46 EvalResultPtr RunPyInferValue(const AnalysisEnginePtr &engine, const AbstractBasePtr &abs_base, 47 const AbstractBasePtrList &args); 48 PrimitivePtr prim_; 49 const StandardPrimitiveImplReg eval_impl_; 50 }; 51 52 using StandardPrimEvaluatorPtr = std::shared_ptr<StandardPrimEvaluator>; 53 54 class PythonPrimEvaluator : public TrivialPrimEvaluator { 55 public: PythonPrimEvaluator(const PrimitivePyPtr primitive)56 explicit PythonPrimEvaluator(const PrimitivePyPtr primitive) 57 : TrivialPrimEvaluator("PythonPrimEvaluator"), prim_py_(primitive) {} 58 ~PythonPrimEvaluator() override = default; 59 MS_DECLARE_PARENT(PythonPrimEvaluator, TrivialPrimEvaluator); 60 EvalResultPtr EvalPrim(const AnalysisEnginePtr &engine, const AbstractBasePtrList &args) override; prim()61 PrimitivePtr prim() { return dyn_cast<Primitive>(prim_py_); } 62 ToString()63 std::string ToString() const override { return identifier_ + prim_py_->name(); } 64 65 private: 66 PrimitivePyPtr prim_py_; 67 }; 68 69 class DoSignatureEvaluator : public Evaluator { 70 public: DoSignatureEvaluator(const PrimitivePtr primitive)71 explicit DoSignatureEvaluator(const PrimitivePtr primitive) : Evaluator("DoSignatureEvaluator"), prim_(primitive) {} 72 ~DoSignatureEvaluator() override = default; 73 EvalResultPtr Run(AnalysisEnginePtr engine, const ConfigPtrList &argrefs, 74 const AnfNodeConfigPtr &out_config = nullptr) override; 75 Eval(AnalysisEnginePtr,const AbstractBasePtrList &,const AnfNodeConfigPtr &)76 EvalResultPtr Eval(AnalysisEnginePtr, const AbstractBasePtrList &, const AnfNodeConfigPtr &) override { 77 MS_LOG(EXCEPTION) << "Eval() should not be called, Run() method should be called"; 78 } 79 80 private: 81 PrimitivePtr prim_; 82 }; 83 84 class UnpackGraphEvaluator : public Evaluator { 85 public: UnpackGraphEvaluator(const PrimitivePtr primitive)86 explicit UnpackGraphEvaluator(const PrimitivePtr primitive) : Evaluator("UnpackGraphEvaluator"), prim_(primitive) {} 87 ~UnpackGraphEvaluator() override = default; 88 EvalResultPtr Run(AnalysisEnginePtr engine, const ConfigPtrList &argrefs, 89 const AnfNodeConfigPtr &out_config = nullptr) override; 90 Eval(AnalysisEnginePtr,const AbstractBasePtrList &,const AnfNodeConfigPtr &)91 EvalResultPtr Eval(AnalysisEnginePtr, const AbstractBasePtrList &, const AnfNodeConfigPtr &) override { 92 MS_LOG(EXCEPTION) << "Eval() should not be called, Run() method should be called"; 93 } 94 95 private: 96 PrimitivePtr prim_; 97 }; 98 99 class MixedPrecisionCastEvaluator : public Evaluator { 100 public: MixedPrecisionCastEvaluator(const PrimitivePtr primitive)101 explicit MixedPrecisionCastEvaluator(const PrimitivePtr primitive) 102 : Evaluator("MixedPrecisionCastEvaluator"), prim_(primitive) {} 103 ~MixedPrecisionCastEvaluator() override = default; 104 EvalResultPtr Run(AnalysisEnginePtr engine, const ConfigPtrList &argrefs, 105 const AnfNodeConfigPtr &out_config = nullptr) override; 106 Eval(AnalysisEnginePtr,const AbstractBasePtrList &,const AnfNodeConfigPtr &)107 EvalResultPtr Eval(AnalysisEnginePtr, const AbstractBasePtrList &, const AnfNodeConfigPtr &) override { 108 MS_LOG(EXCEPTION) << "Eval() should not be called, Run() method should be called"; 109 } 110 111 private: 112 PrimitivePtr prim_; 113 }; 114 115 bool IsInWhiteList(const PrimitivePtr &primitive); 116 117 using ValuePtrList = std::vector<ValuePtr>; 118 using PrimitiveImpl = ValuePtr (*)(const ValuePtrList &); 119 120 class UniformPrimEvaluator : public TrivialPrimEvaluator { 121 public: UniformPrimEvaluator(const FunctionPtr func_desc,PrimitiveImpl impl,bool eval_value,const TypePtr specify_out_type)122 UniformPrimEvaluator(const FunctionPtr func_desc, PrimitiveImpl impl, bool eval_value, const TypePtr specify_out_type) 123 : TrivialPrimEvaluator("UniformPrimEvaluator"), 124 impl_(impl), 125 eval_value_(eval_value), 126 func_desc_(func_desc), 127 nargs_(func_desc_->args().size()), 128 return_value_type_(func_desc_->retval()), 129 specify_out_type_(specify_out_type) { 130 for (size_t i = 0; i < nargs_; ++i) { 131 TypePtr type = func_desc_->args()[i]; 132 if (type_map_[type]) { 133 type_map_[type]->push_back(i); 134 } else { 135 type_map_[type] = std::make_shared<std::vector<size_t>>(); 136 type_map_[type]->push_back(i); 137 } 138 } 139 } 140 ~UniformPrimEvaluator() override = default; 141 MS_DECLARE_PARENT(UniformPrimEvaluator, TrivialPrimEvaluator); 142 143 EvalResultPtr EvalPrim(const AnalysisEnginePtr &engine, const AbstractBasePtrList &args) override; 144 ValuePtr RunImpl(const ValuePtrList &args) const; 145 146 // If eval_value_ is False, return broadened arguments. NormalizeArgs(const AbstractBasePtrList & args_spec_list)147 AbstractBasePtrList NormalizeArgs(const AbstractBasePtrList &args_spec_list) const override { 148 if (!eval_value_) { 149 AbstractBasePtrList broadened_args_spec_list; 150 (void)std::transform(args_spec_list.begin(), args_spec_list.end(), std::back_inserter(broadened_args_spec_list), 151 [](const AbstractBasePtr &arg) -> AbstractBasePtr { return arg->Broaden(); }); 152 return broadened_args_spec_list; 153 } 154 return args_spec_list; 155 } 156 157 private: 158 PrimitiveImpl impl_; 159 bool eval_value_; 160 const FunctionPtr func_desc_; 161 const std::size_t nargs_; 162 const TypePtr return_value_type_; 163 const TypePtr specify_out_type_; 164 std::unordered_map<TypePtr, std::shared_ptr<std::vector<size_t>>, TypeHasher, TypeEqual> type_map_; 165 }; 166 167 PrimEvaluatorMap &GetPrimEvaluatorConstructors(); 168 169 // Check whether type x is a subtype of model. 170 bool IsSubtype(const AbstractBasePtr x, const TypePtr model); 171 172 void ClearPrimEvaluatorMap(); 173 174 py::dict ConvertAbstractToPython(const AbstractBasePtr &abs_base); 175 } // namespace abstract 176 } // namespace mindspore 177 178 #endif // MINDSPORE_CCSRC_PIPELINE_JIT_STATIC_ANALYSIS_PRIM_H_ 179