wehub-resource-sync
154 lines
4.3 KiB
C++
154 lines
4.3 KiB
C++
//
|
|
// FuseGeLu.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/04/28.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <unordered_map>
|
|
|
|
#include "../TemplateMerge.hpp"
|
|
#include "MNN/expr/ExprCreator.hpp"
|
|
#include "MNN_generated.h"
|
|
#include "MergeHelpers.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class FuseGeLu {
|
|
public:
|
|
FuseGeLu();
|
|
private:
|
|
VARP gelu_inputs_;
|
|
};
|
|
|
|
bool IsBinaryPow(EXPRP expr) {
|
|
// return true;
|
|
const Op* op = expr->get();
|
|
if (op == nullptr || op->type() != OpType_BinaryOp) {
|
|
return false;
|
|
}
|
|
return op->main_as_BinaryOp()->opType() == BinaryOpOperation_POW;
|
|
}
|
|
|
|
bool IsTanh(EXPRP expr) {
|
|
const Op* op = expr->get();
|
|
if (op == nullptr) {
|
|
return false;
|
|
}
|
|
if (op->type() == OpType_TanH) {
|
|
return true;
|
|
}
|
|
if (op->type() == OpType_UnaryOp && op->main_as_UnaryOp()->opType() == UnaryOpOperation_TANH) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
FuseGeLu::FuseGeLu() {
|
|
auto match = [this](EXPRP expr) -> bool {
|
|
if (!expr->get() || !helpers::IsBinaryMul(expr)) {
|
|
return false;
|
|
}
|
|
EXPRP x, y, z;
|
|
EXPRP add_1, add_2;
|
|
EXPRP mul_1, mul_2, mul_3;
|
|
EXPRP pow, tanh;
|
|
// x * a
|
|
x = expr->inputs().at(0)->expr().first;
|
|
y = expr->inputs().at(1)->expr().first;
|
|
if (helpers::IsBinaryMul(x)) {
|
|
z = x;
|
|
} else if (helpers::IsBinaryMul(y)) {
|
|
z = y;
|
|
} else {
|
|
return false;
|
|
}
|
|
// x + a
|
|
x = z->inputs().at(0)->expr().first;
|
|
y = z->inputs().at(1)->expr().first;
|
|
if (helpers::IsBinaryAdd(x) && helpers::IsConstant(y)){
|
|
z = x;
|
|
} else if (helpers::IsConstant(x) && helpers::IsBinaryAdd(y)){
|
|
z = y;
|
|
} else {
|
|
return false;
|
|
}
|
|
// tan(x) + a
|
|
x = z->inputs().at(0)->expr().first;
|
|
y = z->inputs().at(1)->expr().first;
|
|
if (helpers::IsConstant(x) && IsTanh(y)) {
|
|
z = y;
|
|
} else if (IsTanh(x) && helpers::IsConstant(y)) {
|
|
z = x;
|
|
} else {
|
|
return false;
|
|
}
|
|
// z = x * a
|
|
z = z->inputs().at(0)->expr().first;
|
|
if (!helpers::IsBinaryMul(z)) {
|
|
return false;
|
|
}
|
|
// z = x + y
|
|
x = z->inputs().at(0)->expr().first;
|
|
y = z->inputs().at(1)->expr().first;
|
|
if (helpers::IsBinaryAdd(x) && helpers::IsConstant(y)){
|
|
z = x;
|
|
} else if (helpers::IsConstant(x) && helpers::IsBinaryAdd(y)) {
|
|
z = y;
|
|
} else {
|
|
return false;
|
|
}
|
|
// z = x * a
|
|
x = z->inputs().at(0)->expr().first;
|
|
y = z->inputs().at(1)->expr().first;
|
|
if (helpers::IsBinaryMul(x)) {
|
|
gelu_inputs_ = z->inputs().at(1);
|
|
z = x;
|
|
} else if (helpers::IsBinaryMul(y)) {
|
|
gelu_inputs_ = z->inputs().at(0);
|
|
z = y;
|
|
} else {
|
|
return false;
|
|
}
|
|
// z = x ^ a
|
|
x = z->inputs().at(0)->expr().first;
|
|
y = z->inputs().at(1)->expr().first;
|
|
if (IsBinaryPow(x) && helpers::IsConstant(y)) {
|
|
z = x;
|
|
} else if (helpers::IsConstant(x) && IsBinaryPow(y)) {
|
|
z = y;
|
|
} else {
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
|
|
auto fold = [this](EXPRP expr) -> bool {
|
|
auto config = Global<modelConfig>::Get();
|
|
auto version = config->targetVersion;
|
|
if (version < 1.2f) {
|
|
// For target version < 1.2 , don't support gelu
|
|
return false;
|
|
}
|
|
std::unique_ptr<OpT> gelu_op(new OpT);
|
|
gelu_op->name = expr->name();
|
|
gelu_op->type = OpType_UnaryOp;
|
|
gelu_op->main.type = OpParameter_UnaryOp;
|
|
gelu_op->main.value = new UnaryOpT;
|
|
auto geluParam = gelu_op->main.AsUnaryOp();
|
|
geluParam->opType = UnaryOpOperation_GELU;
|
|
auto gelu_expr = Variable::create(Expr::create(gelu_op.get(), { gelu_inputs_ }, 1));
|
|
gelu_expr->setName(expr->name());
|
|
Expr::replace(expr, gelu_expr->expr().first);
|
|
return true /*modified*/;
|
|
};
|
|
TemplateMerge::getInstance("Merge").insertTemplate("FuseGeLu", match, fold);
|
|
}
|
|
|
|
static FuseGeLu g_fuse_gelu;
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|