Skip to content

Commit

Permalink
Use a common implementation for has_<simd> methods
Browse files Browse the repository at this point in the history
  • Loading branch information
t20100 committed Feb 8, 2024
1 parent b701ba7 commit 9dd7922
Showing 1 changed file with 14 additions and 28 deletions.
42 changes: 14 additions & 28 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,45 +186,31 @@ def has_cpp14(self) -> bool:
return True
return check_compile_flags(self.__compiler, '-std=c++14', extension='.cc')

def _has_x86_simd(self, *flags) -> bool:
"""Check x86 SIMD availability on host"""
if self.ARCH not in ('X86_32', 'X86_64'):
return False
if not all(has_cpu_flag(flag) for flag in flags):
return False
if self.__compiler.compiler_type == "msvc":
return True
return check_compile_flags(self.__compiler, *(f"-m{flag}" for flag in flags))

def has_sse2(self) -> bool:
"""Check SSE2 availability on host"""
if self.ARCH in ('X86_32', 'X86_64'):
if not has_cpu_flag('sse2'):
return False # SSE2 not available on host
if self.__compiler.compiler_type == "msvc":
return True
return check_compile_flags(self.__compiler, "-msse2")
return False # Disabled by default
return self._has_x86_simd('sse2')

def has_ssse3(self) -> bool:
"""Check SSSE3 availability on host"""
if self.ARCH in ('X86_32', 'X86_64'):
if not has_cpu_flag('ssse3'):
return False # SSSE3 not available on host
if self.__compiler.compiler_type == "msvc":
return True
return check_compile_flags(self.__compiler, "-mssse3")
return False # Disabled by default
return self._has_x86_simd('ssse3')

def has_avx2(self) -> bool:
"""Check AVX2 availability on host"""
if self.ARCH in ('X86_32', 'X86_64'):
if not has_cpu_flag('avx2'):
return False # AVX2 not available on host
if self.__compiler.compiler_type == 'msvc':
return True
return check_compile_flags(self.__compiler, '-mavx2')
return False # Disabled by default
return self._has_x86_simd('avx2')

def has_avx512(self) -> bool:
"""Check AVX512 "F" and "BW" instruction sets availability on host"""
if self.ARCH in ('X86_32', 'X86_64'):
if not (has_cpu_flag('avx512f') and has_cpu_flag('avx512bw')):
return False # AVX512 F and/or BW not available on host
if self.__compiler.compiler_type == 'msvc':
return True
return check_compile_flags(self.__compiler, '-mavx512f', '-mavx512bw')
return False # Disabled by default
return self._has_x86_simd('avx512f', 'avx512bw')

def has_openmp(self) -> bool:
"""Check OpenMP availability on host"""
Expand Down

0 comments on commit 9dd7922

Please sign in to comment.