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,105 @@
+//
+//  VulkanOneHot.cpp
+//  MNN
+//
+//  Created by MNN on 2020/07/20.
+//  Copyright © 2018, Alibaba Group Holding Limited
+//
+
+#include "VulkanOneHot.hpp"
+#include "core/Macro.h"
+namespace MNN {
+struct ConstBuffer {
+    ivec4 inputSize;
+};
+
+VulkanOneHot::VulkanOneHot(const Op* op, Backend* bn, Tensor * tensor) : VulkanBasicExecution(bn) {
+    auto extra = static_cast<VulkanBackend*>(bn);
+    mAxis = op->main_as_OneHotParam()->axis();
+    mConstBuffer = extra->allocUniform();
+    std::string pKey = "glsl_onehot_";
+    if (tensor->getType().code == halide_type_float && extra->useFP16()) {
+        pKey += "FP16_";
+    }
+    pKey += "comp";
+    mPipeline = extra->getPipeline(pKey, {
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
+    });
+    mDescriptorSet.reset(mPipeline->createSet());
+}
+VulkanOneHot::~VulkanOneHot() {
+    auto extra = static_cast<VulkanBackend*>(backend());
+    extra->recycleUniform(mConstBuffer);
+}
+
+ErrorCode VulkanOneHot::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
+                               const VulkanCommandPool::Buffer* cmdBuffer) {
+    auto indices        = inputs[0];
+    auto depthTensor    = inputs[1];
+    auto onValueTensor  = inputs[2];
+    auto offValueTensor = inputs[3];
+
+    int axis = mAxis;
+    if (axis < 0) {
+        axis += outputs[0]->dimensions();
+    }
+    int outerSize = 1;
+    for (int i = 0; i < axis; ++i) {
+        outerSize *= indices->length(i);
+    }
+    const int depth       = outputs[0]->length(axis);
+    const int innerSize   = indices->elementSize() / outerSize;
+    const auto indicesPtr = indices->host<int>();
+    auto total = outputs[0]->elementSize();
+
+    auto dataType    = onValueTensor->getType();
+    auto offDataType = offValueTensor->getType();
+    auto vkBn = static_cast<VulkanBackend*>(backend());
+    MNN_ASSERT(dataType == offDataType);
+    // Set Const Buffer
+    {
+        auto pool = (ConstBuffer*)mConstBuffer->map();
+        ::memset(pool, 0, sizeof(ConstBuffer));
+        pool->inputSize[0]  = outerSize;
+        pool->inputSize[1]  = depth;
+        pool->inputSize[2]  = innerSize;
+        pool->inputSize[3]  = total;
+        mConstBuffer->unmap();
+    }
+
+    // Set Command Buffer
+    {
+        auto outputT = vkBn->getBuffer(outputs[0]);
+        auto inputT = vkBn->getBuffer(indices);
+        auto onT = vkBn->getBuffer(onValueTensor);
+        auto offT = vkBn->getBuffer(offValueTensor);
+        mDescriptorSet->writeBuffer(outputT, 0);
+        mDescriptorSet->writeBuffer(inputT, 1);
+        cmdBuffer->barrierSource(inputT);
+        mDescriptorSet->writeBuffer(onT, 2);
+        mDescriptorSet->writeBuffer(offT, 3);
+        mDescriptorSet->writeBuffer(mConstBuffer->buffer(), 4, mConstBuffer->size());
+        mPipeline->bind(cmdBuffer->get(), mDescriptorSet->get());
+        vkCmdDispatch(cmdBuffer->get(), UP_DIV(total, 256), 1, 1);
+    }
+    return NO_ERROR;
+}
+
+class VulkanOneHotCreator : 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 VulkanOneHot(op, backend, outputs[0]);
+    }
+};
+
+static bool gResistor = []() {
+    VulkanBackend::addCreator(OpType_OneHot, new VulkanOneHotCreator);
+    return true;
+}();
+
+} // namespace MNN