wehub-resource-sync
107 lines
3.8 KiB
C++
107 lines
3.8 KiB
C++
//
|
|
// VulkanPRelu.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "VulkanPRelu.hpp"
|
|
#include "VulkanUnary.hpp"
|
|
#include "core/Macro.h"
|
|
#include "core/TensorUtils.hpp"
|
|
namespace MNN {
|
|
|
|
struct GpuReluParam {
|
|
ivec4 imgSize;
|
|
};
|
|
|
|
//--------------------------Prelu--------------------------//
|
|
VulkanPrelu::VulkanPrelu(Backend *bn, const Op *op, Tensor * tensor) : VulkanBasicExecution(bn) {
|
|
std::vector<VkDescriptorType> types{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER
|
|
};
|
|
auto vulkanBn = static_cast<VulkanBackend *>(bn);
|
|
|
|
bool useFP16 = tensor->getType().code == halide_type_float && vulkanBn->useFP16();
|
|
std::string pKey = "glsl_preluWithChannel_";
|
|
if (useFP16) {
|
|
pKey += "FP16_";
|
|
}
|
|
pKey += "comp";
|
|
mPreluPipeline = vulkanBn->getPipeline(pKey, types);
|
|
const auto prelu = op->main_as_PRelu();
|
|
mGpuPreluParam = vulkanBn->allocUniform();
|
|
int count = ALIGN_UP4(prelu->slope()->size());
|
|
{
|
|
int bytes = useFP16 ? sizeof(int16_t) : sizeof(float);
|
|
std::shared_ptr<VulkanBuffer> slopeBuffer(new VulkanBuffer(
|
|
vulkanBn->getMemoryPool(), false, bytes * count, nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT));
|
|
auto slope = slopeBuffer->map();
|
|
::memset(slope, 0, count * bytes);
|
|
if (useFP16) {
|
|
FLOAT_TO_HALF(prelu->slope()->data(), (int16_t *)slope, prelu->slope()->size());
|
|
} else {
|
|
::memcpy(slope, prelu->slope()->data(), prelu->slope()->size() * sizeof(float));
|
|
}
|
|
slopeBuffer->unmap();
|
|
mSlope = slopeBuffer;
|
|
}
|
|
mDescriptorSet.reset(mPreluPipeline->createSet());
|
|
}
|
|
|
|
VulkanPrelu::~VulkanPrelu() {
|
|
auto extra = static_cast<VulkanBackend*>(backend());
|
|
extra->recycleUniform(mGpuPreluParam);
|
|
}
|
|
|
|
ErrorCode VulkanPrelu::onEncode(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs,
|
|
const VulkanCommandPool::Buffer *cmdBuffer) {
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
|
|
auto preluParam = reinterpret_cast<GpuReluParam *>(mGpuPreluParam->map());
|
|
::memset(preluParam, 0, sizeof(GpuReluParam));
|
|
auto vkBn = static_cast<VulkanBackend *>(backend());
|
|
|
|
const int channelDiv4 = UP_DIV(input->channel(), 4);
|
|
auto planeSize = input->width() * input->height() * input->batch();
|
|
preluParam->imgSize[0] = planeSize;
|
|
preluParam->imgSize[1] = channelDiv4;
|
|
preluParam->imgSize[2] = 1;
|
|
preluParam->imgSize[3] = channelDiv4 * planeSize;
|
|
mGpuPreluParam->unmap();
|
|
auto total = planeSize * channelDiv4;
|
|
|
|
auto vkOutput = vkBn->getBuffer(output);
|
|
auto vkInput = vkBn->getBuffer(input);
|
|
mDescriptorSet->writeBuffer(vkOutput, 0);
|
|
mDescriptorSet->writeBuffer(vkInput, 1);
|
|
mDescriptorSet->writeBuffer(mSlope->buffer(), 2, mSlope->size());
|
|
|
|
mDescriptorSet->writeBuffer(mGpuPreluParam->buffer(), 3, mGpuPreluParam->size());
|
|
mPreluPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(total, 256), 1, 1);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class VulkanReluCreator : 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 (1 == op->main_as_PRelu()->slopeCount()) {
|
|
return new VulkanUnary("RELU", bn, false, op->main_as_PRelu()->slope()->data()[0]);
|
|
}
|
|
return new VulkanPrelu(bn, op, outputs[0]);
|
|
}
|
|
};
|
|
|
|
static bool gr = []() {
|
|
VulkanBackend::addCreator(OpType_PReLU, new VulkanReluCreator);
|
|
return true;
|
|
}();
|
|
|
|
} // namespace MNN
|