[NeoMathEngine] AVX512 add vectorHSwish

Signed-off-by: Kirill Golikov <[email protected]>

Author: favorart

NeoMathEngine/src/CPU/x86/CpuX86MathEngineVectorMathPrivate.h

@@ -1155,7 +1155,10 @@ inline void vectorSigmoid( const float* first, float* result, int vectorSize )
 
 inline void vectorHSwish( const float* first, float* result, int vectorSize )
 {
-	if( CCPUInfo::HasAvxAndFma && vectorSize >= NeoML::Avx2::VectorMathMinSize ) {
+	if( CCPUInfo::HasAvx512 && vectorSize >= NeoML::Avx512::VectorMathMinSize ) {
+		NeoML::Avx512::vectorHSwish( first, result, vectorSize );
+		return;
+	} else if( CCPUInfo::HasAvxAndFma && vectorSize >= NeoML::Avx2::VectorMathMinSize ) {
 		NeoML::Avx2::vectorHSwish( first, result, vectorSize );
 		return;
 	}

NeoMathEngine/src/CPU/x86/avx512/Avx512Functions.h

@@ -23,7 +23,7 @@ namespace NeoML {
 
 namespace Avx512 {
 
-// The minimum vector size recommended for using AVX vector functions
+// The minimum vector size recommended for using AVX512 vector functions
 static constexpr int VectorMathMinSize = 32;
 
 void dataCopy( float* dst, const float* src, int vectorSize );
@@ -40,6 +40,8 @@ void vectorReLU( const float* first, float* result, int vectorSize );
 
 void vectorReLU( const float* first, float* result, int vectorSize, float threshold );
 
+void vectorHSwish( const float* first, float* result, int vectorSize );
+
 } // namespace Avx512
 
 } // namespace NeoML

NeoMathEngine/src/CPU/x86/avx512/Avx512VectorFunctions.cpp

@@ -238,6 +238,65 @@ void vectorReLU( const float* first, float* result, int vectorSize, float thresh
 	}
 }
 
+void vectorHSwish( const float* first, float* result, int vectorSize )
+{
+	const __m512 minusThreeSimd = _mm512_set1_ps( -3.f );
+	const __m512 threeSimd = _mm512_set1_ps( 3.f );
+	const __m512 oneSixthSimd = _mm512_set1_ps( 1.f / 6.f );
+	const __m512 zeroSimd = _mm512_setzero_ps();
+
+	//for( int i = 0; i < nonSseSize; ++i ) {
+	//	if( *first <= -3.f ) {
+	//		*result = 0.f;
+	//	} else if( *first >= 3.f ) {
+	//		*result = *first;
+	//	} else {
+	//		*result = *first * ( 1. / 6. ) * ( *first + 3 );
+	//	}
+	//	++result;
+	//	++first;
+	//}
+
+	while( vectorSize >= AvxBlockSize ) {
+		const __m512 firstSimd = _mm512_loadu_ps( first );
+
+		const __mmask16 middleMask = _mm512_cmp_ps_mask( firstSimd, minusThreeSimd, _CMP_GT_OQ ); // ( first > -3. )
+		const __mmask16 rightMask = _mm512_cmp_ps_mask( firstSimd, threeSimd, _CMP_LT_OQ );       // ( first <  3. )
+
+		const __m512 middleSimd = _mm512_mask_blend_ps( middleMask, zeroSimd/*else*/, firstSimd ); // result = ( first > -3. ) ? first : 0.
+
+		const __m512 resultSimd = _mm512_mask_mul_ps( // result = ( middleMask & rightMask ) ? ( first * ( 1. / 6. ) ) * ( first + 3. ) : middleSimd
+			middleSimd /*else*/,
+			middleMask & rightMask,
+			_mm512_mul_ps( firstSimd, oneSixthSimd ), // ( first * ( 1. / 6. ) ) *
+			_mm512_add_ps( firstSimd, threeSimd ) );  // ( first + 3 )
+
+		_mm512_storeu_ps( result, resultSimd );
+
+		first += AvxBlockSize;
+		result += AvxBlockSize;
+		vectorSize -= AvxBlockSize;
+	}
+
+	if( vectorSize > 0 ) {
+		const __mmask16 mask = AVX512_IO_MASK( vectorSize );
+
+		const __m512 firstSimd = _mm512_mask_loadu_ps( zeroSimd, mask, first );
+
+		const __mmask16 middleMask = _mm512_cmp_ps_mask( firstSimd, minusThreeSimd, _CMP_GT_OQ ); // ( first > -3. )
+		const __mmask16 rightMask = _mm512_cmp_ps_mask( firstSimd, threeSimd, _CMP_LT_OQ );       // ( first <  3. )
+
+		const __m512 middleSimd = _mm512_mask_blend_ps( middleMask, zeroSimd/*else*/, firstSimd ); // result = ( first > -3. ) ? first : 0.
+
+		const __m512 resultSimd = _mm512_mask_mul_ps( // result = ( middleMask & rightMask ) ? ( first * ( 1. / 6. ) ) * ( first + 3. ) : middleSimd
+			middleSimd /*else*/,
+			middleMask & rightMask,
+			_mm512_mul_ps( firstSimd, oneSixthSimd ), // ( first * ( 1. / 6. ) ) *
+			_mm512_add_ps( firstSimd, threeSimd ) );  // ( first + 3 )
+
+		_mm512_mask_store_ps( result, mask, resultSimd );
+	}
+}
 
 } // namespace Avx512