wehub-resource-sync
338 lines
14 KiB
C++
338 lines
14 KiB
C++
#include "VulkanLoop.hpp"
|
|
#include "VulkanBinary.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include <algorithm>
|
|
#include "core/OpCommonUtils.hpp"
|
|
#include "core/Macro.h"
|
|
|
|
namespace MNN {
|
|
|
|
static void _setTensorStack(std::vector<Tensor*>& result, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const LoopParam* loop) {
|
|
if (loop->inputIndexes() != nullptr) {
|
|
for (int i=0; i<loop->inputIndexes()->size(); ++i) {
|
|
result[loop->inputIndexes()->data()[i]] = inputs[i];
|
|
}
|
|
}
|
|
for (int i=0; i<loop->outputIndexes()->size(); ++i) {
|
|
result[loop->outputIndexes()->data()[i]] = outputs[i];
|
|
}
|
|
}
|
|
|
|
class VulkanBatchMatMul : public VulkanBasicExecution {
|
|
public:
|
|
VulkanBatchMatMul(const LoopParam* loop, Backend *bn, Tensor * tensor) : VulkanBasicExecution(bn) {
|
|
mLoop = loop;
|
|
auto vkbackend = static_cast<VulkanBackend*>(bn);
|
|
mParam.reset(new VulkanBuffer(vkbackend->getMemoryPool(), false, sizeof(VulkanBatchMatMulInfo), nullptr, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
|
|
|
auto cmd = loop->commands()->GetAs<RegionCommand>(0);
|
|
mHasBias = cmd->indexes()->size() > 3;
|
|
bool useFP16 = tensor->getType().code == halide_type_float && vkbackend->useFP16();
|
|
|
|
std::string prefix = "glsl_matmulunit_";
|
|
std::string mid = mHasBias ? "HAS_BIAS_" : "";
|
|
std::string postfix = useFP16 ? "FP16_comp" : "comp";
|
|
std::vector<VkDescriptorType> desTypes;
|
|
if (!mHasBias) {
|
|
desTypes = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
|
};
|
|
} else {
|
|
desTypes = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
|
};
|
|
}
|
|
mPipeline = vkbackend->getPipeline(prefix + mid + postfix, desTypes);
|
|
mDescribe.reset(mPipeline->createSet());
|
|
mTensors.resize(mLoop->tensorNumber());
|
|
}
|
|
virtual ~VulkanBatchMatMul() = default;
|
|
virtual ErrorCode onEncode(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) override {
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(0);
|
|
auto size = cmd->size()->data();
|
|
auto AStride = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto BStride = cmd->view()->GetAs<View>(2)->stride()->data();
|
|
auto OStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
int totalSize = mLoop->loopNumber() * size[0] * size[2];
|
|
auto param = reinterpret_cast<VulkanBatchMatMulInfo*>(mParam->map());
|
|
param->size[3] = mLoop->loopNumber();
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
for (int i=0; i<3; ++i) {
|
|
param->size[i] = size[i];
|
|
param->stride_o[i] = OStride[i];
|
|
param->stride_a[i] = AStride[i];
|
|
param->stride_b[i] = BStride[i];
|
|
}
|
|
param->stride_o[3] = cmd->view()->GetAs<View>(0)->offset();
|
|
param->stride_a[3] = cmd->view()->GetAs<View>(1)->offset();
|
|
param->stride_b[3] = cmd->view()->GetAs<View>(2)->offset();
|
|
if (mHasBias) {
|
|
param->stride_c[3] = cmd->view()->GetAs<View>(3)->offset();
|
|
}
|
|
::memcpy(param->step, cmd->steps()->data(), cmd->steps()->size() * sizeof(int));
|
|
::memcpy(param->iter, cmd->iterIndexes()->data(), cmd->iterIndexes()->size() * sizeof(int));
|
|
std::vector<VULKAN_TENSOR> strideBuffers(cmd->indexes()->size());
|
|
for (int i=0; i<cmd->indexes()->size(); ++i) {
|
|
std::get<0>(strideBuffers[i]) = 0;
|
|
if (param->iter[i] >= 0) {
|
|
strideBuffers[i] = vkBn->getBuffer(mTensors[param->iter[i]]);
|
|
}
|
|
}
|
|
mParam->unmap();
|
|
for (int i=0; i<cmd->indexes()->size(); ++i) {
|
|
auto tensor = mTensors[cmd->indexes()->data()[i]];
|
|
mDescribe->writeBuffer(vkBn->getBuffer(tensor), i);
|
|
}
|
|
for (int i=0; i<strideBuffers.size(); ++i) {
|
|
if (0 != std::get<0>(strideBuffers[i])) {
|
|
mDescribe->writeBuffer(strideBuffers[i], cmd->indexes()->size() + i);
|
|
} else {
|
|
mDescribe->writeBuffer(vkBn->getBuffer(inputs[0]), cmd->indexes()->size() + i);
|
|
}
|
|
}
|
|
mDescribe->writeBuffer(mParam->buffer(), cmd->indexes()->size() * 2, mParam->size());
|
|
mPipeline->bind(cmdBuffer->get(), mDescribe->get());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(totalSize,256), 1, 1);
|
|
|
|
return NO_ERROR;
|
|
}
|
|
private:
|
|
const LoopParam* mLoop;
|
|
const VulkanPipeline* mPipeline;
|
|
std::shared_ptr<VulkanBuffer> mParam;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mDescribe;
|
|
std::vector<Tensor*> mTensors;
|
|
bool mHasBias = false;
|
|
};
|
|
|
|
|
|
struct BinaryBroadCastInfo {
|
|
ivec4 srcview0;
|
|
ivec4 srcview1;
|
|
ivec4 dstview;
|
|
ivec4 size;
|
|
};
|
|
|
|
class VulkanBinaryBroadCast : public VulkanBasicExecution {
|
|
public:
|
|
VulkanBinaryBroadCast(const LoopParam* loop, Backend *bn, bool isInt, Tensor * tensor) : VulkanBasicExecution(bn) {
|
|
mLoop = loop;
|
|
auto vkbackend = static_cast<VulkanBackend*>(bn);
|
|
mParam.reset(new VulkanBuffer(vkbackend->getMemoryPool(), false, sizeof(BinaryBroadCastInfo), nullptr, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
|
std::string shaderName;
|
|
if (isInt) {
|
|
shaderName = "glsl_binary_blit_int_" + VulkanBinary::getMidName( mLoop->commands()->GetAs<RegionCommand>(0)->op()) + "_comp";
|
|
} else {
|
|
shaderName = "glsl_binary_blit_" + VulkanBinary::getMidName( mLoop->commands()->GetAs<RegionCommand>(0)->op()) + "_";
|
|
if (tensor->getType().code == halide_type_float && vkbackend->useFP16()) {
|
|
shaderName += "FP16_";
|
|
}
|
|
shaderName += "comp";
|
|
}
|
|
|
|
mPipeline = vkbackend->getPipeline(shaderName, {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
|
});
|
|
mDescribe.reset(mPipeline->createSet());
|
|
mTensors.resize(mLoop->tensorNumber());
|
|
}
|
|
virtual ~VulkanBinaryBroadCast() = default;
|
|
virtual ErrorCode onEncode(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) override {
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(0);
|
|
auto size = cmd->size()->data();
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
auto srcStride0 = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto srcStride1 = cmd->view()->GetAs<View>(2)->stride()->data();
|
|
auto dstStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
int totalSize = size[0] * size[1] * size[2];
|
|
auto param = reinterpret_cast<BinaryBroadCastInfo*>(mParam->map());
|
|
for (int i=0; i<3; ++i) {
|
|
param->size[i] = size[i];
|
|
param->srcview0[i] = srcStride0[i];
|
|
param->srcview1[i] = srcStride1[i];
|
|
param->dstview[i] = dstStride[i];
|
|
}
|
|
param->srcview0[3] = cmd->view()->GetAs<View>(1)->offset();
|
|
param->srcview1[3] = cmd->view()->GetAs<View>(2)->offset();
|
|
param->dstview[3] = cmd->view()->GetAs<View>(0)->offset();
|
|
param->size[3] = size[0] * size[1] * size[2];
|
|
mParam->unmap();
|
|
auto dstTensor = mTensors[cmd->indexes()->data()[0]];
|
|
auto srcTensor = mTensors[cmd->indexes()->data()[1]];
|
|
auto srcTensor1 = mTensors[cmd->indexes()->data()[2]];
|
|
mDescribe->writeBuffer(vkBn->getBuffer(dstTensor), 0);
|
|
mDescribe->writeBuffer(vkBn->getBuffer(srcTensor), 1);
|
|
mDescribe->writeBuffer(vkBn->getBuffer(srcTensor1), 2);
|
|
mDescribe->writeBuffer(mParam->buffer(), 3, mParam->size());
|
|
mPipeline->bind(cmdBuffer->get(), mDescribe->get());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(totalSize,256), 1, 1);
|
|
|
|
return NO_ERROR;
|
|
}
|
|
private:
|
|
const LoopParam* mLoop;
|
|
const VulkanPipeline* mPipeline;
|
|
std::shared_ptr<VulkanBuffer> mParam;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mDescribe;
|
|
std::vector<Tensor*> mTensors;
|
|
};
|
|
struct GatherInfo {
|
|
ivec4 stride;
|
|
ivec4 size;
|
|
ivec4 extent;
|
|
ivec4 step;
|
|
ivec4 iter;
|
|
};
|
|
|
|
class VulkanGather : public VulkanBasicExecution {
|
|
public:
|
|
VulkanGather(const LoopParam* loop, Backend *bn, Tensor * tensor) : VulkanBasicExecution(bn) {
|
|
mLoop = loop;
|
|
auto vkbackend = static_cast<VulkanBackend*>(bn);
|
|
mParam.reset(new VulkanBuffer(vkbackend->getMemoryPool(), false, sizeof(GatherInfo), nullptr, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
|
|
|
std::string pKey = "glsl_blitregion_";
|
|
if (tensor->getType().code == halide_type_float && vkbackend->useFP16()) {
|
|
pKey += "FP16_";
|
|
}
|
|
pKey += "comp";
|
|
|
|
mPipeline = vkbackend->getPipeline(pKey, {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
|
});
|
|
mDescribe.reset(mPipeline->createSet());
|
|
mTensors.resize(mLoop->tensorNumber());
|
|
}
|
|
virtual ~VulkanGather() = default;
|
|
virtual ErrorCode onEncode(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) override {
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(0);
|
|
auto size = cmd->size()->data();
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
auto srcStride = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto dstStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
int totalSize = mLoop->loopNumber() * size[0] * size[1] * size[2];
|
|
VULKAN_TENSOR srcOffsetBuffer;
|
|
std::get<0>(srcOffsetBuffer) = 0;
|
|
VULKAN_TENSOR dstOffsetBuffer;
|
|
std::get<0>(dstOffsetBuffer) = 0;
|
|
auto param = reinterpret_cast<GatherInfo*>(mParam->map());
|
|
for (int i=0; i<3; ++i) {
|
|
param->size[i] = size[i];
|
|
param->stride[i] = srcStride[i];
|
|
param->extent[i] = dstStride[i];
|
|
}
|
|
param->stride[3] = cmd->view()->GetAs<View>(1)->offset();
|
|
param->extent[3] = cmd->view()->GetAs<View>(0)->offset();
|
|
param->size[3] = size[0] * size[1] * size[2];
|
|
param->step[3] = totalSize;
|
|
param->step[0] = cmd->steps()->data()[0];
|
|
param->step[1] = cmd->steps()->data()[1];
|
|
param->iter[0] = 0;
|
|
param->iter[1] = 0;
|
|
auto iterIndex = cmd->iterIndexes()->data();
|
|
if (iterIndex[0] >= 0) {
|
|
dstOffsetBuffer = vkBn->getBuffer(mTensors[iterIndex[0]]);
|
|
param->iter[0] = 1;
|
|
}
|
|
if (iterIndex[1] >= 0) {
|
|
srcOffsetBuffer = vkBn->getBuffer(mTensors[iterIndex[1]]);
|
|
param->iter[1] = 1;
|
|
}
|
|
mParam->unmap();
|
|
auto dstTensor = mTensors[cmd->indexes()->data()[0]];
|
|
auto srcTensor = mTensors[cmd->indexes()->data()[1]];
|
|
mDescribe->writeBuffer(vkBn->getBuffer(dstTensor), 0);
|
|
mDescribe->writeBuffer(vkBn->getBuffer(srcTensor), 1);
|
|
if (std::get<0>(srcOffsetBuffer) != 0) {
|
|
mDescribe->writeBuffer(srcOffsetBuffer, 2);
|
|
} else {
|
|
// Use Invalide buffer
|
|
mDescribe->writeBuffer(vkBn->getBuffer(srcTensor), 2);
|
|
}
|
|
if (std::get<0>(dstOffsetBuffer) != 0) {
|
|
mDescribe->writeBuffer(dstOffsetBuffer, 3);
|
|
} else {
|
|
// Use Invalide buffer
|
|
mDescribe->writeBuffer(vkBn->getBuffer(srcTensor), 3);
|
|
}
|
|
mDescribe->writeBuffer(mParam->buffer(), 4, mParam->size());
|
|
mPipeline->bind(cmdBuffer->get(), mDescribe->get());
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(totalSize,256), 1, 1);
|
|
|
|
return NO_ERROR;
|
|
}
|
|
private:
|
|
const LoopParam* mLoop;
|
|
const VulkanPipeline* mPipeline;
|
|
std::shared_ptr<VulkanBuffer> mParam;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mDescribe;
|
|
std::vector<Tensor*> mTensors;
|
|
};
|
|
|
|
VulkanBasicExecution* VulkanLoop::create(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const Op* op, Backend* bn) {
|
|
auto loop = op->main_as_LoopParam();
|
|
if (nullptr == loop || loop->commands() == nullptr) {
|
|
return nullptr;
|
|
}
|
|
if (nullptr != loop->initCommand()) {
|
|
return nullptr;
|
|
}
|
|
// Make Tensor Stack
|
|
if (1 == loop->commands()->size()) {
|
|
auto cmd = loop->commands()->GetAs<RegionCommand>(0);
|
|
auto subop = cmd->op();
|
|
if (OpType_UnaryOp == subop->type() && nullptr == subop->main() && cmd->fuse() < 0) {
|
|
return new VulkanGather(loop, bn, inputs[0]);
|
|
}
|
|
if (OpType_MatMul == subop->type() && loop->parallel()) {
|
|
return new VulkanBatchMatMul(loop, bn, inputs[0]);
|
|
}
|
|
if (OpType_BinaryOp == subop->type() && cmd->fuse() < 0 && 1 == loop->loopNumber()) {
|
|
bool isInt = inputs[1]->getType().code == halide_type_int;
|
|
return new VulkanBinaryBroadCast(loop, bn, isInt, inputs[0]);
|
|
}
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
class VulkanLoopCreator : 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 VulkanLoop::create(inputs, outputs, op, bn);
|
|
}
|
|
};
|
|
|
|
static bool gResistor = []() {
|
|
VulkanBackend::addCreator(OpType_While, new VulkanLoopCreator);
|
|
return true;
|
|
}();
|
|
|
|
};
|