usearch 0.18.8 → 0.19.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.
@@ -4,9 +4,8 @@ import numpy as np
4
4
 
5
5
 
6
6
  def load_matrix(
7
- filename: str,
8
- start_row: int = 0, count_rows: int = None,
9
- view: bool = False) -> np.ndarray:
7
+ filename: str, start_row: int = 0, count_rows: int = None, view: bool = False
8
+ ) -> np.ndarray:
10
9
  """Read *.ibin, *.bbib, *.hbin, *.fbin, *.dbin files with matrices.
11
10
 
12
11
  :param filename: path to the matrix file
@@ -19,36 +18,37 @@ def load_matrix(
19
18
  """
20
19
  dtype = np.float32
21
20
  scalar_size = 4
22
- if filename.endswith('.fbin'):
21
+ if filename.endswith(".fbin"):
23
22
  dtype = np.float32
24
23
  scalar_size = 4
25
- elif filename.endswith('.dbin'):
24
+ elif filename.endswith(".dbin"):
26
25
  dtype = np.float64
27
26
  scalar_size = 8
28
- elif filename.endswith('.hbin'):
27
+ elif filename.endswith(".hbin"):
29
28
  dtype = np.float16
30
29
  scalar_size = 2
31
- elif filename.endswith('.ibin'):
30
+ elif filename.endswith(".ibin"):
32
31
  dtype = np.int32
33
32
  scalar_size = 4
34
- elif filename.endswith('.bbin'):
33
+ elif filename.endswith(".bbin"):
35
34
  dtype = np.int8
36
35
  scalar_size = 1
37
36
  else:
38
- raise Exception('Unknown file type')
37
+ raise Exception("Unknown file type")
39
38
 
40
- with open(filename, 'rb') as f:
39
+ with open(filename, "rb") as f:
41
40
  rows, cols = np.fromfile(f, count=2, dtype=np.int32)
42
41
  rows = (rows - start_row) if count_rows is None else count_rows
43
42
  row_offset = start_row * scalar_size * cols
44
43
 
45
44
  if view:
46
45
  return np.memmap(
47
- f, dtype=dtype, mode='r', offset=8+row_offset, shape=(rows, cols))
46
+ f, dtype=dtype, mode="r", offset=8 + row_offset, shape=(rows, cols)
47
+ )
48
48
  else:
49
49
  return np.fromfile(
50
- f, count=rows * cols,
51
- dtype=dtype, offset=row_offset).reshape(rows, cols)
50
+ f, count=rows * cols, dtype=dtype, offset=row_offset
51
+ ).reshape(rows, cols)
52
52
 
53
53
 
54
54
  def save_matrix(vectors: np.ndarray, filename: str):
@@ -60,22 +60,22 @@ def save_matrix(vectors: np.ndarray, filename: str):
60
60
  :type filename: str
