usearch 0.4.0 → 0.8.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.clang-tidy +60 -0
- package/.github/workflows/release.yml +42 -16
- package/.github/workflows/update_version.sh +1 -0
- package/.vscode/launch.json +22 -0
- package/.vscode/settings.json +5 -0
- package/.vscode/tasks.json +4 -4
- package/CITATION.cff +11 -0
- package/CMakeLists.txt +108 -95
- package/Cargo.toml +1 -1
- package/Dockerfile +14 -0
- package/README.md +8 -7
- package/VERSION +1 -1
- package/binding.gyp +2 -2
- package/build.rs +2 -2
- package/cmake/config.cmake.in +15 -0
- package/cmake/pkg-config.pc.in +4 -0
- package/cmake/usearchConfigVersion.cmake.in +20 -0
- package/conanfile.py +1 -1
- package/docs/benchmarks.md +11 -2
- package/docs/{Compilation.md → compilation.md} +54 -27
- package/include/usearch/usearch.hpp +664 -390
- package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +17 -16
- package/javascript/lib.cpp +13 -13
- package/javascript/test.js +1 -1
- package/objc/USearchObjective.mm +46 -29
- package/objc/include/USearchObjective.h +18 -26
- package/package.json +1 -1
- package/pyproject.toml +3 -2
- package/python/lib.cpp +179 -120
- package/python/{bench.py → scripts/bench.py} +6 -1
- package/python/scripts/bench_params.py +112 -0
- package/python/{test.py → scripts/test.py} +32 -13
- package/python/usearch/client.py +100 -14
- package/python/usearch/server.py +104 -23
- package/rust/lib.cpp +15 -14
- package/rust/lib.hpp +5 -5
- package/rust/lib.rs +3 -3
- package/setup.py +7 -0
- package/src/CMakeLists.txt +107 -0
- package/src/bench.cpp +45 -15
- package/src/{advanced.hpp → punned.hpp} +227 -199
- package/src/test.cpp +62 -0
- package/swift/Test.swift +1 -1
- package/wolfram/lib.cpp +14 -14
|
@@ -1,37 +1,57 @@
|
|
|
1
1
|
import pytest
|
|
2
2
|
import numpy as np
|
|
3
3
|
|
|
4
|
-
|
|
4
|
+
from usearch.io import load_matrix
|
|
5
5
|
from usearch.index import Index, SetsIndex
|
|
6
6
|
|
|
7
7
|
dimensions = [13, 99, 100, 256]
|
|
8
8
|
sizes = [1, 2, 100]
|
|
9
|
-
|
|
9
|
+
index_types = ['f64', 'f32', 'f16', 'f8']
|
|
10
|
+
numpy_types = [np.float64, np.float32, np.float16, np.byte]
|
|
11
|
+
connectivities = [3, 13, 50]
|
|
10
12
|
|
|
11
13
|
|
|
12
14
|
@pytest.mark.parametrize('ndim', dimensions)
|
|
13
|
-
@pytest.mark.parametrize('
|
|
14
|
-
|
|
15
|
-
|
|
16
|
-
|
|
15
|
+
@pytest.mark.parametrize('index_type', index_types)
|
|
16
|
+
@pytest.mark.parametrize('numpy_type', numpy_types)
|
|
17
|
+
@pytest.mark.parametrize('connectivity', connectivities)
|
|
18
|
+
def test_l2sq(ndim: int, index_type: str, numpy_type: str, connectivity: int):
|
|
19
|
+
index = Index(
|
|
20
|
+
metric='l2sq',
|
|
21
|
+
ndim=ndim,
|
|
22
|
+
dtype=index_type,
|
|
23
|
+
connectivity=connectivity,
|
|
24
|
+
)
|
|
25
|
+
assert index.ndim == ndim
|
|
26
|
+
assert index.connectivity == connectivity
|
|
27
|
+
|
|
28
|
+
vector = np.random.uniform(0, 0.7, (index.ndim)).astype(numpy_type)
|
|
17
29
|
index.add(42, vector)
|
|
18
30
|
matches, distances, count = index.search(vector, 10)
|
|
19
31
|
|
|
32
|
+
assert len(index) == 1
|
|
20
33
|
assert len(matches) == count
|
|
21
34
|
assert len(distances) == count
|
|
22
35
|
assert count == 1
|
|
23
36
|
assert matches[0] == 42
|
|
24
37
|
assert distances[0] == pytest.approx(0, abs=1e-3)
|
|
25
38
|
|
|
39
|
+
index.save('tmp.usearch')
|
|
40
|
+
index.clear()
|
|
41
|
+
assert len(index) == 0
|
|
42
|
+
|
|
43
|
+
index.load('tmp.usearch')
|
|
44
|
+
assert len(index) == 1
|
|
45
|
+
|
|
26
46
|
|
|
27
47
|
@pytest.mark.parametrize('ndim', dimensions)
|
|
28
48
|
@pytest.mark.parametrize('size', sizes)
|
|
29
|
-
@pytest.mark.parametrize('
|
|
30
|
-
|
|
31
|
-
|
|
49
|
+
@pytest.mark.parametrize('index_type', index_types)
|
|
50
|
+
@pytest.mark.parametrize('numpy_type', numpy_types)
|
|
51
|
+
def test_l2sq_batch(ndim: int, size: int, index_type: str, numpy_type: str):
|
|
52
|
+
index = Index(ndim=ndim, metric='l2sq', dtype=index_type)
|
|
32
53
|
|
|
33
|
-
vectors = np.random.uniform(
|
|
34
|
-
0, 0.3, (size, index.ndim)).astype(np.float32)
|
|
54
|
+
vectors = np.random.uniform(0, 0.7, (size, index.ndim)).astype(numpy_type)
|
|
35
55
|
labels = np.array(range(size), dtype=np.longlong)
|
|
36
56
|
|
|
37
57
|
index.add(labels, vectors)
|
|
@@ -84,8 +104,7 @@ def test_sets():
|
|
|
84
104
|
index = SetsIndex()
|
|
85
105
|
index.add(10, np.array([10, 12, 15], dtype=np.uint32))
|
|
86
106
|
index.add(11, np.array([11, 12, 15, 16], dtype=np.uint32))
|
|
87
|
-
results = index.search(np.array([12, 15], dtype=np.uint32),
|
|
88
|
-
|
|
107
|
+
# results = index.search(np.array([12, 15], dtype=np.uint32), 10)
|
|
89
108
|
# assert list(results) == [10, 11]
|
|
90
109
|
|
|
91
110
|
|
package/python/usearch/client.py
CHANGED
|
@@ -1,34 +1,120 @@
|
|
|
1
|
+
from typing import Union, Optional
|
|
2
|
+
|
|
1
3
|
import numpy as np
|
|
2
4
|
from ucall.client import Client
|
|
3
5
|
|
|
4
6
|
|
|
7
|
+
def _vector_to_ascii(vector: np.ndarray) -> Optional[str]:
|
|
8
|
+
if vector.dtype != np.int8 and vector.dtype != np.uint8 and vector.dtype != np.byte:
|
|
9
|
+
return None
|
|
10
|
+
if not np.all((vector >= 0) | (vector <= 100)):
|
|
11
|
+
return None
|
|
12
|
+
|
|
13
|
+
# Let's map [0, 100] to the range from [23, 123],
|
|
14
|
+
# poking 60 and replacing with the 124.
|
|
15
|
+
vector += 23
|
|
16
|
+
vector[vector == 60] = 124
|
|
17
|
+
ascii = str(vector)
|
|
18
|
+
return ascii
|
|
19
|
+
|
|
20
|
+
|
|
5
21
|
class IndexClient:
|
|
6
22
|
|
|
7
23
|
def __init__(self, uri: str = '127.0.0.1', port: int = 8545, use_http: bool = True) -> None:
|
|
8
24
|
self.client = Client(uri=uri, port=port, use_http=use_http)
|
|
9
25
|
|
|
10
|
-
def
|
|
11
|
-
|
|
12
|
-
|
|
26
|
+
def add_one(self, label: int, vector: np.ndarray):
|
|
27
|
+
assert isinstance(label, int)
|
|
28
|
+
assert isinstance(vector, np.ndarray)
|
|
29
|
+
vector = vector.flatten()
|
|
30
|
+
ascii = _vector_to_ascii(vector)
|
|
31
|
+
if ascii:
|
|
32
|
+
self.client.add_ascii(label=label, string=ascii)
|
|
13
33
|
else:
|
|
14
|
-
|
|
34
|
+
self.client.add_one(label=label, vectors=vector)
|
|
35
|
+
|
|
36
|
+
def add_many(self, labels: np.ndarray, vectors: np.ndarray):
|
|
37
|
+
assert isinstance(labels, int)
|
|
38
|
+
assert isinstance(vectors, np.ndarray)
|
|
39
|
+
assert labels.ndim == 1 and vectors.ndim == 2
|
|
40
|
+
assert labels.shape[0] == vectors.shape[0]
|
|
41
|
+
self.client.add_many(labels=labels, vectors=vectors)
|
|
42
|
+
|
|
43
|
+
def add(self, labels: Union[np.ndarray, int], vectors: np.ndarray):
|
|
44
|
+
if isinstance(labels, int) or len(labels) == 1:
|
|
45
|
+
return self.add_one(labels, vectors)
|
|
46
|
+
else:
|
|
47
|
+
return self.add_many(labels, vectors)
|
|
48
|
+
|
|
49
|
+
def search_one(self, vector: np.ndarray, count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
50
|
+
matches: list[dict] = []
|
|
51
|
+
vector = vector.flatten()
|
|
52
|
+
ascii = _vector_to_ascii(vector)
|
|
53
|
+
if ascii:
|
|
54
|
+
matches = self.client.search_ascii(string=ascii, count=count)
|
|
55
|
+
else:
|
|
56
|
+
matches = self.client.search_one(vector=vector, count=count)
|
|
57
|
+
|
|
58
|
+
print(matches.data)
|
|
59
|
+
matches = matches.json
|
|
60
|
+
|
|
61
|
+
labels = np.array((1, count), dtype=np.uint32)
|
|
62
|
+
distances = np.array((1, count), dtype=np.float32)
|
|
63
|
+
counts = np.array((1), dtype=np.uint32)
|
|
64
|
+
for col, result in enumerate(matches):
|
|
65
|
+
labels[0, col] = result['label']
|
|
66
|
+
distances[0, col] = result['distance']
|
|
67
|
+
counts[0] = len(matches)
|
|
15
68
|
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
69
|
+
return labels, distances, counts
|
|
70
|
+
|
|
71
|
+
def search_many(self, vectors: np.ndarray, count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
72
|
+
batch_size: int = vectors.shape[0]
|
|
73
|
+
list_of_matches: list[list[dict]] = self.client.search_many(
|
|
74
|
+
vectors=vectors, count=count)
|
|
75
|
+
|
|
76
|
+
labels = np.array((batch_size, count), dtype=np.uint32)
|
|
77
|
+
distances = np.array((batch_size, count), dtype=np.float32)
|
|
78
|
+
counts = np.array((batch_size), dtype=np.uint32)
|
|
79
|
+
for row, matches in enumerate(list_of_matches):
|
|
80
|
+
for col, result in enumerate(matches):
|
|
81
|
+
labels[row, col] = result['label']
|
|
82
|
+
distances[row, col] = result['distance']
|
|
83
|
+
counts[row] = len(results)
|
|
84
|
+
|
|
85
|
+
return labels, distances, counts
|
|
86
|
+
|
|
87
|
+
def search(self, vectors: np.ndarray, count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
88
|
+
if vectors.ndim == 1 or (vectors.ndim == 2 and vectors.shape[0] == 1):
|
|
89
|
+
return self.search_one(vectors, count)
|
|
90
|
+
else:
|
|
91
|
+
return self.search_many(vectors, count)
|
|
22
92
|
|
|
23
93
|
def __len__(self):
|
|
24
|
-
return self.client.size()
|
|
94
|
+
return self.client.size().json()
|
|
25
95
|
|
|
26
96
|
@property
|
|
27
97
|
def ndim(self):
|
|
28
|
-
return self.client.ndim()
|
|
98
|
+
return self.client.ndim().json()
|
|
29
99
|
|
|
30
100
|
def capacity(self):
|
|
31
|
-
return self.client.capacity()
|
|
101
|
+
return self.client.capacity().json()
|
|
32
102
|
|
|
33
103
|
def connectivity(self):
|
|
34
|
-
return self.client.connectivity()
|
|
104
|
+
return self.client.connectivity().json()
|
|
105
|
+
|
|
106
|
+
def load(self, path: str):
|
|
107
|
+
raise NotImplementedError()
|
|
108
|
+
|
|
109
|
+
def view(self, path: str):
|
|
110
|
+
raise NotImplementedError()
|
|
111
|
+
|
|
112
|
+
def save(self, path: str):
|
|
113
|
+
raise NotImplementedError()
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
if __name__ == '__main__':
|
|
117
|
+
index = IndexClient()
|
|
118
|
+
index.add(42, np.array([0.4] * 256, dtype=np.float32))
|
|
119
|
+
results = index.search(np.array([0.4] * 256, dtype=np.float32), 10)
|
|
120
|
+
print(results)
|
package/python/usearch/server.py
CHANGED
|
@@ -1,32 +1,49 @@
|
|
|
1
|
-
|
|
2
|
-
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
# -*- coding: utf-8 -*-
|
|
3
3
|
|
|
4
|
-
import
|
|
4
|
+
import os
|
|
5
|
+
import argparse
|
|
5
6
|
import numpy as np
|
|
6
|
-
from PIL import Image
|
|
7
7
|
|
|
8
|
+
from ucall.rich_posix import Server
|
|
9
|
+
from usearch.index import Index
|
|
8
10
|
|
|
9
|
-
def serve(ndim: int, metric: str = 'ip'):
|
|
10
11
|
|
|
11
|
-
|
|
12
|
-
|
|
12
|
+
def _results_to_json(results: tuple[np.ndarray, np.ndarray, np.ndarray], row: int) -> list[dict]:
|
|
13
|
+
count = results[2][row]
|
|
14
|
+
labels = results[0][row, :count]
|
|
15
|
+
distances = results[1][row, :count]
|
|
16
|
+
return [{'label': int(l), 'distance': float(d)} for l, d in zip(labels, distances)]
|
|
13
17
|
|
|
14
|
-
@server
|
|
15
|
-
def add_one(label: int, vector: np.array):
|
|
16
|
-
labels = np.array([label], dtype=np.longlong)
|
|
17
|
-
vectors = vector.reshape(vector.shape[0], 1)
|
|
18
|
-
index.add(labels, vectors, copy=True)
|
|
19
18
|
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
19
|
+
def _ascii_to_vector(string: str) -> np.ndarray:
|
|
20
|
+
"""
|
|
21
|
+
WARNING: A dirty performance hack!
|
|
22
|
+
Assuming the `f8` vectors in our implementations are just integers,
|
|
23
|
+
and generally contain scalars in the [0, 100] range, we can transmit
|
|
24
|
+
them as JSON-embedded strings. The only symbols we must avoid are
|
|
25
|
+
the double-quote '"' (code 22) and backslash '\' (code 60).
|
|
26
|
+
Printable ASCII characters are in [20, 126].
|
|
27
|
+
"""
|
|
28
|
+
vector = np.array(string, dtype=np.int8)
|
|
29
|
+
vector[vector == 124] = 60
|
|
30
|
+
vector -= 23
|
|
31
|
+
return vector
|
|
24
32
|
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
33
|
+
|
|
34
|
+
def serve(
|
|
35
|
+
ndim_: int, metric: str = 'ip',
|
|
36
|
+
port: int = 8545, threads: int = 1,
|
|
37
|
+
path: str = 'index.usearch', immutable: bool = False):
|
|
38
|
+
|
|
39
|
+
server = Server(port=port)
|
|
40
|
+
index = Index(ndim=ndim_, metric=metric)
|
|
41
|
+
|
|
42
|
+
if os.path.exists(path):
|
|
43
|
+
if immutable:
|
|
44
|
+
index.view(path)
|
|
45
|
+
else:
|
|
46
|
+
index.load(path)
|
|
30
47
|
|
|
31
48
|
@server
|
|
32
49
|
def size() -> int:
|
|
@@ -44,8 +61,72 @@ def serve(ndim: int, metric: str = 'ip'):
|
|
|
44
61
|
def connectivity() -> int:
|
|
45
62
|
return index.connectivity()
|
|
46
63
|
|
|
47
|
-
server
|
|
64
|
+
@server
|
|
65
|
+
def add_one(label: int, vector: np.ndarray):
|
|
66
|
+
print('adding', label, vector)
|
|
67
|
+
labels = np.array([label], dtype=np.longlong)
|
|
68
|
+
vectors = vector.flatten().reshape(vector.shape[0], 1)
|
|
69
|
+
index.add(labels, vectors)
|
|
70
|
+
|
|
71
|
+
@server
|
|
72
|
+
def add_many(labels: np.ndarray, vectors: np.ndarray):
|
|
73
|
+
labels = labels.astype(np.longlong)
|
|
74
|
+
index.add(labels, vectors, threads=threads)
|
|
75
|
+
|
|
76
|
+
@server
|
|
77
|
+
def search_one(vector: np.ndarray, count: int) -> list[dict]:
|
|
78
|
+
print('search', vector, count)
|
|
79
|
+
vectors = vector.reshape(vector.shape[0], 1)
|
|
80
|
+
results = index.search(vectors, count)
|
|
81
|
+
return _results_to_json(results, 0)
|
|
82
|
+
|
|
83
|
+
@server
|
|
84
|
+
def search_many(vectors: np.ndarray, count: int) -> list[list[dict]]:
|
|
85
|
+
results = index.search(vectors, count)
|
|
86
|
+
return [_results_to_json(results, i) for i in range(vectors.shape[0])]
|
|
87
|
+
|
|
88
|
+
@server
|
|
89
|
+
def add_ascii(label: int, string: str):
|
|
90
|
+
return add_one(label, _ascii_to_vector(string))
|
|
91
|
+
|
|
92
|
+
@server
|
|
93
|
+
def search_ascii(string: str, count: int):
|
|
94
|
+
return search_one(_ascii_to_vector(string), count)
|
|
95
|
+
|
|
96
|
+
try:
|
|
97
|
+
server.run()
|
|
98
|
+
except KeyboardInterrupt:
|
|
99
|
+
if not immutable:
|
|
100
|
+
index.save(path)
|
|
48
101
|
|
|
49
102
|
|
|
50
103
|
if __name__ == '__main__':
|
|
51
|
-
|
|
104
|
+
|
|
105
|
+
parser = argparse.ArgumentParser()
|
|
106
|
+
parser.add_argument('-v', '--verbose', help='log server activity')
|
|
107
|
+
parser.add_argument(
|
|
108
|
+
'--ndim', type=int,
|
|
109
|
+
help='dimensionality of the vectors')
|
|
110
|
+
parser.add_argument(
|
|
111
|
+
'--immutable', type=bool, default=False,
|
|
112
|
+
help='the index can not be updated')
|
|
113
|
+
|
|
114
|
+
parser.add_argument(
|
|
115
|
+
'--metric', type=str, default='ip', choices=['ip', 'cos', 'l2sq', 'haversine'],
|
|
116
|
+
help='distance function to compare vectors')
|
|
117
|
+
parser.add_argument(
|
|
118
|
+
'-p', '--port', type=int, default=8545,
|
|
119
|
+
help='port to open for client connections')
|
|
120
|
+
parser.add_argument(
|
|
121
|
+
'-j', '--threads', type=int, default=1,
|
|
122
|
+
help='number of CPU threads to use')
|
|
123
|
+
parser.add_argument(
|
|
124
|
+
'--path', type=str, default='index.usearch',
|
|
125
|
+
help='where to store the index')
|
|
126
|
+
|
|
127
|
+
args = parser.parse_args()
|
|
128
|
+
assert args.ndim is not None, 'Define the number of dimensions!'
|
|
129
|
+
serve(
|
|
130
|
+
ndim_=args.ndim, metric=args.metric,
|
|
131
|
+
threads=args.threads, port=args.port,
|
|
132
|
+
path=args.path, immutable=args.immutable)
|
package/rust/lib.cpp
CHANGED
|
@@ -4,9 +4,9 @@
|
|
|
4
4
|
using namespace unum::usearch;
|
|
5
5
|
using namespace unum;
|
|
6
6
|
|
|
7
|
-
using
|
|
7
|
+
using punned_t = typename Index::punned_t;
|
|
8
8
|
|
|
9
|
-
Index::Index(std::shared_ptr<
|
|
9
|
+
Index::Index(std::shared_ptr<punned_t> index) : index_(index) {}
|
|
10
10
|
|
|
11
11
|
void Index::add_in_thread(uint32_t label, rust::Slice<float const> vector, size_t thread) const {
|
|
12
12
|
add_config_t config;
|
|
@@ -22,7 +22,8 @@ Matches Index::search_in_thread(rust::Slice<float const> vector, size_t count, s
|
|
|
22
22
|
results.labels.push_back(0), results.distances.push_back(0);
|
|
23
23
|
search_config_t config;
|
|
24
24
|
config.thread = thread;
|
|
25
|
-
results.count =
|
|
25
|
+
results.count =
|
|
26
|
+
index_->search(vector.data(), count, config).dump_to(results.labels.data(), results.distances.data());
|
|
26
27
|
results.labels.truncate(results.count);
|
|
27
28
|
results.distances.truncate(results.count);
|
|
28
29
|
return results;
|
|
@@ -38,7 +39,7 @@ Matches Index::search(rust::Slice<float const> vector, size_t count) const {
|
|
|
38
39
|
results.distances.reserve(count);
|
|
39
40
|
for (size_t i = 0; i != count; ++i)
|
|
40
41
|
results.labels.push_back(0), results.distances.push_back(0);
|
|
41
|
-
results.count = index_->search(vector.data(), count
|
|
42
|
+
results.count = index_->search(vector.data(), count).dump_to(results.labels.data(), results.distances.data());
|
|
42
43
|
results.labels.truncate(results.count);
|
|
43
44
|
results.distances.truncate(results.count);
|
|
44
45
|
return results;
|
|
@@ -57,31 +58,31 @@ void Index::view(rust::Str path) const { index_->view(std::string(path).c_str())
|
|
|
57
58
|
|
|
58
59
|
accuracy_t accuracy(rust::Str quant) { return accuracy_from_name(quant.data(), quant.size()); }
|
|
59
60
|
|
|
60
|
-
std::unique_ptr<Index> wrap(
|
|
61
|
-
std::shared_ptr<
|
|
61
|
+
std::unique_ptr<Index> wrap(punned_t&& index) {
|
|
62
|
+
std::shared_ptr<punned_t> native = std::make_shared<punned_t>(std::move(index));
|
|
62
63
|
return std::unique_ptr<Index>(new Index(native));
|
|
63
64
|
}
|
|
64
65
|
|
|
65
66
|
config_t config(size_t connectivity, size_t exp_add, size_t exp_search) {
|
|
66
67
|
config_t result;
|
|
67
|
-
result.connectivity = connectivity ? connectivity :
|
|
68
|
-
result.expansion_add = exp_add ? exp_add :
|
|
69
|
-
result.expansion_search = exp_search ? exp_search :
|
|
68
|
+
result.connectivity = connectivity ? connectivity : default_connectivity();
|
|
69
|
+
result.expansion_add = exp_add ? exp_add : default_expansion_add();
|
|
70
|
+
result.expansion_search = exp_search ? exp_search : default_expansion_search();
|
|
70
71
|
return result;
|
|
71
72
|
}
|
|
72
73
|
|
|
73
74
|
std::unique_ptr<Index> new_ip(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
|
|
74
|
-
return wrap(
|
|
75
|
+
return wrap(punned_t::ip(dims, accuracy(quant), config(connectivity, exp_add, exp_search)));
|
|
75
76
|
}
|
|
76
77
|
|
|
77
|
-
std::unique_ptr<Index>
|
|
78
|
-
return wrap(
|
|
78
|
+
std::unique_ptr<Index> new_l2sq(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
|
|
79
|
+
return wrap(punned_t::l2sq(dims, accuracy(quant), config(connectivity, exp_add, exp_search)));
|
|
79
80
|
}
|
|
80
81
|
|
|
81
82
|
std::unique_ptr<Index> new_cos(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
|
|
82
|
-
return wrap(
|
|
83
|
+
return wrap(punned_t::cos(dims, accuracy(quant), config(connectivity, exp_add, exp_search)));
|
|
83
84
|
}
|
|
84
85
|
|
|
85
86
|
std::unique_ptr<Index> new_haversine(rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
|
|
86
|
-
return wrap(
|
|
87
|
+
return wrap(punned_t::haversine(accuracy(quant), config(connectivity, exp_add, exp_search)));
|
|
87
88
|
}
|
package/rust/lib.hpp
CHANGED
|
@@ -3,7 +3,7 @@
|
|
|
3
3
|
|
|
4
4
|
#include <memory> // `std::shared_ptr`
|
|
5
5
|
|
|
6
|
-
#include "
|
|
6
|
+
#include "punned.hpp"
|
|
7
7
|
|
|
8
8
|
struct Matches;
|
|
9
9
|
|
|
@@ -15,9 +15,9 @@ class Index {
|
|
|
15
15
|
public:
|
|
16
16
|
using label_t = std::uint32_t;
|
|
17
17
|
using distance_t = unum::usearch::punned_distance_t;
|
|
18
|
-
using
|
|
18
|
+
using punned_t = unum::usearch::punned_gt<label_t>;
|
|
19
19
|
|
|
20
|
-
Index(std::shared_ptr<
|
|
20
|
+
Index(std::shared_ptr<punned_t> index);
|
|
21
21
|
|
|
22
22
|
void reserve(size_t) const;
|
|
23
23
|
|
|
@@ -37,10 +37,10 @@ class Index {
|
|
|
37
37
|
void view(rust::Str path) const;
|
|
38
38
|
|
|
39
39
|
private:
|
|
40
|
-
std::shared_ptr<
|
|
40
|
+
std::shared_ptr<punned_t> index_;
|
|
41
41
|
};
|
|
42
42
|
|
|
43
43
|
std::unique_ptr<Index> new_ip(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search);
|
|
44
|
-
std::unique_ptr<Index>
|
|
44
|
+
std::unique_ptr<Index> new_l2sq(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search);
|
|
45
45
|
std::unique_ptr<Index> new_cos(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search);
|
|
46
46
|
std::unique_ptr<Index> new_haversine(rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search);
|
package/rust/lib.rs
CHANGED
|
@@ -15,7 +15,7 @@ pub mod ffi {
|
|
|
15
15
|
type Index;
|
|
16
16
|
|
|
17
17
|
fn new_ip(dimensions: usize, quantization: &str, connectivity: usize, expansion_add: usize, expansion_search: usize) -> Result<UniquePtr<Index>>;
|
|
18
|
-
fn
|
|
18
|
+
fn new_l2sq(dimensions: usize, quantization: &str, connectivity: usize, expansion_add: usize, expansion_search: usize) -> Result<UniquePtr<Index>>;
|
|
19
19
|
fn new_cos(dimensions: usize, quantization: &str, connectivity: usize, expansion_add: usize, expansion_search: usize) -> Result<UniquePtr<Index>>;
|
|
20
20
|
fn new_haversine(quantization: &str, connectivity: usize, expansion_add: usize, expansion_search: usize) -> Result<UniquePtr<Index>>;
|
|
21
21
|
|
|
@@ -42,7 +42,7 @@ pub mod ffi {
|
|
|
42
42
|
#[cfg(test)]
|
|
43
43
|
mod tests {
|
|
44
44
|
use crate::ffi::new_ip;
|
|
45
|
-
use crate::ffi::
|
|
45
|
+
use crate::ffi::new_l2sq;
|
|
46
46
|
use crate::ffi::new_cos;
|
|
47
47
|
use crate::ffi::new_haversine;
|
|
48
48
|
|
|
@@ -76,7 +76,7 @@ mod tests {
|
|
|
76
76
|
|
|
77
77
|
// Make sure every function is called at least once
|
|
78
78
|
assert!(new_ip(5, &quant, 0, 0, 0).is_ok());
|
|
79
|
-
assert!(
|
|
79
|
+
assert!(new_l2sq(5, &quant, 0, 0, 0).is_ok());
|
|
80
80
|
assert!(new_cos(5, &quant, 0, 0, 0).is_ok());
|
|
81
81
|
assert!(new_haversine(&quant, 0, 0, 0).is_ok());
|
|
82
82
|
}
|
package/setup.py
CHANGED
|
@@ -1,5 +1,6 @@
|
|
|
1
1
|
import os
|
|
2
2
|
import sys
|
|
3
|
+
import platform
|
|
3
4
|
from setuptools import setup
|
|
4
5
|
from pybind11.setup_helpers import Pybind11Extension
|
|
5
6
|
|
|
@@ -10,15 +11,21 @@ link_args = []
|
|
|
10
11
|
if sys.platform == 'linux':
|
|
11
12
|
compile_args.append('-fopenmp')
|
|
12
13
|
link_args.append('-lgomp')
|
|
14
|
+
compile_args.append('-Wno-unknown-pragmas')
|
|
13
15
|
|
|
14
16
|
if sys.platform == 'darwin':
|
|
15
17
|
# MacOS 10.15 orhigher is needed for `aligned_alloc` support.
|
|
16
18
|
# https://github.com/unum-cloud/usearch/actions/runs/4975434891/jobs/8902603392
|
|
17
19
|
compile_args.append('-mmacosx-version-min=10.15')
|
|
20
|
+
compile_args.append('-Wno-unknown-pragmas')
|
|
18
21
|
|
|
19
22
|
if sys.platform == 'win32':
|
|
20
23
|
compile_args = ['/std:c++14', '/O2']
|
|
21
24
|
|
|
25
|
+
if platform.machine() == 'arm64':
|
|
26
|
+
compile_args.append('-march=armv8.2-a+simd+fp16+fp16fml+dotprod')
|
|
27
|
+
elif platform.machine() in ('i386', 'AMD64', 'x86_64'):
|
|
28
|
+
compile_args.append('-march=haswell')
|
|
22
29
|
|
|
23
30
|
ext_modules = [
|
|
24
31
|
Pybind11Extension(
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
cmake_minimum_required(VERSION 3.1)
|
|
2
|
+
|
|
3
|
+
option(USEARCH_USE_OPENMP "Use OpenMP for a thread pool" OFF)
|
|
4
|
+
option(USEARCH_USE_SIMSIMD "Use SimSIMD hardware-accelerated metrics" OFF)
|
|
5
|
+
option(USEARCH_USE_JEMALLOC "Use JeMalloc for faster memory allocations" OFF)
|
|
6
|
+
|
|
7
|
+
# Make "Release" by default
|
|
8
|
+
if(NOT CMAKE_BUILD_TYPE)
|
|
9
|
+
set(CMAKE_BUILD_TYPE Release)
|
|
10
|
+
endif()
|
|
11
|
+
|
|
12
|
+
# List of all possible compiler IDs:
|
|
13
|
+
# https://cmake.org/cmake/help/latest/variable/CMAKE_LANG_COMPILER_ID.html
|
|
14
|
+
if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
|
|
15
|
+
set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fsanitize=address")
|
|
16
|
+
set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -Wall -Wextra -Wno-conversion -Wno-unknown-pragmas")
|
|
17
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=native")
|
|
18
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fmax-errors=1")
|
|
19
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp")
|
|
20
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -pedantic")
|
|
21
|
+
|
|
22
|
+
elseif(CMAKE_CXX_COMPILER_ID MATCHES "Clang")
|
|
23
|
+
set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -Wfatal-errors")
|
|
24
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp")
|
|
25
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -pedantic")
|
|
26
|
+
|
|
27
|
+
elseif(CMAKE_CXX_COMPILER_ID STREQUAL "NVIDIA" OR CMAKE_CXX_COMPILER_ID STREQUAL "NVHPC")
|
|
28
|
+
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr --extended-lambda")
|
|
29
|
+
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -gencode arch=compute_86,code=sm_86")
|
|
30
|
+
set_property(SOURCE bench.cpp PROPERTY LANGUAGE CUDA)
|
|
31
|
+
set_target_properties(bench PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
|
32
|
+
set_target_properties(bench PROPERTIES CUDA_ARCHITECTURES "86")
|
|
33
|
+
|
|
34
|
+
elseif(CMAKE_CXX_COMPILER_ID STREQUAL "IntelLLVM")
|
|
35
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -w")
|
|
36
|
+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -ferror-limit=1")
|
|
37
|
+
|
|
38
|
+
elseif(CMAKE_CXX_COMPILER_ID MATCHES "MSVC")
|
|
39
|
+
set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} /DEBUG")
|
|
40
|
+
set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /O2")
|
|
41
|
+
set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} /O2 /DEBUG")
|
|
42
|
+
else()
|
|
43
|
+
set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -O0")
|
|
44
|
+
set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O3")
|
|
45
|
+
set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -O3 -pg -g")
|
|
46
|
+
endif()
|
|
47
|
+
|
|
48
|
+
if(APPLE)
|
|
49
|
+
set(CMAKE_THREAD_LIBS_INIT "-lpthread")
|
|
50
|
+
set(CMAKE_HAVE_THREADS_LIBRARY 1)
|
|
51
|
+
set(CMAKE_USE_WIN32_THREADS_INIT 0)
|
|
52
|
+
set(CMAKE_USE_PTHREADS_INIT 1)
|
|
53
|
+
set(THREADS_PREFER_PTHREAD_FLAG ON)
|
|
54
|
+
endif()
|
|
55
|
+
|
|
56
|
+
find_package(Threads REQUIRED)
|
|
57
|
+
|
|
58
|
+
if(${USEARCH_USE_OPENMP})
|
|
59
|
+
find_package(OpenMP REQUIRED)
|
|
60
|
+
include_directories(${OPENMP_INCLUDE_DIRS})
|
|
61
|
+
endif()
|
|
62
|
+
|
|
63
|
+
if(${USEARCH_USE_JEMALLOC})
|
|
64
|
+
include(jemalloc)
|
|
65
|
+
endif()
|
|
66
|
+
|
|
67
|
+
include(clipp)
|
|
68
|
+
|
|
69
|
+
set(USEARCH_PUNNED_INCLUDE_DIRS
|
|
70
|
+
"${CMAKE_CURRENT_SOURCE_DIR}/../include"
|
|
71
|
+
"${CMAKE_CURRENT_SOURCE_DIR}/../fp16/include"
|
|
72
|
+
"${CMAKE_CURRENT_SOURCE_DIR}/../simsimd/include"
|
|
73
|
+
"${CMAKE_CURRENT_SOURCE_DIR}/"
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
if(${USEARCH_BUILD_TEST})
|
|
77
|
+
add_executable(test test.cpp)
|
|
78
|
+
target_link_libraries(test PRIVATE Threads::Threads)
|
|
79
|
+
target_include_directories(test PRIVATE ${USEARCH_PUNNED_INCLUDE_DIRS})
|
|
80
|
+
set_target_properties(test PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
|
|
81
|
+
|
|
82
|
+
if(${CMAKE_VERSION} VERSION_EQUAL 3.13 OR ${CMAKE_VERSION} VERSION_GREATER 3.13)
|
|
83
|
+
include(CTest)
|
|
84
|
+
enable_testing()
|
|
85
|
+
add_test(NAME test COMMAND test)
|
|
86
|
+
endif()
|
|
87
|
+
endif()
|
|
88
|
+
|
|
89
|
+
if(${USEARCH_BUILD_BENCHMARK})
|
|
90
|
+
add_executable(bench bench.cpp)
|
|
91
|
+
target_link_libraries(bench PRIVATE Threads::Threads)
|
|
92
|
+
target_include_directories(bench PRIVATE ${USEARCH_PUNNED_INCLUDE_DIRS})
|
|
93
|
+
set_target_properties(bench PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
|
|
94
|
+
|
|
95
|
+
if(${USEARCH_USE_SIMSIMD})
|
|
96
|
+
target_compile_definitions(bench PRIVATE USEARCH_USE_SIMSIMD=1)
|
|
97
|
+
endif()
|
|
98
|
+
|
|
99
|
+
if(${USEARCH_USE_OPENMP})
|
|
100
|
+
target_compile_definitions(bench PRIVATE USEARCH_USE_OPENMP=1)
|
|
101
|
+
target_link_libraries(bench PRIVATE ${OPENMP_LIBRARIES})
|
|
102
|
+
endif()
|
|
103
|
+
|
|
104
|
+
if(${USEARCH_USE_JEMALLOC})
|
|
105
|
+
target_link_libraries(bench PRIVATE ${JEMALLOC_LIBRARIES})
|
|
106
|
+
endif()
|
|
107
|
+
endif()
|