text: // Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //     http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  
  #include "lite/kernels/arm/group_norm_compute.h"
  #include "lite/backends/arm/math/funcs.h"
  #include "lite/core/op_registry.h"
  #include "lite/core/type_system.h"
  
  namespace paddle {
  namespace lite {
  namespace kernels {
  namespace arm {
  
  void GroupNormCompute::PrepareForRun() {}
  
  void GroupNormCompute::Run() {
    auto& param = this->Param<param_t>();
    const float* in = param.x->data<float>();
    const float* scale = param.scale->data<float>();
    const float* bias = param.bias->data<float>();
    float* out = param.out->mutable_data<float>();
    float* saved_mean = param.saved_mean->mutable_data<float>();
    float* saved_variance = param.saved_variance->mutable_data<float>();
    float epsilon = param.epsilon;
    int groups = param.groups;
    int channels = param.channels;
    int n = param.x->dims()[0];
    int c = param.x->dims()[1];
    int ch_per_group = channels / groups;
    int height = param.x->dims()[2];
    int width = param.x->dims()[3];
    int spatial_size = ch_per_group * height * width;
    int ngroup = n * groups;
    int cnt = spatial_size >> 4;
    int remain = spatial_size % 16;
  // compute saved_mean and saved_variance
  #pragma omp parallel for
    for (int n = 0; n < ngroup; ++n) {
      const float* in_p = in + n * spatial_size;
      float sum_spatial = 0.f;
      float summ_spatial = 0.f;
      float32x4_t sum0 = vdupq_n_f32(0.f);
      float32x4_t sum1 = vdupq_n_f32(0.f);
      float32x4_t sum2 = vdupq_n_f32(0.f);
      float32x4_t sum3 = vdupq_n_f32(0.f);
      float32x4_t summ0 = vdupq_n_f32(0.f);
      float32x4_t summ1 = vdupq_n_f32(0.f);
      float32x4_t summ2 = vdupq_n_f32(0.f);
      float32x4_t summ3 = vdupq_n_f32(0.f);
      for (int i = 0; i < cnt; i++) {
        float32x4_t in0 = vld1q_f32(in_p);
        float32x4_t in1 = vld1q_f32(in_p + 4);
        float32x4_t in2 = vld1q_f32(in_p + 8);
        float32x4_t in3 = vld1q_f32(in_p + 12);
        sum0 = vaddq_f32(sum0, in0);
        summ0 = vmlaq_f32(summ0, in0, in0);
        sum1 = vaddq_f32(sum1, in1);
        summ1 = vmlaq_f32(summ1, in1, in1);
        sum2 = vaddq_f32(sum2, in2);
        summ2 = vmlaq_f32(summ2, in2, in2);
        sum3 = vaddq_f32(sum3, in3);
        summ3 = vmlaq_f32(summ3, in3, in3);
        in_p += 16;
      }
      for (int i = 0; i < remain - 3; i += 4) {
        float32x4_t in0 = vld1q_f32(in_p);
        sum1 = vaddq_f32(sum1, in0);
        summ1 = vmlaq_f32(summ1, in0, in0);
        in_p += 4;
      }
      float sum = 0.0;
      float summ = 0.0;
      sum0 = vaddq_f32(sum0, sum1);
      sum2 = vaddq_f32(sum2, sum3);
      summ0 = vaddq_f32(summ0, summ1);
      summ2 = vaddq_f32(summ2, summ3);
      for (int i = 0; i < remain % 4; i++) {
        sum += *in_p;
        summ += (*in_p) * (*in_p);
        in_p++;
      }
      sum0 = vaddq_f32(sum0, sum2);
      summ0 = vaddq_f32(summ0, summ2);
      float32x2_t sum_low = vpadd_f32(vget_low_f32(sum0), vget_high_f32(sum0));
      float32x2_t sum_high = vpadd_f32(vget_low_f32(summ0), vget_high_f32(summ0));
      float32x2_t sum_mix = vpadd_f32(sum_low, sum_high);
      sum += vget_lane_f32(sum_mix, 0);
      summ += vget_lane_f32(sum_mix, 1);
      float mean = sum / spatial_size;
      // float variance = summ / spatial_size - mean * mean;
      // the flolowing code has higher precision than above comment code
      float variance = (summ - mean * mean * spatial_size) / spatial_size;
      float std = 1.f / sqrtf(variance + epsilon);
      saved_mean[n] = mean;
      saved_variance[n] = std;
    }
    int in_size = height * width;
    cnt = in_size >> 4;
    remain = in_size % 16;
  // compute Group_norm result: out = scale * (in - mean) / std + bias
  #pragma omp parallel for
    for (int i = 0; i < ngroup; ++i) {
      const float* in_p = in + i * spatial_size;
      float* out_p = out + i * spatial_size;
      int numc = i % groups;
      numc *= ch_per_group;
      for (int c = 0; c < ch_per_group; c++) {
        int chin = numc + c;
        const float sstd_val = scale[chin] * saved_variance[i];
        const float bias_val = bias[chin];
        const float mean_val = saved_mean[i];
        const float32x4_t vsstd = vdupq_n_f32(sstd_val);
        const float32x4_t vbias = vdupq_n_f32(bias_val);
        const float32x4_t vmean = vdupq_n_f32(mean_val);
        for (int k = 0; k < cnt; k++) {
          float32x4_t in0 = vld1q_f32(in_p);
          float32x4_t in1 = vld1q_f32(in_p + 4);
          float32x4_t in2 = vld1q_f32(in_p + 8);
          float32x4_t in3 = vld1q_f32(in_p + 12);
          float32x4_t submean0 = vsubq_f32(in0, vmean);
          float32x4_t submean1 = vsubq_f32(in1, vmean);
          float32x4_t submean2 = vsubq_f32(in2, vmean);
          float32x4_t submean3 = vsubq_f32(in3, vmean);
          float32x4_t out0 = vmlaq_f32(vbias, submean0, vsstd);
          float32x4_t out1 = vmlaq_f32(vbias, submean1, vsstd);
          float32x4_t out2 = vmlaq_f32(vbias, submean2, vsstd);
          float32x4_t out3 = vmlaq_f32(vbias, submean3, vsstd);
          vst1q_f32(out_p, out0);
          vst1q_f32(out_p + 4, out1);
          vst1q_f32(out_p + 8, out2);
          vst1q_f32(out_p + 12, out3);
          in_p += 16;
          out_p += 16;
        }
        for (int k = 0; k < remain - 3; k += 4) {
          float32x4_t in0 = vld1q_f32(in_p);
          in_p += 4;
          float32x4_t submean0 = vsubq_f32(in0, vmean);
          float32x4_t out0 = vmlaq_f32(vbias, submean0, vsstd);
          vst1q_f32(out_p, out0);
          out_p += 4;
        }
        for (int k = 0; k < remain % 4; k++) {
          *out_p = (*in_p - mean_val) * sstd_val + bias_val;
          in_p++;
          out_p++;
        }
      }
    }
  }
  
  }  // namespace arm
  }  // namespace kernels
  }  // namespace lite
  }  // namespace paddle
  
  REGISTER_LITE_KERNEL(group_norm,
                       kARM,
                       kFloat,
                       kNCHW,
                       paddle::lite::kernels::arm::GroupNormCompute,
                       def)
      .BindInput("X", {LiteType::GetTensorTy(TARGET(kARM))})
      .BindInput("Scale", {LiteType::GetTensorTy(TARGET(kARM))})
      .BindInput("Bias", {LiteType::GetTensorTy(TARGET(kARM))})
      .BindOutput("Y", {LiteType::GetTensorTy(TARGET(kARM))})
      .BindOutput("SavedMean", {LiteType::GetTensorTy(TARGET(kARM))})
      .BindOutput("SavedVariance", {LiteType::GetTensorTy(TARGET(kARM))})
      .Finalize();
  