File size: 5,730 Bytes
1f5470c |
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 |
# 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
def Version(v):
return [int(x) for x in v.split('.')]
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", ...}
"""
# 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
import numpy.distutils.cpuinfo
return "sve" in numpy.distutils.cpuinfo.cpu.info[0].get('Features', "").split()
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 *
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 *
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
has_SVE = "SVE" in instruction_sets
if has_SVE and not loaded:
try:
logger.info("Loading faiss with SVE support.")
from .swigfaiss_sve import *
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:
# 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.")
|