chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,234 @@
|
||||
//
|
||||
// Vulkan.cpp
|
||||
// MNN
|
||||
//
|
||||
// Created by MNN on 2023/07/25.
|
||||
// Copyright © 2018, Alibaba Group Holding Limited
|
||||
//
|
||||
|
||||
#include <stdio.h>
|
||||
#include "VulkanBasicExecution.hpp"
|
||||
#include "core/OpCommonUtils.hpp"
|
||||
namespace MNN {
|
||||
struct ConstBuffer {
|
||||
ivec4 inShape; // inW, inH
|
||||
};
|
||||
|
||||
class VulkanFuse : public VulkanBasicExecution {
|
||||
public:
|
||||
VulkanFuse(const Extra* extra, Backend* bn, int inputSize, int outputSize) : VulkanBasicExecution(bn) {
|
||||
auto vkBn = static_cast<VulkanBackend*>(bn);
|
||||
auto factory = vkBn->getPipelineFactory();
|
||||
mOutputBinding.resize(outputSize);
|
||||
mInputBinding.resize(inputSize);
|
||||
mGroupSize.resize(3);
|
||||
mGlobalSize.resize(3);
|
||||
// Find shader
|
||||
const uint8_t* data = nullptr;
|
||||
size_t dataSize = 0;
|
||||
for (int i=0; i<extra->attr()->size(); ++i) {
|
||||
auto attr = extra->attr()->GetAs<Attribute>(i);
|
||||
if (attr->key()->str() == "spirv") {
|
||||
data = (uint8_t*)attr->tensor()->int8s()->data();
|
||||
dataSize = attr->tensor()->int8s()->size();
|
||||
break;
|
||||
}
|
||||
}
|
||||
for (int i=0; i<extra->attr()->size(); ++i) {
|
||||
auto attr = extra->attr()->GetAs<Attribute>(i);
|
||||
if (attr->key()->str() == "global_size") {
|
||||
// Use Auto set group size
|
||||
auto ptr = attr->tensor()->int32s()->data();
|
||||
mGlobalSize[0] = ptr[0];
|
||||
mGlobalSize[1] = ptr[1];
|
||||
mGlobalSize[2] = ptr[2];
|
||||
mNeedAutoTuning = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// If Has group_size, can't auto tuning
|
||||
for (int i=0; i<extra->attr()->size(); ++i) {
|
||||
auto attr = extra->attr()->GetAs<Attribute>(i);
|
||||
if (attr->key()->str() == "group_size") {
|
||||
auto ptr = attr->tensor()->int32s()->data();
|
||||
mGroupSize[0] = ptr[0];
|
||||
mGroupSize[1] = ptr[1];
|
||||
mGroupSize[2] = ptr[2];
|
||||
mNeedAutoTuning = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<VkDescriptorType> types;
|
||||
int maxIndex = -1;
|
||||
for (int i=0; i<extra->attr()->size(); ++i) {
|
||||
auto attr = extra->attr()->GetAs<Attribute>(i);
|
||||
if (attr->key()->str() == "input") {
|
||||
maxIndex = ALIMAX(maxIndex, attr->i());
|
||||
} else if (attr->key()->str() == "const") {
|
||||
maxIndex = ALIMAX(maxIndex, attr->i());
|
||||
}
|
||||
}
|
||||
types.resize(maxIndex+1);
|
||||
std::vector<std::tuple<int, void*, size_t>> constStoragePtrs;
|
||||
std::vector<std::tuple<int, void*, size_t>> constUniformPtrs;
|
||||
for (int i=0; i<extra->attr()->size(); ++i) {
|
||||
auto attr = extra->attr()->GetAs<Attribute>(i);
|
||||
if (attr->key()->str() == "input") {
|
||||
auto list = attr->list()->i()->data();
|
||||
if (list[1] >= 0) {
|
||||
if (0 == list[0]) {
|
||||
mInputBinding[list[1]] = attr->i();
|
||||
} else {
|
||||
mOutputBinding[list[1]] = attr->i();
|
||||
}
|
||||
}
|
||||
if (attr->b()) {
|
||||
types[attr->i()] = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
|
||||
} else {
|
||||
types[attr->i()] = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (attr->key()->str() == "const") {
|
||||
auto b = attr->tensor();
|
||||
void* result = nullptr;
|
||||
size_t bufferSize = 0;
|
||||
switch (b->dataType()) {
|
||||
case DataType_DT_FLOAT:
|
||||
result = (void*)b->float32s()->Data();
|
||||
bufferSize = b->float32s()->size() * sizeof(float);
|
||||
break;
|
||||
case DataType_DT_INT32:
|
||||
result = (void*)b->int32s()->Data();
|
||||
bufferSize = b->int32s()->size() * sizeof(float);
|
||||
break;
|
||||
default:
|
||||
MNN_ASSERT(false);
|
||||
break;
|
||||
}
|
||||
if (attr->b()) {
|
||||
types[attr->i()] = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
|
||||
constUniformPtrs.emplace_back(std::make_tuple(attr->i(), result, bufferSize));
|
||||
} else {
|
||||
types[attr->i()] = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
|
||||
constStoragePtrs.emplace_back(std::make_tuple(attr->i(), result, bufferSize));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
auto alignSize = vkBn->device().proty().limits.minMemoryMapAlignment;
|
||||
size_t offset = 0;
|
||||
std::shared_ptr<VulkanCommandPool::Buffer> cmdbuffer( vkBn->getPool().allocBuffer());
|
||||
cmdbuffer->begin(0);
|
||||
auto merge = [&](const std::vector<std::tuple<int, void*, size_t>>& constPtrs, VkDescriptorType type) {
|
||||
if (constPtrs.empty()) {
|
||||
return std::make_tuple(std::vector<std::tuple<int, size_t, size_t>>{}, std::shared_ptr<VulkanBuffer>(nullptr), std::shared_ptr<VulkanBuffer>(nullptr));
|
||||
}
|
||||
std::vector<std::tuple<int, size_t, size_t>> mConstOffset;
|
||||
for (auto& constAttr : constPtrs) {
|
||||
auto size = UP_DIV(std::get<2>(constAttr), alignSize) * alignSize;
|
||||
mConstOffset.emplace_back(std::make_tuple(std::get<0>(constAttr), size, offset));
|
||||
offset += size;
|
||||
}
|
||||
std::shared_ptr<VulkanBuffer> hostBuffer(new VulkanBuffer(vkBn->getMemoryPool(), false, offset, nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_SHARING_MODE_EXCLUSIVE, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT));
|
||||
auto ptr = (uint8_t*)hostBuffer->map();
|
||||
for (int i=0; i<constPtrs.size(); ++i) {
|
||||
::memcpy(ptr + std::get<2>(mConstOffset[i]), std::get<1>(constPtrs[i]), std::get<2>(constPtrs[i]));
|
||||
}
|
||||
hostBuffer->unmap();
|
||||
std::shared_ptr<VulkanBuffer> vkBuffer(new VulkanBuffer(vkBn->getMemoryPool(), false, offset, nullptr, type, VK_SHARING_MODE_EXCLUSIVE, 0));
|
||||
VkBufferCopy bufferCopy;
|
||||
bufferCopy.size = offset;
|
||||
bufferCopy.dstOffset = 0;
|
||||
bufferCopy.srcOffset = 0;
|
||||
vkCmdCopyBuffer(cmdbuffer->get(), hostBuffer->buffer(), vkBuffer->buffer(),
|
||||
1, &bufferCopy);
|
||||
return std::make_tuple(mConstOffset, vkBuffer, hostBuffer);
|
||||
};
|
||||
mConstStorageOffset.clear();
|
||||
mConstUniformOffset.clear();
|
||||
auto uniforms = merge(constUniformPtrs, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER);
|
||||
mConstUniformOffset = std::get<0>(uniforms);
|
||||
mConstUniformBuffer = std::get<1>(uniforms);
|
||||
auto storages = merge(constStoragePtrs, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
||||
mConstStorageOffset = std::get<0>(storages);
|
||||
mConstStorageBuffer = std::get<1>(storages);
|
||||
cmdbuffer->end();
|
||||
auto fence = vkBn->getPool().submit(cmdbuffer->get());
|
||||
|
||||
mPipeline = factory->createComputePipeline(data, dataSize, types, std::vector<uint32_t>{});
|
||||
mDescriptorSet = mPipeline->createSet();
|
||||
fence->wait();
|
||||
}
|
||||
virtual ~VulkanFuse() {
|
||||
// Remove set firstly before destroy pipeline
|
||||
mDescriptorSet = nullptr;
|
||||
}
|
||||
virtual ErrorCode onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
||||
const VulkanCommandPool::Buffer* cmdBuffer) override {
|
||||
auto vkBn = static_cast<VulkanBackend*>(backend());
|
||||
for (int i=0; i<inputs.size(); ++i) {
|
||||
mDescriptorSet->writeBuffer(vkBn->getBuffer(inputs[i]), mInputBinding[i]);
|
||||
}
|
||||
for (int i=0; i<outputs.size(); ++i) {
|
||||
mDescriptorSet->writeBuffer(vkBn->getBuffer(outputs[i]), mOutputBinding[i]);
|
||||
}
|
||||
for (auto& iter : mConstStorageOffset) {
|
||||
mDescriptorSet->writeBuffer(mConstStorageBuffer->buffer(), std::get<0>(iter), std::get<1>(iter), std::get<2>(iter));
|
||||
}
|
||||
for (auto& iter : mConstUniformOffset) {
|
||||
mDescriptorSet->writeBuffer(mConstUniformBuffer->buffer(), std::get<0>(iter), std::get<1>(iter), std::get<2>(iter));
|
||||
}
|
||||
if (mNeedAutoTuning) {
|
||||
auto localSize = vkBn->autoTunePipeline(mPipeline.get(), mDescriptorSet, mGlobalSize);
|
||||
mPipeline->changePipeline(localSize);
|
||||
mGroupSize[0] = UP_DIV(mGlobalSize[0], localSize[0]);
|
||||
mGroupSize[1] = UP_DIV(mGlobalSize[1], localSize[1]);
|
||||
mGroupSize[2] = UP_DIV(mGlobalSize[2], localSize[2]);
|
||||
}
|
||||
mPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
|
||||
vkCmdDispatch(cmdBuffer->get(), mGroupSize[0], mGroupSize[1], mGroupSize[2]);
|
||||
return NO_ERROR;
|
||||
}
|
||||
private:
|
||||
std::vector<int> mGroupSize;
|
||||
std::vector<int> mGlobalSize;
|
||||
std::vector<int> mInputBinding;
|
||||
std::vector<int> mOutputBinding;
|
||||
std::shared_ptr<VulkanBuffer> mConstStorageBuffer;
|
||||
std::shared_ptr<VulkanBuffer> mConstUniformBuffer;
|
||||
// Index, offset, size
|
||||
std::vector<std::tuple<int, size_t, size_t>> mConstStorageOffset;
|
||||
std::vector<std::tuple<int, size_t, size_t>> mConstUniformOffset;
|
||||
SharedPtr<VulkanPipeline> mPipeline;
|
||||
SharedPtr<VulkanLayout::DescriptorSet> mDescriptorSet;
|
||||
bool mNeedAutoTuning = false;
|
||||
};
|
||||
|
||||
class VulkanFuseCreator : 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 extra = op->main_as_Extra();
|
||||
if (nullptr == extra) {
|
||||
return nullptr;
|
||||
}
|
||||
if (nullptr == extra->attr()) {
|
||||
return nullptr;
|
||||
}
|
||||
if(extra->type()->str() == "ExtraConvolution2DPrelu"){
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
return new VulkanFuse(extra, backend, (int)inputs.size(), (int)outputs.size());
|
||||
}
|
||||
};
|
||||
|
||||
static bool gResistor = []() {
|
||||
VulkanBackend::addCreator(OpType_Extra, new VulkanFuseCreator);
|
||||
return true;
|
||||
}();
|
||||
|
||||
} // namespace MNN
|
||||
|
||||
Reference in New Issue
Block a user