358 lines
11 KiB
C++
358 lines
11 KiB
C++
/**
|
|
* This is the C++ adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
|
|
*
|
|
* Copyright 2019-2020 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 "backend/optimizer/common/pattern_engine.h"
|
|
|
|
#include <exception>
|
|
#include <iostream>
|
|
#include <functional>
|
|
#include "frontend/optimizer/opt.h"
|
|
#include "ir/anf.h"
|
|
#include "utils/convert_utils_base.h"
|
|
#include "utils/overload.h"
|
|
|
|
namespace mindspore {
|
|
static int GetNextTag() {
|
|
static int kID = 0;
|
|
return kID++;
|
|
}
|
|
|
|
void Var::EnsureTag() {
|
|
if (tag_.length() == 0) {
|
|
std::ostringstream buffer;
|
|
buffer << "_" << GetNextTag();
|
|
tag_ = buffer.str();
|
|
}
|
|
}
|
|
|
|
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);
|
|
return *v1 == *v2;
|
|
}
|
|
|
|
if (lhs->isa<SeqVar>() && rhs->isa<SeqVar>()) {
|
|
SVarPtr v1 = dyn_cast<SeqVar>(lhs);
|
|
SVarPtr v2 = dyn_cast<SeqVar>(rhs);
|
|
return *v1 == *v2;
|
|
}
|
|
return (*lhs == *rhs);
|
|
}
|
|
|
|
std::string SeqVar::ToString() const {
|
|
std::ostringstream buffer;
|
|
buffer << "SeqVar(" << tag() << ", " << subvar_->ToString() << ")";
|
|
return buffer.str();
|
|
}
|
|
|
|
std::ostream &operator<<(std::ostream &os, const VarPtr &var) {
|
|
if (var == nullptr) {
|
|
os << "";
|
|
} else {
|
|
os << var->ToString();
|
|
}
|
|
return os;
|
|
}
|
|
|
|
template <>
|
|
std::ostream &operator<<<VarPtr, BaseRef>(std::ostream &os, const Equiv &equiv) {
|
|
os << "[Equiv]"
|
|
<< "\n";
|
|
for (auto &equiv_item : equiv) {
|
|
auto k = equiv_item.first;
|
|
os << k << ":";
|
|
BaseRef x = equiv_item.second;
|
|
if (utils::isa<AnfNodePtr>(x)) {
|
|
auto node = utils::cast<AnfNodePtr>(x);
|
|
os << "TypeString[" << node->type_name() << "]";
|
|
if (IsValueNode<FuncGraph>(node)) {
|
|
os << "IsValueNodeGraph ";
|
|
}
|
|
os << "type " << node->type_name();
|
|
if (node->isa<ValueNode>()) {
|
|
os << " value " << GetValueNode(node);
|
|
}
|
|
os << " addr: " << node;
|
|
} else if (utils::isa<Named>(x)) {
|
|
os << "Named " << x.ToString().c_str();
|
|
} else if (utils::isa<VarPtr>(x)) {
|
|
os << "TypeString[Var]";
|
|
os << utils::cast<VarPtr>(x);
|
|
} else if (utils::isa<FuncGraphPtr>(x)) {
|
|
os << "TypeString[Graph]";
|
|
}
|
|
os << "\n";
|
|
}
|
|
return os;
|
|
}
|
|
|
|
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);
|
|
MS_LOG(DEBUG) << "TypeString [" + node->type_name() + "]";
|
|
if (node->isa<VarNode>()) {
|
|
MS_LOG(DEBUG) << "IsVarNode " + node->cast<VarNodePtr>()->var_->ToString();
|
|
return node->cast<VarNodePtr>()->var_;
|
|
}
|
|
if (node->isa<ValueNode>()) {
|
|
MS_LOG(DEBUG) << "value " + GetValueNode(node)->ToString() + " addr: " + node->ToString();
|
|
} else {
|
|
MS_LOG(DEBUG) << "type " + node->type_name();
|
|
}
|
|
} else if (utils::isa<Named>(x)) {
|
|
MS_LOG(DEBUG) << "Named " + x.ToString();
|
|
} else if (utils::isa<VectorRef>(x)) {
|
|
MS_LOG(DEBUG) << "VectorRef";
|
|
} else if (utils::isa<VarPtr>(x)) {
|
|
MS_LOG(DEBUG) << "TypeString[Var] " + x.ToString();
|
|
}
|
|
MS_LOG(DEBUG) << "GetVar end: " + x.ToString();
|
|
return x;
|
|
}
|
|
|
|
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)) {
|
|
VarPtr var = utils::cast<VarPtr>(pattern);
|
|
if (var->matches(expr)) {
|
|
(*equiv)[var] = expr;
|
|
MS_LOG(DEBUG) << "pattern is var match: " + pattern.ToString() + ", " + expr.ToString();
|
|
return equiv;
|
|
}
|
|
}
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
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;
|
|
*values_expr = expr_ref;
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
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) {
|
|
values.push_back(GetVar(u));
|
|
return u;
|
|
};
|
|
return fn;
|
|
};
|
|
|
|
visitor_->SetFn(appender_pattern(*values_pattern));
|
|
MS_LOG(DEBUG) << "visit pattern_ref";
|
|
bool success = visitor_->Visit(pattern_ref, nullptr);
|
|
if (!success) {
|
|
return false;
|
|
}
|
|
|
|
auto appender_expr = [](VectorRef &values) {
|
|
std::function<BaseRef(const BaseRef &)> fn = [&](const BaseRef &u) {
|
|
values.push_back(u);
|
|
return u;
|
|
};
|
|
return fn;
|
|
};
|
|
|
|
visitor_->SetFn(appender_expr(*values_expr));
|
|
MS_LOG(DEBUG) << "visit expr_ref";
|
|
return visitor_->Visit(expr_ref, nullptr);
|
|
}
|
|
|
|
static int GetSVarStartIndex(const VectorRef &values) {
|
|
int index = -1;
|
|
int count = 0;
|
|
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";
|
|
return kInvalidVarIndex;
|
|
}
|
|
index = count;
|
|
}
|
|
count++;
|
|
}
|
|
return index;
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
size_t values_pattern_len = values_pattern.size();
|
|
size_t values_expr_len = values_expr.size();
|
|
|
|
if (svar_index == -1) {
|
|
if (values_pattern_len != values_expr_len) {
|
|
MS_LOG(DEBUG) << "Structures of differing size: pattern len " << values_pattern_len << ", expr len "
|
|
<< values_expr_len;
|
|
return nullptr;
|
|
}
|
|
}
|
|
if (values_expr_len < values_pattern_len - 1) {
|
|
MS_LOG(DEBUG) << "invalid size: pattern len " << values_pattern_len << ", expr len " << values_expr_len;
|
|
return nullptr;
|
|
}
|
|
size_t diff = values_expr_len - values_pattern_len + 1;
|
|
for (size_t i = 0; i < values_pattern_len; i++) {
|
|
size_t expr_i = i;
|
|
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, 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], primitive_vars, equiv);
|
|
}
|
|
if (equiv == nullptr) {
|
|
return nullptr;
|
|
}
|
|
}
|
|
return equiv;
|
|
}
|
|
|
|
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);
|
|
MS_LOG(DEBUG) << "GetVar v";
|
|
BaseRef expr_ref = expr;
|
|
|
|
if (equiv == nullptr) {
|
|
MS_LOG(EXCEPTION) << "Equiv pointer is null";
|
|
}
|
|
|
|
MS_LOG(DEBUG) << "Pattern ref " + pattern_ref.ToString() + ", expr ref" + expr_ref.ToString();
|
|
// 1. if pattern_ref is var and already in equiv, replace it.
|
|
if (utils::isa<VarPtr>(pattern_ref)) {
|
|
VarPtr var = utils::cast<VarPtr>(pattern_ref);
|
|
auto iter = equiv->find(var);
|
|
if (iter != equiv->end()) {
|
|
pattern_ref = iter->second;
|
|
}
|
|
}
|
|
|
|
// 2. check equal
|
|
if (eq_(pattern_ref, expr_ref)) {
|
|
return equiv;
|
|
}
|
|
|
|
// 3. match var
|
|
EquivPtr ret_equiv = MatchOnVar(pattern_ref, expr_ref, equiv);
|
|
if (ret_equiv) {
|
|
return ret_equiv;
|
|
}
|
|
|
|
// 4. here the type can be std:vector, std:list,
|
|
// or cnode.
|
|
if (!type_eq_(pattern_ref, expr_ref)) {
|
|
MS_LOG(DEBUG) << "Type mismatch";
|
|
return nullptr;
|
|
}
|
|
|
|
// 5. transfer the Containers by visitor to std::vector
|
|
VectorRef values_pattern;
|
|
VectorRef values_expr;
|
|
if (!ToVector(pattern_ref, expr_ref, &values_pattern, &values_expr)) {
|
|
return nullptr;
|
|
}
|
|
|
|
// 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
|
|
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 {
|
|
MS_EXCEPTION_IF_NULL(equiv);
|
|
MS_LOG(DEBUG) << "-----[in Replace]";
|
|
BaseRef ref = GetVar(pattern);
|
|
BaseRef out;
|
|
bool is_match = false;
|
|
|
|
// w is var
|
|
if (utils::isa<VarPtr>(ref)) {
|
|
const VarPtr &var = utils::cast<VarPtr>(ref);
|
|
auto iter = equiv->find(var);
|
|
if (iter != equiv->end()) {
|
|
out = iter->second;
|
|
is_match = true;
|
|
}
|
|
}
|
|
if (is_match) {
|
|
return out;
|
|
}
|
|
|
|
// visitor to visit the list
|
|
std::function<BaseRef(BaseRef)> fn = [&, this, equiv](const BaseRef &u) { return Replace(u, equiv); };
|
|
|
|
visitor_->SetFn(fn);
|
|
BaseRef visit_out;
|
|
if (!visitor_->Visit(pattern, &visit_out)) {
|
|
return pattern;
|
|
}
|
|
return visit_out;
|
|
}
|
|
} // namespace mindspore
|