Skip to content

Commit 75f8321

Browse files
Neiko2002Neiko2002
committed
Add Uint8 InnerProduct metric and centralize CPU instruction set resolution
- Implement Uint8_InnerProduct distance kernels with AVX-512, AVX2, and Scalar paths - Add AVX-VNNI and AVX512-VNNI CPU feature detection with fallback resolution - Quantize uint8 test datasets directly from synthetic clustered float data - Add comprehensive unit, integration, and 100k regression tests
1 parent ef1990e commit 75f8321

27 files changed

Lines changed: 1138 additions & 200 deletions

‎cpp/deglib/include/deglib/config.h‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,16 @@
77
#define DEGLIB_X86 1
88
#endif
99

10-
// Compile methods with this attribute for AVX-512 functions on GCC/Clang
11-
#if defined(DEGLIB_X86) && (defined(__GNUC__) || defined(__clang__))
12-
#define DEGLIB_TARGET_AVX512 __attribute__((target("avx512f,avx512dq,avx512bw,avx512vpopcntdq,avx2,f16c,fma")))
13-
#else
14-
#define DEGLIB_TARGET_AVX512
15-
#endif
16-
17-
// Compile methods with this attribute for AVX2 F16C functions on GCC/Clang.
18-
// Covers AVX2, F16C, and FMA intrinsics used in fp16_ip.h and fp32_ip.h.
19-
// F16C is assumed available on all CPUs supporting AVX2.
10+
// Compile methods with this attribute for AVX-512 base functions on GCC/Clang
2011
#if defined(DEGLIB_X86) && (defined(__GNUC__) || defined(__clang__))
12+
#define DEGLIB_TARGET_AVX512 __attribute__((target("avx512f,avx512dq,avx512bw,avx512vpopcntdq")))
13+
#define DEGLIB_TARGET_AVX512_VNNI __attribute__((target("avx512f,avx512dq,avx512bw,avx512vpopcntdq,avx512vnni")))
14+
#define DEGLIB_TARGET_AVX2_VNNI __attribute__((target("avx2,f16c,fma,avxvnni")))
2115
#define DEGLIB_TARGET_AVX2 __attribute__((target("avx2,f16c,fma")))
2216
#else
17+
#define DEGLIB_TARGET_AVX512
18+
#define DEGLIB_TARGET_AVX512_VNNI
19+
#define DEGLIB_TARGET_AVX2_VNNI
2320
#define DEGLIB_TARGET_AVX2
2421
#endif
2522

‎cpp/deglib/include/deglib/distance/evp_inner_product.h‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -597,15 +597,10 @@ using DistanceVariant = std::variant<
597597
>;
598598

599599
inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::InstructionSet instruction = deglib::cpu::InstructionSet::Auto) {
600-
if (instruction == deglib::cpu::InstructionSet::Scalar) {
601-
return EvpInnerProduct{};
602-
}
600+
const auto target = deglib::cpu::resolve_instruction_set(instruction);
603601

604602
#if defined(DEGLIB_X86)
605-
if (instruction == deglib::cpu::InstructionSet::AVX512 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx512())) {
606-
if (instruction == deglib::cpu::InstructionSet::AVX512 && !deglib::cpu::has_avx512()) {
607-
throw std::runtime_error("AVX512 instruction set requested, but not supported by CPU");
608-
}
603+
if (target == deglib::cpu::InstructionSet::AVX512) {
609604
if (dim < 512) {
610605
return EvpInnerProduct_AVX512<ResidualMode::TailOnly>{};
611606
} else if (dim < 1024) {
@@ -621,10 +616,7 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
621616
else
622617
return EvpInnerProduct_AVX512<ResidualMode::Full>{};
623618
}
624-
} else if (instruction == deglib::cpu::InstructionSet::AVX2 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx2())) {
625-
if (instruction == deglib::cpu::InstructionSet::AVX2 && !deglib::cpu::has_avx2()) {
626-
throw std::runtime_error("AVX2 instruction set requested, but not supported by CPU");
627-
}
619+
} else if (target == deglib::cpu::InstructionSet::AVX2) {
628620
if (dim < 256) {
629621
return EvpInnerProduct_AVX2<ResidualMode::TailOnly>{};
630622
} else if (dim < 512) {
@@ -641,10 +633,6 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
641633
return EvpInnerProduct_AVX2<ResidualMode::Full>{};
642634
}
643635
}
644-
#else
645-
if (instruction != deglib::cpu::InstructionSet::Auto && instruction != deglib::cpu::InstructionSet::Scalar) {
646-
throw std::runtime_error("Requested SIMD instruction set is not supported on this platform");
647-
}
648636
#endif
649637
return EvpInnerProduct{};
650638
}

