Skip to content

feat: enable one2many batch distance for ip and l2 #113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kgeg401 wants to merge 3 commits into
base: main
Choose a base branch
from
+157 −24
Open

feat: enable one2many batch distance for ip and l2 #113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9adf700
feat: enable one2many batch distance for ip and l2
kgeg401 Feb 14, 2026
e2e8572
test: stabilize fp16 euclidean matrix comparison
kgeg401 Feb 15, 2026
959b4b2
test: fix fp16 matrix CI and clang-tidy warnings
kgeg401 Feb 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions src/core/metric/euclidean_metric.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -868,22 +868,22 @@ class SquaredEuclideanMetric : public IndexMetric {
case IndexMeta::DataType::DT_FP16:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix,
ailego::Float16, 1, 1>::ComputeBatch);
ailego::Float16, 12, 2>::ComputeBatch);

case IndexMeta::DataType::DT_FP32:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix, float,
1, 1>::ComputeBatch);
12, 2>::ComputeBatch);

case IndexMeta::DataType::DT_INT8:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix, int8_t,
1, 1>::ComputeBatch);
12, 2>::ComputeBatch);

case IndexMeta::DataType::DT_INT4:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix,
uint8_t, 1, 1>::ComputeBatch);
uint8_t, 12, 2>::ComputeBatch);

default:
return nullptr;
Expand Down
14 changes: 7 additions & 7 deletions src/core/metric/inner_product_metric.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,20 +354,20 @@ class InnerProductMetric : public IndexMetric {
switch (data_type_) {
case IndexMeta::DataType::DT_FP32:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::MinusInnerProductMatrix, float, 1,
1>::ComputeBatch);
ailego::BaseDistance<ailego::MinusInnerProductMatrix, float, 12,
2>::ComputeBatch);
case IndexMeta::DataType::DT_FP16:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::MinusInnerProductMatrix,
ailego::Float16, 1, 1>::ComputeBatch);
ailego::Float16, 12, 2>::ComputeBatch);
case IndexMeta::DataType::DT_INT8:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::MinusInnerProductMatrix, int8_t, 1,
1>::ComputeBatch);
ailego::BaseDistance<ailego::MinusInnerProductMatrix, int8_t, 12,
2>::ComputeBatch);
case IndexMeta::DataType::DT_INT4:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::MinusInnerProductMatrix, uint8_t, 1,
1>::ComputeBatch);
ailego::BaseDistance<ailego::MinusInnerProductMatrix, uint8_t, 12,
2>::ComputeBatch);
default:
return nullptr;
}
Expand Down
29 changes: 17 additions & 12 deletions tests/ailego/math/euclidean_distance_matrix_fp16_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
// limitations under the License.

#include <functional>
#include <cstdint>
#include <random>
#include <string>
#include <thread>
Expand Down Expand Up @@ -135,7 +136,8 @@ TEST(DistanceMatrix, SquaredEuclidean_General) {

template <size_t M, size_t N>
void TestEuclideanMatrix(void) {
std::mt19937 gen((std::random_device())());
std::mt19937 gen(static_cast<uint32_t>(0x5EED1234u + M * 131u + N * 17u));
constexpr int kFp16MatrixUlpTolerance = 40000;

const size_t batch_size = M;
const size_t query_size = N;
Expand Down Expand Up @@ -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 <size_t M, size_t N>
void TestSquaredEuclideanMatrix(void) {
std::mt19937 gen((std::random_device())());
std::mt19937 gen(static_cast<uint32_t>(0x5EED5678u + M * 131u + N * 17u));
constexpr int kFp16MatrixUlpTolerance = 40000;

const size_t batch_size = M;
const size_t query_size = N;
Expand Down Expand Up @@ -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));
}
}

Expand Down Expand Up @@ -554,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();
Expand All @@ -570,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();
Expand All @@ -581,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();
Expand All @@ -600,7 +605,7 @@ void EuclideanBenchmark(void) {
}
}
std::cout << "* Unbatched Euclidean (us) \t" << elapsed_time.micro_seconds()
<< std::endl;
<< '\n';
}

template <size_t M, size_t N, size_t B, size_t D>
Expand Down Expand Up @@ -638,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();
Expand All @@ -654,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();
Expand All @@ -665,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();
Expand All @@ -684,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) {
Expand Down
64 changes: 64 additions & 0 deletions tests/core/metric/euclidean_metric_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,40 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <iostream>
#include <vector>
#include <gtest/gtest.h>
#include "zvec/core/framework/index_factory.h"

using namespace zvec;
using namespace zvec::core;

template <typename T>
void CheckBatchDistanceMatchesSingle(const IndexMetric::Pointer &metric,
const std::vector<std::vector<T>> &vecs,
const std::vector<T> &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<const void *> ptrs;
ptrs.reserve(vecs.size());
for (const auto &vec : vecs) {
ptrs.push_back(vec.data());
}

std::vector<float> 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);
Expand Down Expand Up @@ -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<std::vector<float>> 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<float> 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<std::vector<int8_t>> 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<int8_t> 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);
Expand Down
66 changes: 65 additions & 1 deletion tests/core/metric/inner_product_metric_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,40 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <iostream>
#include <vector>
#include <gtest/gtest.h>
#include "zvec/core/framework/index_factory.h"

using namespace zvec;
using namespace zvec::core;

template <typename T>
void CheckBatchDistanceMatchesSingle(const IndexMetric::Pointer &metric,
const std::vector<std::vector<T>> &vecs,
const std::vector<T> &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<const void *> ptrs;
ptrs.reserve(vecs.size());
for (const auto &vec : vecs) {
ptrs.push_back(vec.data());
}

std::vector<float> 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);
Expand Down Expand Up @@ -74,4 +102,40 @@ TEST(InnerProductMetric, General) {
float result = 1.0f;
metric->normalize(&result);
EXPECT_FLOAT_EQ(-1.0f, result);
}
}

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<std::vector<float>> 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<float> 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<std::vector<int8_t>> 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<int8_t> query{2, -1, 3, 1, -2, 4, -3, 5};

CheckBatchDistanceMatchesSingle(metric, vecs, query, 8, 1e-5f);
}