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_PI_JIT_CONFIG_H 17 #define MINDSPORE_PI_JIT_CONFIG_H 18 19 #include <set> 20 #include <string> 21 #include "pybind11/pybind11.h" 22 23 namespace mindspore { 24 namespace pijit { 25 namespace py = pybind11; 26 27 class GraphJitConfig { 28 public: 29 enum Options { 30 kBoolConf = 0, 31 kAutoJitCell, 32 kAutoGrad, 33 kAutoJit, 34 kPrintAfterAll, 35 kPrintTraceback, 36 kPrintBB, 37 kPrintCFG, 38 kInterpretCapturedCode, 39 kCompileWithoutCapture, 40 kCompileWithTry, 41 kGuardSpecializeScalar, 42 kGuardSpecializeContainer, 43 kGuardSpecializeTensor, 44 kGuardDetachObject, 45 kPrintGuard, 46 kReuseGraph, 47 kPrintReuseGraph, 48 kAutoCleanCache, 49 kPruneCase, 50 kLoopUnrolling, 51 kInferOnly, 52 kInferPrimitive, 53 kStrictTrace, 54 kPerfStatistics, 55 kLogGraphBreak, 56 kLogPerf, 57 kLogGuardPerf, 58 kTestGraphIR, 59 kEnableOptimizeForAttrItem, 60 kEnableEliminateUnusedOperation, 61 kEnableGeneratorExpressionToTuple, 62 kFeatureBreakAtInlinedFunction, 63 kEnableDynamicShape, 64 kEnableMsApiInfer, 65 kTraceFlag, 66 kSkipException, 67 /* ------------------------------ */ 68 kIntConf, 69 kMaxInlineDepth, 70 kMaxTraceDepth, 71 kMaxPruneCase, 72 kMaxLoopUnrolling, 73 kStaticGraphBytecodeMin, 74 kPerfStatisticsCount, 75 kPerfStatisticsScale10000x, 76 kInferPrimitiveMask, 77 kInferPrimitiveMax, 78 kLimitGraphSize, 79 kLimitGraphCount, 80 kGuardRelaxCount, 81 /* ------------------------------ */ 82 kOptionsCount 83 }; 84 GraphJitConfig(); 85 explicit GraphJitConfig(const py::object &c); GetBoolConfig(Options o)86 bool GetBoolConfig(Options o) const { return o > kBoolConf && o < kIntConf ? bool_conf[o - kBoolConf] : false; } getIntConfig(Options o)87 int getIntConfig(Options o) const { return o > kIntConf && o < kOptionsCount ? int_conf[o - kIntConf] : 0; } allowed_inline_modules()88 const auto &allowed_inline_modules() const { return allowed_inline_modules_; } 89 90 bool ShouldAutoJit(PyFrameObject *f); 91 92 void AddAllowedInlineModules(const std::string &module_name); 93 94 bool SetAutoJitFilter(PyObject *callable); 95 bool AddJitRelaxGuard(PyObject *list); 96 bool AddJitConstexpr(PyObject *callable_list); 97 bool AddJitForbidden(PyObject *callable_list); 98 bool AddAllowedInlineModules(PyObject *str_list); 99 std::string getJitLevel() const; 100 bool AddJitLevel(PyObject *str); 101 102 template <Options o> SetBool(PyObject * value)103 bool SetBool(PyObject *value) { 104 static_assert(o > kBoolConf && o < kIntConf); 105 bool_conf[o - kBoolConf] = value == Py_True; 106 return true; 107 } 108 109 template <Options o> SetInt(PyObject * value)110 bool SetInt(PyObject *value) { 111 static_assert(o > kIntConf && o < kOptionsCount); 112 int res = PyLong_AsLong(value); 113 if (PyErr_Occurred()) { 114 PyErr_Clear(); 115 return false; 116 } 117 int_conf[o - kIntConf] = res; 118 return true; 119 } 120 121 static void ApplyAutoJitCell(); 122 123 private: 124 std::set<std::string> allowed_inline_modules_; 125 int int_conf[kOptionsCount - kIntConf]; 126 bool bool_conf[kIntConf - kBoolConf]; 127 std::string jit_level; 128 }; 129 130 extern GraphJitConfig kPIJitConfigDefault; 131 132 } // namespace pijit 133 } // namespace mindspore 134 135 #endif // MINDSPORE_PI_JIT_CONFIG_H 136