wehub-resource-sync
62 lines
1.6 KiB
C++
62 lines
1.6 KiB
C++
//
|
|
// TorchAddmm.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/05/11.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "MNN_generated.h"
|
|
#include "TorchExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class TorchAddmmTransform : public TorchExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
int alpha = 1;
|
|
int beta = 1;
|
|
|
|
auto extraParam = op->main_as_Extra();
|
|
const int attrSize = extraParam->attr()->size();
|
|
for (int i = 0; i < attrSize; ++i) {
|
|
auto attr = extraParam->attr()->GetAs<Attribute>(i);
|
|
const auto& key = attr->key()->str();
|
|
if (key == "alpha") {
|
|
alpha = attr->i();
|
|
continue;
|
|
}
|
|
if (key == "beta") {
|
|
beta = attr->i();
|
|
continue;
|
|
}
|
|
}
|
|
// X = beta * A + alpha * (B x C)
|
|
auto A = inputs[0];
|
|
auto B = inputs[1];
|
|
auto C = inputs[2];
|
|
if (beta != 1) {
|
|
A = A * _Scalar<int>(beta);
|
|
}
|
|
auto D = _MatMul(B, C);
|
|
if (1 != alpha) {
|
|
D = D * _Scalar<int>(alpha);
|
|
}
|
|
auto X = A + D;
|
|
X->setName(expr->name());
|
|
|
|
return X->expr().first;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
TorchExtraManager::get()->insert("addmm", std::shared_ptr<TorchExtraManager::Transform>(new TorchAddmmTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|