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,138 @@
+//
+//  GeometryLayerNorm.cpp
+//  MNN
+//
+//  Created by MNN on 2020/06/09.
+//  Copyright © 2018, Alibaba Group Holding Limited
+//
+
+#include "geometry/GeometryComputer.hpp"
+#include "geometry/GeometryComputerUtils.hpp"
+#include "core/OpCommonUtils.hpp"
+namespace MNN {
+
+class GeometryLayerNorm : public GeometryComputer {
+public:
+    virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
+                           Context& context, CommandBuffer& res) const override {
+        /* Target: Ensure reduce dimensions must be a sequence subset [-rank,...,rank-1] */
+        auto layernorm          = op->main_as_LayerNorm();
+        if (!layernorm->axis() || op->defaultDimentionFormat() == MNN_DATA_FORMAT_NC4HW4) {
+            std::shared_ptr<Command> cmdP(new Command);
+            auto& cmd = *cmdP;
+            cmd.op      = op;
+            cmd.inputs  = inputs;
+            cmd.outputs = std::move(outputs);
+            res.command.emplace_back(std::move(cmdP));
+            return true;
+        }
+        MNN_ASSERT(1 == outputs.size());
+        MNN_ASSERT(1 == inputs.size());
+        auto reduceDims = layernorm->axis()->data();
+        int reduceDimensionCount = layernorm->axis()->size();
+        auto inputShape = inputs[0]->shape();
+        auto outputShape = outputs[0]->shape();
+        int rank = static_cast<int32_t>(inputShape.size());
+
+        // Case1: Do not need permute
+        bool needPermute = true;
+        if (reduceDims[0] < 0 && reduceDims[reduceDimensionCount - 1] == -1) { // [-r,-r+1...]
+            needPermute = false;
+        }
+        if (reduceDims[reduceDimensionCount - 1] > 0 && reduceDims[reduceDimensionCount - 1] == rank - 1 ) { // [...,r-2,r-1]
+            needPermute = false;
+        }
+        if (reduceDims[0] == 0 && rank == 1) { // reduce dim:[0], input dimensions=1
+            needPermute = false;
+        }
+        std::vector<int> lastdims(reduceDimensionCount);
+        for (int i = 0; i < reduceDimensionCount; ++i) {
+            lastdims[i] = (reduceDims[i] + rank) % rank;
+        }
+        if (false == needPermute) {
+            std::shared_ptr<Command> cmdP(new Command);
+            auto& cmd = *cmdP;
+            cmd.op      = op;
+            cmd.inputs  = {inputs[0]};
+            cmd.outputs = std::move(outputs);
+            res.command.emplace_back(std::move(cmdP));
+            return true;
+        }
+        // Case2 : Need permute
+        int oldorder[MNN_MAX_TENSOR_DIM];
+        int neworder[MNN_MAX_TENSOR_DIM];
+
+        {
+            int di = 0;
+            int idx = 0;
+            while (di < rank) {
+                if (di < lastdims[0] || di > lastdims[reduceDimensionCount - 1]) {
+                    neworder[idx++] = di;
+                }
+                di++;
+            }
+            for (int i = 0; i < reduceDimensionCount; ++i) {
+                neworder[idx++] = lastdims[i];
+            }
+        }
+        {
+            int idx = 0;
+            for (int i = 0; i < rank; ++i) {
+                int j = 0;
+                while (i != neworder[j]) {
+                    ++j;
+                }
+                oldorder[idx++] = j;
+            }
+        }
+        std::vector<int> newshape;
+        for (int i = 0; i < rank; ++i) {
+            newshape.emplace_back(inputShape[neworder[i]]);
+        }
+        
+        std::shared_ptr<Tensor> outputTensorPermute(Tensor::createDevice(newshape, inputs[0]->getType(), inputs[0]->getDimensionType()));
+        res.extras.emplace_back(outputTensorPermute);
+        GeometryComputer::ComputePermuteRegion(inputs[0], outputTensorPermute.get(), neworder, rank);
+        
+        // Create LayerNorm command
+        auto currentInput = outputTensorPermute.get();
+        float epislon = layernorm->epsilon();
+        bool  useRMS  = layernorm->useRMSNorm();
+        std::vector<int64_t> externalData;
+        if (layernorm->external()) {
+            externalData.resize(3);
+            externalData = {layernorm->external()->data()[0], layernorm->external()->data()[1], layernorm->external()->data()[2]};
+        }
+        std::vector<float> gamma, beta;
+        if (layernorm->gamma() && layernorm->beta()) {
+            int   gammaSize = layernorm->gamma()->size();
+            gamma.resize(gammaSize);
+            beta.resize(gammaSize);
+            int   group     = layernorm->group();
+            ::memcpy(gamma.data(), layernorm->gamma()->data(), gammaSize * sizeof(float));
+            ::memcpy(beta.data(), layernorm->beta()->data(), gammaSize * sizeof(float));
+        }
+        std::shared_ptr<Tensor> inputTensorLayernorm(Tensor::createDevice(newshape, inputs[0]->getType(), inputs[0]->getDimensionType()));
+        auto des = TensorUtils::getDescribe(inputTensorLayernorm.get());
+        des->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
+        des->regions    = {TensorUtils::makeFullSlice(currentInput)};
+        res.extras.emplace_back(inputTensorLayernorm);
+        std::shared_ptr<Tensor> outputTensorLayerNorm(Tensor::createDevice(newshape, inputs[0]->getType(), inputs[0]->getDimensionType()));
+        res.extras.emplace_back(outputTensorLayerNorm);
+        {
+            auto cmd = GeometryComputerUtils::makeLayerNorm(inputTensorLayernorm.get(), outputTensorLayerNorm.get(), lastdims, epislon, gamma, beta, externalData, 1, useRMS);
+            res.command.emplace_back(std::move(cmd));
+        }
+        
+        GeometryComputer::ComputePermuteRegion(outputTensorLayerNorm.get(), outputs[0], oldorder, rank);
+        return true;
+    }
+};
+static void _create() {
+    std::shared_ptr<GeometryComputer> comp(new GeometryLayerNorm);
+    GeometryComputer::registerGeometryComputer(comp, {OpType_LayerNorm});
+}
+
+REGISTER_GEOMETRY(GeometryLayerNorm, _create);
+
+} // namespace MNN