diff --git a/src/layer/x86/batchnorm_x86.cpp b/src/layer/x86/batchnorm_x86.cpp index 8a5d8ab6a..96eec7d87 100644 --- a/src/layer/x86/batchnorm_x86.cpp +++ b/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++; + } } }