|
| 1 | + |
| 2 | +/* |
| 3 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 4 | + * or more contributor license agreements. See the NOTICE file |
| 5 | + * distributed with this work for additional information |
| 6 | + * regarding copyright ownership. The ASF licenses this file |
| 7 | + * to you under the Apache License, Version 2.0 (the |
| 8 | + * "License"); you may not use this file except in compliance |
| 9 | + * with the License. You may obtain a copy of the License at |
| 10 | + * |
| 11 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 12 | + * |
| 13 | + * Unless required by applicable law or agreed to in writing, |
| 14 | + * software distributed under the License is distributed on an |
| 15 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 16 | + * KIND, either express or implied. See the License for the |
| 17 | + * specific language governing permissions and limitations |
| 18 | + * under the License. |
| 19 | + */ |
| 20 | +#include <tvm/relay/attrs/nn.h> |
| 21 | +#include <tvm/relay/expr_functor.h> |
| 22 | +#include <tvm/relay/transform.h> |
| 23 | +#include <tvm/runtime/ndarray.h> |
| 24 | + |
| 25 | +#include "../../../qnn/utils.h" |
| 26 | +#include "../../../transforms/pattern_utils.h" |
| 27 | + |
| 28 | +namespace tvm { |
| 29 | +namespace relay { |
| 30 | +namespace contrib { |
| 31 | +namespace cmsisnn { |
| 32 | + |
| 33 | +class ExtractConstantsMutator : public MixedModeMutator { |
| 34 | + public: |
| 35 | + explicit ExtractConstantsMutator(IRModule& mod) : mod_(mod) {} |
| 36 | + |
| 37 | + private: |
| 38 | + String gen_var_name() { return "tvm_var_extract_const_" + std::to_string(var_count_++); } |
| 39 | + |
| 40 | + Expr VisitExpr_(const FunctionNode* func) final { |
| 41 | + Function final_func = GetRef<Function>(func); |
| 42 | + ++func_nesting_level_; |
| 43 | + auto new_body = VisitExpr(func->body); |
| 44 | + --func_nesting_level_; |
| 45 | + if (!new_body.same_as(func->body)) { |
| 46 | + final_func = Function(FreeVars(new_body), new_body, func->ret_type, |
| 47 | + FreeTypeVars(new_body, mod_), func->attrs); |
| 48 | + function_to_constants_.Set(GetRef<Function>(func), constants_within_function_); |
| 49 | + constants_within_function_.clear(); |
| 50 | + } |
| 51 | + return final_func; |
| 52 | + } |
| 53 | + |
| 54 | + Expr Rewrite_(const CallNode* call, const Expr& post) final { |
| 55 | + Expr final_call = post; |
| 56 | + auto* post_call = post.as<CallNode>(); |
| 57 | + if (post_call == nullptr) { |
| 58 | + return final_call; |
| 59 | + } |
| 60 | + |
| 61 | + // Replace Constant arguments with Vars for ML Operators |
| 62 | + // Perform this for non-main Call Nodes only |
| 63 | + if (func_nesting_level_ && call->op.as<OpNode>()) { |
| 64 | + Array<Expr> new_args; |
| 65 | + for (auto& arg : post_call->args) { |
| 66 | + auto* const_arg = arg.as<ConstantNode>(); |
| 67 | + if (const_arg && !const_arg->is_scalar()) { |
| 68 | + Var var_arg = Var(gen_var_name(), const_arg->tensor_type()); |
| 69 | + new_args.push_back(var_arg); |
| 70 | + constants_within_function_.push_back(GetRef<Constant>(const_arg)); |
| 71 | + } else { |
| 72 | + new_args.push_back(arg); |
| 73 | + } |
| 74 | + } |
| 75 | + final_call = Call(call->op, new_args, call->attrs, {}); |
| 76 | + } |
| 77 | + |
| 78 | + // Since the constants are kicked out of partitioned functions |
| 79 | + // a new call to global function is needed |
| 80 | + if (auto* glob_var_node = post_call->op.as<GlobalVarNode>()) { |
| 81 | + auto glob_var = GetRef<GlobalVar>(glob_var_node); |
| 82 | + auto glob_func = Downcast<Function>(mod_->Lookup(glob_var)); |
| 83 | + auto new_glob_func = VisitExpr(glob_func); |
| 84 | + if (!new_glob_func.same_as(glob_func)) { |
| 85 | + mod_->Update(glob_var, Downcast<Function>(new_glob_func)); |
| 86 | + Array<Expr> new_args = post_call->args; |
| 87 | + ICHECK(function_to_constants_.find(glob_func) != function_to_constants_.end()); |
| 88 | + for (auto constant : function_to_constants_.at(glob_func)) { |
| 89 | + new_args.push_back(constant); |
| 90 | + } |
| 91 | + final_call = Call(glob_var, new_args); |
| 92 | + } |
| 93 | + } |
| 94 | + |
| 95 | + // Since the constants are kicked out of the local partitioned functions |
| 96 | + // a new call to local function is needed |
| 97 | + if (auto* func_node = call->op.as<FunctionNode>()) { |
| 98 | + Function func = GetRef<Function>(func_node); |
| 99 | + auto new_func = VisitExpr(func); |
| 100 | + if (!new_func.same_as(func)) { |
| 101 | + Array<Expr> new_args = post_call->args; |
| 102 | + ICHECK(function_to_constants_.find(func) != function_to_constants_.end()); |
| 103 | + for (auto constant : function_to_constants_.at(func)) { |
| 104 | + constants_within_function_.push_back(constant); |
| 105 | + Var var_arg = Var(gen_var_name(), constant->tensor_type()); |
| 106 | + new_args.push_back(var_arg); |
| 107 | + } |
| 108 | + final_call = Call(new_func, new_args); |
| 109 | + } |
| 110 | + } |
| 111 | + |
| 112 | + return final_call; |
| 113 | + } |
| 114 | + |
| 115 | + private: |
| 116 | + /* \brief Updated module where all calls have replaced constants with new variables */ |
| 117 | + IRModule mod_; |
| 118 | + /* \brief Maintains mapping of original function to the replaced constants */ |
| 119 | + Map<Function, Array<Constant>> function_to_constants_; |
| 120 | + /* \brief Constants being kicked out of a function during the function visit */ |
| 121 | + Array<Constant> constants_within_function_; |
| 122 | + /* \brief Keeps track of variables being created */ |
| 123 | + int var_count_ = 0; |
| 124 | + /* \brief Keeps track of function scope */ |
| 125 | + int func_nesting_level_ = 0; |
| 126 | +}; |
| 127 | + |
| 128 | +/*! * \brief Kicks out all constants out of the partitioned function into main() */ |
| 129 | +IRModule ExtractConstants(IRModule mod) { |
| 130 | + String func_name; |
| 131 | + Function func; |
| 132 | + |
| 133 | + auto extract_constants = ExtractConstantsMutator(mod); |
| 134 | + Function main_func = Downcast<Function>(mod->Lookup("main")); |
| 135 | + auto new_main_body = extract_constants.VisitExpr(main_func->body); |
| 136 | + if (!new_main_body.same_as(main_func->body)) { |
| 137 | + auto main_var = mod->GetGlobalVar("main"); |
| 138 | + auto new_main_func = Function(main_func->params, new_main_body, main_func->ret_type, |
| 139 | + main_func->type_params, main_func->attrs); |
| 140 | + mod->Update(main_var, new_main_func); |
| 141 | + } |
| 142 | + return mod; |
| 143 | +} |
| 144 | + |
| 145 | +transform::Pass ExtractConstantsFromPartitionedFunction() { |
| 146 | + runtime::TypedPackedFunc<IRModule(IRModule, transform::PassContext)> pass_func = |
| 147 | + [=](IRModule m, transform::PassContext pc) { return ExtractConstants(m); }; |
| 148 | + return tvm::transform::CreateModulePass(pass_func, 0, "ExtractConstantsFromPartitionedFunction", |
| 149 | + {}); |
| 150 | +} |
| 151 | + |
| 152 | +TVM_REGISTER_GLOBAL("relay.ext.cmsisnn.transform.ExtractConstantsFromPartitionedFunction") |
| 153 | + .set_body_typed([]() { return ExtractConstantsFromPartitionedFunction(); }); |
| 154 | + |
| 155 | +} // namespace cmsisnn |
| 156 | +} // namespace contrib |
| 157 | +} // namespace relay |
| 158 | +} // namespace tvm |
0 commit comments