Add fp16 support to SuperKMeans (#5692)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5692

Accept packed IEEE fp16 rows through `SuperKMeans::train_ex`, widening bounded blocks directly into the existing fp32 rotated workspace. Route raw Float16 IVF training through SuperKMeans while preserving codec, mode, and metric boundaries, and expose the typed input through the low-level Python wrapper.

Add direct, subsampling, multi-block, spherical, validation, IVF, codec-boundary, metric-boundary, and Python parity coverage.

Reviewed By: mnorris11

Differential Revision: D122781442

fbshipit-source-id: fa31ef48ed32b6c7fedfad2dffe5dcef794cb618

faiss/CMakeLists.txt

@@ -636,6 +636,11 @@ endif()
 if(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)" AND NOT MSVC)
   target_compile_definitions(faiss PRIVATE COMPILE_SIMD_ARM_NEON)
   target_sources(faiss PRIVATE ${FAISS_SIMD_NEON_SRC})
+  # GCC requires +fp16 for the vector conversion intrinsics in this source.
+  set_source_files_properties(utils/simd_impl/fp16_kernels_neon.cpp
+    TARGET_DIRECTORY faiss
+    PROPERTIES COMPILE_OPTIONS "-march=armv8.2-a+fp16"
+  )
 endif()
 
 # RVV is the baseline SIMD on rv64 builds; the RVV translation units require

faiss/IndexIVF.cpp

@@ -71,7 +71,6 @@ void train_q1_impl(
     const size_t nlist = q1.nlist;
     const ClusteringParameters& cp = q1.cp;
     Index* clustering_index = q1.clustering_index;
-    const bool float32 = !codec && numeric_type == NumericType::Float32;
 
     // k-means on the training rows, with `assigner` as assignment index
     auto train_clustering = [&](Clustering& clus, Index& assigner) {
@@ -117,8 +116,15 @@ void train_q1_impl(
                 "cp.use_super_kmeans is incompatible with a user-provided "
                 "clustering_index: SuperKMeans assigns with its own index");
         FAISS_THROW_IF_MSG(
-                cp.use_super_kmeans && !float32,
-                "SuperKMeans requires fp32 training data");
+                cp.use_super_kmeans && codec,
+                "SuperKMeans does not support encoded training data");
+        FAISS_THROW_IF_MSG(
+                cp.use_super_kmeans &&
+                        !(metric_type == METRIC_L2 ||
+                          (metric_type == METRIC_INNER_PRODUCT &&
+                           cp.spherical)),
+                "SuperKMeans requires L2 or spherical inner-product "
+                "clustering");
 
         quantizer->reset();
         if (cp.use_super_kmeans) {
@@ -126,7 +132,7 @@ void train_q1_impl(
             static_cast<ClusteringParameters&>(super_cp) = cp;
             SuperKMeans clus(
                     static_cast<int>(d), static_cast<int>(nlist), super_cp);
-            clus.train(n, reinterpret_cast<const float*>(x));
+            clus.train_ex(n, x, numeric_type);
             quantizer->add(nlist, clus.centroids.data());
         } else {
             Clustering clus(static_cast<int>(d), static_cast<int>(nlist), cp);

faiss/IndexIVF.h

@@ -57,7 +57,8 @@ struct Level1Quantizer {
      *
      * Float16 runs k-means natively on the fp16 rows (see
      * Clustering::train_ex). With quantizer_trains_alone == 1 the data is
-     * passed on to quantizer->train_ex(). SuperKMeans requires Float32.
+     * passed on to quantizer->train_ex(). SuperKMeans supports raw Float32
+     * and Float16 input, but not vectors decoded through a codec.
      */
     void train_q1_ex(
             size_t n,

faiss/SuperKMeans.cpp

@@ -276,10 +276,13 @@ void setup_train_state(
         int d,
         int k,
         idx_t n,
-        const float* x) {
-    const size_t line_size = sizeof(float) * static_cast<size_t>(d);
+        const void* x,
+        NumericType numeric_type) {
+    const bool fp16 = numeric_type == NumericType::Float16;
+    const size_t line_size =
+            (fp16 ? sizeof(uint16_t) : sizeof(float)) * static_cast<size_t>(d);
     idx_t nx = n;
-    const uint8_t* x_bytes = reinterpret_cast<const uint8_t*>(x);
+    const uint8_t* x_bytes = static_cast<const uint8_t*>(x);
     std::unique_ptr<uint8_t[]> sampled_x_owner;
     if (static_cast<size_t>(nx) >
         static_cast<size_t>(k) * cp.max_points_per_centroid) {
@@ -298,7 +301,6 @@ void setup_train_state(
         sampled_x_owner.reset(x_new);
         x_bytes = x_new;
     }
-    const float* x_sampled = reinterpret_cast<const float*>(x_bytes);
 
     FAISS_THROW_IF_NOT_MSG(
             nx <= static_cast<idx_t>(std::numeric_limits<int>::max()),
@@ -315,7 +317,27 @@ void setup_train_state(
     }
 
     state.X_tilde.resize(static_cast<size_t>(state.n) * d);
-    state.R->apply_noalloc(state.n, x_sampled, state.X_tilde.data());
+    if (!fp16) {
+        state.R->apply_noalloc(
+                state.n,
+                reinterpret_cast<const float*>(x_bytes),
+                state.X_tilde.data());
+    } else {
+        const size_t block_size =
+                std::min(static_cast<size_t>(state.n), cp.decode_block_size);
+        std::vector<float> decoded(block_size * d);
+        const uint16_t* x16 = reinterpret_cast<const uint16_t*>(x_bytes);
+        for (size_t i0 = 0; i0 < static_cast<size_t>(state.n);
+             i0 += block_size) {
+            const size_t i1 =
+                    std::min(static_cast<size_t>(state.n), i0 + block_size);
+            fp16_to_fp32((i1 - i0) * d, x16 + i0 * d, decoded.data());
+            state.R->apply_noalloc(
+                    static_cast<idx_t>(i1 - i0),
+                    decoded.data(),
+                    state.X_tilde.data() + i0 * d);
+        }
+    }
     sampled_x_owner.reset();
 
     // Forgy init: pick k random rows from the rotated pool as initial
@@ -410,15 +432,36 @@ SuperKMeans::SuperKMeans(int d, int k, const SuperKMeansParameters& cp_in)
 }
 
 void SuperKMeans::train(idx_t n, const float* x) {
+    train_ex(n, x, NumericType::Float32);
+}
+
+void SuperKMeans::train_ex(idx_t n, const void* x, NumericType numeric_type) {
+    FAISS_THROW_IF_NOT_MSG(
+            numeric_type == NumericType::Float32 ||
+                    numeric_type == NumericType::Float16,
+            "SuperKMeans::train_ex: unsupported numeric type");
     FAISS_THROW_IF_NOT_MSG(n > 0, "SuperKMeans: n must be positive");
     FAISS_THROW_IF_MSG(x == nullptr, "SuperKMeans: x must not be null");
     FAISS_THROW_IF_NOT_MSG(
             n >= static_cast<idx_t>(k), "SuperKMeans: n must be >= k");
+    const bool fp16 = numeric_type == NumericType::Float16;
+    FAISS_THROW_IF_NOT_MSG(
+            !fp16 || cp.decode_block_size > 0,
+            "SuperKMeans: decode_block_size must be positive");
     if (cp.check_input_data_for_NaNs) {
-        for (size_t i = 0; i < static_cast<size_t>(n) * d; i++) {
+        const size_t num_values = static_cast<size_t>(n) * d;
+        if (fp16) {
             FAISS_THROW_IF_NOT_MSG(
-                    std::isfinite(x[i]),
+                    detail::fp16_all_finite(
+                            num_values, static_cast<const uint16_t*>(x)),
                     "SuperKMeans: input contains NaN's or Inf's");
+        } else {
+            const float* x32 = static_cast<const float*>(x);
+            for (size_t i = 0; i < num_values; i++) {
+                FAISS_THROW_IF_NOT_MSG(
+                        std::isfinite(x32[i]),
+                        "SuperKMeans: input contains NaN's or Inf's");
+            }
         }
     }
     if (cp.verbose && n < static_cast<idx_t>(k) * cp.min_points_per_centroid) {
@@ -434,7 +477,7 @@ void SuperKMeans::train(idx_t n, const float* x) {
     std::vector<int64_t> labels64;
     SuperKMeansAssignScratch assign_scratch;
     std::vector<float> hassign;
-    setup_train_state(state, labels64, hassign, cp, d, k, n, x);
+    setup_train_state(state, labels64, hassign, cp, d, k, n, x, numeric_type);
 
     iteration_stats.clear();
     iteration_stats.reserve(cp.niter);

faiss/SuperKMeans.h

@@ -89,11 +89,16 @@ struct SuperKMeans {
 
     SuperKMeans(int d, int k, const SuperKMeansParameters& cp = {});
 
-    /// Train on `n` row-major vectors of dimension `d`. Honors the applicable
-    /// ClusteringParameters fields (niter, seed, verbose,
-    /// max_points_per_centroid, use_faster_subsampling,
-    /// check_input_data_for_NaNs).
+    /// Train on `n` row-major fp32 vectors of dimension `d`.
     void train(idx_t n, const float* x);
+
+    /** Train on vectors of the given numeric type.
+     *
+     * Float32 is the same as train(). Float16 takes packed IEEE binary16
+     * vectors and widens bounded blocks while rotating them into the fp32
+     * training workspace. Centroids and all iteration state remain fp32.
+     */
+    void train_ex(idx_t n, const void* x, NumericType numeric_type);
 };
 
 /// Reusable scratch for super_kmeans_assign_iteration; pass one instance across

faiss/python/class_wrappers.py

@@ -141,18 +141,23 @@ def handle_Clustering(the_class):
 
 def handle_SuperKMeans(the_class):
 
-    def replacement_train(self, x):
+    def replacement_train(self, x, *, numeric_type=faiss.Float32):
         """Perform SuperKMeans clustering on a set of vectors.
 
         Parameters
         ----------
         x : array_like
-            Training vectors, shape (n, self.d). `dtype` must be float32.
+            Training vectors, shape (n, self.d).
+        numeric_type : faiss.NumericType, optional
+            Numeric type of the input vectors.
         """
         n, d = x.shape
         assert d == self.d
-        x = np.ascontiguousarray(x, dtype="float32")
-        self.train_c(n, swig_ptr(x))
+        x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type))
+        if numeric_type == faiss.Float32:
+            self.train_c(n, swig_ptr(x))
+        else:
+            self.train_ex(n, swig_ptr(x), numeric_type)
 
     replace_method(the_class, "train", replacement_train)
 

tests/test_ivf_index.cpp

@@ -789,6 +789,62 @@ TEST(IVF, train_float16_matches_float32_on_rounded_input) {
     EXPECT_EQ(half_centroids, float_centroids);
 }
 
+TEST(IVF, train_float16_with_super_kmeans_matches_float32) {
+    constexpr int d = 32;
+    constexpr int n = 96;
+    constexpr int nlist = 4;
+
+    std::vector<float> input(n * d);
+    for (int i = 0; i < n; ++i) {
+        const float center = static_cast<float>(i % nlist) * 8.0f;
+        for (int j = 0; j < d; ++j) {
+            input[static_cast<size_t>(i) * d + j] = center +
+                    static_cast<float>((i * 17 + j * 13) % 31) / 100.0f;
+        }
+    }
+    auto encoded = encode_fp16(input);
+    auto rounded = decode_fp16(encoded);
+
+    faiss::IndexFlatL2 float_quantizer(d);
+    faiss::IndexIVFFlat float_index(&float_quantizer, d, nlist);
+    float_index.cp.seed = 1234;
+    float_index.cp.niter = 3;
+    float_index.cp.min_points_per_centroid = 1;
+    float_index.cp.use_super_kmeans = true;
+    float_index.train(n, rounded.data());
+
+    faiss::IndexFlatL2 half_quantizer(d);
+    faiss::IndexIVFFlat half_index(&half_quantizer, d, nlist);
+    half_index.cp = float_index.cp;
+    half_index.cp.decode_block_size = n;
+    half_index.train_ex(n, encoded.data(), faiss::NumericType::Float16);
+
+    std::vector<float> float_centroids(nlist * d);
+    std::vector<float> half_centroids(nlist * d);
+    float_quantizer.reconstruct_n(0, nlist, float_centroids.data());
+    half_quantizer.reconstruct_n(0, nlist, half_centroids.data());
+    EXPECT_EQ(half_centroids, float_centroids);
+
+    TrackingFp16Codec codec(d);
+    faiss::IndexFlatL2 encoded_quantizer(d);
+    faiss::IndexIVFFlat encoded_index(&encoded_quantizer, d, nlist);
+    encoded_index.cp = float_index.cp;
+    EXPECT_THROW(
+            encoded_index.train_encoded(
+                    n,
+                    reinterpret_cast<const uint8_t*>(encoded.data()),
+                    &codec),
+            faiss::FaissException);
+
+    faiss::IndexFlatIP ip_quantizer(d);
+    faiss::IndexIVFFlat ip_index(
+            &ip_quantizer, d, nlist, faiss::METRIC_INNER_PRODUCT);
+    ip_index.cp = float_index.cp;
+    EXPECT_THROW(
+            ip_index.train_ex(n, encoded.data(), faiss::NumericType::Float16),
+            faiss::FaissException);
+}
+
 TEST(IVF, encoded_training_rejects_non_finite_values) {
     constexpr int d = 2;
     constexpr int n = 4;

tests/test_super_kmeans.py

@@ -17,6 +17,34 @@ class SuperKMeansTest(unittest.TestCase):
         with self.assertRaises(RuntimeError):
             sc.train(x)
 
+    def test_train_float16_matches_rounded_float32(self):
+        d, k, n = 32, 4, 96
+        rows = np.arange(n, dtype="float32")[:, None]
+        dims = np.arange(d, dtype="float32")[None, :]
+        values = (rows % k) * 8.0 + ((rows * 17 + dims * 13) % 31) / 100.0
+        x16 = values.astype("float16")
+        rounded = x16.astype("float32")
+
+        p = faiss.SuperKMeansParameters()
+        p.seed = 1234
+        p.niter = 3
+        p.min_points_per_centroid = 1
+        p.decode_block_size = n
+
+        expected = faiss.SuperKMeans(d, k, p)
+        expected.train(rounded)
+        actual = faiss.SuperKMeans(d, k, p)
+        actual.train(x16, numeric_type=faiss.Float16)
+
+        np.testing.assert_array_equal(
+            faiss.vector_to_array(actual.centroids),
+            faiss.vector_to_array(expected.centroids),
+        )
+        np.testing.assert_array_equal(
+            faiss.vector_to_array(actual.gemm_pruning_rates),
+            faiss.vector_to_array(expected.gemm_pruning_rates),
+        )
+
     def test_objective_decreases_monotonically(self):
         d, k, n = 32, 8, 1000
         x = SyntheticDataset(d, n, 0, 0).get_train()

tests/test_super_kmeans_foundations.cpp

@@ -18,6 +18,7 @@
 #include <faiss/impl/AdSampling.h>
 #include <faiss/impl/PdxLayout.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/fp16.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/simd_impl/super_kmeans_dispatch.h>
 #include <faiss/utils/simd_impl/super_kmeans_kernels.h>
@@ -495,3 +496,164 @@ TEST(SuperKMeansAssignIteration, MultipleTilesMatchBruteForce) {
     }
     EXPECT_GE(matches, n - 10); // <= 5% ADSampling false-prune tolerance
 }
+
+namespace {
+
+std::vector<uint16_t> encode_fp16_values(const std::vector<float>& values) {
+    std::vector<uint16_t> encoded(values.size());
+    for (size_t i = 0; i < values.size(); ++i) {
+        encoded[i] = faiss::encode_fp16(values[i]);
+    }
+    return encoded;
+}
+
+std::vector<float> decode_fp16_values(const std::vector<uint16_t>& values) {
+    std::vector<float> decoded(values.size());
+    for (size_t i = 0; i < values.size(); ++i) {
+        decoded[i] = faiss::decode_fp16(values[i]);
+    }
+    return decoded;
+}
+
+std::vector<float> make_super_kmeans_training_data(int n, int d) {
+    std::vector<float> x(static_cast<size_t>(n) * d);
+    for (int i = 0; i < n; ++i) {
+        const float center = static_cast<float>(i % 4) * 8.0f;
+        for (int j = 0; j < d; ++j) {
+            x[static_cast<size_t>(i) * d + j] = center +
+                    static_cast<float>((i * 17 + j * 13) % 31) / 100.0f;
+        }
+    }
+    return x;
+}
+
+void expect_same_super_kmeans_result(
+        const faiss::SuperKMeans& actual,
+        const faiss::SuperKMeans& expected) {
+    EXPECT_EQ(actual.centroids, expected.centroids);
+    ASSERT_EQ(actual.iteration_stats.size(), expected.iteration_stats.size());
+    for (size_t i = 0; i < actual.iteration_stats.size(); ++i) {
+        EXPECT_EQ(
+                actual.iteration_stats[i].obj, expected.iteration_stats[i].obj);
+        EXPECT_EQ(
+                actual.iteration_stats[i].nsplit,
+                expected.iteration_stats[i].nsplit);
+    }
+    EXPECT_EQ(actual.gemm_pruning_rates, expected.gemm_pruning_rates);
+}
+
+faiss::SuperKMeansParameters small_super_kmeans_params() {
+    faiss::SuperKMeansParameters cp;
+    cp.seed = 1234;
+    cp.niter = 3;
+    cp.min_points_per_centroid = 1;
+    return cp;
+}
+
+} // namespace
+
+TEST(SuperKMeans, TrainExFloat16SingleBlockMatchesRoundedFloat32) {
+    constexpr int n = 96, d = 32, k = 4;
+    auto encoded = encode_fp16_values(make_super_kmeans_training_data(n, d));
+    auto rounded = decode_fp16_values(encoded);
+    auto cp = small_super_kmeans_params();
+    cp.decode_block_size = n;
+
+    faiss::SuperKMeans expected(d, k, cp);
+    expected.train(n, rounded.data());
+
+    faiss::SuperKMeans float32_ex(d, k, cp);
+    float32_ex.train_ex(n, rounded.data(), faiss::NumericType::Float32);
+    expect_same_super_kmeans_result(float32_ex, expected);
+
+    faiss::SuperKMeans actual(d, k, cp);
+    actual.train_ex(n, encoded.data(), faiss::NumericType::Float16);
+    expect_same_super_kmeans_result(actual, expected);
+}
+
+TEST(SuperKMeans, TrainExFloat16SubsamplingMatchesRoundedFloat32) {
+    constexpr int n = 129, d = 32, k = 4;
+    auto encoded = encode_fp16_values(make_super_kmeans_training_data(n, d));
+    auto rounded = decode_fp16_values(encoded);
+    auto cp = small_super_kmeans_params();
+    cp.max_points_per_centroid = 16;
+    cp.decode_block_size = n;
+
+    faiss::SuperKMeans expected(d, k, cp);
+    expected.train(n, rounded.data());
+    faiss::SuperKMeans actual(d, k, cp);
+    actual.train_ex(n, encoded.data(), faiss::NumericType::Float16);
+
+    expect_same_super_kmeans_result(actual, expected);
+}
+
+TEST(SuperKMeans, TrainExFloat16MultiBlockMatchesRoundedFloat32) {
+    constexpr int n = 129, d = 32, k = 4;
+    auto encoded = encode_fp16_values(make_super_kmeans_training_data(n, d));
+    auto rounded = decode_fp16_values(encoded);
+    auto cp = small_super_kmeans_params();
+    cp.decode_block_size = 64;
+
+    faiss::SuperKMeans expected(d, k, cp);
+    expected.train(n, rounded.data());
+    faiss::SuperKMeans actual(d, k, cp);
+    actual.train_ex(n, encoded.data(), faiss::NumericType::Float16);
+
+    ASSERT_EQ(actual.centroids.size(), expected.centroids.size());
+    for (size_t i = 0; i < actual.centroids.size(); ++i) {
+        EXPECT_NEAR(actual.centroids[i], expected.centroids[i], 1e-4f);
+    }
+    ASSERT_EQ(actual.iteration_stats.size(), expected.iteration_stats.size());
+    for (size_t i = 0; i < actual.iteration_stats.size(); ++i) {
+        EXPECT_NEAR(
+                actual.iteration_stats[i].obj,
+                expected.iteration_stats[i].obj,
+                1e-3f);
+        EXPECT_EQ(
+                actual.iteration_stats[i].nsplit,
+                expected.iteration_stats[i].nsplit);
+    }
+
+    cp.spherical = true;
+    faiss::SuperKMeans spherical_expected(d, k, cp);
+    spherical_expected.train(n, rounded.data());
+    faiss::SuperKMeans spherical_actual(d, k, cp);
+    spherical_actual.train_ex(n, encoded.data(), faiss::NumericType::Float16);
+    expect_same_super_kmeans_result(spherical_actual, spherical_expected);
+    for (int i = 0; i < k; ++i) {
+        EXPECT_NEAR(
+                faiss::fvec_norm_L2sqr(
+                        spherical_actual.centroids.data() + i * d, d),
+                1.0f,
+                1e-5f);
+    }
+}
+
+TEST(SuperKMeans, TrainExFloat16ValidatesInput) {
+    constexpr int n = 8, d = 32, k = 2;
+    auto encoded = encode_fp16_values(make_super_kmeans_training_data(n, d));
+    auto cp = small_super_kmeans_params();
+    cp.niter = 1;
+
+    for (uint16_t bad :
+         {uint16_t(0x7e00), uint16_t(0x7c00), uint16_t(0xfc00)}) {
+        auto invalid = encoded;
+        invalid.back() = bad;
+        faiss::SuperKMeans clustering(d, k, cp);
+        EXPECT_THROW(
+                clustering.train_ex(
+                        n, invalid.data(), faiss::NumericType::Float16),
+                faiss::FaissException);
+    }
+
+    cp.decode_block_size = 0;
+    faiss::SuperKMeans zero_block(d, k, cp);
+    EXPECT_THROW(
+            zero_block.train_ex(n, encoded.data(), faiss::NumericType::Float16),
+            faiss::FaissException);
+
+    faiss::SuperKMeans unsupported(d, k, small_super_kmeans_params());
+    EXPECT_THROW(
+            unsupported.train_ex(n, encoded.data(), faiss::NumericType::UInt8),
+            faiss::FaissException);
+}