wehub-resource-sync
83 lines
3.4 KiB
C++
83 lines
3.4 KiB
C++
//
|
|
// VulkanResize.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "VulkanResize.hpp"
|
|
#include "core/Macro.h"
|
|
|
|
namespace MNN {
|
|
struct GpuParam {
|
|
ivec4 inImgSize;
|
|
ivec4 outImgSize;
|
|
vec4 cord;
|
|
};
|
|
|
|
VulkanResize::VulkanResize(Backend* bn, float xScale, float yScale, int resizeType)
|
|
: VulkanBasicExecution(bn), mXScale(xScale), mYScale(yScale) {
|
|
std::vector<VkDescriptorType> VulkanResizeTypes{
|
|
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
|
};
|
|
auto extra = static_cast<VulkanBackend*>(bn);
|
|
if (1 == resizeType) {
|
|
mVulkanResizePipeline = extra->getPipeline(
|
|
"glsl_resizeNearest_comp", VulkanResizeTypes);
|
|
} else if (2 == resizeType) {
|
|
mVulkanResizePipeline = extra->getPipeline(
|
|
"glsl_resizeBilinear_comp", VulkanResizeTypes);
|
|
} else {
|
|
mVulkanResizePipeline = extra->getPipeline(
|
|
"glsl_resizeNearest_NEAREST_ROUND_comp", VulkanResizeTypes);
|
|
}
|
|
mParamBuffer.reset(
|
|
new VulkanBuffer(extra->getMemoryPool(), false, sizeof(GpuParam), nullptr, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
|
}
|
|
VulkanResize::~VulkanResize() {
|
|
}
|
|
|
|
ErrorCode VulkanResize::encodeImpl(Tensor* input, Tensor* output, const float* cords,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) {
|
|
const int channelDiv4 = UP_DIV(input->channel(), 4);
|
|
auto extra = static_cast<VulkanBackend*>(backend());
|
|
|
|
auto VulkanResizeParam = reinterpret_cast<GpuParam*>(mParamBuffer->map());
|
|
::memset(VulkanResizeParam, 0, sizeof(GpuParam));
|
|
VulkanResizeParam->inImgSize[0] = input->width();
|
|
VulkanResizeParam->inImgSize[1] = input->height();
|
|
VulkanResizeParam->inImgSize[2] = channelDiv4;
|
|
VulkanResizeParam->inImgSize[3] = input->batch();
|
|
VulkanResizeParam->outImgSize[0] = output->width();
|
|
VulkanResizeParam->outImgSize[1] = output->height();
|
|
VulkanResizeParam->outImgSize[2] = channelDiv4;
|
|
VulkanResizeParam->outImgSize[3] = output->batch();
|
|
::memcpy(VulkanResizeParam->cord, cords, 4 * sizeof(float));
|
|
mParamBuffer->flush(true, 0, sizeof(GpuParam));
|
|
mParamBuffer->unmap();
|
|
|
|
auto vkOutput = reinterpret_cast<VulkanTensor*>(output->deviceId());
|
|
auto vkInput = reinterpret_cast<VulkanTensor*>(input->deviceId());
|
|
|
|
mDescriptorSet.reset(mVulkanResizePipeline->createSet());
|
|
mDescriptorSet->writeImage(((VulkanTensor*)input->deviceId())->image()->view(), extra->getCommonSampler()->get(),
|
|
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 0);
|
|
mDescriptorSet->writeImage(((VulkanTensor*)output->deviceId())->image()->view(), extra->getCommonSampler()->get(),
|
|
VK_IMAGE_LAYOUT_GENERAL, 1);
|
|
mDescriptorSet->writeBuffer(mParamBuffer->buffer(), 2, mParamBuffer->size());
|
|
mVulkanResizePipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
|
|
|
|
vkInput->image()->barrierRead(cmdBuffer->get());
|
|
vkOutput->image()->barrierWrite(cmdBuffer->get());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(output->width(), 16), UP_DIV(output->height(), 16),
|
|
channelDiv4 * input->batch());
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
|
|
} // namespace MNN
|