root
/
MNN
mirror of https://github.com/alibaba/MNN.git
1
0
Fork 0
Code Issues Actions 1 Packages Projects Releases Wiki Activity
MNN/source/backend/cpu/compute/ConvolutionWinograd.cpp

353 lines
16 KiB
C++
Raw Blame History

//
// ConvolutionWinograd.cpp
// MNN
//
// Created by MNN on 2018/08/20.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include "backend/cpu/compute/ConvolutionWinograd.hpp"
#include <math.h>
#include "backend/cpu/compute/CommonOptFunction.h"
#include "core/Concurrency.h"
#include "backend/cpu/compute/ConvOpt.h"
#include "core/Macro.h"
#include "core/TensorUtils.hpp"
#include "math/WingoradGenerater.hpp"
#ifdef MNN_USE_NEON
#include <arm_neon.h>
#endif
#define CONVOLUTION_WINOGRAD_MAX_UNIT 8
#define CONVOLUTION_WINOGRAD_MIN_UNIT 2
using namespace MNN::Math;
//#define MNN_WINOGRAD_PRINT_REDUCE_RATE
namespace MNN {
ConvolutionWinograd::ConvolutionWinograd(const Convolution2DCommon *convOp, const Tensor *input, const Tensor *output,
Backend *b, const float *originWeight, size_t originWeightSize,
const float *bias, size_t biasSize, int unit)
: MNN::CPUConvolution(convOp, b) {
mBias.reset(Tensor::createDevice<float>({ALIGN_UP4((int)biasSize)}));
mValid = backend()->onAcquireBuffer(mBias.get(), Backend::STATIC);
if (!mValid) {
return;
}
::memset(mBias->host<float>(), 0, mBias->size());
::memcpy(mBias->host<float>(), bias, biasSize * sizeof(float));
mTempBuffer.buffer().type = halide_type_of<float>();
mTransformMidBuffer.buffer().type = halide_type_of<float>();
MNN_ASSERT(mCommon->kernelX() == mCommon->kernelY());
int threadNumber = ((CPUBackend *)backend())->threadNumber();
auto kernelSize = mCommon->kernelY();
WinogradGenerater generator(unit, kernelSize);
int alpha = unit + kernelSize - 1;
int alpha2 = alpha * alpha;
mSourceTransform = WinogradFunction::chooseSourceTransform(alpha, alpha);
mDestTransform = WinogradFunction::chooseDestTransform(alpha, unit);
int srcCount = input->channel();
int outputCount = output->channel();
mTempBuffer.buffer().dim[0].extent = threadNumber;
mTempBuffer.buffer().dim[1].extent = CONVOLUTION_TILED_NUMBER;
mTempBuffer.buffer().dim[2].extent = UP_DIV(srcCount, 4) + UP_DIV(outputCount, 4);
mTempBuffer.buffer().dim[3].extent = 4 * alpha2;
TensorUtils::setLinearLayout(&mTempBuffer);
mTransformMidBuffer.buffer().dim[0].extent = threadNumber;
mTransformMidBuffer.buffer().dim[1].extent = 2;
mTransformMidBuffer.buffer().dim[2].extent = alpha2;
mTransformMidBuffer.buffer().dim[3].extent = 4;
TensorUtils::setLinearLayout(&mTransformMidBuffer);
mA = generator.A();
mB = generator.B();
// Transform Kernel
auto G = generator.G();
std::shared_ptr<Tensor> sourceWeight(Tensor::create<float>(
std::vector<int>{outputCount, srcCount, kernelSize, kernelSize}, (void *)originWeight, Tensor::CAFFE));
mWeight = generator.allocTransformWeight(sourceWeight.get(), 4, 4, false);
mValid = backend()->onAcquireBuffer(mWeight.get(), Backend::STATIC);
if (!mValid) {
return;
}
generator.transformWeight(mWeight.get(), sourceWeight.get());
}
ConvolutionWinograd::~ConvolutionWinograd() {
if (nullptr != mBias) {
backend()->onReleaseBuffer(mBias.get(), Backend::STATIC);
}
if (nullptr != mWeight) {
backend()->onReleaseBuffer(mWeight.get(), Backend::STATIC);
}
}
ErrorCode ConvolutionWinograd::onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
auto input = inputs[0];
auto output = outputs[0];
auto dstUnit = mA->length(1);
auto srcUnit = mA->length(0);
auto srcUnit2 = srcUnit * srcUnit;
auto dstUnit2 = dstUnit * dstUnit;
int ow = output->width();
int oh = output->height();
int iw = input->width();
int ih = input->height();
int ic_4 = UP_DIV(input->channel(), 4);
int dc_4 = UP_DIV(output->channel(), 4);
// MNN_PRINT("%d, %d\n", srcUnit, dstUnit);
int padY = mPadY;
int padX = mPadX;
auto wUnit = UP_DIV(ow, dstUnit);
auto hUnit = UP_DIV(oh, dstUnit);
auto totalCount = wUnit * hUnit;
auto postFunction = mPostFunction;
// MNN_PRINT("ow=%d, oh=%d\n", ow, oh);
int threadNumber = std::max(((CPUBackend *)backend())->threadNumber(), 1);
int tileCount = UP_DIV(totalCount, CONVOLUTION_TILED_NUMBER);
threadNumber = std::min(threadNumber, tileCount);
for (int batchIndex = 0; batchIndex < input->batch(); ++batchIndex) {
auto srcOrigin = input->host<float>() + batchIndex * input->stride(0);
auto dstOrigin = output->host<float>() + batchIndex * output->stride(0);
auto weight = mWeight->host<float>();
auto bias = mBias->host<float>();
auto tFunction = [&](int tId) {
auto _srcOrigin = mTempBuffer.host<float>() + tId * mTempBuffer.stride(0);
auto midBuffer0 = mTransformMidBuffer.host<float>() + tId * mTransformMidBuffer.stride(0);
auto midBuffer1 =
mTransformMidBuffer.host<float>() + tId * mTransformMidBuffer.stride(0) + mTransformMidBuffer.stride(1);
for (int tIndex = (int)tId; tIndex < tileCount; tIndex += threadNumber) {
int xIndex = (int)tIndex * CONVOLUTION_TILED_NUMBER;
int xReamin = totalCount - xIndex;
int xC = xReamin > CONVOLUTION_TILED_NUMBER ? CONVOLUTION_TILED_NUMBER : xReamin;
/*Source Transform Begin*/
{
int sourceZStep = iw * ih * 4;
int dstZStep = xC * 4;
int unitStep = ic_4 * xC * 4;
for (int xi = 0; xi < xC; ++xi) {
auto index = xIndex + xi;
int wIndex = index % wUnit;
int hIndex = index / wUnit;
int srcX = wIndex * dstUnit - padX;
int srcY = hIndex * dstUnit - padY;
int sy = ALIMAX(0, srcY) - srcY;
int ey = ALIMIN(srcY + srcUnit, ih) - srcY;
int sx = ALIMAX(0, srcX) - srcX;
int ex = ALIMIN(srcX + srcUnit, iw) - srcX;
int count = 4 * (ex - sx);
auto dst_x = _srcOrigin + 4 * xi;
auto srcStart = srcOrigin + (srcX + srcY * iw) * 4;
if (ex - sx == srcUnit && ey - sy == srcUnit) {
for (int z = 0; z < ic_4; ++z) {
auto srcZ = srcStart + z * sourceZStep;
// Transform
for (int i = 0; i < srcUnit; ++i) {
mSourceTransform(srcZ + 4 * i * iw, midBuffer1 + 4 * i, 4, 4 * srcUnit);
}
auto dstZ = dst_x + z * dstZStep;
for (int i = 0; i < srcUnit; ++i) {
mSourceTransform(midBuffer1 + 4 * i * srcUnit, dstZ + i * unitStep, 4,
unitStep * srcUnit);
}
}
} else {
for (int z = 0; z < ic_4; ++z) {
// Extract
auto srcZ = srcStart + z * sourceZStep;
::memset(midBuffer0, 0, mTransformMidBuffer.stride(1) * sizeof(float));
if (count > 0) {
for (int yy = sy; yy < ey; ++yy) {
auto dst_yy = midBuffer0 + yy * srcUnit * 4 + sx * 4;
auto src_yy = srcZ + 4 * iw * yy + sx * 4;
::memcpy(dst_yy, src_yy, count * sizeof(float));
}
}
// Transform
for (int i = 0; i < srcUnit; ++i) {
mSourceTransform(midBuffer0 + 4 * i * srcUnit, midBuffer1 + 4 * i, 4, 4 * srcUnit);
}
auto dstZ = dst_x + z * dstZStep;
for (int i = 0; i < srcUnit; ++i) {
mSourceTransform(midBuffer1 + 4 * i * srcUnit, dstZ + i * unitStep, 4,
unitStep * srcUnit);
}
}
}
}
}
/*Source Transform End*/
// Multi
auto _dstOrigin = _srcOrigin + xC * srcUnit2 * ic_4 * 4;
if (xC == CONVOLUTION_TILED_NUMBER) {
for (int i = 0; i < srcUnit2; ++i) {
MNNGemmFloatUnit(_dstOrigin + i * dc_4 * 4 * xC, _srcOrigin + i * ic_4 * 4 * xC,
weight + i * 16 * ic_4 * dc_4, ic_4, xC * 4, dc_4, 0);
}
} else {
for (int i = 0; i < srcUnit2; ++i) {
MNNGemmFloatCommon_4(_dstOrigin + i * dc_4 * 4 * xC, _srcOrigin + i * ic_4 * 4 * xC,
weight + i * 16 * ic_4 * dc_4, ic_4, xC * 4, dc_4, xC, 0);
}
}
/* Dest Transform And Post Treat Begin */
{
int dstZStep = ow * oh * 4;
int srcZStep = xC * 4;
int unitStep = dc_4 * xC * 4;
for (int xi = 0; xi < xC; ++xi) {
auto index = xIndex + xi;
auto srcXi = _dstOrigin + 4 * xi;
int wIndex = index % wUnit;
int hIndex = index / wUnit;
int dstX = wIndex * dstUnit;
int dstY = hIndex * dstUnit;
auto dstStart = dstOrigin + 4 * (dstX + dstY * ow);
int ey = ALIMIN(dstY + dstUnit, oh) - dstY;
int ex = ALIMIN(dstX + dstUnit, ow) - dstX;
int count = ex * 4;
if (ex == dstUnit) {
for (int z = 0; z < dc_4; ++z) {
auto dstZAddr = dstStart + z * dstZStep;
auto srcZ = srcXi + z * srcZStep;
auto biasZ = bias + 4 * z;
// Transform
for (int i = 0; i < srcUnit; ++i) {
mDestTransform(srcZ + i * unitStep, midBuffer0 + i * dstUnit * 4,
srcUnit * unitStep, 4);
}
for (int i = 0; i < ey; ++i) {
auto dstAddr = dstZAddr + i * 4 * ow;
mDestTransform(midBuffer0 + i * 4, dstAddr, 4 * dstUnit, 4);
postFunction(dstAddr, biasZ, dstUnit, 1);
}
}
} else {
for (int z = 0; z < dc_4; ++z) {
auto dstZAddr = dstStart + z * dstZStep;
auto srcZ = srcXi + z * srcZStep;
// Transform
for (int i = 0; i < srcUnit; ++i) {
mDestTransform(srcZ + i * unitStep, midBuffer0 + i * dstUnit * 4,
srcUnit * unitStep, 4);
}
for (int i = 0; i < ey; ++i) {
mDestTransform(midBuffer0 + i * 4, midBuffer1 + i * dstUnit * 4, 4 * dstUnit, 4);
}
// PostTreat
postFunction(midBuffer1, bias + 4 * z, dstUnit2, 1);
for (int yy = 0; yy < ey; ++yy) {
auto dstYAddr = dstZAddr + yy * 4 * ow;
auto srcYAddr = midBuffer1 + yy * 4 * dstUnit;
::memcpy(dstYAddr, srcYAddr, count * sizeof(float));
}
}
}
}
}
/*Dest Transform And Post Treat End*/
}
};
MNN_CONCURRENCY_BEGIN(tId, threadNumber) {
tFunction((int)tId);
}
MNN_CONCURRENCY_END();
}
return NO_ERROR;
}
int ConvolutionWinograd::bestWinogradUnit(const Convolution2DCommon *common, const Tensor *inputTensor,
const Tensor *outputTensor, int threadNumber) {
int ow = outputTensor->width();
int oh = outputTensor->height();
int oc = outputTensor->channel();
int unit2 = UP_DIV(ow * oh, CONVOLUTION_TILED_NUMBER * threadNumber);
int maxUnit = (int)::sqrtf((float)unit2);
maxUnit = std::min(maxUnit, CONVOLUTION_WINOGRAD_MAX_UNIT);
maxUnit = std::max(maxUnit, CONVOLUTION_WINOGRAD_MIN_UNIT);
int ic = inputTensor->channel();
auto kernelSize = common->kernelY();
int unit = CONVOLUTION_WINOGRAD_MIN_UNIT;
float maxRate = 0.0f;
float originCost = (float)ow * oh * (float)ic * oc * kernelSize * kernelSize;
static std::set<int> supportSu{4, 8};
for (int u = CONVOLUTION_WINOGRAD_MIN_UNIT; u <= maxUnit; ++u) {
float su = (float)(u + kernelSize - 1);
if (supportSu.find(su) == supportSu.end()) {
continue;
}
if (nullptr == WinogradFunction::chooseDestTransform((int)su, u)) {
continue;
}
/*Let F(6,3) be choosed when it can speed up from F(2,3) than 0.6*/
float penalty = (su * su) / (float)(kernelSize * kernelSize) * 0.12f;
float winogradCost =
(2 * su * su * ic + su * su * ic * oc + (su + u) * u * oc) * (UP_DIV(ow, u) * UP_DIV(oh, u));
float reduceRate = originCost / winogradCost - penalty;
// MNN_PRINT("ow=%d, oh=%d, %f, %f, winograd unit:%d\n", ow, oh, winogradCost, reduceRate, u);
if (reduceRate > maxRate) {
maxRate = reduceRate;
unit = u;
}
}
if (maxRate < 1.0f) {
return 0;
}
return unit;
}
bool ConvolutionWinograd::canUseWinograd(const Convolution2DCommon *common) {
if (common->kernelY() != common->kernelX() || common->kernelY() <= 1) {
return false;
}
if (common->dilateX() != 1 || common->dilateY() != 1) {
return false;
}
if (common->strideX() != 1 || common->strideY() != 1) {
return false;
}
return true;
}
ErrorCode ConvolutionWinograd::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
CPUConvolution::onResize(inputs, outputs);
// FUNC_PRINT(mA->length(1));
bool success = backend()->onAcquireBuffer(&mTempBuffer, Backend::DYNAMIC);
success = success && (backend()->onAcquireBuffer(&mTransformMidBuffer, Backend::DYNAMIC));
backend()->onReleaseBuffer(&mTempBuffer, Backend::DYNAMIC);
backend()->onReleaseBuffer(&mTransformMidBuffer, Backend::DYNAMIC);
if (!success) {
return OUT_OF_MEMORY;
}
return NO_ERROR;
}
} // namespace MNN
Reference in New Issue View Git Blame Copy Permalink