usearch 0.18.6 → 0.18.8

package/python/scripts/bench.py

@@ -21,12 +21,15 @@ def bench_speed(
     connectivity: int = DEFAULT_CONNECTIVITY,
     expansion_add: int = DEFAULT_EXPANSION_ADD,
     expansion_search: int = DEFAULT_EXPANSION_SEARCH,
+    jit: bool = False,
     train: bool = False,
 ) -> pd.DataFrame:
 
     # Build various indexes:
     indexes = []
     jit_options = [False]
+    if jit:
+        jit_options.append(jit)
     dtype_options = ['f32', 'f16', 'f8']
     for jit, dtype in itertools.product(jit_options, dtype_options):
 
@@ -39,6 +42,10 @@ def bench_speed(
             connectivity=connectivity,
             path='USearch' + ['', '+JIT'][jit] + ':' + dtype,
         )
+        
+        # Skip the cases, where JIT-ing is impossible
+        if jit and not index.jit:
+            continue
         indexes.append(index)
 
     # Add FAISS indexes to the mix:

package/python/usearch/index.py

@@ -184,6 +184,13 @@ class Index:
         :type view: bool, optional
         """
 
+        if metric in MetricKindBitwise:
+            assert dtype is None or dtype == ScalarKind.B1
+            dtype = ScalarKind.B1
+        else:
+            dtype = _normalize_dtype(dtype)
+
+
         if metric is None:
             metric = MetricKind.IP
         elif isinstance(metric, str):
@@ -230,12 +237,6 @@ class Index:
             raise ValueError(
                 'The `metric` must be Numba callback or a `MetricKind`')
 
-        if metric in MetricKindBitwise:
-            assert dtype is None or dtype == ScalarKind.B1
-            dtype = ScalarKind.B1
-        else:
-            dtype = _normalize_dtype(dtype)
-
         self._compiled = _CompiledIndex(
             ndim=ndim,
             metric=self._metric_kind,

package/python/usearch/numba.py

@@ -8,9 +8,12 @@ from typing import Callable
 import numpy as np
 from numba import cfunc, types, carray
 
-from usearch.compiled import MetricKind
+from usearch.compiled import MetricKind, ScalarKind
 
 
+signature_i8args = types.float32(types.CPointer(
+    types.int8), types.CPointer(types.int8))
+
 signature_f16args = types.float32(types.CPointer(
     types.float16), types.CPointer(types.float16))
 
@@ -23,7 +26,7 @@ signature_f64args = types.float32(types.CPointer(
 
 def jit(ndim: int,
         metric: MetricKind = MetricKind.Cos,
-        dtype: np.generic = np.float32) -> Callable:
+        dtype: ScalarKind = ScalarKind.F32) -> Callable:
     """JIT-compiles a distance metric specifically tuned for the target hardware
     and number of dimensions.
 
