[Batchnorm x86] Merge the multiple elempack (#4085)

Co-authored-by: LRY89757 <LRY89757@users.noreply.github.com>
Co-authored-by: nihuini <nihuini@tencent.com>

src/layer/x86/batchnorm_x86.cpp

@@ -34,289 +34,186 @@ BatchNorm_x86::BatchNorm_x86()
 int BatchNorm_x86::forward_inplace(Mat& bottom_top_blob, const Option& opt) const
 {
     int dims = bottom_top_blob.dims;
-
+    int w = bottom_top_blob.w;
+    int h = bottom_top_blob.h;
+    int d = bottom_top_blob.d;
+    int c = bottom_top_blob.c;
     int elempack = bottom_top_blob.elempack;
 
+    if (dims == 1)
+    {
+        float* ptr = bottom_top_blob;
+        const float* aptr = a_data;
+        const float* bptr = b_data;
+
+        const int size = w * elempack;
+
+        int i = 0;
 #if __SSE2__
 #if __AVX__
 #if __AVX512F__
-    if (elempack == 16)
-    {
-        if (dims == 1)
+        for (; i + 15 < size; i += 16)
         {
-            int w = bottom_top_blob.w;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int i = 0; i < w; i++)
-            {
-                float* ptr = (float*)bottom_top_blob + i * 16;
-
-                __m512 _a = _mm512_loadu_ps((const float*)a_data + i * 16);
-                __m512 _b = _mm512_loadu_ps((const float*)b_data + i * 16);
-
-                __m512 _p = _mm512_loadu_ps(ptr);
-                _p = _mm512_fmadd_ps(_p, _b, _a);
-                _mm512_storeu_ps(ptr, _p);
-            }
+            __m512 _p512 = _mm512_loadu_ps(ptr);
+            __m512 _a512 = _mm512_loadu_ps(aptr);
+            __m512 _b512 = _mm512_loadu_ps(bptr);
+            _p512 = _mm512_fmadd_ps(_p512, _b512, _a512);
+            _mm512_storeu_ps(ptr, _p512);
+            ptr += 16;
+            aptr += 16;
+            bptr += 16;
         }
-
-        if (dims == 2)
+#endif // __AVX512F__
+        for (; i + 7 < size; i += 8)
         {
-            int w = bottom_top_blob.w;
-            int h = bottom_top_blob.h;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int i = 0; i < h; i++)
-            {
-                __m512 _a = _mm512_loadu_ps((const float*)a_data + i * 16);
-                __m512 _b = _mm512_loadu_ps((const float*)b_data + i * 16);
-
-                float* ptr = bottom_top_blob.row(i);
-
-                for (int j = 0; j < w; j++)
-                {
-                    __m512 _p = _mm512_loadu_ps(ptr);
-                    _p = _mm512_fmadd_ps(_p, _b, _a);
-                    _mm512_storeu_ps(ptr, _p);
-
-                    ptr += 16;
-                }
-            }
+            __m256 _p256 = _mm256_loadu_ps(ptr);
+            __m256 _a256 = _mm256_loadu_ps(aptr);
+            __m256 _b256 = _mm256_loadu_ps(bptr);
+            _p256 = _mm256_comp_fmadd_ps(_p256, _b256, _a256);
+            _mm256_storeu_ps(ptr, _p256);
+            ptr += 8;
+            aptr += 8;
+            bptr += 8;
         }
-
-        if (dims == 3 || dims == 4)
+#endif // __AVX__
+        for (; i + 3 < size; i += 4)
         {
-            int w = bottom_top_blob.w;
-            int h = bottom_top_blob.h;
-            int d = bottom_top_blob.d;
-            int c = bottom_top_blob.c;
-            int size = w * h * d;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int q = 0; q < c; q++)
-            {
-                __m512 _a = _mm512_loadu_ps((const float*)a_data + q * 16);
-                __m512 _b = _mm512_loadu_ps((const float*)b_data + q * 16);
-
-                float* ptr = bottom_top_blob.channel(q);
-
-                for (int i = 0; i < size; i++)
-                {
-                    __m512 _p = _mm512_loadu_ps(ptr);
-                    _p = _mm512_fmadd_ps(_p, _b, _a);
-                    _mm512_storeu_ps(ptr, _p);
-
-                    ptr += 16;
-                }
-            }
+            __m128 _p128 = _mm_loadu_ps(ptr);
+            __m128 _a128 = _mm_loadu_ps(aptr);
+            __m128 _b128 = _mm_loadu_ps(bptr);
+            _p128 = _mm_comp_fmadd_ps(_p128, _b128, _a128);
+            _mm_storeu_ps(ptr, _p128);
+            ptr += 4;
+            aptr += 4;
+            bptr += 4;
+        }
+#endif // __SSE__
+        for (; i < size; i++)
+        {
+            *ptr = *bptr * *ptr + *aptr;
+            ptr++;
+            aptr++;
+            bptr++;
         }
-
-        return 0;
     }
