Break: New dispatch schema for Arm

ashvardanian · ashvardanian · commit 9bd48ab0710f · 2024-06-16T21:14:31.000Z
This commit add new capability levels for Arm allowing
us to differentiate f16, bf16. and i8-supporting generations
of CPUs, becoming increasingly popular in the datacenter.

This breaks compilation of Rust and Python bindings
due to the "target specific options mismatch".
diff --git a/c/lib.c b/c/lib.c
@@ -110,7 +110,13 @@ SIMSIMD_METRIC_DECLARATION(js, f32, f32)
 SIMSIMD_METRIC_DECLARATION(js, f64, f64)
 
 SIMSIMD_DYNAMIC int simsimd_uses_neon(void) { return (simsimd_capabilities() & simsimd_cap_neon_k) != 0; }
+SIMSIMD_DYNAMIC int simsimd_uses_neon_f16(void) { return (simsimd_capabilities() & simsimd_cap_neon_f16_k) != 0; }
+SIMSIMD_DYNAMIC int simsimd_uses_neon_bf16(void) { return (simsimd_capabilities() & simsimd_cap_neon_bf16_k) != 0; }
+SIMSIMD_DYNAMIC int simsimd_uses_neon_i8(void) { return (simsimd_capabilities() & simsimd_cap_neon_i8_k) != 0; }
 SIMSIMD_DYNAMIC int simsimd_uses_sve(void) { return (simsimd_capabilities() & simsimd_cap_sve_k) != 0; }
+SIMSIMD_DYNAMIC int simsimd_uses_sve_f16(void) { return (simsimd_capabilities() & simsimd_cap_sve_f16_k) != 0; }
+SIMSIMD_DYNAMIC int simsimd_uses_sve_bf16(void) { return (simsimd_capabilities() & simsimd_cap_sve_bf16_k) != 0; }
+SIMSIMD_DYNAMIC int simsimd_uses_sve_i8(void) { return (simsimd_capabilities() & simsimd_cap_sve_i8_k) != 0; }
 SIMSIMD_DYNAMIC int simsimd_uses_haswell(void) { return (simsimd_capabilities() & simsimd_cap_haswell_k) != 0; }
 SIMSIMD_DYNAMIC int simsimd_uses_skylake(void) { return (simsimd_capabilities() & simsimd_cap_skylake_k) != 0; }
 SIMSIMD_DYNAMIC int simsimd_uses_ice(void) { return (simsimd_capabilities() & simsimd_cap_ice_k) != 0; }
