root
/
MNN
mirror of https://github.com/alibaba/MNN.git
1
0
Fork 0
Code Issues Actions 6 Packages Projects Releases Wiki Activity
MNN/source/backend/vulkan/execution/VulkanBinary.cpp

255 lines
11 KiB
C++
Raw Blame History

//
// VulkanBinary.cpp
// MNN
//
// Created by MNN on 2019/01/31.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include "backend/vulkan/execution/VulkanBinary.hpp"
#include "core/Macro.h"
#include "core/TensorUtils.hpp"
#include "core/OpCommonUtils.hpp"
namespace MNN {
struct ConstBuffer {
ivec4 stride00;
ivec4 stride01;
ivec4 stride10;
ivec4 stride11;
ivec4 stride20;
ivec4 stride21;
};
static std::string _getShaderName(const Op* op, bool image) {
std::string prefix = "glsl_binaryBroadcast_";
if (image) {
prefix = "glsl_binaryImage_";
}
std::string posfix = "_comp";
std::string mid = "";
if (op->type() == OpType_Eltwise) {
if (op->main_as_Eltwise()->coeff() != nullptr) {
// Don't support
return "";
}
switch (op->main_as_Eltwise()->type()) {
case EltwiseType_SUB:
mid = "SUB";
break;
case EltwiseType_MAXIMUM:
mid = "VMAX";
break;
case EltwiseType_PROD:
mid = "MUL";
break;
case EltwiseType_SUM:
mid = "ADD";
break;
default:
break;
}
} else if (op->type() == OpType_BinaryOp) {
switch (op->main_as_BinaryOp()->opType()) {
case BinaryOpOperation_ADD:
mid = "ADD";
break;
case BinaryOpOperation_SUB:
mid = "SUB";
break;
case BinaryOpOperation_MAXIMUM:
mid = "VMAX";
break;
case BinaryOpOperation_MINIMUM:
mid = "VMIN";
break;
case BinaryOpOperation_MUL:
mid = "MUL";
break;
case BinaryOpOperation_DIV:
case BinaryOpOperation_REALDIV:
mid = "DIV";
break;
default:
break;
}
}
if (mid.empty()) {
return mid;
}
return prefix + mid + posfix;
}
VulkanBinary::VulkanBinary(const std::string& shaderName, Backend* bn, bool image) : VulkanBasicExecution(bn) {
auto vkBn = static_cast<VulkanBackend*>(bn);
mImage = image;
mConstBuffer = std::make_shared<VulkanBuffer>(vkBn->getMemoryPool(), false, sizeof(ConstBuffer), nullptr,
VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT);
if (image) {
mBinaryPipeline = vkBn->getPipeline(shaderName, {
VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER
});
} else {
mBinaryPipeline = vkBn->getPipeline(shaderName, {
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER
});
}
mDescriptorSet.reset(mBinaryPipeline->createSet());
}
VulkanBinary::~VulkanBinary() {
}
ErrorCode VulkanBinary::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
const VulkanCommandPool::Buffer* cmdBuffer) {
MNN_ASSERT(1 == outputs.size());
auto vkBn = (VulkanBackend*)backend();
const int outputElements = outputs[0]->elementSize();
{
auto input0 = inputs[0];
auto input1 = inputs[1];
auto output = outputs[0];
MNN_ASSERT(input0->getType().code == halide_type_float);
if (!mImage) {
// for buffer input
#define MAX_DIM 6
int dims[MAX_DIM];
int stride[MAX_DIM];
int iStride0[MAX_DIM];
int iStride1[MAX_DIM];
OpCommonUtils::broastCastComputeDim(dims, stride, iStride0, iStride1, input0, input1, output);
auto binaryOpParam = reinterpret_cast<ConstBuffer*>(mConstBuffer->map());
binaryOpParam->stride01[3] = outputElements;
binaryOpParam->stride01[2] = 1;
binaryOpParam->stride01[1] = 1;
binaryOpParam->stride01[0] = dims[5];
binaryOpParam->stride00[3] = dims[4] * binaryOpParam->stride01[0];
binaryOpParam->stride00[2] = dims[3] * binaryOpParam->stride00[3];
binaryOpParam->stride00[1] = dims[2] * binaryOpParam->stride00[2];
binaryOpParam->stride00[0] = dims[1] * binaryOpParam->stride00[1];
::memcpy(binaryOpParam->stride10, iStride0, 4 * sizeof(int));
::memcpy(binaryOpParam->stride11, iStride0 + 4, 2 * sizeof(int));
::memcpy(binaryOpParam->stride20, iStride1, 4 * sizeof(int));
::memcpy(binaryOpParam->stride21, iStride1 + 4, 2 * sizeof(int));
mConstBuffer->unmap();
mDescriptorSet->writeBuffer(reinterpret_cast<VkBuffer>(output->deviceId()), 0, output->size());
mDescriptorSet->writeBuffer(reinterpret_cast<VkBuffer>(input0->deviceId()), 1, input0->size());
mDescriptorSet->writeBuffer(reinterpret_cast<VkBuffer>(input1->deviceId()), 2, input1->size());
mDescriptorSet->writeBuffer(mConstBuffer->buffer(), 3, mConstBuffer->size());
mBinaryPipeline->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(outputElements, 256), 1, 1);
} else {
// for NC4HW4 input
const int iw = input0->width();
const int ih = input0->height();
MNN_ASSERT(input0->dimensions() == input1->dimensions());
const int icDiv4 = UP_DIV(input0->channel(), 4);
auto total = icDiv4 * input0->batch() * iw * ih;
auto binaryOpParam = reinterpret_cast<ConstBuffer*>(mConstBuffer->map());
::memset(binaryOpParam, 0, sizeof(ConstBuffer));
binaryOpParam->stride00[3] = total;
binaryOpParam->stride00[0] = iw;
binaryOpParam->stride00[1] = ih;
binaryOpParam->stride00[2] = icDiv4;
mConstBuffer->unmap();
auto sampler = vkBn->getCommonSampler();
mDescriptorSet->writeImage(reinterpret_cast<VkImageView>(output->deviceId()), sampler->get(),
VK_IMAGE_LAYOUT_GENERAL, 0);
auto input0T = vkBn->findTensor(input0->deviceId());
auto input1T = vkBn->findTensor(input1->deviceId());
cmdBuffer->barrierImage(input0T->image()->get(), VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
cmdBuffer->barrierImage(input1T->image()->get(), VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
mDescriptorSet->writeImage(reinterpret_cast<VkImageView>(input0->deviceId()), sampler->get(),
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 1);
mDescriptorSet->writeImage(reinterpret_cast<VkImageView>(input1->deviceId()), sampler->get(),
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 2);
mDescriptorSet->writeBuffer(mConstBuffer->buffer(), 3, mConstBuffer->size());
mBinaryPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
vkCmdDispatch(cmdBuffer->get(), UP_DIV(total, 256), 1, 1);
}
}
if (inputs.size() > 2) {
mExtraDescriptorSet.clear();
for (int i=2; i<inputs.size(); ++i) {
auto input0 = outputs[0];
auto input1 = inputs[i];
auto output = outputs[0];
std::shared_ptr<VulkanPipeline::DescriptorSet> newSet(mBinaryPipeline->createSet());
mExtraDescriptorSet.push_back(newSet);
if (!mImage) {
newSet->writeBuffer(reinterpret_cast<VkBuffer>(output->deviceId()), 0, output->size());
newSet->writeBuffer(reinterpret_cast<VkBuffer>(input0->deviceId()), 1, input0->size());
newSet->writeBuffer(reinterpret_cast<VkBuffer>(input1->deviceId()), 2, input1->size());
newSet->writeBuffer(mConstBuffer->buffer(), 3, mConstBuffer->size());
mBinaryPipeline->bind(cmdBuffer->get(), newSet->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(outputElements, 256), 1, 1);
} else {
// for NC4HW4 input
const int iw = input0->width();
const int ih = input0->height();
const int icDiv4 = UP_DIV(input0->channel(), 4);
auto total = icDiv4 * input0->batch() * iw * ih;
auto sampler = vkBn->getCommonSampler();
auto input0T = vkBn->findTensor(input0->deviceId());
auto input1T = vkBn->findTensor(input1->deviceId());
cmdBuffer->barrierImage(input0T->image()->get(), VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
cmdBuffer->barrierImage(input1T->image()->get(), VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
newSet->writeImage(reinterpret_cast<VkImageView>(output->deviceId()), sampler->get(),
VK_IMAGE_LAYOUT_GENERAL, 0);
newSet->writeImage(reinterpret_cast<VkImageView>(input0->deviceId()), sampler->get(),
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 1);
newSet->writeImage(reinterpret_cast<VkImageView>(input1->deviceId()), sampler->get(),
VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 2);
newSet->writeBuffer(mConstBuffer->buffer(), 3, mConstBuffer->size());
mBinaryPipeline->bind(cmdBuffer->get(), newSet->get());
vkCmdDispatch(cmdBuffer->get(), UP_DIV(total, 256), 1, 1);
}
}
}
return NO_ERROR;
}
class VulkanBinaryCreator : public VulkanBackend::Creator {
public:
virtual VulkanBasicExecution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const MNN::Op* op,
Backend* backend) const override {
auto input0 = inputs[0];
if (input0->getType().code != halide_type_float) {
return nullptr;
}
auto image = TensorUtils::getDescribe(input0)->dimensionFormat == MNN_DATA_FORMAT_NC4HW4;
auto shader = _getShaderName(op, image);
if (shader.empty()) {
return nullptr;
}
return new VulkanBinary(shader, backend, image);
}
};
static bool gResistor = []() {
VulkanBackend::addCreator(OpType_BinaryOp, new VulkanBinaryCreator);
VulkanBackend::addCreator(OpType_Eltwise, new VulkanBinaryCreator);
return true;
}();
} // namespace MNN
Reference in New Issue View Git Blame Copy Permalink