Skip to content

GitLab

Unverified Commit ff9573c4 authored by Ilia Taraban's avatar Ilia Taraban Committed by GitHub
Browse files

[Misc] Disable BF16 LibXSMM SpMM for AVX2 platforms (#5677)

parent e0d2250e
Show whitespace changes
Inline Side-by-side
Showing with 10 additions and 7 deletions
src/array/cpu/spmm.h
View file @ ff9573c4
...@@ -141,8 +141,10 @@ void SpMMSumCsr( ...@@ -141,8 +141,10 @@ void SpMMSumCsr(
} }
#if !defined(_WIN32) #if !defined(_WIN32)
#ifdef USE_LIBXSMM #ifdef USE_LIBXSMM
const bool no_libxsmm = bcast.use_bcast || int cpu_id = libxsmm_cpuid_x86();
std::is_same<DType, double>::value || const bool no_libxsmm =
bcast.use_bcast || std::is_same<DType, double>::value ||
(std::is_same<DType, BFloat16>::value && cpu_id < LIBXSMM_X86_AVX512) ||
!dgl::runtime::Config::Global()->IsLibxsmmAvailable(); !dgl::runtime::Config::Global()->IsLibxsmmAvailable();
if (!no_libxsmm) { if (!no_libxsmm) {
SpMMSumCsrLibxsmm<IdType, DType, Op>(bcast, csr, ufeat, efeat, out); SpMMSumCsrLibxsmm<IdType, DType, Op>(bcast, csr, ufeat, efeat, out);
...@@ -264,9 +266,10 @@ void SpMMCmpCsr( ...@@ -264,9 +266,10 @@ void SpMMCmpCsr(
} }
#if !defined(_WIN32) #if !defined(_WIN32)
#ifdef USE_LIBXSMM #ifdef USE_LIBXSMM
int cpu_id = libxsmm_cpuid_x86();
const bool no_libxsmm = bcast.use_bcast || const bool no_libxsmm =
std::is_same<DType, double>::value || bcast.use_bcast || std::is_same<DType, double>::value ||
(std::is_same<DType, BFloat16>::value && cpu_id < LIBXSMM_X86_AVX512) ||
!dgl::runtime::Config::Global()->IsLibxsmmAvailable(); !dgl::runtime::Config::Global()->IsLibxsmmAvailable();
if (!no_libxsmm) { if (!no_libxsmm) {
SpMMCmpCsrLibxsmm<IdType, DType, Op, Cmp>( SpMMCmpCsrLibxsmm<IdType, DType, Op, Cmp>(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment