Lichen455
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
mindspore2022/mindspore/ccsrc/pipeline/static_analysis/program_specialize.cc

720 lines
28 KiB
C++
Raw Blame History

/**
* This is the C++ adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
*
* Copyright 2019 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "pipeline/static_analysis/program_specialize.h"
#include <algorithm>
#include <exception>
#include "./common.h"
#include "operator/ops.h"
#include "operator/composite/do_signature.h"
#include "pipeline/static_analysis/abstract_function.h"
#include "utils/graph_utils.h"
#include "utils/log_adapter.h"
#include "utils/profile.h"
#include "debug/trace.h"
namespace mindspore {
namespace abstract {
namespace {
inline AbstractBasePtr GetEvaluatedValueWrap(const AnfNodeConfigPtr &conf) {
if (conf->node()->intermediate_abstract()) {
return conf->node()->intermediate_abstract();
}
return conf->GetEvaluatedValue()->abstract();
}
AnfNodePtr BuildValueNode(const ValuePtr &v, const AbstractBasePtr &abs_base) {
AnfNodePtr value_node = NewValueNode(v);
value_node->set_abstract(abs_base);
MS_LOG(DEBUG) << "Create ValueNode: " << value_node->ToString() << ", with abstract: " << abs_base->ToString();
return value_node;
}
bool IsVisible(FuncGraphPtr fg, const FuncGraphPtr &parent) {
while (fg != nullptr && fg != parent) {
fg = fg->parent();
}
return fg == parent;
}
} // namespace
FuncGraphPtr ProgramSpecializer::Run(const FuncGraphPtr &fg, const AnalysisContextPtr &context) {
MS_EXCEPTION_IF_NULL(fg);
MS_EXCEPTION_IF_NULL(context);
MS_LOG(DEBUG) << "Specialize topmost function graph: " << context->func_graph()->ToString();
return SpecializeFuncGraph(fg, context);
}
FuncGraphPtr ProgramSpecializer::SpecializeFuncGraph(const FuncGraphPtr &fg, const AnalysisContextPtr &context) {
MS_EXCEPTION_IF_NULL(fg);
MS_EXCEPTION_IF_NULL(context);
auto iter = specializations_.find(context->SpecializeKey());
if (iter != specializations_.end()) {
return iter->second->specialized_func_graph();
}
std::shared_ptr<FuncGraphSpecializer> fg_spec = std::make_shared<FuncGraphSpecializer>(this, fg, context);
FuncGraphPtr fg2 = fg_spec->specialized_func_graph();
specializations_[context->SpecializeKey()] = fg_spec;
fg_spec->Run();
return fg2;
}
std::shared_ptr<FuncGraphSpecializer> ProgramSpecializer::GetFuncGraphSpecializer(const AnalysisContextPtr &context) {
MS_EXCEPTION_IF_NULL(context);
auto iter = specializations_.find(context->SpecializeKey());
if (iter != specializations_.end()) {
return iter->second;
}
return nullptr;
}
std::string GetNextCounter() {
static int g_CloneCounter = 1;
std::string str_count = std::to_string(g_CloneCounter);
g_CloneCounter++;
return str_count;
}
FuncGraphSpecializer::FuncGraphSpecializer(ProgramSpecializer *const s, const FuncGraphPtr &fg,
const AnalysisContextPtr &context)
: specializer_(s), func_graph_(fg), context_(context) {
parent_ = s->GetFuncGraphSpecializer(context->parent());
engine_ = s->engine();
cloner_ = SpecializerClone(fg, std::make_shared<TraceSpecialize>(GetNextCounter()));
repl_node_ = cloner_->cloned_node();
specialized_func_graph_ = cloner_->cloned_func_graph()[fg];
todo_.push_back(fg->get_return());
auto ps = fg->parameters();
(void)todo_.insert(todo_.end(), ps.begin(), ps.end());
}
AnfNodePtr FuncGraphSpecializer::ReplicateDisconnectedNode(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
FuncGraphPtr fg = node->func_graph();
if (node->isa<ValueNode>()) {
return node;
}
std::shared_ptr<FuncGraphSpecializer> specializer = shared_from_this();
while (fg != nullptr && fg != specializer->func_graph_) {
specializer = specializer->parent_;
}
// If had replicated, just return that.
auto iter = specializer->repl_node_->find(node);
if (iter != specializer->repl_node_->end()) {
return iter->second;
}
auto new_node = specializer->cloner_->CloneDisconnected(node);
if (node->isa<CNode>()) {
if (!new_node->isa<CNode>()) {
MS_LOG(EXCEPTION) << "new_node must be a CNode, but is " << new_node->DebugString() << ".";
}
auto c_node = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(c_node);
auto inputs = c_node->inputs();
std::vector<AnfNodePtr> new_inputs;
(void)std::transform(inputs.begin(), inputs.end(), std::back_inserter(new_inputs),
[this](const AnfNodePtr &inp) -> AnfNodePtr {
if (inp->isa<ValueNode>()) {
return inp;
}
return ReplicateDisconnectedNode(inp);
});
auto c_new_node = new_node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(c_new_node);
c_new_node->set_inputs(new_inputs);
}
iter = specializer->repl_node_->find(node);
if (iter != specializer->repl_node_->end()) {
if (iter->second == node) {
MS_LOG(EXCEPTION) << "Replicated is same as original node, node: " << node->ToString();
}
} else {
MS_LOG(EXCEPTION) << "Replicate node failed, node: " << node->ToString();
}
return new_node;
}
AnfNodePtr FuncGraphSpecializer::GetReplicatedNode(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
FuncGraphPtr fg = node->func_graph();
std::shared_ptr<FuncGraphSpecializer> specializer = shared_from_this();
while (fg != nullptr && fg != specializer->func_graph_) {
specializer = specializer->parent_;
}
MS_EXCEPTION_IF_NULL(specializer->repl_node_);
auto iter = specializer->repl_node_->find(node);
if (iter != specializer->repl_node_->end()) {
return iter->second;
}
return node;
}
void FuncGraphSpecializer::Run() {
MS_LOG(DEBUG) << "Before run, origin func graph name: " << func_graph_->ToString()
<< ", cloned func graph name: " << specialized_func_graph_->ToString()
<< ", func graph: " << func_graph_->get_return()->DebugString();
FirstPass();
SecondPass();
MS_LOG(DEBUG) << "After run, origin func graph name: " << func_graph_->ToString()
<< ", cloned func graph name: " << specialized_func_graph_->ToString()
<< ", new func graph: " << specialized_func_graph_->get_return()->DebugString();
}
void FuncGraphSpecializer::FirstPass() {
while (todo_.size()) {
AnfNodePtr node = todo_.back();
todo_.pop_back();
if (node->func_graph() == nullptr) {
// do nothing for ValueNode
continue;
}
if (node->func_graph() != func_graph_) {
if (parent_ == nullptr) {
MS_LOG(EXCEPTION) << "Parent must not null NodeInfo: " << trace::GetDebugInfo(node->debug_info());
}
parent_->AddTodoItem(node);
parent_->FirstPass();
AnfNodePtr new_node = parent_->GetReplicatedNode(node);
if (node->isa<CNode>()) {
parent_->ProcessCNode(new_node->cast<CNodePtr>());
}
continue;
}
if (marked_.count(node) > 0) {
continue;
}
(void)marked_.insert(node);
ProcessNode(node);
}
}
// Specialize CNode in func graphs
void FuncGraphSpecializer::SecondPass() {
for (auto &node : BroadFirstSearchGraphCNodes(specialized_func_graph_->get_return())) {
if (node->isa<CNode>()) {
ProcessCNode(node->cast<CNodePtr>());
}
}
}
void FuncGraphSpecializer::ProcessNode(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
ScopeGuard scope_guard(node->scope());
AnfNodeConfigPtr conf = MakeConfig(node);
AnfNodePtr new_node = GetReplicatedNode(node);
MS_EXCEPTION_IF_NULL(new_node);
if (new_node->func_graph() != specialized_func_graph_) {
MS_LOG(EXCEPTION) << "Error in specializer [A] node: " << node->DebugString()
<< ", new_node: " << new_node->DebugString()
<< ", new_node->func_graph(): " << new_node->func_graph()->ToString()
<< ", specialized_func_graph_: " << specialized_func_graph_->ToString();
return;
}
new_node->set_abstract(GetEvaluatedValueWrap(conf));
if (new_node->isa<CNode>() && new_node->abstract()->isa<PartialAbstractClosure>()) {
auto partial_abstract = dyn_cast<PartialAbstractClosure>(new_node->abstract());
if (partial_abstract->node() == node) {
partial_abstract->set_node(new_node);
}
}
MS_LOG(DEBUG) << "Set new_node: " << new_node->ToString() << ", abstract as: " << new_node->abstract()->ToString();
if (node->isa<CNode>()) {
auto attrs = conf->GetEvaluatedValue()->attribute();
auto c_old = node->cast<CNodePtr>();
auto c_new = new_node->cast<CNodePtr>();
auto new_inputs = c_new->inputs();
auto old_inputs = c_old->inputs();
for (size_t i = 0; i < old_inputs.size(); ++i) {
auto node_input = old_inputs[i];
AnfNodeConfigPtr iconf = MakeConfig(node_input);
AbstractBasePtr ival = GetEvaluatedValueWrap(iconf);
// First try to check if node_input can be replaced by a ValueNode. If cannot, then try to check if
// can be replaced by another CNode from anfnode_config_map, otherwise use the replicated node.
AnfNodePtr replace_node = BuildPossibleValueNode(iconf->node(), ival, attrs);
if (replace_node == nullptr) {
replace_node = BuildReplacedNode(iconf);
MS_EXCEPTION_IF_NULL(replace_node);
replace_node->set_abstract(ival);
MS_LOG(DEBUG) << "Set replaced: " << replace_node->ToString() << ", to abstract: " << ival->ToString();
} else {
MS_LOG(DEBUG) << "Build possible value node for node: " << node_input->DebugString()
<< ", ival: " << ival->ToString() << ", replace_node: " << replace_node->ToString();
}
if (new_inputs[i] != replace_node) {
new_inputs[i] = replace_node;
MS_LOG(DEBUG) << "Set new_input[" << i << "] = " << replace_node->DebugString();
}
}
c_new->set_inputs(new_inputs);
}
}
AnfNodePtr FuncGraphSpecializer::BuildReplacedNode(const AnfNodeConfigPtr &conf) {
MS_EXCEPTION_IF_NULL(conf);
auto conf_iter = engine_->anfnode_config_map().find(conf);
AnfNodeConfigPtr new_conf = conf;
while (conf_iter != engine_->anfnode_config_map().end()) {
MS_LOG(DEBUG) << "Origin conf: graph(" << new_conf->node()->func_graph()->ToString() << ", node("
<< new_conf->node()->DebugString() << ")";
new_conf = conf_iter->second;
MS_EXCEPTION_IF_NULL(new_conf);
MS_LOG(DEBUG) << "Replaced conf: graph(" << conf->node()->func_graph()->ToString() << ", node("
<< conf->node()->DebugString() << ")";
(void)ReplicateDisconnectedNode(new_conf->node());
conf_iter = engine_->anfnode_config_map().find(new_conf);
}
todo_.push_back(new_conf->node());
auto repl = GetReplicatedNode(new_conf->node());
if (repl->func_graph()) {
MS_LOG(DEBUG) << "Set repl: graph(" << repl->func_graph()->ToString() << "), node:" << repl->DebugString()
<< ") to replace origin:" << new_conf->node()->DebugString();
} else {
MS_LOG(DEBUG) << "Set repl: graph(nullptr), node(" << repl->DebugString()
<< ") to replace origin: " << new_conf->node()->DebugString();
}
return repl;
}
namespace {
const StringImmPtr kDeadNode = std::make_shared<StringImm>("Dead Node");
const StringImmPtr kPolyNode = std::make_shared<StringImm>("Poly Node");
inline bool CanSpecializeNode(const AnfNodePtr &node) {
if (IsValueNode<FuncGraph>(node) || IsValueNode<MetaFuncGraph>(node) || IsValueNode<Primitive>(node)) {
return true;
}
return false;
}
} // namespace
AnfNodePtr FuncGraphSpecializer::BuildSpecializedNode(const AnfNodePtr &node, const AbstractBasePtr &abs,
const AbstractBasePtrList &argvals) {
MS_EXCEPTION_IF_NULL(abs);
AbstractFunctionPtr real_a = dyn_cast<AbstractFunction>(abs);
MS_EXCEPTION_IF_NULL(real_a);
AbstractFunctionPtr func = real_a->GetUnique();
SpecializeStatusCode errcode;
ScopeGuard scope_guard(node->scope());
AnfNodePtr repl = BuildSpecializedNodeInner(abs, func, argvals, &errcode);
if (repl == nullptr) {
if (errcode == kSpecializeFindUniqueArgvalDead) {
const auto error_dead_node = std::make_shared<AbstractError>(kDeadNode, node);
repl = BuildValueNode(kDeadNode, error_dead_node);
MS_LOG(DEBUG) << "DEAD for node: " << node->DebugString() << ", abstract: " << abs->ToString();
} else if (errcode == kSpecializeFindUniqueArgvalPoly) {
const auto error_poly_node = std::make_shared<AbstractError>(kPolyNode, node);
repl = BuildValueNode(kPolyNode, error_poly_node);
MS_LOG(DEBUG) << "POLY for node: " << node->DebugString() << ", abstract: " << abs->ToString();
} else {
MS_LOG(EXCEPTION) << "Failed to build specialized node, node: " << node->DebugString()
<< ", abstract: " << abs->ToString();
}
}
return repl;
}
AnfNodePtr FuncGraphSpecializer::BuildSpecializedNodeInner(const AbstractBasePtr &abs, const AbstractFunctionPtr &func,
const AbstractBasePtrList &args,
SpecializeStatusCode *errcode) {
MS_EXCEPTION_IF_NULL(abs);
MS_EXCEPTION_IF_NULL(func);
MS_EXCEPTION_IF_NULL(errcode);
*errcode = kSpecializeSuccess;
auto real_func = dyn_cast<TypedPrimitiveAbstractClosure>(func);
if (real_func != nullptr) {
return BuildValueNode(real_func->prim(), abs);
}
EvaluatorPtr eval;
eval = engine_->GetEvaluatorFor(func);
MS_EXCEPTION_IF_NULL(eval);
AbstractBasePtrList argvals = eval->NormalizeArgs(args);
std::pair<AbstractBasePtrList, AbstractBasePtr> result;
SpecializeStatusCode status = FindUniqueArgvals(func, eval, argvals, &result);
if (status != kSpecializeSuccess) {
*errcode = status;
return nullptr;
}
argvals = result.first;
AbstractBasePtr unique_output = result.second;
auto prim_func = dyn_cast<PrimitiveAbstractClosure>(func);
if (prim_func != nullptr) {
auto type_func = std::make_shared<TypedPrimitiveAbstractClosure>(prim_func->prim(), argvals, unique_output);
return BuildValueNode(prim_func->prim(), type_func);
}
if (!eval->isa<BaseFuncGraphEvaluator>()) {
MS_LOG(EXCEPTION) << "Eval is not BaseGraphEvaluator, but " << eval->ToString();
}
auto real_eval = dyn_cast<BaseFuncGraphEvaluator>(eval);
if (func->context() == nullptr) {
MS_LOG(EXCEPTION) << "Func context is nullptr NodeInfo: " << trace::GetDebugInfo(func_graph_->debug_info());
}
AnalysisContextPtr context = real_eval->MakeContext(engine_, argvals);
MS_LOG(DEBUG) << "Specialize function graph: " << context->func_graph()->ToString() << ", args: " << argvals.size()
<< ", graph: " << context->func_graph()->get_return()->DebugString();
FuncGraphPtr v = specializer_->SpecializeFuncGraph(context->func_graph(), context);
return BuildValueNode(v, abs);
}
AnfNodePtr FuncGraphSpecializer::BuildSpecializedParameterNode(const CNodePtr &new_node) {
auto new_inputs = new_node->inputs();
AnfNodePtr func = new_inputs[0];
AbstractBasePtr fnval = new_inputs[0]->abstract();
AbstractBasePtrList args;
auto backed_fnval = fnval;
if (fnval->isa<PartialAbstractClosure>()) {
auto partial_closure = dyn_cast<PartialAbstractClosure>(fnval);
backed_fnval = partial_closure->fn();
args = partial_closure->args();
}
std::transform(new_inputs.cbegin() + 1, new_inputs.cend(), std::back_inserter(args),
[](const AnfNodePtr &inp) { return inp->abstract(); });
ScopeGuard scope_guard(new_node->scope());
auto specialized_node = BuildSpecializedNode(func, backed_fnval, args);
auto wrapped_node = specialized_node;
if (fnval->isa<PartialAbstractClosure>()) {
auto partial_closure = dyn_cast<PartialAbstractClosure>(fnval);
AnfNodePtrList partial_node_list = {BuildValueNode(prim::kPrimPartial, FromValueInside(prim::kPrimPartial)),
specialized_node};
auto anf_node = partial_closure->node();
if (!anf_node->isa<CNode>()) {
MS_LOG(EXCEPTION) << "Must be cnode, but " << anf_node->DebugString();
}
auto cnode = anf_node->cast<CNodePtr>();
if (cnode->size() != partial_closure->args().size() + 2) {
MS_LOG(EXCEPTION) << "Size of cnode: " << cnode->DebugString()
<< " is not equal to 2 added to size of args: " << mindspore::ToString(partial_closure->args());
}
auto attrs = std::make_shared<AttrValueMap>();
for (size_t i = 0; i < partial_closure->args().size(); i++) {
auto old_node = cnode->input(i + 2);
auto possibile_value_node = BuildPossibleValueNode(old_node, partial_closure->args()[i], attrs);
if (possibile_value_node != nullptr) {
partial_node_list.push_back(possibile_value_node);
} else {
if (!(old_node->isa<CNode>() || old_node->isa<Parameter>())) {
MS_LOG(EXCEPTION) << "Old node should be CNode or Parameter, but " << old_node->ToString();
}
partial_node_list.push_back(old_node);
}
}
wrapped_node = new_node->func_graph()->NewCNode(partial_node_list);
wrapped_node->set_abstract(partial_closure);
}
return wrapped_node;
}
const EvaluatorCacheMapPtr &FuncGraphSpecializer::GetEvalCache(const EvaluatorPtr &eval) {
auto cache_iter = evalcaches_.find(eval);
if (cache_iter == evalcaches_.end()) {
evalcaches_[eval] = eval->cache();
return eval->cache();
}
return cache_iter->second;
}
std::pair<AbstractBasePtrList, AbstractBasePtr> FuncGraphSpecializer::BuildFromBroadedArgsVal(
const EvaluatorPtr &eval) {
MS_EXCEPTION_IF_NULL(eval);
std::unordered_set<AbstractBasePtrList, AbstractBasePtrListHasher, AbstractBasePtrListEqual> choices;
EvalResultPtr ret = nullptr;
AbstractBasePtrList broaded_argvals;
for (auto &argvals_map : *evalcaches_[eval]) {
auto argvals = argvals_map.first;
broaded_argvals.clear();
(void)std::transform(argvals.begin(), argvals.end(), std::back_inserter(broaded_argvals),
[](const AbstractBasePtr &arg) -> AbstractBasePtr { return arg->Broaden(); });
(void)choices.insert(broaded_argvals);
MS_LOG(DEBUG) << "Broaded_argvals: " << broaded_argvals.size() << ", " << ::mindspore::ToString(broaded_argvals);
}
if (1 == choices.size()) {
ConfigPtrList args_conf_list;
(void)std::transform(broaded_argvals.begin(), broaded_argvals.end(), std::back_inserter(args_conf_list),
[](AbstractBasePtr v) -> ConfigPtr { return std::make_shared<VirtualConfig>(v); });
// if broaden return null
ret = eval->Run(engine_, args_conf_list, nullptr);
EvaluatorCacheMapPtr real = std::make_shared<EvaluatorCacheMap>();
(*real)[broaded_argvals] = ret;
evalcaches_[eval] = real;
return std::make_pair(broaded_argvals, ret->abstract());
} else {
MS_LOG(DEBUG) << "Choices.size: " << choices.size();
return std::make_pair(AbstractBasePtrList(), nullptr);
}
}
void FuncGraphSpecializer::ProcessCNode(const CNodePtr &new_node) {
MS_EXCEPTION_IF_NULL(new_node);
if (specializer_->seen().count(new_node) > 0) {
return;
}
specializer_->AddSeen(new_node);
auto new_inputs = new_node->inputs();
if (new_inputs.empty()) {
MS_LOG(EXCEPTION) << "Inputs of CNode is empty";
}
AnfNodePtr func = new_inputs[0];
MS_EXCEPTION_IF_NULL(func);
// First element is func so arg start from 1
std::vector<AnfNodePtr> args(new_inputs.begin() + 1, new_inputs.end());
// CNode(CNode(Partial, f, arg1), arg2, ...) --> CNode(f, arg1, arg2, ...)
while (IsPrimitiveCNode(func, prim::kPrimPartial)) {
std::vector<AnfNodePtr> inputs = func->cast<CNodePtr>()->inputs();
// First element is partial, second is func so arg is start from 2
(void)args.insert(args.begin(), inputs.begin() + 2, inputs.end());
func = inputs[1];
}
new_inputs = args;
(void)new_inputs.insert(new_inputs.begin(), func);
AbstractBasePtrList argvals;
MS_EXCEPTION_IF_NULL(new_inputs[0]);
AbstractBasePtr fnval = new_inputs[0]->abstract();
MS_LOG(DEBUG) << "The new_inputs[0] node: pointer: " << new_inputs[0]->ToString() << ", "
<< new_inputs[0]->DebugString() << ", abstract: " << new_inputs[0]->abstract()->ToString();
// First element is func so function arguments start from 1
for (size_t i = 1; i < new_inputs.size(); ++i) {
argvals.push_back(new_inputs[i]->abstract());
MS_LOG(DEBUG) << "The new_inputs[" << i << "] node: pointer: " << new_inputs[i]->ToString() << ", "
<< new_inputs[i]->DebugString() << ", abstract: " << new_inputs[i]->abstract()->ToString();
}
if (!func->isa<ValueNode>()) {
MS_LOG(DEBUG) << func->abstract()->type_name() << " | " << func->abstract()->ToString();
if (func->abstract()->isa<AbstractFunction>() && !func->abstract()->isa<AbstractFuncUnion>()) {
auto func_abs = func->abstract()->cast<AbstractFunctionPtr>();
EvaluatorPtr eval = engine_->GetEvaluatorFor(func_abs);
std::pair<AbstractBasePtrList, AbstractBasePtr> result;
AbstractBasePtrList empty_args;
auto status = FindUniqueArgvals(func_abs, eval, empty_args, &result);
MS_LOG(DEBUG) << "FindUniqueArgvals return status: " << status;
// if a node is a poly node, or an input parameter is a PartialAbstractClosure, expand it early
if (status == kSpecializeFindUniqueArgvalPoly ||
(func->isa<Parameter>() && (func->func_graph()->has_flag(FUNC_GRAPH_FLAG_SPECIALIZE_PARAMETER) ||
func->abstract()->isa<PartialAbstractClosure>()))) {
auto wrapped_node = BuildSpecializedParameterNode(new_node);
new_inputs[0] = wrapped_node;
}
}
}
if (CanSpecializeNode(func)) {
// for primitive node , we build the primitive node with infered attributes in the first pass
// so we do not build replaced node again here in second pass
if (IsValueNode<Primitive>(func)) {
new_inputs[0] = func;
} else {
new_inputs[0] = BuildSpecializedNode(func, fnval, argvals);
}
}
for (size_t i = 0; i < argvals.size();) {
size_t next = i + 1;
if (CanSpecializeNode(args[i])) {
new_inputs[next] = BuildSpecializedNode(args[i], argvals[i], std::vector<AbstractBasePtr>{});
}
i = next;
}
new_node->set_inputs(new_inputs);
}
namespace {
void DumpEvaluatorCache(const EvaluatorCacheMap &evaluator_cache_map, const AbstractBasePtrList &argvals) {
MS_LOG(DEBUG) << "Find unique argvals failed: " << argvals.size() << ", " << argvals << ". Check cache all items.";
int i = 0;
for (const auto &item : evaluator_cache_map) {
MS_LOG(DEBUG) << "evaluator_cache_map[" << i++ << "]: " << item.first;
}
}
bool IsPolyFunc(const AbstractFunctionPtr &func, const AbstractBasePtrList &argvals) {
if (func->isa<PrimitiveAbstractClosure>() && argvals.empty()) {
MS_LOG(DEBUG) << "High order primitive return POLY.";
return true;
}
if (func->isa<MetaFuncGraphAbstractClosure>() && argvals.empty()) {
auto meta_func_graph_wrapper = dyn_cast<MetaFuncGraphAbstractClosure>(func);
auto meta_func_graph = meta_func_graph_wrapper->meta_func_graph();
if (meta_func_graph != nullptr && meta_func_graph->isa<prim::DoSignatureMetaFuncGraph>()) {
auto do_signature = dyn_cast<prim::DoSignatureMetaFuncGraph>(meta_func_graph);
if (do_signature != nullptr && do_signature->function()->isa<Primitive>()) {
MS_LOG(DEBUG) << "High order primitive " << do_signature->function()->ToString() << " return POLY.";
return true;
}
}
}
return false;
}
} // end anonymous namespace
SpecializeStatusCode FuncGraphSpecializer::FindUniqueArgvals(const AbstractFunctionPtr &func, const EvaluatorPtr &eval,
const AbstractBasePtrList &argvals,
std::pair<AbstractBasePtrList, AbstractBasePtr> *result) {
MS_EXCEPTION_IF_NULL(func);
MS_EXCEPTION_IF_NULL(eval);
MS_EXCEPTION_IF_NULL(result);
EvaluatorCacheMap evaluator_cache_map = *eval->cache();
if (evaluator_cache_map.find(argvals) != evaluator_cache_map.end()) {
*result = std::make_pair(argvals, evaluator_cache_map[argvals]->abstract());
return kSpecializeSuccess;
}
DumpEvaluatorCache(evaluator_cache_map, argvals);
const EvaluatorCacheMapPtr &choices = GetEvalCache(eval);
MS_EXCEPTION_IF_NULL(choices);
if (choices->count(argvals)) {
*result = std::make_pair(argvals, (*choices)[argvals]->abstract());
return kSpecializeSuccess;
} else if (choices->size() == 1) {
MS_LOG(DEBUG) << "Evaluator cache has a single item, just use it.";
*result = std::make_pair(choices->begin()->first, choices->begin()->second->abstract());
return kSpecializeSuccess;
} else if (choices->empty()) {
MS_LOG(DEBUG) << "Find DEAD code, it may be optimized in later phase.";
return kSpecializeFindUniqueArgvalDead;
} else {
if (IsPolyFunc(func, argvals)) {
return kSpecializeFindUniqueArgvalPoly;
}
MS_LOG(DEBUG) << "Try to find generalized argvals.";
*result = BuildFromBroadedArgsVal(eval);
if (!result->first.empty()) {
return kSpecializeSuccess;
}
MS_LOG(DEBUG) << "Find POLY code, it may be unused code or unresolved polymorphism.";
return kSpecializeFindUniqueArgvalPoly;
}
}
static PrimitivePtr BuildPrimtiveValueWithAttributes(const PrimitivePtr &prim, const AttrValueMapPtr &attrs) {
auto &prim_attrs = prim->attrs();
bool is_attr_same = true;
for (auto &item : *attrs) {
auto itr = prim_attrs.find(item.first);
if (itr != prim_attrs.end()) {
if (!(*(itr->second) == *(item.second))) {
is_attr_same = false;
break;
}
} else {
is_attr_same = false;
break;
}
}
if (!is_attr_same) {
if (prim->isa<PrimitivePy>()) {
PrimitivePyPtr prim_py = prim->cast<PrimitivePyPtr>();
auto clone_fn = prim_py->GetPyObj().attr("_clone");
py::object new_obj = clone_fn();
auto cloned_prim = new_obj.cast<PrimitivePyPtr>();
for (auto &item : *attrs) {
cloned_prim->AddAttr(item.first, item.second);
}
return cloned_prim;
}
auto cloned_prim = std::make_shared<Primitive>(*prim);
for (auto &item : *attrs) {
cloned_prim->AddAttr(item.first, item.second);
}
return cloned_prim;
}
return prim;
}
AnfNodePtr FuncGraphSpecializer::BuildPossibleValueNode(const AnfNodePtr &origin_node, const AbstractBasePtr &ival,
const AttrValueMapPtr &attrs) {
MS_EXCEPTION_IF_NULL(origin_node);
MS_EXCEPTION_IF_NULL(ival);
AbstractFunctionPtr abs = dyn_cast<AbstractFunction>(ival);
if (abs != nullptr) {
// Cannot build a determinstic ValueNode if there are multiple possible AbstractFunction.
if (abs->isa<AbstractFuncUnion>()) {
return nullptr;
}
ValuePtr value = nullptr;
if (abs->isa<PrimitiveAbstractClosure>()) {
auto real_fn = dyn_cast<PrimitiveAbstractClosure>(abs);
// for primitive, check if the attribute is the same with cnode infererd attribute ,if not, clone a new one
if (attrs != nullptr) {
value = BuildPrimtiveValueWithAttributes(real_fn->prim(), attrs);
} else {
value = real_fn->prim();
}
} else if (abs->isa<MetaFuncGraphAbstractClosure>()) {
auto real_fn = dyn_cast<MetaFuncGraphAbstractClosure>(abs);
value = real_fn->meta_func_graph();
} else if (abs->isa<FuncGraphAbstractClosure>()) {
auto real_fn = dyn_cast<FuncGraphAbstractClosure>(abs);
value = real_fn->func_graph();
} else {
return nullptr;
}
if (!value->isa<FuncGraph>() || value->cast<FuncGraphPtr>()->parent() == nullptr ||
(IsValueNode<FuncGraph>(origin_node) && IsVisible(func_graph_, value->cast<FuncGraphPtr>()->parent()))) {
return BuildValueNode(value, ival);
} else {
return nullptr;
}
} else {
ValuePtr val = ival->BuildValue();
if (val->isa<AnyValue>()) {
return nullptr;
}
// keep primitive 'depend' not to be optimized
if (IsPrimitiveCNode(origin_node, prim::kPrimDepend)) {
return nullptr;
}
return BuildValueNode(val, ival);
}
}
AnfNodeConfigPtr FuncGraphSpecializer::MakeConfig(const AnfNodePtr &node) {
return engine_->MakeConfig(node, context_);
}
} // namespace abstract
} // namespace mindspore
Reference in New Issue View Git Blame Copy Permalink