p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!289 Add cnode mapping after graph match 

Merge pull request !289 from YuJianfeng/find_op
This commit is contained in:
mindspore-ci-bot 2020-04-14 18:51:28 +08:00 committed by Gitee
parent 86911b7cdc e5c67b9088
commit b571fabd77
11 changed files with 182 additions and 186 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
mindspore/ccsrc/pre_activate/ascend/ir_fusion/adam_apply_one_fusion.cc
View File

@ -15,43 +15,9 @@
*/
#include "pre_activate/ascend/ir_fusion/adam_apply_one_fusion.h"
#include "pre_activate/common/helper.h"
#include "utils/utils.h"
namespace mindspore {
namespace opt {
namespace {
void GetAdd0AndAdd1(const AnfNodePtr &sub0, AnfNodePtr *add0, AnfNodePtr *add1) {
MS_EXCEPTION_IF_NULL(sub0);
MS_EXCEPTION_IF_NULL(add0);
MS_EXCEPTION_IF_NULL(add1);
auto sub0_cnode = sub0->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(sub0_cnode);
CheckCNodeInputSize(sub0_cnode, kSubInputNum);
AnfNodePtr mul4 = sub0_cnode->input(2);
MS_EXCEPTION_IF_NULL(mul4);
auto mul4_cnode = mul4->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(mul4_cnode);
CheckCNodeInputSize(mul4_cnode, kMulInputNum);
AnfNodePtr true_div0 = mul4_cnode->input(2);
MS_EXCEPTION_IF_NULL(true_div0);
auto true_div0_cnode = true_div0->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(true_div0_cnode);
CheckCNodeInputSize(true_div0_cnode, kRealDivInputNum);
*add0 = true_div0_cnode->input(1);
AnfNodePtr add2 = true_div0_cnode->input(2);
MS_EXCEPTION_IF_NULL(add2);
auto add2_cnode = add2->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(add2_cnode);
CheckCNodeInputSize(add2_cnode, kAddInputNum);
AnfNodePtr sqrt0 = add2_cnode->input(1);
MS_EXCEPTION_IF_NULL(sqrt0);
auto sqrt0_cnode = sqrt0->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(sqrt0_cnode);
CheckCNodeInputSize(sqrt0_cnode, kSqrtInputNum);
*add1 = sqrt0_cnode->input(1);
}
} // namespace
AnfNodePtr AdamApplyOneFusion::CreateAdamApplyOneNode(const FuncGraphPtr &func_graph, const EquivPtr &equiv) const {
MS_EXCEPTION_IF_NULL(func_graph);
MS_EXCEPTION_IF_NULL(equiv);
@ -79,10 +45,10 @@ const BaseRef AdamApplyOneFusion::DefinePattern() const {
const auto prim_deal_div = std::make_shared<Primitive>(kRealDivOpName);
VectorRef mul2 = VectorRef({prim::kPrimMul, mul_x_input_vars_[2], input_vars_[1]});
VectorRef mul3 = VectorRef({prim::kPrimMul, mul_x_input_vars_[3], VectorRef({prim::kPrimSquare, input_vars_[0]})});
VectorRef sqrt0 = VectorRef({prim_sqrt, VectorRef({prim::kPrimTensorAdd, mul2, mul3})});
VectorRef sqrt0 = VectorRef({prim_sqrt, VectorRef({add1_var_, mul2, mul3})});
VectorRef mul1 = VectorRef({prim::kPrimMul, mul_x_input_vars_[1], input_vars_[0]});
VectorRef mul0 = VectorRef({prim::kPrimMul, mul_x_input_vars_[0], input_vars_[2]});
VectorRef add0 = VectorRef({prim::kPrimTensorAdd, mul0, mul1});
VectorRef add0 = VectorRef({add0_var_, mul0, mul1});
VectorRef true_div0 = VectorRef({prim_deal_div, add0, VectorRef({prim::kPrimTensorAdd, sqrt0, add2_y_})});
return VectorRef({prim::kPrimSub, input_vars_[3], VectorRef({prim::kPrimMul, input_vars_[4], true_div0})});
}
@ -96,10 +62,17 @@ const AnfNodePtr AdamApplyOneFusion::Process(const FuncGraphPtr &func_graph, con
new_node->set_scope(node->scope());
// Set abstract of new node
AbstractBasePtrList new_node_abstract_list;
AnfNodePtr add0 = nullptr;
AnfNodePtr add1 = nullptr;
GetAdd0AndAdd1(node, &add0, &add1);
auto iter_add0 = (*equiv).find(add0_var_);
if (iter_add0 == (*equiv).end()) {
MS_LOG(EXCEPTION) << "The equiv map is expected to contains the add0 var after matched.";
}
auto iter_add1 = (*equiv).find(add1_var_);
if (iter_add1 == (*equiv).end()) {
MS_LOG(EXCEPTION) << "The equiv map is expected to contains the add1 var after matched.";
}
auto add0 = utils::cast<AnfNodePtr>(iter_add0->second);
MS_EXCEPTION_IF_NULL(add0);
auto add1 = utils::cast<AnfNodePtr>(iter_add1->second);
MS_EXCEPTION_IF_NULL(add1);
new_node_abstract_list.push_back(add1->abstract());
new_node_abstract_list.push_back(add0->abstract());

5
mindspore/ccsrc/pre_activate/ascend/ir_fusion/adam_apply_one_fusion.h
View File

@ -19,6 +19,7 @@
#include <vector>
#include <memory>
#include "pre_activate/common/optimizer.h"
#include "utils/utils.h"
namespace mindspore {
namespace opt {
@ -35,6 +36,8 @@ class AdamApplyOneFusion : public PatternProcessPass {
mul_x_input_vars_.push_back(std::make_shared<Var>());
}
add2_y_ = std::make_shared<Var>();
add0_var_ = std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimTensorAdd->name()));
add1_var_ = std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimTensorAdd->name()));
}
~AdamApplyOneFusion() override = default;
@ -46,6 +49,8 @@ class AdamApplyOneFusion : public PatternProcessPass {
std::vector<VarPtr> input_vars_;
std::vector<VarPtr> mul_x_input_vars_;
VarPtr add2_y_;
VarPtr add0_var_;
VarPtr add1_var_;
};
} // namespace opt
} // namespace mindspore

