Add Fastor library
This commit is contained in:
Bassem Girgis
823
noarch/include/Fastor/backend/matmul/matmul_kernels.h
Normal file
823
noarch/include/Fastor/backend/matmul/matmul_kernels.h
Normal file
@@ -0,0 +1,823 @@
|
||||
#ifndef MATMUL_KERNELS_H
|
||||
#define MATMUL_KERNELS_H
|
||||
|
||||
|
||||
#include "Fastor/config/config.h"
|
||||
#include "Fastor/simd_vector/extintrin.h"
|
||||
#include "Fastor/simd_vector/SIMDVector.h"
|
||||
#include "Fastor/meta/tensor_meta.h"
|
||||
|
||||
|
||||
namespace Fastor {
|
||||
|
||||
namespace internal {
|
||||
|
||||
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
|
||||
// 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 std::enable_if<numSIMDCols==1,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_impl(
|
||||
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
||||
const size_t i, const size_t 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 = 0; k < K; ++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 std::enable_if<numSIMDCols==2,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_impl(
|
||||
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
||||
const size_t i, const size_t 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 = 0; k < K; ++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 std::enable_if<numSIMDCols==3,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_impl(
|
||||
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
||||
const size_t i, const size_t 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 = 0; k < K; ++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 std::enable_if<numSIMDCols==4,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_impl(
|
||||
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
||||
const size_t i, const size_t 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 = 0; k < K; ++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 std::enable_if<numSIMDCols==5,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_impl(
|
||||
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
||||
const size_t i, const size_t 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 = 0; k < K; ++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 std::enable_if<numSIMDCols==1,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_scalar_impl(
|
||||
const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c,
|
||||
const size_t i, const size_t 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 = 0; k < K; ++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 std::enable_if<numSIMDCols==1,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_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]) {
|
||||
|
||||
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 = 0; k < K; ++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 std::enable_if<numSIMDCols==1,bool>::type = false>
|
||||
FASTOR_INLINE
|
||||
void interior_block_matmul_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;
|
||||
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 = 0; k < K; ++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 matrix-matrix multiplication for all 2D tensors and
|
||||
// higher order tensor products that can be expressed as gemm
|
||||
// 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>
|
||||
FASTOR_INLINE
|
||||
void _matmul_base(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;
|
||||
|
||||
size_t i = 0;
|
||||
for (; i < M0; i += unrollOuterBlock) {
|
||||
size_t j = 0;
|
||||
for (; j < N0; j += unrollInnerBlock) {
|
||||
interior_block_matmul_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_matmul_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,1>(a,b,c,i,j);
|
||||
}
|
||||
|
||||
// Remaining N - N1 columns
|
||||
for (; j < N; ++j) {
|
||||
interior_block_matmul_scalar_impl<T,V,M,K,N,unrollOuterloop,numSIMDRows,1>(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_matmul_impl<T,V,M,K,N,unrollOuterloop,1,numSIMDCols>(a,b,c,i,j);
|
||||
}
|
||||
|
||||
// Remaining N - N0 columns
|
||||
for (; j < N1; j += V::Size) {
|
||||
V c_ij[unrollOuterloop];
|
||||
for (size_t k = 0; k < K; ++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) {
|
||||
T c_ij[unrollOuterloop] = {};
|
||||
for (size_t k = 0; k < K; ++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
|
||||
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_matmul_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 is the base implementation of matrix-matrix multiplication for all 2D tensors and
|
||||
// higher order tensor products that can be expressed as gemm
|
||||
// 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>
|
||||
FASTOR_INLINE
|
||||
void _matmul_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_matmul_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_matmul_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_matmul_impl<T,V,M,K,N,unrollOuterloop,1,numSIMDCols>(a,b,c,i,j);
|
||||
}
|
||||
|
||||
// Remaining N - N0 columns
|
||||
for (; j < N1; j += V::Size) {
|
||||
V c_ij[unrollOuterloop];
|
||||
for (size_t k = 0; k < K; ++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) {
|
||||
V c_ij[unrollOuterloop];
|
||||
for (size_t k = 0; k < K; ++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
|
||||
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_matmul_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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
// matmul 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>
|
||||
FASTOR_INLINE
|
||||
void _matmul_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) {
|
||||
// This could potentially cost as opposed to directly writing in to c
|
||||
T tmp {};
|
||||
for (size_t k=0; k<K; ++k) {
|
||||
tmp += a[i*K+k]*b[k*N+j];
|
||||
}
|
||||
c[i*N+j] = tmp;
|
||||
}
|
||||
}
|
||||
}
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
// Other variants and slightly older implementations
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
// This is the same implementation as the above case but does not unroll on block sizes and does not require
|
||||
// the registers to be zeroed out but K must be !=1
|
||||
template<typename T, size_t M, size_t K, size_t N>
|
||||
FASTOR_INLINE
|
||||
void _matmul_mkn_square(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;
|
||||
|
||||
// Get 10 parallel independent chains of accumulators for bigger matrices
|
||||
constexpr size_t unrollOuterloop = M >= 64 ? 10UL : (M % 8 == 0 ? 8UL : V::Size);
|
||||
|
||||
// The row index (for a and c) is unrolled using the unrollOuterloop stride. Therefore
|
||||
// the last rows may need special treatment if M is not a multiple of unrollOuterloop.
|
||||
// M0 is the number of rows that can safely be iterated with a stride of
|
||||
// unrollOuterloop.
|
||||
constexpr size_t M0 = M / unrollOuterloop * unrollOuterloop;
|
||||
for (size_t i = 0; i < M0; i += unrollOuterloop) {
|
||||
// The iteration over the column index of b and c uses a stride of V::Size. This
|
||||
// enables row-vector loads (from b) and stores (to c). The matrix storage is
|
||||
// padded accordingly, ensuring correct bounds and alignment.
|
||||
for (size_t j = 0; j < N; j += V::Size) {
|
||||
// This temporary variables are used to accumulate the results of the products
|
||||
// producing the new values for the c matrix. This variable is necessary
|
||||
// because we need a V object for data-parallel accumulation. Storing to c
|
||||
// directly stores to scalar objects and thus would drop the ability for
|
||||
// data-parallel (SIMD) addition.
|
||||
V c_ij[unrollOuterloop];
|
||||
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
||||
c_ij[n] = a[(i + n)*K]*V(&b[j]);
|
||||
}
|
||||
for (size_t k = 1; k < K - 1; ++k) {
|
||||
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
||||
c_ij[n] += a[(i + n)*K+k] * V(&b[k*N+j]);
|
||||
}
|
||||
}
|
||||
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
||||
c_ij[n] += a[(i + n)*K+(K - 1)] * V(&b[(K - 1)*N+j]);
|
||||
c_ij[n].store(&c[(i + n)*N+j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// This is the same implementation as the above case but does not unroll on block sizes and does not require
|
||||
// the registers to be zeroed out but K must be !=1
|
||||
template<typename T, size_t M, size_t K, size_t N>
|
||||
FASTOR_INLINE
|
||||
void _matmul_mkn_non_square(const T * FASTOR_RESTRICT a, const T * FASTOR_RESTRICT b, T * FASTOR_RESTRICT c) {
|
||||
|
||||
// This variant strictly cannot deal outer-product i.e. with K==1
|
||||
|
||||
using V = typename internal::choose_best_simd_type<SIMDVector<T,DEFAULT_ABI>,N>::type;
|
||||
|
||||
// Get 10 parallel independent chains of accumulators for bigger matrices
|
||||
constexpr size_t unrollOuterloop = M < V::Size ? 1UL :
|
||||
(( M >= 64 && K > 10 && N > V::Size ) ? 10UL : (M % 8 == 0 && N > V::Size ? 8UL : V::Size));
|
||||
constexpr bool isPadded = N % V::Size == 0;
|
||||
|
||||
// The row index (for a and c) is unrolled using the unrollOuterloop stride. Therefore
|
||||
// the last rows may need special treatment if M is not a multiple of unrollOuterloop.
|
||||
// M0 is the number of rows that can safely be iterated with a stride of
|
||||
// unrollOuterloop.
|
||||
constexpr size_t M0 = M / unrollOuterloop * unrollOuterloop;
|
||||
constexpr size_t N0 = N / V::Size * V::Size;
|
||||
|
||||
for (size_t i = 0; i < M0; i += unrollOuterloop) {
|
||||
// The iteration over the column index of b and c uses a stride of V::size(). This
|
||||
// enables row-vector loads (from b) and stores (to c). The matrix storage is
|
||||
// padded accordingly, ensuring correct bounds and alignment.
|
||||
size_t j = 0;
|
||||
for (; j < N0; j += V::Size) {
|
||||
// This temporary variables are used to accumulate the results of the products
|
||||
// producing the new values for the c matrix. This variable is necessary
|
||||
// because we need a V object for data-parallel accumulation. Storing to c
|
||||
// directly stores to scalar objects and thus would drop the ability for
|
||||
// data-parallel (SIMD) addition.
|
||||
V c_ij[unrollOuterloop];
|
||||
for (size_t n = 0; n < unrollOuterloop; ++n) { // correct
|
||||
c_ij[n] = a[(i + n)*K]*V(&b[j], isPadded);
|
||||
}
|
||||
for (size_t k = 1; k < K - 1; ++k) { // correct
|
||||
for (size_t n = 0; n < unrollOuterloop; ++n) {
|
||||
c_ij[n] += a[(i + n)*K+k] * V(&b[k*N+j], false);
|
||||
}
|
||||
}
|
||||
for (size_t n = 0; n < unrollOuterloop; ++n) { // correct
|
||||
c_ij[n] += a[(i + n)*K+(K - 1)] * V(&b[(K - 1)*N+j], false);
|
||||
c_ij[n].store(&c[(i + n)*N+j], isPadded);
|
||||
}
|
||||
}
|
||||
|
||||
// Remainder N - N0 columns
|
||||
for (; j < N; ++j) {
|
||||
T c_ij[unrollOuterloop];
|
||||
for (size_t n = 0; n < unrollOuterloop; ++n) { // correct
|
||||
c_ij[n] = a[(i + n)*K]*b[j];
|
||||
}
|
||||
for (size_t k = 1; k < K - 1; ++k) { // correct
|
||||
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) { // correct
|
||||
c_ij[n] += a[(i + n)*K+(K - 1)] * b[(K - 1)*N+j];
|
||||
c[(i + n)*N+j] = c_ij[n];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// This final loop treats the remaining M - M0 rows.
|
||||
size_t j = 0;
|
||||
for (; j < N0; j += V::Size) {
|
||||
V c_ij[M-M0];
|
||||
for (size_t n = M0; n < M; ++n) { // correct
|
||||
c_ij[n - M0] = a[n*K] * V(&b[j], isPadded);
|
||||
}
|
||||
for (size_t k = 1; k < K - 1; ++k) { // correct
|
||||
for (size_t n = M0; n < M; ++n) { // correct
|
||||
c_ij[n - M0] += a[n*K+k] * V(&b[k*N+j], false);
|
||||
}
|
||||
}
|
||||
for (size_t n = M0; n < M; ++n) { // correct
|
||||
c_ij[n - M0] += a[n*K+(K - 1)] * V(&b[(K - 1)*N+j], false);
|
||||
c_ij[n - M0].store(&c[n*N+j], isPadded);
|
||||
}
|
||||
}
|
||||
|
||||
for (; j < N; ++j) {
|
||||
T c_ij[M-M0];
|
||||
for (size_t n = M0; n < M; ++n) { // correct
|
||||
c_ij[n - M0] = a[n*K] * b[j];
|
||||
}
|
||||
for (size_t k = 1; k < K - 1; ++k) { // correct
|
||||
for (size_t n = M0; n < M; ++n) { // correct
|
||||
c_ij[n - M0] += a[n*K+k] * b[k*N+j];
|
||||
}
|
||||
}
|
||||
for (size_t n = M0; n < M; ++n) { // correct
|
||||
c_ij[n - M0] += a[n*K+(K - 1)] * b[(K - 1)*N+j];
|
||||
c[n*N+j] = c_ij[n - M0];
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
|
||||
} // end of namespace internal
|
||||
|
||||
} // end of namespace Fastor
|
||||
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
#include "Fastor/backend/matmul/matmul_mk_smalln.h"
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
//-----------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
#endif // MATMUL_KERNELS_H
|
||||
Reference in New Issue
Block a user