facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Implement NEON-based FINE_SIZE=2 specializations for Index2LevelDecoderImpl and IndexPQDecoder (#5255)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5255

Implement NEON-based FINE_SIZE=2 specializations for Index2LevelDecoderImpl and IndexPQDecoder, previously being compiled from AVX2 code through aarch64 compat layers.

Remove usage of aarch64 compatibility layers, as they are not needed anymore.

Reviewed By: YifanYuan3

Differential Revision: D106693266

fbshipit-source-id: 21fa2c63c20ff7ffd25fc88b877e6331bcc649af
This commit is contained in:
Nicolas De Carli
2026-05-29 10:47:36 -07:00
committed by meta-codesync[bot]
parent f5217d7496
commit aa332cd5f0
3 changed files with 1583 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/cppcontrib/SaDecodeKernels.h
+1 -1
View File
@@ -313,7 +313,7 @@
#ifdef __AVX2__
#include <faiss/cppcontrib/sa_decode/Level2-avx2-inl.h>
#include <faiss/cppcontrib/sa_decode/PQ-avx2-inl.h>
#elif defined(__ARM_NEON)
#elif defined(__ARM_NEON) || defined(__aarch64__)
#include <faiss/cppcontrib/sa_decode/Level2-neon-inl.h>
#include <faiss/cppcontrib/sa_decode/PQ-neon-inl.h>
#else
faiss/cppcontrib/sa_decode/Level2-neon-inl.h
+890
View File
@@ -142,6 +142,57 @@ inline float32x4x2_t elementaryBlock8x1bAccum(
return {result0, result1};
}
// Processes 8 float values.
// Returns {
// val[0] = {*coarse[0..1] + *fine0[0..1], *coarse[2..3] + *fine1[0..1]};
// val[1] = {*coarse[4..5] + *fine2[0..1], *coarse[6..7] + *fine3[0..1]};
// }
inline float32x4x2_t elementaryBlock2x4b(
const float* const __restrict coarse,
const float* const __restrict fine0,
const float* const __restrict fine1,
const float* const __restrict fine2,
const float* const __restrict fine3) {
const auto fine0Value = vld1_f32(fine0);
const auto fine1Value = vld1_f32(fine1);
const auto fine2Value = vld1_f32(fine2);
const auto fine3Value = vld1_f32(fine3);
const auto coarseValue0 = vld1q_f32(coarse);
const auto coarseValue1 = vld1q_f32(coarse + 4);
const auto fineResult0 = vcombine_f32(fine0Value, fine1Value);
const auto fineResult1 = vcombine_f32(fine2Value, fine3Value);
return {vaddq_f32(fineResult0, coarseValue0),
vaddq_f32(fineResult1, coarseValue1)};
}
// Processes 8 float values.
// Returns {
// val[0] = existingValue.val[0] + weight * (*coarse[0..1] + *fine0[0..1],
// *coarse[2..3] + *fine1[0..1]); val[1] = existingValue.val[1] + weight *
// (*coarse[4..5] + *fine2[0..1], *coarse[6..7] + *fine3[0..1]);
// }
inline float32x4x2_t elementaryBlock2x4bAccum(
const float* const __restrict coarse,
const float* const __restrict fine0,
const float* const __restrict fine1,
const float* const __restrict fine2,
const float* const __restrict fine3,
const float weight,
const float32x4x2_t existingValue) {
const auto combinedValue =
elementaryBlock2x4b(coarse, fine0, fine1, fine2, fine3);
const auto weightNeon = vdupq_n_f32(weight);
const auto result0 =
vfmaq_f32(existingValue.val[0], weightNeon, combinedValue.val[0]);
const auto result1 =
vfmaq_f32(existingValue.val[1], weightNeon, combinedValue.val[1]);
return {result0, result1};
}
// The following code uses template-based for-loop unrolling,
// because the compiler does not do that on its own as needed.
// The idea is the following:
@@ -168,12 +219,846 @@ template <
intptr_t COARSE_BITS,
intptr_t FINE_BITS,
intptr_t CPOS,
bool FINE_SIZE_EQ_2 = FINE_SIZE == 2,
bool FINE_SIZE_EQ_4 = FINE_SIZE == 4,
bool QPOS_LEFT_GE_8 = (FINE_SIZE - CPOS % FINE_SIZE >= 8),
bool QPOS_LEFT_GE_4 = (FINE_SIZE - CPOS % FINE_SIZE >= 4),
bool DIM_EQ_CPOS = DIM == CPOS>
struct Index2LevelDecoderImpl;
template <
intptr_t DIM,
intptr_t COARSE_SIZE,
intptr_t COARSE_BITS,
intptr_t FINE_BITS,
intptr_t CPOS,
bool QPOS_LEFT_GE_8,
bool QPOS_LEFT_GE_4>
struct Index2LevelDecoderImpl<
DIM,
COARSE_SIZE,
2,
COARSE_BITS,
FINE_BITS,
CPOS,
true,
false,
QPOS_LEFT_GE_8,
QPOS_LEFT_GE_4,
false> {
static constexpr intptr_t FINE_SIZE = 2;
static constexpr intptr_t coarseCentroidIdx = CPOS / COARSE_SIZE;
static constexpr intptr_t coarseCentroidOffset = CPOS % COARSE_SIZE;
static constexpr intptr_t fineCentroidIdx = CPOS / FINE_SIZE;
static constexpr intptr_t fineCentroidOffset = CPOS % FINE_SIZE;
static constexpr intptr_t QPOS_LEFT = FINE_SIZE - fineCentroidOffset;
// coarse quantizer storage
static constexpr intptr_t COARSE_TABLE_BYTES = (1 << COARSE_BITS);
// coarse quantizer bytes start from 0
// fine quantizer bytes start from N_COARSE_ELEMENTS_BYTES
static constexpr intptr_t N_COARSE_ELEMENTS = DIM / COARSE_SIZE;
static constexpr intptr_t N_COARSE_ELEMENTS_BITS =
N_COARSE_ELEMENTS * COARSE_BITS;
static constexpr intptr_t N_COARSE_ELEMENTS_BYTES =
(N_COARSE_ELEMENTS_BITS + 7) / 8;
static constexpr intptr_t FINE_TABLE_BYTES = (1 << FINE_BITS);
// process 1 sample
static void store(
const float* const __restrict pqCoarseCentroids0,
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
float* const __restrict outputStore) {
// coarse quantizer
const uint8_t* const __restrict coarse0 = code0;
// fine quantizer
const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t coarseCode0 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse0);
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const auto storeValue = elementaryBlock2x4b(
pqCoarseCentroids0 +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset);
vst1q_f32(outputStore + CPOS, storeValue.val[0]);
vst1q_f32(outputStore + CPOS + 4, storeValue.val[1]);
// next
Index2LevelDecoderImpl<
DIM,
COARSE_SIZE,
FINE_SIZE,
COARSE_BITS,
FINE_BITS,
CPOS + 8>::
store(pqCoarseCentroids0, pqFineCentroids0, code0, outputStore);
}
// process 1 sample
static void accum(
const float* const __restrict pqCoarseCentroids0,
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
const float weight0,
float* const __restrict outputAccum) {
// coarse quantizer
const uint8_t* const __restrict coarse0 = code0;
// fine quantizer
const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t coarseCode0 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse0);
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids0 +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
Index2LevelDecoderImpl<
DIM,
COARSE_SIZE,
FINE_SIZE,
COARSE_BITS,
FINE_BITS,
CPOS + 8>::
accum(pqCoarseCentroids0,
pqFineCentroids0,
code0,
weight0,
outputAccum);
}
// Process 2 samples.
// Each code uses its own coarse pq centroids table and fine pq centroids
// table.
static void accum(
const float* const __restrict pqCoarseCentroids0,
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
const float weight0,
const float* const __restrict pqCoarseCentroids1,
const float* const __restrict pqFineCentroids1,
const uint8_t* const __restrict code1,
const float weight1,
float* const __restrict outputAccum) {
// coarse quantizer
const uint8_t* const __restrict coarse0 = code0;
const uint8_t* const __restrict coarse1 = code1;
// fine quantizer
const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES;
const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t coarseCode0 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse0);
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t coarseCode1 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse1);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids0 +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids1 +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
Index2LevelDecoderImpl<
DIM,
COARSE_SIZE,
FINE_SIZE,
COARSE_BITS,
FINE_BITS,
CPOS + 8>::
accum(pqCoarseCentroids0,
pqFineCentroids0,
code0,
weight0,
pqCoarseCentroids1,
pqFineCentroids1,
code1,
weight1,
outputAccum);
}
// Process 2 samples.
// Coarse pq centroids table and fine pq centroids table are shared among
// codes.
static void accum(
const float* const __restrict pqCoarseCentroids,
const float* const __restrict pqFineCentroids,
const uint8_t* const __restrict code0,
const float weight0,
const uint8_t* const __restrict code1,
const float weight1,
float* const __restrict outputAccum) {
// coarse quantizer
const uint8_t* const __restrict coarse0 = code0;
const uint8_t* const __restrict coarse1 = code1;
// fine quantizer
const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES;
const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t coarseCode0 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse0);
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t coarseCode1 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse1);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
Index2LevelDecoderImpl<
DIM,
COARSE_SIZE,
FINE_SIZE,
COARSE_BITS,
FINE_BITS,
CPOS + 8>::
accum(pqCoarseCentroids,
pqFineCentroids,
code0,
weight0,
code1,
weight1,
outputAccum);
}
// Process 3 samples.
// Each code uses its own coarse pq centroids table and fine pq centroids
// table.
static void accum(
const float* const __restrict pqCoarseCentroids0,
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
const float weight0,
const float* const __restrict pqCoarseCentroids1,
const float* const __restrict pqFineCentroids1,
const uint8_t* const __restrict code1,
const float weight1,
const float* const __restrict pqCoarseCentroids2,
const float* const __restrict pqFineCentroids2,
const uint8_t* const __restrict code2,
const float weight2,
float* const __restrict outputAccum) {
// coarse quantizer
const uint8_t* const __restrict coarse0 = code0;
const uint8_t* const __restrict coarse1 = code1;
const uint8_t* const __restrict coarse2 = code2;
// fine quantizer
const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES;
const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES;
const uint8_t* const __restrict fine2 = code2 + N_COARSE_ELEMENTS_BYTES;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t coarseCode0 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse0);
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t coarseCode1 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse1);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
const intptr_t coarseCode2 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse2);
const intptr_t fineCode2a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine2);
const intptr_t fineCode2b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine2);
const intptr_t fineCode2c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine2);
const intptr_t fineCode2d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine2);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids0 +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids1 +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids2 +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode2) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids2 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode2a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids2 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode2b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids2 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode2c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids2 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode2d) *
FINE_SIZE +
fineCentroidOffset,
weight2,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
Index2LevelDecoderImpl<
DIM,
COARSE_SIZE,
FINE_SIZE,
COARSE_BITS,
FINE_BITS,
CPOS + 8>::
accum(pqCoarseCentroids0,
pqFineCentroids0,
code0,
weight0,
pqCoarseCentroids1,
pqFineCentroids1,
code1,
weight1,
pqCoarseCentroids2,
pqFineCentroids2,
code2,
weight2,
outputAccum);
}
// Process 3 samples.
// Coarse pq centroids table and fine pq centroids table are shared among
// codes.
static void accum(
const float* const __restrict pqCoarseCentroids,
const float* const __restrict pqFineCentroids,
const uint8_t* const __restrict code0,
const float weight0,
const uint8_t* const __restrict code1,
const float weight1,
const uint8_t* const __restrict code2,
const float weight2,
float* const __restrict outputAccum) {
// coarse quantizer
const uint8_t* const __restrict coarse0 = code0;
const uint8_t* const __restrict coarse1 = code1;
const uint8_t* const __restrict coarse2 = code2;
// fine quantizer
const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES;
const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES;
const uint8_t* const __restrict fine2 = code2 + N_COARSE_ELEMENTS_BYTES;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t coarseCode0 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse0);
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t coarseCode1 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse1);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
const intptr_t coarseCode2 = detail::
UintReader<DIM, COARSE_SIZE, COARSE_BITS, coarseCentroidIdx>::
get(coarse2);
const intptr_t fineCode2a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine2);
const intptr_t fineCode2b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine2);
const intptr_t fineCode2c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine2);
const intptr_t fineCode2d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine2);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
existingValue = elementaryBlock2x4bAccum(
pqCoarseCentroids +
(coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode2) *
COARSE_SIZE +
coarseCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode2a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode2b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode2c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode2d) *
FINE_SIZE +
fineCentroidOffset,
weight2,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
Index2LevelDecoderImpl<
DIM,
COARSE_SIZE,
FINE_SIZE,
COARSE_BITS,
FINE_BITS,
CPOS + 8>::
accum(pqCoarseCentroids,
pqFineCentroids,
code0,
weight0,
code1,
weight1,
code2,
weight2,
outputAccum);
}
};
template <
intptr_t DIM,
intptr_t COARSE_SIZE,
@@ -189,6 +1074,7 @@ struct Index2LevelDecoderImpl<
COARSE_BITS,
FINE_BITS,
CPOS,
false,
true,
QPOS_LEFT_GE_8,
QPOS_LEFT_GE_4,
@@ -829,6 +1715,7 @@ struct Index2LevelDecoderImpl<
FINE_BITS,
CPOS,
false,
false,
true,
true,
false> {
@@ -1353,6 +2240,7 @@ struct Index2LevelDecoderImpl<
CPOS,
false,
false,
false,
true,
false> {
static constexpr intptr_t coarseCentroidIdx = CPOS / COARSE_SIZE;
@@ -1856,6 +2744,7 @@ template <
intptr_t FINE_SIZE,
intptr_t COARSE_BITS,
intptr_t FINE_BITS,
bool FINE_SIZE_EQ_2,
bool FINE_SIZE_EQ_4,
bool QPOS_LEFT_GE_8,
bool QPOS_LEFT_GE_4>
@@ -1866,6 +2755,7 @@ struct Index2LevelDecoderImpl<
COARSE_BITS,
FINE_BITS,
DIM,
FINE_SIZE_EQ_2,
FINE_SIZE_EQ_4,
QPOS_LEFT_GE_8,
QPOS_LEFT_GE_4,
faiss/cppcontrib/sa_decode/PQ-neon-inl.h
+692
View File
@@ -122,6 +122,49 @@ inline float32x4x2_t elementaryBlock8x1bAccum(
return {result0, result1};
}
// Processes 8 float values.
// Returns {
// val[0] = {*fine0[0..1], *fine1[0..1]};
// val[1] = {*fine2[0..1], *fine3[0..1]};
// }
inline float32x4x2_t elementaryBlock2x4b(
const float* const __restrict fine0,
const float* const __restrict fine1,
const float* const __restrict fine2,
const float* const __restrict fine3) {
const auto fine0Value = vld1_f32(fine0);
const auto fine1Value = vld1_f32(fine1);
const auto fine2Value = vld1_f32(fine2);
const auto fine3Value = vld1_f32(fine3);
const auto result0 = vcombine_f32(fine0Value, fine1Value);
const auto result1 = vcombine_f32(fine2Value, fine3Value);
return {result0, result1};
}
// Processes 8 float values.
// Returns {
// val[0] = existingValue.val[0] + weight * {*fine0[0..1], *fine1[0..1]};
// val[1] = existingValue.val[1] + weight * {*fine2[0..1], *fine3[0..1]};
// }
inline float32x4x2_t elementaryBlock2x4bAccum(
const float* const __restrict fine0,
const float* const __restrict fine1,
const float* const __restrict fine2,
const float* const __restrict fine3,
const float weight,
const float32x4x2_t existingValue) {
const auto fineValue = elementaryBlock2x4b(fine0, fine1, fine2, fine3);
const auto weightNeon = vdupq_n_f32(weight);
const auto result0 =
vfmaq_f32(existingValue.val[0], weightNeon, fineValue.val[0]);
const auto result1 =
vfmaq_f32(existingValue.val[1], weightNeon, fineValue.val[1]);
return {result0, result1};
}
// The following code uses template-based for-loop unrolling,
// because the compiler does not do that on its own as needed.
// The idea is the following:
@@ -146,12 +189,656 @@ template <
intptr_t FINE_SIZE,
intptr_t FINE_BITS,
intptr_t CPOS,
bool FINE_SIZE_EQ_2 = FINE_SIZE == 2,
bool FINE_SIZE_EQ_4 = FINE_SIZE == 4,
bool QPOS_LEFT_GE_8 = (FINE_SIZE - CPOS % FINE_SIZE >= 8),
bool QPOS_LEFT_GE_4 = (FINE_SIZE - CPOS % FINE_SIZE >= 4),
bool DIM_EQ_CPOS = DIM == CPOS>
struct IndexPQDecoderImpl;
template <
intptr_t DIM,
intptr_t CPOS,
intptr_t FINE_BITS,
bool QPOS_LEFT_GE_8,
bool QPOS_LEFT_GE_4>
struct IndexPQDecoderImpl<
DIM,
2,
FINE_BITS,
CPOS,
true,
false,
QPOS_LEFT_GE_8,
QPOS_LEFT_GE_4,
false> {
static constexpr intptr_t FINE_SIZE = 2;
static constexpr intptr_t fineCentroidIdx = CPOS / FINE_SIZE;
static constexpr intptr_t fineCentroidOffset = CPOS % FINE_SIZE;
static constexpr intptr_t QPOS_LEFT = FINE_SIZE - fineCentroidOffset;
static constexpr intptr_t FINE_TABLE_BYTES = (1 << FINE_BITS);
// process 1 sample
static void store(
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
float* const __restrict outputStore) {
// fine quantizer
const uint8_t* const __restrict fine0 = code0;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const auto storeValue = elementaryBlock2x4b(
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset);
vst1q_f32(outputStore + CPOS, storeValue.val[0]);
vst1q_f32(outputStore + CPOS + 4, storeValue.val[1]);
// next
IndexPQDecoderImpl<DIM, FINE_SIZE, FINE_BITS, CPOS + 8>::store(
pqFineCentroids0, code0, outputStore);
}
// process 1 sample
static void accum(
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
const float weight0,
float* const __restrict outputAccum) {
// fine quantizer
const uint8_t* const __restrict fine0 = code0;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
IndexPQDecoderImpl<DIM, FINE_SIZE, FINE_BITS, CPOS + 8>::accum(
pqFineCentroids0, code0, weight0, outputAccum);
}
// Process 2 samples.
// Each code uses its own fine pq centroids table.
static void accum(
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
const float weight0,
const float* const __restrict pqFineCentroids1,
const uint8_t* const __restrict code1,
const float weight1,
float* const __restrict outputAccum) {
// fine quantizer
const uint8_t* const __restrict fine0 = code0;
const uint8_t* const __restrict fine1 = code1;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqFineCentroids1 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
IndexPQDecoderImpl<DIM, FINE_SIZE, FINE_BITS, CPOS + 8>::accum(
pqFineCentroids0,
code0,
weight0,
pqFineCentroids1,
code1,
weight1,
outputAccum);
}
// Process 2 samples.
// Fine pq centroids table is shared among codes.
static void accum(
const float* const __restrict pqFineCentroids,
const uint8_t* const __restrict code0,
const float weight0,
const uint8_t* const __restrict code1,
const float weight1,
float* const __restrict outputAccum) {
// fine quantizer
const uint8_t* const __restrict fine0 = code0;
const uint8_t* const __restrict fine1 = code1;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
IndexPQDecoderImpl<DIM, FINE_SIZE, FINE_BITS, CPOS + 8>::accum(
pqFineCentroids, code0, weight0, code1, weight1, outputAccum);
}
// Process 3 samples.
// Each code uses its own fine pq centroids table.
static void accum(
const float* const __restrict pqFineCentroids0,
const uint8_t* const __restrict code0,
const float weight0,
const float* const __restrict pqFineCentroids1,
const uint8_t* const __restrict code1,
const float weight1,
const float* const __restrict pqFineCentroids2,
const uint8_t* const __restrict code2,
const float weight2,
float* const __restrict outputAccum) {
// fine quantizer
const uint8_t* const __restrict fine0 = code0;
const uint8_t* const __restrict fine1 = code1;
const uint8_t* const __restrict fine2 = code2;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
const intptr_t fineCode2a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine2);
const intptr_t fineCode2b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine2);
const intptr_t fineCode2c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine2);
const intptr_t fineCode2d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine2);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqFineCentroids0 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids0 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqFineCentroids1 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids1 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
existingValue = elementaryBlock2x4bAccum(
pqFineCentroids2 +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode2a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids2 +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode2b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids2 +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode2c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids2 +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode2d) *
FINE_SIZE +
fineCentroidOffset,
weight2,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
IndexPQDecoderImpl<DIM, FINE_SIZE, FINE_BITS, CPOS + 8>::accum(
pqFineCentroids0,
code0,
weight0,
pqFineCentroids1,
code1,
weight1,
pqFineCentroids2,
code2,
weight2,
outputAccum);
}
// Process 3 samples.
// Fine pq centroids table is shared among codes.
static void accum(
const float* const __restrict pqFineCentroids,
const uint8_t* const __restrict code0,
const float weight0,
const uint8_t* const __restrict code1,
const float weight1,
const uint8_t* const __restrict code2,
const float weight2,
float* const __restrict outputAccum) {
// fine quantizer
const uint8_t* const __restrict fine0 = code0;
const uint8_t* const __restrict fine1 = code1;
const uint8_t* const __restrict fine2 = code2;
// process chunks, 2 float
// but 8 floats per loop
const intptr_t fineCode0a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine0);
const intptr_t fineCode0b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine0);
const intptr_t fineCode0c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine0);
const intptr_t fineCode0d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine0);
const intptr_t fineCode1a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine1);
const intptr_t fineCode1b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine1);
const intptr_t fineCode1c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine1);
const intptr_t fineCode1d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine1);
const intptr_t fineCode2a = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 0>::get(
fine2);
const intptr_t fineCode2b = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 1>::get(
fine2);
const intptr_t fineCode2c = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 2>::get(
fine2);
const intptr_t fineCode2d = detail::
UintReader<DIM, FINE_SIZE, FINE_BITS, fineCentroidIdx + 3>::get(
fine2);
auto existingValue0 = vld1q_f32(outputAccum + CPOS);
auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4);
auto existingValue = elementaryBlock2x4bAccum(
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode0a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode0b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode0c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode0d) *
FINE_SIZE +
fineCentroidOffset,
weight0,
{existingValue0, existingValue1});
existingValue = elementaryBlock2x4bAccum(
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode1a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode1b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode1c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode1d) *
FINE_SIZE +
fineCentroidOffset,
weight1,
existingValue);
existingValue = elementaryBlock2x4bAccum(
pqFineCentroids +
((fineCentroidIdx + 0) * FINE_TABLE_BYTES +
fineCode2a) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 1) * FINE_TABLE_BYTES +
fineCode2b) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 2) * FINE_TABLE_BYTES +
fineCode2c) *
FINE_SIZE +
fineCentroidOffset,
pqFineCentroids +
((fineCentroidIdx + 3) * FINE_TABLE_BYTES +
fineCode2d) *
FINE_SIZE +
fineCentroidOffset,
weight2,
existingValue);
vst1q_f32(outputAccum + CPOS, existingValue.val[0]);
vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]);
// next
IndexPQDecoderImpl<DIM, FINE_SIZE, FINE_BITS, CPOS + 8>::accum(
pqFineCentroids,
code0,
weight0,
code1,
weight1,
code2,
weight2,
outputAccum);
}
};
template <
intptr_t DIM,
intptr_t CPOS,
@@ -163,6 +850,7 @@ struct IndexPQDecoderImpl<
4,
FINE_BITS,
CPOS,
false,
true,
QPOS_LEFT_GE_8,
QPOS_LEFT_GE_4,
@@ -609,6 +1297,7 @@ struct IndexPQDecoderImpl<
FINE_BITS,
CPOS,
false,
false,
true,
true,
false> {
@@ -938,6 +1627,7 @@ struct IndexPQDecoderImpl<
CPOS,
false,
false,
false,
true,
false> {
static constexpr intptr_t fineCentroidIdx = CPOS / FINE_SIZE;
@@ -1252,6 +1942,7 @@ template <
intptr_t DIM,
intptr_t FINE_SIZE,
intptr_t FINE_BITS,
bool FINE_SIZE_EQ_2,
bool FINE_SIZE_EQ_4,
bool QPOS_LEFT_GE_8,
bool QPOS_LEFT_GE_4>
@@ -1260,6 +1951,7 @@ struct IndexPQDecoderImpl<
FINE_SIZE,
FINE_BITS,
DIM,
FINE_SIZE_EQ_2,
FINE_SIZE_EQ_4,
QPOS_LEFT_GE_8,
QPOS_LEFT_GE_4,