MNN/backupcode/cpubackend/CPUInnerProduct.cpp

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