usearch 2.21.3 → 2.25.0

package/numkong/include/numkong/sparse/neon.h

@@ -0,0 +1,288 @@
+/**
+ *  @brief NEON-accelerated Sparse Vector Operations.
+ *  @file include/numkong/sparse/neon.h
+ *  @author Ash Vardanian
+ *  @date February 6, 2026
+ *
+ *  @sa include/numkong/sparse.h
+ */
+#ifndef NK_SPARSE_NEON_H
+#define NK_SPARSE_NEON_H
+
+#if NK_TARGET_ARM64_
+#if NK_TARGET_NEON
+
+#include "numkong/types.h"
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if defined(__clang__)
+#pragma clang attribute push(__attribute__((target("arch=armv8-a"))), apply_to = function)
+#elif defined(__GNUC__)
+#pragma GCC push_options
+#pragma GCC target("arch=armv8-a")
+#endif
+
+NK_INTERNAL uint32x4_t nk_intersect_u32x4_neon_(uint32x4_t a, uint32x4_t b) {
+    uint32x4_t b_rot1 = vextq_u32(b, b, 1);
+    uint32x4_t b_rot2 = vextq_u32(b, b, 2);
+    uint32x4_t b_rot3 = vextq_u32(b, b, 3);
+    uint32x4_t matches_rot0 = vceqq_u32(a, b);
+    uint32x4_t matches_rot1 = vceqq_u32(a, b_rot1);
+    uint32x4_t matches_rot2 = vceqq_u32(a, b_rot2);
+    uint32x4_t matches_rot3 = vceqq_u32(a, b_rot3);
+    uint32x4_t matches = vorrq_u32(vorrq_u32(matches_rot0, matches_rot1), vorrq_u32(matches_rot2, matches_rot3));
+    return matches;
+}
+
+NK_INTERNAL uint16x8_t nk_intersect_u16x8_neon_(uint16x8_t a, uint16x8_t b) {
+    uint16x8_t b_rot1 = vextq_u16(b, b, 1);
+    uint16x8_t b_rot2 = vextq_u16(b, b, 2);
+    uint16x8_t b_rot3 = vextq_u16(b, b, 3);
+    uint16x8_t b_rot4 = vextq_u16(b, b, 4);
+    uint16x8_t b_rot5 = vextq_u16(b, b, 5);
+    uint16x8_t b_rot6 = vextq_u16(b, b, 6);
+    uint16x8_t b_rot7 = vextq_u16(b, b, 7);
+    uint16x8_t matches_rot0 = vceqq_u16(a, b);
+    uint16x8_t matches_rot1 = vceqq_u16(a, b_rot1);
+    uint16x8_t matches_rot2 = vceqq_u16(a, b_rot2);
+    uint16x8_t matches_rot3 = vceqq_u16(a, b_rot3);
+    uint16x8_t matches_rot4 = vceqq_u16(a, b_rot4);
+    uint16x8_t matches_rot5 = vceqq_u16(a, b_rot5);
+    uint16x8_t matches_rot6 = vceqq_u16(a, b_rot6);
+    uint16x8_t matches_rot7 = vceqq_u16(a, b_rot7);
+    uint16x8_t matches = vorrq_u16(
+        vorrq_u16(vorrq_u16(matches_rot0, matches_rot1), vorrq_u16(matches_rot2, matches_rot3)),
+        vorrq_u16(vorrq_u16(matches_rot4, matches_rot5), vorrq_u16(matches_rot6, matches_rot7)));
+    return matches;
+}
+
+NK_PUBLIC void nk_sparse_intersect_u16_neon( //
+    nk_u16_t const *a, nk_u16_t const *b,    //
+    nk_size_t a_length, nk_size_t b_length,  //
+    nk_u16_t *result, nk_size_t *count) {
+
+    // NEON lacks compress-store, so fall back to serial for result output
+    if (result) {
+        nk_sparse_intersect_u16_serial(a, b, a_length, b_length, result, count);
+        return;
+    }
+
+#if NK_ALLOW_ISA_REDIRECT
+    // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
+    if (a_length < 32 && b_length < 32) {
+        nk_sparse_intersect_u16_serial(a, b, a_length, b_length, result, count);
+        return;
+    }
+#endif
+
+    nk_u16_t const *const a_end = a + a_length;
+    nk_u16_t const *const b_end = b + b_length;
+    nk_b128_vec_t a_vec, b_vec;
+    uint16x8_t c_counts_u16x8 = vdupq_n_u16(0);
+
+    while (a + 8 <= a_end && b + 8 <= b_end) {
+        a_vec.u16x8 = vld1q_u16(a);
+        b_vec.u16x8 = vld1q_u16(b);
+
+        // Intersecting registers with `nk_intersect_u16x8_neon_` involves a lot of shuffling
+        // and comparisons, so we want to avoid it if the slices don't overlap at all..
+        nk_u16_t a_min;
+        nk_u16_t a_max = a_vec.u16s[7];
+        nk_u16_t b_min = b_vec.u16s[0];
+        nk_u16_t b_max = b_vec.u16s[7];
+
+        // If the slices don't overlap, advance the appropriate pointer
+        while (a_max < b_min && a + 16 <= a_end) {
+            a += 8;
+            a_vec.u16x8 = vld1q_u16(a);
+            a_max = a_vec.u16s[7];
+        }
+        a_min = a_vec.u16s[0];
+        while (b_max < a_min && b + 16 <= b_end) {
+            b += 8;
+            b_vec.u16x8 = vld1q_u16(b);
+            b_max = b_vec.u16s[7];
+        }
+        b_min = b_vec.u16s[0];
+
+        // Transform match-masks into "ones", accumulate them between the cycles,
+        // and merge all together in the end.
+        uint16x8_t a_matches = nk_intersect_u16x8_neon_(a_vec.u16x8, b_vec.u16x8);
+        c_counts_u16x8 = vaddq_u16(c_counts_u16x8, vandq_u16(a_matches, vdupq_n_u16(1)));
+
+        // Use `vclz_u32` to compute leading zeros for both `a_step` and `b_step` in parallel.
+        // Narrow comparison masks from 128→64→32 bits, pack both into a `uint32x2_t`.
+        uint16x8_t a_inrange_u16x8 = vcleq_u16(a_vec.u16x8, vdupq_n_u16(b_max));
+        uint16x8_t b_inrange_u16x8 = vcleq_u16(b_vec.u16x8, vdupq_n_u16(a_max));
+        uint8x8_t a_narrow_u8x8 = vmovn_u16(a_inrange_u16x8);
+        uint8x8_t b_narrow_u8x8 = vmovn_u16(b_inrange_u16x8);
+        uint8x8_t packed_u8x8 = vshrn_n_u16(vreinterpretq_u16_u8(vcombine_u8(a_narrow_u8x8, b_narrow_u8x8)), 4);
+        uint32x2_t clz_u32x2 = vclz_u32(vreinterpret_u32_u8(packed_u8x8));
+        a += (32 - vget_lane_u32(clz_u32x2, 0)) / 4;
+        b += (32 - vget_lane_u32(clz_u32x2, 1)) / 4;
+    }
+
+    nk_size_t tail_count = 0;
+    nk_sparse_intersect_u16_serial(a, b, a_end - a, b_end - b, 0, &tail_count);
+    *count = tail_count + (nk_size_t)vaddvq_u16(c_counts_u16x8);
+}
+
+NK_PUBLIC void nk_sparse_intersect_u32_neon( //
+    nk_u32_t const *a, nk_u32_t const *b,    //
+    nk_size_t a_length, nk_size_t b_length,  //
+    nk_u32_t *result, nk_size_t *count) {
+
+    // NEON lacks compress-store, so fall back to serial for result output
+    if (result) {
+        nk_sparse_intersect_u32_serial(a, b, a_length, b_length, result, count);
+        return;
+    }
+
+#if NK_ALLOW_ISA_REDIRECT
+    // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
+    if (a_length < 32 && b_length < 32) {
+        nk_sparse_intersect_u32_serial(a, b, a_length, b_length, result, count);
+        return;
+    }
+#endif
+
+    nk_u32_t const *const a_end = a + a_length;
+    nk_u32_t const *const b_end = b + b_length;
+    nk_b128_vec_t a_vec, b_vec;
+    uint32x4_t c_counts_u32x4 = vdupq_n_u32(0);
+
+    while (a + 4 <= a_end && b + 4 <= b_end) {
+        a_vec.u32x4 = vld1q_u32(a);
+        b_vec.u32x4 = vld1q_u32(b);
+
+        // Intersecting registers with `nk_intersect_u32x4_neon_` involves a lot of shuffling
+        // and comparisons, so we want to avoid it if the slices don't overlap at all..
+        nk_u32_t a_min;
+        nk_u32_t a_max = a_vec.u32s[3];
+        nk_u32_t b_min = b_vec.u32s[0];
+        nk_u32_t b_max = b_vec.u32s[3];
+
+        // If the slices don't overlap, advance the appropriate pointer
+        while (a_max < b_min && a + 8 <= a_end) {
+            a += 4;
+            a_vec.u32x4 = vld1q_u32(a);
+            a_max = a_vec.u32s[3];
+        }
+        a_min = a_vec.u32s[0];
+        while (b_max < a_min && b + 8 <= b_end) {
+            b += 4;
+            b_vec.u32x4 = vld1q_u32(b);
+            b_max = b_vec.u32s[3];
+        }
+        b_min = b_vec.u32s[0];
+
+        // Transform match-masks into "ones", accumulate them between the cycles,
+        // and merge all together in the end.
+        uint32x4_t a_matches = nk_intersect_u32x4_neon_(a_vec.u32x4, b_vec.u32x4);
+        c_counts_u32x4 = vaddq_u32(c_counts_u32x4, vandq_u32(a_matches, vdupq_n_u32(1)));
+
+        uint32x4_t a_inrange_u32x4 = vcleq_u32(a_vec.u32x4, vdupq_n_u32(b_max));
+        uint32x4_t b_inrange_u32x4 = vcleq_u32(b_vec.u32x4, vdupq_n_u32(a_max));
+        uint8x8_t packed_u8x8 = vmovn_u16(vcombine_u16(vmovn_u32(a_inrange_u32x4), vmovn_u32(b_inrange_u32x4)));
+        uint32x2_t clz_u32x2 = vclz_u32(vreinterpret_u32_u8(packed_u8x8));
+        a += (32 - vget_lane_u32(clz_u32x2, 0)) / 8;
+        b += (32 - vget_lane_u32(clz_u32x2, 1)) / 8;
+    }
+
+    nk_size_t tail_count = 0;
+    nk_sparse_intersect_u32_serial(a, b, a_end - a, b_end - b, 0, &tail_count);
+    *count = tail_count + (nk_size_t)vaddvq_u32(c_counts_u32x4);
+}
+
+NK_INTERNAL uint64x2_t nk_intersect_u64x2_neon_(uint64x2_t a, uint64x2_t b) {
+    uint64x2_t b_rot1 = vextq_u64(b, b, 1);
+    uint64x2_t matches_rot0 = vceqq_u64(a, b);
+    uint64x2_t matches_rot1 = vceqq_u64(a, b_rot1);
+    uint64x2_t matches = vorrq_u64(matches_rot0, matches_rot1);
+    return matches;
+}
+
+NK_PUBLIC void nk_sparse_intersect_u64_neon( //
+    nk_u64_t const *a, nk_u64_t const *b,    //
+    nk_size_t a_length, nk_size_t b_length,  //
+    nk_u64_t *result, nk_size_t *count) {
+
+    // NEON lacks compress-store, so fall back to serial for result output
+    if (result) {
+        nk_sparse_intersect_u64_serial(a, b, a_length, b_length, result, count);
+        return;
+    }
+
+#if NK_ALLOW_ISA_REDIRECT
+    // The baseline implementation for very small arrays (2 registers or less) can be quite simple:
+    if (a_length < 8 && b_length < 8) {
+        nk_sparse_intersect_u64_serial(a, b, a_length, b_length, result, count);
+        return;
+    }
+#endif
+
+    nk_u64_t const *const a_end = a + a_length;
+    nk_u64_t const *const b_end = b + b_length;
+    nk_b128_vec_t a_vec, b_vec;
+    uint64x2_t c_counts_u64x2 = vdupq_n_u64(0);
+
+    while (a + 2 <= a_end && b + 2 <= b_end) {
+        a_vec.u64x2 = vld1q_u64(a);
+        b_vec.u64x2 = vld1q_u64(b);
+
+        // Intersecting registers with `nk_intersect_u64x2_neon_` involves comparisons,
+        // so we want to avoid it if the slices don't overlap at all.
+        nk_u64_t a_min;
+        nk_u64_t a_max = a_vec.u64s[1];
+        nk_u64_t b_min = b_vec.u64s[0];
+        nk_u64_t b_max = b_vec.u64s[1];
+
+        // If the slices don't overlap, advance the appropriate pointer
+        while (a_max < b_min && a + 4 <= a_end) {
+            a += 2;
+            a_vec.u64x2 = vld1q_u64(a);
+            a_max = a_vec.u64s[1];
+        }
+        a_min = a_vec.u64s[0];
+        while (b_max < a_min && b + 4 <= b_end) {
+            b += 2;
+            b_vec.u64x2 = vld1q_u64(b);
+            b_max = b_vec.u64s[1];
+        }
+        b_min = b_vec.u64s[0];
+
+        // Now we are likely to have some overlap, so we can intersect the registers
+        // Transform match-masks into "ones", accumulate them between the cycles,
+        // and merge all together in the end.
+        uint64x2_t a_matches = nk_intersect_u64x2_neon_(a_vec.u64x2, b_vec.u64x2);
+        c_counts_u64x2 = vaddq_u64(c_counts_u64x2, vandq_u64(a_matches, vdupq_n_u64(1)));
+
+        uint64x2_t a_inrange_u64x2 = vcleq_u64(a_vec.u64x2, vdupq_n_u64(b_max));
+        uint64x2_t b_inrange_u64x2 = vcleq_u64(b_vec.u64x2, vdupq_n_u64(a_max));
+        uint16x4_t packed_u16x4 = vmovn_u32(vcombine_u32(vmovn_u64(a_inrange_u64x2), vmovn_u64(b_inrange_u64x2)));
+        uint32x2_t clz_u32x2 = vclz_u32(vreinterpret_u32_u16(packed_u16x4));
+        a += (32 - vget_lane_u32(clz_u32x2, 0)) / 16;
+        b += (32 - vget_lane_u32(clz_u32x2, 1)) / 16;
+    }
+
+    nk_size_t tail_count = 0;
+    nk_sparse_intersect_u64_serial(a, b, a_end - a, b_end - b, 0, &tail_count);
+    *count = tail_count + (nk_size_t)vaddvq_u64(c_counts_u64x2);
+}
+
+#if defined(__clang__)
+#pragma clang attribute pop
+#elif defined(__GNUC__)
+#pragma GCC pop_options
+#endif
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_TARGET_NEON
+#endif // NK_TARGET_ARM64_
+#endif // NK_SPARSE_NEON_H

