mirror of https://github.com/alibaba/MNN.git
84 lines
3.6 KiB
C++
84 lines
3.6 KiB
C++
//
|
|
// VulkanSpatialProduct.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "backend/vulkan/execution/VulkanSpatialProduct.hpp"
|
|
#include "core/Macro.h"
|
|
|
|
namespace MNN {
|
|
struct GpuParam {
|
|
ivec4 imgsize;
|
|
};
|
|
|
|
VulkanSpatialProduct::VulkanSpatialProduct(const Op* op, Backend* bn) : VulkanBasicExecution(bn) {
|
|
std::vector<VkDescriptorType> VulkanSpatialProductTypes{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
|
|
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
|
|
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
|
};
|
|
auto extra = static_cast<VulkanBackend*>(bn);
|
|
mVulkanSpatialProductPipeline =
|
|
extra->getPipeline("glsl_SpatialProduct_comp",
|
|
/*glsl_SpatialProduct_comp, glsl_SpatialProduct_comp_len,*/ VulkanSpatialProductTypes);
|
|
mParamBuffer = std::make_shared<VulkanBuffer>(extra->getMemoryPool(), false, sizeof(GpuParam), nullptr,
|
|
VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT);
|
|
mSampler = extra->getCommonSampler();
|
|
}
|
|
VulkanSpatialProduct::~VulkanSpatialProduct() {
|
|
}
|
|
ErrorCode VulkanSpatialProduct::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) {
|
|
auto input0 = inputs[0];
|
|
auto input1 = inputs[1];
|
|
MNN_ASSERT(input1->channel() == 1);
|
|
MNN_ASSERT(input0->width() == input1->width());
|
|
MNN_ASSERT(input0->height() == input1->height());
|
|
auto output = outputs[0];
|
|
const int channelDiv4 = UP_DIV(input0->channel(), 4);
|
|
auto VulkanSpatialProductParam = reinterpret_cast<GpuParam*>(mParamBuffer->map());
|
|
::memset(VulkanSpatialProductParam, 0, sizeof(GpuParam));
|
|
VulkanSpatialProductParam->imgsize[0] = input0->width();
|
|
VulkanSpatialProductParam->imgsize[1] = input0->height();
|
|
VulkanSpatialProductParam->imgsize[2] = channelDiv4;
|
|
VulkanSpatialProductParam->imgsize[3] = 0;
|
|
mParamBuffer->flush(true, 0, sizeof(GpuParam));
|
|
mParamBuffer->unmap();
|
|
|
|
mDescriptorSet.reset(mVulkanSpatialProductPipeline->createSet());
|
|
|
|
mDescriptorSet->writeImage(reinterpret_cast<VkImageView>(output->deviceId()), mSampler->get(),
|
|
VK_IMAGE_LAYOUT_GENERAL, 0);
|
|
mDescriptorSet->writeImage(reinterpret_cast<VkImageView>(input0->deviceId()), mSampler->get(),
|
|
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 1);
|
|
mDescriptorSet->writeImage(reinterpret_cast<VkImageView>(input1->deviceId()), mSampler->get(),
|
|
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 2);
|
|
mDescriptorSet->writeBuffer(mParamBuffer->buffer(), 3, mParamBuffer->size());
|
|
|
|
mVulkanSpatialProductPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
|
|
cmdBuffer->barrierSource(reinterpret_cast<VkBuffer>(input0->deviceId()), 0, input0->size());
|
|
cmdBuffer->barrierSource(reinterpret_cast<VkBuffer>(input1->deviceId()), 0, input1->size());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(input0->width(), 8), UP_DIV(input0->height(), 8),
|
|
channelDiv4 * input0->batch());
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class VulkanSpatialProductCreator : 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 {
|
|
return new VulkanSpatialProduct(op, bn);
|
|
}
|
|
};
|
|
|
|
static bool gResistor = []() {
|
|
VulkanBackend::addCreator(OpType_SpatialProduct, new VulkanSpatialProductCreator);
|
|
return true;
|
|
}();
|
|
|
|
} // namespace MNN
|