54
mindspore/ccsrc/pre_activate/ascend/ir_fusion/adam_apply_one_with_decay_rule.cc
View File

@ -17,48 +17,13 @@
#include <memory>
#include <vector>
#include <tuple>
#include "session/anf_runtime_algorithm.h"
#include "ir/primitive.h"
#include "utils/utils.h"
#include "pre_activate/common/helper.h"
namespace mindspore {
namespace opt {
namespace {
std::tuple<AnfNodePtr, AnfNodePtr> GetAdd0Add1Node(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
auto sub0 = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(sub0);
auto mul5_anf = sub0->input(2);
MS_EXCEPTION_IF_NULL(mul5_anf);
auto mul5 = mul5_anf->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(mul5);
auto add3_anf = mul5->input(2);
MS_EXCEPTION_IF_NULL(add3_anf);
auto add3 = add3_anf->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(add3);
auto real_div0_anf = add3->input(1);
MS_EXCEPTION_IF_NULL(real_div0_anf);
auto real_div0 = real_div0_anf->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(real_div0);
auto add0_anf = real_div0->input(1);
MS_EXCEPTION_IF_NULL(add0_anf);
auto add2_anf = real_div0->input(2);
MS_EXCEPTION_IF_NULL(add2_anf);
auto add2 = add2_anf->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(add2);
auto sqrt0_anf = add2->input(1);
MS_EXCEPTION_IF_NULL(sqrt0_anf);
auto sqrt0 = sqrt0_anf->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(sqrt0);
auto add1_anf = sqrt0->input(1);
MS_EXCEPTION_IF_NULL(add1_anf);
return std::make_tuple(add0_anf, add1_anf);
}
} // namespace
std::vector<AnfNodePtr> AdamApplyOneWithDecayRule::GetFusionNodeInputs(const EquivPtr &equiv) const {
MS_EXCEPTION_IF_NULL(equiv);
auto input0 = utils::cast<AnfNodePtr>((*equiv)[input0_]);
@ -82,10 +47,10 @@ const BaseRef AdamApplyOneWithDecayRule::DefinePattern() const {
VectorRef mul0_pattern({prim::kPrimMul, mul0_x_, input2_});
VectorRef mul1_pattern({prim::kPrimMul, mul1_x_, input0_});
VectorRef square0_pattern({prim::kPrimSquare, input0_});
VectorRef add0_pattern({prim::kPrimTensorAdd, mul0_pattern, mul1_pattern});
VectorRef add0_pattern({add0_var_, mul0_pattern, mul1_pattern});
VectorRef mul2_pattern({prim::kPrimMul, mul2_x_, input1_});
VectorRef mul3_pattern({prim::kPrimMul, mul3_x_, square0_pattern});
VectorRef add1_pattern({prim::kPrimTensorAdd, mul2_pattern, mul3_pattern});
VectorRef add1_pattern({add1_var_, mul2_pattern, mul3_pattern});
VectorRef sqrt0_pattern({sqrt, add1_pattern});
VectorRef add2_pattern({prim::kPrimTensorAdd, sqrt0_pattern, add2_y_});
VectorRef mul4_pattern({prim::kPrimMul, mul4_x_, input3_});
@ -107,9 +72,18 @@ const AnfNodePtr AdamApplyOneWithDecayRule::Process(const FuncGraphPtr &graph, c
MS_EXCEPTION_IF_NULL(fusion_node);
fusion_node->set_scope(node->scope());
AnfNodePtr add0 = nullptr;
AnfNodePtr add1 = nullptr;
std::tie(add0, add1) = GetAdd0Add1Node(node);
auto iter_add0 = (*equiv).find(add0_var_);
if (iter_add0 == (*equiv).end()) {
MS_LOG(EXCEPTION) << "The equiv map is expected to contains the add0 var after matched.";
}
auto iter_add1 = (*equiv).find(add1_var_);
if (iter_add1 == (*equiv).end()) {
MS_LOG(EXCEPTION) << "The equiv map is expected to contains the add1 var after matched.";
}
auto add0 = utils::cast<AnfNodePtr>(iter_add0->second);
MS_EXCEPTION_IF_NULL(add0);
auto add1 = utils::cast<AnfNodePtr>(iter_add1->second);
MS_EXCEPTION_IF_NULL(add1);
auto types = {AnfAlgo::GetOutputInferDataType(add1, 0), AnfAlgo::GetOutputInferDataType(add0, 0),
AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {AnfAlgo::GetOutputInferShape(add1, 0), AnfAlgo::GetOutputInferShape(add0, 0),

5
mindspore/ccsrc/pre_activate/ascend/ir_fusion/adam_apply_one_with_decay_rule.h
View File

@ -19,6 +19,7 @@
#include <vector>
#include <memory>
#include "pre_activate/common/optimizer.h"
#include "utils/utils.h"
namespace mindspore {
namespace opt {
class AdamApplyOneWithDecayRule : public PatternProcessPass {
@ -36,6 +37,8 @@ class AdamApplyOneWithDecayRule : public PatternProcessPass {
mul3_x_ = std::make_shared<Var>();
mul4_x_ = std::make_shared<Var>();
add2_y_ = std::make_shared<Var>();
add0_var_ = std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimTensorAdd->name()));
add1_var_ = std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimTensorAdd->name()));
}
~AdamApplyOneWithDecayRule() override = default;
const BaseRef DefinePattern() const override;
@ -54,6 +57,8 @@ class AdamApplyOneWithDecayRule : public PatternProcessPass {
VarPtr mul3_x_;
VarPtr mul4_x_;
VarPtr add2_y_;
VarPtr add0_var_;
VarPtr add1_var_;
};
} // namespace opt
} // namespace mindspore

35
mindspore/ccsrc/pre_activate/ascend/ir_fusion/lamb_next_right_rule.cc
View File

@ -16,36 +16,9 @@
#include "pre_activate/ascend/ir_fusion/lamb_next_right_rule.h"
#include <vector>
#include "pre_activate/common/helper.h"
#include "utils/utils.h"
namespace mindspore {
namespace opt {
namespace {
AnfNodePtr GetAdd1Node(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
auto add2_cnode = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(add2_cnode);
if (add2_cnode->inputs().size() != kAddInputNum) {
MS_LOG(ERROR) << "The input size of Add2 is not equal to " << kAddInputNum;
}
AnfNodePtr sqrt0 = add2_cnode->input(1);
MS_EXCEPTION_IF_NULL(sqrt0);
auto sqrt0_cnode = sqrt0->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(sqrt0_cnode);
if (sqrt0_cnode->inputs().size() != kSqrtInputNum) {
MS_LOG(ERROR) << "The input size of Sqrt0 is not equal to " << kSqrtInputNum;
}
AnfNodePtr real_div1 = sqrt0_cnode->input(1);
MS_EXCEPTION_IF_NULL(real_div1);
auto real_div1_cnode = real_div1->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(real_div1_cnode);
if (real_div1_cnode->inputs().size() != kMulInputNum) {
MS_LOG(ERROR) << "The input size of RealDiv1 is not equal to " << kMulInputNum;
}
return real_div1_cnode->input(1);
}
} // namespace
AnfNodePtr LambNextRightRule::CreateLambNextRightNode(const FuncGraphPtr &func_graph, const EquivPtr &equiv) const {
MS_EXCEPTION_IF_NULL(func_graph);
MS_EXCEPTION_IF_NULL(equiv);
@ -79,7 +52,7 @@ const BaseRef LambNextRightRule::DefinePattern() const {
const auto prim_sqrt = std::make_shared<Primitive>(kSqrtOpName);
MS_EXCEPTION_IF_NULL(prim_sqrt);
VectorRef mul3 = VectorRef({prim::kPrimMul, mul3_x_, VectorRef({prim::kPrimSquare, input0_})});
VectorRef add1 = VectorRef({prim::kPrimTensorAdd, VectorRef({prim::kPrimMul, mul2_x_, input1_}), mul3});
VectorRef add1 = VectorRef({add1_var_, VectorRef({prim::kPrimMul, mul2_x_, input1_}), mul3});
return VectorRef(
{prim::kPrimTensorAdd, VectorRef({prim_sqrt, VectorRef({prim::kPrimMul, add1, true_div1_recip_})}), add2_y_});
}
@ -91,7 +64,11 @@ const AnfNodePtr LambNextRightRule::Process(const FuncGraphPtr &func_graph, cons
auto new_node = CreateLambNextRightNode(func_graph, equiv);
MS_EXCEPTION_IF_NULL(new_node);
// Set abstract of new node
AnfNodePtr add1 = GetAdd1Node(node);
auto iter_add1 = (*equiv).find(add1_var_);
if (iter_add1 == (*equiv).end()) {
MS_LOG(EXCEPTION) << "The equiv map is expected to contains the add1 var after matched.";
}
auto add1 = utils::cast<AnfNodePtr>(iter_add1->second);
MS_EXCEPTION_IF_NULL(add1);
AbstractBasePtrList new_node_abstract_list;
new_node_abstract_list.push_back(add1->abstract());

6
mindspore/ccsrc/pre_activate/ascend/ir_fusion/lamb_next_right_rule.h
View File

@ -18,6 +18,8 @@
#include <memory>
#include "pre_activate/common/optimizer.h"
#include "utils/utils.h"
namespace mindspore {
namespace opt {
class LambNextRightRule : public PatternProcessPass {
@ -29,7 +31,8 @@ class LambNextRightRule : public PatternProcessPass {
mul2_x_(std::make_shared<Var>()),
mul3_x_(std::make_shared<Var>()),
true_div1_recip_(std::make_shared<Var>()),
add2_y_(std::make_shared<Var>()) {}
add2_y_(std::make_shared<Var>()),
add1_var_(std::make_shared<Var>(std::make_shared<Primitive>(prim::kPrimTensorAdd->name()))) {}
~LambNextRightRule() override = default;
const BaseRef DefinePattern() const override;
@ -44,6 +47,7 @@ class LambNextRightRule : public PatternProcessPass {
VarPtr mul3_x_;
VarPtr true_div1_recip_;
VarPtr add2_y_;
VarPtr add1_var_;
};
} // namespace opt
} // namespace mindspore

30
mindspore/ccsrc/pre_activate/common/optimizer.cc
View File

@ -30,7 +30,8 @@
namespace mindspore {
namespace opt {
namespace {
AnfNodePtr HandleSexpVector(const BaseRef &sexp, const BaseRef &graph, bool multigraph);
AnfNodePtr HandleSexpVector(const BaseRef &sexp, const BaseRef &graph, PrimitiveVarMap *primitive_vars,
bool multigraph);
ValueNodePtr CreateValueNodeWithSexp(const BaseRef &sexp) {
if (utils::isa<int>(sexp)) {
@ -71,12 +72,20 @@ VarNodePtr CreateVarNodeWithSexp(const BaseRef &sexp, const BaseRef &graph) {
return nullptr;
}
AnfNodePtr SexpToNode(const BaseRef &sexp, const BaseRef &graph, bool multigraph = false) {
AnfNodePtr SexpToNode(const BaseRef &sexp, const BaseRef &graph, PrimitiveVarMap *primitive_vars,
bool multigraph = false) {
MS_LOG(DEBUG) << "SexpToNode sexp: " + sexp.ToString() + ", graph " + graph.ToString();
MS_EXCEPTION_IF_NULL(primitive_vars);
if (utils::isa<VectorRef>(sexp)) {
return HandleSexpVector(sexp, graph, multigraph);
return HandleSexpVector(sexp, graph, primitive_vars, multigraph);
}
if (utils::isa<VarPtr>(sexp)) {
auto var_ptr = utils::cast<VarPtr>(sexp);
MS_EXCEPTION_IF_NULL(var_ptr);
if (var_ptr->primitive()) {
(*primitive_vars)[var_ptr->primitive()] = var_ptr;
return NewValueNode(var_ptr->primitive());
}
return CreateVarNodeWithSexp(sexp, graph);
}
if (utils::isa<AnfNodePtr>(sexp)) {
@ -89,13 +98,14 @@ AnfNodePtr SexpToNode(const BaseRef &sexp, const BaseRef &graph, bool multigraph
return value_node;
}
AnfNodePtr HandleSexpVector(const BaseRef &sexp, const BaseRef &graph, bool multigraph) {
AnfNodePtr HandleSexpVector(const BaseRef &sexp, const BaseRef &graph, PrimitiveVarMap *primitive_vars,
bool multigraph) {
MS_LOG(DEBUG) << "HandleSexpVector sexp: " + sexp.ToString() + ", graph " + graph.ToString();
std::vector<AnfNodePtr> input_nodes;
const auto &tuple = utils::cast<VectorRef>(sexp);
if (multigraph && utils::isa<VarPtr>(graph)) {
for (auto &x : tuple) {
AnfNodePtr node = SexpToNode(x, std::make_shared<Var>("G"), true);
AnfNodePtr node = SexpToNode(x, std::make_shared<Var>("G"), primitive_vars, true);
input_nodes.push_back(node);
}
VarPtr var_ptr = utils::cast<VarPtr>(graph);
@ -103,7 +113,7 @@ AnfNodePtr HandleSexpVector(const BaseRef &sexp, const BaseRef &graph, bool mult
}
for (auto &x : tuple) {
AnfNodePtr node = SexpToNode(x, graph, multigraph);
AnfNodePtr node = SexpToNode(x, graph, primitive_vars, multigraph);
input_nodes.push_back(node);
}
return CreateCNodeWithGraph(input_nodes, graph);
@ -166,7 +176,8 @@ PatternProcessPass::PatternProcessPass(const std::string &name, bool multigraph)
multigraph_(multigraph),
pattern_engine_(PatternEngine(std::make_shared<DefaultVisitor>(),
std::function<bool(const BaseRef &, const BaseRef &)>(AnfEqual),
std::function<bool(const BaseRef &, const BaseRef &)>(CNodeTypeEqual))) {}
std::function<bool(const BaseRef &, const BaseRef &)>(CNodeTypeEqual))),
primitive_vars_(std::make_shared<PrimitiveVarMap>()) {}
const BaseRef PatternProcessPass::DefinePattern() const {
VarPtr X = std::make_shared<Var>();
@ -176,7 +187,7 @@ const BaseRef PatternProcessPass::DefinePattern() const {
void PatternProcessPass::Build() {
VarPtr fg = std::make_shared<Var>("RootG");
BaseRef pattern = std::move(DefinePattern());
pattern_ = SexpToNode(pattern, fg, multigraph_);
pattern_ = SexpToNode(pattern, fg, primitive_vars_.get(), multigraph_);
}
AnfNodePtr PatternProcessPass::Run(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
@ -185,7 +196,8 @@ AnfNodePtr PatternProcessPass::Run(const FuncGraphPtr &func_graph, const AnfNode
}
auto empty_equiv = std::make_shared<Equiv>();
EquivPtr equiv = pattern_engine_.Match(pattern_, node, empty_equiv);
MS_EXCEPTION_IF_NULL(primitive_vars_);
EquivPtr equiv = pattern_engine_.Match(pattern_, node, *primitive_vars_, empty_equiv);
if (equiv != nullptr && !equiv->empty()) {
return Process(func_graph, node, equiv);
}

2
mindspore/ccsrc/pre_activate/common/optimizer.h
View File

@ -19,6 +19,7 @@
#include <memory>
#include <string>
#include <vector>
#include <unordered_map>
#include "ir/anf.h"
#include "ir/func_graph.h"
@ -46,6 +47,7 @@ class PatternProcessPass : public NodePass {
AnfNodePtr pattern_ = nullptr;
bool multigraph_ = true;
PatternEngine pattern_engine_;
PrimitiveVarMapPtr primitive_vars_;
};
class GraphOptimizer {

79
mindspore/ccsrc/pre_activate/common/pattern_engine.cc
View File

@ -42,7 +42,7 @@ void Var::EnsureTag() {
}
}
bool operator==(const VarPtr& lhs, const VarPtr& rhs) {
bool operator==(const VarPtr &lhs, const VarPtr &rhs) {
if (lhs->isa<CondVar>() && rhs->isa<CondVar>()) {
CondVarPtr v1 = dyn_cast<CondVar>(lhs);
CondVarPtr v2 = dyn_cast<CondVar>(rhs);
@ -63,7 +63,7 @@ std::string SeqVar::ToString() const {
return buffer.str();
}
std::ostream& operator<<(std::ostream& os, const VarPtr& var) {
std::ostream &operator<<(std::ostream &os, const VarPtr &var) {
if (var == nullptr) {
os << "";
} else {
@ -73,10 +73,10 @@ std::ostream& operator<<(std::ostream& os, const VarPtr& var) {
}
template <>
std::ostream& operator<<<VarPtr, BaseRef>(std::ostream& os, const Equiv& equiv) {
std::ostream &operator<<<VarPtr, BaseRef>(std::ostream &os, const Equiv &equiv) {
os << "[Equiv]"
<< "\n";
for (auto& equiv_item : equiv) {
for (auto &equiv_item : equiv) {
auto k = equiv_item.first;
os << k << ":";
BaseRef x = equiv_item.second;
@ -104,7 +104,7 @@ std::ostream& operator<<<VarPtr, BaseRef>(std::ostream& os, const Equiv& equiv)
return os;
}
static BaseRef GetVar(const BaseRef& x) {
static BaseRef GetVar(const BaseRef &x) {
MS_LOG(DEBUG) << "getVar start :%s" + x.ToString();
if (utils::isa<AnfNodePtr>(x)) {
auto node = utils::cast<AnfNodePtr>(x);
@ -129,7 +129,7 @@ static BaseRef GetVar(const BaseRef& x) {
return x;
}
EquivPtr MatchOnVar(const BaseRef& pattern, const BaseRef& expr, EquivPtr equiv) {
EquivPtr MatchOnVar(const BaseRef &pattern, const BaseRef &expr, EquivPtr equiv) {
MS_LOG(DEBUG) << "MatchOnVar pattern " + pattern.ToString() + " expr: " + expr.ToString();
MS_EXCEPTION_IF_NULL(equiv);
if (utils::isa<VarPtr>(pattern)) {
@ -144,8 +144,8 @@ EquivPtr MatchOnVar(const BaseRef& pattern, const BaseRef& expr, EquivPtr equiv)
return nullptr;
}
bool PatternEngine::ToVector(const VectorRef& pattern_ref, const VectorRef& expr_ref, VectorRef* const values_pattern,
VectorRef* const values_expr) const {
bool PatternEngine::ToVector(const VectorRef &pattern_ref, const VectorRef &expr_ref, VectorRef *const values_pattern,
VectorRef *const values_expr) const {
MS_EXCEPTION_IF_NULL(values_expr);
if (utils::isa<SeqPtr>(pattern_ref)) {
*values_pattern = pattern_ref;
@ -155,12 +155,12 @@ bool PatternEngine::ToVector(const VectorRef& pattern_ref, const VectorRef& expr
return false;
}
bool PatternEngine::ToVector(const BaseRef& pattern_ref, const BaseRef& expr_ref, VectorRef* const values_pattern,
VectorRef* const values_expr) const {
bool PatternEngine::ToVector(const BaseRef &pattern_ref, const BaseRef &expr_ref, VectorRef *const values_pattern,
VectorRef *const values_expr) const {
MS_EXCEPTION_IF_NULL(values_expr);
// visitor to visite the list
auto appender_pattern = [](VectorRef& values) {
std::function<BaseRef(const BaseRef&)> fn = [&](const BaseRef& u) {
auto appender_pattern = [](VectorRef &values) {
std::function<BaseRef(const BaseRef &)> fn = [&](const BaseRef &u) {
values.push_back(GetVar(u));
return u;
};
@ -174,8 +174,8 @@ bool PatternEngine::ToVector(const BaseRef& pattern_ref, const BaseRef& expr_ref
return false;
}
auto appender_expr = [](VectorRef& values) {
std::function<BaseRef(const BaseRef&)> fn = [&](const BaseRef& u) {
auto appender_expr = [](VectorRef &values) {
std::function<BaseRef(const BaseRef &)> fn = [&](const BaseRef &u) {
values.push_back(u);
return u;
};
@ -187,10 +187,10 @@ bool PatternEngine::ToVector(const BaseRef& pattern_ref, const BaseRef& expr_ref
return visitor_->Visit(expr_ref, nullptr);
}
static int GetSVarStartIndex(const VectorRef& values) {
static int GetSVarStartIndex(const VectorRef &values) {
int index = -1;
int count = 0;
for (auto& value : values) {
for (auto &value : values) {
if (utils::isa<VarPtr>(value) && utils::cast<VarPtr>(value)->isa<SeqVar>()) {
if (index != -1) {
MS_LOG(DEBUG) << "Multiple SVars in sequence";
@ -203,7 +203,35 @@ static int GetSVarStartIndex(const VectorRef& values) {
return index;
}
EquivPtr PatternEngine::AlignSVar(const VectorRef& values_pattern, const VectorRef& values_expr, EquivPtr equiv) const {
void UpdateEquivMap(const VectorRef &values_pattern, const BaseRef &expr_ref, const PrimitiveVarMap &primitive_vars,
EquivPtr equiv) {
if (equiv == nullptr || values_pattern.empty() || !utils::isa<AnfNodePtr>(values_pattern[0]) ||
!utils::isa<AnfNodePtr>(expr_ref)) {
return;
}
auto real_node = utils::cast<AnfNodePtr>(expr_ref);
MS_EXCEPTION_IF_NULL(real_node);
if (!real_node->isa<CNode>()) {
return;
}
auto prim_node = utils::cast<AnfNodePtr>(values_pattern[0]);
MS_EXCEPTION_IF_NULL(prim_node);
if (!IsValueNode<Primitive>(prim_node)) {
return;
}
ValuePtr value = GetValueNode(prim_node);
MS_EXCEPTION_IF_NULL(value);
auto prim = value->cast<PrimitivePtr>();
MS_EXCEPTION_IF_NULL(prim);
auto iter = primitive_vars.find(prim);
if (iter == primitive_vars.end()) {
return;
}
(*equiv)[iter->second] = real_node;
}
EquivPtr PatternEngine::AlignSVar(const VectorRef &values_pattern, const VectorRef &values_expr,
const PrimitiveVarMap &primitive_vars, EquivPtr equiv) const {
int svar_index = GetSVarStartIndex(values_pattern);
if (svar_index == kInvalidVarIndex) {
return nullptr;
@ -229,12 +257,12 @@ EquivPtr PatternEngine::AlignSVar(const VectorRef& values_pattern, const VectorR
if (svar_index != -1 && i == IntToSize(svar_index)) {
auto seq =
std::vector<BaseRef>(values_expr.begin() + svar_index, values_expr.begin() + svar_index + SizeToInt(diff));
equiv = Match(values_pattern[svar_index], seq, equiv);
equiv = Match(values_pattern[svar_index], seq, primitive_vars, equiv);
} else {
if (svar_index != -1 && i > IntToSize(svar_index)) {
expr_i = i + diff - 1;
}
equiv = Match(values_pattern[i], values_expr[expr_i], equiv);
equiv = Match(values_pattern[i], values_expr[expr_i], primitive_vars, equiv);
}
if (equiv == nullptr) {
return nullptr;
@ -243,7 +271,8 @@ EquivPtr PatternEngine::AlignSVar(const VectorRef& values_pattern, const VectorR
return equiv;
}
EquivPtr PatternEngine::Match(const BaseRef& pattern, const BaseRef& expr, EquivPtr equiv) const {
EquivPtr PatternEngine::Match(const BaseRef &pattern, const BaseRef &expr, const PrimitiveVarMap &primitive_vars,
EquivPtr equiv) const {
MS_LOG(DEBUG) << "-----[in Match]";
MS_LOG(DEBUG) << "GetVar w";
BaseRef pattern_ref = GetVar(pattern);
@ -292,10 +321,12 @@ EquivPtr PatternEngine::Match(const BaseRef& pattern, const BaseRef& expr, Equiv
// 6. if any svar in both side, find the SeqVar index,
// try to pack the Var s in std::vector to a Seq and match elements one by one.
// check svar
return AlignSVar(values_pattern, values_expr, equiv);
equiv = AlignSVar(values_pattern, values_expr, primitive_vars, equiv);
UpdateEquivMap(values_pattern, expr_ref, primitive_vars, equiv);
return equiv;
}
BaseRef PatternEngine::Replace(const BaseRef& pattern, const EquivPtr& equiv) const {
BaseRef PatternEngine::Replace(const BaseRef &pattern, const EquivPtr &equiv) const {
MS_EXCEPTION_IF_NULL(equiv);
MS_LOG(DEBUG) << "-----[in Replace]";
BaseRef ref = GetVar(pattern);
@ -304,7 +335,7 @@ BaseRef PatternEngine::Replace(const BaseRef& pattern, const EquivPtr& equiv) co
// w is var
if (utils::isa<VarPtr>(ref)) {
const VarPtr& var = utils::cast<VarPtr>(ref);
const VarPtr &var = utils::cast<VarPtr>(ref);
auto iter = equiv->find(var);
if (iter != equiv->end()) {
out = iter->second;
@ -316,7 +347,7 @@ BaseRef PatternEngine::Replace(const BaseRef& pattern, const EquivPtr& equiv) co
}
// visitor to visit the list
std::function<BaseRef(BaseRef)> fn = [&, this, equiv](const BaseRef& u) { return Replace(u, equiv); };
std::function<BaseRef(BaseRef)> fn = [&, this, equiv](const BaseRef &u) { return Replace(u, equiv); };
visitor_->SetFn(fn);
BaseRef visit_out;

71
mindspore/ccsrc/pre_activate/common/pattern_engine.h
View File

@ -31,6 +31,7 @@
#include <map>
#include <stdexcept>
#include <list>
#include <utility>
#include "pre_activate/common/visit.h"
#include "ir/base.h"
@ -44,16 +45,19 @@ using CondVarPtr = std::shared_ptr<CondVar>;
using SVarPtr = std::shared_ptr<SeqVar>;
const int kInvalidVarIndex = -2;
using ConditionFunc = std::function<bool(const BaseRef&)>;
using ConditionFunc = std::function<bool(const BaseRef &)>;
// Base wildcard variable which could match any anf node.
class Var : public Base {
friend class VarHasher;
public:
explicit Var(const std::string& tag = "") : tag_(tag) { EnsureTag(); }
Var(const Var& other) : Base(other), tag_(other.tag_) {}
virtual Var& operator=(const Var& other) {
explicit Var(std::string tag = "") : tag_(std::move(tag)), primitive_(nullptr) { EnsureTag(); }
explicit Var(const PrimitivePtr &primitive, std::string tag = "") : tag_(std::move(tag)), primitive_(primitive) {
EnsureTag();
}
Var(const Var &other) : Base(other), tag_(other.tag_) {}
virtual Var &operator=(const Var &other) {
if (&other == this) {
return *this;
}
@ -63,12 +67,13 @@ class Var : public Base {
~Var() override = default;
MS_DECLARE_PARENT(Var, Base);
virtual bool matches(const BaseRef&) { return true; }
virtual bool matches(const BaseRef &) { return true; }
virtual bool operator==(const Var& other) const { return tag_ == other.tag_; }
bool operator!=(const Var& other) const { return !(&other == this); }
virtual bool operator==(const Var &other) const { return tag_ == other.tag_; }
bool operator!=(const Var &other) const { return !(&other == this); }
std::string tag() const { return tag_; }
PrimitivePtr primitive() const { return primitive_; }
std::string ToString() const override {
std::ostringstream buffer;
buffer << "Var(" << tag_ << ")";
@ -80,12 +85,13 @@ class Var : public Base {
void EnsureTag();
std::string tag_;
PrimitivePtr primitive_;
};
// VarNode means variable node, a subclass of AnfNode
class VarNode : public AnfNode {
public:
VarNode(const VarPtr& value, const FuncGraphPtr& func_graph) : AnfNode(func_graph), var_(value) {}
VarNode(const VarPtr &value, const FuncGraphPtr &func_graph) : AnfNode(func_graph), var_(value) {}
~VarNode() override = default;
MS_DECLARE_PARENT(VarNode, AnfNode);
@ -95,16 +101,16 @@ using VarNodePtr = std::shared_ptr<VarNode>;
class VarHasher {
public:
std::size_t operator()(const Var& var) const { return var.hash(); }
std::size_t operator()(const Var &var) const { return var.hash(); }
};
// Condition Var, match an anf node when condition function return true.
class CondVar : public Var {
public:
explicit CondVar(const ConditionFunc& cond) : cond_fn_(cond) {}
explicit CondVar(const ConditionFunc &cond) : cond_fn_(cond) {}
~CondVar() override = default;
MS_DECLARE_PARENT(CondVar, Var);
bool matches(const BaseRef& value) override {
bool matches(const BaseRef &value) override {
MS_LOG(DEBUG) << "CondVarPtr match: " + value.ToString();
if (utils::isa<Var>(value)) {
return false;
@ -124,55 +130,60 @@ class SeqVar : public Var {
~SeqVar() override = default;
MS_DECLARE_PARENT(SeqVar, Var);
explicit SeqVar(const VarPtr subvar) : subvar_(nullptr) { subvar_ = subvar; }
bool matches(const BaseRef& value) override {
bool matches(const BaseRef &value) override {
// match Seq.
if (utils::isa<Seq>(value)) {
const Seq& seq = utils::cast<Seq>(value);
return std::all_of(seq.begin(), seq.end(), [this](const BaseRef& v) {
const Seq &seq = utils::cast<Seq>(value);
return std::all_of(seq.begin(), seq.end(), [this](const BaseRef &v) {
auto eq = subvar_->matches(v);
return eq;
});
}
return false;
}
bool operator==(const SeqVar& other) const { return *subvar_ == *other.subvar_; }
bool operator==(const SeqVar &other) const { return *subvar_ == *other.subvar_; }
std::string ToString() const override;
private:
VarPtr subvar_;
};
bool operator==(const VarPtr& lhs, const VarPtr& rhs);
bool operator==(const VarPtr &lhs, const VarPtr &rhs);
inline bool operator!=(const VarPtr& lhs, const VarPtr& rhs) { return !(lhs == rhs); }
inline bool operator!=(const VarPtr &lhs, const VarPtr &rhs) { return !(lhs == rhs); }
std::ostream& operator<<(std::ostream& os, const VarPtr& var);
std::ostream &operator<<(std::ostream &os, const VarPtr &var);
using Equiv = std::map<VarPtr, BaseRef>;
using EquivPtr = std::shared_ptr<Equiv>;
using PrimitiveVarMap = std::unordered_map<PrimitivePtr, VarPtr>;
using PrimitiveVarMapPtr = std::shared_ptr<PrimitiveVarMap>;
inline bool DefaultTypeEq(const BaseRef& x, const BaseRef& y) { return x.type() == y.type(); }
inline bool DefaultTypeEq(const BaseRef &x, const BaseRef &y) { return x.type() == y.type(); }
class PatternEngine {
public:
PatternEngine(const std::shared_ptr<Visitor>& visitor, const std::function<bool(const BaseRef&, const BaseRef&)>& eq,
const std::function<bool(const BaseRef&, const BaseRef&)>& type_eq = DefaultTypeEq)
PatternEngine(const std::shared_ptr<Visitor> &visitor,
const std::function<bool(const BaseRef &, const BaseRef &)> &eq,
const std::function<bool(const BaseRef &, const BaseRef &)> &type_eq = DefaultTypeEq)
: visitor_(visitor), eq_(eq), type_eq_(type_eq) {}
~PatternEngine() = default;
EquivPtr Match(const BaseRef& pattern, const BaseRef& expr, EquivPtr equiv) const;
EquivPtr Match(const BaseRef &pattern, const BaseRef &expr, const PrimitiveVarMap &primitive_vars,
EquivPtr equiv) const;
// Replace pattern with equivalent
BaseRef Replace(const BaseRef& pattern, const EquivPtr& equiv) const;
BaseRef Replace(const BaseRef &pattern, const EquivPtr &equiv) const;
private:
EquivPtr AlignSVar(const VectorRef& values_pattern, const VectorRef& values_expr, EquivPtr equiv) const;
bool ToVector(const BaseRef& pattern, const BaseRef& expr, VectorRef* const values_pattern,
VectorRef* const values_expr) const;
bool ToVector(const VectorRef& pattern_ref, const VectorRef& expr_ref, VectorRef* const values_pattern,
VectorRef* const values_expr) const;
EquivPtr AlignSVar(const VectorRef &values_pattern, const VectorRef &values_expr,
const PrimitiveVarMap &primitive_vars, EquivPtr equiv) const;
bool ToVector(const BaseRef &pattern, const BaseRef &expr, VectorRef *const values_pattern,
VectorRef *const values_expr) const;
bool ToVector(const VectorRef &pattern_ref, const VectorRef &expr_ref, VectorRef *const values_pattern,
VectorRef *const values_expr) const;
std::shared_ptr<Visitor> visitor_;
std::function<bool(const BaseRef&, const BaseRef&)> eq_;
std::function<bool(const BaseRef&, const BaseRef&)> type_eq_;
std::function<bool(const BaseRef &, const BaseRef &)> eq_;
std::function<bool(const BaseRef &, const BaseRef &)> type_eq_;
};
} // namespace mindspore
namespace std {

30
tests/ut/cpp/pre_activate/common/pattern_engine_test.cc
View File

@ -40,6 +40,7 @@ class TestMatchEngine : public UT::Common {
public:
PatternEngine TU;
EquivPtr equiv_null;
PrimitiveVarMap primitive_vars_null;
};
TEST_F(TestMatchEngine, Var) {
@ -106,30 +107,30 @@ TEST_F(TestMatchEngine, MatchRaw_Var) {
// common
equiv_null->clear();
d = TU.Match(v1, 1, equiv_null);
d = TU.Match(v1, 1, primitive_vars_null, equiv_null);
ASSERT_EQ((*d)[v1], 1);
equiv_null->clear();
(*equiv_null)[v1] = v2;
d = TU.Match(v1, 1, equiv_null);
d = TU.Match(v1, 1, primitive_vars_null, equiv_null);
ASSERT_EQ(d->count(v2), std::size_t(1));
ASSERT_EQ((*d)[v2], 1);
equiv_null->clear();
(*equiv_null)[v1] = v2;
(*equiv_null)[v3] = 1;
d = TU.Match(v1, 1, equiv_null);
d = TU.Match(v1, 1, primitive_vars_null, equiv_null);
ASSERT_EQ(d->count(v2), std::size_t(1));
ASSERT_EQ((*d)[v2], 1);
equiv_null->clear();
d = TU.Match(VectorRef({v1}), VectorRef({1}), equiv_null);
d = TU.Match(VectorRef({v1}), VectorRef({1}), primitive_vars_null, equiv_null);
ASSERT_EQ(d->size(), std::size_t(1));
ASSERT_EQ(d->count(v1), std::size_t(1));
ASSERT_EQ((*d)[v1], 1);
equiv_null->clear();
ASSERT_EQ(TU.Match(1, 2, equiv_null), nullptr);
ASSERT_EQ(TU.Match(1, 2, primitive_vars_null, equiv_null), nullptr);
}
TEST_F(TestMatchEngine, MatchRaw_SVar) {
@ -139,22 +140,22 @@ TEST_F(TestMatchEngine, MatchRaw_SVar) {
EquivPtr d;
equiv_null->clear();
d = TU.Match(VectorRef({sv1}), VectorRef({1, 2}), equiv_null);
d = TU.Match(VectorRef({sv1}), VectorRef({1, 2}), primitive_vars_null, equiv_null);
ASSERT_EQ(d->size(), std::size_t(1));
ASSERT_EQ(d->count(sv1), std::size_t(1));
ASSERT_EQ(utils::cast<Seq>((*d)[sv1]), Seq({1, 2}));
equiv_null->clear();
d = TU.Match(VectorRef({v1, sv1}), VectorRef({1, 2}), equiv_null);
d = TU.Match(VectorRef({v1, sv1}), VectorRef({1, 2}), primitive_vars_null, equiv_null);
ASSERT_EQ(d->size(), std::size_t(2));
ASSERT_EQ(utils::cast<Seq>((*d)[sv1]), Seq({2}));
equiv_null->clear();
ASSERT_EQ(TU.Match(VectorRef({sv1, sv2}), VectorRef({1, 2}), equiv_null), nullptr);
ASSERT_EQ(TU.Match(VectorRef({sv1, sv2}), VectorRef({1, 2}), primitive_vars_null, equiv_null), nullptr);
equiv_null->clear();
(*equiv_null)[sv1] = std::make_shared<Seq>(PatternListType{1, 2});
d = TU.Match(VectorRef({v1, sv1}), VectorRef({1, 1, 2}), equiv_null);
d = TU.Match(VectorRef({v1, sv1}), VectorRef({1, 1, 2}), primitive_vars_null, equiv_null);
ASSERT_EQ(d->size(), std::size_t(2));
ASSERT_EQ((*d)[v1], 1);
}
@ -167,13 +168,13 @@ TEST_F(TestMatchEngine, Match) {
EquivPtr d;
equiv_null->clear();
d = TU.Match(VectorRef({v1, v1, v2}), VectorRef({1, 1, 2}), equiv_null);
d = TU.Match(VectorRef({v1, v1, v2}), VectorRef({1, 1, 2}), primitive_vars_null, equiv_null);
ASSERT_EQ(d->size(), std::size_t(2));
ASSERT_EQ((*d)[v1], 1);
ASSERT_EQ((*d)[v2], 2);
equiv_null->clear();
d = TU.Match(static_cast<int>(1), static_cast<float>(1), equiv_null);
d = TU.Match(static_cast<int>(1), static_cast<float>(1), primitive_vars_null, equiv_null);
ASSERT_EQ(d, nullptr);
}
@ -197,18 +198,19 @@ TEST_F(TestMatchEngine, Match_CondVar) {
EquivPtr d;
equiv_null->clear();
d = TU.Match(VectorRef({vf, vn}), VectorRef({static_cast<float>(1.0), -1}), equiv_null);
d = TU.Match(VectorRef({vf, vn}), VectorRef({static_cast<float>(1.0), -1}), primitive_vars_null, equiv_null);
ASSERT_GE(d->size(), std::size_t(0));
auto vfn = (*d)[vf];
ASSERT_EQ((*d)[vf], static_cast<float>(1.0));
ASSERT_EQ((*d)[vn], -1);
equiv_null->clear();
d = TU.Match(VectorRef({vf, vn}), VectorRef({1, static_cast<float>(-1.0)}), equiv_null);
d = TU.Match(VectorRef({vf, vn}), VectorRef({1, static_cast<float>(-1.0)}), primitive_vars_null, equiv_null);
ASSERT_EQ(d, nullptr);
equiv_null->clear();
d = TU.Match(VectorRef({vf, vn}), VectorRef({static_cast<float>(1.0), static_cast<int>(1)}), equiv_null);
d = TU.Match(VectorRef({vf, vn}), VectorRef({static_cast<float>(1.0), static_cast<int>(1)}), primitive_vars_null,
equiv_null);
ASSERT_EQ(d, nullptr);
}