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,114 @@
+//
+//  VulkanTopKV2.cpp
+//  MNN
+//
+//  Vulkan buffer-mode implementation of TopKV2.
+//
+
+#include "VulkanTopKV2.hpp"
+#include "core/Macro.h"
+#include "core/TensorUtils.hpp"
+
+namespace MNN {
+
+struct TopKV2ConstBuffer {
+    int rowSize;
+    int k;
+    int numRows;
+    int pad;
+};
+
+VulkanTopKV2::VulkanTopKV2(const Op* op, Backend* bn, int k, Tensor* input) : VulkanBasicExecution(bn) {
+    MNN_ASSERT(k > 0);
+    MNN_ASSERT(input != nullptr);
+    auto vkBn = (VulkanBackend*)backend();
+    mK = k;
+    mLargest = true;
+    auto param = op->main_as_TopKV2();
+    if (nullptr != param) {
+        mLargest = param->largest();
+    }
+
+    mConstBuffer = vkBn->allocUniform();
+    std::vector<VkDescriptorType> types{
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
+    };
+
+    std::string pKey = "glsl_topkv2_";
+    if (mLargest) {
+        pKey += "SORT_DESC_";
+    }
+    if (input->getType().code == halide_type_float && vkBn->useFP16()) {
+        pKey += "FP16_";
+    }
+    pKey += "comp";
+
+    mPipeline = vkBn->getPipeline(pKey, types);
+    mDescriptorSet.reset(mPipeline->createSet());
+}
+
+VulkanTopKV2::~VulkanTopKV2() {
+    auto vkBn = (VulkanBackend*)backend();
+    vkBn->recycleUniform(mConstBuffer);
+}
+
+ErrorCode VulkanTopKV2::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
+                                  const VulkanCommandPool::Buffer* cmdBuffer) {
+    auto input = inputs[0];
+    auto outputValue = outputs[0];
+    auto outputIndex = outputs[1];
+
+    const int rowSize = input->length(input->dimensions() - 1);
+    if (rowSize <= 0) {
+        return NO_ERROR;
+    }
+    const int numRows = input->elementSize() / rowSize;
+    const int k = mK;
+
+    auto vkBn = static_cast<VulkanBackend*>(backend());
+
+    // Set GPU params
+    auto topkParam = reinterpret_cast<TopKV2ConstBuffer*>(mConstBuffer->map());
+    topkParam->rowSize = rowSize;
+    topkParam->k = k;
+    topkParam->numRows = numRows;
+    topkParam->pad = 0;
+    mConstBuffer->unmap();
+
+    // Bind buffers
+    mDescriptorSet->writeBuffer(vkBn->getBuffer(outputValue), 0);
+    mDescriptorSet->writeBuffer(vkBn->getBuffer(outputIndex), 1);
+    mDescriptorSet->writeBuffer(vkBn->getBuffer(input), 2);
+    mDescriptorSet->writeBuffer(mConstBuffer->buffer(), 3, mConstBuffer->size());
+
+    // Dispatch: x=1 (workgroup handles 128 threads internally), y=numRows, z=1
+    mPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
+    vkCmdDispatch(cmdBuffer->get(), 1, numRows, 1);
+
+    return NO_ERROR;
+}
+
+class VulkanTopKV2Creator : 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 (inputs.size() < 2 || outputs.size() != 2) {
+            return nullptr;
+        }
+        if (TensorUtils::getDescribe(inputs[0])->dimensionFormat == MNN_DATA_FORMAT_NC4HW4) {
+            return nullptr;
+        }
+        const int k = inputs[1]->host<int32_t>()[0];
+        return new VulkanTopKV2(op, backend, k, inputs[0]);
+    }
+};
+
+static bool gResistor = []() {
+    VulkanBackend::addCreator(OpType_TopKV2, new VulkanTopKV2Creator);
+    return true;
+}();
+
+} // namespace MNN