Spaces:
Running
Running
File size: 3,697 Bytes
2a0bc63 |
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 |
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from packaging.version import Version
import platform
import subprocess
import logging
import os
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
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", ...}
"""
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}
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")
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 = any("AVX512" in x.upper() for x in instruction_sets)
if has_AVX512:
try:
logger.info("Loading faiss with AVX512 support.")
from .swigfaiss_avx512 import *
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 *
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
if not loaded:
# we import * so that the symbol X can be accessed as faiss.X
logger.info("Loading faiss.")
from .swigfaiss import *
logger.info("Successfully loaded faiss.")
|