From 9adf700efe300a66da448482006181053c58d0a8 Mon Sep 17 00:00:00 2001 From: kgeg401 Date: Sat, 14 Feb 2026 13:08:24 -0500 Subject: [PATCH 1/3] feat: enable one2many batch distance for ip and l2 --- src/core/metric/euclidean_metric.cc | 8 +-- src/core/metric/inner_product_metric.cc | 14 ++-- tests/core/metric/euclidean_metric_test.cc | 64 ++++++++++++++++++ .../core/metric/inner_product_metric_test.cc | 66 ++++++++++++++++++- 4 files changed, 140 insertions(+), 12 deletions(-) diff --git a/src/core/metric/euclidean_metric.cc b/src/core/metric/euclidean_metric.cc index a1a8d598..d679c753 100644 --- a/src/core/metric/euclidean_metric.cc +++ b/src/core/metric/euclidean_metric.cc @@ -868,22 +868,22 @@ class SquaredEuclideanMetric : public IndexMetric { case IndexMeta::DataType::DT_FP16: return reinterpret_cast( ailego::BaseDistance::ComputeBatch); + ailego::Float16, 12, 2>::ComputeBatch); case IndexMeta::DataType::DT_FP32: return reinterpret_cast( ailego::BaseDistance::ComputeBatch); + 12, 2>::ComputeBatch); case IndexMeta::DataType::DT_INT8: return reinterpret_cast( ailego::BaseDistance::ComputeBatch); + 12, 2>::ComputeBatch); case IndexMeta::DataType::DT_INT4: return reinterpret_cast( ailego::BaseDistance::ComputeBatch); + uint8_t, 12, 2>::ComputeBatch); default: return nullptr; diff --git a/src/core/metric/inner_product_metric.cc b/src/core/metric/inner_product_metric.cc index 8ef0a11b..e4a609c8 100644 --- a/src/core/metric/inner_product_metric.cc +++ b/src/core/metric/inner_product_metric.cc @@ -354,20 +354,20 @@ class InnerProductMetric : public IndexMetric { switch (data_type_) { case IndexMeta::DataType::DT_FP32: return reinterpret_cast( - ailego::BaseDistance::ComputeBatch); + ailego::BaseDistance::ComputeBatch); case IndexMeta::DataType::DT_FP16: return reinterpret_cast( ailego::BaseDistance::ComputeBatch); + ailego::Float16, 12, 2>::ComputeBatch); case IndexMeta::DataType::DT_INT8: return reinterpret_cast( - ailego::BaseDistance::ComputeBatch); + ailego::BaseDistance::ComputeBatch); case IndexMeta::DataType::DT_INT4: return reinterpret_cast( - ailego::BaseDistance::ComputeBatch); + ailego::BaseDistance::ComputeBatch); default: return nullptr; } diff --git a/tests/core/metric/euclidean_metric_test.cc b/tests/core/metric/euclidean_metric_test.cc index c0c3a361..0dda0172 100644 --- a/tests/core/metric/euclidean_metric_test.cc +++ b/tests/core/metric/euclidean_metric_test.cc @@ -12,12 +12,40 @@ // See the License for the specific language governing permissions and // limitations under the License. #include +#include #include #include "zvec/core/framework/index_factory.h" using namespace zvec; using namespace zvec::core; +template +void CheckBatchDistanceMatchesSingle(const IndexMetric::Pointer &metric, + const std::vector> &vecs, + const std::vector &query, size_t dim, + float tolerance) { + auto single_distance = metric->distance(); + auto batch_distance = metric->batch_distance(); + ASSERT_TRUE(single_distance); + ASSERT_TRUE(batch_distance); + + std::vector ptrs; + ptrs.reserve(vecs.size()); + for (const auto &vec : vecs) { + ptrs.push_back(vec.data()); + } + + std::vector batch_result(vecs.size(), 0.0f); + batch_distance(ptrs.data(), query.data(), vecs.size(), dim, + batch_result.data()); + + for (size_t i = 0; i < vecs.size(); ++i) { + float single_result = 0.0f; + single_distance(vecs[i].data(), query.data(), dim, &single_result); + EXPECT_NEAR(single_result, batch_result[i], tolerance); + } +} + TEST(SquaredEuclideanMetric, General) { auto metric = IndexFactory::CreateMetric("SquaredEuclidean"); EXPECT_TRUE(metric); @@ -81,6 +109,42 @@ TEST(SquaredEuclideanMetric, General) { EXPECT_FLOAT_EQ(1.0f, result); } +TEST(SquaredEuclideanMetric, BatchDistanceMatchesSingleFp32) { + auto metric = IndexFactory::CreateMetric("SquaredEuclidean"); + ASSERT_TRUE(metric); + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, 8); + ASSERT_EQ(0, metric->init(meta, ailego::Params())); + + std::vector> vecs{ + {1.0f, 2.0f, 3.0f, 4.0f, 1.5f, -2.0f, 0.5f, -1.0f}, + {0.0f, -1.0f, 1.0f, 2.0f, -0.5f, 4.0f, -2.5f, 3.0f}, + {-3.0f, 2.5f, 1.5f, 0.0f, 3.0f, -1.0f, 2.0f, 1.0f}, + }; + std::vector query{2.0f, -1.0f, 0.5f, 3.0f, 1.0f, -2.0f, 4.0f, -1.5f}; + + CheckBatchDistanceMatchesSingle(metric, vecs, query, 8, 1e-5f); +} + +TEST(SquaredEuclideanMetric, BatchDistanceMatchesSingleInt8) { + auto metric = IndexFactory::CreateMetric("SquaredEuclidean"); + ASSERT_TRUE(metric); + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_INT8, 8); + ASSERT_EQ(0, metric->init(meta, ailego::Params())); + + std::vector> vecs{ + {1, 2, 3, 4, 5, 6, 7, 8}, + {-1, 0, 2, -3, 4, -5, 6, -7}, + {8, 7, 6, 5, 4, 3, 2, 1}, + }; + std::vector query{2, -1, 3, 1, -2, 4, -3, 5}; + + CheckBatchDistanceMatchesSingle(metric, vecs, query, 8, 1e-5f); +} + TEST(EuclideanMetric, General) { auto metric = IndexFactory::CreateMetric("Euclidean"); EXPECT_TRUE(metric); diff --git a/tests/core/metric/inner_product_metric_test.cc b/tests/core/metric/inner_product_metric_test.cc index 8a0c521c..805239d9 100644 --- a/tests/core/metric/inner_product_metric_test.cc +++ b/tests/core/metric/inner_product_metric_test.cc @@ -12,12 +12,40 @@ // See the License for the specific language governing permissions and // limitations under the License. #include +#include #include #include "zvec/core/framework/index_factory.h" using namespace zvec; using namespace zvec::core; +template +void CheckBatchDistanceMatchesSingle(const IndexMetric::Pointer &metric, + const std::vector> &vecs, + const std::vector &query, size_t dim, + float tolerance) { + auto single_distance = metric->distance(); + auto batch_distance = metric->batch_distance(); + ASSERT_TRUE(single_distance); + ASSERT_TRUE(batch_distance); + + std::vector ptrs; + ptrs.reserve(vecs.size()); + for (const auto &vec : vecs) { + ptrs.push_back(vec.data()); + } + + std::vector batch_result(vecs.size(), 0.0f); + batch_distance(ptrs.data(), query.data(), vecs.size(), dim, + batch_result.data()); + + for (size_t i = 0; i < vecs.size(); ++i) { + float single_result = 0.0f; + single_distance(vecs[i].data(), query.data(), dim, &single_result); + EXPECT_NEAR(single_result, batch_result[i], tolerance); + } +} + TEST(InnerProductMetric, General) { auto metric = IndexFactory::CreateMetric("InnerProduct"); ASSERT_TRUE(metric); @@ -74,4 +102,40 @@ TEST(InnerProductMetric, General) { float result = 1.0f; metric->normalize(&result); EXPECT_FLOAT_EQ(-1.0f, result); -} \ No newline at end of file +} + +TEST(InnerProductMetric, BatchDistanceMatchesSingleFp32) { + auto metric = IndexFactory::CreateMetric("InnerProduct"); + ASSERT_TRUE(metric); + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, 8); + ASSERT_EQ(0, metric->init(meta, ailego::Params())); + + std::vector> vecs{ + {1.0f, 2.0f, 3.0f, 4.0f, 1.5f, -2.0f, 0.5f, -1.0f}, + {0.0f, -1.0f, 1.0f, 2.0f, -0.5f, 4.0f, -2.5f, 3.0f}, + {-3.0f, 2.5f, 1.5f, 0.0f, 3.0f, -1.0f, 2.0f, 1.0f}, + }; + std::vector query{2.0f, -1.0f, 0.5f, 3.0f, 1.0f, -2.0f, 4.0f, -1.5f}; + + CheckBatchDistanceMatchesSingle(metric, vecs, query, 8, 1e-5f); +} + +TEST(InnerProductMetric, BatchDistanceMatchesSingleInt8) { + auto metric = IndexFactory::CreateMetric("InnerProduct"); + ASSERT_TRUE(metric); + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_INT8, 8); + ASSERT_EQ(0, metric->init(meta, ailego::Params())); + + std::vector> vecs{ + {1, 2, 3, 4, 5, 6, 7, 8}, + {-1, 0, 2, -3, 4, -5, 6, -7}, + {8, 7, 6, 5, 4, 3, 2, 1}, + }; + std::vector query{2, -1, 3, 1, -2, 4, -3, 5}; + + CheckBatchDistanceMatchesSingle(metric, vecs, query, 8, 1e-5f); +} From e2e857266ded8f87f9c25872da37107922858095 Mon Sep 17 00:00:00 2001 From: kgeg401 Date: Sat, 14 Feb 2026 20:54:51 -0500 Subject: [PATCH 2/3] test: stabilize fp16 euclidean matrix comparison --- .../math/euclidean_distance_matrix_fp16_test.cc | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc b/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc index 6d647c3c..0c7cc9e6 100644 --- a/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc +++ b/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include #include #include @@ -135,7 +136,8 @@ TEST(DistanceMatrix, SquaredEuclidean_General) { template void TestEuclideanMatrix(void) { - std::mt19937 gen((std::random_device())()); + std::mt19937 gen(static_cast(0x5EED1234u + M * 131u + N * 17u)); + constexpr int kFp16MatrixUlpTolerance = 20000; const size_t batch_size = M; const size_t query_size = N; @@ -174,13 +176,15 @@ void TestEuclideanMatrix(void) { for (size_t i = 0; i < batch_size * query_size; ++i) { // EXPECT_FLOAT_EQ(result1[i], result2[i]); - EXPECT_TRUE(MathHelper::IsAlmostEqual(result1[i], result2[i], 10000)); + EXPECT_TRUE(MathHelper::IsAlmostEqual( + result1[i], result2[i], kFp16MatrixUlpTolerance)); } } template void TestSquaredEuclideanMatrix(void) { - std::mt19937 gen((std::random_device())()); + std::mt19937 gen(static_cast(0x5EED5678u + M * 131u + N * 17u)); + constexpr int kFp16MatrixUlpTolerance = 20000; const size_t batch_size = M; const size_t query_size = N; @@ -219,7 +223,8 @@ void TestSquaredEuclideanMatrix(void) { for (size_t i = 0; i < batch_size * query_size; ++i) { // EXPECT_FLOAT_EQ(result1[i], result2[i]); - EXPECT_TRUE(MathHelper::IsAlmostEqual(result1[i], result2[i], 10000)); + EXPECT_TRUE(MathHelper::IsAlmostEqual( + result1[i], result2[i], kFp16MatrixUlpTolerance)); } } From 959b4b22d91b943e5de7489fd95c9ed013696a37 Mon Sep 17 00:00:00 2001 From: kgeg401 Date: Sat, 14 Feb 2026 21:36:51 -0500 Subject: [PATCH 3/3] test: fix fp16 matrix CI and clang-tidy warnings --- .../euclidean_distance_matrix_fp16_test.cc | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc b/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc index 0c7cc9e6..f2d00a9c 100644 --- a/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc +++ b/tests/ailego/math/euclidean_distance_matrix_fp16_test.cc @@ -137,7 +137,7 @@ TEST(DistanceMatrix, SquaredEuclidean_General) { template void TestEuclideanMatrix(void) { std::mt19937 gen(static_cast(0x5EED1234u + M * 131u + N * 17u)); - constexpr int kFp16MatrixUlpTolerance = 20000; + constexpr int kFp16MatrixUlpTolerance = 40000; const size_t batch_size = M; const size_t query_size = N; @@ -184,7 +184,7 @@ void TestEuclideanMatrix(void) { template void TestSquaredEuclideanMatrix(void) { std::mt19937 gen(static_cast(0x5EED5678u + M * 131u + N * 17u)); - constexpr int kFp16MatrixUlpTolerance = 20000; + constexpr int kFp16MatrixUlpTolerance = 40000; const size_t batch_size = M; const size_t query_size = N; @@ -559,7 +559,7 @@ void EuclideanBenchmark(void) { std::cout << "# (" << IntelIntrinsics() << ") FP16 " << dimension << "d, " << batch_size << " * " << query_size << " * " << block_size - << std::endl; + << '\n'; // 1 Batched Euclidean elapsed_time.reset(); @@ -575,7 +575,7 @@ void EuclideanBenchmark(void) { } } std::cout << "* 1 Batched Euclidean (us) \t" << elapsed_time.micro_seconds() - << std::endl; + << '\n'; // N Batched Euclidean elapsed_time.reset(); @@ -586,7 +586,7 @@ void EuclideanBenchmark(void) { matrix_batch, &query2[0], dimension, results); } std::cout << "* N Batched Euclidean (us) \t" << elapsed_time.micro_seconds() - << std::endl; + << '\n'; // Unbatched Euclidean elapsed_time.reset(); @@ -605,7 +605,7 @@ void EuclideanBenchmark(void) { } } std::cout << "* Unbatched Euclidean (us) \t" << elapsed_time.micro_seconds() - << std::endl; + << '\n'; } template @@ -643,7 +643,7 @@ void SquaredEuclideanBenchmark(void) { std::cout << "# (" << IntelIntrinsics() << ") FP16 " << dimension << "d, " << batch_size << " * " << query_size << " * " << block_size - << std::endl; + << '\n'; // 1 Batched Euclidean elapsed_time.reset(); @@ -659,7 +659,7 @@ void SquaredEuclideanBenchmark(void) { } } std::cout << "* 1 Batched SquaredEuclidean (us) \t" - << elapsed_time.micro_seconds() << std::endl; + << elapsed_time.micro_seconds() << '\n'; // N Batched Euclidean elapsed_time.reset(); @@ -670,7 +670,7 @@ void SquaredEuclideanBenchmark(void) { matrix_batch, &query2[0], dimension, results); } std::cout << "* N Batched SquaredEuclidean (us) \t" - << elapsed_time.micro_seconds() << std::endl; + << elapsed_time.micro_seconds() << '\n'; // Unbatched Euclidean elapsed_time.reset(); @@ -689,7 +689,7 @@ void SquaredEuclideanBenchmark(void) { } } std::cout << "* Unbatched SquaredEuclidean (us) \t" - << elapsed_time.micro_seconds() << std::endl; + << elapsed_time.micro_seconds() << '\n'; } TEST(DistanceMatrix, DISABLED_Euclidean_Benchmark) {