wehub-resource-sync
252 lines
11 KiB
C++
252 lines
11 KiB
C++
#include "VulkanRaster.hpp"
|
|
#include "VulkanMatMul.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include <algorithm>
|
|
#include "core/OpCommonUtils.hpp"
|
|
#include "core/Macro.h"
|
|
namespace MNN {
|
|
static void writeNCHW(int* dims, Tensor* origin) {
|
|
int w = std::max(origin->width(), 1);
|
|
int h = std::max(origin->height(), 1);
|
|
int b = origin->batch();
|
|
dims[0] = w;
|
|
dims[1] = h;
|
|
dims[2] = origin->channel();
|
|
dims[3] = b;
|
|
}
|
|
struct SamplerInfo {
|
|
ivec4 stride;//stride[3] + offset
|
|
ivec4 size;//size[3] + totalSize
|
|
ivec4 extent;//dstStride[3]+dstOffset
|
|
ivec4 imageSize;// srcwh and dstwh
|
|
ivec2 depth;//c4 for src and dst
|
|
};
|
|
|
|
static void writeSamplerInfo(SamplerInfo& info, const Tensor::InsideDescribe::Region& sampler) {
|
|
int sizeTotal = 1;
|
|
for (int i=0; i<3; ++i) {
|
|
info.size[i] = sampler.size[i];
|
|
info.stride[i] = sampler.src.stride[i];
|
|
info.extent[i] = sampler.dst.stride[i];
|
|
sizeTotal *= info.size[i];
|
|
}
|
|
info.size[3] = sizeTotal;
|
|
info.stride[3] = sampler.src.offset;
|
|
info.extent[3] = sampler.dst.offset;
|
|
}
|
|
|
|
void VulkanRaster::onEncodeFast(const Tensor* input, const Tensor* output, const VulkanCommandPool::Buffer *cmdBuffer, bool zero) {
|
|
auto des = TensorUtils::getDescribe(input);
|
|
mBlitImages.resize(des->regions.size());
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
auto dstTensor = reinterpret_cast<VulkanTensor*>(output->deviceId());
|
|
bool isInt = output->getType().code == halide_type_int || output->getType().code == halide_type_uint;
|
|
if (zero) {
|
|
std::string fillName = isInt ? "glsl_fill_image_INT_comp" : "glsl_fill_image_comp";
|
|
auto fillPipeline =
|
|
vkBn->getPipeline(fillName, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER});
|
|
struct FillImage {
|
|
vec4 value;
|
|
ivec4 imageSize;
|
|
};
|
|
FillImage uniformInfo;
|
|
::memset(&uniformInfo, 0, sizeof(FillImage));
|
|
auto image = dstTensor->image();
|
|
uniformInfo.imageSize[0] = image->width();
|
|
uniformInfo.imageSize[1] = image->height();
|
|
uniformInfo.imageSize[2] = 0;
|
|
uniformInfo.imageSize[3] = image->width() * image->height();
|
|
std::shared_ptr<VulkanBuffer> uniform(new VulkanBuffer(vkBn->getMemoryPool(), false, sizeof(FillImage), &uniformInfo, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
|
mExtraUniform.emplace_back(uniform);
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> des(fillPipeline->createSet());
|
|
des->writeImage(image->view(), vkBn->getCommonSampler()->get(), VK_IMAGE_LAYOUT_GENERAL, 0);
|
|
des->writeBuffer(uniform->buffer(), 1, uniform->size());
|
|
auto totalSize = UP_DIV(uniformInfo.imageSize[3], 256);
|
|
mExtraDescribes.emplace_back(des);
|
|
fillPipeline->bind(cmdBuffer->get(), des->get());
|
|
image->barrierWrite(cmdBuffer->get());
|
|
vkCmdDispatch(cmdBuffer->get(), totalSize, 1, 1);
|
|
}
|
|
|
|
std::string blitName = isInt ? "glsl_blit_image_INT_comp" : "glsl_blit_image_comp";
|
|
auto blitPipeline =
|
|
vkBn->getPipeline(blitName, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER});
|
|
|
|
for (int i=0; i< des->regions.size(); ++i) {
|
|
auto& slice = des->regions[i];
|
|
Tensor::InsideDescribe::Region newRegion;
|
|
OpCommonUtils::turnToPackRegion(slice, newRegion, output, 4);
|
|
// TODO: Find better way
|
|
newRegion.dst.offset /= 4;
|
|
newRegion.src.offset /= 4;
|
|
auto& dst = mBlitImages[i];
|
|
SamplerInfo info;
|
|
writeSamplerInfo(info, newRegion);
|
|
auto nhwcSrc = VulkanTensor::tensorShapeFormat(slice.origin);
|
|
auto nhwcDst = VulkanTensor::tensorShapeFormat(output);
|
|
info.imageSize[0] = nhwcSrc[2];
|
|
info.imageSize[1] = nhwcSrc[1];
|
|
info.imageSize[2] = nhwcDst[2];
|
|
info.imageSize[3] = nhwcDst[1];
|
|
info.depth[0] = UP_DIV(nhwcSrc[3], 4);
|
|
info.depth[1] = UP_DIV(nhwcDst[3], 4);
|
|
auto total = info.size[0] * info.size[1] * info.size[2];
|
|
auto group = UP_DIV(total, 256);
|
|
dst.describe.reset(blitPipeline->createSet());
|
|
dst.uniform.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, sizeof(SamplerInfo), &info, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
|
auto srcTensor = reinterpret_cast<VulkanTensor*>(slice.origin->deviceId());
|
|
dst.describe->writeImage(srcTensor->image()->view(), vkBn->getCommonSampler()->get(), VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 1);
|
|
dst.describe->writeImage(dstTensor->image()->view(), vkBn->getCommonSampler()->get(), VK_IMAGE_LAYOUT_GENERAL, 0);
|
|
dst.describe->writeBuffer(dst.uniform->buffer(), 2, dst.uniform->size());
|
|
srcTensor->image()->barrierRead(cmdBuffer->get());
|
|
dstTensor->image()->barrierWrite(cmdBuffer->get());
|
|
blitPipeline->bind(cmdBuffer->get(), dst.describe->get());
|
|
vkCmdDispatch(cmdBuffer->get(), group, 1, 1);
|
|
}
|
|
}
|
|
|
|
|
|
ErrorCode VulkanRaster::onEncode(const std::vector<Tensor *> &___inputs, const std::vector<Tensor *> &outputs,
|
|
const VulkanCommandPool::Buffer *cmdBuffer) {
|
|
MNN_ASSERT(outputs.size() == 1);
|
|
auto output = outputs[0];
|
|
OpCommonUtils::rasterInputReset(___inputs, outputs[0]);
|
|
auto des = TensorUtils::getDescribe(output);
|
|
auto outputDes = TensorUtils::getDescribe(output);
|
|
bool needZero = !TensorUtils::regionIsFull(output);
|
|
|
|
/** Alloc Begin*/
|
|
mInputBuffers.clear();
|
|
mOutputBuffer.buffer = nullptr;
|
|
mBlits.resize(des->regions.size());
|
|
mBlitImages.clear();
|
|
mExtraUniform.clear();
|
|
mExtraDescribes.clear();
|
|
if (outputDes->dimensionFormat == MNN_DATA_FORMAT_NC4HW4) {
|
|
// TODO: Optimize it
|
|
bool fast = true;
|
|
for (int i=0; i< des->regions.size(); ++i) {
|
|
auto& slice = des->regions[i];
|
|
if (TensorUtils::getDescribe(slice.origin)->dimensionFormat != MNN_DATA_FORMAT_NC4HW4) {
|
|
fast = false;
|
|
break;
|
|
}
|
|
if (!OpCommonUtils::canBlitFast(slice, output)) {
|
|
fast = false;
|
|
break;
|
|
}
|
|
}
|
|
if (fast) {
|
|
onEncodeFast(output, output, cmdBuffer, needZero);
|
|
return NO_ERROR;
|
|
}
|
|
}
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
std::vector<VkDescriptorType> nchwConvertTypes{VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER};
|
|
for (auto& slice : des->regions) {
|
|
auto origin = slice.origin;
|
|
if (mInputBuffers.find(origin)!=mInputBuffers.end()) {
|
|
continue;
|
|
}
|
|
MNN_ASSERT(origin->deviceId() != 0);
|
|
int bufferSize = origin->getType().bytes();
|
|
for (int i=0; i<origin->dimensions(); ++i) {
|
|
bufferSize *= origin->length(i);
|
|
}
|
|
ConvertInfo info;
|
|
info.buffer.reset(new VulkanBuffer(vkBn->getDynamicMemoryPool(),
|
|
false, bufferSize, nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT));
|
|
info.convert.reset(new VulkanImageConverter(vkBn));
|
|
mInputBuffers.insert(std::make_pair(origin, std::move(info)));
|
|
}
|
|
{
|
|
int bufferSize = output->getType().bytes();
|
|
for (int i=0; i<output->dimensions(); ++i) {
|
|
bufferSize *= output->length(i);
|
|
}
|
|
mOutputBuffer.convert.reset(new VulkanImageConverter(vkBn));
|
|
mOutputBuffer.buffer.reset(new VulkanBuffer(vkBn->getDynamicMemoryPool(), false, bufferSize, nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT));
|
|
}
|
|
for (auto& iter : mInputBuffers) {
|
|
iter.second.buffer->release();
|
|
}
|
|
if (nullptr != mOutputBuffer.buffer) {
|
|
mOutputBuffer.buffer->release();
|
|
}
|
|
bool isInt = output->getType().code == halide_type_int || output->getType().code == halide_type_uint;
|
|
std::string blitName = isInt ? "glsl_blit_INT_comp" : "glsl_blit_comp";
|
|
auto blitPipeline =
|
|
vkBn->getPipeline(blitName, {VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER});
|
|
for (int i=0; i<mBlits.size(); ++i) {
|
|
auto& origin = des->regions[i];
|
|
auto& dst = mBlits[i];
|
|
SamplerInfo info;
|
|
writeSamplerInfo(info, origin);
|
|
auto total = info.size[0] * info.size[1] * info.size[2];
|
|
dst.workGroup[2] = 1;
|
|
dst.workGroup[1] = 1;
|
|
dst.workGroup[0] = UP_DIV(total, 256);
|
|
dst.pipeline = blitPipeline;
|
|
dst.describe.reset(blitPipeline->createSet());
|
|
dst.uniform.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, sizeof(info), &info, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
|
auto srcIter = mInputBuffers.find(origin.origin);
|
|
dst.srcBuffer = srcIter->second.buffer->buffer();
|
|
dst.srcBufferSize = srcIter->second.buffer->size();
|
|
dst.dstBuffer = mOutputBuffer.buffer->buffer();
|
|
dst.dstBufferSize = mOutputBuffer.buffer->size();
|
|
}
|
|
if (needZero) {
|
|
mZero.dstBuffer = mOutputBuffer.buffer->buffer();
|
|
mZero.dstBufferSize = mOutputBuffer.buffer->size();
|
|
}
|
|
/** Alloc End*/
|
|
|
|
/** Encode Begin*/
|
|
// Convert NC4HW4 image to buffer
|
|
for (auto& iter : mInputBuffers) {
|
|
auto& info = iter.second;
|
|
info.convert->encodeTensorToBuffer(iter.first, info.buffer->buffer(), info.buffer->size(), 0, VulkanImageConverter::getTensorLinearFormat(iter.first), cmdBuffer);
|
|
cmdBuffer->barrierSource(info.buffer->buffer(), 0, info.buffer->size());
|
|
}
|
|
//Blit
|
|
if (needZero) {
|
|
vkCmdFillBuffer(cmdBuffer->get(), mZero.dstBuffer, 0, mZero.dstBufferSize, 0);
|
|
cmdBuffer->barrierSource(mZero.dstBuffer, 0, mZero.dstBufferSize, VulkanCommandPool::Buffer::WRITE_WRITE);
|
|
}
|
|
for (auto& info : mBlits) {
|
|
info.describe->writeBuffer(info.dstBuffer, 0, info.dstBufferSize);
|
|
info.describe->writeBuffer(info.srcBuffer, 1, info.srcBufferSize);
|
|
info.describe->writeBuffer(info.uniform->buffer(), 2, info.uniform->size());
|
|
info.pipeline->bind(cmdBuffer->get(), info.describe->get());
|
|
vkCmdDispatch(cmdBuffer->get(), info.workGroup[0], info.workGroup[1], info.workGroup[2]);
|
|
}
|
|
|
|
// Convert buffer to NC4HW4 image
|
|
{
|
|
auto& info = mOutputBuffer;
|
|
info.convert->encodeBufferToTensor(info.buffer->buffer(), output, info.buffer->size(), 0, VulkanImageConverter::getTensorLinearFormat(output), cmdBuffer);
|
|
}
|
|
/** Encode End*/
|
|
return NO_ERROR;
|
|
}
|
|
class VulkanRasterCreator : public VulkanBackend::Creator {
|
|
public:
|
|
virtual VulkanBasicExecution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const MNN::Op* op, Backend* bn) const override {
|
|
if (outputs[0]->getType().bytes() < 4) {
|
|
return nullptr;
|
|
}
|
|
return new VulkanRaster(bn);
|
|
}
|
|
};
|
|
|
|
|
|
static bool gResistor = []() {
|
|
VulkanBackend::addCreator(OpType_Raster, new VulkanRasterCreator);
|
|
return true;
|
|
}();
|
|
|
|
};
|