‎cpp/deglib/include/deglib/distance/fp16_ip.h‎

Lines changed: 10 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -45,15 +45,16 @@ class InnerProductFP16 {
4545

4646
#if defined(DEGLIB_X86)
4747
DEGLIB_TARGET_AVX2 inline static float fp16_hsum256(__m256 s) {
48-
__m128 sum128 = _mm_add_ps(_mm256_extractf128_ps(s, 0), _mm256_extractf128_ps(s, 1));
49-
alignas(32) float f[4];
50-
_mm_store_ps(f, sum128);
51-
return f[0] + f[1] + f[2] + f[3];
48+
__m128 sum128 = _mm_add_ps(_mm256_castps256_ps128(s), _mm256_extractf128_ps(s, 1));
49+
__m128 shuf = _mm_movehdup_ps(sum128);
50+
__m128 sums = _mm_add_ps(sum128, shuf);
51+
shuf = _mm_movehl_ps(shuf, sums);
52+
sums = _mm_add_ss(sums, shuf);
53+
return _mm_cvtss_f32(sums);
5254
}
5355

5456
DEGLIB_TARGET_AVX512 inline static float fp16_hsum512(__m512 s) {
55-
__m256 sum256 = _mm256_add_ps(_mm512_extractf32x8_ps(s, 0), _mm512_extractf32x8_ps(s, 1));
56-
return fp16_hsum256(sum256);
57+
return _mm512_reduce_add_ps(s);
5758
}
5859

5960
// -------------------------------------------------------------------
@@ -367,15 +368,10 @@ using DistanceVariant = std::variant<
367368
>;
368369

369370
inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::InstructionSet instruction = deglib::cpu::InstructionSet::Auto) {
370-
if (instruction == deglib::cpu::InstructionSet::Scalar) {
371-
return InnerProductFP16{};
372-
}
371+
const auto target = deglib::cpu::resolve_instruction_set(instruction);
373372

374373
#if defined(DEGLIB_X86)
375-
if (instruction == deglib::cpu::InstructionSet::AVX512 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx512())) {
376-
if (instruction == deglib::cpu::InstructionSet::AVX512 && !deglib::cpu::has_avx512()) {
377-
throw std::runtime_error("AVX512 instruction set requested, but not supported by CPU");
378-
}
374+
if (target == deglib::cpu::InstructionSet::AVX512) {
379375
if (dim < 16) {
380376
return InnerProductFP16_AVX512<ResidualMode::TailOnly>{};
381377
} else if (dim < 32) {
@@ -391,10 +387,7 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
391387
else
392388
return InnerProductFP16_AVX512<ResidualMode::Full>{};
393389
}
394-
} else if (instruction == deglib::cpu::InstructionSet::AVX2 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx2())) {
395-
if (instruction == deglib::cpu::InstructionSet::AVX2 && !deglib::cpu::has_avx2()) {
396-
throw std::runtime_error("AVX2 instruction set requested, but not supported by CPU");
397-
}
390+
} else if (target == deglib::cpu::InstructionSet::AVX2) {
398391
if (dim < 8) {
399392
return InnerProductFP16_AVX2<ResidualMode::TailOnly>{};
400393
} else if (dim < 16) {
@@ -411,10 +404,6 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
411404
return InnerProductFP16_AVX2<ResidualMode::Full>{};
412405
}
413406
}
414-
#else
415-
if (instruction != deglib::cpu::InstructionSet::Auto && instruction != deglib::cpu::InstructionSet::Scalar) {
416-
throw std::runtime_error("Requested SIMD instruction set is not supported on this platform");
417-
}
418407
#endif
419408

