API fallback and python bindings refactoring

bindings/python/include/svs/python/conversion.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2023 Intel Corporation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+// pybind
+#include <pybind11/pybind11.h>
+
+namespace svs::python::conversion {
+void wrap(pybind11::module& m);
+} // namespace svs::python::conversion

bindings/python/include/svs/python/core.h

@@ -27,6 +27,14 @@
 #include "svs/lib/meta.h"
 #include "svs/lib/misc.h"
 
+#include "svs/fallback/fallback.h"
+
+#ifdef USE_PROPRIETARY
+
+#include "svs/fallback/fallback_python.h"
+
+#endif // USE_PROPRIETARY
+
 // pybind
 #include <pybind11/pybind11.h>
 
@@ -118,6 +126,22 @@ class UnspecializedGraphLoader {
 using DistanceL2 = svs::distance::DistanceL2;
 using DistanceIP = svs::distance::DistanceIP;
 
+/////
+///// LVQ
+/////
+
+// Compressors - online compression of existing data
+using LVQReloader = svs::quantization::lvq::Reload;
+using LVQ = svs::quantization::lvq::ProtoLVQLoader<Allocator>;
+
+/////
+///// LeanVec
+/////
+
+// Dimensionality reduction using LeanVec
+using LeanVecReloader = svs::leanvec::Reload;
+using LeanVec = svs::leanvec::ProtoLeanVecLoader<Allocator>;
+
 namespace core {
 void wrap(pybind11::module& m);
 } // namespace core

bindings/python/include/svs/python/dispatch.h

@@ -49,3 +49,83 @@ struct svs::lib::DispatchConverter<
         return To{object.context().get_directory()};
     }
 };
+
+template <
+    size_t Primary,
+    size_t Residual,
+    size_t Extent,
+    svs::quantization::lvq::LVQPackingStrategy Strategy>
+struct svs::lib::DispatchConverter<
+    svs::lib::SerializedObject,
+    svs::quantization::lvq::LVQLoader<
+        Primary,
+        Residual,
+        Extent,
+        Strategy,
+        svs::python::RebindAllocator<std::byte>>> {
+    using To = svs::quantization::lvq::LVQLoader<
+        Primary,
+        Residual,
+        Extent,
+        Strategy,
+        svs::python::RebindAllocator<std::byte>>;
+
+    using LVQStrategyDispatch = svs::quantization::lvq::LVQStrategyDispatch;
+
+    static int64_t match(const svs::lib::SerializedObject& object) {
+        // TODO: Use a LoadTable directly instead of forcing reparsing every time.
+        auto ex = svs::lib::try_load<svs::quantization::lvq::Matcher>(object);
+        if (!ex) {
+            return svs::lib::invalid_match;
+        }
+
+        return svs::quantization::lvq::overload_score<Primary, Residual, Extent, Strategy>(
+            ex.value(), LVQStrategyDispatch::Auto
+        );
+    }
+
+    static To convert(const svs::lib::SerializedObject& object) {
+        return To{
+            svs::quantization::lvq::Reload{std::move(object.context().get_directory())},
+            0,
+            svs::python::RebindAllocator<std::byte>()};
+    }
+};
+
+template <typename PrimaryKind, typename SecondaryKind, size_t LeanVecDims, size_t Extent>
+struct svs::lib::DispatchConverter<
+    svs::lib::SerializedObject,
+    svs::leanvec::LeanVecLoader<
+        PrimaryKind,
+        SecondaryKind,
+        LeanVecDims,
+        Extent,
+        svs::python::RebindAllocator<std::byte>>> {
+    using To = leanvec::LeanVecLoader<
+        PrimaryKind,
+        SecondaryKind,
+        LeanVecDims,
+        Extent,
+        svs::python::RebindAllocator<std::byte>>;
+
+    static int64_t match(const svs::lib::SerializedObject& object) {
+        // TODO: Use a LoadTable directly instead of forcing reparsing every time.
+        auto ex = svs::lib::try_load<svs::leanvec::Matcher>(object);
+        if (!ex) {
+            return svs::lib::invalid_match;
+        }
+
+        return svs::leanvec::
+            overload_score<PrimaryKind, SecondaryKind, LeanVecDims, Extent>(ex.value());
+    }
+
+    static To convert(const svs::lib::SerializedObject& object) {
+        return To{
+            leanvec::Reload{object.context().get_directory()},
+            LeanVecDims, // TODO: This is a hack for now. Since we're reloading, it doesn't
+                         // matter.
+            std::nullopt,
+            0,
+            svs::python::RebindAllocator<std::byte>()};
+    }
+};

bindings/python/include/svs/python/dynamic_vamana.h

@@ -33,5 +33,50 @@ template <typename F> void for_standard_specializations(F&& f) {
 #undef X
 }
 
