1188 lines
42 KiB
C++
1188 lines
42 KiB
C++
#ifndef MATMUL_SPECIALISATIONS_KERNELS_H
|
|
#define MATMUL_SPECIALISATIONS_KERNELS_H
|
|
|
|
|
|
// Forward declare
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
namespace internal {
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M==N && M==8 && std::is_same<T,float>::value,bool>::type = 0 >
|
|
FASTOR_INLINE
|
|
void _matmul8k8_float(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out);
|
|
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M==N && M==8 && std::is_same<T,double>::value,bool>::type = 0 >
|
|
FASTOR_INLINE
|
|
void _matmul8k8_double(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out);
|
|
|
|
} // internal
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
|
|
#ifdef FASTOR_SSE2_IMPL
|
|
// (2xk) x (kx2) matrices
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<(M!=K && M==N && M==2 && std::is_same<T,double>::value),bool>::type = 0>
|
|
FASTOR_INLINE
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m128d out_row0 = ZEROPD;
|
|
__m128d out_row1 = ZEROPD;
|
|
|
|
for (size_t i=0; i<K; ++i) {
|
|
__m128d brow = _mm_loadu_pd(&b[i*2]);
|
|
#ifndef FASTOR_FMA_IMPL
|
|
// row 0
|
|
__m128d a_vec0 = _mm_set1_pd(a[i]);
|
|
out_row0 = _mm_add_pd(out_row0,_mm_mul_pd(a_vec0,brow));
|
|
// row 1
|
|
__m128d a_vec1 = _mm_set1_pd(a[K+i]);
|
|
out_row1 = _mm_add_pd(out_row1,_mm_mul_pd(a_vec1,brow));
|
|
#else
|
|
// row 0
|
|
__m128d a_vec0 = _mm_set1_pd(a[i]);
|
|
out_row0 = _mm_fmadd_pd(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m128d a_vec1 = _mm_set1_pd(a[K+i]);
|
|
out_row1 = _mm_fmadd_pd(a_vec1,brow,out_row1);
|
|
#endif
|
|
}
|
|
_mm_storeu_pd(out ,out_row0);
|
|
_mm_storeu_pd(out+2,out_row1);
|
|
}
|
|
|
|
// (2xk) x (kx2) matrices
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<(M!=K && M==N && M==2 && std::is_same<T,float>::value),bool>::type = 0>
|
|
FASTOR_INLINE
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m128 out_row0 = ZEROPS;
|
|
__m128 out_row1 = ZEROPS;
|
|
|
|
__m128 brow = ZEROPS;
|
|
for (size_t i=0; i<K; i++) {
|
|
// __m128 brow = _mm_loadu_ps(&b[i*2]);
|
|
brow = _mm_loadl_pi(brow,(__m64*)&b[i*2]);
|
|
#ifndef FASTOR_FMA_IMPL
|
|
// row 0
|
|
__m128 a_vec0 = _mm_set1_ps(a[i]);
|
|
out_row0 = _mm_add_ps(out_row0,_mm_mul_ps(a_vec0,brow));
|
|
// row 1
|
|
__m128 a_vec1 = _mm_set1_ps(a[K+i]);
|
|
out_row1 = _mm_add_ps(out_row1,_mm_mul_ps(a_vec1,brow));
|
|
#else
|
|
// row 0
|
|
__m128 a_vec0 = _mm_set1_ps(a[i]);
|
|
out_row0 = _mm_fmadd_ps(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m128 a_vec1 = _mm_set1_ps(a[K+i]);
|
|
out_row1 = _mm_fmadd_ps(a_vec1,brow,out_row1);
|
|
#endif
|
|
}
|
|
_mm_storeu_ps(out,_mm_shuffle_ps(out_row0,out_row1,_MM_SHUFFLE(1,0,1,0)));
|
|
}
|
|
#else
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M!=K && M==N && M==2,bool>::type = 0>
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
internal::_matmul_base<T,M,K,N>(a,b,out);
|
|
}
|
|
#endif
|
|
|
|
#ifdef FASTOR_AVX_IMPL
|
|
// (3xk) x (kx3) matrices
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<(M!=K && M==N && M==3 && std::is_same<T,double>::value),bool>::type = 0>
|
|
FASTOR_INLINE
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m256d out_row0 = VZEROPD;
|
|
__m256d out_row1 = VZEROPD;
|
|
__m256d out_row2 = VZEROPD;
|
|
|
|
for (size_t i=0; i<K; ++i) {
|
|
__m256d brow = _mm256_loadul3_pd(&b[i*3]);
|
|
#ifndef FASTOR_FMA_IMPL
|
|
// row 0
|
|
__m256d a_vec0 = _mm256_set1_pd(a[i]);
|
|
out_row0 = _mm256_add_pd(out_row0,_mm256_mul_pd(a_vec0,brow));
|
|
// row 1
|
|
__m256d a_vec1 = _mm256_set1_pd(a[K+i]);
|
|
out_row1 = _mm256_add_pd(out_row1,_mm256_mul_pd(a_vec1,brow));
|
|
// row 2
|
|
__m256d a_vec2 = _mm256_set1_pd(a[2*K+i]);
|
|
out_row2 = _mm256_add_pd(out_row2,_mm256_mul_pd(a_vec2,brow));
|
|
#else
|
|
// row 0
|
|
__m256d a_vec0 = _mm256_set1_pd(a[i]);
|
|
out_row0 = _mm256_fmadd_pd(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m256d a_vec1 = _mm256_set1_pd(a[K+i]);
|
|
out_row1 = _mm256_fmadd_pd(a_vec1,brow,out_row1);
|
|
// row 2
|
|
__m256d a_vec2 = _mm256_set1_pd(a[2*K+i]);
|
|
out_row2 = _mm256_fmadd_pd(a_vec2,brow,out_row2);
|
|
#endif
|
|
}
|
|
_mm256_storeu_pd (out,out_row0);
|
|
_mm256_storeu_pd (out+3,out_row1);
|
|
_mm256_storeul3_pd(out+6,out_row2);
|
|
}
|
|
|
|
#else
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M!=K && M==N && M==3 && std::is_same<T,double>::value,bool>::type = 0>
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
internal::_matmul_base<T,M,K,N>(a,b,out);
|
|
}
|
|
#endif
|
|
|
|
#ifdef FASTOR_SSE2_IMPL
|
|
// (3xk) x (kx3) matrices
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<(M!=K && M==N && M==3 && std::is_same<T,float>::value),bool>::type = 0>
|
|
FASTOR_INLINE
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m128 out_row0 = ZEROPS;
|
|
__m128 out_row1 = ZEROPS;
|
|
__m128 out_row2 = ZEROPS;
|
|
|
|
for (size_t i=0; i<K; ++i) {
|
|
__m128 brow = _mm_loadul3_ps(&b[i*3]);
|
|
#ifndef FASTOR_FMA_IMPL
|
|
// row 0
|
|
__m128 a_vec0 = _mm_set1_ps(a[i]);
|
|
out_row0 = _mm_add_ps(out_row0,_mm_mul_ps(a_vec0,brow));
|
|
// row 1
|
|
__m128 a_vec1 = _mm_set1_ps(a[K+i]);
|
|
out_row1 = _mm_add_ps(out_row1,_mm_mul_ps(a_vec1,brow));
|
|
// row 2
|
|
__m128 a_vec2 = _mm_set1_ps(a[2*K+i]);
|
|
out_row2 = _mm_add_ps(out_row2,_mm_mul_ps(a_vec2,brow));
|
|
#else
|
|
// row 0
|
|
__m128 a_vec0 = _mm_set1_ps(a[i]);
|
|
out_row0 = _mm_fmadd_ps(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m128 a_vec1 = _mm_set1_ps(a[K+i]);
|
|
out_row1 = _mm_fmadd_ps(a_vec1,brow,out_row1);
|
|
// row 2
|
|
__m128 a_vec2 = _mm_set1_ps(a[2*K+i]);
|
|
out_row2 = _mm_fmadd_ps(a_vec2,brow,out_row2);
|
|
#endif
|
|
|
|
}
|
|
_mm_storeu_ps (out,out_row0);
|
|
_mm_storeu_ps (&out[3],out_row1);
|
|
_mm_storeul3_ps(&out[6],out_row2);
|
|
}
|
|
#else
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<(M!=K && M==N && M==3 && std::is_same<T,float>::value),bool>::type = 0>
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
internal::_matmul_base<T,M,K,N>(a,b,out);
|
|
}
|
|
#endif
|
|
|
|
#ifdef FASTOR_AVX_IMPL
|
|
// (4xk) x (kx4) matrices
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<(M!=K && M==N && M==4 && std::is_same<T,double>::value),bool>::type = 0>
|
|
FASTOR_INLINE
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m256d out_row0 = VZEROPD;
|
|
__m256d out_row1 = VZEROPD;
|
|
__m256d out_row2 = VZEROPD;
|
|
__m256d out_row3 = VZEROPD;
|
|
|
|
for (size_t i=0; i<K; ++i) {
|
|
__m256d brow = _mm256_loadu_pd(&b[i*4]);
|
|
#ifndef FASTOR_FMA_IMPL
|
|
// row 0
|
|
__m256d a_vec0 = _mm256_set1_pd(a[i]);
|
|
out_row0 = _mm256_add_pd(out_row0,_mm256_mul_pd(a_vec0,brow));
|
|
// row 1
|
|
__m256d a_vec1 = _mm256_set1_pd(a[K+i]);
|
|
out_row1 = _mm256_add_pd(out_row1,_mm256_mul_pd(a_vec1,brow));
|
|
// row 2
|
|
__m256d a_vec2 = _mm256_set1_pd(a[2*K+i]);
|
|
out_row2 = _mm256_add_pd(out_row2,_mm256_mul_pd(a_vec2,brow));
|
|
// row 3
|
|
__m256d a_vec3 = _mm256_set1_pd(a[3*K+i]);
|
|
out_row3 = _mm256_add_pd(out_row3,_mm256_mul_pd(a_vec3,brow));
|
|
#else
|
|
// row 0
|
|
__m256d a_vec0 = _mm256_set1_pd(a[i]);
|
|
out_row0 = _mm256_fmadd_pd(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m256d a_vec1 = _mm256_set1_pd(a[K+i]);
|
|
out_row1 = _mm256_fmadd_pd(a_vec1,brow,out_row1);
|
|
// row 2
|
|
__m256d a_vec2 = _mm256_set1_pd(a[2*K+i]);
|
|
out_row2 = _mm256_fmadd_pd(a_vec2,brow,out_row2);
|
|
// row 3
|
|
__m256d a_vec3 = _mm256_set1_pd(a[3*K+i]);
|
|
out_row3 = _mm256_fmadd_pd(a_vec3,brow,out_row3);
|
|
#endif
|
|
}
|
|
_mm256_storeu_pd(&out[0] , out_row0);
|
|
_mm256_storeu_pd(&out[4] , out_row1);
|
|
_mm256_storeu_pd(&out[8] , out_row2);
|
|
_mm256_storeu_pd(&out[12], out_row3);
|
|
}
|
|
#else
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M!=K && M==N && M==4 && std::is_same<T,double>::value,bool>::type = 0>
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
internal::_matmul_base<T,M,K,N>(a,b,out);
|
|
}
|
|
#endif
|
|
|
|
#ifdef FASTOR_SSE2_IMPL
|
|
// (4xk) x (kx4) matrices
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<(M!=K && M==N && M==4 && std::is_same<T,float>::value),bool>::type = 0>
|
|
FASTOR_INLINE
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m128 out_row0 = ZEROPS;
|
|
__m128 out_row1 = ZEROPS;
|
|
__m128 out_row2 = ZEROPS;
|
|
__m128 out_row3 = ZEROPS;
|
|
|
|
for (size_t i=0; i<K; ++i) {
|
|
__m128 brow = _mm_loadu_ps(&b[i*4]);
|
|
#ifndef FASTOR_FMA_IMPL
|
|
// row 0
|
|
__m128 a_vec0 = _mm_set1_ps(a[i]);
|
|
out_row0 = _mm_add_ps(out_row0,_mm_mul_ps(a_vec0,brow));
|
|
// row 1
|
|
__m128 a_vec1 = _mm_set1_ps(a[K+i]);
|
|
out_row1 = _mm_add_ps(out_row1,_mm_mul_ps(a_vec1,brow));
|
|
// row 2
|
|
__m128 a_vec2 = _mm_set1_ps(a[2*K+i]);
|
|
out_row2 = _mm_add_ps(out_row2,_mm_mul_ps(a_vec2,brow));
|
|
// row 3
|
|
__m128 a_vec3 = _mm_set1_ps(a[3*K+i]);
|
|
out_row3 = _mm_add_ps(out_row3,_mm_mul_ps(a_vec3,brow));
|
|
#else
|
|
// row 0
|
|
__m128 a_vec0 = _mm_set1_ps(a[i]);
|
|
out_row0 = _mm_fmadd_ps(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m128 a_vec1 = _mm_set1_ps(a[K+i]);
|
|
out_row1 = _mm_fmadd_ps(a_vec1,brow,out_row1);
|
|
// row 2
|
|
__m128 a_vec2 = _mm_set1_ps(a[2*K+i]);
|
|
out_row2 = _mm_fmadd_ps(a_vec2,brow,out_row2);
|
|
// row 3
|
|
__m128 a_vec3 = _mm_set1_ps(a[3*K+i]);
|
|
out_row3 = _mm_fmadd_ps(a_vec3,brow,out_row3);
|
|
#endif
|
|
}
|
|
_mm_storeu_ps(out,out_row0);
|
|
_mm_storeu_ps(&out[4],out_row1);
|
|
_mm_storeu_ps(&out[8],out_row2);
|
|
_mm_storeu_ps(&out[12],out_row3);
|
|
}
|
|
#else
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M!=K && M==N && M==4 && std::is_same<T,float>::value,bool>::type = 0>
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
internal::_matmul_base<T,M,K,N>(a,b,out);
|
|
}
|
|
#endif
|
|
|
|
|
|
// (8xk) x (kx8) matrices
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M!=K && M==N && M==8 && std::is_same<T,double>::value,bool>::type = 0>
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
#ifdef FASTOR_AVX512_IMPL
|
|
FASTOR_IF_CONSTEXPR(K<=64)
|
|
internal::_matmul8k8_double<T,M,K,N>(a,b,out);
|
|
else
|
|
internal::_matmul_mk_smalln<T,M,K,N>(a,b,out);
|
|
#else
|
|
internal::_matmul_mk_smalln<T,M,K,N>(a,b,out);
|
|
#endif
|
|
}
|
|
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M!=K && M==N && M==8 && std::is_same<T,float>::value,bool>::type = 0>
|
|
void _matmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
#ifdef FASTOR_AVX_IMPL
|
|
internal::_matmul8k8_float<T,M,K,N>(a,b,out);
|
|
#else
|
|
internal::_matmul_base<T,M,K,N>(a,b,out);
|
|
#endif
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
#ifdef FASTOR_SSE2_IMPL
|
|
template<>
|
|
FASTOR_INLINE
|
|
void _matmul<float,2,2,2>(const float * FASTOR_RESTRICT a, const float * FASTOR_RESTRICT b, float * FASTOR_RESTRICT out) {
|
|
|
|
// Ivy 17 - Sky 16 [12 if fused] OPS
|
|
__m128 ar = _mm_loadu_ps(a);
|
|
__m128 br = _mm_loadu_ps(b);
|
|
__m128 ar0 = _mm_shuffle_ps(ar,ar,_MM_SHUFFLE(2,2,0,0));
|
|
__m128 ar1 = _mm_shuffle_ps(ar,ar,_MM_SHUFFLE(3,3,1,1));
|
|
__m128 br0 = _mm_shuffle_ps(br,br,_MM_SHUFFLE(1,0,1,0));
|
|
__m128 br1 = _mm_shuffle_ps(br,br,_MM_SHUFFLE(3,2,3,2));
|
|
__m128 res = _mm_add_ps(_mm_mul_ps(ar0,br0),_mm_mul_ps(ar1,br1));
|
|
_mm_storeu_ps(out,res);
|
|
}
|
|
|
|
template<>
|
|
FASTOR_INLINE
|
|
void _matmul<float,3,3,3>(const float * FASTOR_RESTRICT a, const float * FASTOR_RESTRICT b, float * FASTOR_RESTRICT out) {
|
|
|
|
// Ivy 63 - HW 45 - Sky 36 OPS
|
|
// Fully vectorised approach that reduces
|
|
// (27 scalar mul + 18 scalar add) to
|
|
// (9 SSE mul + 6 SEE add) [non-FMA]
|
|
// (3 SSE mul + 6 SEE FMA) [ FMA]
|
|
|
|
__m128 bmm0 = _mm_loadu_ps (b );
|
|
__m128 bmm1 = _mm_loadu_ps (b+3);
|
|
__m128 bmm2 = _mm_loadul3_ps(b+6);
|
|
|
|
__m128 omm0(_mm_mul_ps(_mm_set1_ps(a[0]),bmm0));
|
|
__m128 omm1(_mm_mul_ps(_mm_set1_ps(a[3]),bmm0));
|
|
__m128 omm2(_mm_mul_ps(_mm_set1_ps(a[6]),bmm0));
|
|
#ifndef FASTOR_FMA_IMPL
|
|
omm0 = _mm_add_ps(omm0,_mm_mul_ps(_mm_set1_ps(a[1]),bmm1));
|
|
omm1 = _mm_add_ps(omm1,_mm_mul_ps(_mm_set1_ps(a[4]),bmm1));
|
|
omm2 = _mm_add_ps(omm2,_mm_mul_ps(_mm_set1_ps(a[7]),bmm1));
|
|
|
|
omm0 = _mm_add_ps(omm0,_mm_mul_ps(_mm_set1_ps(a[2]),bmm2));
|
|
omm1 = _mm_add_ps(omm1,_mm_mul_ps(_mm_set1_ps(a[5]),bmm2));
|
|
omm2 = _mm_add_ps(omm2,_mm_mul_ps(_mm_set1_ps(a[8]),bmm2));
|
|
#else
|
|
omm0 = _mm_fmadd_ps(_mm_set1_ps(a[1]),bmm1,omm0);
|
|
omm1 = _mm_fmadd_ps(_mm_set1_ps(a[4]),bmm1,omm1);
|
|
omm2 = _mm_fmadd_ps(_mm_set1_ps(a[7]),bmm1,omm2);
|
|
|
|
omm0 = _mm_fmadd_ps(_mm_set1_ps(a[2]),bmm2,omm0);
|
|
omm1 = _mm_fmadd_ps(_mm_set1_ps(a[5]),bmm2,omm1);
|
|
omm2 = _mm_fmadd_ps(_mm_set1_ps(a[8]),bmm2,omm2);
|
|
#endif
|
|
|
|
_mm_storeu_ps (out , omm0);
|
|
_mm_storeu_ps (out+3, omm1);
|
|
_mm_storeul3_ps(out+6, omm2);
|
|
}
|
|
|
|
template<>
|
|
FASTOR_INLINE void _matmul<float,4,4,4>(const float * FASTOR_RESTRICT a, const float * FASTOR_RESTRICT b, float * FASTOR_RESTRICT out) {
|
|
|
|
// Ivy 116 - HW 80 - Sky 64 OPS
|
|
// Fully vectorised approach that reduces
|
|
// (64 scalar mul + 48 scalar add) to
|
|
// (16 SSE mul + 12 SEE add) [non-FMA]
|
|
// (4 SSE mul + 12 SEE FMA) [ FMA]
|
|
|
|
__m128 bmm0 = _mm_loadu_ps (b );
|
|
__m128 bmm1 = _mm_loadu_ps (b+4 );
|
|
__m128 bmm2 = _mm_loadu_ps (b+8 );
|
|
__m128 bmm3 = _mm_loadu_ps (b+12);
|
|
|
|
__m128 omm0(_mm_mul_ps(_mm_set1_ps(a[0 ]),bmm0));
|
|
__m128 omm1(_mm_mul_ps(_mm_set1_ps(a[4 ]),bmm0));
|
|
__m128 omm2(_mm_mul_ps(_mm_set1_ps(a[8 ]),bmm0));
|
|
__m128 omm3(_mm_mul_ps(_mm_set1_ps(a[12]),bmm0));
|
|
#ifndef FASTOR_FMA_IMPL
|
|
omm0 = _mm_add_ps(omm0,_mm_mul_ps(_mm_set1_ps(a[1 ]),bmm1));
|
|
omm1 = _mm_add_ps(omm1,_mm_mul_ps(_mm_set1_ps(a[5 ]),bmm1));
|
|
omm2 = _mm_add_ps(omm2,_mm_mul_ps(_mm_set1_ps(a[9 ]),bmm1));
|
|
omm3 = _mm_add_ps(omm3,_mm_mul_ps(_mm_set1_ps(a[13]),bmm1));
|
|
|
|
omm0 = _mm_add_ps(omm0,_mm_mul_ps(_mm_set1_ps(a[2 ]),bmm2));
|
|
omm1 = _mm_add_ps(omm1,_mm_mul_ps(_mm_set1_ps(a[6 ]),bmm2));
|
|
omm2 = _mm_add_ps(omm2,_mm_mul_ps(_mm_set1_ps(a[10]),bmm2));
|
|
omm3 = _mm_add_ps(omm3,_mm_mul_ps(_mm_set1_ps(a[14]),bmm2));
|
|
|
|
omm0 = _mm_add_ps(omm0,_mm_mul_ps(_mm_set1_ps(a[3 ]),bmm3));
|
|
omm1 = _mm_add_ps(omm1,_mm_mul_ps(_mm_set1_ps(a[7 ]),bmm3));
|
|
omm2 = _mm_add_ps(omm2,_mm_mul_ps(_mm_set1_ps(a[11]),bmm3));
|
|
omm3 = _mm_add_ps(omm3,_mm_mul_ps(_mm_set1_ps(a[15]),bmm3));
|
|
#else
|
|
omm0 = _mm_fmadd_ps(_mm_set1_ps(a[1 ]),bmm1,omm0);
|
|
omm1 = _mm_fmadd_ps(_mm_set1_ps(a[5 ]),bmm1,omm1);
|
|
omm2 = _mm_fmadd_ps(_mm_set1_ps(a[9 ]),bmm1,omm2);
|
|
omm3 = _mm_fmadd_ps(_mm_set1_ps(a[13]),bmm1,omm3);
|
|
|
|
omm0 = _mm_fmadd_ps(_mm_set1_ps(a[2 ]),bmm2,omm0);
|
|
omm1 = _mm_fmadd_ps(_mm_set1_ps(a[6 ]),bmm2,omm1);
|
|
omm2 = _mm_fmadd_ps(_mm_set1_ps(a[10]),bmm2,omm2);
|
|
omm3 = _mm_fmadd_ps(_mm_set1_ps(a[14]),bmm2,omm3);
|
|
|
|
omm0 = _mm_fmadd_ps(_mm_set1_ps(a[3 ]),bmm3,omm0);
|
|
omm1 = _mm_fmadd_ps(_mm_set1_ps(a[7 ]),bmm3,omm1);
|
|
omm2 = _mm_fmadd_ps(_mm_set1_ps(a[11]),bmm3,omm2);
|
|
omm3 = _mm_fmadd_ps(_mm_set1_ps(a[15]),bmm3,omm3);
|
|
#endif
|
|
|
|
_mm_storeu_ps (out , omm0);
|
|
_mm_storeu_ps (out+4 , omm1);
|
|
_mm_storeu_ps (out+8 , omm2);
|
|
_mm_storeu_ps (out+12, omm3);
|
|
}
|
|
#endif
|
|
|
|
#ifdef FASTOR_AVX_IMPL
|
|
namespace internal {
|
|
// This is the common interface for 8k8 matmul and not only for 888 so do not
|
|
// make it specific to 888 floats
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M==N && M==8 && std::is_same<T,float>::value,bool>::type>
|
|
FASTOR_INLINE
|
|
void _matmul8k8_float(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m256 out_row0 = VZEROPS;
|
|
__m256 out_row1 = VZEROPS;
|
|
__m256 out_row2 = VZEROPS;
|
|
__m256 out_row3 = VZEROPS;
|
|
__m256 out_row4 = VZEROPS;
|
|
__m256 out_row5 = VZEROPS;
|
|
__m256 out_row6 = VZEROPS;
|
|
__m256 out_row7 = VZEROPS;
|
|
|
|
for (size_t i=0; i<K; ++i) {
|
|
__m256 brow = _mm256_load_ps(&b[i*8]);
|
|
#ifndef FASTOR_FMA_IMPL
|
|
// row 0
|
|
__m256 a_vec0 = _mm256_set1_ps(a[i]);
|
|
out_row0 = _mm256_add_ps(out_row0,_mm256_mul_ps(a_vec0,brow));
|
|
// row 1
|
|
__m256 a_vec1 = _mm256_set1_ps(a[K+i]);
|
|
out_row1 = _mm256_add_ps(out_row1,_mm256_mul_ps(a_vec1,brow));
|
|
// row 2
|
|
__m256 a_vec2 = _mm256_set1_ps(a[2*K+i]);
|
|
out_row2 = _mm256_add_ps(out_row2,_mm256_mul_ps(a_vec2,brow));
|
|
// row 3
|
|
__m256 a_vec3 = _mm256_set1_ps(a[3*K+i]);
|
|
out_row3 = _mm256_add_ps(out_row3,_mm256_mul_ps(a_vec3,brow));
|
|
// row 4
|
|
__m256 a_vec4 = _mm256_set1_ps(a[4*K+i]);
|
|
out_row4 = _mm256_add_ps(out_row4,_mm256_mul_ps(a_vec4,brow));
|
|
// row 5
|
|
__m256 a_vec5 = _mm256_set1_ps(a[5*K+i]);
|
|
out_row5 = _mm256_add_ps(out_row5,_mm256_mul_ps(a_vec5,brow));
|
|
// row 6
|
|
__m256 a_vec6 = _mm256_set1_ps(a[6*K+i]);
|
|
out_row6 = _mm256_add_ps(out_row6,_mm256_mul_ps(a_vec6,brow));
|
|
// row 7
|
|
__m256 a_vec7 = _mm256_set1_ps(a[7*K+i]);
|
|
out_row7 = _mm256_add_ps(out_row7,_mm256_mul_ps(a_vec7,brow));
|
|
#else
|
|
// row 0
|
|
__m256 a_vec0 = _mm256_set1_ps(a[i]);
|
|
out_row0 = _mm256_fmadd_ps(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m256 a_vec1 = _mm256_set1_ps(a[K+i]);
|
|
out_row1 = _mm256_fmadd_ps(a_vec1,brow,out_row1);
|
|
// row 2
|
|
__m256 a_vec2 = _mm256_set1_ps(a[2*K+i]);
|
|
out_row2 = _mm256_fmadd_ps(a_vec2,brow,out_row2);
|
|
// row 3
|
|
__m256 a_vec3 = _mm256_set1_ps(a[3*K+i]);
|
|
out_row3 = _mm256_fmadd_ps(a_vec3,brow,out_row3);
|
|
// row 4
|
|
__m256 a_vec4 = _mm256_set1_ps(a[4*K+i]);
|
|
out_row4 = _mm256_fmadd_ps(a_vec4,brow,out_row4);
|
|
// row 5
|
|
__m256 a_vec5 = _mm256_set1_ps(a[5*K+i]);
|
|
out_row5 = _mm256_fmadd_ps(a_vec5,brow,out_row5);
|
|
// row 6
|
|
__m256 a_vec6 = _mm256_set1_ps(a[6*K+i]);
|
|
out_row6 = _mm256_fmadd_ps(a_vec6,brow,out_row6);
|
|
// row 7
|
|
__m256 a_vec7 = _mm256_set1_ps(a[7*K+i]);
|
|
out_row7 = _mm256_fmadd_ps(a_vec7,brow,out_row7);
|
|
#endif
|
|
}
|
|
_mm256_store_ps(out,out_row0);
|
|
_mm256_store_ps(&out[8],out_row1);
|
|
_mm256_store_ps(&out[16],out_row2);
|
|
_mm256_store_ps(&out[24],out_row3);
|
|
_mm256_store_ps(&out[32],out_row4);
|
|
_mm256_store_ps(&out[40],out_row5);
|
|
_mm256_store_ps(&out[48],out_row6);
|
|
_mm256_store_ps(&out[56],out_row7);
|
|
}
|
|
|
|
} // internal
|
|
|
|
template<>
|
|
FASTOR_INLINE void _matmul<float,8,8,8>(const float * FASTOR_RESTRICT a, const float * FASTOR_RESTRICT b, float * FASTOR_RESTRICT out) {
|
|
internal::_matmul8k8_float<float,8,8,8>(a,b,out);
|
|
return;
|
|
}
|
|
#endif
|
|
|
|
|
|
#ifdef FASTOR_AVX512_IMPL
|
|
namespace internal {
|
|
template<typename T, size_t M, size_t K, size_t N,
|
|
typename std::enable_if<M==N && M==8 && std::is_same<T,double>::value,bool>::type>
|
|
FASTOR_INLINE
|
|
void _matmul8k8_double(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
__m512d out_row0 = _mm512_setzero_pd();
|
|
__m512d out_row1 = _mm512_setzero_pd();
|
|
__m512d out_row2 = _mm512_setzero_pd();
|
|
__m512d out_row3 = _mm512_setzero_pd();
|
|
__m512d out_row4 = _mm512_setzero_pd();
|
|
__m512d out_row5 = _mm512_setzero_pd();
|
|
__m512d out_row6 = _mm512_setzero_pd();
|
|
__m512d out_row7 = _mm512_setzero_pd();
|
|
|
|
for (size_t i=0; i<K; ++i) {
|
|
__m512d brow = _mm512_loadu_pd(&b[i*8]);
|
|
// row 0
|
|
__m512d a_vec0 = _mm512_set1_pd(a[i]);
|
|
out_row0 = _mm512_fmadd_pd(a_vec0,brow,out_row0);
|
|
// row 1
|
|
__m512d a_vec1 = _mm512_set1_pd(a[K+i]);
|
|
out_row1 = _mm512_fmadd_pd(a_vec1,brow,out_row1);
|
|
// row 2
|
|
__m512d a_vec2 = _mm512_set1_pd(a[2*K+i]);
|
|
out_row2 = _mm512_fmadd_pd(a_vec2,brow,out_row2);
|
|
// row 3
|
|
__m512d a_vec3 = _mm512_set1_pd(a[3*K+i]);
|
|
out_row3 = _mm512_fmadd_pd(a_vec3,brow,out_row3);
|
|
// row 4
|
|
__m512d a_vec4 = _mm512_set1_pd(a[4*K+i]);
|
|
out_row4 = _mm512_fmadd_pd(a_vec4,brow,out_row4);
|
|
// row 5
|
|
__m512d a_vec5 = _mm512_set1_pd(a[5*K+i]);
|
|
out_row5 = _mm512_fmadd_pd(a_vec5,brow,out_row5);
|
|
// row 6
|
|
__m512d a_vec6 = _mm512_set1_pd(a[6*K+i]);
|
|
out_row6 = _mm512_fmadd_pd(a_vec6,brow,out_row6);
|
|
// row 7
|
|
__m512d a_vec7 = _mm512_set1_pd(a[7*K+i]);
|
|
out_row7 = _mm512_fmadd_pd(a_vec7,brow,out_row7);
|
|
}
|
|
_mm512_storeu_pd(out,out_row0);
|
|
_mm512_storeu_pd(&out[8],out_row1);
|
|
_mm512_storeu_pd(&out[16],out_row2);
|
|
_mm512_storeu_pd(&out[24],out_row3);
|
|
_mm512_storeu_pd(&out[32],out_row4);
|
|
_mm512_storeu_pd(&out[40],out_row5);
|
|
_mm512_storeu_pd(&out[48],out_row6);
|
|
_mm512_storeu_pd(&out[56],out_row7);
|
|
}
|
|
} // internal
|
|
|
|
template<>
|
|
FASTOR_INLINE void _matmul<double,8,8,8>(const double * FASTOR_RESTRICT a, const double * FASTOR_RESTRICT b, double * FASTOR_RESTRICT out) {
|
|
internal::_matmul8k8_double<double,8,8,8>(a,b,out);
|
|
return;
|
|
}
|
|
#endif
|
|
|
|
|
|
#ifdef FASTOR_SSE2_IMPL
|
|
template<>
|
|
FASTOR_INLINE
|
|
void _matmul<double,2,2,2>(const double * FASTOR_RESTRICT a, const double * FASTOR_RESTRICT b, double * FASTOR_RESTRICT out) {
|
|
|
|
// Ivy 26 - HW 20 - Sky 16 OPS
|
|
|
|
__m128d bmm0 = _mm_loadu_pd (b );
|
|
__m128d bmm1 = _mm_loadu_pd (b+2);
|
|
|
|
__m128d omm0(_mm_mul_pd(_mm_set1_pd(a[0]),bmm0));
|
|
__m128d omm1(_mm_mul_pd(_mm_set1_pd(a[2]),bmm0));
|
|
#ifndef FASTOR_FMA_IMPL
|
|
omm0 = _mm_add_pd(omm0,_mm_mul_pd(_mm_set1_pd(a[1]),bmm1));
|
|
omm1 = _mm_add_pd(omm1,_mm_mul_pd(_mm_set1_pd(a[3]),bmm1));
|
|
#else
|
|
omm0 = _mm_fmadd_pd(_mm_set1_pd(a[1]),bmm1,omm0);
|
|
omm1 = _mm_fmadd_pd(_mm_set1_pd(a[3]),bmm1,omm1);
|
|
#endif
|
|
_mm_storeu_pd (out , omm0);
|
|
_mm_storeu_pd (out+2, omm1);
|
|
}
|
|
#endif
|
|
|
|
#ifdef FASTOR_AVX_IMPL
|
|
template<>
|
|
FASTOR_INLINE
|
|
void _matmul<double,3,3,3>(const double * FASTOR_RESTRICT a, const double * FASTOR_RESTRICT b, double * FASTOR_RESTRICT out) {
|
|
|
|
// Ivy 63 - HW 45 - Sky 36 OPS
|
|
// Fully vectorised approach that reduces
|
|
// (27 scalar mul + 18 scalar add) to
|
|
// (9 AVX mul + 6 AVX add) [non-FMA]
|
|
// (3 AVX mul + 6 AVX FMA) [ FMA]
|
|
|
|
__m256d bmm0 = _mm256_loadu_pd (b );
|
|
__m256d bmm1 = _mm256_loadu_pd (b+3);
|
|
__m256d bmm2 = _mm256_loadul3_pd(b+6);
|
|
|
|
__m256d omm0(_mm256_mul_pd(_mm256_set1_pd(a[0]),bmm0));
|
|
__m256d omm1(_mm256_mul_pd(_mm256_set1_pd(a[3]),bmm0));
|
|
__m256d omm2(_mm256_mul_pd(_mm256_set1_pd(a[6]),bmm0));
|
|
#ifndef FASTOR_FMA_IMPL
|
|
omm0 = _mm256_add_pd(omm0,_mm256_mul_pd(_mm256_set1_pd(a[1]),bmm1));
|
|
omm1 = _mm256_add_pd(omm1,_mm256_mul_pd(_mm256_set1_pd(a[4]),bmm1));
|
|
omm2 = _mm256_add_pd(omm2,_mm256_mul_pd(_mm256_set1_pd(a[7]),bmm1));
|
|
|
|
omm0 = _mm256_add_pd(omm0,_mm256_mul_pd(_mm256_set1_pd(a[2]),bmm2));
|
|
omm1 = _mm256_add_pd(omm1,_mm256_mul_pd(_mm256_set1_pd(a[5]),bmm2));
|
|
omm2 = _mm256_add_pd(omm2,_mm256_mul_pd(_mm256_set1_pd(a[8]),bmm2));
|
|
#else
|
|
omm0 = _mm256_fmadd_pd(_mm256_set1_pd(a[1]),bmm1,omm0);
|
|
omm1 = _mm256_fmadd_pd(_mm256_set1_pd(a[4]),bmm1,omm1);
|
|
omm2 = _mm256_fmadd_pd(_mm256_set1_pd(a[7]),bmm1,omm2);
|
|
|
|
omm0 = _mm256_fmadd_pd(_mm256_set1_pd(a[2]),bmm2,omm0);
|
|
omm1 = _mm256_fmadd_pd(_mm256_set1_pd(a[5]),bmm2,omm1);
|
|
omm2 = _mm256_fmadd_pd(_mm256_set1_pd(a[8]),bmm2,omm2);
|
|
#endif
|
|
|
|
_mm256_storeu_pd (out , omm0);
|
|
_mm256_storeu_pd (out+3, omm1);
|
|
_mm256_storeul3_pd(out+6, omm2);
|
|
}
|
|
|
|
template<>
|
|
FASTOR_INLINE
|
|
void _matmul<double,4,4,4>(const double * FASTOR_RESTRICT a, const double * FASTOR_RESTRICT b, double * FASTOR_RESTRICT out) {
|
|
|
|
// Ivy 116 - HW 80 - Sky 64 OPS
|
|
// Fully vectorised approach that reduces
|
|
// (64 scalar mul + 48 scalar add) to
|
|
// (16 SSE mul + 12 SEE add) [non-FMA]
|
|
// (4 SSE mul + 12 SEE FMA) [ FMA]
|
|
|
|
__m256d bmm0 = _mm256_loadu_pd (b );
|
|
__m256d bmm1 = _mm256_loadu_pd (b+4 );
|
|
__m256d bmm2 = _mm256_loadu_pd (b+8 );
|
|
__m256d bmm3 = _mm256_loadu_pd (b+12);
|
|
|
|
__m256d omm0(_mm256_mul_pd(_mm256_set1_pd(a[0 ]),bmm0));
|
|
__m256d omm1(_mm256_mul_pd(_mm256_set1_pd(a[4 ]),bmm0));
|
|
__m256d omm2(_mm256_mul_pd(_mm256_set1_pd(a[8 ]),bmm0));
|
|
__m256d omm3(_mm256_mul_pd(_mm256_set1_pd(a[12]),bmm0));
|
|
#ifndef FASTOR_FMA_IMPL
|
|
omm0 = _mm256_add_pd(omm0,_mm256_mul_pd(_mm256_set1_pd(a[1 ]),bmm1));
|
|
omm1 = _mm256_add_pd(omm1,_mm256_mul_pd(_mm256_set1_pd(a[5 ]),bmm1));
|
|
omm2 = _mm256_add_pd(omm2,_mm256_mul_pd(_mm256_set1_pd(a[9 ]),bmm1));
|
|
omm3 = _mm256_add_pd(omm3,_mm256_mul_pd(_mm256_set1_pd(a[13]),bmm1));
|
|
|
|
omm0 = _mm256_add_pd(omm0,_mm256_mul_pd(_mm256_set1_pd(a[2 ]),bmm2));
|
|
omm1 = _mm256_add_pd(omm1,_mm256_mul_pd(_mm256_set1_pd(a[6 ]),bmm2));
|
|
omm2 = _mm256_add_pd(omm2,_mm256_mul_pd(_mm256_set1_pd(a[10]),bmm2));
|
|
omm3 = _mm256_add_pd(omm3,_mm256_mul_pd(_mm256_set1_pd(a[14]),bmm2));
|
|
|
|
omm0 = _mm256_add_pd(omm0,_mm256_mul_pd(_mm256_set1_pd(a[3 ]),bmm3));
|
|
omm1 = _mm256_add_pd(omm1,_mm256_mul_pd(_mm256_set1_pd(a[7 ]),bmm3));
|
|
omm2 = _mm256_add_pd(omm2,_mm256_mul_pd(_mm256_set1_pd(a[11]),bmm3));
|
|
omm3 = _mm256_add_pd(omm3,_mm256_mul_pd(_mm256_set1_pd(a[15]),bmm3));
|
|
#else
|
|
omm0 = _mm256_fmadd_pd(_mm256_set1_pd(a[1 ]),bmm1,omm0);
|
|
omm1 = _mm256_fmadd_pd(_mm256_set1_pd(a[5 ]),bmm1,omm1);
|
|
omm2 = _mm256_fmadd_pd(_mm256_set1_pd(a[9 ]),bmm1,omm2);
|
|
omm3 = _mm256_fmadd_pd(_mm256_set1_pd(a[13]),bmm1,omm3);
|
|
|
|
omm0 = _mm256_fmadd_pd(_mm256_set1_pd(a[2 ]),bmm2,omm0);
|
|
omm1 = _mm256_fmadd_pd(_mm256_set1_pd(a[6 ]),bmm2,omm1);
|
|
omm2 = _mm256_fmadd_pd(_mm256_set1_pd(a[10]),bmm2,omm2);
|
|
omm3 = _mm256_fmadd_pd(_mm256_set1_pd(a[14]),bmm2,omm3);
|
|
|
|
omm0 = _mm256_fmadd_pd(_mm256_set1_pd(a[3 ]),bmm3,omm0);
|
|
omm1 = _mm256_fmadd_pd(_mm256_set1_pd(a[7 ]),bmm3,omm1);
|
|
omm2 = _mm256_fmadd_pd(_mm256_set1_pd(a[11]),bmm3,omm2);
|
|
omm3 = _mm256_fmadd_pd(_mm256_set1_pd(a[15]),bmm3,omm3);
|
|
#endif
|
|
|
|
_mm256_storeu_pd (out , omm0);
|
|
_mm256_storeu_pd (out+4 , omm1);
|
|
_mm256_storeu_pd (out+8 , omm2);
|
|
_mm256_storeu_pd (out+12, omm3);
|
|
}
|
|
#endif
|
|
|
|
|
|
|
|
|
|
|
|
//! Matrix-vector multiplication
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
// Don't call this function directly as it's name is unconventional
|
|
// It gets called from within matmul anyway so always call matmul
|
|
namespace internal {
|
|
|
|
template<typename T, size_t M, size_t N>
|
|
FASTOR_INLINE
|
|
void _matvecmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
using V = typename choose_best_simd_type<SIMDVector<T,DEFAULT_ABI>,N>::type;
|
|
constexpr size_t unrollOuterloop = 8UL;
|
|
constexpr size_t M0 = M / unrollOuterloop * unrollOuterloop;
|
|
|
|
// Number of columns of c (N) that can be safely unrolled with V::Size
|
|
constexpr size_t N1 = N / V::Size * V::Size;
|
|
|
|
constexpr bool isAAligned = false;
|
|
constexpr bool isBAligned = false;
|
|
|
|
size_t i=0;
|
|
FASTOR_IF_CONSTEXPR(N < V::Size) {
|
|
|
|
for (; i<ROUND_DOWN(M,2UL); i+=2UL) {
|
|
V omm0, omm1;
|
|
size_t j = 0;
|
|
for (; j< N1; j+=V::Size) {
|
|
V amm0(&a[(i )*N+j],false);
|
|
V amm1(&a[(i+1UL)*N+j],false);
|
|
V bmm0(&b[j],false);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
}
|
|
T out_s0 = 0;
|
|
T out_s1 = 0;
|
|
for (; j< N; ++j) {
|
|
out_s0 += a[i*N+j]*b[j];
|
|
out_s1 += a[(i+1)*N+j]*b[j];
|
|
}
|
|
out[i] = omm0.sum() + out_s0;
|
|
out[i+1] = omm1.sum() + out_s1;
|
|
}
|
|
|
|
for (; i<M; ++i) {
|
|
V omm0;
|
|
size_t j = 0;
|
|
for (; j< N1; j+=V::Size) {
|
|
const V amm0(&a[i*N+j],false);
|
|
const V bmm0(&b[j],false);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
}
|
|
T out_s0 = 0;
|
|
for (; j< N; j+=1) {
|
|
out_s0 += a[i*N+j]*b[j];
|
|
}
|
|
out[i] = omm0.sum() + out_s0;
|
|
}
|
|
return;
|
|
}
|
|
|
|
else
|
|
{
|
|
|
|
for (; i<M0; i+=unrollOuterloop) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
V omm1(V(&a[(i+1UL)*N],isAAligned)*bmm0);
|
|
V omm2(V(&a[(i+2UL)*N],isAAligned)*bmm0);
|
|
V omm3(V(&a[(i+3UL)*N],isAAligned)*bmm0);
|
|
V omm4(V(&a[(i+4UL)*N],isAAligned)*bmm0);
|
|
V omm5(V(&a[(i+5UL)*N],isAAligned)*bmm0);
|
|
V omm6(V(&a[(i+6UL)*N],isAAligned)*bmm0);
|
|
V omm7(V(&a[(i+7UL)*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
const V amm1(&a[(i+1UL)*N+j],isAAligned);
|
|
const V amm2(&a[(i+2UL)*N+j],isAAligned);
|
|
const V amm3(&a[(i+3UL)*N+j],isAAligned);
|
|
const V amm4(&a[(i+4UL)*N+j],isAAligned);
|
|
const V amm5(&a[(i+5UL)*N+j],isAAligned);
|
|
const V amm6(&a[(i+6UL)*N+j],isAAligned);
|
|
const V amm7(&a[(i+7UL)*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
omm2 = fmadd(amm2,bmm0,omm2);
|
|
omm3 = fmadd(amm3,bmm0,omm3);
|
|
omm4 = fmadd(amm4,bmm0,omm4);
|
|
omm5 = fmadd(amm5,bmm0,omm5);
|
|
omm6 = fmadd(amm6,bmm0,omm6);
|
|
omm7 = fmadd(amm7,bmm0,omm7);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
out[i+1UL] = omm1.sum();
|
|
out[i+2UL] = omm2.sum();
|
|
out[i+3UL] = omm3.sum();
|
|
out[i+4UL] = omm4.sum();
|
|
out[i+5UL] = omm5.sum();
|
|
out[i+6UL] = omm6.sum();
|
|
out[i+7UL] = omm7.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
out[i+1UL] += a[(i+1UL)*N+j]*bmm0;
|
|
out[i+2UL] += a[(i+2UL)*N+j]*bmm0;
|
|
out[i+3UL] += a[(i+3UL)*N+j]*bmm0;
|
|
out[i+4UL] += a[(i+4UL)*N+j]*bmm0;
|
|
out[i+5UL] += a[(i+5UL)*N+j]*bmm0;
|
|
out[i+6UL] += a[(i+6UL)*N+j]*bmm0;
|
|
out[i+7UL] += a[(i+7UL)*N+j]*bmm0;
|
|
}
|
|
}
|
|
|
|
FASTOR_IF_CONSTEXPR (M - M0 == 7) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
V omm1(V(&a[(i+1UL)*N],isAAligned)*bmm0);
|
|
V omm2(V(&a[(i+2UL)*N],isAAligned)*bmm0);
|
|
V omm3(V(&a[(i+3UL)*N],isAAligned)*bmm0);
|
|
V omm4(V(&a[(i+4UL)*N],isAAligned)*bmm0);
|
|
V omm5(V(&a[(i+5UL)*N],isAAligned)*bmm0);
|
|
V omm6(V(&a[(i+6UL)*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
const V amm1(&a[(i+1UL)*N+j],isAAligned);
|
|
const V amm2(&a[(i+2UL)*N+j],isAAligned);
|
|
const V amm3(&a[(i+3UL)*N+j],isAAligned);
|
|
const V amm4(&a[(i+4UL)*N+j],isAAligned);
|
|
const V amm5(&a[(i+5UL)*N+j],isAAligned);
|
|
const V amm6(&a[(i+6UL)*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
omm2 = fmadd(amm2,bmm0,omm2);
|
|
omm3 = fmadd(amm3,bmm0,omm3);
|
|
omm4 = fmadd(amm4,bmm0,omm4);
|
|
omm5 = fmadd(amm5,bmm0,omm5);
|
|
omm6 = fmadd(amm6,bmm0,omm6);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
out[i+1UL] = omm1.sum();
|
|
out[i+2UL] = omm2.sum();
|
|
out[i+3UL] = omm3.sum();
|
|
out[i+4UL] = omm4.sum();
|
|
out[i+5UL] = omm5.sum();
|
|
out[i+6UL] = omm6.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
out[i+1UL] += a[(i+1UL)*N+j]*bmm0;
|
|
out[i+2UL] += a[(i+2UL)*N+j]*bmm0;
|
|
out[i+3UL] += a[(i+3UL)*N+j]*bmm0;
|
|
out[i+4UL] += a[(i+4UL)*N+j]*bmm0;
|
|
out[i+5UL] += a[(i+5UL)*N+j]*bmm0;
|
|
out[i+6UL] += a[(i+6UL)*N+j]*bmm0;
|
|
}
|
|
}
|
|
|
|
else FASTOR_IF_CONSTEXPR (M - M0 == 6) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
V omm1(V(&a[(i+1UL)*N],isAAligned)*bmm0);
|
|
V omm2(V(&a[(i+2UL)*N],isAAligned)*bmm0);
|
|
V omm3(V(&a[(i+3UL)*N],isAAligned)*bmm0);
|
|
V omm4(V(&a[(i+4UL)*N],isAAligned)*bmm0);
|
|
V omm5(V(&a[(i+5UL)*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
const V amm1(&a[(i+1UL)*N+j],isAAligned);
|
|
const V amm2(&a[(i+2UL)*N+j],isAAligned);
|
|
const V amm3(&a[(i+3UL)*N+j],isAAligned);
|
|
const V amm4(&a[(i+4UL)*N+j],isAAligned);
|
|
const V amm5(&a[(i+5UL)*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
omm2 = fmadd(amm2,bmm0,omm2);
|
|
omm3 = fmadd(amm3,bmm0,omm3);
|
|
omm4 = fmadd(amm4,bmm0,omm4);
|
|
omm5 = fmadd(amm5,bmm0,omm5);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
out[i+1UL] = omm1.sum();
|
|
out[i+2UL] = omm2.sum();
|
|
out[i+3UL] = omm3.sum();
|
|
out[i+4UL] = omm4.sum();
|
|
out[i+5UL] = omm5.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
out[i+1UL] += a[(i+1UL)*N+j]*bmm0;
|
|
out[i+2UL] += a[(i+2UL)*N+j]*bmm0;
|
|
out[i+3UL] += a[(i+3UL)*N+j]*bmm0;
|
|
out[i+4UL] += a[(i+4UL)*N+j]*bmm0;
|
|
out[i+5UL] += a[(i+5UL)*N+j]*bmm0;
|
|
}
|
|
}
|
|
|
|
else FASTOR_IF_CONSTEXPR (M - M0 == 5) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
V omm1(V(&a[(i+1UL)*N],isAAligned)*bmm0);
|
|
V omm2(V(&a[(i+2UL)*N],isAAligned)*bmm0);
|
|
V omm3(V(&a[(i+3UL)*N],isAAligned)*bmm0);
|
|
V omm4(V(&a[(i+4UL)*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
const V amm1(&a[(i+1UL)*N+j],isAAligned);
|
|
const V amm2(&a[(i+2UL)*N+j],isAAligned);
|
|
const V amm3(&a[(i+3UL)*N+j],isAAligned);
|
|
const V amm4(&a[(i+4UL)*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
omm2 = fmadd(amm2,bmm0,omm2);
|
|
omm3 = fmadd(amm3,bmm0,omm3);
|
|
omm4 = fmadd(amm4,bmm0,omm4);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
out[i+1UL] = omm1.sum();
|
|
out[i+2UL] = omm2.sum();
|
|
out[i+3UL] = omm3.sum();
|
|
out[i+4UL] = omm4.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
out[i+1UL] += a[(i+1UL)*N+j]*bmm0;
|
|
out[i+2UL] += a[(i+2UL)*N+j]*bmm0;
|
|
out[i+3UL] += a[(i+3UL)*N+j]*bmm0;
|
|
out[i+4UL] += a[(i+4UL)*N+j]*bmm0;
|
|
}
|
|
}
|
|
|
|
else FASTOR_IF_CONSTEXPR (M - M0 == 4) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
V omm1(V(&a[(i+1UL)*N],isAAligned)*bmm0);
|
|
V omm2(V(&a[(i+2UL)*N],isAAligned)*bmm0);
|
|
V omm3(V(&a[(i+3UL)*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
const V amm1(&a[(i+1UL)*N+j],isAAligned);
|
|
const V amm2(&a[(i+2UL)*N+j],isAAligned);
|
|
const V amm3(&a[(i+3UL)*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
omm2 = fmadd(amm2,bmm0,omm2);
|
|
omm3 = fmadd(amm3,bmm0,omm3);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
out[i+1UL] = omm1.sum();
|
|
out[i+2UL] = omm2.sum();
|
|
out[i+3UL] = omm3.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
out[i+1UL] += a[(i+1UL)*N+j]*bmm0;
|
|
out[i+2UL] += a[(i+2UL)*N+j]*bmm0;
|
|
out[i+3UL] += a[(i+3UL)*N+j]*bmm0;
|
|
}
|
|
}
|
|
|
|
else FASTOR_IF_CONSTEXPR (M - M0 == 3) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
V omm1(V(&a[(i+1UL)*N],isAAligned)*bmm0);
|
|
V omm2(V(&a[(i+2UL)*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
const V amm1(&a[(i+1UL)*N+j],isAAligned);
|
|
const V amm2(&a[(i+2UL)*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
omm2 = fmadd(amm2,bmm0,omm2);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
out[i+1UL] = omm1.sum();
|
|
out[i+2UL] = omm2.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
out[i+1UL] += a[(i+1UL)*N+j]*bmm0;
|
|
out[i+2UL] += a[(i+2UL)*N+j]*bmm0;
|
|
}
|
|
}
|
|
|
|
else FASTOR_IF_CONSTEXPR (M - M0 == 2) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
V omm1(V(&a[(i+1UL)*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
const V amm1(&a[(i+1UL)*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
omm1 = fmadd(amm1,bmm0,omm1);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
out[i+1UL] = omm1.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
out[i+1UL] += a[(i+1UL)*N+j]*bmm0;
|
|
}
|
|
}
|
|
|
|
else FASTOR_IF_CONSTEXPR (M - M0 == 1) {
|
|
|
|
const V bmm0(&b[0],isBAligned);
|
|
|
|
V omm0(V(&a[(i )*N],isAAligned)*bmm0);
|
|
|
|
size_t j = V::Size;
|
|
for (; j< N1; j+=V::Size) {
|
|
|
|
const V bmm0(&b[j],isBAligned);
|
|
|
|
const V amm0(&a[(i )*N+j],isAAligned);
|
|
|
|
omm0 = fmadd(amm0,bmm0,omm0);
|
|
}
|
|
|
|
out[i ] = omm0.sum();
|
|
|
|
for (; j< N; ++j) {
|
|
const T bmm0(b[j]);
|
|
out[i ] += a[(i )*N+j]*bmm0;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
|
|
#ifdef FASTOR_SSE2_IMPL
|
|
template<>
|
|
FASTOR_INLINE void _matmul<float,2,2,1>(const float * FASTOR_RESTRICT a, const float * FASTOR_RESTRICT b, float * FASTOR_RESTRICT out) {
|
|
// 11 OPS
|
|
__m128 amm = _mm_loadu_ps(a);
|
|
__m128 bmm = ZEROPS;
|
|
bmm = _mm_loadl_pi(bmm, (__m64*)b);
|
|
|
|
__m128 res = _mm_mul_ps(amm,_mm_movelh_ps(bmm,bmm));
|
|
__m128 res2 = _mm_shuffle_ps(res ,res ,_MM_SHUFFLE(2,3,0,1));
|
|
res2 = _mm_add_ps(res,res2);
|
|
__m128 res3 = _mm_shuffle_ps(res2,res2,_MM_SHUFFLE(3,1,2,0));
|
|
_mm_storel_pi((__m64*) out,res3);
|
|
}
|
|
|
|
template<>
|
|
FASTOR_INLINE void _matmul<float,3,3,1>(const float * FASTOR_RESTRICT a, const float * FASTOR_RESTRICT b, float * FASTOR_RESTRICT out) {
|
|
// IVY/HW 47 OPS
|
|
// Mask loads on AVX and AVX512 otherwise 3 single loads each
|
|
__m128 amm0 = _mm_loadul3_ps(a);
|
|
__m128 amm1 = _mm_loadul3_ps(&a[3]);
|
|
__m128 amm2 = _mm_loadul3_ps(&a[6]);
|
|
__m128 bmm = _mm_loadul3_ps(b);
|
|
|
|
out[0] = _mm_sum_ps(_mm_mul_ps(amm0,bmm));
|
|
out[1] = _mm_sum_ps(_mm_mul_ps(amm1,bmm));
|
|
out[2] = _mm_sum_ps(_mm_mul_ps(amm2,bmm));
|
|
}
|
|
|
|
template<>
|
|
FASTOR_INLINE void _matmul<double,2,2,1>(const double * FASTOR_RESTRICT a, const double * FASTOR_RESTRICT b, double * FASTOR_RESTRICT out) {
|
|
// IVY/HW/SKY 19 OPS
|
|
__m128d amm0 = _mm_loadu_pd(a );
|
|
__m128d amm1 = _mm_loadu_pd(a+2);
|
|
__m128d bmm0 = _mm_loadu_pd(b );
|
|
|
|
__m128d omm0 = _mm_mul_pd(amm0,bmm0);
|
|
__m128d omm1 = _mm_mul_pd(amm1,bmm0);
|
|
omm0 = _mm_add_pd(omm0, _mm_reverse_pd(omm0));
|
|
omm1 = _mm_add_pd(omm1, _mm_reverse_pd(omm1));
|
|
|
|
_mm_storeu_pd(out,_mm_shuffle_pd(omm0,omm1,0x1));
|
|
}
|
|
#endif
|
|
|
|
|
|
#ifdef FASTOR_AVX_IMPL
|
|
template<>
|
|
FASTOR_INLINE void _matmul<double,3,3,1>(const double * FASTOR_RESTRICT a, const double * FASTOR_RESTRICT b, double * FASTOR_RESTRICT out) {
|
|
// SKY 49 - ICE 87 OPS
|
|
// Mask loads on AVX and AVX512
|
|
__m256d amm0 = _mm256_loadul3_pd(a);
|
|
__m256d amm1 = _mm256_loadul3_pd(&a[3]);
|
|
__m256d amm2 = _mm256_loadul3_pd(&a[6]);
|
|
__m256d bmm = _mm256_loadul3_pd(b);
|
|
|
|
out[0] = _mm256_sum_pd(_mm256_mul_pd(amm0,bmm));
|
|
out[1] = _mm256_sum_pd(_mm256_mul_pd(amm1,bmm));
|
|
out[2] = _mm256_sum_pd(_mm256_mul_pd(amm2,bmm));
|
|
}
|
|
#endif
|
|
|
|
#endif // MATMUL_SPECIALISATIONS_KERNELS_H
|