420409
return InnerProductFP16{};

‎cpp/deglib/include/deglib/distance/fp32_ip.h‎

Lines changed: 10 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -66,15 +66,16 @@ class InnerProductFloat {
6666
// -------------------------------------------------------------------
6767

6868
DEGLIB_TARGET_AVX2 inline static float fp32_hsum256(__m256 s) {
69-
__m128 sum128 = _mm_add_ps(_mm256_extractf128_ps(s, 0), _mm256_extractf128_ps(s, 1));
70-
alignas(32) float f[4];
71-
_mm_store_ps(f, sum128);
72-
return f[0] + f[1] + f[2] + f[3];
69+
__m128 sum128 = _mm_add_ps(_mm256_castps256_ps128(s), _mm256_extractf128_ps(s, 1));
70+
__m128 shuf = _mm_movehdup_ps(sum128);
71+
__m128 sums = _mm_add_ps(sum128, shuf);
72+
shuf = _mm_movehl_ps(shuf, sums);
73+
sums = _mm_add_ss(sums, shuf);
74+
return _mm_cvtss_f32(sums);
7375
}
7476

7577
DEGLIB_TARGET_AVX512 inline static float fp32_hsum512(__m512 s) {
76-
__m256 sum256 = _mm256_add_ps(_mm512_extractf32x8_ps(s, 0), _mm512_extractf32x8_ps(s, 1));
77-
return fp32_hsum256(sum256);
78+
return _mm512_reduce_add_ps(s);
7879
}
7980

8081
template <ResidualMode Mode = ResidualMode::Full>
@@ -344,15 +345,10 @@ using DistanceVariant = std::variant<
344345
>;
345346

346347
inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::InstructionSet instruction = deglib::cpu::InstructionSet::Auto) {
347-
if (instruction == deglib::cpu::InstructionSet::Scalar) {
348-
return InnerProductFloat{};
349-
}
348+
const auto target = deglib::cpu::resolve_instruction_set(instruction);
350349

351350
#if defined(DEGLIB_X86)
352-
if (instruction == deglib::cpu::InstructionSet::AVX512 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx512())) {
353-
if (instruction == deglib::cpu::InstructionSet::AVX512 && !deglib::cpu::has_avx512()) {
354-
throw std::runtime_error("AVX512 instruction set requested, but not supported by CPU");
355-
}
351+
if (target == deglib::cpu::InstructionSet::AVX512) {
356352
if (dim < 16) {
357353
return InnerProductFloat_AVX512<ResidualMode::TailOnly>{};
358354
} else if (dim < 32) {
@@ -368,10 +364,7 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
368364
else
369365
return InnerProductFloat_AVX512<ResidualMode::Full>{};
370366
}
371-
} else if (instruction == deglib::cpu::InstructionSet::AVX2 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx2())) {
372-
if (instruction == deglib::cpu::InstructionSet::AVX2 && !deglib::cpu::has_avx2()) {
373-
throw std::runtime_error("AVX2 instruction set requested, but not supported by CPU");
374-
}
367+
} else if (target == deglib::cpu::InstructionSet::AVX2) {
375368
if (dim < 8) {
376369
return InnerProductFloat_AVX2<ResidualMode::TailOnly>{};
377370
} else if (dim < 16) {
@@ -388,10 +381,6 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
388381
return InnerProductFloat_AVX2<ResidualMode::Full>{};
389382
}
390383
}
391-
#else
392-
if (instruction != deglib::cpu::InstructionSet::Auto && instruction != deglib::cpu::InstructionSet::Scalar) {
393-
throw std::runtime_error("Requested SIMD instruction set is not supported on this platform");
394-
}
395384
#endif
396385

