Repository navigation
Expand file tree
/
Copy pathcpu.py
More file actions
59 lines (42 loc) · 1.71 KB
/
Copy pathcpu.py
File metadata and controls
59 lines (42 loc) · 1.71 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import enum
import deglib_cpp.cpu as cpp_cpu
class InstructionSet(enum.IntEnum):
"""
CPU instruction set used for vector distance calculations.
- ``Auto`` (0): Automatically selects the fastest available instruction set supported by the host CPU.
- ``Scalar`` (1): Standard non-vectorized fallback implementation.
- ``AVX2`` (2): Uses 256-bit vector extensions when available.
- ``AVX2_VNNI`` (3): Uses 256-bit vector extensions with VNNI acceleration when available.
- ``AVX512`` (4): Uses 512-bit vector extensions when available.
- ``AVX512_VNNI`` (5): Uses 512-bit vector extensions with VNNI acceleration when available.
"""
Auto = cpp_cpu.InstructionSet.Auto
Scalar = cpp_cpu.InstructionSet.Scalar
AVX2 = cpp_cpu.InstructionSet.AVX2
AVX2_VNNI = cpp_cpu.InstructionSet.AVX2_VNNI
AVX512 = cpp_cpu.InstructionSet.AVX512
AVX512_VNNI = cpp_cpu.InstructionSet.AVX512_VNNI
def has_avx2() -> bool:
"""
Check whether the host CPU supports AVX2 extensions.
:return: True if AVX2 is supported, False otherwise.
"""
return cpp_cpu.has_avx2()
def has_avx_vnni() -> bool:
"""
Check whether the host CPU supports AVX-VNNI (256-bit) extensions.
:return: True if AVX-VNNI is supported, False otherwise.
"""
return cpp_cpu.has_avx_vnni()
def has_avx512() -> bool:
"""
Check whether the host CPU supports AVX-512 extensions.
:return: True if AVX-512 is supported, False otherwise.
"""
return cpp_cpu.has_avx512()
def has_avx512_vnni() -> bool:
"""
Check whether the host CPU supports AVX-512 VNNI extensions.
:return: True if AVX-512 VNNI is supported, False otherwise.
"""
return cpp_cpu.has_avx512_vnni()