chore: import upstream snapshot with attribution

2026-07-13 13:33:03 +08:00
commit 5b57521aa1
8226 changed files with 3425766 additions and 0 deletions
@@ -0,0 +1,129 @@
+//
+//  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;
+}();
+
+}