61
61
  """
62
62
  dtype = np.float32
63
- if filename.endswith('.fbin'):
63
+ if filename.endswith(".fbin"):
64
64
  dtype = np.float32
65
- elif filename.endswith('.dbin'):
65
+ elif filename.endswith(".dbin"):
66
66
  dtype = np.float64
67
- elif filename.endswith('.hbin'):
67
+ elif filename.endswith(".hbin"):
68
68
  dtype = np.float16
69
- elif filename.endswith('.ibin'):
69
+ elif filename.endswith(".ibin"):
70
70
  dtype = np.int32
71
- elif filename.endswith('.bbin'):
71
+ elif filename.endswith(".bbin"):
72
72
  dtype = np.int8
73
73
  else:
74
- raise Exception('Unknown file type')
74
+ raise Exception("Unknown file type")
75
75
 
76
- assert len(vectors.shape) == 2, 'Input array must have 2 dimensions'
77
- with open(filename, 'wb') as f:
76
+ assert len(vectors.shape) == 2, "Input array must have 2 dimensions"
77
+ with open(filename, "wb") as f:
78
78
  count, dim = vectors.shape
79
- f.write(struct.pack('<i', count))
80
- f.write(struct.pack('<i', dim))
79
+ f.write(struct.pack("<i", count))
80
+ f.write(struct.pack("<i", dim))
81
81
  vectors.astype(dtype).flatten().tofile(f)
@@ -7,7 +7,6 @@ from index import Index, Matches, Label
7
7
 
8
8
 
9
9
  class bidict(dict):
10
-
11
10
  def __init__(self, *args, **kwargs):
12
11
  super(bidict, self).__init__(*args, **kwargs)
13
12
  self.inverse = {}
@@ -37,10 +36,10 @@ def index_in_array(array: np.ndarray, item: Label) -> int:
37
36
 
38
37
 
39
38
  def semantic_join(a: Index, b: Index, resolution: int = 10) -> bidict:
40
- """Performs a Semantic Join, matching one entry from `a` with
39
+ """Performs a Semantic Join, matching one entry from `a` with
41
40
  one entry from `b`, converging towards a Stable Marriage.
42
41
  Assuming the collections can be different in size, classical solution
43
- doesn't provide
42
+ doesn't provide
44
43
 
45
44
  :param resolution: Approximate matches per member to consider, defaults to 10
46
45
  :type resolution: int, optional
@@ -65,7 +64,6 @@ def semantic_join(a: Index, b: Index, resolution: int = 10) -> bidict:
65
64
  idle_cycles: int = 0
66
65
  hopeless_men_count = len(men) - len(women)
67
66
  while len(free_men) > hopeless_men_count:
68
-
69
67
  # In the worst case scenario we may need to match more candidates for every
70
68
  # remaining man. This, however, may drastically increase the runtime.
71
69
  if len(free_men) == idle_cycles:
@@ -3,34 +3,34 @@
3
3
  # Python tooling, linters, and static analyzers. It also embeds JIT
4
4
  # into the primary `Index` class, connecting USearch with Numba.
5
5
  from math import sqrt
6
- from typing import Callable
7
6
 
8
- import numpy as np
9
7
  from numba import cfunc, types, carray
10
8
 
11
- from usearch.compiled import MetricKind, ScalarKind
9
+ from usearch.index import MetricKind, ScalarKind, MetricSignature, CompiledMetric
12
10
 
13
11
 
14
- signature_i8args = types.float32(types.CPointer(
15
- types.int8), types.CPointer(types.int8))
12
+ signature_i8args = types.float32(types.CPointer(types.int8), types.CPointer(types.int8))
16
13
 
17
- signature_f16args = types.float32(types.CPointer(
18
- types.float16), types.CPointer(types.float16))
14
+ signature_f16args = types.float32(
15
+ types.CPointer(types.float16), types.CPointer(types.float16)
16
+ )
19
17
 
20
- signature_f32args = types.float32(types.CPointer(
21
- types.float32), types.CPointer(types.float32))
18
+ signature_f32args = types.float32(
19
+ types.CPointer(types.float32), types.CPointer(types.float32)
20
+ )
22
21
 
23
- signature_f64args = types.float32(types.CPointer(
24
- types.float64), types.CPointer(types.float64))
22
+ signature_f64args = types.float32(
23
+ types.CPointer(types.float64), types.CPointer(types.float64)
24
+ )
25
25
 
26
26
 
27
- def jit(ndim: int,
28
- metric: MetricKind = MetricKind.Cos,
29
- dtype: ScalarKind = ScalarKind.F32) -> Callable:
27
+ def jit(
28
+ ndim: int, metric: MetricKind = MetricKind.Cos, dtype: ScalarKind = ScalarKind.F32
29
+ ) -> CompiledMetric:
30
30
  """JIT-compiles a distance metric specifically tuned for the target hardware
31
31
  and number of dimensions.
32
32
 
33
- Uses Numba `cfunc` functionality, annotating it with Numba `types` instead
33
+ Uses Numba `cfunc` functionality, annotating it with Numba `types` instead
34
34
  of `ctypes` to support half-precision.
35
35
  https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
