wehub-resource-sync
130 lines
4.5 KiB
C++
130 lines
4.5 KiB
C++
//
|
||
// VulkanArgMax.cpp
|
||
// MNN
|
||
//
|
||
// Created by MNN on 2024/08/20.
|
||
// Copyright © 2018, Alibaba Group Holding Limited
|
||
//
|
||
|
||
#include "VulkanArgMax.hpp"
|
||
|
||
namespace MNN {
|
||
|
||
struct GpuArgMaxParam {
|
||
ivec4 size; // inside, mid, outside, 0
|
||
};
|
||
|
||
VulkanArgMax::VulkanArgMax(const Op* op, Backend* bn) : VulkanBasicExecution(bn) {
|
||
auto vkBn = (VulkanBackend *)backend();
|
||
|
||
mAxis = op->main_as_ArgMax()->axis();
|
||
|
||
std::vector<VkDescriptorType> types{
|
||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
||
};
|
||
if (op->type() == OpType_ArgMax) {
|
||
mArgMaxPipeline =
|
||
vkBn->getPipeline("glsl_argmax_comp", types);
|
||
} else {
|
||
MNN_ASSERT(op->type() == OpType_ArgMin);
|
||
mArgMaxPipeline =
|
||
vkBn->getPipeline("glsl_argmax_ARGMIN_comp", types);
|
||
}
|
||
|
||
mGpuArgMaxParam.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, sizeof(GpuArgMaxParam), nullptr, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT));
|
||
mDescriptorSet.reset(mArgMaxPipeline->createSet());
|
||
}
|
||
|
||
VulkanArgMax::~VulkanArgMax() {
|
||
}
|
||
|
||
// set descriptorSet， including output, input and GPU param
|
||
ErrorCode VulkanArgMax::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
||
const VulkanCommandPool::Buffer* cmdBuffer) {
|
||
auto vkBn = (VulkanBackend*)backend();
|
||
auto input = inputs[0];
|
||
auto output = outputs[0];
|
||
|
||
// set GPU param
|
||
auto axis = mAxis;
|
||
if (axis < 0) {
|
||
axis = input->dimensions() + axis;
|
||
}
|
||
int inside = 1;
|
||
int outside = 1;
|
||
int mid = input->length(axis);
|
||
for (int i=0; i<axis; ++i) {
|
||
outside *= input->length(i);
|
||
}
|
||
for (int i=axis+1; i<input->dimensions(); ++i) {
|
||
inside *= input->length(i);
|
||
}
|
||
auto total = outside * inside;
|
||
|
||
auto Argmax = reinterpret_cast<GpuArgMaxParam *>(mGpuArgMaxParam->map());
|
||
Argmax->size[0] = inside;
|
||
Argmax->size[1] = mid;
|
||
Argmax->size[2] = outside;
|
||
Argmax->size[3] = 0;
|
||
mGpuArgMaxParam->unmap();
|
||
|
||
// set necessary storages, set descriptorSet and bind commandBuffer
|
||
{
|
||
int bufferSizeSource = sizeof(float);
|
||
for (int i=0; i<input->dimensions(); ++i) {
|
||
bufferSizeSource *= input->length(i);
|
||
}
|
||
mSource.buffer.reset(new VulkanBuffer(vkBn->getDynamicMemoryPool(), false, bufferSizeSource, nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT));
|
||
mSource.convert.reset(new VulkanImageConverter(vkBn));
|
||
}
|
||
{
|
||
int bufferSizeOutput = sizeof(float);
|
||
for (int i=0; i<output->dimensions(); ++i) {
|
||
bufferSizeOutput *= output->length(i);
|
||
}
|
||
mOutput.convert.reset(new VulkanImageConverter(vkBn));
|
||
mOutput.buffer.reset(new VulkanBuffer(vkBn->getDynamicMemoryPool(), false, bufferSizeOutput, nullptr, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT));
|
||
}
|
||
|
||
mSource.convert->encodeTensorToBuffer(input, mSource.buffer->buffer(), mSource.buffer->size(), 0, VulkanImageConverter::getTensorLinearFormat(input), cmdBuffer);
|
||
|
||
mDescriptorSet->writeBuffer(mOutput.buffer->buffer(), 0, mOutput.buffer->size());
|
||
mDescriptorSet->writeBuffer(mSource.buffer->buffer(), 1, mSource.buffer->size());
|
||
mDescriptorSet->writeBuffer(mGpuArgMaxParam->buffer(), 2, mGpuArgMaxParam->size());
|
||
|
||
cmdBuffer->barrierSource(mSource.buffer->buffer(), 0, mSource.buffer->size());
|
||
|
||
mArgMaxPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
|
||
vkCmdDispatch(cmdBuffer->get(), UP_DIV(total, 256), 1, 1);
|
||
|
||
cmdBuffer->barrierSource(mOutput.buffer->buffer(), 0, mOutput.buffer->size());
|
||
mOutput.convert->encodeBufferToTensor(mOutput.buffer->buffer(), output, mOutput.buffer->size(), 0, VulkanImageConverter::getTensorLinearFormat(output), cmdBuffer);
|
||
{
|
||
mSource.buffer->release();
|
||
mOutput.buffer->release();
|
||
}
|
||
return NO_ERROR;
|
||
}
|
||
|
||
class VulkanArgMaxCreator : 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 {
|
||
if (TensorUtils::getDescribe(inputs[0])->dimensionFormat == MNN_DATA_FORMAT_NC4HW4) {
|
||
// Don't support legency version
|
||
return nullptr;
|
||
}
|
||
return new VulkanArgMax(op, backend);
|
||
}
|
||
};
|
||
|
||
static bool gResistor = []() {
|
||
VulkanBackend::addCreator(OpType_ArgMax, new VulkanArgMaxCreator);
|
||
VulkanBackend::addCreator(OpType_ArgMin, new VulkanArgMaxCreator);
|
||
return true;
|
||
}();
|
||
|
||
}
|