CMSIS-NN: Sum over 32 bits for Average pool

1. Sum is done over 32 bits instead of the existing
16 bits for Cortex-M CPUs with DSP extension.

2. Variable names are re-named to be uniform with
max pool.

Change-Id: If3a10eae51cd2679c8190fea503ce47a62439bf6
diff --git a/CMSIS/NN/Source/PoolingFunctions/arm_avgpool_s8.c b/CMSIS/NN/Source/PoolingFunctions/arm_avgpool_s8.c
index 5d26653..048d32b 100755
--- a/CMSIS/NN/Source/PoolingFunctions/arm_avgpool_s8.c
+++ b/CMSIS/NN/Source/PoolingFunctions/arm_avgpool_s8.c
@@ -21,8 +21,8 @@
  * Title:        arm_avgpool_s8.c
  * Description:  Pooling function implementations
  *
- * $Date:        May 29,2020
- * $Revision:    V.2.0.1
+ * $Date:        September 3,2020
+ * $Revision:    V.2.0.2
  *
  * Target Processor:  Cortex-M CPUs
  *
@@ -31,23 +31,24 @@
 #include "arm_math.h"
 #include "arm_nnfunctions.h"
 
+#if defined(ARM_MATH_DSP) && !defined(ARM_MATH_MVEI)
 
-#if defined (ARM_MATH_DSP) && !defined (ARM_MATH_MVEI)
-
-static void buffer_scale_back_q15_to_q7_and_clamp(q15_t *buffer, q7_t *target, uint16_t length,
-                                                  uint16_t count,const int act_min, const int act_max)
+static void scale_q31_to_q7_and_clamp(const q31_t *buffer,
+                                      q7_t *target,
+                                      int32_t length,
+                                      const int32_t count,
+                                      const int act_min,
+                                      const int act_max)
 {
-  int i;
-  int sum;
-
-  for (i = 0; i < length; i++)
+  const int half_count = count / 2;
+  for (int i = 0; i < length; i++)
   {
-    sum = buffer[i] > 0 ? (buffer[i] + count / 2) / count : (buffer[i] - count / 2) / count;
-
+    int32_t sum = buffer[i] > 0 ? (buffer[i] + half_count) : (buffer[i] - half_count);
+    sum = sum / count;
     sum = MAX(sum, act_min);
     sum = MIN(sum, act_max);
 
-    target[i] = (q7_t) (sum);
+    target[i] = (q7_t)sum;
   }
 }
 #endif
