804 lines
34 KiB
C++
804 lines
34 KiB
C++
#ifndef MATMUL_KERNELS2_H
|
|
#define MATMUL_KERNELS2_H
|
|
|
|
|
|
#include "Fastor/config/config.h"
|
|
#include "Fastor/simd_vector/extintrin.h"
|
|
#include "Fastor/simd_vector/SIMDVector.h"
|
|
#include "Fastor/meta/meta.h"
|
|
#include "Fastor/meta/tensor_meta.h"
|
|
|
|
|
|
namespace Fastor {
|
|
|
|
namespace internal {
|
|
|
|
// TRMM implementation of Fastor - matrix-matrix multiplication when either or both operands are
|
|
// lower or upper triangular. The matrices do not need to be square and trapezoidal cases are also
|
|
// covered. For big matrices the speed-up is 2X or even better over matmul for when one operand is
|
|
// is triangular and nearly 4X for when both operands are triangular.
|
|
// For small matrices due to aggressive unrolling for SIMD the matrices cannot be exactly traversed
|
|
// within their triangular part(s) and a bit of the non-triangular part(s) need(s) to be loaded as well
|
|
// hence, the performance may not be exactly 2X/4X over the general matmul case
|
|
|
|
|
|
// The functions here are exact replica of those in matmul_kernels.h and will be eventually
|
|
// merged together as these variants have no associated overhead for the general case of matmul
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
/*
|
|
|
|
For triangular matmul only the iteration span of K is modified
|
|
using the following logic [lt = Lower, ut=Upper]
|
|
|
|
// if lhs == lt
|
|
const size_t kfirst = 0;
|
|
const size_t klast = min(i+1,K); // or min(i+unrollOuterloop,K);
|
|
|
|
// if lhs == ut
|
|
const size_t kfirst = i;
|
|
const size_t klast = K;
|
|
|
|
// if rhs == lt
|
|
const size_t kfirst = j;
|
|
const size_t klast = K;
|
|
|
|
// if rhs == ut
|
|
const size_t kfirst = 0;
|
|
const size_t klast = min(j+1,K); // or min(j+unrollOuterloop,K);
|
|
|
|
|
|
// both lower
|
|
const size_t kfirst = j;
|
|
const size_t klast = min(i+1,K); // or min(i+unrollOuterloop,K);
|
|
|
|
// if lhs == lt && rhs == ut
|
|
const size_t kfirst = 0;
|
|
const size_t klast = min(min(i+1,j+1),K); // or min(min(i+unrollOuterloop,j+unrollInnerloop),K);
|
|
|
|
// if lhs == ut && rhs == lt
|
|
const size_t kfirst = max(i,j);
|
|
const size_t klast = K;
|
|
|
|
// if both upper
|
|
const size_t kfirst = i;
|
|
const size_t klast = min(j+1,K); // or min(j+unrollInnerloop,K);
|
|
|
|
*/
|
|
|
|
template<typename T, T K, T unrollOuterloop=1,T unrollInnerloop=1, typename LhsType = UpLoType::General, typename RhsType = UpLoType::General>
|
|
constexpr FASTOR_INLINE T find_kfirst(const T i, const T j) {
|
|
return is_same_v_<LhsType,UpLoType::Lower> || is_same_v_<LhsType,UpLoType::General> ?
|
|
( is_same_v_<RhsType,UpLoType::Lower> ? j : 0UL ) :
|
|
(is_same_v_<LhsType,UpLoType::Upper> ? ( is_same_v_<RhsType,UpLoType::Lower> ? internal::max_(i,j) : i ) : 0UL );
|
|
}
|
|
template<typename T, T K, T unrollOuterloop=1,T unrollInnerloop=1, typename LhsType = UpLoType::General, typename RhsType = UpLoType::General>
|
|
constexpr FASTOR_INLINE T find_klast(const T i, const T j) {
|
|
return is_same_v_<LhsType,UpLoType::Lower> ?
|
|
( is_same_v_<RhsType,UpLoType::Upper> ? internal::min_(internal::min_(i+unrollOuterloop,j+unrollInnerloop),K) : internal::min_(i+unrollOuterloop,K) ) :
|
|
(is_same_v_<LhsType,UpLoType::Upper> || (is_same_v_<LhsType,UpLoType::General>) ?
|
|
( is_same_v_<RhsType,UpLoType::Upper> ? internal::min_(j+unrollInnerloop,K) : K ) : K );
|
|
}
|
|
|
|
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
|
|
// A set of helper functions for the inner blocks of matmul. Almost all compilers (GCC/CLang/Intel)
|
|
// unroll the inner-most loop (on unrollOuterloop)
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
template<typename T, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==1,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
V c_ij[unrollOuterloop*numSIMDCols];
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
|
|
const V bmm0(&b[k*N+j],false);
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const V amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].store(&c[(i+ii*unrollOuterloop+n)*N+j],false);
|
|
}
|
|
}
|
|
}
|
|
|
|
template<typename T, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==2,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
V c_ij[unrollOuterloop*numSIMDCols];
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
|
|
const V bmm0(&b[k*N+j],false);
|
|
const V bmm1(&b[k*N+j+V::Size],false);
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const V amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
c_ij[n+unrollOuterloop] = fmadd(amm0,bmm1,c_ij[n+unrollOuterloop]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].store(&c[(i+ii*unrollOuterloop+n)*N+j],false);
|
|
c_ij[n+unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+V::Size],false);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
template<typename T, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==3,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
V c_ij[unrollOuterloop*numSIMDCols];
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
|
|
const V bmm0(&b[k*N+j],false);
|
|
const V bmm1(&b[k*N+j+V::Size],false);
|
|
const V bmm2(&b[k*N+j+2*V::Size],false);
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const V amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
c_ij[n+unrollOuterloop] = fmadd(amm0,bmm1,c_ij[n+unrollOuterloop]);
|
|
c_ij[n+2*unrollOuterloop] = fmadd(amm0,bmm2,c_ij[n+2*unrollOuterloop]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].store(&c[(i+ii*unrollOuterloop+n)*N+j],false);
|
|
c_ij[n+unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+V::Size],false);
|
|
c_ij[n+2*unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+2*V::Size],false);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
template<typename T, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==4,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
V c_ij[unrollOuterloop*numSIMDCols];
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
|
|
const V bmm0(&b[k*N+j],false);
|
|
const V bmm1(&b[k*N+j+V::Size],false);
|
|
const V bmm2(&b[k*N+j+2*V::Size],false);
|
|
const V bmm3(&b[k*N+j+3*V::Size],false);
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const V amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
c_ij[n+unrollOuterloop] = fmadd(amm0,bmm1,c_ij[n+unrollOuterloop]);
|
|
c_ij[n+2*unrollOuterloop] = fmadd(amm0,bmm2,c_ij[n+2*unrollOuterloop]);
|
|
c_ij[n+3*unrollOuterloop] = fmadd(amm0,bmm3,c_ij[n+3*unrollOuterloop]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].store(&c[(i+ii*unrollOuterloop+n)*N+j],false);
|
|
c_ij[n+unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+V::Size],false);
|
|
c_ij[n+2*unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+2*V::Size],false);
|
|
c_ij[n+3*unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+3*V::Size],false);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
template<typename T, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==5,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
V c_ij[unrollOuterloop*numSIMDCols];
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
|
|
const V bmm0(&b[k*N+j],false);
|
|
const V bmm1(&b[k*N+j+V::Size],false);
|
|
const V bmm2(&b[k*N+j+2*V::Size],false);
|
|
const V bmm3(&b[k*N+j+3*V::Size],false);
|
|
const V bmm4(&b[k*N+j+4*V::Size],false);
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const V amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
c_ij[n+unrollOuterloop] = fmadd(amm0,bmm1,c_ij[n+unrollOuterloop]);
|
|
c_ij[n+2*unrollOuterloop] = fmadd(amm0,bmm2,c_ij[n+2*unrollOuterloop]);
|
|
c_ij[n+3*unrollOuterloop] = fmadd(amm0,bmm3,c_ij[n+3*unrollOuterloop]);
|
|
c_ij[n+4*unrollOuterloop] = fmadd(amm0,bmm3,c_ij[n+4*unrollOuterloop]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].store(&c[(i+ii*unrollOuterloop+n)*N+j],false);
|
|
c_ij[n+unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+V::Size],false);
|
|
c_ij[n+2*unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+2*V::Size],false);
|
|
c_ij[n+3*unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+3*V::Size],false);
|
|
c_ij[n+4*unrollOuterloop].store(&c[(i+ii*unrollOuterloop+n)*N+j+4*V::Size],false);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
template<typename T, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==1,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_scalar_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
T c_ij[unrollOuterloop*numSIMDCols] = {};
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
const T bmm0(b[k*N+j]);
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const T amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] += amm0*bmm0;
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c[(i+ii*unrollOuterloop+n)*N+j] = c_ij[n];
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
template<typename T, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==1,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_mask_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j, const int (&maska)[V::Size]) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
V c_ij[unrollOuterloop*numSIMDCols];
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
|
|
const V bmm0(maskload<V>(&b[k*N+j],maska));
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const V amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
maskstore(&c[(i+ii*unrollOuterloop+n)*N+j],maska,c_ij[n]);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
template<typename T, typename MaskT, typename V, size_t M, size_t K, size_t N, size_t unrollOuterloop, size_t numSIMDRows, size_t numSIMDCols,
|
|
typename LhsType = UpLoType::General, typename RhsType = UpLoType::General,
|
|
typename std::enable_if<numSIMDCols==1,bool>::type = false>
|
|
FASTOR_INLINE
|
|
void interior_block_tmatmul_mask_impl(
|
|
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
|
const size_t i, const size_t j, const MaskT mask) {
|
|
|
|
V bmm0;
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop*numSIMDRows,numSIMDCols*V::Size,LhsType,RhsType>(i,j);
|
|
|
|
for (size_t ii = 0; ii < numSIMDRows; ++ii) {
|
|
|
|
V c_ij[unrollOuterloop*numSIMDCols];
|
|
// Loop over columns of a (rows of b)
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
|
|
bmm0.mask_load(&b[k*N+j],mask,false);
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
|
|
const V amm0 = a[(i+ii*unrollOuterloop+n)*K+k];
|
|
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
}
|
|
}
|
|
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].mask_store(&c[(i+ii*unrollOuterloop+n)*N+j],mask,false);
|
|
}
|
|
}
|
|
}
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
// This is the base implementation of triangular matrix-matrix multiplication for all 2D tensors and
|
|
// higher order tensor products that can be expressed as trmm
|
|
// The function uses two level unrolling one based on block sizes and one based on register widths
|
|
// with any remainder left treated in a scalar fashion
|
|
template<typename T, size_t M, size_t K, size_t N, typename LhsType = UpLoType::General, typename RhsType = UpLoType::General>
|
|
FASTOR_INLINE
|
|
void _tmatmul_base(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c) {
|
|
|
|
using V = choose_best_simd_t<SIMDVector<T,DEFAULT_ABI>,N>;
|
|
|
|
// This parameter can be adjusted and does not need to be 4UL/8UL etc
|
|
// constexpr size_t unrollOuterloop = M % 5UL == 0 ? 5UL : 4UL;
|
|
constexpr size_t unrollOuterloop = 4UL;
|
|
|
|
#ifndef FASTOR_MATMUL_OUTER_BLOCK_SIZE
|
|
// Unroll the rows of (a and c) (M) by [numSIMDRows * V::Size]
|
|
constexpr size_t numSIMDRows = M % (unrollOuterloop * 3UL) == 0 ? 3UL : (M < 2UL*V::Size ? 1UL : 2UL);
|
|
#else
|
|
constexpr size_t numSIMDRows = FASTOR_MATMUL_OUTER_BLOCK_SIZE;
|
|
#endif
|
|
#ifndef FASTOR_MATMUL_INNER_BLOCK_SIZE
|
|
// Unroll the columns of (b and c) (N) by [numSIMDCols * V::Size]
|
|
constexpr size_t numSIMDCols = (N % (V::Size * 3UL) == 0 && M % (V::Size * 3UL) == 0 && N > 24UL) ? 3UL : 2UL;
|
|
#else
|
|
constexpr size_t numSIMDCols = FASTOR_MATMUL_INNER_BLOCK_SIZE;
|
|
#endif
|
|
|
|
// The goal is to get 10 parallel independent chains of accumulators
|
|
// to saturate the pipeline by having a completely unrolled block of
|
|
// [(unrollOuterloop) * (numSIMDCols)] at a time. A minimum value of
|
|
// unrollOuterloop=4 ensures a minimum of 8 independent parallel chains
|
|
// while a maximum of 12 i.e. for numSIMDCols=2 and numSIMDCols=3 respectively.
|
|
// However, most recent X86/64 architectures can do 2 FMAs per load so
|
|
// so unrolling with numSIMDCols > 2 is not beneficial
|
|
|
|
constexpr size_t unrollOuterBlock = numSIMDRows*unrollOuterloop;
|
|
// Number of rows of c (M) that can be safely unrolled with this block size.
|
|
constexpr size_t M0 = M / unrollOuterBlock * unrollOuterBlock;
|
|
|
|
constexpr size_t unrollInnerBlock = numSIMDCols*V::Size;
|
|
// Number of columns of c (N) that can be safely unrolled with this block size
|
|
constexpr size_t N0 = N / unrollInnerBlock * unrollInnerBlock;
|
|
|
|
// Number of columns of c (N) that can be safely unrolled with V::Size
|
|
constexpr size_t N1 = N / V::Size * V::Size;
|
|
|
|
size_t i = 0;
|
|
for (; i < M0; i += unrollOuterBlock) {
|
|
size_t j = 0;
|
|
for (; j < N0; j += unrollInnerBlock) {
|
|
interior_block_tmatmul_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,numSIMDCols,LhsType,RhsType>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N0 columns
|
|
for (; j < N1; j += V::Size) {
|
|
interior_block_tmatmul_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,1,LhsType,RhsType>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N1 columns
|
|
for (; j < N; ++j) {
|
|
interior_block_tmatmul_scalar_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,1,LhsType,RhsType>(a,b,c,i,j);
|
|
}
|
|
|
|
}
|
|
|
|
// The remaining M-M0 rows are now unrolled yet again by unrollOuterloop.
|
|
// This is necessary as for small sizes the earlier block loop may not be
|
|
// triggered if the size of the block is bigger than the number of rows of
|
|
// (a and c) i.e. M
|
|
constexpr size_t M1 = (M / unrollOuterloop * unrollOuterloop);
|
|
for (; i < M1; i += unrollOuterloop) {
|
|
size_t j = 0;
|
|
for (; j < N0; j += unrollInnerBlock) {
|
|
interior_block_tmatmul_impl<T,V,M,K,N,unrollOuterloop,1,numSIMDCols,LhsType,RhsType>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N0 columns
|
|
for (; j < N1; j += V::Size) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop,V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop,V::Size,LhsType,RhsType>(i,j);
|
|
|
|
V c_ij[unrollOuterloop];
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n] = fmadd(V(a[(i + n)*K+k]), V(&b[k*N+j],false), c_ij[n]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].store(&c[(i + n)*N+j],false);
|
|
}
|
|
}
|
|
|
|
// Remaining N - N1 columns
|
|
for (; j < N; ++j) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop,1,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop,1,LhsType,RhsType>(i,j);
|
|
|
|
T c_ij[unrollOuterloop] = {};
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n] += a[(i + n)*K+k] * b[k*N+j];
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c[(i + n)*N+j] = c_ij[n];
|
|
}
|
|
}
|
|
}
|
|
|
|
// Now treat the remaining M-M1 rows - here the klast - kfirst range is not used
|
|
// so the implementation is exactly the same as matmul_base
|
|
FASTOR_IF_CONSTEXPR (M-M1 > 0) {
|
|
// Hack to get around zero length array issue
|
|
constexpr size_t MM1 = M-M1 != 0 ? M-M1 : 1;
|
|
size_t j = 0;
|
|
for (; j < N0; j += unrollInnerBlock) {
|
|
// If MM1==0 the function never gets invoked anyway
|
|
interior_block_tmatmul_impl<T,V,M,K,N,MM1,1,numSIMDCols>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N0 columns
|
|
for (; j < N1; j += V::Size) {
|
|
|
|
V c_ij[MM1];
|
|
for (size_t k = 0; k < K; ++k) {
|
|
for (size_t n = M1; n < M; ++n) {
|
|
c_ij[n-M1] = fmadd(V(a[n*K+k]), V(&b[k*N+j],false), c_ij[n-M1]);
|
|
c_ij[n-M1].store(&c[n*N+j],false);
|
|
}
|
|
}
|
|
for (size_t n = M1; n < M; ++n) {
|
|
c_ij[n-M1].store(&c[n*N+j],false);
|
|
}
|
|
}
|
|
|
|
// Remaining N - N1 columns
|
|
for (; j < N; ++j) {
|
|
|
|
T c_ij[MM1] = {};
|
|
for (size_t k = 0; k < K; ++k) {
|
|
for (size_t n = M1; n < M; ++n) {
|
|
c_ij[n-M1] += a[n*K+k] * b[k*N+j];
|
|
c[n*N+j] = c_ij[n-M1];
|
|
}
|
|
}
|
|
for (size_t n = M1; n < M; ++n) {
|
|
c[n*N+j] = c_ij[n-M1];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
|
|
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
// This the base implementation of triangular matrix-matrix multiplication for all 2D tensors and
|
|
// higher order tensor products that can be expressed as trmm
|
|
// The function uses two level unrolling one based on block sizes and one based on register widths
|
|
// with any remainder left treated in vector mode with masked and conditional load/stores.
|
|
// Note that conditional load/store requires at least AVX intrinsics
|
|
template<typename T, size_t M, size_t K, size_t N, typename LhsType = UpLoType::General, typename RhsType = UpLoType::General>
|
|
FASTOR_INLINE
|
|
void _tmatmul_base_masked(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c) {
|
|
|
|
using V = typename internal::choose_best_simd_type<SIMDVector<T,DEFAULT_ABI>,N>::type;
|
|
|
|
// This parameter can be adjusted and does not need to be 4UL/8UL etc
|
|
// constexpr size_t unrollOuterloop = M % 5UL == 0 ? 5UL : 4UL;
|
|
constexpr size_t unrollOuterloop = 4UL;
|
|
|
|
#ifndef FASTOR_MATMUL_OUTER_BLOCK_SIZE
|
|
// Unroll the rows of (a and c) (M) by [numSIMDRows * V::Size]
|
|
constexpr size_t numSIMDRows = M % (unrollOuterloop * 3UL) == 0 ? 3UL : (M < 2UL*V::Size ? 1UL : 2UL);
|
|
#else
|
|
constexpr size_t numSIMDRows = FASTOR_MATMUL_OUTER_BLOCK_SIZE;
|
|
#endif
|
|
#ifndef FASTOR_MATMUL_INNER_BLOCK_SIZE
|
|
// Unroll the columns of (b and c) (N) by [numSIMDCols * V::Size]
|
|
constexpr size_t numSIMDCols = (N % (V::Size * 3UL) == 0 && M % (V::Size * 3UL) == 0 && N > 24UL) ? 3UL : 2UL;
|
|
#else
|
|
constexpr size_t numSIMDCols = FASTOR_MATMUL_INNER_BLOCK_SIZE;
|
|
#endif
|
|
|
|
// The goal is to get 10 parallel independent chains of accumulators
|
|
// to saturate the pipeline by having a completely unrolled block of
|
|
// [(unrollOuterloop) * (numSIMDCols)] at a time. A minimum value of
|
|
// unrollOuterloop=4 ensures a minimum of 8 independent parallel chains
|
|
// while a maximum of 12 i.e. for numSIMDCols=2 and numSIMDCols=3 respectively.
|
|
// However, most recent X86/64 architectures can do 2 FMAs per load so
|
|
// so unrolling with numSIMDCols > 2 is not beneficial
|
|
|
|
constexpr size_t unrollOuterBlock = numSIMDRows*unrollOuterloop;
|
|
// Number of rows of c (M) that can be safely unrolled with this block size.
|
|
constexpr size_t M0 = M / unrollOuterBlock * unrollOuterBlock;
|
|
|
|
constexpr size_t unrollInnerBlock = numSIMDCols*V::Size;
|
|
// Number of columns of c (N) that can be safely unrolled with this block size
|
|
constexpr size_t N0 = N / unrollInnerBlock * unrollInnerBlock;
|
|
|
|
// Number of columns of c (N) that can be safely unrolled with V::Size
|
|
constexpr size_t N1 = N / V::Size * V::Size;
|
|
|
|
int maska[V::Size];
|
|
std::fill(maska,&maska[V::Size], -1);
|
|
for (size_t jj=0; jj < V::Size - (N-N1); ++jj) maska[jj] = 0;
|
|
#ifdef FASTOR_HAS_AVX512_MASKS
|
|
const auto mask = array_to_mask(maska);
|
|
#endif
|
|
|
|
size_t i = 0;
|
|
for (; i < M0; i += unrollOuterBlock) {
|
|
size_t j = 0;
|
|
for (; j < N0; j += unrollInnerBlock) {
|
|
interior_block_tmatmul_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,numSIMDCols>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N0 columns
|
|
for (; j < N1; j += V::Size) {
|
|
interior_block_tmatmul_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,1>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N1 columns
|
|
for (; j < N; j+= N-N1) {
|
|
#ifdef FASTOR_HAS_AVX512_MASKS
|
|
interior_block_matmul_mask_impl<T,decltype(mask),V,M,K,N,unrollOuterloop,numSIMDRows,1>(a,b,c,i,j,mask);
|
|
#else
|
|
interior_block_matmul_mask_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,1>(a,b,c,i,j,maska);
|
|
#endif
|
|
}
|
|
}
|
|
|
|
// The remaining M-M0 rows are now unrolled yet again by unrollOuterloop.
|
|
// This is necessary as for small sizes the earlier block loop may not be
|
|
// triggered if the size of the block is bigger than the number of rows of
|
|
// (a and c) i.e. M
|
|
constexpr size_t M1 = (M / unrollOuterloop * unrollOuterloop);
|
|
for (; i < M1; i += unrollOuterloop) {
|
|
size_t j = 0;
|
|
for (; j < N0; j += unrollInnerBlock) {
|
|
interior_block_tmatmul_impl<T,V,M,K,N,unrollOuterloop,1,numSIMDCols>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N0 columns
|
|
for (; j < N1; j += V::Size) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop,V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop,V::Size,LhsType,RhsType>(i,j);
|
|
|
|
V c_ij[unrollOuterloop];
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n] = fmadd(V(a[(i + n)*K+k]), V(&b[k*N+j],false), c_ij[n]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
c_ij[n].store(&c[(i + n)*N+j],false);
|
|
}
|
|
}
|
|
|
|
// Remaining N - N1 columns
|
|
for (; j < N; j+=N-N1) {
|
|
|
|
const size_t kfirst = find_kfirst<size_t,K,unrollOuterloop,V::Size,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,unrollOuterloop,V::Size,LhsType,RhsType>(i,j);
|
|
|
|
V c_ij[unrollOuterloop];
|
|
for (size_t k = kfirst; k < klast; ++k) {
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
#ifdef FASTOR_HAS_AVX512_MASKS
|
|
V bmm0; bmm0.mask_load(&b[k*N+j],mask);
|
|
#else
|
|
const V bmm0(maskload<V>(&b[k*N+j],maska));
|
|
#endif
|
|
const V amm0 = a[(i + n)*K+k];
|
|
c_ij[n] = fmadd(amm0,bmm0,c_ij[n]);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
|
#ifdef FASTOR_HAS_AVX512_MASKS
|
|
c_ij[n].mask_store(&c[(i+n)*N+j],mask,false);
|
|
#else
|
|
maskstore(&c[(i+n)*N+j],maska,c_ij[n]);
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
|
|
// Now treat the remaining M-M1 rows - here the klast - kfirst range is not used
|
|
// so the implementation is exactly the same as matmul_base
|
|
FASTOR_IF_CONSTEXPR (M-M1 > 0) {
|
|
// Hack to get around zero length array issue
|
|
constexpr size_t MM1 = M-M1 != 0 ? M-M1 : 1;
|
|
size_t j = 0;
|
|
for (; j < N0; j += unrollInnerBlock) {
|
|
// If MM1==0 the function never gets invoked anyway
|
|
interior_block_tmatmul_impl<T,V,M,K,N,MM1,1,numSIMDCols>(a,b,c,i,j);
|
|
}
|
|
|
|
// Remaining N - N0 columns
|
|
for (; j < N1; j += V::Size) {
|
|
V c_ij[MM1];
|
|
for (size_t k = 0; k < K; ++k) {
|
|
for (size_t n = M1; n < M; ++n) {
|
|
c_ij[n-M1] = fmadd(V(a[n*K+k]), V(&b[k*N+j],false), c_ij[n-M1]);
|
|
c_ij[n-M1].store(&c[n*N+j],false);
|
|
}
|
|
}
|
|
for (size_t n = M1; n < M; ++n) {
|
|
c_ij[n-M1].store(&c[n*N+j],false);
|
|
}
|
|
}
|
|
|
|
// Remaining N - N1 columns
|
|
for (; j < N; j+=N-N1) {
|
|
V c_ij[MM1] = {};
|
|
for (size_t k = 0; k < K; ++k) {
|
|
for (size_t n = M1; n < M; ++n) {
|
|
#ifdef FASTOR_HAS_AVX512_MASKS
|
|
V bmm0; bmm0.mask_load(&b[k*N+j],mask);
|
|
#else
|
|
const V bmm0(maskload<V>(&b[k*N+j],maska));
|
|
#endif
|
|
const V amm0 = a[n*K+k];
|
|
c_ij[n-M1] = fmadd(amm0,bmm0,c_ij[n-M1]);
|
|
}
|
|
}
|
|
for (size_t n = M1; n < M; ++n) {
|
|
#ifdef FASTOR_HAS_AVX512_MASKS
|
|
c_ij[n-M1].mask_store(&c[n*N+j],mask,false);
|
|
#else
|
|
maskstore(&c[n*N+j],maska,c_ij[n-M1]);
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
}
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
|
|
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
// tmatmul kernel for non-fundamental types
|
|
// The assumption here is that non-fundamental types are not SIMD vectorisable for instance
|
|
// Tensor<std::vector<T>,3,3> or Tensor<Tensor<...>,...> plus they cannot fuse [do fused-add-multiply]
|
|
// so operations like [c += a*b] or potentially [c = c + a*b] might introduce multiple copies in
|
|
// the inner most loops of matmul
|
|
template<typename T, size_t M, size_t K, size_t N, typename LhsType = UpLoType::General, typename RhsType = UpLoType::General>
|
|
FASTOR_INLINE
|
|
void _tmatmul_base_non_primitive(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c) {
|
|
// There is no SIMD here as V::Size == 1 anyway
|
|
// No outer loop unrolling otherwise the innermost loop
|
|
// will create unnecessary temporaries
|
|
for (size_t i=0; i<M; ++i) {
|
|
// V::Size == 1 so this loop can't be unrolled
|
|
for (size_t j=0; j<N; ++j) {
|
|
const size_t kfirst = find_kfirst<size_t,K,1,1,LhsType,RhsType>(i,j);
|
|
const size_t klast = find_klast <size_t,K,1,1,LhsType,RhsType>(i,j);
|
|
// This could potentially cost as opposed to directly writing in to c
|
|
T tmp {};
|
|
for (size_t k=kfirst; k<klast; ++k) {
|
|
tmp += a[i*K+k]*b[k*N+j];
|
|
}
|
|
c[i*N+j] = tmp;
|
|
}
|
|
}
|
|
}
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
|
|
} // end of namespace internal
|
|
|
|
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
// Backend tmatmul function
|
|
template<typename T, size_t M, size_t K, size_t N, typename LhsType = UpLoType::General, typename RhsType = UpLoType::General>
|
|
FASTOR_INLINE
|
|
void _tmatmul(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT out) {
|
|
|
|
// Non-primitive types
|
|
FASTOR_IF_CONSTEXPR (!is_primitive_v_<T>) {
|
|
internal::_tmatmul_base_non_primitive<T,M,K,N,LhsType,RhsType>(a,b,out);
|
|
return;
|
|
}
|
|
|
|
// Use specialised kernels
|
|
#if defined(FASTOR_AVX2_IMPL) || defined(FASTOR_HAS_AVX512_MASKS)
|
|
|
|
using nativeV = SIMDVector<T,DEFAULT_ABI>;
|
|
using V = choose_best_simd_t<nativeV,N>;
|
|
|
|
FASTOR_IF_CONSTEXPR(N % V::Size <= 1UL) {
|
|
internal::_tmatmul_base<T,M,K,N,LhsType,RhsType>(a,b,out);
|
|
return;
|
|
}
|
|
else {
|
|
internal::_tmatmul_base_masked<T,M,K,N,LhsType,RhsType>(a,b,out);
|
|
return;
|
|
}
|
|
#else
|
|
internal::_tmatmul_base<T,M,K,N,LhsType,RhsType>(a,b,out);
|
|
return;
|
|
#endif
|
|
}
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
//-----------------------------------------------------------------------------------------------------------
|
|
} // end of namespace Fastor
|
|
|
|
|
|
#endif // MATMUL_KERNELS_H
|