diff options
Diffstat (limited to 'src/PackAMatrix.cc')
-rw-r--r-- | src/PackAMatrix.cc | 10 |
1 files changed, 8 insertions, 2 deletions
diff --git a/src/PackAMatrix.cc b/src/PackAMatrix.cc index 143e11d..5fabf97 100644 --- a/src/PackAMatrix.cc +++ b/src/PackAMatrix.cc @@ -34,7 +34,8 @@ PackAMatrix<T, accT>::PackAMatrix( if (!cpuinfo_initialize()) { throw std::runtime_error("Failed to initialize cpuinfo!"); } - if ((!fbgemmHasAvx512Support() && !fbgemmHasAvx2Support())) { + if ((!fbgemmHasAvx512VnniSupport() && !fbgemmHasAvx512Support() && + !fbgemmHasAvx2Support())) { assert(0 && "unknown architecure"); } @@ -43,7 +44,12 @@ PackAMatrix<T, accT>::PackAMatrix( BaseType::bcol_ = params->KCB; row_interleave_B_ = params->ROW_INTERLEAVE; } else { - if (fbgemmHasAvx512Support()) { + if (fbgemmHasAvx512VnniSupport()) { + BaseType::brow_ = PackingTraits<T, accT, inst_set_t::avx512_vnni>::MCB; + BaseType::bcol_ = PackingTraits<T, accT, inst_set_t::avx512_vnni>::KCB; + row_interleave_B_ = + PackingTraits<T, accT, inst_set_t::avx512_vnni>::ROW_INTERLEAVE; + } else if (fbgemmHasAvx512Support()) { BaseType::brow_ = PackingTraits<T, accT, inst_set_t::avx512>::MCB; BaseType::bcol_ = PackingTraits<T, accT, inst_set_t::avx512>::KCB; row_interleave_B_ = |