Misc. packet math cleanups.

libeigen/eigen!2070

Co-authored-by: Rasmus Munk Larsen <rmlarsen@google.com>
This commit is contained in:
Rasmus Munk Larsen
2025-11-08 21:57:20 +00:00
parent 4cb0776f8e
commit 23a5482fc0
11 changed files with 211 additions and 289 deletions

View File

@@ -57,12 +57,12 @@ struct default_packet_traits {
HasConj = 1,
HasSetLinear = 1,
HasSign = 1,
HasAbsDiff = 1,
// By default, the nearest integer functions (rint, round, floor, ceil, trunc) are enabled for all scalar and packet
// types
HasRound = 1,
HasArg = 0,
HasAbsDiff = 0,
// This flag is used to indicate whether packet comparison is supported.
// pcmp_eq and pcmp_lt should be defined for it to be true.
HasCmp = 0,
@@ -116,6 +116,7 @@ struct packet_traits : default_packet_traits {
enum {
HasAdd = 0,
HasSub = 0,
HasAbsDiff = 0,
HasMul = 0,
HasNegate = 0,
HasAbs = 0,
@@ -130,17 +131,18 @@ struct packet_traits : default_packet_traits {
template <typename T>
struct packet_traits<const T> : packet_traits<T> {};
struct default_unpacket_traits {
enum { vectorizable = false, masked_load_available = false, masked_store_available = false };
};
template <typename T>
struct unpacket_traits {
struct unpacket_traits : default_unpacket_traits {
typedef T type;
typedef T half;
typedef typename numext::get_integer_by_size<sizeof(T)>::signed_type integer_packet;
enum {
size = 1,
alignment = alignof(T),
vectorizable = false,
masked_load_available = false,
masked_store_available = false
};
};
@@ -747,9 +749,15 @@ EIGEN_DEVICE_FUNC inline Packet pldexp(const Packet& a, const Packet& exponent)
/** \internal \returns the min of \a a and \a b (coeff-wise) */
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pabsdiff(const Packet& a, const Packet& b) {
EIGEN_DEVICE_FUNC inline std::enable_if_t<NumTraits<typename unpacket_traits<Packet>::type>::IsInteger, Packet>
pabsdiff(const Packet& a, const Packet& b) {
return pselect(pcmp_lt(a, b), psub(b, a), psub(a, b));
}
template <typename Packet>
EIGEN_DEVICE_FUNC inline std::enable_if_t<!NumTraits<typename unpacket_traits<Packet>::type>::IsInteger, Packet>
pabsdiff(const Packet& a, const Packet& b) {
return pabs(psub(a, b));
}
/** \internal \returns a packet version of \a *from, from must be properly aligned */
template <typename Packet>

View File

@@ -250,7 +250,6 @@ struct packet_traits<uint32_t> : default_packet_traits {
HasDiv = 0,
HasNegate = 0,
HasSqrt = 0,
HasCmp = 1,
HasMin = 1,
@@ -277,12 +276,9 @@ struct packet_traits<uint64_t> : default_packet_traits {
AlignedOnScalar = 1,
size = 4,
// HasMin = 0,
// HasMax = 0,
HasDiv = 0,
HasTranspose = 0,
HasNegate = 0,
HasSqrt = 0,
HasCmp = 1,
HasShift = 1
};

View File

@@ -81,10 +81,10 @@ class gemm_class {
Index m;
const Index n, k, ldc;
const Index inc;
const Scalar *alpha;
const Scalar* alpha;
const Scalar *a, *b;
Scalar *c;
Scalar* c;
const bool is_alpha1;
const bool is_beta0;
@@ -92,26 +92,26 @@ class gemm_class {
const Index a_stride, b_stride;
const Index a_off, b_off;
EIGEN_ALWAYS_INLINE void prefetch_a(const Scalar *a_addr) {
_mm_prefetch((char *)(a_prefetch_size + a_addr - a_shift), _MM_HINT_T0);
EIGEN_ALWAYS_INLINE void prefetch_a(const Scalar* a_addr) {
_mm_prefetch((char*)(a_prefetch_size + a_addr - a_shift), _MM_HINT_T0);
}
EIGEN_ALWAYS_INLINE void prefetch_b(const Scalar *b_addr) {
_mm_prefetch((char *)(b_prefetch_size + b_addr - b_shift), _MM_HINT_T0);
EIGEN_ALWAYS_INLINE void prefetch_b(const Scalar* b_addr) {
_mm_prefetch((char*)(b_prefetch_size + b_addr - b_shift), _MM_HINT_T0);
}
EIGEN_ALWAYS_INLINE void prefetch_x(const Scalar *x_addr) { _mm_prefetch((char *)(x_addr - a_shift), _MM_HINT_T2); }
EIGEN_ALWAYS_INLINE void prefetch_x(const Scalar* x_addr) { _mm_prefetch((char*)(x_addr - a_shift), _MM_HINT_T2); }
EIGEN_ALWAYS_INLINE void prefetch_c(const Scalar *c_addr) {
EIGEN_ALWAYS_INLINE void prefetch_c(const Scalar* c_addr) {
#if defined(__PRFCHW__) && __PRFCHW__ == 1
_m_prefetchw((void *)c_addr);
_m_prefetchw((void*)c_addr);
#else
_mm_prefetch((char *)c_addr, _MM_HINT_T0);
_mm_prefetch((char*)c_addr, _MM_HINT_T0);
#endif
}
template <int nelems>
EIGEN_ALWAYS_INLINE void a_load(vec &a_reg, const Scalar *a_addr) {
EIGEN_ALWAYS_INLINE void a_load(vec& a_reg, const Scalar* a_addr) {
switch (nelems * sizeof(*a_addr) * 8) {
default:
case 512 * 3:
@@ -124,13 +124,13 @@ class gemm_class {
a_reg = ploadu<vec>(a_addr);
break;
case 256 * 1:
a_reg = preinterpret<vec>(_mm512_broadcast_f64x4(ploadu<Packet4d>(reinterpret_cast<const double *>(a_addr))));
a_reg = preinterpret<vec>(_mm512_broadcast_f64x4(ploadu<Packet4d>(reinterpret_cast<const double*>(a_addr))));
break;
case 128 * 1:
a_reg = preinterpret<vec>(_mm512_broadcast_f32x4(ploadu<Packet4f>(reinterpret_cast<const float *>(a_addr))));
a_reg = preinterpret<vec>(_mm512_broadcast_f32x4(ploadu<Packet4f>(reinterpret_cast<const float*>(a_addr))));
break;
case 64 * 1:
a_reg = preinterpret<vec>(pload1<Packet8d>(reinterpret_cast<const double *>(a_addr)));
a_reg = preinterpret<vec>(pload1<Packet8d>(reinterpret_cast<const double*>(a_addr)));
break;
case 32 * 1:
a_reg = pload1<vec>(a_addr);
@@ -138,10 +138,10 @@ class gemm_class {
}
}
EIGEN_ALWAYS_INLINE void b_load(vec &b_reg, const Scalar *b_addr) { b_reg = pload1<vec>(b_addr); }
EIGEN_ALWAYS_INLINE void b_load(vec& b_reg, const Scalar* b_addr) { b_reg = pload1<vec>(b_addr); }
template <int nelems>
EIGEN_ALWAYS_INLINE void c_store(Scalar *mem, vec &src) {
EIGEN_ALWAYS_INLINE void c_store(Scalar* mem, vec& src) {
if (is_unit_inc) {
switch (nelems * sizeof(*mem) * 8) {
default:
@@ -196,7 +196,7 @@ class gemm_class {
}
template <int nelems>
EIGEN_ALWAYS_INLINE void vaddm(vec &dst, const Scalar *mem, vec &src, vec &reg) {
EIGEN_ALWAYS_INLINE void vaddm(vec& dst, const Scalar* mem, vec& src, vec& reg) {
if (is_unit_inc) {
switch (nelems * sizeof(*mem) * 8) {
default:
@@ -263,7 +263,7 @@ class gemm_class {
}
}
EIGEN_STRONG_INLINE void vfmadd(vec &dst, const vec &src1, const vec &src2) {
EIGEN_STRONG_INLINE void vfmadd(vec& dst, const vec& src1, const vec& src2) {
dst = pmadd(src1, src2, dst);
#if (EIGEN_COMP_GNUC != 0) || (EIGEN_COMP_CLANG != 0)
@@ -273,7 +273,7 @@ class gemm_class {
}
template <int nelems>
EIGEN_ALWAYS_INLINE void vfmaddm(vec &dst, const Scalar *mem, vec &src, vec &scale, vec &reg) {
EIGEN_ALWAYS_INLINE void vfmaddm(vec& dst, const Scalar* mem, vec& src, vec& scale, vec& reg) {
if (is_unit_inc) {
switch (nelems * sizeof(*mem) * 8) {
default:
@@ -350,16 +350,16 @@ class gemm_class {
}
template <int j, int endX, int i, int endY, int nelems>
EIGEN_ALWAYS_INLINE std::enable_if_t<(j > endX) || (i > endY)> a_loads(const Scalar *ao) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(j > endX) || (i > endY)> a_loads(const Scalar* ao) {
EIGEN_UNUSED_VARIABLE(ao);
}
template <int j, int endX, int i, int endY, int nelems>
EIGEN_ALWAYS_INLINE std::enable_if_t<(j <= endX) && (i <= endY)> a_loads(const Scalar *ao) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(j <= endX) && (i <= endY)> a_loads(const Scalar* ao) {
if (j < endX) {
if (i < endY) {
auto &a_reg = zmm[a_regs[i + (j % 2) * 3]];
const Scalar *a_addr = ao + nelems * j + nelems_in_cache_line * i - a_shift;
auto& a_reg = zmm[a_regs[i + (j % 2) * 3]];
const Scalar* a_addr = ao + nelems * j + nelems_in_cache_line * i - a_shift;
a_load<nelems>(a_reg, a_addr);
a_loads<j, endX, i + 1, endY, nelems>(ao);
@@ -370,8 +370,8 @@ class gemm_class {
}
template <int un, int max_b_unroll, int i, int um_vecs, int a_unroll, int b_unroll>
EIGEN_ALWAYS_INLINE std::enable_if_t<(un > max_b_unroll) || (i > um_vecs)> prefetch_cs(const Scalar *co1,
const Scalar *co2) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(un > max_b_unroll) || (i > um_vecs)> prefetch_cs(const Scalar* co1,
const Scalar* co2) {
EIGEN_UNUSED_VARIABLE(co1);
EIGEN_UNUSED_VARIABLE(co2);
}
@@ -391,13 +391,13 @@ class gemm_class {
*/
template <int un, int max_b_unroll, int i, int um_vecs, int a_unroll, int b_unroll>
EIGEN_ALWAYS_INLINE std::enable_if_t<(un <= max_b_unroll) && (i <= um_vecs)> prefetch_cs(Scalar *&co1, Scalar *&co2) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(un <= max_b_unroll) && (i <= um_vecs)> prefetch_cs(Scalar*& co1, Scalar*& co2) {
if (un < max_b_unroll) {
if (b_unroll >= un + 1) {
if (un == 4 && i == 0) co2 = co1 + 4 * ldc;
if (i < um_vecs) {
Scalar *co = (un + 1 <= 4) ? co1 : co2;
Scalar* co = (un + 1 <= 4) ? co1 : co2;
auto co_off = (un % 4) * ldc + a_unroll - 1 + i * nelems_in_cache_line * sizeof *co;
prefetch_c(co + co_off);
@@ -414,16 +414,16 @@ class gemm_class {
// load_c
template <int i, int um_vecs, int idx, int nelems>
EIGEN_ALWAYS_INLINE std::enable_if_t<(i > um_vecs)> scale_load_c(const Scalar *cox, vec &alpha_reg) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(i > um_vecs)> scale_load_c(const Scalar* cox, vec& alpha_reg) {
EIGEN_UNUSED_VARIABLE(cox);
EIGEN_UNUSED_VARIABLE(alpha_reg);
}
template <int i, int um_vecs, int idx, int nelems>
EIGEN_ALWAYS_INLINE std::enable_if_t<(i <= um_vecs)> scale_load_c(const Scalar *cox, vec &alpha_reg) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(i <= um_vecs)> scale_load_c(const Scalar* cox, vec& alpha_reg) {
if (i < um_vecs) {
auto &c_reg = zmm[c_regs[i + idx * 3]];
auto &c_load_reg = zmm[c_load_regs[i % 3]];
auto& c_reg = zmm[c_regs[i + idx * 3]];
auto& c_load_reg = zmm[c_load_regs[i % 3]];
auto c_mem = cox;
if (is_unit_inc)
c_mem += i * nelems_in_cache_line;
@@ -443,14 +443,14 @@ class gemm_class {
// store_c
template <int i, int um_vecs, int idx, int nelems>
EIGEN_ALWAYS_INLINE std::enable_if_t<(i > um_vecs)> write_c(Scalar *cox) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(i > um_vecs)> write_c(Scalar* cox) {
EIGEN_UNUSED_VARIABLE(cox);
}
template <int i, int um_vecs, int idx, int nelems>
EIGEN_ALWAYS_INLINE std::enable_if_t<(i <= um_vecs)> write_c(Scalar *cox) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(i <= um_vecs)> write_c(Scalar* cox) {
if (i < um_vecs) {
auto &c_reg = zmm[c_regs[i + idx * 3]];
auto& c_reg = zmm[c_regs[i + idx * 3]];
auto c_mem = cox;
if (is_unit_inc)
c_mem += i * nelems_in_cache_line;
@@ -495,20 +495,20 @@ class gemm_class {
*/
template <int pow, int a_unroll, int idx>
EIGEN_ALWAYS_INLINE void c_update_1count(Scalar *&cox) {
EIGEN_ALWAYS_INLINE void c_update_1count(Scalar*& cox) {
if (pow >= 4) cox += ldc;
const int um_vecs = numext::div_ceil(a_unroll, nelems_in_cache_line);
auto &alpha_reg = zmm[alpha_load_reg];
auto& alpha_reg = zmm[alpha_load_reg];
scale_load_c<0, um_vecs, idx, a_unroll>(cox, alpha_reg);
write_c<0, um_vecs, idx, a_unroll>(cox);
}
template <int pow, int a_unroll>
EIGEN_ALWAYS_INLINE void c_update_1pow(Scalar *&co1, Scalar *&co2) {
EIGEN_ALWAYS_INLINE void c_update_1pow(Scalar*& co1, Scalar*& co2) {
constexpr int idx = pow / 2;
Scalar *&cox = idx == 0 ? co1 : co2;
Scalar*& cox = idx == 0 ? co1 : co2;
constexpr int max_count = (pow + 1) / 2;
static_assert(max_count <= 4, "Unsupported max_count.");
@@ -520,8 +520,8 @@ class gemm_class {
}
template <int max_b_unroll, int a_unroll, int b_unroll>
EIGEN_ALWAYS_INLINE void c_update(Scalar *&co1, Scalar *&co2) {
auto &alpha_reg = zmm[alpha_load_reg];
EIGEN_ALWAYS_INLINE void c_update(Scalar*& co1, Scalar*& co2) {
auto& alpha_reg = zmm[alpha_load_reg];
co2 = co1 + ldc;
if (!is_alpha1) alpha_reg = pload1<vec>(alpha);
@@ -542,8 +542,8 @@ class gemm_class {
// compute
template <int um, int um_vecs, int idx, int uk, bool fetch_x, bool ktail>
EIGEN_ALWAYS_INLINE std::enable_if_t<(um > um_vecs)> compute(const Scalar *ao, const Scalar *bo, int &fetchA_idx,
int &fetchB_idx, vec &b_reg) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(um > um_vecs)> compute(const Scalar* ao, const Scalar* bo, int& fetchA_idx,
int& fetchB_idx, vec& b_reg) {
EIGEN_UNUSED_VARIABLE(ao);
EIGEN_UNUSED_VARIABLE(bo);
EIGEN_UNUSED_VARIABLE(fetchA_idx);
@@ -552,11 +552,11 @@ class gemm_class {
}
template <int um, int um_vecs, int idx, int uk, bool fetch_x, bool ktail>
EIGEN_ALWAYS_INLINE std::enable_if_t<(um <= um_vecs)> compute(const Scalar *ao, const Scalar *bo, int &fetchA_idx,
int &fetchB_idx, vec &b_reg) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(um <= um_vecs)> compute(const Scalar* ao, const Scalar* bo, int& fetchA_idx,
int& fetchB_idx, vec& b_reg) {
if (um < um_vecs) {
auto &c_reg = zmm[c_regs[um + idx * 3]];
auto &a_reg = zmm[a_regs[um + (uk % 2) * 3]];
auto& c_reg = zmm[c_regs[um + idx * 3]];
auto& a_reg = zmm[a_regs[um + (uk % 2) * 3]];
vfmadd(c_reg, a_reg, b_reg);
@@ -578,25 +578,25 @@ class gemm_class {
// load_a
template <int um, int um_vecs, int uk, int nelems, bool ktail>
EIGEN_ALWAYS_INLINE std::enable_if_t<(um > um_vecs)> load_a(const Scalar *ao) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(um > um_vecs)> load_a(const Scalar* ao) {
EIGEN_UNUSED_VARIABLE(ao);
}
template <int um, int um_vecs, int uk, int nelems, bool ktail>
EIGEN_ALWAYS_INLINE std::enable_if_t<(um <= um_vecs)> load_a(const Scalar *ao) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(um <= um_vecs)> load_a(const Scalar* ao) {
if (um < um_vecs) {
auto &a_reg = zmm[a_regs[um + (uk % 2) * 3]];
const Scalar *a_addr = ao + nelems * (1 + !ktail * !use_less_a_regs + uk) + nelems_in_cache_line * um - a_shift;
auto& a_reg = zmm[a_regs[um + (uk % 2) * 3]];
const Scalar* a_addr = ao + nelems * (1 + !ktail * !use_less_a_regs + uk) + nelems_in_cache_line * um - a_shift;
a_load<nelems>(a_reg, a_addr);
load_a<um + 1, um_vecs, uk, nelems, ktail>(ao);
}
}
template <int uk, int pow, int count, int um_vecs, int b_unroll, bool ktail, bool fetch_x, bool c_fetch>
EIGEN_ALWAYS_INLINE std::enable_if_t<(count > (pow + 1) / 2)> innerkernel_1pow(const Scalar *&aa,
const Scalar *const &ao,
const Scalar *const &bo, Scalar *&co2,
int &fetchA_idx, int &fetchB_idx) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(count > (pow + 1) / 2)> innerkernel_1pow(const Scalar*& aa,
const Scalar* const& ao,
const Scalar* const& bo, Scalar*& co2,
int& fetchA_idx, int& fetchB_idx) {
EIGEN_UNUSED_VARIABLE(aa);
EIGEN_UNUSED_VARIABLE(ao);
EIGEN_UNUSED_VARIABLE(bo);
@@ -606,14 +606,14 @@ class gemm_class {
}
template <int uk, int pow, int count, int um_vecs, int b_unroll, bool ktail, bool fetch_x, bool c_fetch>
EIGEN_ALWAYS_INLINE std::enable_if_t<(count <= (pow + 1) / 2)> innerkernel_1pow(const Scalar *&aa,
const Scalar *const &ao,
const Scalar *const &bo, Scalar *&co2,
int &fetchA_idx, int &fetchB_idx) {
EIGEN_ALWAYS_INLINE std::enable_if_t<(count <= (pow + 1) / 2)> innerkernel_1pow(const Scalar*& aa,
const Scalar* const& ao,
const Scalar* const& bo, Scalar*& co2,
int& fetchA_idx, int& fetchB_idx) {
const int idx = (pow / 2) + count;
if (count < (pow + 1) / 2) {
auto &b_reg = zmm[b_regs[idx % 2]];
auto& b_reg = zmm[b_regs[idx % 2]];
if (fetch_x && uk == 3 && idx == 0) prefetch_x(aa);
if (fetch_x && uk == 3 && idx == 4) aa += 8;
@@ -621,7 +621,7 @@ class gemm_class {
if (b_unroll >= pow) {
compute<0, um_vecs, idx, uk, fetch_x, ktail>(ao, bo, fetchA_idx, fetchB_idx, b_reg);
const Scalar *b_addr = bo + b_unroll * uk + idx + 1 + (b_unroll > 1) * !use_less_b_regs - b_shift;
const Scalar* b_addr = bo + b_unroll * uk + idx + 1 + (b_unroll > 1) * !use_less_b_regs - b_shift;
b_load(b_reg, b_addr);
}
@@ -643,8 +643,8 @@ class gemm_class {
template <int uk, int max_b_unroll, int a_unroll, int b_unroll, bool ktail, bool fetch_x, bool c_fetch,
bool no_a_preload = false>
EIGEN_ALWAYS_INLINE void innerkernel_1uk(const Scalar *&aa, const Scalar *const &ao, const Scalar *const &bo,
Scalar *&co2, int &fetchA_idx, int &fetchB_idx) {
EIGEN_ALWAYS_INLINE void innerkernel_1uk(const Scalar*& aa, const Scalar* const& ao, const Scalar* const& bo,
Scalar*& co2, int& fetchA_idx, int& fetchB_idx) {
const int um_vecs = numext::div_ceil(a_unroll, nelems_in_cache_line);
if (max_b_unroll >= 1)
@@ -701,7 +701,7 @@ class gemm_class {
template <int a_unroll, int b_unroll, int k_factor, int max_b_unroll, int max_k_factor, bool c_fetch,
bool no_a_preload = false>
EIGEN_ALWAYS_INLINE void innerkernel(const Scalar *&aa, const Scalar *&ao, const Scalar *&bo, Scalar *&co2) {
EIGEN_ALWAYS_INLINE void innerkernel(const Scalar*& aa, const Scalar*& ao, const Scalar*& bo, Scalar*& co2) {
int fetchA_idx = 0;
int fetchB_idx = 0;
@@ -731,7 +731,7 @@ class gemm_class {
}
template <int a_unroll, int b_unroll, int max_b_unroll>
EIGEN_ALWAYS_INLINE void kloop(const Scalar *&aa, const Scalar *&ao, const Scalar *&bo, Scalar *&co1, Scalar *&co2) {
EIGEN_ALWAYS_INLINE void kloop(const Scalar*& aa, const Scalar*& ao, const Scalar*& bo, Scalar*& co1, Scalar*& co2) {
const int um_vecs = numext::div_ceil(a_unroll, nelems_in_cache_line);
if (!use_less_a_regs && k > 1)
a_loads<0, 2, 0, um_vecs, a_unroll>(ao);
@@ -795,7 +795,7 @@ class gemm_class {
}
template <int a_unroll, int b_unroll, int max_b_unroll>
EIGEN_ALWAYS_INLINE void nloop(const Scalar *&aa, const Scalar *&ao, const Scalar *&bo, Scalar *&co1, Scalar *&co2) {
EIGEN_ALWAYS_INLINE void nloop(const Scalar*& aa, const Scalar*& ao, const Scalar*& bo, Scalar*& co1, Scalar*& co2) {
// Set A matrix pointer.
ao = a + a_off * a_unroll;
@@ -812,9 +812,9 @@ class gemm_class {
}
template <int a_unroll, int max_a_unroll, int max_b_unroll>
EIGEN_ALWAYS_INLINE void mloop(const Scalar *&ao, const Scalar *&bo, Scalar *&co1, Scalar *&co2) {
EIGEN_ALWAYS_INLINE void mloop(const Scalar*& ao, const Scalar*& bo, Scalar*& co1, Scalar*& co2) {
// Set prefetch A pointers.
const Scalar *aa = a + a_unroll * a_stride;
const Scalar* aa = a + a_unroll * a_stride;
// Set C matrix pointers.
co1 = c;
@@ -856,10 +856,10 @@ class gemm_class {
a -= -a_shift;
b -= -b_shift;
const Scalar *ao = nullptr;
const Scalar *bo = nullptr;
Scalar *co1 = nullptr;
Scalar *co2 = nullptr;
const Scalar* ao = nullptr;
const Scalar* bo = nullptr;
Scalar* co1 = nullptr;
Scalar* co2 = nullptr;
// Main m-loop.
for (; m >= max_a_unroll; m -= max_a_unroll) mloop<max_a_unroll, max_a_unroll, max_b_unroll>(ao, bo, co1, co2);
@@ -883,8 +883,8 @@ class gemm_class {
}
}
gemm_class(Index m_, Index n_, Index k_, Index ldc_, Index inc_, const Scalar *alpha_, const Scalar *a_,
const Scalar *b_, Scalar *c_, bool is_alpha1_, bool is_beta0_, Index a_stride_, Index b_stride_,
gemm_class(Index m_, Index n_, Index k_, Index ldc_, Index inc_, const Scalar* alpha_, const Scalar* a_,
const Scalar* b_, Scalar* c_, bool is_alpha1_, bool is_beta0_, Index a_stride_, Index b_stride_,
Index a_off_, Index b_off_)
: m(m_),
n(n_),
@@ -937,8 +937,8 @@ class gemm_class {
// max_a_unroll: 24, 16, 8, 4, 2, 1
// max_b_unroll: 8, 4, 2, 1
template <typename Scalar, int max_a_unroll, int max_b_unroll, bool is_alpha1, bool is_beta0, bool is_unit_inc>
EIGEN_DONT_INLINE void gemm_kern_avx512(Index m, Index n, Index k, Scalar *alpha, const Scalar *a, const Scalar *b,
Scalar *c, Index ldc, Index inc = 1, Index a_stride = -1, Index b_stride = -1,
EIGEN_DONT_INLINE void gemm_kern_avx512(Index m, Index n, Index k, Scalar* alpha, const Scalar* a, const Scalar* b,
Scalar* c, Index ldc, Index inc = 1, Index a_stride = -1, Index b_stride = -1,
Index a_off = 0, Index b_off = 0) {
if (a_stride == -1) a_stride = k;
if (b_stride == -1) b_stride = k;
@@ -972,13 +972,13 @@ struct gemm_pack_rhs<Scalar, Index, DataMapper, 8, ColMajor, Conjugate, PanelMod
typedef typename packet_traits<Scalar>::type Packet;
typedef typename DataMapper::LinearMapper LinearMapper;
enum { PacketSize = packet_traits<Scalar>::size };
EIGEN_DONT_INLINE void operator()(Scalar *blockB, const DataMapper &rhs, Index depth, Index cols, Index stride = 0,
EIGEN_DONT_INLINE void operator()(Scalar* blockB, const DataMapper& rhs, Index depth, Index cols, Index stride = 0,
Index offset = 0);
};
template <typename Scalar, typename Index, typename DataMapper, bool Conjugate, bool PanelMode>
EIGEN_DONT_INLINE void gemm_pack_rhs<Scalar, Index, DataMapper, 8, ColMajor, Conjugate, PanelMode>::operator()(
Scalar *blockB, const DataMapper &rhs, Index depth, Index cols, Index stride, Index offset) {
Scalar* blockB, const DataMapper& rhs, Index depth, Index cols, Index stride, Index offset) {
constexpr int nr = 8;
EIGEN_ASM_COMMENT("EIGEN PRODUCT PACK RHS COLMAJOR");
EIGEN_UNUSED_VARIABLE(stride);
@@ -1106,7 +1106,7 @@ struct gemm_pack_rhs<Scalar, Index, DataMapper, 8, RowMajor, Conjugate, PanelMod
HalfPacketSize = unpacket_traits<HalfPacket>::size,
QuarterPacketSize = unpacket_traits<QuarterPacket>::size
};
EIGEN_DONT_INLINE void operator()(Scalar *blockB, const DataMapper &rhs, Index depth, Index cols, Index stride = 0,
EIGEN_DONT_INLINE void operator()(Scalar* blockB, const DataMapper& rhs, Index depth, Index cols, Index stride = 0,
Index offset = 0) {
constexpr int nr = 8;
EIGEN_ASM_COMMENT("EIGEN PRODUCT PACK RHS ROWMAJOR");
@@ -1205,33 +1205,32 @@ struct gemm_pack_rhs<Scalar, Index, DataMapper, 8, RowMajor, Conjugate, PanelMod
template <typename Scalar, typename Index, typename DataMapper, int mr, bool ConjugateLhs, bool ConjugateRhs>
struct gebp_kernel<Scalar, Scalar, Index, DataMapper, mr, 8, ConjugateLhs, ConjugateRhs> {
EIGEN_ALWAYS_INLINE void operator()(const DataMapper &res, const Scalar *blockA, const Scalar *blockB, Index rows,
EIGEN_ALWAYS_INLINE void operator()(const DataMapper& res, const Scalar* blockA, const Scalar* blockB, Index rows,
Index depth, Index cols, Scalar alpha, Index strideA = -1, Index strideB = -1,
Index offsetA = 0, Index offsetB = 0);
};
template <typename Scalar, typename Index, typename DataMapper, int mr, bool ConjugateLhs, bool ConjugateRhs>
EIGEN_ALWAYS_INLINE void gebp_kernel<Scalar, Scalar, Index, DataMapper, mr, 8, ConjugateLhs, ConjugateRhs>::operator()(
const DataMapper &res, const Scalar *blockA, const Scalar *blockB, Index rows, Index depth, Index cols,
const DataMapper& res, const Scalar* blockA, const Scalar* blockB, Index rows, Index depth, Index cols,
Scalar alpha, Index strideA, Index strideB, Index offsetA, Index offsetB) {
if (res.incr() == 1) {
if (alpha == 1) {
gemm_kern_avx512<Scalar, mr, 8, true, false, true>(rows, cols, depth, &alpha, blockA, blockB,
(Scalar *)res.data(), res.stride(), res.incr(), strideA,
strideB, offsetA, offsetB);
gemm_kern_avx512<Scalar, mr, 8, true, false, true>(rows, cols, depth, &alpha, blockA, blockB, (Scalar*)res.data(),
res.stride(), res.incr(), strideA, strideB, offsetA, offsetB);
} else {
gemm_kern_avx512<Scalar, mr, 8, false, false, true>(rows, cols, depth, &alpha, blockA, blockB,
(Scalar *)res.data(), res.stride(), res.incr(), strideA,
(Scalar*)res.data(), res.stride(), res.incr(), strideA,
strideB, offsetA, offsetB);
}
} else {
if (alpha == 1) {
gemm_kern_avx512<Scalar, mr, 8, true, false, false>(rows, cols, depth, &alpha, blockA, blockB,
(Scalar *)res.data(), res.stride(), res.incr(), strideA,
(Scalar*)res.data(), res.stride(), res.incr(), strideA,
strideB, offsetA, offsetB);
} else {
gemm_kern_avx512<Scalar, mr, 8, false, false, false>(rows, cols, depth, &alpha, blockA, blockB,
(Scalar *)res.data(), res.stride(), res.incr(), strideA,
(Scalar*)res.data(), res.stride(), res.incr(), strideA,
strideB, offsetA, offsetB);
}
}

View File

@@ -100,7 +100,6 @@ struct packet_traits<half> : default_packet_traits {
HasCos = EIGEN_FAST_MATH,
HasTanh = EIGEN_FAST_MATH,
HasErf = EIGEN_FAST_MATH,
HasBlend = 0
};
};
#endif
@@ -118,7 +117,6 @@ struct packet_traits<float> : default_packet_traits {
HasMin = 1,
HasMax = 1,
HasConj = 1,
HasBlend = 1,
HasSin = EIGEN_FAST_MATH,
HasCos = EIGEN_FAST_MATH,
HasACos = 1,
@@ -151,7 +149,6 @@ struct packet_traits<double> : default_packet_traits {
Vectorizable = 1,
AlignedOnScalar = 1,
size = 8,
HasBlend = 1,
HasSqrt = 1,
HasRsqrt = 1,
HasCbrt = 1,
@@ -176,7 +173,7 @@ template <>
struct packet_traits<int> : default_packet_traits {
typedef Packet16i type;
typedef Packet8i half;
enum { Vectorizable = 1, AlignedOnScalar = 1, HasBlend = 0, HasCmp = 1, HasDiv = 1, size = 16 };
enum { Vectorizable = 1, AlignedOnScalar = 1, HasCmp = 1, HasDiv = 1, size = 16 };
};
template <>
@@ -2512,7 +2509,6 @@ struct packet_traits<bfloat16> : default_packet_traits {
Vectorizable = 1,
AlignedOnScalar = 1,
size = 16,
HasBlend = 0,
HasInsert = 1,
HasSin = EIGEN_FAST_MATH,
HasCos = EIGEN_FAST_MATH,

View File

@@ -59,7 +59,6 @@ struct packet_traits<half> : default_packet_traits {
HasCos = EIGEN_FAST_MATH,
HasTanh = EIGEN_FAST_MATH,
HasErf = 0, // EIGEN_FAST_MATH,
HasBlend = 0
};
};

View File

@@ -121,7 +121,7 @@ int64_t avx512_trsm_cutoff(int64_t L2Size, int64_t N, double L2Cap) {
* Used by gemmKernel for the case A/B row-major and C col-major.
*/
template <typename Scalar, typename vec, int64_t unrollM, int64_t unrollN, bool remM, bool remN>
EIGEN_ALWAYS_INLINE void transStoreC(PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> &zmm, Scalar *C_arr,
EIGEN_ALWAYS_INLINE void transStoreC(PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS>& zmm, Scalar* C_arr,
int64_t LDC, int64_t remM_ = 0, int64_t remN_ = 0) {
EIGEN_UNUSED_VARIABLE(remN_);
EIGEN_UNUSED_VARIABLE(remM_);
@@ -221,7 +221,7 @@ EIGEN_ALWAYS_INLINE void transStoreC(PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_
* handleKRem: Handle arbitrary K? This is not needed for trsm.
*/
template <typename Scalar, bool isARowMajor, bool isCRowMajor, bool isAdd, bool handleKRem>
void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t N, int64_t K, int64_t LDA, int64_t LDB,
void gemmKernel(Scalar* A_arr, Scalar* B_arr, Scalar* C_arr, int64_t M, int64_t N, int64_t K, int64_t LDA, int64_t LDB,
int64_t LDC) {
using urolls = unrolls::gemm<Scalar, isAdd>;
constexpr int64_t U3 = urolls::PacketSize * 3;
@@ -264,8 +264,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
}
}
if (M - i >= 4) { // Note: this block assumes EIGEN_AVX_MAX_NUM_ROW = 8. Should be removed otherwise
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<3, 4>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -294,8 +294,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 4;
}
if (M - i >= 2) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<3, 2>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -324,8 +324,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 2;
}
if (M - i > 0) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<3, 1>(zmm);
{
@@ -387,8 +387,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
}
}
if (M - i >= 4) { // Note: this block assumes EIGEN_AVX_MAX_NUM_ROW = 8. Should be removed otherwise
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<2, 4>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -417,8 +417,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 4;
}
if (M - i >= 2) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<2, 2>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -447,8 +447,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 2;
}
if (M - i > 0) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<2, 1>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -508,8 +508,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
}
}
if (M - i >= 4) { // Note: this block assumes EIGEN_AVX_MAX_NUM_ROW = 8. Should be removed otherwise
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<1, 4>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -538,8 +538,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 4;
}
if (M - i >= 2) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<1, 2>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -568,8 +568,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 2;
}
if (M - i > 0) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<1, 1>(zmm);
{
@@ -603,8 +603,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
constexpr int64_t EIGEN_AVX_MAX_B_LOAD = EIGEN_AVX_B_LOAD_SETS * 1;
int64_t i = 0;
for (; i < M_; i += EIGEN_AVX_MAX_NUM_ROW) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<1, EIGEN_AVX_MAX_NUM_ROW>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -632,8 +632,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
}
}
if (M - i >= 4) { // Note: this block assumes EIGEN_AVX_MAX_NUM_ROW = 8. Should be removed otherwise
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<1, 4>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -662,8 +662,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 4;
}
if (M - i >= 2) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<1, 2>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -692,8 +692,8 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
i += 2;
}
if (M - i > 0) {
Scalar *A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar *B_t = &B_arr[0 * LDB + j];
Scalar* A_t = &A_arr[idA<isARowMajor>(i, 0, LDA)];
Scalar* B_t = &B_arr[0 * LDB + j];
PacketBlock<vec, EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS> zmm;
urolls::template setzero<1, 1>(zmm);
for (int64_t k = 0; k < K_; k += EIGEN_AVX_MAX_K_UNROL) {
@@ -732,7 +732,7 @@ void gemmKernel(Scalar *A_arr, Scalar *B_arr, Scalar *C_arr, int64_t M, int64_t
* The B matrix (RHS) is assumed to be row-major
*/
template <typename Scalar, typename vec, int64_t unrollM, bool isARowMajor, bool isFWDSolve, bool isUnitDiag>
EIGEN_ALWAYS_INLINE void triSolveKernel(Scalar *A_arr, Scalar *B_arr, int64_t K, int64_t LDA, int64_t LDB) {
EIGEN_ALWAYS_INLINE void triSolveKernel(Scalar* A_arr, Scalar* B_arr, int64_t K, int64_t LDA, int64_t LDB) {
static_assert(unrollM <= EIGEN_AVX_MAX_NUM_ROW, "unrollM should be equal to EIGEN_AVX_MAX_NUM_ROW");
using urolls = unrolls::trsm<Scalar>;
constexpr int64_t U3 = urolls::PacketSize * 3;
@@ -782,7 +782,7 @@ EIGEN_ALWAYS_INLINE void triSolveKernel(Scalar *A_arr, Scalar *B_arr, int64_t K,
* The B matrix (RHS) is assumed to be row-major
*/
template <typename Scalar, bool isARowMajor, bool isFWDSolve, bool isUnitDiag>
void triSolveKernelLxK(Scalar *A_arr, Scalar *B_arr, int64_t M, int64_t K, int64_t LDA, int64_t LDB) {
void triSolveKernelLxK(Scalar* A_arr, Scalar* B_arr, int64_t M, int64_t K, int64_t LDA, int64_t LDB) {
// Note: this assumes EIGEN_AVX_MAX_NUM_ROW = 8. Unrolls should be adjusted
// accordingly if EIGEN_AVX_MAX_NUM_ROW is smaller.
using vec = typename std::conditional<std::is_same<Scalar, float>::value, vecFullFloat, vecFullDouble>::type;
@@ -813,7 +813,7 @@ void triSolveKernelLxK(Scalar *A_arr, Scalar *B_arr, int64_t M, int64_t K, int64
*
*/
template <typename Scalar, bool toTemp = true, bool remM = false>
EIGEN_ALWAYS_INLINE void copyBToRowMajor(Scalar *B_arr, int64_t LDB, int64_t K, Scalar *B_temp, int64_t LDB_,
EIGEN_ALWAYS_INLINE void copyBToRowMajor(Scalar* B_arr, int64_t LDB, int64_t K, Scalar* B_temp, int64_t LDB_,
int64_t remM_ = 0) {
EIGEN_UNUSED_VARIABLE(remM_);
using urolls = unrolls::transB<Scalar>;
@@ -900,7 +900,7 @@ EIGEN_ALWAYS_INLINE void copyBToRowMajor(Scalar *B_arr, int64_t LDB, int64_t K,
*/
template <typename Scalar, bool isARowMajor = true, bool isBRowMajor = true, bool isFWDSolve = true,
bool isUnitDiag = false>
void triSolve(Scalar *A_arr, Scalar *B_arr, int64_t M, int64_t numRHS, int64_t LDA, int64_t LDB) {
void triSolve(Scalar* A_arr, Scalar* B_arr, int64_t M, int64_t numRHS, int64_t LDA, int64_t LDB) {
constexpr int64_t psize = packet_traits<Scalar>::size;
/**
* The values for kB, numM were determined experimentally.
@@ -919,7 +919,7 @@ void triSolve(Scalar *A_arr, Scalar *B_arr, int64_t M, int64_t numRHS, int64_t L
constexpr int64_t numM = 8 * EIGEN_AVX_MAX_NUM_ROW;
int64_t sizeBTemp = 0;
Scalar *B_temp = NULL;
Scalar* B_temp = NULL;
EIGEN_IF_CONSTEXPR(!isBRowMajor) {
/**
* If B is col-major, we copy it to a fixed-size temporary array of size at most ~numM*kB and
@@ -929,7 +929,7 @@ void triSolve(Scalar *A_arr, Scalar *B_arr, int64_t M, int64_t numRHS, int64_t L
sizeBTemp = (((std::min(kB, numRHS) + psize - 1) / psize + 4) * psize) * numM;
}
EIGEN_IF_CONSTEXPR(!isBRowMajor) B_temp = (Scalar *)handmade_aligned_malloc(sizeof(Scalar) * sizeBTemp, 64);
EIGEN_IF_CONSTEXPR(!isBRowMajor) B_temp = (Scalar*)handmade_aligned_malloc(sizeof(Scalar) * sizeBTemp, 64);
for (int64_t k = 0; k < numRHS; k += kB) {
int64_t bK = numRHS - k > kB ? kB : numRHS - k;
@@ -1070,19 +1070,19 @@ struct trsmKernelR;
template <typename Index, int Mode, int TriStorageOrder>
struct trsmKernelR<float, Index, Mode, false, TriStorageOrder, 1, true> {
static void kernel(Index size, Index otherSize, const float *_tri, Index triStride, float *_other, Index otherIncr,
static void kernel(Index size, Index otherSize, const float* _tri, Index triStride, float* _other, Index otherIncr,
Index otherStride);
};
template <typename Index, int Mode, int TriStorageOrder>
struct trsmKernelR<double, Index, Mode, false, TriStorageOrder, 1, true> {
static void kernel(Index size, Index otherSize, const double *_tri, Index triStride, double *_other, Index otherIncr,
static void kernel(Index size, Index otherSize, const double* _tri, Index triStride, double* _other, Index otherIncr,
Index otherStride);
};
template <typename Index, int Mode, int TriStorageOrder>
EIGEN_DONT_INLINE void trsmKernelR<float, Index, Mode, false, TriStorageOrder, 1, true>::kernel(
Index size, Index otherSize, const float *_tri, Index triStride, float *_other, Index otherIncr,
Index size, Index otherSize, const float* _tri, Index triStride, float* _other, Index otherIncr,
Index otherStride) {
EIGEN_UNUSED_VARIABLE(otherIncr);
#ifdef EIGEN_RUNTIME_NO_MALLOC
@@ -1093,12 +1093,12 @@ EIGEN_DONT_INLINE void trsmKernelR<float, Index, Mode, false, TriStorageOrder, 1
}
#endif
triSolve<float, TriStorageOrder != RowMajor, true, (Mode & Lower) != Lower, (Mode & UnitDiag) != 0>(
const_cast<float *>(_tri), _other, size, otherSize, triStride, otherStride);
const_cast<float*>(_tri), _other, size, otherSize, triStride, otherStride);
}
template <typename Index, int Mode, int TriStorageOrder>
EIGEN_DONT_INLINE void trsmKernelR<double, Index, Mode, false, TriStorageOrder, 1, true>::kernel(
Index size, Index otherSize, const double *_tri, Index triStride, double *_other, Index otherIncr,
Index size, Index otherSize, const double* _tri, Index triStride, double* _other, Index otherIncr,
Index otherStride) {
EIGEN_UNUSED_VARIABLE(otherIncr);
#ifdef EIGEN_RUNTIME_NO_MALLOC
@@ -1109,7 +1109,7 @@ EIGEN_DONT_INLINE void trsmKernelR<double, Index, Mode, false, TriStorageOrder,
}
#endif
triSolve<double, TriStorageOrder != RowMajor, true, (Mode & Lower) != Lower, (Mode & UnitDiag) != 0>(
const_cast<double *>(_tri), _other, size, otherSize, triStride, otherStride);
const_cast<double*>(_tri), _other, size, otherSize, triStride, otherStride);
}
#endif // (EIGEN_USE_AVX512_TRSM_R_KERNELS)
@@ -1121,19 +1121,19 @@ struct trsmKernelL;
template <typename Index, int Mode, int TriStorageOrder>
struct trsmKernelL<float, Index, Mode, false, TriStorageOrder, 1, true> {
static void kernel(Index size, Index otherSize, const float *_tri, Index triStride, float *_other, Index otherIncr,
static void kernel(Index size, Index otherSize, const float* _tri, Index triStride, float* _other, Index otherIncr,
Index otherStride);
};
template <typename Index, int Mode, int TriStorageOrder>
struct trsmKernelL<double, Index, Mode, false, TriStorageOrder, 1, true> {
static void kernel(Index size, Index otherSize, const double *_tri, Index triStride, double *_other, Index otherIncr,
static void kernel(Index size, Index otherSize, const double* _tri, Index triStride, double* _other, Index otherIncr,
Index otherStride);
};
template <typename Index, int Mode, int TriStorageOrder>
EIGEN_DONT_INLINE void trsmKernelL<float, Index, Mode, false, TriStorageOrder, 1, true>::kernel(
Index size, Index otherSize, const float *_tri, Index triStride, float *_other, Index otherIncr,
Index size, Index otherSize, const float* _tri, Index triStride, float* _other, Index otherIncr,
Index otherStride) {
EIGEN_UNUSED_VARIABLE(otherIncr);
#ifdef EIGEN_RUNTIME_NO_MALLOC
@@ -1144,12 +1144,12 @@ EIGEN_DONT_INLINE void trsmKernelL<float, Index, Mode, false, TriStorageOrder, 1
}
#endif
triSolve<float, TriStorageOrder == RowMajor, false, (Mode & Lower) == Lower, (Mode & UnitDiag) != 0>(
const_cast<float *>(_tri), _other, size, otherSize, triStride, otherStride);
const_cast<float*>(_tri), _other, size, otherSize, triStride, otherStride);
}
template <typename Index, int Mode, int TriStorageOrder>
EIGEN_DONT_INLINE void trsmKernelL<double, Index, Mode, false, TriStorageOrder, 1, true>::kernel(
Index size, Index otherSize, const double *_tri, Index triStride, double *_other, Index otherIncr,
Index size, Index otherSize, const double* _tri, Index triStride, double* _other, Index otherIncr,
Index otherStride) {
EIGEN_UNUSED_VARIABLE(otherIncr);
#ifdef EIGEN_RUNTIME_NO_MALLOC
@@ -1160,7 +1160,7 @@ EIGEN_DONT_INLINE void trsmKernelL<double, Index, Mode, false, TriStorageOrder,
}
#endif
triSolve<double, TriStorageOrder == RowMajor, false, (Mode & Lower) == Lower, (Mode & UnitDiag) != 0>(
const_cast<double *>(_tri), _other, size, otherSize, triStride, otherStride);
const_cast<double*>(_tri), _other, size, otherSize, triStride, otherStride);
}
#endif // EIGEN_USE_AVX512_TRSM_L_KERNELS

View File

@@ -84,7 +84,6 @@ struct packet_traits<float> : default_packet_traits {
HasIGammac = 1,
HasBetaInc = 1,
HasBlend = 0,
HasFloor = 1,
HasCmp = EIGEN_HAS_GPU_DEVICE_FUNCTIONS
};
@@ -117,7 +116,6 @@ struct packet_traits<double> : default_packet_traits {
HasGammaSampleDerAlpha = 1,
HasIGammac = 1,
HasBetaInc = 1,
HasBlend = 0,
};
};

View File

@@ -157,10 +157,7 @@ struct packet_traits<float> : default_packet_traits {
HasMax = 1,
HasConj = 0,
HasSetLinear = 0,
HasBlend = 0,
HasDiv = 0,
HasSin = 0,
HasCos = 0,
HasACos = 0,

View File

@@ -200,7 +200,6 @@ struct packet_traits<float> : default_packet_traits {
HasTanh = EIGEN_FAST_MATH,
HasErf = EIGEN_FAST_MATH,
HasErfc = EIGEN_FAST_MATH,
HasBlend = 1,
HasSign = 0 // The manually vectorized version is slightly slower for SSE.
};
};
@@ -230,7 +229,6 @@ struct packet_traits<double> : default_packet_traits {
HasCbrt = 1,
HasATan = 1,
HasATanh = 1,
HasBlend = 1
};
};
template <>
@@ -245,7 +243,6 @@ struct packet_traits<int> : default_packet_traits {
HasCmp = 1,
HasDiv = 1,
HasShift = 1,
HasBlend = 1
};
};
template <>
@@ -257,11 +254,9 @@ struct packet_traits<uint32_t> : default_packet_traits {
AlignedOnScalar = 1,
size = 4,
HasDiv = 0,
HasNegate = 0,
HasCmp = 1,
HasShift = 1,
HasBlend = 1
};
};
template <>
@@ -273,10 +268,8 @@ struct packet_traits<int64_t> : default_packet_traits {
AlignedOnScalar = 1,
size = 2,
HasDiv = 0,
HasCmp = 1,
HasShift = 1,
HasBlend = 1
};
};
#endif

View File

@@ -30,7 +30,7 @@ namespace Eigen {
namespace internal {
template <int has_blend, int lengths>
template <int lengths>
struct sycl_packet_traits : default_packet_traits {
enum {
Vectorizable = 1,
@@ -60,7 +60,6 @@ struct sycl_packet_traits : default_packet_traits {
HasIGamma = 0,
HasIGammac = 0,
HasBetaInc = 0,
HasBlend = has_blend,
// This flag is used to indicate whether packet comparison is supported.
// pcmp_eq, pcmp_lt and pcmp_le should be defined for it to be true.
HasCmp = 1,
@@ -78,19 +77,19 @@ struct sycl_packet_traits : default_packet_traits {
};
#ifdef SYCL_DEVICE_ONLY
#define SYCL_PACKET_TRAITS(packet_type, has_blend, unpacket_type, lengths) \
template <> \
struct packet_traits<unpacket_type> : sycl_packet_traits<has_blend, lengths> { \
typedef packet_type type; \
typedef packet_type half; \
#define SYCL_PACKET_TRAITS(packet_type, unpacket_type, lengths) \
template <> \
struct packet_traits<unpacket_type> : sycl_packet_traits<lengths> { \
typedef packet_type type; \
typedef packet_type half; \
};
SYCL_PACKET_TRAITS(cl::sycl::cl_half8, 1, Eigen::half, 8)
SYCL_PACKET_TRAITS(cl::sycl::cl_half8, 1, const Eigen::half, 8)
SYCL_PACKET_TRAITS(cl::sycl::cl_float4, 1, float, 4)
SYCL_PACKET_TRAITS(cl::sycl::cl_float4, 1, const float, 4)
SYCL_PACKET_TRAITS(cl::sycl::cl_double2, 0, double, 2)
SYCL_PACKET_TRAITS(cl::sycl::cl_double2, 0, const double, 2)
SYCL_PACKET_TRAITS(cl::sycl::cl_half8, Eigen::half, 8)
SYCL_PACKET_TRAITS(cl::sycl::cl_half8, const Eigen::half, 8)
SYCL_PACKET_TRAITS(cl::sycl::cl_float4, float, 4)
SYCL_PACKET_TRAITS(cl::sycl::cl_float4, const float, 4)
SYCL_PACKET_TRAITS(cl::sycl::cl_double2, double, 2)
SYCL_PACKET_TRAITS(cl::sycl::cl_double2, const double, 2)
#undef SYCL_PACKET_TRAITS
// Make sure this is only available when targeting a GPU: we don't want to
@@ -135,14 +134,14 @@ template <typename PacketReturnType, int PacketSize>
struct PacketWrapper {
typedef typename ::Eigen::internal::unpacket_traits<PacketReturnType>::type Scalar;
template <typename Index>
EIGEN_DEVICE_FUNC static Scalar scalarize(Index, PacketReturnType &) {
EIGEN_DEVICE_FUNC static Scalar scalarize(Index, PacketReturnType&) {
eigen_assert(false && "THERE IS NO PACKETIZE VERSION FOR THE CHOSEN TYPE");
abort();
}
EIGEN_DEVICE_FUNC static PacketReturnType convert_to_packet_type(Scalar in, Scalar) {
return ::Eigen::internal::template plset<PacketReturnType>(in);
}
EIGEN_DEVICE_FUNC static void set_packet(PacketReturnType, Scalar *) {
EIGEN_DEVICE_FUNC static void set_packet(PacketReturnType, Scalar*) {
eigen_assert(false && "THERE IS NO PACKETIZE VERSION FOR THE CHOSEN TYPE");
abort();
}
@@ -153,7 +152,7 @@ template <typename PacketReturnType>
struct PacketWrapper<PacketReturnType, 4> {
typedef typename ::Eigen::internal::unpacket_traits<PacketReturnType>::type Scalar;
template <typename Index>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Scalar scalarize(Index index, PacketReturnType &in) {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Scalar scalarize(Index index, PacketReturnType& in) {
switch (index) {
case 0:
return in.x();
@@ -174,7 +173,7 @@ struct PacketWrapper<PacketReturnType, 4> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static PacketReturnType convert_to_packet_type(Scalar in, Scalar other) {
return PacketReturnType(in, other, other, other);
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void set_packet(PacketReturnType &lhs, Scalar *rhs) {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void set_packet(PacketReturnType& lhs, Scalar* rhs) {
lhs = PacketReturnType(rhs[0], rhs[1], rhs[2], rhs[3]);
}
};
@@ -183,20 +182,20 @@ template <typename PacketReturnType>
struct PacketWrapper<PacketReturnType, 1> {
typedef typename ::Eigen::internal::unpacket_traits<PacketReturnType>::type Scalar;
template <typename Index>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Scalar scalarize(Index, PacketReturnType &in) {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Scalar scalarize(Index, PacketReturnType& in) {
return in;
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static PacketReturnType convert_to_packet_type(Scalar in, Scalar) {
return PacketReturnType(in);
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void set_packet(PacketReturnType &lhs, Scalar *rhs) { lhs = rhs[0]; }
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void set_packet(PacketReturnType& lhs, Scalar* rhs) { lhs = rhs[0]; }
};
template <typename PacketReturnType>
struct PacketWrapper<PacketReturnType, 2> {
typedef typename ::Eigen::internal::unpacket_traits<PacketReturnType>::type Scalar;
template <typename Index>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Scalar scalarize(Index index, PacketReturnType &in) {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Scalar scalarize(Index index, PacketReturnType& in) {
switch (index) {
case 0:
return in.x();
@@ -213,7 +212,7 @@ struct PacketWrapper<PacketReturnType, 2> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static PacketReturnType convert_to_packet_type(Scalar in, Scalar other) {
return PacketReturnType(in, other);
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void set_packet(PacketReturnType &lhs, Scalar *rhs) {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void set_packet(PacketReturnType& lhs, Scalar* rhs) {
lhs = PacketReturnType(rhs[0], rhs[1]);
}
};

View File

@@ -31,13 +31,9 @@ using Packet16i = detail::VectorType<int32_t, 16>;
using Packet8l = detail::VectorType<int64_t, 8>;
// --- packet_traits specializations ---
template <>
struct packet_traits<float> : default_packet_traits {
using type = Packet16f;
using half = Packet16f;
struct generic_float_packet_traits : default_packet_traits {
enum {
Vectorizable = 1,
size = 16,
AlignedOnScalar = 1,
HasAdd = 1,
HasSub = 1,
@@ -46,7 +42,8 @@ struct packet_traits<float> : default_packet_traits {
HasNegate = 1,
HasAbs = 1,
HasRound = 1,
HasMinMax = 1,
HasMin = 1,
HasMax = 1,
HasCmp = 1,
HasSet1 = 1,
HasCast = 1,
@@ -80,12 +77,24 @@ struct packet_traits<float> : default_packet_traits {
};
template <>
struct packet_traits<double> : default_packet_traits {
struct packet_traits<float> : generic_float_packet_traits {
using type = Packet16f;
using half = Packet16f;
enum {
size = 16,
};
};
template <>
struct packet_traits<double> : generic_float_packet_traits {
using type = Packet8d;
using half = Packet8d;
enum { size = 8, HasACos = 0, HasASin = 0 };
};
struct generic_integer_packet_traits : default_packet_traits {
enum {
Vectorizable = 1,
size = 8,
AlignedOnScalar = 1,
HasAdd = 1,
HasSub = 1,
@@ -93,157 +102,85 @@ struct packet_traits<double> : default_packet_traits {
HasDiv = 1,
HasNegate = 1,
HasAbs = 1,
HasRound = 1,
HasMinMax = 1,
HasMin = 1,
HasMax = 1,
HasCmp = 1,
HasSet1 = 1,
HasCast = 1,
HasBitwise = 1,
HasRedux = 1,
HasSign = 1,
// Set remaining to 0
HasRound = 1,
HasSqrt = 0,
HasRsqrt = 0,
HasReciprocal = 0,
HasArg = 0,
HasConj = 1,
// Math functions
HasReciprocal = 1,
HasSin = 1,
HasCos = 1,
HasACos = 0,
HasASin = 0,
HasATan = 1,
HasATanh = 1,
HasLog = 1,
HasLog1p = 1,
HasExpm1 = 1,
HasExp = 1,
HasPow = 1,
HasNdtri = 1,
HasBessel = 1,
HasSqrt = 1,
HasRsqrt = 1,
HasCbrt = 1,
HasTanh = 1,
HasErf = 1,
HasErfc = 1
HasExp = 0,
HasLog = 0,
HasSin = 0,
HasCos = 0,
};
};
template <>
struct packet_traits<int32_t> : default_packet_traits {
struct packet_traits<int32_t> : generic_integer_packet_traits {
using type = Packet16i;
using half = Packet16i;
enum {
Vectorizable = 1,
size = 16,
AlignedOnScalar = 1,
HasAdd = 1,
HasSub = 1,
HasMul = 1,
HasDiv = 1,
HasNegate = 1,
HasAbs = 1,
HasMinMax = 1,
HasCmp = 1,
HasSet1 = 1,
HasCast = 1,
HasBitwise = 1,
HasRedux = 1,
// Set remaining to 0
HasRound = 1,
HasSqrt = 0,
HasRsqrt = 0,
HasReciprocal = 0,
HasArg = 0,
HasConj = 1,
HasExp = 0,
HasLog = 0,
HasSin = 0,
HasCos = 0,
};
};
template <>
struct packet_traits<int64_t> : default_packet_traits {
struct packet_traits<int64_t> : generic_integer_packet_traits {
using type = Packet8l;
using half = Packet8l;
enum {
Vectorizable = 1,
size = 8,
AlignedOnScalar = 1,
HasAdd = 1,
HasSub = 1,
HasMul = 1,
HasDiv = 1,
HasNegate = 1,
HasAbs = 1,
HasMinMax = 1,
HasCmp = 1,
HasSet1 = 1,
HasCast = 1,
HasBitwise = 1,
HasRedux = 1,
// Set remaining to 0
HasRound = 1,
HasSqrt = 0,
HasRsqrt = 0,
HasReciprocal = 0,
HasArg = 0,
HasConj = 1,
HasExp = 0,
HasLog = 0,
HasSin = 0,
HasCos = 0,
};
};
// --- unpacket_traits specializations ---
struct generic_unpacket_traits : default_unpacket_traits {
enum {
alignment = EIGEN_GENERIC_VECTOR_SIZE_BYTES,
vectorizable = true,
};
};
template <>
struct unpacket_traits<Packet16f> {
struct unpacket_traits<Packet16f> : generic_unpacket_traits {
using type = float;
using half = Packet16f;
using integer_packet = Packet16i;
enum {
size = 16,
alignment = EIGEN_GENERIC_VECTOR_SIZE_BYTES,
vectorizable = true,
masked_load_available = false,
masked_store_available = false
};
};
template <>
struct unpacket_traits<Packet8d> {
struct unpacket_traits<Packet8d> : generic_unpacket_traits {
using type = double;
using half = Packet8d;
using integer_packet = Packet8l;
enum {
size = 8,
alignment = EIGEN_GENERIC_VECTOR_SIZE_BYTES,
vectorizable = true,
masked_load_available = false,
masked_store_available = false
};
};
template <>
struct unpacket_traits<Packet16i> {
struct unpacket_traits<Packet16i> : generic_unpacket_traits {
using type = int32_t;
using half = Packet16i;
enum {
size = 16,
alignment = EIGEN_GENERIC_VECTOR_SIZE_BYTES,
vectorizable = true,
masked_load_available = false,
masked_store_available = false
};
};
template <>
struct unpacket_traits<Packet8l> {
struct unpacket_traits<Packet8l> : generic_unpacket_traits {
using type = int64_t;
using half = Packet8l;
enum {
size = 8,
alignment = EIGEN_GENERIC_VECTOR_SIZE_BYTES,
vectorizable = true,
masked_load_available = false,
masked_store_available = false
};
};
@@ -494,7 +431,7 @@ EIGEN_CLANG_PACKET_BITWISE_FLOAT(Packet8d, detail::pcast_double_to_long, detail:
template <> \
EIGEN_STRONG_INLINE PACKET_TYPE pselect<PACKET_TYPE>(const PACKET_TYPE& mask, const PACKET_TYPE& a, \
const PACKET_TYPE& b) { \
return __builtin_elementwise_abs(mask) == 0 ? b : a; \
return mask != 0 ? a : b; \
}
EIGEN_CLANG_PACKET_ELEMENTWISE(Packet16f)
@@ -589,11 +526,11 @@ EIGEN_CLANG_PACKET_MADD(Packet8d)
EIGEN_STRONG_INLINE PACKET_TYPE pgather<typename unpacket_traits<PACKET_TYPE>::type, PACKET_TYPE>( \
const unpacket_traits<PACKET_TYPE>::type* from, Index stride) { \
constexpr int size = unpacket_traits<PACKET_TYPE>::size; \
unpacket_traits<PACKET_TYPE>::type arr[size]; \
PACKET_TYPE result; \
for (int i = 0; i < size; ++i) { \
arr[i] = from[i * stride]; \
result[i] = from[i * stride]; \
} \
return *reinterpret_cast<PACKET_TYPE*>(arr); \
return result; \
}
EIGEN_CLANG_PACKET_SCATTER_GATHER(Packet16f)