cpufeature fails to be installed on some non-x86 CPUs
#279 opened on 2023年3月1日
Repository metrics
- Stars
- (8,248 stars)
- PR merge metrics
- (30d に merged PR はありません)
説明
We use the cpufeature module to detect if a CPU supports AVX512, however, cpufeature fails to be installed on some non-x86 CPUs (e.g., ARMs).
We need to remove this dependency and use code like this to detect if the CPU supports AVX512:
def is_fast_bfloat16_supported() -> bool:
if os.name != "posix" or platform.processor() != "x86_64":
return False
try:
eax = _run_x86_assembler(
b"\xB8\x07\x00\x00\x00" # mov eax, 7
b"\xB9\x01\x00\x00\x00" # mov ecx, 1
b"\x0f\xa2" # cpuid
b"\xC3" # ret
)
avx512_bf16 = bool(eax & (1 << 5))
return amx_bf16 or avx512_bf16
except:
logger.warning("Failed to check for the CPU's native bfloat16 support:", exc_info=True)
return False
def _run_x86_assembler(machine_code: bytes) -> int:
memory = mmap.mmap(
-1,
len(machine_code),
flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS,
prot=mmap.PROT_WRITE | mmap.PROT_READ | mmap.PROT_EXEC,
)
try:
memory.write(machine_code)
address = ctypes.addressof(ctypes.c_int.from_buffer(memory))
functype = ctypes.CFUNCTYPE(ctypes.c_uint32)
func = functype(address)
return func()
finally:
memory.close()
Note that the code above checks for the presence of AVX512_BF16 instruction set, but actually pure AVX512 is enough. This needs to be changed before integrating this code, see https://en.wikichip.org/wiki/x86/avx-512 for the relevant instruction params.
Before merging, it's necessary to test on (a) x86-64 CPUs without AVX512, (b) x86-64 CPUs with AVX512, (c) ARM CPUs (e.g. available on Oracle free-tier cloud or other places).