wehub-resource-sync
55 lines
1.5 KiB
C++
55 lines
1.5 KiB
C++
//
|
|
// FoldExpandDimsConst.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/12/10.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../TemplateMerge.hpp"
|
|
#include "MNN/expr/ExprCreator.hpp"
|
|
#include "MNN_generated.h"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class FoldExpandDimsConst {
|
|
public:
|
|
FoldExpandDimsConst();
|
|
};
|
|
|
|
FoldExpandDimsConst::FoldExpandDimsConst() {
|
|
auto match = [this](EXPRP expr) -> bool {
|
|
// Check the current op is `ExpandDims` op.
|
|
if (!expr->get() || expr->get()->type() != OpType_ExpandDims ||
|
|
expr->inputs().size() < 2) {
|
|
return false;
|
|
}
|
|
// Check the second input is `Constant` op.
|
|
VARP axis = expr->inputs().at(1);
|
|
if (axis->expr().first->inputType() != VARP::CONSTANT) {
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
|
|
auto fold = [this](EXPRP expr) -> bool {
|
|
VARP axis = expr->inputs().at(1);
|
|
int axis_val = axis->readMap<int>()[0];
|
|
std::unique_ptr<OpT> expand_dims_op(expr->get()->UnPack());
|
|
expand_dims_op->main.AsExpandDims()->axis = axis_val;
|
|
|
|
auto expand_dims = Expr::create(expand_dims_op.get(), // NOLINT
|
|
{expr->inputs().at(0)});
|
|
Expr::replace(expr, expand_dims);
|
|
return true;
|
|
};
|
|
|
|
TemplateMerge::getInstance("Merge").insertTemplate("FoldExpandDimsConst", match, fold);
|
|
}
|
|
|
|
static FoldExpandDimsConst g_fold_expand_dims_const;
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|