Merge pull request #28104 from nklskyoy:rms-norm

RMSNorm: reference cpu impl #28104

https://onnx.ai/onnx/operators/onnx__RMSNormalization.html

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [ ] The PR is proposed to the proper branch
- [ ] There is a reference to the original bug report and related work
- [ ] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [ ] The feature is well documented and sample code can be built with the project CMake

modules/dnn/include/opencv2/dnn/all_layers.hpp

@@ -1544,6 +1544,13 @@ CV__DNN_INLINE_NS_BEGIN
         static Ptr<TopK2Layer> create(const LayerParams &params);
     };
 
+
+    class CV_EXPORTS RMSNormLayer : public Layer
+    {
+    public:
+        static Ptr<RMSNormLayer> create(const LayerParams& params);
+    };
+
 //! @}
 //! @}
 CV__DNN_INLINE_NS_END

modules/dnn/src/init.cpp

@@ -195,6 +195,7 @@ void initializeLayerFactory()
     CV_DNN_REGISTER_LAYER_CLASS(Gather2,        Gather2Layer);
     CV_DNN_REGISTER_LAYER_CLASS(GatherElements, GatherElementsLayer);
     CV_DNN_REGISTER_LAYER_CLASS(LayerNormalization, LayerNormLayer);
+    CV_DNN_REGISTER_LAYER_CLASS(RMSNormalization, RMSNormLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Expand,         ExpandLayer);
     CV_DNN_REGISTER_LAYER_CLASS(InstanceNormalization, InstanceNormLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Attention,      AttentionLayer);

modules/dnn/src/layers/cpu_kernels/fast_norm.cpp

@@ -42,7 +42,7 @@ void fastNorm(const Mat &input, Mat &output, float epsilon, size_t normalized_ax
     parallel_for_(Range(0, loops), fn, nstripes);
 }
 
