Improve: bf16 API for index_dense (#561)

Closes #553

Co-authored-by: anirudlappathi <anirud.lappathi+anirudlappathi@users.noreply.github.com>

Author: anirudlappathi

cpp/test.cpp

@@ -872,16 +872,18 @@ void test_absurd(std::size_t dimensions, std::size_t connectivity, std::size_t e
  * @param dataset_count Number of vectors in the dataset.
  * @param queries_count Number of query vectors.
  * @param wanted_count Number of top matches required from each query.
+ * @tparam scalar_at Data type of the elements in the vectors.
  */
+template <typename scalar_at>
 void test_exact_search(std::size_t dataset_count, std::size_t queries_count, std::size_t wanted_count) {
     std::size_t dimensions = 32;
     metric_punned_t metric(dimensions, metric_kind_t::cos_k);
 
     std::random_device rd;
     std::mt19937 gen(rd());
     std::uniform_real_distribution<> dis(0.0, 1.0);
-    std::vector<float> dataset(dataset_count * dimensions);
-    std::generate(dataset.begin(), dataset.end(), [&] { return dis(gen); });
+    std::vector<scalar_at> dataset(dataset_count * dimensions);
+    std::generate(dataset.begin(), dataset.end(), [&] { return static_cast<scalar_at>(dis(gen)); });
 
     exact_search_t search;
     auto results = search(                                                        //
@@ -1099,6 +1101,7 @@ template <typename key_at, typename slot_at> void test_replacing_update() {
 int main(int, char**) {
     test_uint40();
     test_cosine<float, std::int64_t, uint40_t>(10, 10);
+    test_cosine<bf16_t, std::int64_t, uint40_t>(10, 10);
 
     // Test plugins, like K-Means clustering.
     {
@@ -1121,8 +1124,10 @@ int main(int, char**) {
     std::printf("Testing exact search\n");
     for (std::size_t dataset_count : {10, 100})
         for (std::size_t queries_count : {1, 10})
-            for (std::size_t wanted_count : {1, 5})
-                test_exact_search(dataset_count, queries_count, wanted_count);
+            for (std::size_t wanted_count : {1, 5}) {
+                test_exact_search<float>(dataset_count, queries_count, wanted_count);
+                test_exact_search<bf16_t>(dataset_count, queries_count, wanted_count);
+            }
 
     // Make sure the initializers and the algorithms can work with inadequately small values.
     // Be warned - this combinatorial explosion of tests produces close to __500'000__ tests!
@@ -1149,6 +1154,8 @@ int main(int, char**) {
             test_cosine<float, std::int64_t, slot32_t>(collection_size, dimensions);
             std::printf("- Indexing %zu vectors with cos: <float, std::int64_t, uint40_t> \n", collection_size);
             test_cosine<float, std::int64_t, uint40_t>(collection_size, dimensions);
+            std::printf("- Indexing %zu vectors with cos: <bf16, std::int64_t, uint40_t> \n", collection_size);
+            test_cosine<bf16_t, std::int64_t, uint40_t>(collection_size, dimensions);
         }
 
     // Test with binary vectors

include/usearch/index_dense.hpp

@@ -760,36 +760,42 @@ class index_dense_gt {
     add_result_t add(vector_key_t key, b1x8_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from.b1x8); }
     add_result_t add(vector_key_t key, i8_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from.i8); }
     add_result_t add(vector_key_t key, f16_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from.f16); }
+    add_result_t add(vector_key_t key, bf16_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from.bf16); }
     add_result_t add(vector_key_t key, f32_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from.f32); }
     add_result_t add(vector_key_t key, f64_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from.f64); }
 
     search_result_t search(b1x8_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, dummy_predicate_t {}, thread, exact, casts_.from.b1x8); }
     search_result_t search(i8_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, dummy_predicate_t {}, thread, exact, casts_.from.i8); }
     search_result_t search(f16_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, dummy_predicate_t {}, thread, exact, casts_.from.f16); }
+    search_result_t search(bf16_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, dummy_predicate_t {}, thread, exact, casts_.from.bf16); }
     search_result_t search(f32_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, dummy_predicate_t {}, thread, exact, casts_.from.f32); }
     search_result_t search(f64_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, dummy_predicate_t {}, thread, exact, casts_.from.f64); }
 
     template <typename predicate_at> search_result_t filtered_search(b1x8_t const* vector, std::size_t wanted, predicate_at&& predicate, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, std::forward<predicate_at>(predicate), thread, exact, casts_.from.b1x8); }
     template <typename predicate_at> search_result_t filtered_search(i8_t const* vector, std::size_t wanted, predicate_at&& predicate, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, std::forward<predicate_at>(predicate), thread, exact, casts_.from.i8); }
     template <typename predicate_at> search_result_t filtered_search(f16_t const* vector, std::size_t wanted, predicate_at&& predicate, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, std::forward<predicate_at>(predicate), thread, exact, casts_.from.f16); }