package/numkong/include/numkong/sparse/serial.h

@@ -0,0 +1,117 @@
+/**
+ *  @brief Serial Sparse Vector Operations.
+ *  @file include/numkong/sparse/serial.h
+ *  @author Ash Vardanian
+ *  @date February 6, 2026
+ *
+ *  @sa include/numkong/sparse.h
+ */
+#ifndef NK_SPARSE_SERIAL_H
+#define NK_SPARSE_SERIAL_H
+
+#include "numkong/types.h"
+#include "numkong/cast/serial.h" // `nk_bf16_to_f32_serial`, `nk_assign_from_to_`
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#define nk_define_sparse_intersect_(input_type)                                                                      \
+    NK_PUBLIC nk_size_t nk_sparse_intersect_##input_type##_galloping_search_(                                        \
+        nk_##input_type##_t const *array, nk_size_t start, nk_size_t length, nk_##input_type##_t val) {              \
+        nk_size_t low = start;                                                                                       \
+        nk_size_t high = start + 1;                                                                                  \
+        while (high < length && array[high] < val) {                                                                 \
+            low = high;                                                                                              \
+            high = (2 * high < length) ? 2 * high : length;                                                          \
+        }                                                                                                            \
+        while (low < high) {                                                                                         \
+            nk_size_t mid = low + (high - low) / 2;                                                                  \
+            if (array[mid] < val) { low = mid + 1; }                                                                 \
+            else { high = mid; }                                                                                     \
+        }                                                                                                            \
+        return low;                                                                                                  \
+    }                                                                                                                \
+    NK_PUBLIC nk_size_t nk_sparse_intersect_##input_type##_linear_scan_(                                             \
+        nk_##input_type##_t const *a, nk_##input_type##_t const *b, nk_size_t a_length, nk_size_t b_length,          \
+        nk_##input_type##_t *result) {                                                                               \
+        nk_size_t intersection_size = 0;                                                                             \
+        nk_size_t i = 0, j = 0;                                                                                      \
+        while (i != a_length && j != b_length) {                                                                     \
+            nk_##input_type##_t ai = a[i];                                                                           \
+            nk_##input_type##_t bj = b[j];                                                                           \
+            if (ai == bj) {                                                                                          \
+                if (result) result[intersection_size] = ai;                                                          \
+                intersection_size++;                                                                                 \
+            }                                                                                                        \
+            i += ai <= bj;                                                                                           \
+            j += ai >= bj;                                                                                           \
+        }                                                                                                            \
+        return intersection_size;                                                                                    \
+    }                                                                                                                \
+    NK_PUBLIC void nk_sparse_intersect_##input_type##_serial(                                                        \
+        nk_##input_type##_t const *shorter, nk_##input_type##_t const *longer, nk_size_t shorter_length,             \
+        nk_size_t longer_length, nk_##input_type##_t *result, nk_size_t *count) {                                    \
+        /* Swap arrays if necessary, as we want "longer" to be larger than "shorter" */                              \
+        if (longer_length < shorter_length) {                                                                        \
+            nk_##input_type##_t const *temp = shorter;                                                               \
+            shorter = longer;                                                                                        \
+            longer = temp;                                                                                           \
+            nk_size_t temp_length = shorter_length;                                                                  \
+            shorter_length = longer_length;                                                                          \
+            longer_length = temp_length;                                                                             \
+        }                                                                                                            \
+                                                                                                                     \
+        /* Use the accurate implementation if galloping is not beneficial */                                         \
+        if (longer_length < 64 * shorter_length) {                                                                   \
+            *count = nk_sparse_intersect_##input_type##_linear_scan_(shorter, longer, shorter_length, longer_length, \
+                                                                     result);                                        \
+            return;                                                                                                  \
+        }                                                                                                            \
+                                                                                                                     \
+        /* Perform galloping, shrinking the target range */                                                          \
+        nk_size_t intersection_size = 0;                                                                             \
+        nk_size_t j = 0;                                                                                             \
+        for (nk_size_t i = 0; i < shorter_length; ++i) {                                                             \
+            nk_##input_type##_t shorter_i = shorter[i];                                                              \
+            j = nk_sparse_intersect_##input_type##_galloping_search_(longer, j, longer_length, shorter_i);           \
+            if (j < longer_length && longer[j] == shorter_i) {                                                       \
+                if (result) result[intersection_size] = shorter_i;                                                   \
+                intersection_size++;                                                                                 \
+            }                                                                                                        \
+        }                                                                                                            \
+        *count = intersection_size;                                                                                  \
+    }
+
+#define nk_define_sparse_dot_(input_type, weight_type, accumulator_type, load_and_convert)                 \
+    NK_PUBLIC void nk_sparse_dot_##input_type##weight_type##_serial(                                       \
+        nk_##input_type##_t const *a, nk_##input_type##_t const *b, nk_##weight_type##_t const *a_weights, \
+        nk_##weight_type##_t const *b_weights, nk_size_t a_length, nk_size_t b_length,                     \
+        nk_##accumulator_type##_t *product) {                                                              \
+        nk_##accumulator_type##_t weights_product = 0, awi, bwi;                                           \
+        nk_size_t i = 0, j = 0;                                                                            \
+        while (i != a_length && j != b_length) {                                                           \
+            nk_##input_type##_t ai = a[i];                                                                 \
+            nk_##input_type##_t bj = b[j];                                                                 \
+            int matches = ai == bj;                                                                        \
+            load_and_convert(a_weights + i, &awi);                                                         \
+            load_and_convert(b_weights + j, &bwi);                                                         \
+            weights_product += (nk_##accumulator_type##_t)matches * awi * bwi;                             \
+            i += ai < bj;                                                                                  \
+            j += ai >= bj;                                                                                 \
+        }                                                                                                  \
+        *product = weights_product;                                                                        \
+    }
+
+nk_define_sparse_intersect_(u16) // nk_sparse_intersect_u16_serial
+nk_define_sparse_intersect_(u32) // nk_sparse_intersect_u32_serial
+nk_define_sparse_intersect_(u64) // nk_sparse_intersect_u64_serial
+
+nk_define_sparse_dot_(u16, bf16, f32, nk_bf16_to_f32_serial) // nk_sparse_dot_u16bf16_serial
+nk_define_sparse_dot_(u32, f32, f64, nk_assign_from_to_)     // nk_sparse_dot_u32f32_serial
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_SPARSE_SERIAL_H