+template <typename F> void for_compressed_specializations(F&& f) {
+    using Sequential = svs::quantization::lvq::Sequential;
+#define X(Dist, Primary, Residual, Strategy, N) \
+    f.template operator()<Dist, Primary, Residual, Strategy, N>()
+    // Sequential
+    X(DistanceL2, 4, 0, Sequential, Dynamic);
+    X(DistanceIP, 4, 0, Sequential, Dynamic);
+    X(DistanceL2, 4, 4, Sequential, Dynamic);
+    X(DistanceIP, 4, 4, Sequential, Dynamic);
+    X(DistanceL2, 4, 8, Sequential, Dynamic);
+    X(DistanceIP, 4, 8, Sequential, Dynamic);
+    X(DistanceL2, 8, 0, Sequential, Dynamic);
+    X(DistanceIP, 8, 0, Sequential, Dynamic);
+
+    // Turbo
+    using Turbo16x8 = svs::quantization::lvq::Turbo<16, 8>;
+    X(DistanceL2, 4, 0, Turbo16x8, Dynamic);
+    X(DistanceIP, 4, 0, Turbo16x8, Dynamic);
+    X(DistanceL2, 4, 4, Turbo16x8, Dynamic);
+    X(DistanceIP, 4, 4, Turbo16x8, Dynamic);
+    X(DistanceL2, 4, 8, Turbo16x8, Dynamic);
+    X(DistanceIP, 4, 8, Turbo16x8, Dynamic);
+#undef X
+}
+
+template <typename F> void for_leanvec_specializations(F&& f) {
+#define X(Dist, Primary, Secondary, L, N) \
+    f.template operator()<Dist, Primary, Secondary, L, N>()
+    X(DistanceL2, svs::Float16, svs::Float16, Dynamic, Dynamic);
+    X(DistanceIP, svs::Float16, svs::Float16, Dynamic, Dynamic);
+
+    X(DistanceL2, svs::leanvec::UsingLVQ<8>, svs::Float16, Dynamic, Dynamic);
+    X(DistanceIP, svs::leanvec::UsingLVQ<8>, svs::Float16, Dynamic, Dynamic);
+
+    X(DistanceL2, svs::leanvec::UsingLVQ<8>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic);
+    X(DistanceIP, svs::leanvec::UsingLVQ<8>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic);
+
+    X(DistanceL2, svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic);
+    X(DistanceIP, svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic);
+
+    X(DistanceL2, svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<4>, Dynamic, Dynamic);
+    X(DistanceIP, svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<4>, Dynamic, Dynamic);
+#undef X
+}
+
 void wrap(pybind11::module& m);
 } // namespace svs::python::dynamic_vamana

bindings/python/include/svs/python/vamana.h

@@ -89,6 +89,126 @@ template <typename F> void for_standard_specializations(F&& f) {
 #undef XN
 #undef X
 }
