wehub-resource-sync
99 lines
3.4 KiB
C++
99 lines
3.4 KiB
C++
//
|
|
// FuseLayerNorm.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2024/01/29.
|
|
// 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 FuseLayerNormRMSGamma {
|
|
public:
|
|
FuseLayerNormRMSGamma();
|
|
|
|
private:
|
|
std::vector<int> axis_var_;
|
|
VARP x_var_;
|
|
VARP epsilon_var_;
|
|
};
|
|
FuseLayerNormRMSGamma::FuseLayerNormRMSGamma() {
|
|
auto match = [this](EXPRP expr) -> bool {
|
|
if (!expr->get() || !helpers::IsBinaryMul(expr)) {
|
|
return false;
|
|
}
|
|
auto mulexpr = expr->inputs().at(1)->expr().first;
|
|
auto conexpr = expr->inputs().at(0)->expr().first;
|
|
if (helpers::IsLayerNorm(mulexpr) && helpers::IsConstant(conexpr)) {
|
|
std::unique_ptr<OpT> op(mulexpr->get()->UnPack());
|
|
auto params = op->main.AsLayerNorm();
|
|
if (!params->gamma.empty() || !params->beta.empty()) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
if (helpers::IsLayerNorm(conexpr) && helpers::IsConstant(mulexpr)) {
|
|
std::unique_ptr<OpT> op(conexpr->get()->UnPack());
|
|
auto params = op->main.AsLayerNorm();
|
|
if (!params->gamma.empty() || !params->beta.empty()) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
return false;
|
|
};
|
|
|
|
auto fold = [this](EXPRP expr) -> bool {
|
|
std::unique_ptr<MNN::LayerNormT> layer_norm(new MNN::LayerNormT);
|
|
auto mulexpr = expr->inputs().at(1)->expr().first;
|
|
auto conexpr = expr->inputs().at(0)->expr().first;
|
|
int k = 0;
|
|
if (helpers::IsConstant(mulexpr) && helpers::IsLayerNorm(conexpr)) {
|
|
k = 1;
|
|
}
|
|
|
|
|
|
mulexpr = expr->inputs().at(1-k)->expr().first;
|
|
std::unique_ptr<OpT> op(mulexpr->get()->UnPack());
|
|
auto params = op->main.AsLayerNorm();
|
|
std::unique_ptr<MNN::LayerNormT> layernorm(new MNN::LayerNormT);
|
|
layernorm->axis = params->axis;
|
|
layernorm->epsilon = params->epsilon;
|
|
layernorm->useRMSNorm = params->useRMSNorm;
|
|
auto gammaVar = expr->inputs().at(k);
|
|
auto gammasize = gammaVar->getInfo()->size;
|
|
// if (expr->inputs().at(1 - k)->getInfo() == nullptr) {
|
|
// return false;
|
|
// }
|
|
// auto reducesize = expr->inputs().at(1 - k)->getInfo()->dim[expr->inputs().at(1 - k)->getInfo()->dim.size() - 1];
|
|
// if (reducesize != gammasize) {
|
|
// return false;
|
|
// }
|
|
layernorm->gamma.resize(gammasize);
|
|
::memcpy(layernorm->gamma.data(), gammaVar->readMap<float>(), gammasize * sizeof(float));
|
|
layernorm->beta.resize(gammasize);
|
|
std::unique_ptr<OpT> newOp(new OpT);
|
|
newOp->name = mulexpr->name();
|
|
newOp->type = OpType_LayerNorm;
|
|
newOp->main.type = OpParameter_LayerNorm;
|
|
newOp->main.value = layernorm.release();
|
|
|
|
EXPRP layer_norm_expr = Expr::create(newOp.get(), mulexpr->inputs(), 1);
|
|
layer_norm_expr->setName(expr->name());
|
|
Expr::replace(expr, layer_norm_expr);
|
|
return true;
|
|
|
|
};
|
|
TemplateMerge::getInstance("Merge").insertTemplate("FuseLayerNormRMSGamma", match, fold);
|
|
}
|
|
static FuseLayerNormRMSGamma g_fuse_layer_norm_rms_gamma;
|
|
} // namespace Express
|
|
} // namespace MNN
|