-
Notifications
You must be signed in to change notification settings - Fork 4.5k
Expand file tree
/
Copy pathloader.py
More file actions
228 lines (197 loc) · 8.21 KB
/
Copy pathloader.py
File metadata and controls
228 lines (197 loc) · 8.21 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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import platform
import subprocess
import logging
import os
import sys
from packaging.version import Version
def supported_instruction_sets():
"""
Returns the set of supported CPU features, see
https://github.com/numpy/numpy/blob/master/numpy/core/src/common/npy_cpu_features.h # noqa: E501
for the list of features that this set may contain per architecture.
Example:
>>> supported_instruction_sets() # for x86
{"SSE2", "AVX2", "AVX512", ...}
>>> supported_instruction_sets() # for PPC
{"VSX", "VSX2", ...}
>>> supported_instruction_sets() # for ARM
{"NEON", "ASIMD", ...}
"""
# Old numpy.core._multiarray_umath.__cpu_features__ doesn't support Arm SVE,
# so let's read Features in numpy.distutils.cpuinfo and search 'sve' entry
def is_sve_supported():
if platform.machine() != "aarch64":
return False
# Currently SVE is only supported on Linux
if platform.system() != "Linux":
return False
# Numpy 2.0 supports SVE detection by __cpu_features__, so just skip
import numpy
if Version(numpy.__version__) >= Version("2.0"):
return False
# platform-dependent legacy fallback using numpy.distutils.cpuinfo
try:
import numpy.distutils.cpuinfo
return (
"sve"
in numpy.distutils.cpuinfo.cpu.info[0]
.get("Features", "")
.split()
)
except ImportError:
# check if SVE is supported by checking the auxval
# using values defined as:
# #define AT_HWCAP 16
# #define HWCAP_SVE (1 << 22)
return bool(
__import__("ctypes").CDLL(None).getauxval(16) & (1 << 22)
)
import numpy
if Version(numpy.__version__) >= Version("1.19"):
# use private API as next-best thing until numpy/numpy#18058 is solved
from numpy._core._multiarray_umath import __cpu_features__
# __cpu_features__ is a dictionary with CPU features
# as keys, and True / False as values
supported = {k for k, v in __cpu_features__.items() if v}
if is_sve_supported():
supported.add("SVE")
for f in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r"):
supported.discard(f)
return supported
# platform-dependent legacy fallback before numpy 1.19, no windows
if platform.system() == "Darwin":
if (
subprocess.check_output(["/usr/sbin/sysctl", "hw.optional.avx2_0"])[
-1
]
== "1"
):
return {"AVX2"}
elif platform.system() == "Linux":
import numpy.distutils.cpuinfo
result = set()
if "avx2" in numpy.distutils.cpuinfo.cpu.info[0].get("flags", ""):
result.add("AVX2")
if "avx512" in numpy.distutils.cpuinfo.cpu.info[0].get("flags", ""):
result.add("AVX512")
if "avx512_fp16" in numpy.distutils.cpuinfo.cpu.info[0].get(
"flags", ""
):
# avx512_fp16 is supported starting SPR
result.add("AVX512_SPR")
if is_sve_supported():
result.add("SVE")
for f in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r"):
result.discard(f)
return result
return set()
logger = logging.getLogger(__name__)
instruction_sets = None
# try to load optimization level from env variable
opt_env_variable_name = "FAISS_OPT_LEVEL"
opt_level = os.environ.get(opt_env_variable_name, None)
if opt_level is None:
logger.debug(
f"Environment variable {opt_env_variable_name} is not set, "
"so let's pick the instruction set according to the current CPU"
)
instruction_sets = supported_instruction_sets()
else:
logger.debug(f"Using {opt_level} as an instruction set.")
instruction_sets = set()
instruction_sets.add(opt_level)
loaded = False
has_AVX512_SPR = any("AVX512_SPR" in x.upper() for x in instruction_sets)
if has_AVX512_SPR:
try:
logger.info("Loading faiss with AVX512-SPR support.")
from .swigfaiss_avx512_spr import * # noqa: F401,F403
logger.info("Successfully loaded faiss with AVX512-SPR support.")
loaded = True
except ImportError as e:
logger.info(
f"Could not load library with AVX512-SPR support due to:\n{e!r}"
)
# reset so that we load without AVX512 below
loaded = False
has_AVX512 = any("AVX512" in x.upper() for x in instruction_sets)
if has_AVX512 and not loaded:
try:
logger.info("Loading faiss with AVX512 support.")
from .swigfaiss_avx512 import * # noqa: F401,F403
logger.info("Successfully loaded faiss with AVX512 support.")
loaded = True
except ImportError as e:
logger.info(
f"Could not load library with AVX512 support due to:\n{e!r}"
)
# reset so that we load without AVX512 below
loaded = False
has_AVX2 = "AVX2" in instruction_sets
if has_AVX2 and not loaded:
try:
logger.info("Loading faiss with AVX2 support.")
from .swigfaiss_avx2 import * # noqa: F401,F403
logger.info("Successfully loaded faiss with AVX2 support.")
loaded = True
except ImportError as e:
logger.info(f"Could not load library with AVX2 support due to:\n{e!r}")
# reset so that we load without AVX2 below
loaded = False
has_SVE = "SVE" in instruction_sets
if has_SVE and not loaded:
try:
logger.info("Loading faiss with SVE support.")
from .swigfaiss_sve import * # noqa: F401,F403
logger.info("Successfully loaded faiss with SVE support.")
loaded = True
except ImportError as e:
logger.info(f"Could not load library with SVE support due to:\n{e!r}")
# reset so that we load without SVE below
loaded = False
if not loaded:
try:
# we import * so that the symbol X can be accessed as faiss.X
logger.info("Loading faiss.")
from .swigfaiss import * # noqa: F401,F403
logger.info("Successfully loaded faiss.")
except ModuleNotFoundError:
formatted_ins_sets = ", ".join(supported_instruction_sets())
message = (
f"No module named 'faiss.swigfaiss' found. To fix this, you must "
f"do both of the following:\n"
f"A) Set the correct FAISS_OPT_LEVEL value when executing "
f"'cmake'.\n"
f"B) Build the correct SWIG wrapper.\n\n"
f"These are the supported instruction sets on your system:\n"
f"{formatted_ins_sets}\n"
f"- If 'AVX512_SPR' (case insensitive) is supported on your "
f"system, you can set the FAISS_OPT_LEVEL=avx512_spr "
f"to build the SWIG wrapper with 'AVX512-SPR' support.\n"
f"You will have to build the 'swigfaiss_avx512_spr' "
f"target in this case.\n"
f"- If 'AVX512' (case insensitive) is supported on your system, "
f"you can set the FAISS_OPT_LEVEL=avx512 to build the SWIG wrapper "
f"with 'AVX512' support.\n"
f"You will have to build the 'swigfaiss_avx512' target in this "
f"case.\n"
f"- If 'AVX2' (case sensitive) is supported on your system, you "
f"can set the FAISS_OPT_LEVEL=AVX2 to build the SWIG wrapper "
f"with 'AVX2' support.\n"
f"You will have to build the 'swigfaiss_avx2' target in this "
f"case.\n"
f"- If 'SVE' (case sensitive) is supported on your system, you can "
f"set the FAISS_OPT_LEVEL=SVE to build the SWIG wrapper with "
f"'SVE' support.\n"
f"You will have to build the 'swigfaiss_sve' target in this "
f"case.\n"
f"- If none of the above instruction sets are supported on your "
f"system, you can execute 'cmake' without setting the "
f"FAISS_OPT_LEVEL variable and build the 'swigfaiss' target."
)
logger.error(message)
sys.exit(1)