-#endif // __AVX512F__
 
-    if (elempack == 8)
+    if (dims == 2)
     {
-        if (dims == 1)
-        {
-            int w = bottom_top_blob.w;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int i = 0; i < w; i++)
-            {
-                float* ptr = (float*)bottom_top_blob + i * 8;
-
-                __m256 _a = _mm256_loadu_ps((const float*)a_data + i * 8);
-                __m256 _b = _mm256_loadu_ps((const float*)b_data + i * 8);
+        const int size = w * elempack;
 
-                __m256 _p = _mm256_loadu_ps(ptr);
-                _p = _mm256_comp_fmadd_ps(_p, _b, _a);
-                _mm256_storeu_ps(ptr, _p);
-            }
-        }
-
-        if (dims == 2)
+        #pragma omp parallel for num_threads(opt.num_threads)
+        for (int i = 0; i < h; i++)
         {
-            int w = bottom_top_blob.w;
-            int h = bottom_top_blob.h;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int i = 0; i < h; i++)
-            {
-                __m256 _a = _mm256_loadu_ps((const float*)a_data + i * 8);
-                __m256 _b = _mm256_loadu_ps((const float*)b_data + i * 8);
-
-                float* ptr = bottom_top_blob.row(i);
-
-                for (int j = 0; j < w; j++)
-                {
-                    __m256 _p = _mm256_loadu_ps(ptr);
-                    _p = _mm256_comp_fmadd_ps(_p, _b, _a);
-                    _mm256_storeu_ps(ptr, _p);
-
-                    ptr += 8;
-                }
-            }
-        }
+            float* ptr = bottom_top_blob.row(i);
+            float a = a_data[i];
+            float b = b_data[i];
 
-        if (dims == 3 || dims == 4)
-        {
-            int w = bottom_top_blob.w;
-            int h = bottom_top_blob.h;
-            int d = bottom_top_blob.d;
-            int c = bottom_top_blob.c;
-            int size = w * h * d;
+#if __SSE2__
+            __m128 _a128 = (elempack == 4) ? _mm_loadu_ps((const float*)a_data + i * 4) : _mm_set1_ps(a);
+            __m128 _b128 = (elempack == 4) ? _mm_loadu_ps((const float*)b_data + i * 4) : _mm_set1_ps(b);
+#if __AVX__
+            __m256 _a256 = (elempack == 8) ? _mm256_loadu_ps((const float*)a_data + i * 8) : _mm256_insertf128_ps(_mm256_castps128_ps256(_a128), _a128, 1);
+            __m256 _b256 = (elempack == 8) ? _mm256_loadu_ps((const float*)b_data + i * 8) : _mm256_insertf128_ps(_mm256_castps128_ps256(_b128), _b128, 1);
+#if __AVX512F__
+            __m512 _a512 = (elempack == 16) ? _mm512_loadu_ps((const float*)a_data + i * 16) : _mm512_insertf32x8(_mm512_castps256_ps512(_a256), _a256, 1);
+            __m512 _b512 = (elempack == 16) ? _mm512_loadu_ps((const float*)b_data + i * 16) : _mm512_insertf32x8(_mm512_castps256_ps512(_b256), _b256, 1);
+#endif // __AVX512F__
+#endif // __AVX__
+#endif // __SSE2__
 
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int q = 0; q < c; q++)
+            int j = 0;
+#if __SSE2__
+#if __AVX__
+#if __AVX512F__
+            for (; j + 15 < size; j += 16)
             {
-                __m256 _a = _mm256_loadu_ps((const float*)a_data + q * 8);
-                __m256 _b = _mm256_loadu_ps((const float*)b_data + q * 8);
-
-                float* ptr = bottom_top_blob.channel(q);
-
-                for (int i = 0; i < size; i++)
-                {
-                    __m256 _p = _mm256_loadu_ps(ptr);
-                    _p = _mm256_comp_fmadd_ps(_p, _b, _a);
-                    _mm256_storeu_ps(ptr, _p);
-
-                    ptr += 8;
-                }
+                __m512 _p512 = _mm512_loadu_ps(ptr);
+                _p512 = _mm512_fmadd_ps(_p512, _b512, _a512);
+                _mm512_storeu_ps(ptr, _p512);
+                ptr += 16;
             }