+    template <typename predicate_at> search_result_t filtered_search(bf16_t const* vector, std::size_t wanted, predicate_at&& predicate, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, std::forward<predicate_at>(predicate), thread, exact, casts_.from.bf16); }
     template <typename predicate_at> search_result_t filtered_search(f32_t const* vector, std::size_t wanted, predicate_at&& predicate, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, std::forward<predicate_at>(predicate), thread, exact, casts_.from.f32); }
     template <typename predicate_at> search_result_t filtered_search(f64_t const* vector, std::size_t wanted, predicate_at&& predicate, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, std::forward<predicate_at>(predicate), thread, exact, casts_.from.f64); }
 
     std::size_t get(vector_key_t key, b1x8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to.b1x8); }
     std::size_t get(vector_key_t key, i8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to.i8); }
     std::size_t get(vector_key_t key, f16_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to.f16); }
+    std::size_t get(vector_key_t key, bf16_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to.bf16); }
     std::size_t get(vector_key_t key, f32_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to.f32); }
     std::size_t get(vector_key_t key, f64_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to.f64); }
 
     cluster_result_t cluster(b1x8_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from.b1x8); }
     cluster_result_t cluster(i8_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from.i8); }
     cluster_result_t cluster(f16_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from.f16); }
+    cluster_result_t cluster(bf16_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from.bf16); }
     cluster_result_t cluster(f32_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from.f32); }
     cluster_result_t cluster(f64_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from.f64); }
 
     aggregated_distances_t distance_between(vector_key_t key, b1x8_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to.b1x8); }
     aggregated_distances_t distance_between(vector_key_t key, i8_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to.i8); }
     aggregated_distances_t distance_between(vector_key_t key, f16_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to.f16); }
+    aggregated_distances_t distance_between(vector_key_t key, bf16_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to.bf16); }
     aggregated_distances_t distance_between(vector_key_t key, f32_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to.f32); }
     aggregated_distances_t distance_between(vector_key_t key, f64_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to.f64); }
     // clang-format on

include/usearch/index_plugins.hpp

@@ -581,6 +581,11 @@ class bf16_bits_t {
         uint16_ = f32_to_bf16(v / bf16_to_f32(uint16_));
         return *this;
     }
+
+    inline bf16_bits_t& operator=(float v) noexcept {
+        uint16_ = f32_to_bf16(v);
+        return *this;
+    }
 };
 
 /**
@@ -1223,6 +1228,7 @@ struct casts_punned_t {
         cast_punned_t b1x8{};
         cast_punned_t i8{};
         cast_punned_t f16{};
+        cast_punned_t bf16{};
         cast_punned_t f32{};
         cast_punned_t f64{};
 
@@ -1231,7 +1237,7 @@ struct casts_punned_t {
             case scalar_kind_t::f64_k: return f64;
             case scalar_kind_t::f32_k: return f32;
             case scalar_kind_t::f16_k: return f16;
-            case scalar_kind_t::bf16_k: return f16;
+            case scalar_kind_t::bf16_k: return bf16;
             case scalar_kind_t::i8_k: return i8;
             case scalar_kind_t::b1x8_k: return b1x8;
             default: return nullptr;
@@ -1246,12 +1252,14 @@ struct casts_punned_t {
         result.from.b1x8 = &cast_gt<b1x8_t, scalar_at>::try_;
         result.from.i8 = &cast_gt<i8_t, scalar_at>::try_;
         result.from.f16 = &cast_gt<f16_t, scalar_at>::try_;
+        result.from.bf16 = &cast_gt<bf16_t, scalar_at>::try_;
         result.from.f32 = &cast_gt<f32_t, scalar_at>::try_;
         result.from.f64 = &cast_gt<f64_t, scalar_at>::try_;
 
         result.to.b1x8 = &cast_gt<scalar_at, b1x8_t>::try_;
         result.to.i8 = &cast_gt<scalar_at, i8_t>::try_;
         result.to.f16 = &cast_gt<scalar_at, f16_t>::try_;
+        result.to.bf16 = &cast_gt<scalar_at, bf16_t>::try_;
         result.to.f32 = &cast_gt<scalar_at, f32_t>::try_;
         result.to.f64 = &cast_gt<scalar_at, f64_t>::try_;