234 lines
12 KiB
C++
234 lines
12 KiB
C++
//
|
|
// VulkanConvolutionImpl.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "VulkanConvolutionImpl.hpp"
|
|
#include "core/Macro.h"
|
|
#include "VulkanConvolution.hpp"
|
|
#include "VulkanConvolutionWinograd.hpp"
|
|
#include "VulkanMatMul.hpp"
|
|
#include "VulkanConvolution1x1.hpp"
|
|
//#define MNN_OPEN_TIME_TRACE
|
|
#include <MNN/AutoTime.hpp>
|
|
namespace MNN {
|
|
//#define VULKAN_IM2COL_GEMM_UNIT 512
|
|
static void writeParameters(VulkanMatMul::Reorder::nchwBuffer& parameters, int co, int ci, int kh, int kw) {
|
|
parameters.size[0] = co;
|
|
parameters.size[1] = ci;
|
|
parameters.size[2] = kh;
|
|
parameters.size[3] = kw;
|
|
parameters.stride[0] = ci * kh * kw;
|
|
parameters.stride[1] = kh * kw;
|
|
parameters.stride[2] = kw;
|
|
parameters.stride[3] = 1;
|
|
}
|
|
class VulkanConvolutionIm2Col : public VulkanBasicExecution {
|
|
public:
|
|
|
|
VulkanConvolutionIm2Col(VulkanBackend* backend, const Convolution2DCommon* convOption, const float* weightPtr,
|
|
const float* biasPtr, int ci, int co) : VulkanBasicExecution(backend), mConvCommonOption(convOption) {
|
|
auto kw = convOption->kernelX();
|
|
auto kh = convOption->kernelY();
|
|
if (nullptr != weightPtr) {
|
|
// Static weight
|
|
VulkanMatMul::Reorder reorder(backend, true);
|
|
VulkanMatMul::Reorder::nchwBuffer parameters;
|
|
writeParameters(parameters, co, ci, kh, kw);
|
|
mKernel = VulkanMatrixMultier4x4::createKernel(backend, nullptr, ALIGN_UP4(ci) * kh * kw, co, 1);
|
|
auto weightSize = ci * co * kh * kw;
|
|
std::shared_ptr<VulkanBuffer> tempBuffer(new VulkanBuffer(backend->getMemoryPool(), false, weightSize*sizeof(float), nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT));
|
|
auto tempWeightBuffer = tempBuffer->map();
|
|
::memcpy(tempWeightBuffer, weightPtr, weightSize * sizeof(float));
|
|
tempBuffer->unmap();
|
|
std::shared_ptr<VulkanBuffer> tempBuffer2(new VulkanBuffer(backend->getMemoryPool(), false, reorder.computeMiddleBufferSize(co, kh, kw, ci) *sizeof(float), nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT));
|
|
// TODO: Move to init buffer
|
|
std::shared_ptr<VulkanCommandPool::Buffer> cmdBuffer(backend->getPool().allocBuffer());
|
|
cmdBuffer->begin(0);
|
|
reorder.encode(tempBuffer->buffer(), tempBuffer->size(), tempBuffer2->buffer()
|
|
, tempBuffer2->size(), mKernel.get(), cmdBuffer.get(), parameters);
|
|
mKernel->barrierRead(cmdBuffer->get());
|
|
cmdBuffer->end();
|
|
backend->getPool().submitAndWait(cmdBuffer->get());
|
|
}
|
|
mMultiCreator = [ci, kh, kw, co, backend, this]() {
|
|
auto multi = std::make_shared<VulkanMatrixMultier4x4>(backend, nullptr, ALIGN_UP4(ci) * kh * kw, co, 1, mKernel);
|
|
return multi;
|
|
};
|
|
std::vector<VkDescriptorType> im2Coltypes{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER};
|
|
if (kw == 1 && kh == 1 && convOption->padX() == 0 && convOption->padY() == 0) {
|
|
mIm2Col =
|
|
backend->getPipeline("glsl_im2col1x1_comp", /* glsl_im2col1x1_comp, glsl_im2col1x1_comp_len,*/ im2Coltypes);
|
|
} else {
|
|
mIm2Col = backend->getPipeline("glsl_im2col_comp", /*glsl_im2col_comp, glsl_im2col_comp_len,*/ im2Coltypes);
|
|
}
|
|
std::vector<VkDescriptorType> Col2imTypes{
|
|
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
|
|
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER};
|
|
auto macro = VulkanConvolutionCommon::getPostTreatMacro(convOption);
|
|
mCol2Im = backend->getPipeline("glsl_col2Im_" + macro + "comp", Col2imTypes);
|
|
|
|
mSampler = backend->getCommonSampler();
|
|
if (nullptr != biasPtr) {
|
|
// Static bias
|
|
mBias = std::make_shared<VulkanImage>(backend->getMemoryPool(), false, UP_DIV(co, 4), 1);
|
|
auto tempBias = std::make_shared<VulkanBuffer>(backend->getMemoryPool(), false, sizeof(float) * ALIGN_UP4(co));
|
|
auto bias = tempBias->map();
|
|
::memset(bias, 0, sizeof(float) * ALIGN_UP4(co));
|
|
::memcpy(bias, biasPtr, sizeof(float) * co);
|
|
tempBias->unmap();
|
|
backend->copyBufferToImage(tempBias.get(), mBias.get(), VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
|
|
}
|
|
}
|
|
~VulkanConvolutionIm2Col() {
|
|
// Do nothing
|
|
}
|
|
virtual ErrorCode onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) override {
|
|
auto src = inputs[0];
|
|
auto dst = outputs[0];
|
|
const int icDiv4 = UP_DIV(src->channel(), 4);
|
|
const int ocDiv4 = UP_DIV(dst->channel(), 4);
|
|
auto vkBn = (VulkanBackend*)backend();
|
|
int limit = vkBn->proty().limits.maxImageDimension2D * 4;
|
|
#ifdef VULKAN_IM2COL_GEMM_UNIT
|
|
limit = VULKAN_IM2COL_GEMM_UNIT;
|
|
#endif
|
|
if (limit < dst->width()) {
|
|
MNN_ERROR("Don't support width too large feature: %d x %d, limit = %d\n", dst->width(), dst->height(), limit);
|
|
return NOT_SUPPORT;
|
|
}
|
|
int batchLoopNumber = 1;
|
|
int heightLoopNumber = 1;
|
|
int unitHeight = dst->height();
|
|
int unitBatch = dst->batch();
|
|
auto area = dst->width() * dst->height();
|
|
if (limit < area) {
|
|
batchLoopNumber = dst->batch();
|
|
unitBatch = 1;
|
|
unitHeight = limit / dst->width();
|
|
heightLoopNumber = UP_DIV(dst->height(), unitHeight);
|
|
} else if (limit < area * dst->batch()) {
|
|
unitBatch = limit / area;
|
|
batchLoopNumber = UP_DIV(dst->batch(), unitBatch);
|
|
}
|
|
int loopNumber = batchLoopNumber * heightLoopNumber;
|
|
mConvParams.resize(loopNumber);
|
|
mMultilers.resize(loopNumber);
|
|
mIm2ColSet.resize(loopNumber);
|
|
mCol2ImSet.resize(loopNumber);
|
|
reinterpret_cast<VulkanTensor*>(src->deviceId())->image()->barrierRead(cmdBuffer->get());
|
|
reinterpret_cast<VulkanTensor*>(dst->deviceId())->image()->barrierWrite(cmdBuffer->get());
|
|
|
|
for (int i=0; i<batchLoopNumber; ++i) {
|
|
int batchOffset = i * unitBatch;
|
|
int currentBatch = dst->batch() - batchOffset;
|
|
if (currentBatch > unitBatch) {
|
|
currentBatch = unitBatch;
|
|
}
|
|
for (int j=0; j<heightLoopNumber; ++j) {
|
|
int heightOffset = j * unitHeight;
|
|
int currentHeight = dst->height() - heightOffset;
|
|
if (currentHeight > unitHeight) {
|
|
currentHeight = unitHeight;
|
|
}
|
|
auto index = i * heightLoopNumber + j;
|
|
auto totalNumberInput = currentBatch * icDiv4 * dst->width() * currentHeight;
|
|
auto totalNumberOutput = currentBatch * ocDiv4 * dst->width() * currentHeight;
|
|
mConvParams[index] = std::make_shared<VulkanBuffer>(vkBn->getMemoryPool(), false,
|
|
sizeof(VulkanConvolutionCommon::ConvolutionParameter), nullptr,
|
|
VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT);
|
|
{
|
|
auto convCons = reinterpret_cast<VulkanConvolutionCommon::ConvolutionParameter*>(mConvParams[index]->map());
|
|
VulkanConvolutionCommon::writeParameter(convCons, mConvCommonOption, src, dst);
|
|
convCons->offset[0] = batchOffset;
|
|
convCons->offset[1] = heightOffset;
|
|
convCons->outputSize[3] = currentBatch;
|
|
convCons->outputSize[1] = currentHeight;
|
|
mConvParams[index]->unmap();
|
|
}
|
|
mIm2ColSet[index].reset(mIm2Col->createSet());
|
|
mCol2ImSet[index].reset(mCol2Im->createSet());
|
|
mMultilers[index] = mMultiCreator();
|
|
mMultilers[index]->prepare(static_cast<VulkanBackend*>(backend())->getInitCommandBuffer(), dst->width() * currentHeight * currentBatch);
|
|
auto mMultiler = mMultilers[index].get();
|
|
if (true) {
|
|
auto colImage = mMultiler->source();
|
|
// Barrier
|
|
mIm2ColSet[index]->writeImage(colImage->view(), mSampler->get(), VK_IMAGE_LAYOUT_GENERAL, 0);
|
|
mIm2ColSet[index]->writeImage((reinterpret_cast<VulkanTensor*>(src->deviceId()))->image()->view(), mSampler->get(),
|
|
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 1);
|
|
mIm2ColSet[index]->writeBuffer(mConvParams[index]->buffer(), 2, mConvParams[index]->size());
|
|
mIm2Col->bind(cmdBuffer->get(), mIm2ColSet[index]->get());
|
|
|
|
colImage->barrierWrite(cmdBuffer->get());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(totalNumberInput, VulkanConvolutionCommon::gImage2ColLocal),
|
|
1, 1);
|
|
}
|
|
mMultilers[index]->compute(cmdBuffer);
|
|
if (true) {
|
|
auto dstImage = mMultiler->dest();
|
|
mCol2ImSet[index]->writeImage(dstImage->view(), mSampler->get(), VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 0);
|
|
mCol2ImSet[index]->writeImage((reinterpret_cast<VulkanTensor*>(dst->deviceId()))->image()->view(), mSampler->get(),
|
|
VK_IMAGE_LAYOUT_GENERAL, 1);
|
|
|
|
mCol2ImSet[index]->writeImage(mBias->view(), mSampler->get(), VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 2);
|
|
mCol2ImSet[index]->writeBuffer(mConvParams[index]->buffer(), 3, mConvParams[index]->size());
|
|
mCol2Im->bind(cmdBuffer->get(), mCol2ImSet[index]->get());
|
|
|
|
dstImage->barrierRead(cmdBuffer->get());
|
|
mBias->barrierRead(cmdBuffer->get());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(totalNumberOutput, VulkanConvolutionCommon::gImage2ColLocal),
|
|
1, 1);
|
|
}
|
|
}
|
|
}
|
|
return NO_ERROR;
|
|
}
|
|
private:
|
|
const VulkanPipeline* mIm2Col;
|
|
const VulkanPipeline* mCol2Im;
|
|
const VulkanSampler* mSampler;
|
|
|
|
std::shared_ptr<VulkanImage> mBias;
|
|
std::shared_ptr<VulkanImage> mKernel;
|
|
const Convolution2DCommon* mConvCommonOption;
|
|
std::vector<std::shared_ptr<VulkanLayout::DescriptorSet>> mCol2ImSet;
|
|
std::vector<std::shared_ptr<VulkanLayout::DescriptorSet>> mIm2ColSet;
|
|
std::vector<std::shared_ptr<VulkanBuffer>> mConvParams;
|
|
std::vector<std::shared_ptr<VulkanMatrixMultier4x4>> mMultilers;
|
|
std::function<std::shared_ptr<VulkanMatrixMultier4x4>()> mMultiCreator;
|
|
};
|
|
|
|
VulkanBasicExecution* VulkanConvolutionImpl::create(VulkanBackend* backend, const Convolution2DCommon* convOption,
|
|
const std::vector<Tensor*>& inputs, const Tensor* output,
|
|
const float* weightPtr, const float* biasPtr, int ci, int co) {
|
|
AUTOTIME;
|
|
if (inputs.size() > 1) {
|
|
return new VulkanConvolutionIm2Col(backend, convOption, weightPtr, biasPtr, ci, co);
|
|
}
|
|
auto imageLimit = backend->proty().limits.maxImageDimension2D;
|
|
if (VulkanConvolutionWinograd::support(convOption)) {
|
|
if (output->width() >= 4 && output->height() >= 4 && output->batch() == 1) {
|
|
return new VulkanConvolutionWinograd(backend, convOption, weightPtr, biasPtr, ci, co);
|
|
}
|
|
}
|
|
if (ALIGN_UP4(ci) * convOption->kernelX() * convOption->kernelY() > imageLimit) {
|
|
return nullptr;
|
|
}
|
|
|
|
if (convOption->kernelX() == 1 && convOption->kernelY() == 1 &&
|
|
convOption->strideX() == 1 && convOption->strideY() == 1 &&
|
|
inputs[0]->width() == output->width() && inputs[0]->height() == output->height()) {
|
|
return new VulkanConvolution1x1(backend, convOption, weightPtr, biasPtr, ci, co);
|
|
}
|
|
|
|
return new VulkanConvolutionIm2Col(backend, convOption, weightPtr, biasPtr, ci, co);
|
|
}
|
|
|
|
} // namespace MNN
|