-        }
-
-        return 0;
-    }
-#endif // __AVX__
-
-    if (elempack == 4)
-    {
-        if (dims == 1)
-        {
-            int w = bottom_top_blob.w;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int i = 0; i < w; i++)
+#endif // __AVX512F__
+            for (; j + 7 < size; j += 8)
             {
-                float* ptr = (float*)bottom_top_blob + i * 4;
-
-                __m128 _a = _mm_load_ps((const float*)a_data + i * 4);
-                __m128 _b = _mm_load_ps((const float*)b_data + i * 4);
-
-                __m128 _p = _mm_load_ps(ptr);
-                _p = _mm_mul_ps(_p, _b);
-                _p = _mm_add_ps(_p, _a);
-                _mm_store_ps(ptr, _p);
+                __m256 _p256 = _mm256_loadu_ps(ptr);
+                _p256 = _mm256_comp_fmadd_ps(_p256, _b256, _a256);
+                _mm256_storeu_ps(ptr, _p256);
+                ptr += 8;
             }
-        }
-
-        if (dims == 2)
-        {
-            int w = bottom_top_blob.w;
-            int h = bottom_top_blob.h;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int i = 0; i < h; i++)
+#endif // __AVX__
+            for (; j + 3 < size; j += 4)
             {
-                __m128 _a = _mm_load_ps((const float*)a_data + i * 4);
-                __m128 _b = _mm_load_ps((const float*)b_data + i * 4);
-
-                float* ptr = bottom_top_blob.row(i);
-
-                for (int j = 0; j < w; j++)
-                {
-                    __m128 _p = _mm_load_ps(ptr);
-                    _p = _mm_mul_ps(_p, _b);
-                    _p = _mm_add_ps(_p, _a);
-                    _mm_store_ps(ptr, _p);
-
-                    ptr += 4;
-                }
+                __m128 _p128 = _mm_loadu_ps(ptr);
+                _p128 = _mm_comp_fmadd_ps(_p128, _b128, _a128);
+                _mm_storeu_ps(ptr, _p128);
+                ptr += 4;
             }
-        }
-
-        if (dims == 3 || dims == 4)
-        {
-            int w = bottom_top_blob.w;
-            int h = bottom_top_blob.h;
-            int d = bottom_top_blob.d;
-            int c = bottom_top_blob.c;
-            int size = w * h * d;
-
-            #pragma omp parallel for num_threads(opt.num_threads)
-            for (int q = 0; q < c; q++)
+#endif // __SSE__
+            for (; j < size; j++)
             {
-                __m128 _a = _mm_load_ps((const float*)a_data + q * 4);
-                __m128 _b = _mm_load_ps((const float*)b_data + q * 4);
-
-                float* ptr = bottom_top_blob.channel(q);
-
-                for (int i = 0; i < size; i++)
-                {
-                    __m128 _p = _mm_load_ps(ptr);
-                    _p = _mm_mul_ps(_p, _b);
-                    _p = _mm_add_ps(_p, _a);
-                    _mm_store_ps(ptr, _p);
-
-                    ptr += 4;
-                }
+                *ptr = b * *ptr + a;
+                ptr++;
             }
         }
-
-        return 0;
     }
-#endif // __SSE2__
-
-    if (dims != 3 && dims != 4)
-        return BatchNorm::forward_inplace(bottom_top_blob, opt);
 