@@ -79,53 +80,52 @@
                           q7_t *dst)
 {
   (void)ctx;
-  const int32_t dim_src_height = input_dims->h;
-  const int32_t dim_src_width = input_dims->w;
-  const int32_t dim_dst_height = output_dims->h;
-  const int32_t dim_dst_width = output_dims->w;
-  const int32_t stride_height = pool_params->stride.h;
-  const int32_t stride_width = pool_params->stride.w;
-  const int32_t dim_kernel_height = filter_dims->h;
-  const int32_t dim_kernel_width = filter_dims->w;
-  const int32_t padding_height = pool_params->padding.h;
-  const int32_t padding_width = pool_params->padding.w;
+  const int32_t input_y = input_dims->h;
+  const int32_t input_x = input_dims->w;
+  const int32_t output_y = output_dims->h;
+  const int32_t output_x = output_dims->w;
+  const int32_t stride_y = pool_params->stride.h;
+  const int32_t stride_x = pool_params->stride.w;
+  const int32_t kernel_y = filter_dims->h;
+  const int32_t kernel_x = filter_dims->w;
+  const int32_t pad_y = pool_params->padding.h;
+  const int32_t pad_x = pool_params->padding.w;
   const int32_t act_min = pool_params->activation.min;
   const int32_t act_max = pool_params->activation.max;
   const int32_t ch_src = input_dims->c;
 
-  int32_t   i_x, i_y;
-  int32_t   k_x, k_y;
+  int32_t i_x, i_y;
+  int32_t k_x, k_y;
 
-  for (i_y = 0; i_y < dim_dst_height; i_y++)
+  for (i_y = 0; i_y < output_y; i_y++)
   {
-    for (i_x = 0; i_x < dim_dst_width; i_x++)
+    for (i_x = 0; i_x < output_x; i_x++)
     {
 
-      int32_t   k_y_start,k_y_end;
-      int32_t   k_x_start,k_x_end;
-      int32_t   chCnt;
-      const int8_t    *pTmp, *pTmpInner;
-      int8_t    *pDst;
+      int32_t k_y_start, k_y_end;
+      int32_t k_x_start, k_x_end;
+      int32_t chCnt;
+      const int8_t *pTmp, *pTmpInner;
+      int8_t *pDst;
 
-      k_y_start = MAX(0, i_y * stride_height - padding_height);
-      k_y_end = MIN(i_y * stride_height - padding_height + dim_kernel_height,dim_src_height);
+      k_y_start = MAX(0, i_y * stride_y - pad_y);
+      k_y_end = MIN(i_y * stride_y - pad_y + kernel_y, input_y);
 
-      k_x_start = MAX(0,i_x * stride_width - padding_width);
-      k_x_end = MIN(i_x * stride_width - padding_width + dim_kernel_width, dim_src_width);
-
+      k_x_start = MAX(0, i_x * stride_x - pad_x);
+      k_x_end = MIN(i_x * stride_x - pad_x + kernel_x, input_x);
 
       pTmp = src;
-      pDst = &dst[ch_src * (i_x + i_y * dim_dst_width)];
+      pDst = &dst[ch_src * (i_x + i_y * output_x)];
 
       chCnt = ch_src >> 4;
-      while(chCnt > 0)
+      while (chCnt > 0)
       {
-        int32x4_t sumV1,sumV2,sumV3,sumV4;
+        int32x4_t sumV1, sumV2, sumV3, sumV4;
 
         int8x16_t tempV;
         int16x8_t tempVLO, tempVHI;
         int32x4_t tempVLOLO, tempVLOHI, tempVHILO, tempVHIHI;
-        int32_t   count = 0;
+        int32_t count = 0;
 
         sumV1 = vdupq_n_s32(0);
         sumV2 = vdupq_n_s32(0);
@@ -136,8 +136,8 @@
         {
           for (k_x = k_x_start; k_x < k_x_end; k_x++)
           {
-            pTmpInner = pTmp + (ch_src * (k_x + k_y * dim_src_width));
-            tempV = vldrbq_s8 (pTmpInner);
+            pTmpInner = pTmp + (ch_src * (k_x + k_y * input_x));
+            tempV = vldrbq_s8(pTmpInner);
 
             tempVLO = vmovlbq_s8(tempV);
             tempVHI = vmovltq_s8(tempV);
@@ -148,16 +148,15 @@
             tempVHILO = vmovlbq_s16(tempVHI);
             tempVHIHI = vmovltq_s16(tempVHI);
 
-            sumV1 = vaddq_s32(sumV1,tempVLOLO);
-            sumV2 = vaddq_s32(sumV2,tempVLOHI);
-            sumV3 = vaddq_s32(sumV3,tempVHILO);
-            sumV4 = vaddq_s32(sumV4,tempVHIHI);
+            sumV1 = vaddq_s32(sumV1, tempVLOLO);
+            sumV2 = vaddq_s32(sumV2, tempVLOHI);
+            sumV3 = vaddq_s32(sumV3, tempVHILO);
+            sumV4 = vaddq_s32(sumV4, tempVHIHI);
 
             count++;
           }
         }
 
-
         sumV1[0] = sumV1[0] > 0 ? (sumV1[0] + count / 2) / count : (sumV1[0] - count / 2) / count;
         sumV1[1] = sumV1[1] > 0 ? (sumV1[1] + count / 2) / count : (sumV1[1] - count / 2) / count;
         sumV1[2] = sumV1[2] > 0 ? (sumV1[2] + count / 2) / count : (sumV1[2] - count / 2) / count;
@@ -190,34 +189,33 @@
         sumV4 = vmaxq_s32(sumV4, vdupq_n_s32(act_min));
         sumV4 = vminq_s32(sumV4, vdupq_n_s32(act_max));
 
-        tempVLO = vmovnbq_s32(tempVLO,sumV1);
-        tempVLO = vmovntq_s32(tempVLO,sumV2);
+        tempVLO = vmovnbq_s32(tempVLO, sumV1);
+        tempVLO = vmovntq_s32(tempVLO, sumV2);
 
-        tempVHI = vmovnbq_s32(tempVHI,sumV3);
-        tempVHI = vmovntq_s32(tempVHI,sumV4);
+        tempVHI = vmovnbq_s32(tempVHI, sumV3);
+        tempVHI = vmovntq_s32(tempVHI, sumV4);
 
+        tempV = vmovnbq_s16(tempV, tempVLO);
+        tempV = vmovntq_s16(tempV, tempVHI);
 
-        tempV = vmovnbq_s16(tempV,tempVLO);
-        tempV = vmovntq_s16(tempV,tempVHI);
-
-        vstrbq_s8(pDst,tempV);
+        vstrbq_s8(pDst, tempV);
         pDst += 16;
 
-        chCnt --;
+        chCnt--;
         pTmp += 16;
       }
 
       chCnt = ch_src & 0xF;
-      while(chCnt > 0)
+      while (chCnt > 0)
       {
-        int32_t       sum = 0;
-        int32_t       count = 0;
+        int32_t sum = 0;
+        int32_t count = 0;
 
         for (k_y = k_y_start; k_y < k_y_end; k_y++)
         {
           for (k_x = k_x_start; k_x < k_x_end; k_x++)
           {
-            sum += pTmp[ch_src * (k_x + k_y * dim_src_width)];
+            sum += pTmp[ch_src * (k_x + k_y * input_x)];
             count++;
           }
         }
@@ -227,7 +225,7 @@
 
         *pDst++ = sum;
 
-        chCnt --;
+        chCnt--;
         pTmp++;
       }
     }
@@ -244,61 +242,65 @@
                           const cmsis_nn_dims *output_dims,
                           q7_t *dst)
 {
-  const int32_t dim_src_height = input_dims->h;
-  const int32_t dim_src_width = input_dims->w;
-  const int32_t dim_dst_height = output_dims->h;
-  const int32_t dim_dst_width = output_dims->w;
-  const int32_t stride_height = pool_params->stride.h;
-  const int32_t stride_width = pool_params->stride.w;
-  const int32_t dim_kernel_height = filter_dims->h;
-  const int32_t dim_kernel_width = filter_dims->w;
-  const int32_t padding_height = pool_params->padding.h;
-  const int32_t padding_width = pool_params->padding.w;
+  const int32_t input_y = input_dims->h;
+  const int32_t input_x = input_dims->w;
+  const int32_t output_y = output_dims->h;
+  const int32_t output_x = output_dims->w;
+  const int32_t stride_y = pool_params->stride.h;
+  const int32_t stride_x = pool_params->stride.w;
+  const int32_t kernel_y = filter_dims->h;
+  const int32_t kernel_x = filter_dims->w;
+  const int32_t pad_y = pool_params->padding.h;
+  const int32_t pad_x = pool_params->padding.w;
   const int32_t act_min = pool_params->activation.min;
   const int32_t act_max = pool_params->activation.max;
   const int32_t ch_src = input_dims->c;
-  q15_t *bufferA = (q15_t *)ctx->buf;
+  q31_t *buffer = (q31_t *)ctx->buf;
 
-#if defined (ARM_MATH_DSP)
+#if defined(ARM_MATH_DSP)
 
-  /* Run the following code for Cortex-M4 and Cortex-M7
+  /* Run the following code for CPU's with DSP extension
    */
-  int32_t k_x, k_y, i_x, i_y;
-
-  for (i_y = 0; i_y < dim_dst_height; i_y++)
+  for (int i_y = 0, idx_y = -pad_y; i_y < output_y; idx_y += stride_y, i_y++)
   {
-    for (i_x = 0; i_x < dim_dst_width; i_x++)
+    for (int i_x = 0, idx_x = -pad_x; i_x < output_x; idx_x += stride_x, i_x++)
     {
-      /* Condition for kernel start dimension: (base_idx_<x,y> + kernel_<x,y>_start) >= 0 */
-      const int32_t base_idx_y = (i_y * stride_height) - padding_height;
-      const int32_t base_idx_x = (i_x * stride_width) - padding_width;
-      const int32_t kernel_y_start = MAX(0, -base_idx_y);
-      const int32_t kernel_x_start = MAX(0, -base_idx_x);
+      /* Condition for kernel start dimension:
+                (base_idx_<x,y> + kernel_<x,y>_start) >= 0 */
+      const int32_t kernel_y_start = MAX(0, -idx_y);
+      const int32_t kernel_x_start = MAX(0, -idx_x);
 
-      /* Condition for kernel end dimension: (base_idx_<x,y> + kernel_<x,y>_end) < dim_src_<width,height> */
-      const int32_t kernel_y_end = MIN(dim_kernel_height, dim_src_height - base_idx_y);
-      const int32_t kernel_x_end = MIN(dim_kernel_width, dim_src_width - base_idx_x);
+      /* Condition for kernel end dimension:
+             (base_idx_<x,y> + kernel_<x,y>_end) < dim_src_<width,height> */
+      const int32_t kernel_y_end = MIN(kernel_y, input_y - idx_y);
+      const int32_t kernel_x_end = MIN(kernel_x, input_x - idx_x);
 
       int count = 0;
 
-      for (k_y = kernel_y_start; k_y < kernel_y_end; k_y++)
+      for (int k_y = kernel_y_start; k_y < kernel_y_end; k_y++)
       {
-        for (k_x = kernel_x_start; k_x < kernel_x_end; k_x++)
+        for (int k_x = kernel_x_start; k_x < kernel_x_end; k_x++)
         {
-          const q7_t *start = src + ch_src * (k_x + base_idx_x + (k_y + base_idx_y) * dim_src_width);
+          const q7_t *start = src + ch_src * (k_x + idx_x + (k_y + idx_y) * input_x);
 
           if (count == 0)
           {
-            arm_q7_to_q15_no_shift(start, bufferA, ch_src);
+            for (int i = 0; i < ch_src; i++)
+            {
+              buffer[i] = start[i];
+            }
           }
           else
           {
-            arm_nn_accumulate_q7_to_q15(bufferA, start, ch_src);
+            for (int i = 0; i < ch_src; i++)
+            {
+              buffer[i] = __QADD(start[i], buffer[i]);
+            }
           }
           count++;
         }
       }
-      buffer_scale_back_q15_to_q7_and_clamp(bufferA, dst, ch_src, count, act_min, act_max);
+      scale_q31_to_q7_and_clamp(buffer, dst, ch_src, count, act_min, act_max);
       dst += ch_src;
     }
   }
@@ -306,25 +308,25 @@
 
   /* Reference C code adapted from CMSIS-NN arm_avepool_q7_HWC.
    */
-  (void)bufferA;
+  (void)buffer;
   int16_t i_ch_in, i_x, i_y;
   int16_t k_x, k_y;
 
-  for (i_y = 0; i_y < dim_dst_height; i_y++)
+  for (i_y = 0; i_y < output_y; i_y++)
   {
-    for (i_x = 0; i_x < dim_dst_width; i_x++)
+    for (i_x = 0; i_x < output_x; i_x++)
     {
       for (i_ch_in = 0; i_ch_in < ch_src; i_ch_in++)
       {
         int sum = 0;
         int count = 0;
-        for (k_y = i_y * stride_height - padding_height; k_y < i_y * stride_height - padding_height + dim_kernel_height; k_y++)
+        for (k_y = i_y * stride_y - pad_y; k_y < i_y * stride_y - pad_y + kernel_y; k_y++)
         {
-          for (k_x = i_x * stride_width - padding_width; k_x < i_x * stride_width - padding_width + dim_kernel_width; k_x++)
+          for (k_x = i_x * stride_x - pad_x; k_x < i_x * stride_x - pad_x + kernel_x; k_x++)
           {
-            if (k_y >= 0 && k_x >= 0 && k_y < dim_src_height && k_x < dim_src_width)
+            if (k_y >= 0 && k_x >= 0 && k_y < input_y && k_x < input_x)
             {
-              sum += src[i_ch_in + ch_src * (k_x + k_y * dim_src_width)];
+              sum += src[i_ch_in + ch_src * (k_x + k_y * input_x)];
               count++;
             }
           }
@@ -333,7 +335,7 @@
         sum = MAX(sum, act_min);
         sum = MIN(sum, act_max);
 
-        dst[i_ch_in + ch_src * (i_x + i_y * dim_dst_width)] = sum;
+        dst[i_ch_in + ch_src * (i_x + i_y * output_x)] = sum;
       }
     }
   }
@@ -344,13 +346,13 @@
 
 #endif /* ARM_MATH_MVEI */
 
-int32_t arm_avgpool_s8_get_buffer_size(const int dim_dst_width,
+int32_t arm_avgpool_s8_get_buffer_size(const int output_x,
                                        const int ch_src)
 {
-  (void)dim_dst_width;
+  (void)output_x;
 
 #if defined(ARM_MATH_DSP) && !defined(ARM_MATH_MVEI)
-  return (ch_src * sizeof(int16_t));
+  return (ch_src * sizeof(int32_t));
 #else
   (void)ch_src;
   return 0;