+
+// Compressed search specializations.
+// Pattern:
+// DistanceType, Primary, Residual, Dimensionality, Strategy, EnableBuild
+#define X(Dist, P, R, N, S, B) f.template operator()<Dist, P, R, N, S, B>()
+template <typename F> void lvq_specialize_4x0(const F& f) {
+    using Sequential = svs::quantization::lvq::Sequential;
+    using Turbo = svs::quantization::lvq::Turbo<16, 8>;
+
+    // Sequential
+    X(DistanceL2, 4, 0, Dynamic, Sequential, true);
+    X(DistanceIP, 4, 0, Dynamic, Sequential, true);
+    X(DistanceCosineSimilarity, 4, 0, Dynamic, Sequential, true);
+    // Turbo
+    X(DistanceL2, 4, 0, Dynamic, Turbo, true);
+    X(DistanceIP, 4, 0, Dynamic, Turbo, true);
+    X(DistanceCosineSimilarity, 4, 0, Dynamic, Turbo, true);
+}
+
+template <typename F> void lvq_specialize_4x4(const F& f) {
+    using Sequential = svs::quantization::lvq::Sequential;
+    using Turbo = svs::quantization::lvq::Turbo<16, 8>;
+
+    // Sequential
+    X(DistanceL2, 4, 4, Dynamic, Sequential, true);
+    X(DistanceIP, 4, 4, Dynamic, Sequential, true);
+    X(DistanceCosineSimilarity, 4, 4, Dynamic, Sequential, true);
+    // Turbo
+    X(DistanceL2, 4, 4, Dynamic, Turbo, true);
+    X(DistanceIP, 4, 4, Dynamic, Turbo, true);
+    X(DistanceCosineSimilarity, 4, 4, Dynamic, Turbo, true);
+}
+
+template <typename F> void lvq_specialize_4x8(const F& f) {
+    using Sequential = svs::quantization::lvq::Sequential;
+    using Turbo = svs::quantization::lvq::Turbo<16, 8>;
+
+    // Sequential
+    X(DistanceL2, 4, 8, Dynamic, Sequential, true);
+    X(DistanceIP, 4, 8, Dynamic, Sequential, true);
+    X(DistanceCosineSimilarity, 4, 8, Dynamic, Sequential, true);
+    // Turbo
+    X(DistanceL2, 4, 8, Dynamic, Turbo, true);
+    X(DistanceIP, 4, 8, Dynamic, Turbo, true);
+    X(DistanceCosineSimilarity, 4, 8, Dynamic, Turbo, true);
+}
+
+template <typename F> void lvq_specialize_8x0(const F& f) {
+    using Sequential = svs::quantization::lvq::Sequential;
+    using Turbo = svs::quantization::lvq::Turbo<16, 4>;
+
+    // Sequential
+    X(DistanceL2, 8, 0, Dynamic, Sequential, true);
+    X(DistanceIP, 8, 0, Dynamic, Sequential, true);
+    X(DistanceCosineSimilarity, 8, 0, Dynamic, Sequential, true);
+    // Turbo
+    X(DistanceL2, 8, 0, Dynamic, Turbo, true);
+    X(DistanceIP, 8, 0, Dynamic, Turbo, true);
+    X(DistanceCosineSimilarity, 8, 0, Dynamic, Turbo, true);
+}
+
+template <typename F> void lvq_specialize_8x8(const F& f) {
+    using Sequential = svs::quantization::lvq::Sequential;
+    X(DistanceL2, 8, 8, Dynamic, Sequential, false);
+    X(DistanceIP, 8, 8, Dynamic, Sequential, false);
+    X(DistanceCosineSimilarity, 8, 8, Dynamic, Sequential, false);
+}
+
+template <typename F> void compressed_specializations(F&& f) {
+    lvq_specialize_4x0(f);
+    lvq_specialize_4x4(f);
+    lvq_specialize_4x8(f);
+    lvq_specialize_8x0(f);
+    lvq_specialize_8x8(f);
+}
+#undef X
+
+// LeanVec specializations.
+// Pattern:
+// Primary, Secondary, LeanVec Dimensionality, Dimensionality, DistanceType
+#define X(P, S, L, N, D) f.template operator()<P, S, L, N, D>()
+template <typename F> void leanvec_specialize_unc_unc(const F& f) {
+    X(float, float, Dynamic, Dynamic, DistanceL2);
+    X(float, float, Dynamic, Dynamic, DistanceIP);
+    X(float, float, Dynamic, Dynamic, DistanceCosineSimilarity);
+
+    X(svs::Float16, svs::Float16, Dynamic, Dynamic, DistanceL2);
+    X(svs::Float16, svs::Float16, Dynamic, Dynamic, DistanceIP);
+    X(svs::Float16, svs::Float16, Dynamic, Dynamic, DistanceCosineSimilarity);
+}
+
+template <typename F> void leanvec_specialize_lvq_unc(const F& f) {
+    X(svs::leanvec::UsingLVQ<8>, svs::Float16, Dynamic, Dynamic, DistanceL2);
+    X(svs::leanvec::UsingLVQ<8>, svs::Float16, Dynamic, Dynamic, DistanceIP);
+    X(svs::leanvec::UsingLVQ<8>, svs::Float16, Dynamic, Dynamic, DistanceCosineSimilarity);
+}
+
+template <typename F> void leanvec_specialize_lvq_lvq(const F& f) {
+    // clang-format off
+    X(svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<4>, Dynamic, Dynamic, DistanceL2);
+    X(svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<4>, Dynamic, Dynamic, DistanceIP);
+    X(svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<4>, Dynamic, Dynamic, DistanceCosineSimilarity);
+
+    X(svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic, DistanceL2);
+    X(svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic, DistanceIP);
+    X(svs::leanvec::UsingLVQ<4>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic, DistanceCosineSimilarity);
+
+    X(svs::leanvec::UsingLVQ<8>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic, DistanceL2);
+    X(svs::leanvec::UsingLVQ<8>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic, DistanceIP);
+    X(svs::leanvec::UsingLVQ<8>, svs::leanvec::UsingLVQ<8>, Dynamic, Dynamic, DistanceCosineSimilarity);
+    // clang-format on
+}
+
+template <typename F> void leanvec_specializations(F&& f) {
+    leanvec_specialize_unc_unc(f);
+    leanvec_specialize_lvq_unc(f);
+    leanvec_specialize_lvq_lvq(f);
+}
+#undef X
+
 } // namespace vamana_specializations
 
 namespace vamana {