Skip to content

Commit 21e20f3

Browse files
yurymalkovyurymalkov
authored
Merge pull request #367 from nmslib/develop
Update master to 0.6.1
2 parents 14cabd0 + 2ebbc2c commit 21e20f3

File tree

9 files changed

+228
-34
lines changed

9 files changed

+228
-34
lines changed

‎README.md

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,14 @@ Header-only C++ HNSW implementation with python bindings.
33

44
**NEWS:**
55

6-
**version 0.6**
6+
7+
**version 0.6.1**
8+
9+
* Thanks to ([@tony-kuo](https://github.com/tony-kuo)) hnswlib AVX512 and AVX builds are not backwards-compatible with older SSE and non-AVX512 architectures.
10+
* Thanks to ([@psobot](https://github.com/psobot)) there is now a sencible message instead of segfault when passing a scalar to get_items.
11+
* Thanks to ([@urigoren](https://github.com/urigoren)) hnswlib has a lazy index creation python wrapper.
12+
13+
**version 0.6.0**
714
* Thanks to ([@dyashuni](https://github.com/dyashuni)) hnswlib now uses github actions for CI, there is a search speedup in some scenarios with deletions. `unmark_deleted(label)` is now also a part of the python interface (note now it throws an exception for double deletions).
815
* Thanks to ([@slice4e](https://github.com/slice4e)) we now support AVX512; thanks to ([@LTLA](https://github.com/LTLA)) the cmake interface for the lib is now updated.
916
* Thanks to ([@alonre24](https://github.com/alonre24)) we now have a python bindings for brute-force (and examples for recall tuning: [TESTING_RECALL.md](TESTING_RECALL.md).

‎examples/git_tester.py

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,34 @@
11
from pydriller import Repository
22
import os
33
import datetime
4-
os.system("cp examples/speedtest.py examples/speedtest2.py")
5-
for commit in Repository('.', from_tag="v0.5.2").traverse_commits():
6-
print(commit.hash)
7-
print(commit.msg)
4+
os.system("cp examples/speedtest.py examples/speedtest2.py") # the file has to be outside of git
5+
for idx, commit in enumerate(Repository('.', from_tag="v0.6.0").traverse_commits()):
6+
name=commit.msg.replace('\n', ' ').replace('\r', ' ')
7+
print(idx, commit.hash, name)
8+
9+
10+
11+
for commit in Repository('.', from_tag="v0.6.0").traverse_commits():
12+
13+
name=commit.msg.replace('\n', ' ').replace('\r', ' ')
14+
print(commit.hash, name)
815

916
os.system(f"git checkout {commit.hash}; rm -rf build; ")
10-
os.system("python -m pip install .")
11-
os.system(f'python examples/speedtest2.py -n "{commit.msg}" -d 4 -t 1')
12-
os.system(f'python examples/speedtest2.py -n "{commit.msg}" -d 64 -t 1')
13-
os.system(f'python examples/speedtest2.py -n "{commit.msg}" -d 128 -t 1')
14-
os.system(f'python examples/speedtest2.py -n "{commit.msg}" -d 4 -t 24')
15-
os.system(f'python examples/speedtest2.py -n "{commit.msg}" -d 128 -t 24')
17+
print("\n\n--------------------\n\n")
18+
ret=os.system("python -m pip install .")
19+
print(ret)
20+
21+
if ret != 0:
22+
print ("build failed!!!!")
23+
print ("build failed!!!!")
24+
print ("build failed!!!!")
25+
print ("build failed!!!!")
26+
continue
27+
28+
os.system(f'python examples/speedtest2.py -n "{name}" -d 4 -t 1')
29+
os.system(f'python examples/speedtest2.py -n "{name}" -d 64 -t 1')
30+
os.system(f'python examples/speedtest2.py -n "{name}" -d 128 -t 1')
31+
os.system(f'python examples/speedtest2.py -n "{name}" -d 4 -t 24')
32+
os.system(f'python examples/speedtest2.py -n "{name}" -d 128 -t 24')
33+
1634

‎hnswlib/hnswlib.h

Lines changed: 76 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,25 @@
1515
#ifdef _MSC_VER
1616
#include <intrin.h>
1717
#include <stdexcept>
18+
#include "cpu_x86.h"
19+
void cpu_x86::cpuid(int32_t out[4], int32_t eax, int32_t ecx) {
20+
__cpuidex(out, eax, ecx);
21+
}
22+
__int64 xgetbv(unsigned int x) {
23+
return _xgetbv(x);
24+
}
1825
#else
1926
#include <x86intrin.h>
27+
#include <cpuid.h>
28+
#include <stdint.h>
29+
void cpuid(int32_t cpuInfo[4], int32_t eax, int32_t ecx) {
30+
__cpuid_count(eax, ecx, cpuInfo[0], cpuInfo[1], cpuInfo[2], cpuInfo[3]);
31+
}
32+
uint64_t xgetbv(unsigned int index) {
33+
uint32_t eax, edx;
34+
__asm__ __volatile__("xgetbv" : "=a"(eax), "=d"(edx) : "c"(index));
35+
return ((uint64_t)edx << 32) | eax;
36+
}
2037
#endif
2138

2239
#if defined(USE_AVX512)
@@ -30,6 +47,65 @@
3047
#define PORTABLE_ALIGN32 __declspec(align(32))
3148
#define PORTABLE_ALIGN64 __declspec(align(64))
3249
#endif
50+
51+
// Adapted from https://github.com/Mysticial/FeatureDetector
52+
#define _XCR_XFEATURE_ENABLED_MASK 0
53+
54+
bool AVXCapable() {
55+
int cpuInfo[4];
56+
57+
// CPU support
58+
cpuid(cpuInfo, 0, 0);
59+
int nIds = cpuInfo[0];
60+
61+
bool HW_AVX = false;
62+
if (nIds >= 0x00000001) {
63+
cpuid(cpuInfo, 0x00000001, 0);
64+
HW_AVX = (cpuInfo[2] & ((int)1 << 28)) != 0;
65+
}
66+
67+
// OS support
68+
cpuid(cpuInfo, 1, 0);
69+
70+
bool osUsesXSAVE_XRSTORE = (cpuInfo[2] & (1 << 27)) != 0;
71+
bool cpuAVXSuport = (cpuInfo[2] & (1 << 28)) != 0;
72+
73+
bool avxSupported = false;
74+
if (osUsesXSAVE_XRSTORE && cpuAVXSuport) {
75+
uint64_t xcrFeatureMask = xgetbv(_XCR_XFEATURE_ENABLED_MASK);
76+
avxSupported = (xcrFeatureMask & 0x6) == 0x6;
77+
}
78+
return HW_AVX && avxSupported;
79+
}
80+
81+
bool AVX512Capable() {
82+
if (!AVXCapable()) return false;
83+
84+
int cpuInfo[4];
85+
86+
// CPU support
87+
cpuid(cpuInfo, 0, 0);
88+
int nIds = cpuInfo[0];
89+
90+
bool HW_AVX512F = false;
91+
if (nIds >= 0x00000007) { // AVX512 Foundation
92+
cpuid(cpuInfo, 0x00000007, 0);
93+
HW_AVX512F = (cpuInfo[1] & ((int)1 << 16)) != 0;
94+
}
95+
96+
// OS support
97+
cpuid(cpuInfo, 1, 0);
98+
99+
bool osUsesXSAVE_XRSTORE = (cpuInfo[2] & (1 << 27)) != 0;
100+
bool cpuAVXSuport = (cpuInfo[2] & (1 << 28)) != 0;
101+
102+
bool avx512Supported = false;
103+
if (osUsesXSAVE_XRSTORE && cpuAVXSuport) {
104+
uint64_t xcrFeatureMask = xgetbv(_XCR_XFEATURE_ENABLED_MASK);
105+
avx512Supported = (xcrFeatureMask & 0xe6) == 0xe6;
106+
}
107+
return HW_AVX512F && avx512Supported;
108+
}
33109
#endif
34110

35111
#include <queue>
@@ -108,7 +184,6 @@ namespace hnswlib {
108184

109185
return result;
110186
}
111-
112187
}
113188

114189
#include "space_l2.h"

‎hnswlib/space_ip.h

Lines changed: 31 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ namespace hnswlib {
1818

1919
// Favor using AVX if available.
2020
static float
21-
InnerProductSIMD4Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
21+
InnerProductSIMD4ExtAVX(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
2222
float PORTABLE_ALIGN32 TmpRes[8];
2323
float *pVect1 = (float *) pVect1v;
2424
float *pVect2 = (float *) pVect2v;
@@ -64,10 +64,12 @@ namespace hnswlib {
6464
return 1.0f - sum;
6565
}
6666

67-
#elif defined(USE_SSE)
67+
#endif
68+
69+
#if defined(USE_SSE)
6870

6971
static float
70-
InnerProductSIMD4Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
72+
InnerProductSIMD4ExtSSE(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
7173
float PORTABLE_ALIGN32 TmpRes[8];
7274
float *pVect1 = (float *) pVect1v;
7375
float *pVect2 = (float *) pVect2v;
@@ -128,7 +130,7 @@ namespace hnswlib {
128130
#if defined(USE_AVX512)
129131

130132
static float
131-
InnerProductSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
133+
InnerProductSIMD16ExtAVX512(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
132134
float PORTABLE_ALIGN64 TmpRes[16];
133135
float *pVect1 = (float *) pVect1v;
134136
float *pVect2 = (float *) pVect2v;
@@ -157,10 +159,12 @@ namespace hnswlib {
157159
return 1.0f - sum;
158160
}
159161

160-
#elif defined(USE_AVX)
162+
#endif
163+
164+
#if defined(USE_AVX)
161165

162166
static float
163-
InnerProductSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
167+
InnerProductSIMD16ExtAVX(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
164168
float PORTABLE_ALIGN32 TmpRes[8];
165169
float *pVect1 = (float *) pVect1v;
166170
float *pVect2 = (float *) pVect2v;
@@ -195,10 +199,12 @@ namespace hnswlib {
195199
return 1.0f - sum;
196200
}
197201

198-
#elif defined(USE_SSE)
202+
#endif
203+
204+
#if defined(USE_SSE)
199205

200206
static float
201-
InnerProductSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
207+
InnerProductSIMD16ExtSSE(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
202208
float PORTABLE_ALIGN32 TmpRes[8];
203209
float *pVect1 = (float *) pVect1v;
204210
float *pVect2 = (float *) pVect2v;
@@ -245,6 +251,9 @@ namespace hnswlib {
245251
#endif
246252

247253
#if defined(USE_SSE) || defined(USE_AVX) || defined(USE_AVX512)
254+
DISTFUNC<float> InnerProductSIMD16Ext = InnerProductSIMD16ExtSSE;
255+
DISTFUNC<float> InnerProductSIMD4Ext = InnerProductSIMD4ExtSSE;
256+
248257
static float
249258
InnerProductSIMD16ExtResiduals(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
250259
size_t qty = *((size_t *) qty_ptr);
@@ -283,6 +292,20 @@ namespace hnswlib {
283292
InnerProductSpace(size_t dim) {
284293
fstdistfunc_ = InnerProduct;
285294
#if defined(USE_AVX) || defined(USE_SSE) || defined(USE_AVX512)
295+
#if defined(USE_AVX512)
296+
if (AVX512Capable())
297+
InnerProductSIMD16Ext = InnerProductSIMD16ExtAVX512;
298+
else if (AVXCapable())
299+
InnerProductSIMD16Ext = InnerProductSIMD16ExtAVX;
300+
#elif defined(USE_AVX)
301+
if (AVXCapable())
302+
InnerProductSIMD16Ext = InnerProductSIMD16ExtAVX;
303+
#endif
304+
#if defined(USE_AVX)
305+
if (AVXCapable())
306+
InnerProductSIMD4Ext = InnerProductSIMD4ExtAVX;
307+
#endif
308+
286309
if (dim % 16 == 0)
287310
fstdistfunc_ = InnerProductSIMD16Ext;
288311
else if (dim % 4 == 0)

‎hnswlib/space_l2.h

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ namespace hnswlib {
2323

2424
// Favor using AVX512 if available.
2525
static float
26-
L2SqrSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
26+
L2SqrSIMD16ExtAVX512(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
2727
float *pVect1 = (float *) pVect1v;
2828
float *pVect2 = (float *) pVect2v;
2929
size_t qty = *((size_t *) qty_ptr);
@@ -52,12 +52,13 @@ namespace hnswlib {
5252

5353
return (res);
5454
}
55+
#endif
5556

56-
#elif defined(USE_AVX)
57+
#if defined(USE_AVX)
5758

5859
// Favor using AVX if available.
5960
static float
60-
L2SqrSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
61+
L2SqrSIMD16ExtAVX(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
6162
float *pVect1 = (float *) pVect1v;
6263
float *pVect2 = (float *) pVect2v;
6364
size_t qty = *((size_t *) qty_ptr);
@@ -89,10 +90,12 @@ namespace hnswlib {
8990
return TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + TmpRes[5] + TmpRes[6] + TmpRes[7];
9091
}
9192

92-
#elif defined(USE_SSE)
93+
#endif
94+
95+
#if defined(USE_SSE)
9396

9497
static float
95-
L2SqrSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
98+
L2SqrSIMD16ExtSSE(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
9699
float *pVect1 = (float *) pVect1v;
97100
float *pVect2 = (float *) pVect2v;
98101
size_t qty = *((size_t *) qty_ptr);
@@ -141,6 +144,8 @@ namespace hnswlib {
141144
#endif
142145

143146
#if defined(USE_SSE) || defined(USE_AVX) || defined(USE_AVX512)
147+
DISTFUNC<float> L2SqrSIMD16Ext = L2SqrSIMD16ExtSSE;
148+
144149
static float
145150
L2SqrSIMD16ExtResiduals(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
146151
size_t qty = *((size_t *) qty_ptr);
@@ -156,7 +161,7 @@ namespace hnswlib {
156161
#endif
157162

158163

159-
#ifdef USE_SSE
164+
#if defined(USE_SSE)
160165
static float
161166
L2SqrSIMD4Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
162167
float PORTABLE_ALIGN32 TmpRes[8];
@@ -208,7 +213,17 @@ namespace hnswlib {
208213
public:
209214
L2Space(size_t dim) {
210215
fstdistfunc_ = L2Sqr;
211-
#if defined(USE_SSE) || defined(USE_AVX) || defined(USE_AVX512)
216+
#if defined(USE_SSE) || defined(USE_AVX) || defined(USE_AVX512)
217+
#if defined(USE_AVX512)
218+
if (AVX512Capable())
219+
L2SqrSIMD16Ext = L2SqrSIMD16ExtAVX512;
220+
else if (AVXCapable())
221+
L2SqrSIMD16Ext = L2SqrSIMD16ExtAVX;
222+
#elif defined(USE_AVX)
223+
if (AVXCapable())
224+
L2SqrSIMD16Ext = L2SqrSIMD16ExtAVX;
225+
#endif
226+
212227
if (dim % 16 == 0)
213228
fstdistfunc_ = L2SqrSIMD16Ext;
214229
else if (dim % 4 == 0)
@@ -217,7 +232,7 @@ namespace hnswlib {
217232
fstdistfunc_ = L2SqrSIMD16ExtResiduals;
218233
else if (dim > 4)
219234
fstdistfunc_ = L2SqrSIMD4ExtResiduals;
220-
#endif
235+
#endif
221236
dim_ = dim;
222237
data_size_ = dim * sizeof(float);
223238
}

‎python_bindings/LazyIndex.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import hnswlib
2+
"""
3+
A python wrapper for lazy indexing, preserves the same api as hnswlib.Index but initializes the index only after adding items for the first time with `add_items`.
4+
"""
5+
class LazyIndex(hnswlib.Index):
6+
def __init__(self, space, dim,max_elements=1024, ef_construction=200, M=16):
7+
super().__init__(space, dim)
8+
self.init_max_elements=max_elements
9+
self.init_ef_construction=ef_construction
10+
self.init_M=M
11+
def init_index(self, max_elements=0,M=0,ef_construction=0):
12+
if max_elements>0:
13+
self.init_max_elements=max_elements
14+
if ef_construction>0:
15+
self.init_ef_construction=ef_construction
16+
if M>0:
17+
self.init_M=M
18+
super().init_index(self.init_max_elements, self.init_M, self.init_ef_construction)
19+
def add_items(self, data, ids=None, num_threads=-1):
20+
if self.max_elements==0:
21+
self.init_index()
22+
return super().add_items(data,ids, num_threads)
23+
def get_items(self, ids=None):
24+
if self.max_elements==0:
25+
return []
26+
return super().get_items(ids)
27+
def knn_query(self, data,k=1, num_threads=-1):
28+
if self.max_elements==0:
29+
return [], []
30+
return super().knn_query(data, k, num_threads)
31+
def resize_index(self, size):
32+
if self.max_elements==0:
33+
return self.init_index(size)
34+
else:
35+
return super().resize_index(size)
36+
def set_ef(self, ef):
37+
if self.max_elements==0:
38+
self.init_ef_construction=ef
39+
return
40+
super().set_ef(ef)
41+
def get_max_elements(self):
42+
return self.max_elements
43+
def get_current_count(self):
44+
return self.element_count

0 commit comments

Comments
 (0)