397386
return InnerProductFloat{};

‎cpp/deglib/include/deglib/distance/fp32_l2.h‎

Lines changed: 10 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -64,15 +64,16 @@ class L2Float {
6464
// -------------------------------------------------------------------
6565

6666
DEGLIB_TARGET_AVX2 inline static float fp32_l2_hsum256(__m256 s) {
67-
__m128 sum128 = _mm_add_ps(_mm256_extractf128_ps(s, 0), _mm256_extractf128_ps(s, 1));
68-
alignas(32) float f[4];
69-
_mm_store_ps(f, sum128);
70-
return f[0] + f[1] + f[2] + f[3];
67+
__m128 sum128 = _mm_add_ps(_mm256_castps256_ps128(s), _mm256_extractf128_ps(s, 1));
68+
__m128 shuf = _mm_movehdup_ps(sum128);
69+
__m128 sums = _mm_add_ps(sum128, shuf);
70+
shuf = _mm_movehl_ps(shuf, sums);
71+
sums = _mm_add_ss(sums, shuf);
72+
return _mm_cvtss_f32(sums);
7173
}
7274

7375
DEGLIB_TARGET_AVX512 inline static float fp32_l2_hsum512(__m512 s) {
74-
__m256 sum256 = _mm256_add_ps(_mm512_extractf32x8_ps(s, 0), _mm512_extractf32x8_ps(s, 1));
75-
return fp32_l2_hsum256(sum256);
76+
return _mm512_reduce_add_ps(s);
7677
}
7778
template <ResidualMode Mode = ResidualMode::Full>
7879
class L2Float_AVX512 {
@@ -344,15 +345,10 @@ using DistanceVariant = std::variant<
344345
>;
345346

346347
inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::InstructionSet instruction = deglib::cpu::InstructionSet::Auto) {
347-
if (instruction == deglib::cpu::InstructionSet::Scalar) {
348-
return L2Float{};
349-
}
348+
const auto target = deglib::cpu::resolve_instruction_set(instruction);
350349

351350
#if defined(DEGLIB_X86)
352-
if (instruction == deglib::cpu::InstructionSet::AVX512 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx512())) {
353-
if (instruction == deglib::cpu::InstructionSet::AVX512 && !deglib::cpu::has_avx512()) {
354-
throw std::runtime_error("AVX512 instruction set requested, but not supported by CPU");
355-
}
351+
if (target == deglib::cpu::InstructionSet::AVX512) {
356352
if (dim < 16) {
357353
return L2Float_AVX512<ResidualMode::TailOnly>{};
358354
} else if (dim < 32) {
@@ -368,10 +364,7 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
368364
else
369365
return L2Float_AVX512<ResidualMode::Full>{};
370366
}
371-
} else if (instruction == deglib::cpu::InstructionSet::AVX2 || (instruction == deglib::cpu::InstructionSet::Auto && deglib::cpu::has_avx2())) {
372-
if (instruction == deglib::cpu::InstructionSet::AVX2 && !deglib::cpu::has_avx2()) {
373-
throw std::runtime_error("AVX2 instruction set requested, but not supported by CPU");
374-
}
367+
} else if (target == deglib::cpu::InstructionSet::AVX2) {
375368
if (dim < 8) {
376369
return L2Float_AVX2<ResidualMode::TailOnly>{};
377370
} else if (dim < 16) {
@@ -388,10 +381,6 @@ inline DistanceVariant select_dist(const size_t dim, const deglib::cpu::Instruct
388381
return L2Float_AVX2<ResidualMode::Full>{};
389382
}
390383
}
391-
#else
392-
if (instruction != deglib::cpu::InstructionSet::Auto && instruction != deglib::cpu::InstructionSet::Scalar) {
393-
throw std::runtime_error("Requested SIMD instruction set is not supported on this platform");
394-
}
395384
#endif
396385

397386
return L2Float{};

0 commit comments

Comments
 (0)