mirror of https://github.com/alibaba/MNN.git
92 lines
3.6 KiB
C++
92 lines
3.6 KiB
C++
//
|
|
// CPUInnerProduct.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2018/08/02.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "backend/cpu/CPUInnerProduct.hpp"
|
|
#include "core/AutoStorage.h"
|
|
#include "backend/cpu/CPUConvolution.hpp"
|
|
#include "backend/cpu/compute/CommonOptFunction.h"
|
|
#include "backend/cpu/compute/ConvOpt.h"
|
|
#include "core/Macro.h"
|
|
|
|
namespace MNN {
|
|
|
|
class CPUInnerProductExecutor : public Execution {
|
|
public:
|
|
CPUInnerProductExecutor(Backend *bn, const MNN::Op *op) : Execution(bn) {
|
|
auto parameter = op->main_as_InnerProduct();
|
|
int outputCount = parameter->outputCount();
|
|
int srcCount = parameter->weight()->size() / outputCount;
|
|
mWeight.reset(CPUConvolution::reorderWeightSize(srcCount, outputCount, 1, 4, 4));
|
|
if (mWeight.get() == nullptr) {
|
|
mValid = false;
|
|
return;
|
|
}
|
|
mWeight.clear();
|
|
AutoStorage<float> cache(mWeight.size());
|
|
CPUConvolution::reorderWeight(mWeight.get(), parameter->weight()->data(), srcCount, outputCount, 1, cache.get());
|
|
mBias.reset(ALIGN_UP4(outputCount));
|
|
mBias.clear();
|
|
::memcpy(mBias.get(), parameter->bias()->data(), parameter->bias()->size() * sizeof(float));
|
|
mInputPad.reset(new Tensor(2));
|
|
mOutputPad.reset(new Tensor(2));
|
|
}
|
|
virtual ~CPUInnerProductExecutor() = default;
|
|
virtual ErrorCode onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) override {
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
mOutputPad->buffer().dim[1].extent = ALIGN_UP4(output->buffer().dim[1].extent);
|
|
mOutputPad->buffer().dim[0].extent = output->buffer().dim[0].extent;
|
|
mInputPad->buffer().dim[1].extent = ALIGN_UP4(input->buffer().dim[1].extent);
|
|
mInputPad->buffer().dim[0].extent = input->buffer().dim[0].extent;
|
|
|
|
backend()->onAcquireBuffer(mOutputPad.get(), Backend::DYNAMIC);
|
|
backend()->onAcquireBuffer(mInputPad.get(), Backend::DYNAMIC);
|
|
|
|
backend()->onReleaseBuffer(mOutputPad.get(), Backend::DYNAMIC);
|
|
backend()->onReleaseBuffer(mInputPad.get(), Backend::DYNAMIC);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
virtual ErrorCode onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) override {
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
|
|
auto originSource = input->host<float>();
|
|
int srcDepthQuad = mInputPad->buffer().dim[1].extent / 4;
|
|
int dstDepthQuad = mOutputPad->buffer().dim[1].extent / 4;
|
|
auto width = mInputPad->buffer().dim[0].extent;
|
|
|
|
auto source = mInputPad->host<float>();
|
|
MNNPackC4(source, originSource, width, input->buffer().dim[1].extent);
|
|
auto dest = mOutputPad->host<float>();
|
|
MNNGemmFloatCommon_4(dest, source, mWeight.get(), srcDepthQuad, 4 * width, dstDepthQuad, width, 0);
|
|
|
|
MNNAddBias(dest, mBias.get(), width, dstDepthQuad);
|
|
auto originDest = output->host<float>();
|
|
MNNUnpackC4(originDest, dest, width, output->buffer().dim[1].extent);
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
private:
|
|
AutoStorage<float> mWeight;
|
|
AutoStorage<float> mBias;
|
|
|
|
std::unique_ptr<Tensor> mInputPad;
|
|
std::unique_ptr<Tensor> mOutputPad;
|
|
};
|
|
|
|
Execution *CPUInnerProductCreator::onCreate(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs,
|
|
const MNN::Op *op, Backend *backend) const {
|
|
return new CPUInnerProductExecutor(backend, op);
|
|
}
|
|
|
|
REGISTER_CPU_OP_CREATOR(CPUInnerProductCreator, OpType_InnerProduct);
|
|
|
|
} // namespace MNN
|