-void fastNorm(const Mat &input, const Mat &scale, Mat &output, float epsilon, size_t normalized_axis) {
+void fastNorm(const Mat &input, const Mat &scale, Mat &output, float epsilon, size_t normalized_axis, bool recenter) {
     const auto input_shape = shape(input);
     CV_CheckLT(normalized_axis, input_shape.size(), "fastNorm: axis out of range");
 
@@ -61,7 +61,8 @@ void fastNorm(const Mat &input, const Mat &scale, Mat &output, float epsilon, si
             float mean = 0.f, mean_square = 0.f;
             for (int j = 0; j < norm_size; j++) {
                 float v = x[j];
-                mean += v;
+                if (recenter)
+                    mean += v;
                 mean_square += v * v;
             }
 

modules/dnn/src/layers/cpu_kernels/fast_norm.hpp

@@ -13,7 +13,7 @@ namespace cv { namespace dnn {
 void fastNorm(const Mat &input, Mat &output, float epsilon, size_t normalized_axis = 0, bool normalize_variance = true);
 
 // Normalization speedup by multi-threading with absent bias. Mainly for LayerNormalization.
-void fastNorm(const Mat &input, const Mat &scale, Mat &output, float epsilon, size_t normalized_axis = 0);
+void fastNorm(const Mat &input, const Mat &scale, Mat &output, float epsilon, size_t normalized_axis = 0, bool recenter=true);
 
 // Normalization speedup by multi-threading with scale and bias. Mainly for LayerNormalization.
 void fastNorm(const Mat &input, const Mat &scale, const Mat &bias, Mat &output, float epsilon, size_t normalized_axis = 0);

modules/dnn/src/layers/rms_norm.cpp

@@ -0,0 +1,94 @@
+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+#include "../precomp.hpp"
+#include "layers_common.hpp"
+#include "cpu_kernels/fast_norm.hpp"
+
+namespace cv { namespace dnn {
+
+// https://github.com/onnx/onnx/blob/main/docs/Operators.md#LayerNormalization
+class RMSNormLayerImpl CV_FINAL : public RMSNormLayer
+{
+#ifdef HAVE_OPENCL
+    UMat weight_umat, bias_umat;
+#endif
+
+public:
+    int axis0;
+
+    RMSNormLayerImpl(const LayerParams& params)
+    {
+        setParamsFrom(params);
+
+        // standard attr
+        axis = params.get<int>("axis", -1);
+        epsilon = params.get<float>("epsilon", 1e-5);
+    }
+
+    virtual bool supportBackend(int backendId) CV_OVERRIDE
+    {
+        return backendId == DNN_BACKEND_OPENCV;
+    }
+
+    virtual bool getMemoryShapes(const std::vector<MatShape> &inputs,
+                                 const int requiredOutputs,
+                                 std::vector<MatShape> &outputs,
+                                 std::vector<MatShape> &internals) const CV_OVERRIDE
+    {
+        const int n_inputs = inputs.size();
+        CV_Check(n_inputs, n_inputs == 2, "RMSNorm: require two (x, scale) inputs");
+
+        auto x_shape = inputs[0];
+        auto scale_shape = inputs[1];
+        const int normalized_axis = normalize_axis(axis, static_cast<int>(x_shape.size()));
+        const int x_ndims = static_cast<int>(x_shape.size());
+
+        CV_CheckTrue(scale_shape.size() == x_ndims - normalized_axis,
+            "RMSNorm: scale shape should match normalized shape");
+
+        for (int i = 0; i + normalized_axis < x_ndims; ++i)
+        {
+            CV_CheckTrue(
+                x_shape[i + normalized_axis] == scale_shape[i],
+                "RMSNorm: scale shape should match normalized shape");
+        }
+
+        outputs.assign(1, inputs[0]);
+        return true;
+    }
+
+    void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE
+    {
+        CV_TRACE_FUNCTION();
+        CV_TRACE_ARG_VALUE(name, "name", name.c_str());
+
+        if (inputs_arr.depth() == CV_16F)
+        {
+            forward_fallback(inputs_arr, outputs_arr, internals_arr);
+            return;
+        }
+
+        std::vector<Mat> inputs, outputs;
+        inputs_arr.getMatVector(inputs);
+        outputs_arr.getMatVector(outputs);
+
+        const auto &input = inputs[0];
+        const auto &scale = inputs[1];
+        auto &output = outputs[0];
+
+        axis = normalize_axis(axis, input.dims);
+        fastNorm(input, scale, output, epsilon, static_cast<size_t>(axis), false);
+    }
+private:
+    int axis;
+    float epsilon;
+};
+
+Ptr<RMSNormLayer> RMSNormLayer::create(const LayerParams& params)
+{
+    return makePtr<RMSNormLayerImpl>(params);
+}
+
+}} // cv::dnn

modules/dnn/src/onnx/onnx_importer2.cpp

@@ -247,7 +247,7 @@ protected:
     void parseBitShift             (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseBitwise              (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseBitwiseNot           (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
-    void parseRotaryEmbedding      (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
+    void parseRMSNormalization     (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);    void parseRotaryEmbedding      (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     // Domain: com.microsoft
     // URL: https://github.com/microsoft/onnxruntime/blob/master/docs/ContribOperators.md
     void parseAttention            (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
@@ -2014,6 +2014,11 @@ void ONNXImporter2::parseQuantizeLinear(LayerParams& layerParams, const opencv_o
     addLayer(layerParams, node_proto);
 }
 
+void ONNXImporter2::parseRMSNormalization(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
+{
+    addLayer(layerParams, node_proto);
+}
+
 // BUG: https://github.com/opencv/opencv/issues/26310
 /*void ONNXImporter2::parseQConv(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto_)
 {

modules/dnn/test/test_onnx_conformance_layer_filter__openvino.inl.hpp

@@ -2764,6 +2764,44 @@ CASE(test_rotary_embedding_with_rotary_dim)
     SKIP;
 CASE(test_rotary_embedding_with_rotary_dim_expanded)
     SKIP;
+CASE(test_rms_normalization_2d_axis0)
+    SKIP;
+CASE(test_rms_normalization_2d_axis1)
+    SKIP;
+CASE(test_rms_normalization_2d_axis_negative_1)
+    SKIP;
+CASE(test_rms_normalization_2d_axis_negative_2)
+    SKIP;
+CASE(test_rms_normalization_3d_axis0_epsilon)
+    SKIP;
+CASE(test_rms_normalization_3d_axis1_epsilon)
+    SKIP;
+CASE(test_rms_normalization_3d_axis2_epsilon)
+    SKIP;
+CASE(test_rms_normalization_3d_axis_negative_1_epsilon)
+    SKIP;
+CASE(test_rms_normalization_3d_axis_negative_2_epsilon)
+    SKIP;
+CASE(test_rms_normalization_3d_axis_negative_3_epsilon)
+    SKIP;
+CASE(test_rms_normalization_4d_axis0)
+    SKIP;
+CASE(test_rms_normalization_4d_axis1)
+    SKIP;
+CASE(test_rms_normalization_4d_axis2)
+    SKIP;
+CASE(test_rms_normalization_4d_axis3)
+    SKIP;
+CASE(test_rms_normalization_4d_axis_negative_1)
+    SKIP;
+CASE(test_rms_normalization_4d_axis_negative_2)
+    SKIP;
+CASE(test_rms_normalization_4d_axis_negative_3)
+    SKIP;
+CASE(test_rms_normalization_4d_axis_negative_4)
+    SKIP;
+CASE(test_rms_normalization_default_axis)
+    SKIP;
 #if SKIP_SET_1
     SKIP;
 #endif

modules/dnn/test/test_onnx_conformance_layer_parser_denylist.inl.hpp

@@ -534,43 +534,24 @@
 "test_resize_upsample_sizes_nearest_not_smaller",
 "test_reversesequence_batch",  // Issue:: Parser: Can't create layer "onnx_node_output_0!y" of type "ReverseSequence" in function 'getLayerInstance'
 "test_reversesequence_time",  // ---- same as above ---
-"test_rms_normalization_2d_axis0",
 "test_rms_normalization_2d_axis0_expanded",
-"test_rms_normalization_2d_axis1",
 "test_rms_normalization_2d_axis1_expanded",
-"test_rms_normalization_2d_axis_negative_1",
 "test_rms_normalization_2d_axis_negative_1_expanded",
-"test_rms_normalization_2d_axis_negative_2",
 "test_rms_normalization_2d_axis_negative_2_expanded",
-"test_rms_normalization_3d_axis0_epsilon",
 "test_rms_normalization_3d_axis0_epsilon_expanded",
-"test_rms_normalization_3d_axis1_epsilon",
 "test_rms_normalization_3d_axis1_epsilon_expanded",
-"test_rms_normalization_3d_axis2_epsilon",
 "test_rms_normalization_3d_axis2_epsilon_expanded",
-"test_rms_normalization_3d_axis_negative_1_epsilon",
 "test_rms_normalization_3d_axis_negative_1_epsilon_expanded",
-"test_rms_normalization_3d_axis_negative_2_epsilon",
 "test_rms_normalization_3d_axis_negative_2_epsilon_expanded",
-"test_rms_normalization_3d_axis_negative_3_epsilon",
 "test_rms_normalization_3d_axis_negative_3_epsilon_expanded",
-"test_rms_normalization_4d_axis0",
 "test_rms_normalization_4d_axis0_expanded",
-"test_rms_normalization_4d_axis1",
 "test_rms_normalization_4d_axis1_expanded",
-"test_rms_normalization_4d_axis2",
 "test_rms_normalization_4d_axis2_expanded",
-"test_rms_normalization_4d_axis3",
 "test_rms_normalization_4d_axis3_expanded",
-"test_rms_normalization_4d_axis_negative_1",
 "test_rms_normalization_4d_axis_negative_1_expanded",
-"test_rms_normalization_4d_axis_negative_2",
 "test_rms_normalization_4d_axis_negative_2_expanded",
-"test_rms_normalization_4d_axis_negative_3",
 "test_rms_normalization_4d_axis_negative_3_expanded",
-"test_rms_normalization_4d_axis_negative_4",
 "test_rms_normalization_4d_axis_negative_4_expanded",
-"test_rms_normalization_default_axis",
 "test_rms_normalization_default_axis_expanded",
 "test_rnn_seq_length",  // Issue:: Parser:  Can't create layer "onnx_node_output_1!Y_h" of type "RNN" in function 'getLayerInstance'
 "test_roialign_aligned_false",  // Issue:: Parser: Layer does not exist (RoiAlign)