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 @@
+//
+//  VulkanSoftmax.cpp
+//  MNN
+//
+//  Created by MNN on 2019/01/31.
+//  Copyright © 2018, Alibaba Group Holding Limited
+//
+
+#include "VulkanSoftmax.hpp"
+#include "core/Macro.h"
+#include "core/TensorUtils.hpp"
+
+namespace MNN {
+
+struct ConstBuffer {
+    int w;
+    int h;
+    int c;
+};
+
+VulkanSoftmax::VulkanSoftmax(const Op* op, Backend* bn) : VulkanBasicExecution(bn) {
+    const auto softmaxParam = op->main_as_Axis();
+    mAxis                   = softmaxParam->axis();
+    auto vkBn = (VulkanBackend*)backend();
+    mConstBuffer = std::make_shared<VulkanBuffer>(vkBn->getMemoryPool(), false, sizeof(ConstBuffer), nullptr,
+                                                  VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT);
+    std::vector<VkDescriptorType> types{
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER
+    };
+    // Modified logic to support FP16
+    std::string shaderName = "glsl_softmaxHeight_NHWC_";
+    if (vkBn->useFP16()) {
+        shaderName += "FP16_";
+    }
+    shaderName += "comp";
+    
+    mSoftmaxPipeline =
+        vkBn->getPipeline(shaderName, types);
+    mDescriptorSet.reset(mSoftmaxPipeline->createSet());
+}
+
+VulkanSoftmax::~VulkanSoftmax() {
+}
+
+ErrorCode VulkanSoftmax::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
+                                  const VulkanCommandPool::Buffer* cmdBuffer) {
+    auto input  = inputs[0];
+    auto output = outputs[0];
+    auto extra = (VulkanBackend*)backend();
+    auto inputFormat = TensorUtils::getDescribe(input)->dimensionFormat;
+    auto axis = mAxis;
+    if (axis < 0) {
+        axis = input->dimensions() + axis;
+    }
+    if (inputFormat == MNN_DATA_FORMAT_NC4HW4) {
+        mConvert.reset(new ConvertComponent);
+        mConvert->mTempInputTensor.reset(new Tensor(input, Tensor::CAFFE, false));
+        mConvert->mTempOutputTensor.reset(new Tensor(output, Tensor::CAFFE, false));
+        auto res = backend()->onAcquireBuffer(mConvert->mTempInputTensor.get(), Backend::DYNAMIC);
+        if (!res) {
+            return OUT_OF_MEMORY;
+        }
+        res = backend()->onAcquireBuffer(mConvert->mTempOutputTensor.get(), Backend::DYNAMIC);
+        if (!res) {
+            return OUT_OF_MEMORY;
+        }
+        mConvert->mInputConvert = VulkanRaster::create(mConvert->mTempInputTensor.get(), backend());
+        mConvert->mOutputConvert = VulkanRaster::create(mConvert->mTempOutputTensor.get(), backend());
+        TensorUtils::getDescribe(mConvert->mTempInputTensor.get())->regions = {TensorUtils::makeFullSlice(inputs[0])};
+        TensorUtils::getDescribe(outputs[0])->regions = {TensorUtils::makeFullSlice(mConvert->mTempOutputTensor.get())};
+        input = mConvert->mTempInputTensor.get();
+        output = mConvert->mTempOutputTensor.get();
+        
+        mConvert->mInputConvert.exe->onEncode({}, {mConvert->mTempInputTensor.get()}, cmdBuffer);
+        auto inputBuffer = extra->getTensorBuffer(input);
+        auto inputBufferSize = extra->getTensorSize(input);
+        cmdBuffer->barrierSource(inputBuffer.first->buffer(), inputBuffer.second, inputBufferSize);
+    }
+    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<output->dimensions(); ++i) {
+        inside *= input->length(i);
+    }
+    // gpu param
+    {
+        auto softmax = reinterpret_cast<ConstBuffer*>(mConstBuffer->map());
+        ::memset(softmax, 0, sizeof(ConstBuffer));
+        softmax->w = inside;
+        softmax->h = mid;
+        softmax->c = outside;
+        mConstBuffer->unmap();
+    }
+    auto outputBuffer = extra->getTensorBuffer(output);
+    auto inputBuffer = extra->getTensorBuffer(input);
+    mDescriptorSet->writeBuffer(outputBuffer.first->buffer(), 0, extra->getTensorSize(output), outputBuffer.second);
+    mDescriptorSet->writeBuffer(inputBuffer.first->buffer(), 1, extra->getTensorSize(input), inputBuffer.second);
+    mDescriptorSet->writeBuffer(mConstBuffer->buffer(), 2, mConstBuffer->size());
+    mSoftmaxPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
+    vkCmdDispatch(cmdBuffer->get(), UP_DIV(outside, 8), UP_DIV(inside, 8), 1);
+    if (inputFormat == MNN_DATA_FORMAT_NC4HW4) {
+        cmdBuffer->barrierSource(outputBuffer.first->buffer(), outputBuffer.second, extra->getTensorSize(output));
+        mConvert->mOutputConvert.exe->onEncode({}, {outputs[0]}, cmdBuffer);
+        backend()->onReleaseBuffer(mConvert->mTempInputTensor.get(), Backend::DYNAMIC);
+        backend()->onReleaseBuffer(mConvert->mTempOutputTensor.get(), Backend::DYNAMIC);
+    }
+
+    return NO_ERROR;
+}
+
+class VulkanSoftmaxCreator : 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 {
+        return new VulkanSoftmax(op, backend);
+    }
+};
+
+static bool gResistor = []() {
+    VulkanBackend::addCreator(OpType_Softmax, new VulkanSoftmaxCreator);
+    return true;
+}();
+
+} // namespace MNN