use py-cpuinfo to detect SIMD_WIDTH in platform-independent way (#1616)

Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
This commit is contained in:
Victor 2021-12-10 18:31:55 -08:00 коммит произвёл GitHub
Родитель feb6afb049
Коммит 64c2946a23
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: 4AEE18F83AFDEB23
2 изменённых файлов: 18 добавлений и 22 удалений

Просмотреть файл

@ -266,44 +266,39 @@ class OpBuilder(ABC):
return [x for x in args if len(x) > 0]
def cpu_arch(self):
if not self.command_exists('lscpu'):
self.warning(
f"{self.name} attempted to query 'lscpu' to detect the CPU architecture. "
"However, 'lscpu' does not appear to exist on "
"your system, will fall back to use -march=native.")
return '-march=native'
try:
from cpuinfo import get_cpu_info
except ImportError as e:
self.warning(f"{self.name} to failed import cpu_info:\n{e}\n"
"will fall back to use -march=native")
return "-march=native"
cpu_info = get_cpu_info()
result = subprocess.check_output('lscpu', shell=True)
result = result.decode('utf-8').strip().lower()
if 'ppc64le' in result:
if cpu_info['arch'].startswith('PPC_'):
# gcc does not provide -march on PowerPC, use -mcpu instead
return '-mcpu=native'
return '-march=native'
def simd_width(self):
if not self.command_exists('lscpu'):
self.warning(
f"{self.name} attempted to query 'lscpu' to detect the existence "
"of AVX instructions. However, 'lscpu' does not appear to exist on "
"your system, will fall back to non-vectorized execution.")
try:
from cpuinfo import get_cpu_info
except ImportError as e:
self.warning(f"{self.name} to failed import cpu_info:\n{e}\n"
"will fall back to non-vectorized execution.")
return '-D__SCALAR__'
try:
result = subprocess.check_output('lscpu', shell=True)
result = result.decode('utf-8').strip().lower()
cpu_info = get_cpu_info()
except Exception as e:
print(
f"{WARNING} {self.name} SIMD_WIDTH cannot be recognized due to {str(e)}!"
)
return '-D__SCALAR__'
if 'genuineintel' in result:
if 'avx512' in result:
if cpu_info['arch'] == 'X86_64':
if 'avx512' in cpu_info['flags']:
return '-D__AVX512__'
elif 'avx2' in result:
return '-D__AVX256__'
elif 'authenticamd' in result:
if 'avx2' in result:
elif 'avx2' in cpu_info['flags']:
return '-D__AVX256__'
return '-D__SCALAR__'

Просмотреть файл

@ -3,5 +3,6 @@ ninja
numpy
packaging
psutil
py-cpuinfo
torch
tqdm