usearch 2.15.2 → 2.16.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.
Files changed (54) hide show
  1. package/README.md +1 -1
  2. package/binding.gyp +1 -1
  3. package/c/CMakeLists.txt +10 -3
  4. package/cpp/README.md +1 -1
  5. package/cpp/bench.cpp +2 -2
  6. package/cpp/test.cpp +16 -0
  7. package/include/usearch/index.hpp +44 -17
  8. package/include/usearch/index_dense.hpp +51 -89
  9. package/include/usearch/index_plugins.hpp +468 -43
  10. package/javascript/lib.cpp +17 -5
  11. package/javascript/usearch.test.js +205 -34
  12. package/package.json +1 -1
  13. package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
  14. package/prebuilds/linux-arm64/usearch.node +0 -0
  15. package/prebuilds/linux-x64/usearch.node +0 -0
  16. package/simsimd/.git-blame-ignore-revs +2 -0
  17. package/simsimd/CODE_OF_CONDUCT.md +128 -0
  18. package/simsimd/CONTRIBUTING.md +84 -26
  19. package/simsimd/Cargo.lock +1 -1
  20. package/simsimd/Cargo.toml +4 -4
  21. package/simsimd/MANIFEST.in +3 -0
  22. package/simsimd/Package.swift +15 -3
  23. package/simsimd/README.md +534 -177
  24. package/simsimd/binding.gyp +6 -2
  25. package/simsimd/build.rs +13 -1
  26. package/simsimd/c/lib.c +266 -53
  27. package/simsimd/include/simsimd/binary.h +321 -112
  28. package/simsimd/include/simsimd/curved.h +803 -0
  29. package/simsimd/include/simsimd/dot.h +694 -521
  30. package/simsimd/include/simsimd/elementwise.h +1437 -0
  31. package/simsimd/include/simsimd/geospatial.h +9 -3
  32. package/simsimd/include/simsimd/mesh.h +69 -0
  33. package/simsimd/include/simsimd/probability.h +234 -206
  34. package/simsimd/include/simsimd/simsimd.h +1476 -617
  35. package/simsimd/include/simsimd/sparse.h +1428 -0
  36. package/simsimd/include/simsimd/spatial.h +1556 -742
  37. package/simsimd/include/simsimd/types.h +276 -132
  38. package/simsimd/javascript/lib.c +1 -2
  39. package/simsimd/javascript/simsimd.ts +4 -4
  40. package/simsimd/package.json +4 -4
  41. package/simsimd/rust/lib.rs +179 -1
  42. package/simsimd/scripts/bench.cxx +1070 -0
  43. package/simsimd/scripts/bench.ipynb +187 -0
  44. package/simsimd/{javascript → scripts}/bench.js +1 -1
  45. package/simsimd/{rust/benches/sqeuclidean.rs → scripts/bench_sqeuclidean.rs} +6 -6
  46. package/simsimd/scripts/bench_vectors.py +700 -0
  47. package/simsimd/scripts/bench_vectors_live.py +213 -0
  48. package/simsimd/scripts/test.c +155 -0
  49. package/simsimd/{javascript → scripts}/test.mjs +2 -2
  50. package/simsimd/scripts/test.py +1410 -0
  51. package/simsimd/swift/{source/SimSIMD.swift → SimSIMD.swift} +6 -2
  52. package/simsimd/swift/{tests/SimSIMDTests.swift → Test.swift} +1 -1
  53. /package/simsimd/{rust/benches/cosine.rs → scripts/bench_cosine.rs} +0 -0
  54. /package/simsimd/{rust/benches/native.rs → scripts/bench_native.rs} +0 -0