diff --git a/cpp/bench.cxx b/cpp/bench.cxx
@@ -242,7 +242,13 @@ int main(int argc, char** argv) {
     std::printf("\n");
     std::printf("Run-time settings:\n");
     std::printf("- Arm NEON support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_k) != 0]);
+    std::printf("- Arm NEON F16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_f16_k) != 0]);
+    std::printf("- Arm NEON BF16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_bf16_k) != 0]);
+    std::printf("- Arm NEON I8 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_i8_k) != 0]);
     std::printf("- Arm SVE support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_k) != 0]);
+    std::printf("- Arm SVE F16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_f16_k) != 0]);
+    std::printf("- Arm SVE BF16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_bf16_k) != 0]);
+    std::printf("- Arm SVE I8 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_i8_k) != 0]);
     std::printf("- x86 Haswell support enabled: %s\n", flags[(runtime_caps & simsimd_cap_haswell_k) != 0]);
     std::printf("- x86 Skylake support enabled: %s\n", flags[(runtime_caps & simsimd_cap_skylake_k) != 0]);
     std::printf("- x86 Ice Lake support enabled: %s\n", flags[(runtime_caps & simsimd_cap_ice_k) != 0]);
diff --git a/cpp/test.c b/cpp/test.c
@@ -36,7 +36,13 @@ void print_capabilities(void) {
     printf("\n");
     printf("Run-time settings:\n");
     printf("- Arm NEON support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_k) != 0]);
+    printf("- Arm NEON F16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_f16_k) != 0]);
+    printf("- Arm NEON BF16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_bf16_k) != 0]);
+    printf("- Arm NEON I8 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_neon_i8_k) != 0]);
     printf("- Arm SVE support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_k) != 0]);
+    printf("- Arm SVE F16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_f16_k) != 0]);
+    printf("- Arm SVE BF16 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_bf16_k) != 0]);
+    printf("- Arm SVE I8 support enabled: %s\n", flags[(runtime_caps & simsimd_cap_sve_i8_k) != 0]);
     printf("- x86 Haswell support enabled: %s\n", flags[(runtime_caps & simsimd_cap_haswell_k) != 0]);
     printf("- x86 Skylake support enabled: %s\n", flags[(runtime_caps & simsimd_cap_skylake_k) != 0]);
     printf("- x86 Ice Lake support enabled: %s\n", flags[(runtime_caps & simsimd_cap_ice_k) != 0]);
diff --git a/include/simsimd/simsimd.h b/include/simsimd/simsimd.h
diff --git a/include/simsimd/spatial.h b/include/simsimd/spatial.h
@@ -300,6 +300,7 @@ SIMSIMD_PUBLIC void simsimd_cos_f16_neon(simsimd_f16_t const* a, simsimd_f16_t c
 
 SIMSIMD_PUBLIC void simsimd_cos_bf16_neon(simsimd_bf16_t const* a, simsimd_bf16_t const* b, simsimd_size_t n,
                                           simsimd_distance_t* result) {
+    // TODO: Redo with BFMMLA - vbfmmlaq_f32
     float32x4_t ab_high_vec = vdupq_n_f32(0), ab_low_vec = vdupq_n_f32(0);
     float32x4_t a2_high_vec = vdupq_n_f32(0), a2_low_vec = vdupq_n_f32(0);
     float32x4_t b2_high_vec = vdupq_n_f32(0), b2_low_vec = vdupq_n_f32(0);
@@ -440,6 +441,7 @@ SIMSIMD_PUBLIC void simsimd_cos_i8_neon(simsimd_i8_t const* a, simsimd_i8_t cons
     //     b2_vec = vaddq_s32(b2_vec, vaddq_s32(vmovl_s16(vget_high_s16(b2_part_vec)), //
     //                                          vmovl_s16(vget_low_s16(b2_part_vec))));
     // }
+    // TODO: Redo with MMLA: vmmlaq_s32
     for (; i + 16 <= n; i += 16) {
         int8x16_t a_vec = vld1q_s8(a + i);
         int8x16_t b_vec = vld1q_s8(b + i);
diff --git a/include/simsimd/types.h b/include/simsimd/types.h
@@ -87,28 +87,6 @@
 
 // Compiling for x86: SIMSIMD_TARGET_SKYLAKE, SIMSIMD_TARGET_ICE, SIMSIMD_TARGET_SAPPHIRE
 //
-// It's important to provide fine-grained controls over AVX512 families, as they are very fragmented:
-// - Intel Skylake servers: F, CD, VL, DQ, BW
-// - Intel Cascade Lake workstations: F, CD, VL, DQ, BW, VNNI
-//      > In other words, it extends Skylake with VNNI support
-// - Intel Sunny Cove (Ice Lake) servers:
-//        F, CD, VL, DQ, BW, VNNI, VPOPCNTDQ, IFMA, VBMI, VAES, GFNI, VBMI2, BITALG, VPCLMULQDQ
-// - AMD Zen4 (Genoa):
-//        F, CD, VL, DQ, BW, VNNI, VPOPCNTDQ, IFMA, VBMI, VAES, GFNI, VBMI2, BITALG, VPCLMULQDQ, BF16
-//      > In other words, it extends Sunny Cove with BF16 support
-// - Golden Cove (Sapphire Rapids): extends Zen4 and Sunny Cove with FP16 support
-//
-// Intel Palm Cove was an irrelevant intermediate release extending Skylake with IFMA and VBMI.
-// Intel Willow Cove was an irrelevant intermediate release extending Sunny Cove with VP2INTERSECT,
-// that aren't supported by any other CPU built to date... and those are only available in Tiger Lake laptops.
-// Intel Cooper Lake was the only intermediary platform, that supported BF16, but not FP16.
-// It's mostly used in 4-socket and 8-socket high-memory configurations.
-//
-// In practical terms, it makes sense to differentiate only 3 AVX512 generations:
-// 1. Skylake (pre 2019): supports single-precision dot-products.
-// 2. Ice Lake (2019-2021): advanced integer algorithms.
-// 3. Sapphire Rapids (2023+): advanced mixed-precision float processing.
-//
 // To list all available macros for x86, take a recent compiler, like GCC 12 and run:
 //      gcc-12 -march=sapphirerapids -dM -E - < /dev/null | egrep "SSE|AVX" | sort
 // On Arm machines you may want to check for other flags:
@@ -199,14 +177,13 @@ typedef unsigned long long simsimd_u64_t;
 typedef simsimd_u64_t simsimd_size_t;
 typedef simsimd_f64_t simsimd_distance_t;
 
-#if !defined(SIMSIMD_NATIVE_F16) || SIMSIMD_NATIVE_F16
-/**
- *  @brief  Half-precision floating-point type.
+/*  @brief  Half-precision floating-point type.
  *
- *  - GCC or Clang on 64-bit ARM: `__fp16`, may require `-mfp16-format` option.
+ *  - GCC or Clang on 64-bit Arm: `__fp16`, may require `-mfp16-format` option.
  *  - GCC or Clang on 64-bit x86: `_Float16`.
  *  - Default: `unsigned short`.
  */
+#if !defined(SIMSIMD_NATIVE_F16) || SIMSIMD_NATIVE_F16
 #if (defined(__GNUC__) || defined(__clang__)) && (defined(__ARM_ARCH) || defined(__aarch64__)) &&                      \
     (defined(__ARM_FP16_FORMAT_IEEE))
 #if !defined(SIMSIMD_NATIVE_F16)
@@ -228,6 +205,12 @@ typedef _Float16 simsimd_f16_t;
 typedef unsigned short simsimd_f16_t;
 #endif
 
+/*  @brief  Half-precision "brain" floating-point type.
+ *
+ *  - GCC or Clang on 64-bit Arm: `__bf16`, may require `-mbf16-format` option.
+ *  - GCC or Clang on 64-bit x86: `bfloat16_t`.
+ *  - Default: `unsigned short`.
+ */
 #if !defined(SIMSIMD_NATIVE_BF16)
 #define SIMSIMD_NATIVE_BF16 1
 #endif
diff --git a/python/lib.c b/python/lib.c
@@ -181,10 +181,20 @@ static PyObject* api_enable_capability(PyObject* self, PyObject* args) {
 
     if (same_string(cap_name, "neon")) {
         static_capabilities |= simsimd_cap_neon_k;
+    } else if (same_string(cap_name, "neon_f16")) {
+        static_capabilities |= simsimd_cap_neon_f16_k;
+    } else if (same_string(cap_name, "neon_bf16")) {
+        static_capabilities |= simsimd_cap_neon_bf16_k;
+    } else if (same_string(cap_name, "neon_i8")) {
+        static_capabilities |= simsimd_cap_neon_i8_k;
     } else if (same_string(cap_name, "sve")) {
         static_capabilities |= simsimd_cap_sve_k;
-    } else if (same_string(cap_name, "sve2")) {
-        static_capabilities |= simsimd_cap_sve2_k;
+    } else if (same_string(cap_name, "sve_f16")) {
+        static_capabilities |= simsimd_cap_sve_f16_k;
+    } else if (same_string(cap_name, "sve_bf16")) {
+        static_capabilities |= simsimd_cap_sve_bf16_k;
+    } else if (same_string(cap_name, "sve_i8")) {
+        static_capabilities |= simsimd_cap_sve_i8_k;
     } else if (same_string(cap_name, "haswell")) {
         static_capabilities |= simsimd_cap_haswell_k;
     } else if (same_string(cap_name, "skylake")) {
@@ -214,10 +224,20 @@ static PyObject* api_disable_capability(PyObject* self, PyObject* args) {
 
     if (same_string(cap_name, "neon")) {
         static_capabilities &= ~simsimd_cap_neon_k;
+    } else if (same_string(cap_name, "neon_f16")) {
+        static_capabilities &= ~simsimd_cap_neon_f16_k;
+    } else if (same_string(cap_name, "neon_bf16")) {
+        static_capabilities &= ~simsimd_cap_neon_bf16_k;
+    } else if (same_string(cap_name, "neon_i8")) {
+        static_capabilities &= ~simsimd_cap_neon_i8_k;
     } else if (same_string(cap_name, "sve")) {
         static_capabilities &= ~simsimd_cap_sve_k;
-    } else if (same_string(cap_name, "sve2")) {
-        static_capabilities &= ~simsimd_cap_sve2_k;
+    } else if (same_string(cap_name, "sve_f16")) {
+        static_capabilities &= ~simsimd_cap_sve_f16_k;
+    } else if (same_string(cap_name, "sve_bf16")) {
+        static_capabilities &= ~simsimd_cap_sve_bf16_k;
+    } else if (same_string(cap_name, "sve_i8")) {
+        static_capabilities &= ~simsimd_cap_sve_i8_k;
     } else if (same_string(cap_name, "haswell")) {
         static_capabilities &= ~simsimd_cap_haswell_k;
     } else if (same_string(cap_name, "skylake")) {
diff --git a/python/test.py b/python/test.py
@@ -74,8 +74,13 @@ def test_capabilities_list():
     """Tests the visibility of hardware capabilities."""
     assert "serial" in simd.get_capabilities()
     assert "neon" in simd.get_capabilities()
+    assert "neon_f16" in simd.get_capabilities()
+    assert "neon_bf16" in simd.get_capabilities()
+    assert "neon_i8" in simd.get_capabilities()
     assert "sve" in simd.get_capabilities()
-    assert "sve2" in simd.get_capabilities()
+    assert "sve_f16" in simd.get_capabilities()
+    assert "sve_bf16" in simd.get_capabilities()
+    assert "sve_i8" in simd.get_capabilities()
     assert "haswell" in simd.get_capabilities()
     assert "ice" in simd.get_capabilities()
     assert "skylake" in simd.get_capabilities()
diff --git a/rust/lib.rs b/rust/lib.rs
@@ -87,7 +87,13 @@ extern "C" {
     fn simsimd_kl_f64(a: *const f64, b: *const f64, c: usize, d: *mut Distance);
 
     fn simsimd_uses_neon() -> i32;
+    fn simsimd_uses_neon_f16() -> i32;
+    fn simsimd_uses_neon_bf16() -> i32;
+    fn simsimd_uses_neon_i8() -> i32;
     fn simsimd_uses_sve() -> i32;
+    fn simsimd_uses_sve_f16() -> i32;
+    fn simsimd_uses_sve_bf16() -> i32;
+    fn simsimd_uses_sve_i8() -> i32;
     fn simsimd_uses_haswell() -> i32;
     fn simsimd_uses_skylake() -> i32;
     fn simsimd_uses_ice() -> i32;
@@ -109,10 +115,34 @@ pub mod capabilties {
         unsafe { crate::simsimd_uses_neon() != 0 }
     }
 
+    pub fn uses_neon_f16() -> bool {
+        unsafe { crate::simsimd_uses_neon_f16() != 0 }
+    }
+
+    pub fn uses_neon_bf16() -> bool {
+        unsafe { crate::simsimd_uses_neon_bf16() != 0 }
+    }
+
+    pub fn uses_neon_i8() -> bool {
+        unsafe { crate::simsimd_uses_neon_i8() != 0 }
+    }
+
     pub fn uses_sve() -> bool {
         unsafe { crate::simsimd_uses_sve() != 0 }
     }
 
+    pub fn uses_sve_f16() -> bool {
+        unsafe { crate::simsimd_uses_sve_f16() != 0 }
+    }
+
+    pub fn uses_sve_bf16() -> bool {
+        unsafe { crate::simsimd_uses_sve_bf16() != 0 }
+    }
+
+    pub fn uses_sve_i8() -> bool {
+        unsafe { crate::simsimd_uses_sve_i8() != 0 }
+    }
+
     pub fn uses_haswell() -> bool {
         unsafe { crate::simsimd_uses_haswell() != 0 }
     }