@@ -31,22 +34,25 @@ def jit(ndim: int,
     of `ctypes` to support half-precision.
     https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
     """
-    normalize_dtype = {
+    assert isinstance(metric, MetricKind)
+    assert isinstance(dtype, ScalarKind)
+
+    numba_supported_types = (
+        ScalarKind.F8,
         # Half-precision is still unsupported
         # https://github.com/numba/numba/issues/4402
-        # 'f16': np.float16,
-        'f32': np.float32,
-        'f64': np.float64,
-    }
-    if dtype in normalize_dtype.keys():
-        dtype = normalize_dtype[dtype]
-    if dtype not in normalize_dtype.values():
+        # ScalarKind.F16: np.float16,
+        ScalarKind.F32,
+        ScalarKind.F64,
+    )
+    if dtype not in numba_supported_types:
         return None
 
     scalar_kind_to_accumulator_type = {
-        np.float16: types.float32,
-        np.float32: types.float32,
-        np.float64: types.float64,
+        ScalarKind.F8: types.int32,
+        ScalarKind.F16: types.float16,
+        ScalarKind.F32: types.float32,
+        ScalarKind.F64: types.float64,
     }
     accumulator = scalar_kind_to_accumulator_type[dtype]
 
@@ -87,9 +93,10 @@ def jit(ndim: int,
         return types.float32(ab_delta_sq)
 
     scalar_kind_to_signature = {
-        np.float16: signature_f16args,
-        np.float32: signature_f32args,
-        np.float64: signature_f64args,
+        ScalarKind.F8: signature_i8args,
+        ScalarKind.F16: signature_f16args,
+        ScalarKind.F32: signature_f32args,
+        ScalarKind.F64: signature_f64args,
     }
 
     metric_kind_to_function = {
@@ -98,4 +105,7 @@ def jit(ndim: int,
         MetricKind.L2sq: numba_l2sq,
     }
 
+    if dtype == ScalarKind.F8 and metric == MetricKind.IP:
+        metric = MetricKind.Cos
+
     return cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])

package/simsimd/include/simsimd/simsimd.h

@@ -185,19 +185,33 @@ inline static simsimd_f32_t simsimd_cos_f16x4neon(simsimd_f16_t const* a, simsim
 
 inline static simsimd_f32_t simsimd_cos_i8x16neon(int8_t const* a, int8_t const* b, size_t d) {
 #if defined(__ARM_NEON)
-    int8x16_t const* a8x16 = reinterpret_cast<int8x16_t const*>(a);
-    int8x16_t const* b8x16 = reinterpret_cast<int8x16_t const*>(b);
-
     int32x4_t ab_vec = vdupq_n_s32(0);
     int32x4_t a2_vec = vdupq_n_s32(0);
     int32x4_t b2_vec = vdupq_n_s32(0);
-    for (size_t i = 0; i != d / 16; ++i) {
-        int8x16_t a_vec = a8x16[i];
-        int8x16_t b_vec = b8x16[i];
+
+#if 0 // This 128-bit `vdot_s32` intrinsic is often unavailable, so we use the 64-bit `vdot_s32`.
+    for (size_t i = 0; i != d; i += 16) {
+        int8x16_t a_vec = vld1q_s8(a + i);
+        int8x16_t b_vec = vld1q_s8(b + i);
         ab_vec = vdotq_s32(ab_vec, a_vec, b_vec);
         a2_vec = vdotq_s32(a2_vec, a_vec, a_vec);
         b2_vec = vdotq_s32(b2_vec, b_vec, b_vec);
     }
+#else
+    for (size_t i = 0; i != d; i += 8) {
+        int16x8_t a_vec = vmovl_s8(vld1_s8(a + i));
+        int16x8_t b_vec = vmovl_s8(vld1_s8(b + i));
+        int16x8_t ab_part_vec = vmulq_s16(a_vec, b_vec);
+        int16x8_t a2_part_vec = vmulq_s16(a_vec, a_vec);
+        int16x8_t b2_part_vec = vmulq_s16(b_vec, b_vec);
+        ab_vec =
+            vaddq_s32(ab_vec, vaddq_s32(vmovl_s16(vget_high_s16(ab_part_vec)), vmovl_s16(vget_low_s16(ab_part_vec))));
+        a2_vec =
+            vaddq_s32(a2_vec, vaddq_s32(vmovl_s16(vget_high_s16(a2_part_vec)), vmovl_s16(vget_low_s16(a2_part_vec))));
+        b2_vec =
+            vaddq_s32(b2_vec, vaddq_s32(vmovl_s16(vget_high_s16(b2_part_vec)), vmovl_s16(vget_low_s16(b2_part_vec))));
+    }
+#endif
 
     int32x2_t ab_part = vadd_s32(vget_high_s32(ab_vec), vget_low_s32(ab_vec));
     int32_t ab = vget_lane_s32(vpadd_s32(ab_part, ab_part), 0);