@@ -0,0 +1,1428 @@
1
+ /**
2
+ * @file sparse.h
3
+ * @brief SIMD-accelerated functions for Sparse Vectors.
4
+ * @author Ash Vardanian
5
+ * @date March 21, 2024
6
+ *
7
+ * Contains:
8
+ * - Set Intersection ~ Jaccard Distance
9
+ * - Sparse Dot Products, outputting the count and weighted product
10
+ *
11
+ * For datatypes:
12
+ * - u16: for vocabularies under 64 thousand tokens
13
+ * - u32: for vocabularies under 4 billion tokens
14
+ * - u16 indicies + i16 weights: for weighted word counts
15
+ * - u16 indicies + bf16 weights: for sparse matrices
16
+ *
17
+ * For hardware architectures:
18
+ * - x86: Ice Lake, Turin
19
+ * - Arm: SVE2
20
+ *
21
+ * Interestingly, to implement sparse distances and products, the most important function
22
+ * is analogous to `std::set_intersection`, that outputs the intersection of two sorted
23
+ * sequences. The naive implementation of that function would look like:
24
+ *
25
+ * std::size_t intersection_size = 0;
26
+ * while (i != a_length && j != b_length) {
27
+ * scalar_t ai = a[i], bj = b[j];
28
+ * intersection_size += ai == bj;
29
+ * i += ai < bj;
30
+ * j += ai >= bj;
31
+ * }
32
+ *
33
+ * Assuming we are dealing with sparse arrays, most of the time we are just evaluating
34
+ * branches and skipping entries. So what if we could skip multiple entries at a time
35
+ * searching for the next chunk, where an intersection is possible. For weighted arrays:
36
+ *
37
+ * double product = 0;
38
+ * while (i != a_length && j != b_length) {
39
+ * scalar_t ai = a[i], bj = b[j];
40
+ * product += ai == bj ? a_weights[i] * b_weights[j] : 0;
41
+ * i += ai < bj;
42
+ * j += ai >= bj;
43
+ * }
44
+ *
45
+ * x86 intrinsics: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/
46
+ * Arm intrinsics: https://developer.arm.com/architectures/instruction-sets/intrinsics/
47
+ */
48
+ #ifndef SIMSIMD_SPARSE_H
49
+ #define SIMSIMD_SPARSE_H
50
+
51
+ #include "types.h"
52
+
53
+ #ifdef __cplusplus
54
+ extern "C" {
55
+ #endif
56
+
57
+ /* Implements the serial set intersection algorithm, similar to `std::set_intersection in C++ STL`,
58
+ * but uses clever galloping logic, if the arrays significantly differ in size.
59
+ */
60
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_serial( //
61
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
62
+ simsimd_size_t a_length, simsimd_size_t b_length, //
63
+ simsimd_distance_t *results);
64
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_serial( //
65
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
66
+ simsimd_size_t a_length, simsimd_size_t b_length, //
67
+ simsimd_distance_t *results);
68
+ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_serial( //
69
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
70
+ simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
71
+ simsimd_size_t a_length, simsimd_size_t b_length, //
72
+ simsimd_distance_t *results);
73
+ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_serial( //
74
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
75
+ simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
76
+ simsimd_size_t a_length, simsimd_size_t b_length, //
77
+ simsimd_distance_t *results);
78
+
79
+ /* Implements the most naive set intersection algorithm, similar to `std::set_intersection in C++ STL`,
80
+ * naively enumerating the elements of two arrays.
81
+ */
82
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_accurate( //
83
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
84
+ simsimd_size_t a_length, simsimd_size_t b_length, //
85
+ simsimd_distance_t *results);
86
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_accurate( //
87
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
88
+ simsimd_size_t a_length, simsimd_size_t b_length, //
89
+ simsimd_distance_t *results);
90
+ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_accurate( //
91
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
92
+ simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
93
+ simsimd_size_t a_length, simsimd_size_t b_length, //
94
+ simsimd_distance_t *results);
95
+ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_accurate( //
96
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
97
+ simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
98
+ simsimd_size_t a_length, simsimd_size_t b_length, //
99
+ simsimd_distance_t *results);
100
+
101
+ /* SIMD-powered backends for Arm SVE, mostly using 32-bit arithmetic over variable-length platform-defined word sizes.
102
+ * Designed for Arm Graviton 3, Microsoft Cobalt, as well as Nvidia Grace and newer Ampere Altra CPUs.
103
+ */
104
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_sve2( //
105
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
106
+ simsimd_size_t a_length, simsimd_size_t b_length, //
107
+ simsimd_distance_t *results);
108
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_sve2( //
109
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
110
+ simsimd_size_t a_length, simsimd_size_t b_length, //
111
+ simsimd_distance_t *results);
112
+ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2( //
113
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
114
+ simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
115
+ simsimd_size_t a_length, simsimd_size_t b_length, //
116
+ simsimd_distance_t *results);
117
+ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2( //
118
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
119
+ simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
120
+ simsimd_size_t a_length, simsimd_size_t b_length, //
121
+ simsimd_distance_t *results);
122
+
123
+ /* SIMD-powered backends for various generations of AVX512 CPUs.
124
+ * Skylake is handy, as it supports masked loads and other operations, avoiding the need for the tail loop.
125
+ * Ice Lake, however, is needed even for the most basic kernels to perform integer matching.
126
+ */
127
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_ice( //
128
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
129
+ simsimd_size_t a_length, simsimd_size_t b_length, //
130
+ simsimd_distance_t *results);
131
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_ice( //
132
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
133
+ simsimd_size_t a_length, simsimd_size_t b_length, //
134
+ simsimd_distance_t *results);
135
+
136
+ /* SIMD-powered backends for AMD Turin CPUs with cheap VP2INTERSECT instructions.
137
+ * On the Intel side, only mobile Tiger Lake support them, but have prohibitively high latency.
138
+ */
139
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_turin( //
140
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
141
+ simsimd_size_t a_length, simsimd_size_t b_length, //
142
+ simsimd_distance_t *results);
143
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_turin( //
144
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
145
+ simsimd_size_t a_length, simsimd_size_t b_length, //
146
+ simsimd_distance_t *results);
147
+ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_turin( //
148
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
149
+ simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
150
+ simsimd_size_t a_length, simsimd_size_t b_length, //
151
+ simsimd_distance_t *results);
152
+ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_turin( //
153
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
154
+ simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
155
+ simsimd_size_t a_length, simsimd_size_t b_length, //
156
+ simsimd_distance_t *results);
157
+
158
+ #define SIMSIMD_MAKE_INTERSECT_LINEAR(name, input_type, counter_type) \
159
+ SIMSIMD_PUBLIC void simsimd_intersect_##input_type##_##name( \
160
+ simsimd_##input_type##_t const *a, simsimd_##input_type##_t const *b, simsimd_size_t a_length, \
161
+ simsimd_size_t b_length, simsimd_distance_t *result) { \
162
+ simsimd_##counter_type##_t intersection_size = 0; \
163
+ simsimd_size_t i = 0, j = 0; \
164
+ while (i != a_length && j != b_length) { \
165
+ simsimd_##input_type##_t ai = a[i]; \
166
+ simsimd_##input_type##_t bj = b[j]; \
167
+ intersection_size += ai == bj; \
168
+ i += ai < bj; \
169
+ j += ai >= bj; \
170
+ } \
171
+ *result = intersection_size; \
172
+ }
173
+
174
+ SIMSIMD_MAKE_INTERSECT_LINEAR(accurate, u16, size) // simsimd_intersect_u16_accurate
175
+ SIMSIMD_MAKE_INTERSECT_LINEAR(accurate, u32, size) // simsimd_intersect_u32_accurate
176
+
177
+ #define SIMSIMD_MAKE_INTERSECT_WEIGHTED(name, variation, input_type, counter_type, weight_type, accumulator_type, \
178
+ load_and_convert) \
179
+ SIMSIMD_PUBLIC void simsimd_##variation##_##input_type##_##name( \
180
+ simsimd_##input_type##_t const *a, simsimd_##input_type##_t const *b, \
181
+ simsimd_##weight_type##_t const *a_weights, simsimd_##weight_type##_t const *b_weights, \
182
+ simsimd_size_t a_length, simsimd_size_t b_length, simsimd_distance_t *results) { \
183
+ simsimd_##counter_type##_t intersection_size = 0; \
184
+ simsimd_##accumulator_type##_t weights_product = 0; \
185
+ simsimd_size_t i = 0, j = 0; \
186
+ while (i != a_length && j != b_length) { \
187
+ simsimd_##input_type##_t ai = a[i]; \
188
+ simsimd_##input_type##_t bj = b[j]; \
189
+ int matches = ai == bj; \
190
+ simsimd_##counter_type##_t awi = load_and_convert(a_weights + i); \
191
+ simsimd_##counter_type##_t bwi = load_and_convert(b_weights + i); \
192
+ weights_product += matches * awi * bwi; \
193
+ intersection_size += matches; \
194
+ i += ai < bj; \
195
+ j += ai >= bj; \
196
+ } \
197
+ results[0] = intersection_size; \
198
+ results[1] = weights_product; \
199
+ }
200
+
201
+ SIMSIMD_MAKE_INTERSECT_WEIGHTED(accurate, spdot_counts, u16, size, i16, i64,
202
+ SIMSIMD_DEREFERENCE) // simsimd_spdot_counts_u16_accurate
203
+ SIMSIMD_MAKE_INTERSECT_WEIGHTED(accurate, spdot_weights, u16, size, bf16, f64,
204
+ SIMSIMD_BF16_TO_F32) // simsimd_spdot_weights_u16_accurate
205
+
206
+ #define SIMSIMD_MAKE_INTERSECT_GALLOPING(name, input_type, counter_type) \
207
+ SIMSIMD_PUBLIC simsimd_size_t simsimd_galloping_search_##input_type(simsimd_##input_type##_t const *array, \
208
+ simsimd_size_t start, simsimd_size_t length, \
209
+ simsimd_##input_type##_t val) { \
210
+ simsimd_size_t low = start; \
211
+ simsimd_size_t high = start + 1; \
212
+ while (high < length && array[high] < val) { \
213
+ low = high; \
214
+ high = (2 * high < length) ? 2 * high : length; \
215
+ } \
216
+ while (low < high) { \
217
+ simsimd_size_t mid = low + (high - low) / 2; \
218
+ if (array[mid] < val) { low = mid + 1; } \
219
+ else { high = mid; } \
220
+ } \
221
+ return low; \
222
+ } \
223
+ \
224
+ SIMSIMD_PUBLIC void simsimd_intersect_##input_type##_##name( \
225
+ simsimd_##input_type##_t const *shorter, simsimd_##input_type##_t const *longer, \
226
+ simsimd_size_t shorter_length, simsimd_size_t longer_length, simsimd_distance_t *result) { \
227
+ /* Swap arrays if necessary, as we want "longer" to be larger than "shorter" */ \
228
+ if (longer_length < shorter_length) { \
229
+ simsimd_##input_type##_t const *temp = shorter; \
230
+ shorter = longer; \
231
+ longer = temp; \
232
+ simsimd_size_t temp_length = shorter_length; \
233
+ shorter_length = longer_length; \
234
+ longer_length = temp_length; \
235
+ } \
236
+ \
237
+ /* Use the accurate implementation if galloping is not beneficial */ \
238
+ if (longer_length < 64 * shorter_length) { \
239
+ simsimd_intersect_##input_type##_accurate(shorter, longer, shorter_length, longer_length, result); \
240
+ return; \
241
+ } \
242
+ \
243
+ /* Perform galloping, shrinking the target range */ \
244
+ simsimd_##counter_type##_t intersection_size = 0; \
245
+ simsimd_size_t j = 0; \
246
+ for (simsimd_size_t i = 0; i < shorter_length; ++i) { \
247
+ simsimd_##input_type##_t shorter_i = shorter[i]; \
248
+ j = simsimd_galloping_search_##input_type(longer, j, longer_length, shorter_i); \
249
+ if (j < longer_length && longer[j] == shorter_i) { intersection_size++; } \
250
+ } \
251
+ *result = intersection_size; \
252
+ }
253
+
254
+ SIMSIMD_MAKE_INTERSECT_GALLOPING(serial, u16, size) // simsimd_intersect_u16_serial
255
+ SIMSIMD_MAKE_INTERSECT_GALLOPING(serial, u32, size) // simsimd_intersect_u32_serial
256
+ SIMSIMD_MAKE_INTERSECT_WEIGHTED(serial, spdot_counts, u16, size, i16, i32,
257
+ SIMSIMD_DEREFERENCE) // simsimd_spdot_counts_u16_serial
258
+ SIMSIMD_MAKE_INTERSECT_WEIGHTED(serial, spdot_weights, u16, size, bf16, f32,
259
+ SIMSIMD_BF16_TO_F32) // simsimd_spdot_weights_u16_serial
260
+
261
+ /* The AVX-512 implementations are inspired by the "Faster-Than-Native Alternatives
262
+ * for x86 VP2INTERSECT Instructions" paper by Guille Diez-Canas, 2022.
263
+ *
264
+ * https://github.com/mozonaut/vp2intersect
265
+ * https://arxiv.org/pdf/2112.06342.pdf
266
+ *
267
+ * For R&D purposes, it's important to keep the following latencies in mind:
268
+ *
269
+ * - `_mm512_permutex_epi64` - needs F - 3 cycles latency
270
+ * - `_mm512_shuffle_epi8` - needs BW - 1 cycle latency
271
+ * - `_mm512_permutexvar_epi16` - needs BW - 4-6 cycles latency
272
+ * - `_mm512_permutexvar_epi8` - needs VBMI - 3 cycles latency
273
+ */
274
+ #if _SIMSIMD_TARGET_X86
275
+ #if SIMSIMD_TARGET_ICE
276
+ #pragma GCC push_options
277
+ #pragma GCC target("avx2", "avx512f", "avx512vl", "bmi2", "lzcnt", "popcnt", "avx512bw", "avx512vbmi2")
278
+ #pragma clang attribute push(__attribute__((target("avx2,avx512f,avx512vl,bmi2,lzcnt,popcnt,avx512bw,avx512vbmi2"))), \
279
+ apply_to = function)
280
+
281
+ /**
282
+ * @brief Analogous to `_mm512_2intersect_epi16_mask`, but compatible with Ice Lake CPUs,
283
+ * slightly faster than the native Tiger Lake implementation, but returns only one mask.
284
+ */
285
+ SIMSIMD_INTERNAL simsimd_u32_t _simsimd_intersect_u16x32_ice(__m512i a, __m512i b) {
286
+ __m512i a1 = _mm512_alignr_epi32(a, a, 4);
287
+ __m512i a2 = _mm512_alignr_epi32(a, a, 8);
288
+ __m512i a3 = _mm512_alignr_epi32(a, a, 12);
289
+
290
+ __m512i b1 = _mm512_shuffle_epi32(b, _MM_PERM_ADCB);
291
+ __m512i b2 = _mm512_shuffle_epi32(b, _MM_PERM_BADC);
292
+ __m512i b3 = _mm512_shuffle_epi32(b, _MM_PERM_CBAD);
293
+
294
+ __m512i b01 = _mm512_shrdi_epi32(b, b, 16);
295
+ __m512i b11 = _mm512_shrdi_epi32(b1, b1, 16);
296
+ __m512i b21 = _mm512_shrdi_epi32(b2, b2, 16);
297
+ __m512i b31 = _mm512_shrdi_epi32(b3, b3, 16);
298
+
299
+ __mmask32 nm00 = _mm512_cmpneq_epi16_mask(a, b);
300
+ __mmask32 nm01 = _mm512_cmpneq_epi16_mask(a1, b);
301
+ __mmask32 nm02 = _mm512_cmpneq_epi16_mask(a2, b);
302
+ __mmask32 nm03 = _mm512_cmpneq_epi16_mask(a3, b);
303
+
304
+ __mmask32 nm10 = _mm512_mask_cmpneq_epi16_mask(nm00, a, b01);
305
+ __mmask32 nm11 = _mm512_mask_cmpneq_epi16_mask(nm01, a1, b01);
306
+ __mmask32 nm12 = _mm512_mask_cmpneq_epi16_mask(nm02, a2, b01);
307
+ __mmask32 nm13 = _mm512_mask_cmpneq_epi16_mask(nm03, a3, b01);
308
+
309
+ __mmask32 nm20 = _mm512_mask_cmpneq_epi16_mask(nm10, a, b1);
310
+ __mmask32 nm21 = _mm512_mask_cmpneq_epi16_mask(nm11, a1, b1);
311
+ __mmask32 nm22 = _mm512_mask_cmpneq_epi16_mask(nm12, a2, b1);
312
+ __mmask32 nm23 = _mm512_mask_cmpneq_epi16_mask(nm13, a3, b1);
313
+
314
+ __mmask32 nm30 = _mm512_mask_cmpneq_epi16_mask(nm20, a, b11);
315
+ __mmask32 nm31 = _mm512_mask_cmpneq_epi16_mask(nm21, a1, b11);
316
+ __mmask32 nm32 = _mm512_mask_cmpneq_epi16_mask(nm22, a2, b11);
317
+ __mmask32 nm33 = _mm512_mask_cmpneq_epi16_mask(nm23, a3, b11);
318
+
319
+ __mmask32 nm40 = _mm512_mask_cmpneq_epi16_mask(nm30, a, b2);
320
+ __mmask32 nm41 = _mm512_mask_cmpneq_epi16_mask(nm31, a1, b2);
321
+ __mmask32 nm42 = _mm512_mask_cmpneq_epi16_mask(nm32, a2, b2);
322
+ __mmask32 nm43 = _mm512_mask_cmpneq_epi16_mask(nm33, a3, b2);
323
+
324
+ __mmask32 nm50 = _mm512_mask_cmpneq_epi16_mask(nm40, a, b21);
325
+ __mmask32 nm51 = _mm512_mask_cmpneq_epi16_mask(nm41, a1, b21);
326
+ __mmask32 nm52 = _mm512_mask_cmpneq_epi16_mask(nm42, a2, b21);
327
+ __mmask32 nm53 = _mm512_mask_cmpneq_epi16_mask(nm43, a3, b21);
328
+
329
+ __mmask32 nm60 = _mm512_mask_cmpneq_epi16_mask(nm50, a, b3);
330
+ __mmask32 nm61 = _mm512_mask_cmpneq_epi16_mask(nm51, a1, b3);
331
+ __mmask32 nm62 = _mm512_mask_cmpneq_epi16_mask(nm52, a2, b3);
332
+ __mmask32 nm63 = _mm512_mask_cmpneq_epi16_mask(nm53, a3, b3);
333
+
334
+ __mmask32 nm70 = _mm512_mask_cmpneq_epi16_mask(nm60, a, b31);
335
+ __mmask32 nm71 = _mm512_mask_cmpneq_epi16_mask(nm61, a1, b31);
336
+ __mmask32 nm72 = _mm512_mask_cmpneq_epi16_mask(nm62, a2, b31);
337
+ __mmask32 nm73 = _mm512_mask_cmpneq_epi16_mask(nm63, a3, b31);
338
+
339
+ return ~(simsimd_u32_t)(nm70 & simsimd_u32_rol(nm71, 8) & simsimd_u32_rol(nm72, 16) & simsimd_u32_ror(nm73, 8));
340
+ }
341
+
342
+ /**
343
+ * @brief Analogous to `_mm512_2intersect_epi32`, but compatible with Ice Lake CPUs,
344
+ * slightly faster than the native Tiger Lake implementation, but returns only one mask.
345
+ *
346
+ * Some latencies to keep in mind:
347
+ *
348
+ * - `_mm512_shuffle_epi32` - "VPSHUFD (ZMM, ZMM, I8)":
349
+ * - 1 cycle latency on Ice Lake: 1*p5
350
+ * - 1 cycle latency on Genoa: 1*FP123
351
+ * - `_mm512_mask_cmpneq_epi32_mask` - "VPCMPD (K, ZMM, ZMM, I8)":
352
+ * - 3 cycle latency on Ice Lake: 1*p5
353
+ * - 1 cycle latency on Genoa: 1*FP01
354
+ * - `_mm512_alignr_epi32` - "VPALIGNR (ZMM, ZMM, ZMM, I8)":
355
+ * - 1 cycle latency on Ice Lake: 1*p5
356
+ * - 2 cycle latency on Genoa: 1*FP12
357
+ * - `_mm512_conflict_epi32` - "VPCONFLICTD (ZMM, ZMM)":
358
+ * - up to 26 cycles latency on Ice Lake: 11*p0+9*p05+17*p5
359
+ * - up to 7 cycle latency on Genoa: 1*FP01+1*FP12
360
+ */
361
+ SIMSIMD_INTERNAL simsimd_u16_t _simsimd_intersect_u32x16_ice(__m512i a, __m512i b) {
362
+ __m512i a1 = _mm512_alignr_epi32(a, a, 4);
363
+ __m512i b1 = _mm512_shuffle_epi32(b, _MM_PERM_ADCB);
364
+ __mmask16 nm00 = _mm512_cmpneq_epi32_mask(a, b);
365
+
366
+ __m512i a2 = _mm512_alignr_epi32(a, a, 8);
367
+ __m512i a3 = _mm512_alignr_epi32(a, a, 12);
368
+ __mmask16 nm01 = _mm512_cmpneq_epi32_mask(a1, b);
369
+ __mmask16 nm02 = _mm512_cmpneq_epi32_mask(a2, b);
370
+
371
+ __mmask16 nm03 = _mm512_cmpneq_epi32_mask(a3, b);
372
+ __mmask16 nm10 = _mm512_mask_cmpneq_epi32_mask(nm00, a, b1);
373
+ __mmask16 nm11 = _mm512_mask_cmpneq_epi32_mask(nm01, a1, b1);
374
+
375
+ __m512i b2 = _mm512_shuffle_epi32(b, _MM_PERM_BADC);
376
+ __mmask16 nm12 = _mm512_mask_cmpneq_epi32_mask(nm02, a2, b1);
377
+ __mmask16 nm13 = _mm512_mask_cmpneq_epi32_mask(nm03, a3, b1);
378
+ __mmask16 nm20 = _mm512_mask_cmpneq_epi32_mask(nm10, a, b2);
379
+
380
+ __m512i b3 = _mm512_shuffle_epi32(b, _MM_PERM_CBAD);
381
+ __mmask16 nm21 = _mm512_mask_cmpneq_epi32_mask(nm11, a1, b2);
382
+ __mmask16 nm22 = _mm512_mask_cmpneq_epi32_mask(nm12, a2, b2);
383
+ __mmask16 nm23 = _mm512_mask_cmpneq_epi32_mask(nm13, a3, b2);
384
+
385
+ __mmask16 nm0 = _mm512_mask_cmpneq_epi32_mask(nm20, a, b3);
386
+ __mmask16 nm1 = _mm512_mask_cmpneq_epi32_mask(nm21, a1, b3);
387
+ __mmask16 nm2 = _mm512_mask_cmpneq_epi32_mask(nm22, a2, b3);
388
+ __mmask16 nm3 = _mm512_mask_cmpneq_epi32_mask(nm23, a3, b3);
389
+
390
+ return ~(simsimd_u16_t)(nm0 & simsimd_u16_rol(nm1, 4) & simsimd_u16_rol(nm2, 8) & simsimd_u16_ror(nm3, 4));
391
+ }
392
+
393
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_ice( //
394
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
395
+ simsimd_size_t a_length, simsimd_size_t b_length, //
396
+ simsimd_distance_t *results) {
397
+
398
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
399
+ if (a_length < 64 && b_length < 64) {
400
+ simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
401
+ return;
402
+ }
403
+
404
+ simsimd_u16_t const *const a_end = a + a_length;
405
+ simsimd_u16_t const *const b_end = b + b_length;
406
+ simsimd_size_t c = 0;
407
+ union vec_t {
408
+ __m512i zmm;
409
+ simsimd_u16_t u16[32];
410
+ simsimd_u8_t u8[64];
411
+ } a_vec, b_vec;
412
+
413
+ while (a + 32 < a_end && b + 32 < b_end) {
414
+ a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
415
+ b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
416
+
417
+ // Intersecting registers with `_simsimd_intersect_u16x32_ice` involves a lot of shuffling
418
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
419
+ simsimd_u16_t a_min;
420
+ simsimd_u16_t a_max = a_vec.u16[31];
421
+ simsimd_u16_t b_min = b_vec.u16[0];
422
+ simsimd_u16_t b_max = b_vec.u16[31];
423
+
424
+ // If the slices don't overlap, advance the appropriate pointer
425
+ while (a_max < b_min && a + 64 < a_end) {
426
+ a += 32;
427
+ a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
428
+ a_max = a_vec.u16[31];
429
+ }
430
+ a_min = a_vec.u16[0];
431
+ while (b_max < a_min && b + 64 < b_end) {
432
+ b += 32;
433
+ b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
434
+ b_max = b_vec.u16[31];
435
+ }
436
+ b_min = b_vec.u16[0];
437
+
438
+ // Now we are likely to have some overlap, so we can intersect the registers
439
+ __mmask32 a_matches = _simsimd_intersect_u16x32_ice(a_vec.zmm, b_vec.zmm);
440
+
441
+ // The paper also contained a very nice procedure for exporting the matches,
442
+ // but we don't need it here:
443
+ // _mm512_mask_compressstoreu_epi16(c, a_matches, a_vec);
444
+ c += _mm_popcnt_u32(a_matches); // MSVC has no `_popcnt32`
445
+
446
+ __m512i a_last_broadcasted = _mm512_set1_epi16(*(short const *)&a_max);
447
+ __m512i b_last_broadcasted = _mm512_set1_epi16(*(short const *)&b_max);
448
+ __mmask32 a_step_mask = _mm512_cmple_epu16_mask(a_vec.zmm, b_last_broadcasted);
449
+ __mmask32 b_step_mask = _mm512_cmple_epu16_mask(b_vec.zmm, a_last_broadcasted);
450
+ a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask);
451
+ b += 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
452
+ }
453
+
454
+ simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
455
+ *results += c;
456
+ }
457
+
458
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_ice( //
459
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
460
+ simsimd_size_t a_length, simsimd_size_t b_length, //
461
+ simsimd_distance_t *results) {
462
+
463
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
464
+ if (a_length < 32 && b_length < 32) {
465
+ simsimd_intersect_u32_serial(a, b, a_length, b_length, results);
466
+ return;
467
+ }
468
+
469
+ simsimd_u32_t const *const a_end = a + a_length;
470
+ simsimd_u32_t const *const b_end = b + b_length;
471
+ simsimd_size_t c = 0;
472
+ union vec_t {
473
+ __m512i zmm;
474
+ simsimd_u32_t u32[16];
475
+ simsimd_u8_t u8[64];
476
+ } a_vec, b_vec;
477
+
478
+ while (a + 16 < a_end && b + 16 < b_end) {
479
+ a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
480
+ b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
481
+
482
+ // Intersecting registers with `_simsimd_intersect_u32x16_ice` involves a lot of shuffling
483
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
484
+ simsimd_u32_t a_min;
485
+ simsimd_u32_t a_max = a_vec.u32[15];
486
+ simsimd_u32_t b_min = b_vec.u32[0];
487
+ simsimd_u32_t b_max = b_vec.u32[15];
488
+
489
+ // If the slices don't overlap, advance the appropriate pointer
490
+ while (a_max < b_min && a + 32 < a_end) {
491
+ a += 16;
492
+ a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
493
+ a_max = a_vec.u32[15];
494
+ }
495
+ a_min = a_vec.u32[0];
496
+ while (b_max < a_min && b + 32 < b_end) {
497
+ b += 16;
498
+ b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
499
+ b_max = b_vec.u32[15];
500
+ }
501
+ b_min = b_vec.u32[0];
502
+
503
+ // Now we are likely to have some overlap, so we can intersect the registers
504
+ __mmask16 a_matches = _simsimd_intersect_u32x16_ice(a_vec.zmm, b_vec.zmm);
505
+
506
+ // The paper also contained a very nice procedure for exporting the matches,
507
+ // but we don't need it here:
508
+ // _mm512_mask_compressstoreu_epi32(c, a_matches, a_vec);
509
+ c += _mm_popcnt_u32(a_matches); // MSVC has no `_popcnt32`
510
+
511
+ __m512i a_last_broadcasted = _mm512_set1_epi32(*(int const *)&a_max);
512
+ __m512i b_last_broadcasted = _mm512_set1_epi32(*(int const *)&b_max);
513
+ __mmask16 a_step_mask = _mm512_cmple_epu32_mask(a_vec.zmm, b_last_broadcasted);
514
+ __mmask16 b_step_mask = _mm512_cmple_epu32_mask(b_vec.zmm, a_last_broadcasted);
515
+ a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask);
516
+ b += 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
517
+ }
518
+
519
+ simsimd_intersect_u32_serial(a, b, a_end - a, b_end - b, results);
520
+ *results += c;
521
+ }
522
+
523
+ #pragma clang attribute pop
524
+ #pragma GCC pop_options
525
+ #endif // SIMSIMD_TARGET_ICE
526
+
527
+ #if SIMSIMD_TARGET_TURIN
528
+ #pragma GCC push_options
529
+ #pragma GCC target("avx2", "avx512f", "avx512vl", "bmi2", "lzcnt", "popcnt", "avx512bw", "avx512vbmi2", "avx512bf16", \
530
+ "avx512vnni", "avx512vp2intersect")
531
+ #pragma clang attribute push( \
532
+ __attribute__((target( \
533
+ "avx2,avx512f,avx512vl,bmi2,lzcnt,popcnt,avx512bw,avx512vbmi2,avx512bf16,avx512vnni,avx512vp2intersect"))), \
534
+ apply_to = function)
535
+
536
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_turin( //
537
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
538
+ simsimd_size_t a_length, simsimd_size_t b_length, //
539
+ simsimd_distance_t *results) {
540
+
541
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
542
+ if (a_length < 64 && b_length < 64) {
543
+ simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
544
+ return;
545
+ }
546
+
547
+ //! There is no such thing as `_mm512_2intersect_epi16`, only the 32-bit variant!
548
+ //! So instead of jumping through 32 entries at a time, like on Ice Lake, we will
549
+ //! step through 16 entries at a time.
550
+ simsimd_u16_t const *const a_end = a + a_length;
551
+ simsimd_u16_t const *const b_end = b + b_length;
552
+ simsimd_size_t c = 0;
553
+ union vec_t {
554
+ __m256i ymm;
555
+ simsimd_u16_t u16[16];
556
+ simsimd_u8_t u8[32];
557
+ } a_vec, b_vec;
558
+
559
+ while (a + 16 < a_end && b + 16 < b_end) {
560
+ a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
561
+ b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
562
+
563
+ // Intersecting registers with `_mm512_2intersect_epi16_mask` involves a lot of shuffling
564
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
565
+ simsimd_u16_t a_min;
566
+ simsimd_u16_t a_max = a_vec.u16[15];
567
+ simsimd_u16_t b_min = b_vec.u16[0];
568
+ simsimd_u16_t b_max = b_vec.u16[15];
569
+
570
+ // If the slices don't overlap, advance the appropriate pointer
571
+ while (a_max < b_min && a + 32 < a_end) {
572
+ a += 16;
573
+ a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
574
+ a_max = a_vec.u16[15];
575
+ }
576
+ a_min = a_vec.u16[0];
577
+ while (b_max < a_min && b + 32 < b_end) {
578
+ b += 16;
579
+ b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
580
+ b_max = b_vec.u16[15];
581
+ }
582
+ b_min = b_vec.u16[0];
583
+
584
+ // Now we are likely to have some overlap, so we can intersect the registers
585
+ __m512i a_i32_vec = _mm512_cvtepu16_epi32(a_vec.ymm);
586
+ __m512i b_i32_vec = _mm512_cvtepu16_epi32(b_vec.ymm);
587
+ __mmask16 a_matches_any_in_b, b_matches_any_in_a;
588
+ _mm512_2intersect_epi32(a_i32_vec, b_i32_vec, &a_matches_any_in_b, &b_matches_any_in_a);
589
+
590
+ // The paper also contained a very nice procedure for exporting the matches,
591
+ // but we don't need it here:
592
+ // _mm512_mask_compressstoreu_epi16(c, a_matches_any_in_b, a_vec);
593
+ c += _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
594
+
595
+ __m256i a_last_broadcasted = _mm256_set1_epi16(*(short const *)&a_max);
596
+ __m256i b_last_broadcasted = _mm256_set1_epi16(*(short const *)&b_max);
597
+ __mmask16 a_step_mask = _mm256_cmple_epu16_mask(a_vec.ymm, b_last_broadcasted);
598
+ __mmask16 b_step_mask = _mm256_cmple_epu16_mask(b_vec.ymm, a_last_broadcasted);
599
+ a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask); //? Is this correct? Needs testing!
600
+ b += 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
601
+ }
602
+
603
+ simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
604
+ *results += c;
605
+ }
606
+
607
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_turin( //
608
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
609
+ simsimd_size_t a_length, simsimd_size_t b_length, //
610
+ simsimd_distance_t *results) {
611
+
612
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
613
+ if (a_length < 32 && b_length < 32) {
614
+ simsimd_intersect_u32_serial(a, b, a_length, b_length, results);
615
+ return;
616
+ }
617
+
618
+ simsimd_u32_t const *const a_end = a + a_length;
619
+ simsimd_u32_t const *const b_end = b + b_length;
620
+ simsimd_size_t c = 0;
621
+ union vec_t {
622
+ __m512i zmm;
623
+ simsimd_u32_t u32[16];
624
+ simsimd_u8_t u8[64];
625
+ } a_vec, b_vec;
626
+
627
+ while (a + 16 < a_end && b + 16 < b_end) {
628
+ a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
629
+ b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
630
+
631
+ // Intersecting registers with `_mm512_2intersect_epi32` involves a lot of shuffling
632
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
633
+ simsimd_u32_t a_min;
634
+ simsimd_u32_t a_max = a_vec.u32[15];
635
+ simsimd_u32_t b_min = b_vec.u32[0];
636
+ simsimd_u32_t b_max = b_vec.u32[15];
637
+
638
+ // If the slices don't overlap, advance the appropriate pointer
639
+ while (a_max < b_min && a + 32 < a_end) {
640
+ a += 16;
641
+ a_vec.zmm = _mm512_loadu_si512((__m512i const *)a);
642
+ a_max = a_vec.u32[15];
643
+ }
644
+ a_min = a_vec.u32[0];
645
+ while (b_max < a_min && b + 32 < b_end) {
646
+ b += 16;
647
+ b_vec.zmm = _mm512_loadu_si512((__m512i const *)b);
648
+ b_max = b_vec.u32[15];
649
+ }
650
+ b_min = b_vec.u32[0];
651
+
652
+ // Now we are likely to have some overlap, so we can intersect the registers
653
+ __mmask16 a_matches_any_in_b, b_matches_any_in_a;
654
+ _mm512_2intersect_epi32(a_vec.zmm, b_vec.zmm, &a_matches_any_in_b, &b_matches_any_in_a);
655
+
656
+ // The paper also contained a very nice procedure for exporting the matches,
657
+ // but we don't need it here:
658
+ // _mm512_mask_compressstoreu_epi32(c, a_matches_any_in_b, a_vec);
659
+ c += _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
660
+
661
+ __m512i a_last_broadcasted = _mm512_set1_epi32(*(int const *)&a_max);
662
+ __m512i b_last_broadcasted = _mm512_set1_epi32(*(int const *)&b_max);
663
+ __mmask16 a_step_mask = _mm512_cmple_epu32_mask(a_vec.zmm, b_last_broadcasted);
664
+ __mmask16 b_step_mask = _mm512_cmple_epu32_mask(b_vec.zmm, a_last_broadcasted);
665
+ a += 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask);
666
+ b += 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
667
+ }
668
+
669
+ simsimd_intersect_u32_serial(a, b, a_end - a, b_end - b, results);
670
+ *results += c;
671
+ }
672
+
673
+ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_turin( //
674
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
675
+ simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
676
+ simsimd_size_t a_length, simsimd_size_t b_length, //
677
+ simsimd_distance_t *results) {
678
+
679
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
680
+ if (a_length < 64 && b_length < 64) {
681
+ simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
682
+ return;
683
+ }
684
+
685
+ //! There is no such thing as `_mm512_2intersect_epi16`, only the 32-bit variant!
686
+ //! So instead of jumping through 32 entries at a time, like on Ice Lake, we will
687
+ //! step through 16 entries at a time.
688
+ simsimd_u16_t const *const a_end = a + a_length;
689
+ simsimd_u16_t const *const b_end = b + b_length;
690
+ simsimd_size_t intersection_size = 0;
691
+ union vec_t {
692
+ __m256i ymm;
693
+ __m256 ymmps;
694
+ simsimd_u16_t u16[16];
695
+ simsimd_u8_t u8[32];
696
+ } a_vec, b_vec, product_vec;
697
+ product_vec.ymmps = _mm256_setzero_ps();
698
+
699
+ while (a + 16 < a_end && b + 16 < b_end) {
700
+ a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
701
+ b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
702
+
703
+ // Intersecting registers with `_mm512_2intersect_epi16_mask` involves a lot of shuffling
704
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
705
+ simsimd_u16_t a_min;
706
+ simsimd_u16_t a_max = a_vec.u16[15];
707
+ simsimd_u16_t b_min = b_vec.u16[0];
708
+ simsimd_u16_t b_max = b_vec.u16[15];
709
+
710
+ // If the slices don't overlap, advance the appropriate pointer
711
+ while (a_max < b_min && a + 32 < a_end) {
712
+ a += 16, a_weights += 16;
713
+ a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
714
+ a_max = a_vec.u16[15];
715
+ }
716
+ a_min = a_vec.u16[0];
717
+ while (b_max < a_min && b + 32 < b_end) {
718
+ b += 16, b_weights += 16;
719
+ b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
720
+ b_max = b_vec.u16[15];
721
+ }
722
+ b_min = b_vec.u16[0];
723
+
724
+ // Now we are likely to have some overlap, so we can intersect the registers
725
+ __m512i a_i32_vec = _mm512_cvtepu16_epi32(a_vec.ymm);
726
+ __m512i b_i32_vec = _mm512_cvtepu16_epi32(b_vec.ymm);
727
+ __mmask16 a_matches_any_in_b, b_matches_any_in_a;
728
+ _mm512_2intersect_epi32(a_i32_vec, b_i32_vec, &a_matches_any_in_b, &b_matches_any_in_a);
729
+
730
+ // The paper also contained a very nice procedure for exporting the matches,
731
+ // but we don't need it here:
732
+ // _mm512_mask_compressstoreu_epi16(intersection_size, a_matches_any_in_b, a_vec);
733
+ int a_matches_count_in_b = _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
734
+ intersection_size += a_matches_count_in_b;
735
+
736
+ // Load and shift all the relevant weights to the start of the vector before doing the dot product
737
+ if (a_matches_count_in_b) {
738
+ __m256i a_weights_vec = _mm256_lddqu_si256((__m256i const *)a_weights);
739
+ a_weights_vec = _mm256_maskz_compress_epi16(a_matches_any_in_b, a_weights_vec);
740
+ __m256i b_weights_vec = _mm256_lddqu_si256((__m256i const *)b_weights);
741
+ b_weights_vec = _mm256_maskz_compress_epi16(b_matches_any_in_a, b_weights_vec);
742
+ product_vec.ymmps = _mm256_dpbf16_ps(product_vec.ymmps, (__m256bh)a_weights_vec, (__m256bh)b_weights_vec);
743
+ }
744
+
745
+ __m256i a_last_broadcasted = _mm256_set1_epi16(*(short const *)&a_max);
746
+ __m256i b_last_broadcasted = _mm256_set1_epi16(*(short const *)&b_max);
747
+ __mmask16 a_step_mask = _mm256_cmple_epu16_mask(a_vec.ymm, b_last_broadcasted);
748
+ __mmask16 b_step_mask = _mm256_cmple_epu16_mask(b_vec.ymm, a_last_broadcasted);
749
+ int a_step = 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask); //? Is this correct? Needs testing!
750
+ int b_step = 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
751
+ a += a_step, a_weights += a_step;
752
+ b += b_step, b_weights += b_step;
753
+ }
754
+
755
+ simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
756
+ *results += intersection_size;
757
+ }
758
+
759
+ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_turin( //
760
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
761
+ simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
762
+ simsimd_size_t a_length, simsimd_size_t b_length, //
763
+ simsimd_distance_t *results) {
764
+
765
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
766
+ if (a_length < 64 && b_length < 64) {
767
+ simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
768
+ return;
769
+ }
770
+
771
+ //! There is no such thing as `_mm512_2intersect_epi16`, only the 32-bit variant!
772
+ //! So instead of jumping through 32 entries at a time, like on Ice Lake, we will
773
+ //! step through 16 entries at a time.
774
+ simsimd_u16_t const *const a_end = a + a_length;
775
+ simsimd_u16_t const *const b_end = b + b_length;
776
+ simsimd_size_t intersection_size = 0;
777
+ union vec_t {
778
+ __m256i ymm;
779
+ simsimd_u16_t u16[16];
780
+ simsimd_u8_t u8[32];
781
+ } a_vec, b_vec, product_vec;
782
+ product_vec.ymm = _mm256_setzero_si256();
783
+
784
+ while (a + 16 < a_end && b + 16 < b_end) {
785
+ a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
786
+ b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
787
+
788
+ // Intersecting registers with `_mm512_2intersect_epi16_mask` involves a lot of shuffling
789
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
790
+ simsimd_u16_t a_min;
791
+ simsimd_u16_t a_max = a_vec.u16[15];
792
+ simsimd_u16_t b_min = b_vec.u16[0];
793
+ simsimd_u16_t b_max = b_vec.u16[15];
794
+
795
+ // If the slices don't overlap, advance the appropriate pointer
796
+ while (a_max < b_min && a + 32 < a_end) {
797
+ a += 16, a_weights += 16;
798
+ a_vec.ymm = _mm256_lddqu_si256((__m256i const *)a);
799
+ a_max = a_vec.u16[15];
800
+ }
801
+ a_min = a_vec.u16[0];
802
+ while (b_max < a_min && b + 32 < b_end) {
803
+ b += 16, b_weights += 16;
804
+ b_vec.ymm = _mm256_lddqu_si256((__m256i const *)b);
805
+ b_max = b_vec.u16[15];
806
+ }
807
+ b_min = b_vec.u16[0];
808
+
809
+ // Now we are likely to have some overlap, so we can intersect the registers
810
+ __m512i a_i32_vec = _mm512_cvtepu16_epi32(a_vec.ymm);
811
+ __m512i b_i32_vec = _mm512_cvtepu16_epi32(b_vec.ymm);
812
+ __mmask16 a_matches_any_in_b, b_matches_any_in_a;
813
+ _mm512_2intersect_epi32(a_i32_vec, b_i32_vec, &a_matches_any_in_b, &b_matches_any_in_a);
814
+
815
+ // The paper also contained a very nice procedure for exporting the matches,
816
+ // but we don't need it here:
817
+ // _mm512_mask_compressstoreu_epi16(intersection_size, a_matches_any_in_b, a_vec);
818
+ int a_matches_count_in_b = _mm_popcnt_u32(a_matches_any_in_b); // MSVC has no `_popcnt32`
819
+ intersection_size += a_matches_count_in_b;
820
+
821
+ // Load and shift all the relevant weights to the start of the vector before doing the dot product
822
+ if (a_matches_count_in_b) {
823
+ __m256i a_weights_vec = _mm256_lddqu_si256((__m256i const *)a_weights);
824
+ a_weights_vec = _mm256_maskz_compress_epi16(a_matches_any_in_b, a_weights_vec);
825
+ __m256i b_weights_vec = _mm256_lddqu_si256((__m256i const *)b_weights);
826
+ b_weights_vec = _mm256_maskz_compress_epi16(b_matches_any_in_a, b_weights_vec);
827
+ product_vec.ymm = _mm256_dpwssds_epi32(product_vec.ymm, a_weights_vec, b_weights_vec);
828
+ }
829
+
830
+ __m256i a_last_broadcasted = _mm256_set1_epi16(*(short const *)&a_max);
831
+ __m256i b_last_broadcasted = _mm256_set1_epi16(*(short const *)&b_max);
832
+ __mmask16 a_step_mask = _mm256_cmple_epu16_mask(a_vec.ymm, b_last_broadcasted);
833
+ __mmask16 b_step_mask = _mm256_cmple_epu16_mask(b_vec.ymm, a_last_broadcasted);
834
+ int a_step = 32 - _lzcnt_u32((simsimd_u32_t)a_step_mask); //? Is this correct? Needs testing!
835
+ int b_step = 32 - _lzcnt_u32((simsimd_u32_t)b_step_mask);
836
+ a += a_step, a_weights += a_step;
837
+ b += b_step, b_weights += b_step;
838
+ }
839
+
840
+ simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
841
+ *results += intersection_size;
842
+ }
843
+
844
+ #pragma clang attribute pop
845
+ #pragma GCC pop_options
846
+ #endif // SIMSIMD_TARGET_TURIN
847
+ #endif // _SIMSIMD_TARGET_X86
848
+
849
+ #if _SIMSIMD_TARGET_ARM
850
+ #if SIMSIMD_TARGET_NEON
851
+ #pragma GCC push_options
852
+ #pragma GCC target("arch=armv8.2-a")
853
+ #pragma clang attribute push(__attribute__((target("arch=armv8.2-a"))), apply_to = function)
854
+
855
+ /**
856
+ * @brief Uses `vshrn` to produce a bitmask, similar to `movemask` in SSE.
857
+ * https://community.arm.com/arm-community-blogs/b/infrastructure-solutions-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon
858
+ */
859
+ SIMSIMD_INTERNAL simsimd_u64_t _simsimd_u8_to_u4_neon(uint8x16_t vec) {
860
+ return vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(vec), 4)), 0);
861
+ }
862
+
863
+ SIMSIMD_INTERNAL int _simsimd_clz_u64(simsimd_u64_t x) {
864
+ // On GCC and Clang use the builtin, otherwise use the generic implementation
865
+ #if defined(__GNUC__) || defined(__clang__)
866
+ return __builtin_clzll(x);
867
+ #else
868
+ int n = 0;
869
+ while ((x & 0x8000000000000000ull) == 0) n++, x <<= 1;
870
+ return n;
871
+ #endif
872
+ }
873
+
874
+ SIMSIMD_INTERNAL uint32x4_t _simsimd_intersect_u32x4_neon(uint32x4_t a, uint32x4_t b) {
875
+ uint32x4_t b1 = vextq_u32(b, b, 1);
876
+ uint32x4_t b2 = vextq_u32(b, b, 2);
877
+ uint32x4_t b3 = vextq_u32(b, b, 3);
878
+ uint32x4_t nm00 = vceqq_u32(a, b);
879
+ uint32x4_t nm01 = vceqq_u32(a, b1);
880
+ uint32x4_t nm02 = vceqq_u32(a, b2);
881
+ uint32x4_t nm03 = vceqq_u32(a, b3);
882
+ uint32x4_t nm = vorrq_u32(vorrq_u32(nm00, nm01), vorrq_u32(nm02, nm03));
883
+ return nm;
884
+ }
885
+
886
+ SIMSIMD_INTERNAL uint16x8_t _simsimd_intersect_u16x8_neon(uint16x8_t a, uint16x8_t b) {
887
+ uint16x8_t b1 = vextq_u16(b, b, 1);
888
+ uint16x8_t b2 = vextq_u16(b, b, 2);
889
+ uint16x8_t b3 = vextq_u16(b, b, 3);
890
+ uint16x8_t b4 = vextq_u16(b, b, 4);
891
+ uint16x8_t b5 = vextq_u16(b, b, 5);
892
+ uint16x8_t b6 = vextq_u16(b, b, 6);
893
+ uint16x8_t b7 = vextq_u16(b, b, 7);
894
+ uint16x8_t nm00 = vceqq_u16(a, b);
895
+ uint16x8_t nm01 = vceqq_u16(a, b1);
896
+ uint16x8_t nm02 = vceqq_u16(a, b2);
897
+ uint16x8_t nm03 = vceqq_u16(a, b3);
898
+ uint16x8_t nm04 = vceqq_u16(a, b4);
899
+ uint16x8_t nm05 = vceqq_u16(a, b5);
900
+ uint16x8_t nm06 = vceqq_u16(a, b6);
901
+ uint16x8_t nm07 = vceqq_u16(a, b7);
902
+ uint16x8_t nm = vorrq_u16(vorrq_u16(vorrq_u16(nm00, nm01), vorrq_u16(nm02, nm03)),
903
+ vorrq_u16(vorrq_u16(nm04, nm05), vorrq_u16(nm06, nm07)));
904
+ return nm;
905
+ }
906
+
907
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_neon( //
908
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
909
+ simsimd_size_t a_length, simsimd_size_t b_length, //
910
+ simsimd_distance_t *results) {
911
+
912
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
913
+ if (a_length < 32 && b_length < 32) {
914
+ simsimd_intersect_u16_serial(a, b, a_length, b_length, results);
915
+ return;
916
+ }
917
+
918
+ simsimd_u16_t const *const a_end = a + a_length;
919
+ simsimd_u16_t const *const b_end = b + b_length;
920
+ union vec_t {
921
+ uint16x8_t u16x8;
922
+ simsimd_u16_t u16[8];
923
+ simsimd_u8_t u8[16];
924
+ } a_vec, b_vec, c_counts_vec;
925
+ c_counts_vec.u16x8 = vdupq_n_u16(0);
926
+
927
+ while (a + 8 < a_end && b + 8 < b_end) {
928
+ a_vec.u16x8 = vld1q_u16(a);
929
+ b_vec.u16x8 = vld1q_u16(b);
930
+
931
+ // Intersecting registers with `_simsimd_intersect_u16x8_neon` involves a lot of shuffling
932
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
933
+ simsimd_u16_t a_min;
934
+ simsimd_u16_t a_max = a_vec.u16[7];
935
+ simsimd_u16_t b_min = b_vec.u16[0];
936
+ simsimd_u16_t b_max = b_vec.u16[7];
937
+
938
+ // If the slices don't overlap, advance the appropriate pointer
939
+ while (a_max < b_min && a + 16 < a_end) {
940
+ a += 8;
941
+ a_vec.u16x8 = vld1q_u16(a);
942
+ a_max = a_vec.u16[7];
943
+ }
944
+ a_min = a_vec.u16[0];
945
+ while (b_max < a_min && b + 16 < b_end) {
946
+ b += 8;
947
+ b_vec.u16x8 = vld1q_u16(b);
948
+ b_max = b_vec.u16[7];
949
+ }
950
+ b_min = b_vec.u16[0];
951
+
952
+ // Now we are likely to have some overlap, so we can intersect the registers.
953
+ // We can do it by performing a population count at every cycle, but it's not the cheapest in terms of cycles.
954
+ //
955
+ // simsimd_u64_t a_matches = __builtin_popcountll(
956
+ // _simsimd_u8_to_u4_neon(vreinterpretq_u8_u16(
957
+ // _simsimd_intersect_u16x8_neon(a_vec.u16x8, b_vec.u16x8))));
958
+ // c += a_matches / 8;
959
+ //
960
+ // Alternatively, we can we can transform match-masks into "ones", accumulate them between the cycles,
961
+ // and merge all together in the end.
962
+ uint16x8_t a_matches = _simsimd_intersect_u16x8_neon(a_vec.u16x8, b_vec.u16x8);
963
+ c_counts_vec.u16x8 = vaddq_u16(c_counts_vec.u16x8, vandq_u16(a_matches, vdupq_n_u16(1)));
964
+
965
+ // Counting leading zeros is tricky. On Arm we can use inline Assembly to get the result,
966
+ // but MSVC doesn't support that:
967
+ //
968
+ // SIMSIMD_INTERNAL int _simsimd_clz_u64(simsimd_u64_t value) {
969
+ // simsimd_u64_t result;
970
+ // __asm__("clz %x0, %x1" : "=r"(result) : "r"(value));
971
+ // return (int)result;
972
+ // }
973
+ //
974
+ // Alternatively, we can use the `vclz_u32` NEON intrinsic.
975
+ // It will compute the leading zeros number for both `a_step` and `b_step` in parallel.
976
+ uint16x8_t a_last_broadcasted = vdupq_n_u16(a_max);
977
+ uint16x8_t b_last_broadcasted = vdupq_n_u16(b_max);
978
+ simsimd_u64_t a_step = _simsimd_clz_u64(_simsimd_u8_to_u4_neon( //
979
+ vreinterpretq_u8_u16(vcleq_u16(a_vec.u16x8, b_last_broadcasted))));
980
+ simsimd_u64_t b_step = _simsimd_clz_u64(_simsimd_u8_to_u4_neon( //
981
+ vreinterpretq_u8_u16(vcleq_u16(b_vec.u16x8, a_last_broadcasted))));
982
+ a += (64 - a_step) / 8;
983
+ b += (64 - b_step) / 8;
984
+ }
985
+
986
+ simsimd_intersect_u16_serial(a, b, a_end - a, b_end - b, results);
987
+ *results += vaddvq_u16(c_counts_vec.u16x8);
988
+ }
989
+
990
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_neon( //
991
+ simsimd_u32_t const *a, simsimd_u32_t const *b, //
992
+ simsimd_size_t a_length, simsimd_size_t b_length, //
993
+ simsimd_distance_t *results) {
994
+
995
+ // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
996
+ if (a_length < 32 && b_length < 32) {
997
+ simsimd_intersect_u32_serial(a, b, a_length, b_length, results);
998
+ return;
999
+ }
1000
+
1001
+ simsimd_u32_t const *const a_end = a + a_length;
1002
+ simsimd_u32_t const *const b_end = b + b_length;
1003
+ union vec_t {
1004
+ uint32x4_t u32x4;
1005
+ simsimd_u32_t u32[4];
1006
+ simsimd_u8_t u8[16];
1007
+ } a_vec, b_vec, c_counts_vec;
1008
+ c_counts_vec.u32x4 = vdupq_n_u32(0);
1009
+
1010
+ while (a + 4 < a_end && b + 4 < b_end) {
1011
+ a_vec.u32x4 = vld1q_u32(a);
1012
+ b_vec.u32x4 = vld1q_u32(b);
1013
+
1014
+ // Intersecting registers with `_simsimd_intersect_u32x4_neon` involves a lot of shuffling
1015
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
1016
+ simsimd_u32_t a_min;
1017
+ simsimd_u32_t a_max = a_vec.u32[3];
1018
+ simsimd_u32_t b_min = b_vec.u32[0];
1019
+ simsimd_u32_t b_max = b_vec.u32[3];
1020
+
1021
+ // If the slices don't overlap, advance the appropriate pointer
1022
+ while (a_max < b_min && a + 8 < a_end) {
1023
+ a += 4;
1024
+ a_vec.u32x4 = vld1q_u32(a);
1025
+ a_max = a_vec.u32[3];
1026
+ }
1027
+ a_min = a_vec.u32[0];
1028
+ while (b_max < a_min && b + 8 < b_end) {
1029
+ b += 4;
1030
+ b_vec.u32x4 = vld1q_u32(b);
1031
+ b_max = b_vec.u32[3];
1032
+ }
1033
+ b_min = b_vec.u32[0];
1034
+
1035
+ // Now we are likely to have some overlap, so we can intersect the registers
1036
+ // We can do it by performing a population count at every cycle, but it's not the cheapest in terms of cycles.
1037
+ //
1038
+ // simsimd_u64_t a_matches = __builtin_popcountll(
1039
+ // _simsimd_u8_to_u4_neon(vreinterpretq_u8_u32(
1040
+ // _simsimd_intersect_u32x4_neon(a_vec.u32x4, b_vec.u32x4))));
1041
+ // c += a_matches / 16;
1042
+ //
1043
+ // Alternatively, we can we can transform match-masks into "ones", accumulate them between the cycles,
1044
+ // and merge all together in the end.
1045
+ uint32x4_t a_matches = _simsimd_intersect_u32x4_neon(a_vec.u32x4, b_vec.u32x4);
1046
+ c_counts_vec.u32x4 = vaddq_u32(c_counts_vec.u32x4, vandq_u32(a_matches, vdupq_n_u32(1)));
1047
+
1048
+ uint32x4_t a_last_broadcasted = vdupq_n_u32(a_max);
1049
+ uint32x4_t b_last_broadcasted = vdupq_n_u32(b_max);
1050
+ simsimd_u64_t a_step = _simsimd_clz_u64(_simsimd_u8_to_u4_neon( //
1051
+ vreinterpretq_u8_u32(vcleq_u32(a_vec.u32x4, b_last_broadcasted))));
1052
+ simsimd_u64_t b_step = _simsimd_clz_u64(_simsimd_u8_to_u4_neon( //
1053
+ vreinterpretq_u8_u32(vcleq_u32(b_vec.u32x4, a_last_broadcasted))));
1054
+ a += (64 - a_step) / 16;
1055
+ b += (64 - b_step) / 16;
1056
+ }
1057
+
1058
+ simsimd_intersect_u32_serial(a, b, a_end - a, b_end - b, results);
1059
+ *results += vaddvq_u32(c_counts_vec.u32x4);
1060
+ }
1061
+
1062
+ #pragma clang attribute pop
1063
+ #pragma GCC pop_options
1064
+ #endif // SIMSIMD_TARGET_NEON
1065
+
1066
+ /* SVE2 introduces many new integer-oriented instructions, extending some of the NEON functionality
1067
+ * to variable-length SVE registers. Those include "compare multiple" intrinsics:
1068
+ *
1069
+ * - `svmatch[_u16]` that matches each scalar in first vector against all members of a 128-bit lane in the second.
1070
+ * - `svhistcnt[_s32]_z` does something similar, performing an inclusive prefix scan.
1071
+ * - `svtbx[_u16]` does extended table lookup
1072
+ *
1073
+ * Other notable instructions:
1074
+ *
1075
+ * - `DUP`: Broadcast indexed predicate element
1076
+ * https://developer.arm.com/documentation/ddi0602/2021-06/SVE-Instructions/DUP--predicate---Broadcast-indexed-predicate-element-?lang=en
1077
+ * - `SCLAMP` and `UCLAMP`: clamp values, i.e. combined min+max
1078
+ * https://developer.arm.com/documentation/ddi0602/2021-06/SVE-Instructions/SCLAMP--Signed-clamp-to-minimum-maximum-vector-?lang=en
1079
+ * https://developer.arm.com/documentation/ddi0602/2021-06/SVE-Instructions/UCLAMP--Unsigned-clamp-to-minimum-maximum-vector-?lang=en
1080
+ * - `TBLQ`: Table lookup quadword
1081
+ * https://developer.arm.com/documentation/ddi0602/2022-12/SVE-Instructions/TBLQ--Programmable-table-lookup-within-each-quadword-vector-segment--zeroing--?lang=en
1082
+ *
1083
+ * Great resources for SVE2 intrinsics:
1084
+ *
1085
+ * > ARM’s Scalable Vector Extensions: A Critical Look at SVE2 For Integer Workloads
1086
+ * https://gist.github.com/zingaburga/805669eb891c820bd220418ee3f0d6bd
1087
+ */
1088
+ #if SIMSIMD_TARGET_SVE2
1089
+ #pragma GCC push_options
1090
+ #pragma GCC target("arch=armv8.2-a+sve+sve2")
1091
+ #pragma clang attribute push(__attribute__((target("arch=armv8.2-a+sve+sve2"))), apply_to = function)
1092
+
1093
+ SIMSIMD_PUBLIC void simsimd_intersect_u16_sve2( //
1094
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
1095
+ simsimd_size_t a_length,
1096
+ simsimd_size_t b_length, //
1097
+ simsimd_distance_t *results) {
1098
+
1099
+ // A single SVE lane is 128 bits wide, so one lane fits 8 values.
1100
+ simsimd_size_t const register_size = svcnth();
1101
+ simsimd_size_t const lanes_count = register_size / 8;
1102
+ simsimd_size_t a_idx = 0, b_idx = 0;
1103
+ simsimd_size_t c = 0;
1104
+
1105
+ while (a_idx < a_length && b_idx < b_length) {
1106
+ // Load `a_member` and broadcast it, load `b_members_vec` from memory
1107
+ svbool_t a_progress = svwhilelt_b16_u64(a_idx, a_length);
1108
+ svbool_t b_progress = svwhilelt_b16_u64(b_idx, b_length);
1109
+ svuint16_t a_vec = svld1_u16(a_progress, a + a_idx);
1110
+ svuint16_t b_vec = svld1_u16(b_progress, b + b_idx);
1111
+
1112
+ // Intersecting registers with `svmatch_u16` involves a lot of shuffling
1113
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
1114
+ simsimd_u16_t a_min;
1115
+ simsimd_u16_t a_max = svlastb(a_progress, a_vec);
1116
+ simsimd_u16_t b_min = svlasta(svpfalse_b(), b_vec);
1117
+ simsimd_u16_t b_max = svlastb(b_progress, b_vec);
1118
+
1119
+ // If the slices don't overlap, advance the appropriate pointer
1120
+ while (a_max < b_min && (a_idx + register_size) < a_length) {
1121
+ a_idx += register_size;
1122
+ a_progress = svwhilelt_b16_u64(a_idx, a_length);
1123
+ a_vec = svld1_u16(a_progress, a + a_idx);
1124
+ a_max = svlastb(a_progress, a_vec);
1125
+ }
1126
+ a_min = svlasta(svpfalse_b(), a_vec);
1127
+ while (b_max < a_min && (b_idx + register_size) < b_length) {
1128
+ b_idx += register_size;
1129
+ b_progress = svwhilelt_b16_u64(b_idx, b_length);
1130
+ b_vec = svld1_u16(b_progress, b + b_idx);
1131
+ b_max = svlastb(b_progress, b_vec);
1132
+ }
1133
+ b_min = svlasta(svpfalse_b(), b_vec);
1134
+
1135
+ // Before we evaluate the intersection size, obfurscating the order in `b_vec`,
1136
+ // let's estimate how much we will need to advance the pointers afterwards.
1137
+ // For that, we don't even need to broadcast the values in SVE, as the whole
1138
+ // register can be compared against a scalar:
1139
+ //
1140
+ // svuint16_t a_last_broadcasted = svdup_n_u16(a_max);
1141
+ // svuint16_t b_last_broadcasted = svdup_n_u16(b_max);
1142
+ svbool_t a_mask = svcmple_n_u16(a_progress, a_vec, b_max);
1143
+ svbool_t b_mask = svcmple_n_u16(b_progress, b_vec, a_max);
1144
+ simsimd_u64_t a_step = svcntp_b16(a_progress, a_mask);
1145
+ simsimd_u64_t b_step = svcntp_b16(b_progress, b_mask);
1146
+
1147
+ // Compare `a_vec` with each lane of `b_vec`
1148
+ svbool_t equal_mask = svmatch_u16(a_progress, a_vec, b_vec);
1149
+ for (simsimd_size_t i = 1; i < lanes_count; i++) {
1150
+ b_vec = svext_u16(b_vec, b_vec, 8);
1151
+ equal_mask = svorr_z(svptrue_b16(), equal_mask, svmatch_u16(a_progress, a_vec, b_vec));
1152
+ }
1153
+ simsimd_size_t equal_count = svcntp_b16(svptrue_b16(), equal_mask);
1154
+
1155
+ // Advance
1156
+ a_idx += a_step;
1157
+ b_idx += b_step;
1158
+ c += equal_count;
1159
+ }
1160
+ *results = c;
1161
+ }
1162
+
1163
+ SIMSIMD_PUBLIC void simsimd_intersect_u32_sve2(simsimd_u32_t const *a, simsimd_u32_t const *b, simsimd_size_t a_length,
1164
+ simsimd_size_t b_length, simsimd_distance_t *results) {
1165
+
1166
+ // A single SVE lane is 128 bits wide, so one lane fits 4 values.
1167
+ simsimd_size_t const register_size = svcntw();
1168
+ simsimd_size_t const lanes_count = register_size / 4;
1169
+ simsimd_size_t a_idx = 0, b_idx = 0;
1170
+ simsimd_size_t c = 0;
1171
+
1172
+ while (a_idx < a_length && b_idx < b_length) {
1173
+ // Load `a_member` and broadcast it, load `b_members_vec` from memory
1174
+ svbool_t a_progress = svwhilelt_b32_u64(a_idx, a_length);
1175
+ svbool_t b_progress = svwhilelt_b32_u64(b_idx, b_length);
1176
+ svuint32_t a_vec = svld1_u32(a_progress, a + a_idx);
1177
+ svuint32_t b_vec = svld1_u32(b_progress, b + b_idx);
1178
+
1179
+ // Intersecting registers with `svmatch_u16` involves a lot of shuffling
1180
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
1181
+ simsimd_u32_t a_min;
1182
+ simsimd_u32_t a_max = svlastb(a_progress, a_vec);
1183
+ simsimd_u32_t b_min = svlasta(svpfalse_b(), b_vec);
1184
+ simsimd_u32_t b_max = svlastb(b_progress, b_vec);
1185
+
1186
+ // If the slices don't overlap, advance the appropriate pointer
1187
+ while (a_max < b_min && (a_idx + register_size) < a_length) {
1188
+ a_idx += register_size;
1189
+ a_progress = svwhilelt_b32_u64(a_idx, a_length);
1190
+ a_vec = svld1_u32(a_progress, a + a_idx);
1191
+ a_max = svlastb(a_progress, a_vec);
1192
+ }
1193
+ a_min = svlasta(svpfalse_b(), a_vec);
1194
+ while (b_max < a_min && (b_idx + register_size) < b_length) {
1195
+ b_idx += register_size;
1196
+ b_progress = svwhilelt_b32_u64(b_idx, b_length);
1197
+ b_vec = svld1_u32(b_progress, b + b_idx);
1198
+ b_max = svlastb(b_progress, b_vec);
1199
+ }
1200
+ b_min = svlasta(svpfalse_b(), b_vec);
1201
+
1202
+ // Before we evaluate the intersection size, obfurscating the order in `b_vec`,
1203
+ // let's estimate how much we will need to advance the pointers afterwards.
1204
+ // For that, we don't even need to broadcast the values in SVE, as the whole
1205
+ // register can be compared against a scalar:
1206
+ //
1207
+ // svuint32_t a_last_broadcasted = svdup_n_u32(a_max);
1208
+ // svuint32_t b_last_broadcasted = svdup_n_u32(b_max);
1209
+ svbool_t a_mask = svcmple_n_u32(a_progress, a_vec, b_max);
1210
+ svbool_t b_mask = svcmple_n_u32(b_progress, b_vec, a_max);
1211
+ simsimd_u64_t a_step = svcntp_b32(a_progress, a_mask);
1212
+ simsimd_u64_t b_step = svcntp_b32(b_progress, b_mask);
1213
+
1214
+ // Comparing `a_vec` with each lane of `b_vec` can't be done with `svmatch`,
1215
+ // the same way as in `simsimd_intersect_u16_sve2`, as that instruction is only
1216
+ // available for 8-bit and 16-bit integers.
1217
+ //
1218
+ // svbool_t equal_mask = svpfalse_b();
1219
+ // for (simsimd_size_t i = 0; i < register_size; i++) {
1220
+ // equal_mask = svorr_z(svptrue_b32(), equal_mask, svcmpeq_u32(a_progress, a_vec, b_vec));
1221
+ // b_vec = svext_u32(b_vec, b_vec, 1);
1222
+ // }
1223
+ // simsimd_size_t equal_count = svcntp_b32(a_progress, equal_mask);
1224
+ //
1225
+ // Alternatively, one can use histogram instructions, like `svhistcnt_u32_z`.
1226
+ // They practically compute the prefix-matching count, which is equivalent to
1227
+ // the lower triangle of the row-major intersection matrix.
1228
+ // To compute the upper triangle, we can reverse (with `svrev_b32`) the order of
1229
+ // elements and repeat the operation, accumulating the results for top and bottom.
1230
+ // Let's look at 4x element registers as an example:
1231
+ //
1232
+ // ⊐ α = {A, B, C, D}, β = {X, Y, Z, W}:
1233
+ //
1234
+ // hist(α, β): hist(α_rev, β_rev):
1235
+ //
1236
+ // X Y Z W W Z Y X
1237
+ // A 1 0 0 0 D 1 0 0 0
1238
+ // B 1 1 0 0 C 1 1 0 0
1239
+ // C 1 1 1 0 B 1 1 1 0
1240
+ // D 1 1 1 1 A 1 1 1 1
1241
+ //
1242
+ svuint32_t hist_lower = svhistcnt_u32_z(a_progress, a_vec, b_vec);
1243
+ svuint32_t a_rev_vec = svrev_u32(a_vec);
1244
+ svuint32_t b_rev_vec = svrev_u32(b_vec);
1245
+ svuint32_t hist_upper = svrev_u32(svhistcnt_u32_z(svptrue_b32(), a_rev_vec, b_rev_vec));
1246
+ svuint32_t hist = svorr_u32_x(a_progress, hist_lower, hist_upper);
1247
+ svbool_t equal_mask = svcmpne_n_u32(a_progress, hist, 0);
1248
+ simsimd_size_t equal_count = svcntp_b32(a_progress, equal_mask);
1249
+
1250
+ // Advance
1251
+ a_idx += a_step;
1252
+ b_idx += b_step;
1253
+ c += equal_count;
1254
+ }
1255
+ *results = c;
1256
+ }
1257
+
1258
+ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2( //
1259
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
1260
+ simsimd_i16_t const *a_weights, simsimd_i16_t const *b_weights, //
1261
+ simsimd_size_t a_length, simsimd_size_t b_length, //
1262
+ simsimd_distance_t *results) {
1263
+
1264
+ // A single SVE lane is 128 bits wide, so one lane fits 8 values.
1265
+ simsimd_size_t const register_size = svcnth();
1266
+ simsimd_size_t const lanes_count = register_size / 8;
1267
+ simsimd_size_t a_idx = 0, b_idx = 0;
1268
+ svint64_t product_vec = svdupq_n_s64(0, 0);
1269
+ simsimd_size_t intersection_size = 0;
1270
+
1271
+ while (a_idx < a_length && b_idx < b_length) {
1272
+ // Load `a_member` and broadcast it, load `b_members_vec` from memory
1273
+ svbool_t a_progress = svwhilelt_b16_u64(a_idx, a_length);
1274
+ svbool_t b_progress = svwhilelt_b16_u64(b_idx, b_length);
1275
+ svuint16_t a_vec = svld1_u16(a_progress, a + a_idx);
1276
+ svuint16_t b_vec = svld1_u16(b_progress, b + b_idx);
1277
+
1278
+ // Intersecting registers with `svmatch_u16` involves a lot of shuffling
1279
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
1280
+ simsimd_u16_t a_min;
1281
+ simsimd_u16_t a_max = svlastb(a_progress, a_vec);
1282
+ simsimd_u16_t b_min = svlasta(svpfalse_b(), b_vec);
1283
+ simsimd_u16_t b_max = svlastb(b_progress, b_vec);
1284
+
1285
+ // If the slices don't overlap, advance the appropriate pointer
1286
+ while (a_max < b_min && (a_idx + register_size) < a_length) {
1287
+ a_idx += register_size;
1288
+ a_progress = svwhilelt_b16_u64(a_idx, a_length);
1289
+ a_vec = svld1_u16(a_progress, a + a_idx);
1290
+ a_max = svlastb(a_progress, a_vec);
1291
+ }
1292
+ a_min = svlasta(svpfalse_b(), a_vec);
1293
+ while (b_max < a_min && (b_idx + register_size) < b_length) {
1294
+ b_idx += register_size;
1295
+ b_progress = svwhilelt_b16_u64(b_idx, b_length);
1296
+ b_vec = svld1_u16(b_progress, b + b_idx);
1297
+ b_max = svlastb(b_progress, b_vec);
1298
+ }
1299
+ b_min = svlasta(svpfalse_b(), b_vec);
1300
+
1301
+ // Before we evaluate the intersection size, obfurscating the order in `b_vec`,
1302
+ // let's estimate how much we will need to advance the pointers afterwards.
1303
+ // For that, we don't even need to broadcast the values in SVE, as the whole
1304
+ // register can be compared against a scalar:
1305
+ //
1306
+ // svuint16_t a_last_broadcasted = svdup_n_u16(a_max);
1307
+ // svuint16_t b_last_broadcasted = svdup_n_u16(b_max);
1308
+ svbool_t a_mask = svcmple_n_u16(a_progress, a_vec, b_max);
1309
+ svbool_t b_mask = svcmple_n_u16(b_progress, b_vec, a_max);
1310
+ simsimd_u64_t a_step = svcntp_b16(a_progress, a_mask);
1311
+ simsimd_u64_t b_step = svcntp_b16(b_progress, b_mask);
1312
+
1313
+ // Compare `a_vec` with each lane of `b_vec`
1314
+ svint16_t a_weights_vec = svld1_s16(a_progress, a_weights + a_idx);
1315
+ svint16_t b_weights_vec = svld1_s16(b_progress, b_weights + b_idx);
1316
+ for (simsimd_size_t i = 0; i < lanes_count; i++) {
1317
+ svbool_t equal_mask = svmatch_u16(a_progress, a_vec, b_vec);
1318
+ svint16_t b_equal_weights_vec = svsel_s16(equal_mask, b_weights_vec, svdup_n_s16(0.f));
1319
+ product_vec = svdot_s64(product_vec, a_weights_vec, b_equal_weights_vec);
1320
+ b_vec = svext_u16(b_vec, b_vec, 8);
1321
+ intersection_size += svcntp_b16(svptrue_b16(), equal_mask);
1322
+ }
1323
+
1324
+ // Advance
1325
+ a_idx += a_step;
1326
+ b_idx += b_step;
1327
+ }
1328
+ results[0] = (simsimd_distance_t)intersection_size;
1329
+ results[1] = svaddv_s64(svptrue_b64(), product_vec);
1330
+ }
1331
+
1332
+ #pragma clang attribute pop
1333
+ #pragma GCC pop_options
1334
+ #endif // SIMSIMD_TARGET_SVE2
1335
+
1336
+ #if SIMSIMD_TARGET_SVE2 && SIMSIMD_TARGET_SVE_BF16
1337
+ #pragma GCC push_options
1338
+ #pragma GCC target("arch=armv8.6-a+sve+sve2+bf16")
1339
+ #pragma clang attribute push(__attribute__((target("arch=armv8.6-a+sve+sve2+bf16"))), apply_to = function)
1340
+
1341
+ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2( //
1342
+ simsimd_u16_t const *a, simsimd_u16_t const *b, //
1343
+ simsimd_bf16_t const *a_weights, simsimd_bf16_t const *b_weights, //
1344
+ simsimd_size_t a_length, simsimd_size_t b_length, //
1345
+ simsimd_distance_t *results) {
1346
+
1347
+ // A single SVE lane is 128 bits wide, so one lane fits 8 values.
1348
+ simsimd_size_t const register_size = svcnth();
1349
+ simsimd_size_t const lanes_count = register_size / 8;
1350
+ simsimd_size_t a_idx = 0, b_idx = 0;
1351
+ svfloat32_t product_vec = svdupq_n_f32(0.f, 0.f, 0.f, 0.f);
1352
+ simsimd_size_t intersection_size = 0;
1353
+
1354
+ while (a_idx < a_length && b_idx < b_length) {
1355
+ // Load `a_member` and broadcast it, load `b_members_vec` from memory
1356
+ svbool_t a_progress = svwhilelt_b16_u64(a_idx, a_length);
1357
+ svbool_t b_progress = svwhilelt_b16_u64(b_idx, b_length);
1358
+ svuint16_t a_vec = svld1_u16(a_progress, a + a_idx);
1359
+ svuint16_t b_vec = svld1_u16(b_progress, b + b_idx);
1360
+
1361
+ // Intersecting registers with `svmatch_u16` involves a lot of shuffling
1362
+ // and comparisons, so we want to avoid it if the slices don't overlap at all..
1363
+ simsimd_u16_t a_min;
1364
+ simsimd_u16_t a_max = svlastb(a_progress, a_vec);
1365
+ simsimd_u16_t b_min = svlasta(svpfalse_b(), b_vec);
1366
+ simsimd_u16_t b_max = svlastb(b_progress, b_vec);
1367
+
1368
+ // If the slices don't overlap, advance the appropriate pointer
1369
+ while (a_max < b_min && (a_idx + register_size) < a_length) {
1370
+ a_idx += register_size;
1371
+ a_progress = svwhilelt_b16_u64(a_idx, a_length);
1372
+ a_vec = svld1_u16(a_progress, a + a_idx);
1373
+ a_max = svlastb(a_progress, a_vec);
1374
+ }
1375
+ a_min = svlasta(svpfalse_b(), a_vec);
1376
+ while (b_max < a_min && (b_idx + register_size) < b_length) {
1377
+ b_idx += register_size;
1378
+ b_progress = svwhilelt_b16_u64(b_idx, b_length);
1379
+ b_vec = svld1_u16(b_progress, b + b_idx);
1380
+ b_max = svlastb(b_progress, b_vec);
1381
+ }
1382
+ b_min = svlasta(svpfalse_b(), b_vec);
1383
+
1384
+ // Before we evaluate the intersection size, obfurscating the order in `b_vec`,
1385
+ // let's estimate how much we will need to advance the pointers afterwards.
1386
+ // For that, we don't even need to broadcast the values in SVE, as the whole
1387
+ // register can be compared against a scalar:
1388
+ //
1389
+ // svuint16_t a_last_broadcasted = svdup_n_u16(a_max);
1390
+ // svuint16_t b_last_broadcasted = svdup_n_u16(b_max);
1391
+ svbool_t a_mask = svcmple_n_u16(a_progress, a_vec, b_max);
1392
+ svbool_t b_mask = svcmple_n_u16(b_progress, b_vec, a_max);
1393
+ simsimd_u64_t a_step = svcntp_b16(a_progress, a_mask);
1394
+ simsimd_u64_t b_step = svcntp_b16(b_progress, b_mask);
1395
+
1396
+ // Compare `a_vec` with each lane of `b_vec`
1397
+ svbfloat16_t a_weights_vec = svld1_bf16(a_progress, (__bf16 const *)a_weights + a_idx);
1398
+ svbfloat16_t b_weights_vec = svld1_bf16(b_progress, (__bf16 const *)b_weights + b_idx);
1399
+ for (simsimd_size_t i = 0; i < lanes_count; i++) {
1400
+ svbool_t equal_mask = svmatch_u16(a_progress, a_vec, b_vec);
1401
+ //! The `svsel_bf16` intrinsic is broken in many compilers, not returning the correct type.
1402
+ //! So we reinterprete floats as integers and apply `svsel_s16`, but the `svreinterpret_s16_bs16`
1403
+ //! and `svreinterpret_bf16_s16` are not always properly defined!
1404
+ svint16_t b_equal_weights_vec =
1405
+ svsel_s16(equal_mask, svreinterpret_s16_bf16(b_weights_vec), svdup_n_s16(0));
1406
+ product_vec = svbfdot_f32(product_vec, a_weights_vec, svreinterpret_bf16_s16(b_equal_weights_vec));
1407
+ b_vec = svext_u16(b_vec, b_vec, 8);
1408
+ intersection_size += svcntp_b16(svptrue_b16(), equal_mask);
1409
+ }
1410
+
1411
+ // Advance
1412
+ a_idx += a_step;
1413
+ b_idx += b_step;
1414
+ }
1415
+ results[0] = (simsimd_distance_t)intersection_size;
1416
+ results[1] = svaddv_f32(svptrue_b32(), product_vec);
1417
+ }
1418
+
1419
+ #pragma clang attribute pop
1420
+ #pragma GCC pop_options
1421
+ #endif // SIMSIMD_TARGET_SVE2 && SIMSIMD_TARGET_SVE_BF16
1422
+ #endif // _SIMSIMD_TARGET_ARM
1423
+
1424
+ #ifdef __cplusplus
1425
+ }
1426
+ #endif
1427
+
1428
+ #endif