36
36
  """
@@ -88,8 +88,7 @@ def jit(ndim: int,
88
88
  b_array = carray(b, ndim)
89
89
  ab_delta_sq = accumulator(0)
90
90
  for i in range(ndim):
91
- ab_delta_sq += (a_array[i] - b_array[i]) * \
92
- (a_array[i] - b_array[i])
91
+ ab_delta_sq += (a_array[i] - b_array[i]) * (a_array[i] - b_array[i])
93
92
  return types.float32(ab_delta_sq)
94
93
 
95
94
  scalar_kind_to_signature = {
@@ -108,4 +107,9 @@ def jit(ndim: int,
108
107
  if dtype == ScalarKind.F8 and metric == MetricKind.IP:
109
108
  metric = MetricKind.Cos
110
109
 
111
- return cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])
110
+ pointer = cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])
111
+ return CompiledMetric(
112
+ pointer=pointer,
113
+ kind=metric,
114
+ signature=MetricSignature.ArrayArray,
115
+ )
@@ -26,10 +26,13 @@ def _ascii_to_vector(string: str) -> np.ndarray:
26
26
 
27
27
 
28
28
  def serve(
29
- ndim_: int, metric: str = 'ip',
30
- port: int = 8545, threads: int = 1,
31
- path: str = 'index.usearch', immutable: bool = False):
32
-
29
+ ndim_: int,
30
+ metric: str = "ip",
31
+ port: int = 8545,
32
+ threads: int = 1,
33
+ path: str = "index.usearch",
34
+ immutable: bool = False,
35
+ ):
33
36
  server = Server(port=port)
34
37
  index = Index(ndim=ndim_, metric=metric)
35
38
 
@@ -57,7 +60,7 @@ def serve(
57
60
 
58
61
  @server
59
62
  def add_one(label: int, vector: np.ndarray):
60
- print('adding', label, vector)
63
+ print("adding", label, vector)
61
64
  labels = np.array([label], dtype=Label)
62
65
  vectors = vector.flatten().reshape(vector.shape[0], 1)
63
66
  index.add(labels, vectors)
@@ -68,7 +71,7 @@ def serve(
68
71
 
69
72
  @server
70
73
  def search_one(vector: np.ndarray, count: int) -> List[dict]:
71
- print('search', vector, count)
74
+ print("search", vector, count)
72
75
  vectors = vector.reshape(vector.shape[0], 1)
73
76
  results: Matches = index.search(vectors, count)
74
77
  return results.to_list()
@@ -93,33 +96,42 @@ def serve(
93
96
  index.save(path)
94
97
 
95
98
 
96
- if __name__ == '__main__':
97
-
99
+ if __name__ == "__main__":
98
100
  parser = argparse.ArgumentParser()
99
- parser.add_argument('-v', '--verbose', help='log server activity')
100
- parser.add_argument(
101
- '--ndim', type=int,
102
- help='dimensionality of the vectors')
101
+ parser.add_argument("-v", "--verbose", help="log server activity")
102
+ parser.add_argument("--ndim", type=int, help="dimensionality of the vectors")
103
103
  parser.add_argument(
104
- '--immutable', type=bool, default=False,
105
- help='the index can not be updated')
104
+ "--immutable", type=bool, default=False, help="the index can not be updated"
105
+ )
106
106
 
107
107
  parser.add_argument(
108
- '--metric', type=str, default='ip', choices=['ip', 'cos', 'l2sq', 'haversine'],
109
- help='distance function to compare vectors')
108
+ "--metric",
109
+ type=str,
110
+ default="ip",
111
+ choices=["ip", "cos", "l2sq", "haversine"],
112
+ help="distance function to compare vectors",
113
+ )
110
114
  parser.add_argument(
111
- '-p', '--port', type=int, default=8545,
112
- help='port to open for client connections')
115
+ "-p",
116
+ "--port",
117
+ type=int,
118
+ default=8545,
119
+ help="port to open for client connections",
120
+ )
113
121
  parser.add_argument(
114
- '-j', '--threads', type=int, default=1,
115
- help='number of CPU threads to use')
122
+ "-j", "--threads", type=int, default=1, help="number of CPU threads to use"
123
+ )
116
124
  parser.add_argument(
117
- '--path', type=str, default='index.usearch',
118
- help='where to store the index')
125
+ "--path", type=str, default="index.usearch", help="where to store the index"
126
+ )
119
127
 
120
128
  args = parser.parse_args()
121
- assert args.ndim is not None, 'Define the number of dimensions!'
129
+ assert args.ndim is not None, "Define the number of dimensions!"
122
130
  serve(
123
- ndim_=args.ndim, metric=args.metric,
124
- threads=args.threads, port=args.port,
125
- path=args.path, immutable=args.immutable)
131
+ ndim_=args.ndim,
132
+ metric=args.metric,
133
+ threads=args.threads,
134
+ port=args.port,
135
+ path=args.path,
136
+ immutable=args.immutable,
137
+ )
@@ -63,7 +63,8 @@
63
63
  "streambuf": "cpp",
64
64
  "thread": "cpp",
65
65
  "typeinfo": "cpp",
66
- "complex": "cpp"
66
+ "complex": "cpp",
67
+ "immintrin.h": "c"
67
68
  },
68
69
  "cSpell.words": [
69
70
  "dtype",
@@ -52,7 +52,7 @@ endif()
52
52
 
53
53
  # Define our header-only library.
54
54
  add_library(simsimd INTERFACE)
55
- target_sources(simsimd INTERFACE include/simsimd/simsimd.hpp)
55
+ target_sources(simsimd INTERFACE include/simsimd/simsimd.h)
56
56
  target_include_directories(simsimd INTERFACE "${PROJECT_SOURCE_DIR}/include")
57
57
 
58
58
  # Build the benchmarks if needed.
package/simsimd/README.md CHANGED
@@ -66,8 +66,20 @@ With AVX2 and 256-bit registers the results should be better, but would be less
66
66
 
67
67
  ---
68
68
 
69
- To replicate on your hardware, please run:
69
+ To replicate on your hardware, please run following on Linux:
70
70
 
71
71
  ```sh
72
- cmake -DCMAKE_BUILD_TYPE=Release -DUNSW_BUILD_BENCHMARKS=1 -B ./build && make -C ./build && ./build/simsimd_bench
72
+ cmake -DCMAKE_BUILD_TYPE=Release -DSIMSIMD_BUILD_BENCHMARKS=1 -B ./build && make -C ./build && ./build/simsimd_bench
73
73
  ```
74
+
75
+ MacOS:
76
+
77
+ ```sh
78
+ brew install llvm
79
+ cmake -B ./build \
80
+ -DCMAKE_C_COMPILER="/opt/homebrew/opt/llvm/bin/clang" \
81
+ -DCMAKE_CXX_COMPILER="/opt/homebrew/opt/llvm/bin/clang++" \
82
+ -DSIMSIMD_BUILD_BENCHMARKS=1 \
83
+ && \
84
+ make -C ./build -j && ./build/simsimd_bench
85
+ ```
package/simsimd/bench.cxx CHANGED
@@ -1,46 +1,35 @@
1
+ #include <thread>
2
+
1
3
  #include <benchmark/benchmark.h>
2
4
 
3
- #include <simsimd/simsimd.hpp>
5
+ #include <simsimd/simsimd.h>
4
6
 
5
7
  namespace bm = benchmark;
6
- using namespace av::simsimd;
7
8
 
8
- static constexpr std::size_t threads_k = 64;
9
- static constexpr std::size_t time_k = 100;
9
+ static const std::size_t threads_k = std::thread::hardware_concurrency();
10
+ static const std::size_t time_k = 10;
10
11
 
11
- template <typename metric_at, typename scalar_at, std::size_t bytes_per_vector_ak,
12
+ template <typename scalar_at, typename metric_at, std::size_t bytes_per_vector_ak = 256,
12
13
  std::size_t dimensions_ak = bytes_per_vector_ak / sizeof(scalar_at)> //
13
- static void measure(bm::State& state) {
14
+ static void measure(bm::State& state, metric_at metric) {
14
15
 
15
- constexpr std::size_t buffer_size_k = bytes_per_vector_ak / sizeof(scalar_at);
16
- alignas(64) scalar_at a[buffer_size_k]{};
17
- alignas(64) scalar_at b[buffer_size_k]{};
18
- scalar_at c{};
16
+ alignas(64) scalar_at a[dimensions_ak]{};
17
+ alignas(64) scalar_at b[dimensions_ak]{};
18
+ float c{};
19
19
 
20
- std::fill_n(a, buffer_size_k, static_cast<scalar_at>(1));
21
- std::fill_n(b, buffer_size_k, static_cast<scalar_at>(2));
20
+ std::fill_n(a, dimensions_ak, static_cast<scalar_at>(1));
21
+ std::fill_n(b, dimensions_ak, static_cast<scalar_at>(2));
22
22
 
23
23
  for (auto _ : state)
24
- bm::DoNotOptimize((c = metric_at{}(a, b, dimensions_ak)));
24
+ bm::DoNotOptimize((c = metric(a, b, dimensions_ak)));
25
25
 
26
26
  state.SetBytesProcessed(state.iterations() * bytes_per_vector_ak * 2u);
27
27
  state.SetItemsProcessed(state.iterations());
28
28
  }
29
29
 
30
- BENCHMARK_TEMPLATE(measure, euclidean_distance_t, f32_t, 32)->Threads(threads_k)->MinTime(time_k);
31
- BENCHMARK_TEMPLATE(measure, euclidean_distance_t, f32_t, 256)->Threads(threads_k)->MinTime(time_k);
32
- BENCHMARK_TEMPLATE(measure, euclidean_distance_t, f16_t, 32)->Threads(threads_k)->MinTime(time_k);
33
- BENCHMARK_TEMPLATE(measure, euclidean_distance_t, f16_t, 256)->Threads(threads_k)->MinTime(time_k);
34
-
35
- BENCHMARK_TEMPLATE(measure, cosine_similarity_t, f32_t, 32)->Threads(threads_k)->MinTime(time_k);
36
- BENCHMARK_TEMPLATE(measure, cosine_similarity_t, f32_t, 256)->Threads(threads_k)->MinTime(time_k);
37
- BENCHMARK_TEMPLATE(measure, cosine_similarity_f32x4k_t, f32_t, 32)->Threads(threads_k)->MinTime(time_k);
38
- BENCHMARK_TEMPLATE(measure, cosine_similarity_f32x4k_t, f32_t, 256)->Threads(threads_k)->MinTime(time_k);
39
-
40
- BENCHMARK_TEMPLATE(measure, hamming_bits_distance_t, u64_t, 32, 256)->Threads(threads_k)->MinTime(time_k);
41
- BENCHMARK_TEMPLATE(measure, hamming_bits_distance_t, u64_t, 256, 2048)->Threads(threads_k)->MinTime(time_k);
42
- BENCHMARK_TEMPLATE(measure, hamming_bits_distance_u1x128k_t, u64_t, 32, 256)->Threads(threads_k)->MinTime(time_k);
43
- BENCHMARK_TEMPLATE(measure, hamming_bits_distance_u1x128k_t, u64_t, 256, 2048)->Threads(threads_k)->MinTime(time_k);
30
+ template <typename scalar_at, typename metric_at> void register_(char const* name, metric_at distance_func) {
31
+ bm::RegisterBenchmark(name, measure<scalar_at, metric_at>, distance_func)->Threads(threads_k)->MinTime(time_k);
32
+ }
44
33
 
45
34
  int main(int argc, char** argv) {
46
35
 
@@ -76,6 +65,34 @@ int main(int argc, char** argv) {
76
65
  bm::Initialize(&argc, argv);
77
66
  if (bm::ReportUnrecognizedArguments(argc, argv))
78
67
  return 1;
68
+
69
+ #if defined(__ARM_FEATURE_SVE)
70
+ register_<simsimd_f32_t>("dot_f32sve", simsimd_dot_f32sve);
71
+ register_<simsimd_f32_t>("cos_f32sve", simsimd_cos_f32sve);
72
+ register_<simsimd_f32_t>("l2sq_f32sve", simsimd_l2sq_f32sve);
73
+ register_<std::int16_t>("l2sq_f16sve", simsimd_l2sq_f16sve);
74
+ register_<std::uint8_t>("hamming_b1x8sve", simsimd_hamming_b1x8sve);
75
+ register_<std::uint8_t>("hamming_b1x128sve", simsimd_hamming_b1x128sve);
76
+ #endif
77
+
78
+ #if defined(__ARM_NEON)
79
+ register_<simsimd_f32_t>("dot_f32x4neon", simsimd_dot_f32x4neon);
80
+ register_<std::int16_t>("cos_f16x4neon", simsimd_cos_f16x4neon);
81
+ register_<std::int8_t>("cos_i8x16neon", simsimd_cos_i8x16neon);
82
+ register_<simsimd_f32_t>("cos_f32x4neon", simsimd_cos_f32x4neon);
83
+ #endif
84
+
85
+ #if defined(__AVX2__)
86
+ register_<simsimd_f32_t>("dot_f32x4avx2", simsimd_dot_f32x4avx2);
87
+ register_<std::int8_t>("dot_i8x16avx2", simsimd_dot_i8x16avx2);
88
+ register_<simsimd_f32_t>("cos_f32x4avx2", simsimd_cos_f32x4avx2);
89
+ #endif
90
+
91
+ #if defined(__AVX512F__)
92
+ register_<std::int16_t>("cos_f16x16avx512", simsimd_cos_f16x16avx512);
93
+ register_<std::uint8_t>("hamming_b1x128avx512", simsimd_hamming_b1x128avx512);
94
+ #endif
95
+
79
96
  bm::RunSpecifiedBenchmarks();
80
97
  bm::Shutdown();
81
98
  return 0;
@@ -10,7 +10,8 @@
10
10
  */
11
11
 
12
12
  #pragma once
13
- #include <math.h> // `sqrt`
13
+ #include <math.h> // `sqrt`
14
+ #include <stddef.h> // `size_t`
14
15
 
15
16
  #if defined(__AVX2__)
16
17
  #include <immintrin.h>