-    int w = bottom_top_blob.w;
-    int h = bottom_top_blob.h;
-    int d = bottom_top_blob.d;
-    // int c = bottom_top_blob.c;
-    int size = w * h * d;
-
-    #pragma omp parallel for num_threads(opt.num_threads)
-    for (int q = 0; q < channels; q++)
+    if (dims == 3 || dims == 4)
     {
-        float* ptr = bottom_top_blob.channel(q);
+        const int size = w * h * d * elempack;
 
-        float a = a_data[q];
-        float b = b_data[q];
+        #pragma omp parallel for num_threads(opt.num_threads)
+        for (int q = 0; q < c; q++)
+        {
+            float* ptr = bottom_top_blob.channel(q);
+            float a = a_data[q];
+            float b = b_data[q];
 
-        int i = 0;
 #if __SSE2__
+            __m128 _a128 = (elempack == 4) ? _mm_loadu_ps((const float*)a_data + q * 4) : _mm_set1_ps(a);
+            __m128 _b128 = (elempack == 4) ? _mm_loadu_ps((const float*)b_data + q * 4) : _mm_set1_ps(b);
 #if __AVX__
-        __m256 _a256 = _mm256_set1_ps(a);
-        __m256 _b256 = _mm256_set1_ps(b);
-        for (; i + 7 < size; i += 8)
-        {
-            __m256 _p = _mm256_loadu_ps(ptr);
-            _p = _mm256_comp_fmadd_ps(_p, _b256, _a256);
-            _mm256_storeu_ps(ptr, _p);
-            ptr += 8;
-        }
-        __m128 _a128 = _mm256_castps256_ps128(_a256);
-        __m128 _b128 = _mm256_castps256_ps128(_b256);
-#else
-        __m128 _a128 = _mm_set1_ps(a);
-        __m128 _b128 = _mm_set1_ps(b);
+            __m256 _a256 = (elempack == 8) ? _mm256_loadu_ps((const float*)a_data + q * 8) : _mm256_insertf128_ps(_mm256_castps128_ps256(_a128), _a128, 1);
+            __m256 _b256 = (elempack == 8) ? _mm256_loadu_ps((const float*)b_data + q * 8) : _mm256_insertf128_ps(_mm256_castps128_ps256(_b128), _b128, 1);
+#if __AVX512F__
+            __m512 _a512 = (elempack == 16) ? _mm512_loadu_ps((const float*)a_data + q * 16) : _mm512_insertf32x8(_mm512_castps256_ps512(_a256), _a256, 1);
+            __m512 _b512 = (elempack == 16) ? _mm512_loadu_ps((const float*)b_data + q * 16) : _mm512_insertf32x8(_mm512_castps256_ps512(_b256), _b256, 1);
+#endif // __AVX512F__
 #endif // __AVX__
-        for (; i + 3 < size; i += 4)
-        {
-            __m128 _p = _mm_load_ps(ptr);
-            _p = _mm_mul_ps(_p, _b128);
-            _p = _mm_add_ps(_p, _a128);
-            _mm_store_ps(ptr, _p);
-            ptr += 4;
-        }
 #endif // __SSE2__
-        for (; i < size; i++)
-        {
-            *ptr = b * *ptr + a;
 
-            ptr++;
+            int i = 0;
+#if __SSE2__
+#if __AVX__
+#if __AVX512F__
+            for (; i + 15 < size; i += 16)
+            {
+                __m512 _p512 = _mm512_loadu_ps(ptr);
+                _p512 = _mm512_fmadd_ps(_p512, _b512, _a512);
+                _mm512_storeu_ps(ptr, _p512);
+                ptr += 16;
+            }
+#endif // __AVX512F__
+            for (; i + 7 < size; i += 8)
+            {
+                __m256 _p256 = _mm256_loadu_ps(ptr);
+                _p256 = _mm256_comp_fmadd_ps(_p256, _b256, _a256);
+                _mm256_storeu_ps(ptr, _p256);
+                ptr += 8;
+            }
+#endif // __AVX__
+            for (; i + 3 < size; i += 4)
+            {
+                __m128 _p128 = _mm_loadu_ps(ptr);
+                _p128 = _mm_comp_fmadd_ps(_p128, _b128, _a128);
+                _mm_storeu_ps(ptr, _p128);
+                ptr += 4;
+            }
+#endif // __SSE__
+            for (; i < size; i++)
+            {
+                *ptr = b * *ptr + a;
+                ptr++;
+            }
         }
     }